Numpyro: Sequential Sampling Strategy

[ross-h1]

Hi all,

I have a sequential problem where I need to build an initial model and then update this as new data arrives (eg. every month for monthly data) I know that sequential MCMC is a whole separate field but for a few reasons (mostly its generality) I want to find an efficient way of using NUTS in this context. Instead of the brute force method of re running the whole sampling from scratch on an expanding window of data (too slow), I’m considering:

1. Re using parameter estimates and mass matrix from initial sampler run to start the sampler on the expanded window of data and running fewer sampler iterations.

What’s the best way to get the (dense) mass matrix and initialisation values from run t in to the sampler for run t+1?

Or:

1. Use posterior estimates from first run of the sampler as priors in the second run, and only the new data. Pymc3 has an ‘Interpolated’ Distribution type which samples from an fitted empirical histogram of data to do this. (Example here https://docs.pymc.io/notebooks/updating_priors.html) I guess this could be implemented in numpyro? Alternatively I guess I can use my prior distributions re parameterised based on samples from previous fit...

Again I want to use mass matrix and parameter estimates from previous fits, so need to set up the sampler with these somehow?

[fehiepsi]

Hi @ross-h1, your requests are interesting! For parameter estimates, I believe you can simply calculate it from posterior samples. For the mass matrix, you can do the same (by calculating (co)variance of your samples) or extract from the hidden attribute (_last_state).

mcmc._last_state.adapt_state.inverse_mass_matrix

To specify initial values, you can use init_to_values initialization strategy. We currently do not expose initial mass matrix to NUTS/HMC class (though it is exposed in the functional hmc api). The reason is we want to wait for a user request before doing so. :)

re running the whole sampling from scratch

Tip: If your model does not have local latent variables, you can turn jit_model_args in MCMC to True to avoid recompiling your model.

Interpolated Distribution

In case you want to have this distribution, could you help us create a new feature request? It would be a very good first issue for NumPyro contributors. :)

[ross-h1]

Sounds like I have two requests then!

Had looked at hmc class, but i think the health warning in the docs is well founded!

So:

1. Expose Initial Mass Matrix in NUTS class

2. ’Interpolated’ distribution, although perhaps ‘Empirical’ may be a better name?

Very happy to help create a formal request!

On 6 Feb 2020, at 14:32, Du Phan [e-mail] wrote:


Hi @ross-h1, your requests are interesting! For parameter estimates, I believe you can simply calculate it from posterior samples. For the mass matrix, you can do the same (by calculating (co)variance of your samples) or extract from the hidden attribute (_last_state).

mcmc._last_state.adapt_state.inverse_mass_matrix
To specify initial values, you can use init_to_values initialization strategy. We currently do not expose initial mass matrix to NUTS/HMC class (though it is exposed in the functional hmc api). The reason is we want to wait for a user request before doing so. :)

re running the whole sampling from scratch

Tip: If your model does not have local latent variables, you can turn jit_model_args in MCMC to True to avoid recompiling your model.

Interpolated Distribution

In case you want to have this distribution, could you help us create a new feature request? It would be a very good first issue for NumPyro contributors. :)

—
You are receiving this because you were mentioned.
Reply to this email directly, view it on GitHub, or unsubscribe.

[ross-h1]

Also, I think the mass matrix is calculated on unconstrained HMC parameters, not on the constrained versions returned in samples?

So need to get at these somehow…🤔

The mass matrix does seem to have a significant effect on sampling speed where there are correlations in the posterior, so worthwhile being able to tweak I think, especially for sequential sampling where the posterior changes very little from epoch to epoch. I start my model with 120 observations, and increase by 1 every month, so it’s less than 1% change every iteration… seems very wasteful to estimate the mass matrix every run of the sampler...

On 6 Feb 2020, at 14:32, Du Phan [e-mail] wrote:

mcmc._last_state.adapt_state.inverse_mass_matrix

[fehiepsi]

the mass matrix is calculated on unconstrained HMC parameters

It is a very good point! Yeah, the better option would be to extract from mcmc._last_state.

’Interpolated’ distribution, although perhaps ‘Empirical’ may be a better name?

I am not 100% sure but it seems that Interpolated is a continuous distribution, while Empirical is a discrete one.

[ross-h1]

Yep, ok. Would be great to be able to build a new mass matrix though, so could perhaps update this after each set of samples, perhaps using some sort of exponential weighting… then plug back in for the next set of samples. This would provide maximum flexibility, and hopefully make NUTS more viable for sequential problems.

The mass matrix is only updated during warmup I think? This way could get the benefit of more samples contributing to the covariance matrix.

Also, how is warmup run using the hmc interface? Although its specified as a parameter in init_kernel, I don’t see how to actually run the warmup iterations and therefore update the mass matrix…

On 6 Feb 2020, at 17:31, Du Phan [e-mail] wrote:

the mass matrix is calculated on unconstrained HMC parameters

It is a very good point! Yeah, the better option would be to extract from mcmc._last_state.

—
You are receiving this because you were mentioned.
Reply to this email directly, view it on GitHub https://github.com/pyro-ppl/numpyro/issues/534?email_source=notifications&email_token=AF5RD3IBLYKRSAFL4BVOMATRBRCNVA5CNFSM4KQ47NK2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOELACROQ#issuecomment-583018682, or unsubscribe https://github.com/notifications/unsubscribe-auth/AF5RD3KDVMGQGFIBGWIP53DRBRCNVANCNFSM4KQ47NKQ.

[fehiepsi]

@ross-h1 You can run warmup using fori_collect in this example. To get the last warmup state, you can use fori_collect(num_warmup, num_warmup, ..., return_last_val=True) (sounds complicated... :D).

The mass matrix is only updated during warmup I think?

Yes.

[neerajprad]

I think we would like to support (2) for active learning. @thangbui and @karalets have also brought up this issue before. I'm not the best informed on this topic, but the Interpolated distribution in pymc3 sounds promising. I'm not sure how well that works in high dimensions.

ps - Thanks @theofanis for alerting me that I hadn't mentioned the right person, and sorry for spamming you unnecessarily.

[theofanis]

I hate it to disrupt any discussion here, but I think I was mentioned in a comment here by mistake, which is totally fine, but @neerajprad probably wanted to notify someone else, so it's only fair to drop by and mention it and help you guys reach a resolution faster 😄.

Feel free to delete my comment - or I will gladly do it - after you see this and mention the person you initially intended to.

[vanAmsterdam]

Tip: If your model does not have local latent variables, you can turn jit_model_args in MCMC to True to avoid recompiling your model.

this has been here for some time, but can you explain why local latent variables mess with the jitting speed-up?

[fehiepsi]

why local latent variables mess with the jitting speed-up?

Good point. I can't recall why I said so. The requirement should be: model args are jittable (e.g. they are numpy.ndarray or a tuple/dict of them).

[atlasd]

@ross-h1 did you figure out how to do this? I want to do the same thing, looking for an example.

[fehiepsi]

@atlasd This is supported in the last release. Please see an example here and let us know what you think.

[MrinankSharma]

@fehiepsi hey, i'm also requesting that the initial mass matrix could be exposed! By the way, do you know how I can get the correspondence between elements of the mass matrix and the actual variables in my model?

[fehiepsi]

how I can get the correspondence between elements of the mass matrix and the actual variables in my model

Good question! The correspondence is a bit ambiguous (depending on how jax converts a dictionary to a flatten array). That's why we haven't exposed inverse_mass_matrix to the api. To specify the initial inverse mass matrix, you can use MCMC.init to get the initial state, then you can replace the fields adapt_state.inverse_mass_matrix, adapt_state.mass_matrix_sqrt by your specified mass matrix. Then before running MCMC, setting post_warmup_state to that modified initial state.

To get the ordering, you can use z_flat, unravel_fn = jax.flatten_util.ravel_pytree(a_dict_of_latent_variables), then call unravel_fn(inverse_mass_matrix).

[MrinankSharma]

how I can get the correspondence between elements of the mass matrix and the actual variables in my model

Good question! The correspondence is a bit ambiguous (depending on how jax converts a dictionary to a flatten array). That's why we haven't exposed inverse_mass_matrix to the api. To specify the initial inverse mass matrix, you can use MCMC.init to get the initial state, then you can replace the fields adapt_state.inverse_mass_matrix, adapt_state.mass_matrix_sqrt by your specified mass matrix. Then before running MCMC, setting post_warmup_state to that modified initial state.

To get the ordering, you can use z_flat, unravel_fn = jax.flatten_util.ravel_pytree(a_dict_of_latent_variables), then call unravel_fn(inverse_mass_matrix).

Thanks for the quick response :)

I'm not sure where to get the dictionary of latent variables from either to be honest. For example, I know that numpyro will convert my variables into unconstrained versions, so I'd guess I can't pass some of the posterior samples.

[fehiepsi]

I see. It will be stored in the z field of the initial state (obtained from kernel's init method, e.g. NUTS.init(...)).

[MrinankSharma]

I see. It will be stored in the z field of the initial state (obtained from kernel's init method, e.g. NUTS.init(...)).

Thanks! managed to get it to work

[MrinankSharma]

how I can get the correspondence between elements of the mass matrix and the actual variables in my model

Good question! The correspondence is a bit ambiguous (depending on how jax converts a dictionary to a flatten array). That's why we haven't exposed inverse_mass_matrix to the api. To specify the initial inverse mass matrix, you can use MCMC.init to get the initial state, then you can replace the fields adapt_state.inverse_mass_matrix, adapt_state.mass_matrix_sqrt by your specified mass matrix. Then before running MCMC, setting post_warmup_state to that modified initial state.

To get the ordering, you can use z_flat, unravel_fn = jax.flatten_util.ravel_pytree(a_dict_of_latent_variables), then call unravel_fn(inverse_mass_matrix).

another question -- actually it seems that HMCAdaptState doesn't allow for item assignment, so I can't directly set the matrix.

[fehiepsi]

You can use _replace method to change field values of a namedtuple. In this case, you can do state = state._replace(adapt_state=state.adapt_state._replace(inverse_mass_matrix=...))