Repository: pyro-ppl/numpyro

Issue:
Sequential Sampling Strategy
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:

2. 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?


Discussion:
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](http://num.pyro.ai/en/stable/utilities.html#init-to-value) initialization strategy. We currently do not expose initial mass matrix to NUTS/HMC class (though it is exposed in the functional [hmc](http://num.pyro.ai/en/stable/mcmc.html#numpyro.infer.mcmc.hmc.init_kernel) 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](http://pyro.ai/examples/svi_part_ii.html#Subsampling-when-there-are-only-local-random-variables), you can turn `jit_model_args` in [MCMC](http://num.pyro.ai/en/stable/mcmc.html#numpyro.infer.mcmc.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. :)
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 <notifications@github.com> 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.

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 <notifications@github.com> wrote:
> 
> mcmc._last_state.adapt_state.inverse_mass_matrix


> 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.
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 <notifications@github.com> 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>.
> 


@ross-h1 You can run warmup using `fori_collect` in [this example](https://github.com/pyro-ppl/numpyro/blob/master/examples/stochastic_volatility.py#L91). 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.
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.
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.

> Tip: If your model does not have [local latent variables](http://pyro.ai/examples/svi_part_ii.html#Subsampling-when-there-are-only-local-random-variables), you can turn `jit_model_args` in [MCMC](http://num.pyro.ai/en/stable/mcmc.html#numpyro.infer.mcmc.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?
> 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).

Resolving patch:
diff --git a/numpyro/distributions/discrete.py b/numpyro/distributions/discrete.py
--- a/numpyro/distributions/discrete.py
+++ b/numpyro/distributions/discrete.py
@@ -53,7 +53,7 @@
 
 
 def _to_probs_bernoulli(logits):
-    return 1 / (1 + jnp.exp(-logits))
+    return expit(logits)
 
 
 def _to_logits_bernoulli(probs):
diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py
--- a/numpyro/infer/hmc.py
+++ b/numpyro/infer/hmc.py
@@ -121,12 +121,10 @@ def hmc(potential_fn=None, potential_fn_gen=None, kinetic_fn=None, algo='NUTS'):
 
         >>> true_coefs = jnp.array([1., 2., 3.])
         >>> data = random.normal(random.PRNGKey(2), (2000, 3))
-        >>> dim = 3
         >>> labels = dist.Bernoulli(logits=(true_coefs * data).sum(-1)).sample(random.PRNGKey(3))
         >>>
         >>> def model(data, labels):
-        ...     coefs_mean = jnp.zeros(dim)
-        ...     coefs = numpyro.sample('beta', dist.Normal(coefs_mean, jnp.ones(3)))
+        ...     coefs = numpyro.sample('coefs', dist.Normal(jnp.zeros(3), jnp.ones(3)))
         ...     intercept = numpyro.sample('intercept', dist.Normal(0., 10.))
         ...     return numpyro.sample('y', dist.Bernoulli(logits=(coefs * data + intercept).sum(-1)), obs=labels)
         >>>
@@ -137,7 +135,7 @@ def hmc(potential_fn=None, potential_fn_gen=None, kinetic_fn=None, algo='NUTS'):
         ...                         num_warmup=300)
         >>> samples = fori_collect(0, 500, sample_kernel, hmc_state,
         ...                        transform=lambda state: model_info.postprocess_fn(state.z))
-        >>> print(jnp.mean(samples['beta'], axis=0))  # doctest: +SKIP
+        >>> print(jnp.mean(samples['coefs'], axis=0))  # doctest: +SKIP
         [0.9153987 2.0754058 2.9621222]
     """
     if kinetic_fn is None:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -382,11 +382,42 @@ def _compile(self, rng_key, *args, extra_fields=(), init_params=None, **kwargs):
         except TypeError:
             pass
 
+    @property
+    def post_warmup_state(self):
+        """
+        The state before the sampling phase. If this attribute is not None,
+        :meth:`run` will skip the warmup phase and start with the state
+        specified in this attribute.
+
+        .. note:: This attribute can be used to sequentially draw MCMC samples. For example,
+
+            .. code-block:: python
+
+                mcmc = MCMC(NUTS(model), 100, 100)
+                mcmc.run(random.PRNGKey(0))
+                first_100_samples = mcmc.get_samples()
+                mcmc.post_warmup_state = mcmc.last_state
+                mcmc.run(mcmc.post_warmup_state.rng_key)  # or mcmc.run(random.PRNGKey(1))
+                second_100_samples = mcmc.get_samples()
+        """
+        return self._warmup_state
+
+    @post_warmup_state.setter
+    def post_warmup_state(self, state):
+        self._warmup_state = state
+
+    @property
+    def last_state(self):
+        """
+        The final MCMC state at the end of the sampling phase.
+        """
+        return self._last_state
+
     def warmup(self, rng_key, *args, extra_fields=(), collect_warmup=False, init_params=None, **kwargs):
         """
-        Run the MCMC warmup adaptation phase. After this call, the :meth:`run` method
-        will skip the warmup adaptation phase. To run `warmup` again for the new data,
-        it is required to run :meth:`warmup` again.
+        Run the MCMC warmup adaptation phase. After this call, `self.warmup_state` will be set
+        and the :meth:`run` method will skip the warmup adaptation phase. To run `warmup` again
+        for the new data, it is required to run :meth:`warmup` again.
 
         :param random.PRNGKey rng_key: Random number generator key to be used for the sampling.
         :param args: Arguments to be provided to the :meth:`numpyro.infer.mcmc.MCMCKernel.init` method.


Test patch:
diff --git a/test/test_mcmc.py b/test/test_mcmc.py
--- a/test/test_mcmc.py
+++ b/test/test_mcmc.py
@@ -115,6 +115,7 @@ def model(data):
     kernel = NUTS(model=model)
     mcmc = MCMC(kernel, num_warmup=num_warmup, num_samples=num_samples)
     mcmc.warmup(random.PRNGKey(2), data, collect_warmup=True)
+    assert mcmc.post_warmup_state is not None
     warmup_samples = mcmc.get_samples()
     mcmc.run(random.PRNGKey(3), data)
     samples = mcmc.get_samples()
@@ -122,6 +123,12 @@ def model(data):
     assert len(samples['loc']) == num_samples
     assert_allclose(jnp.mean(samples['loc'], 0), true_coef, atol=0.05)
 
+    mcmc.post_warmup_state = mcmc.last_state
+    mcmc.run(random.PRNGKey(3), data)
+    samples = mcmc.get_samples()
+    assert len(samples['loc']) == num_samples
+    assert_allclose(jnp.mean(samples['loc'], 0), true_coef, atol=0.05)
+
 
 def test_improper_normal():
     true_coef = 0.9
@@ -293,7 +300,7 @@ def model(data):
     mcmc1.warmup(random.PRNGKey(2), data)
     mcmc1.run(random.PRNGKey(3), data)
     check_close(mcmc1.get_samples(), mcmc.get_samples(), atol=1e-4, rtol=1e-4)
-    check_close(mcmc1._warmup_state, mcmc._warmup_state, atol=1e-4, rtol=1e-4)
+    check_close(mcmc1.post_warmup_state, mcmc.post_warmup_state, atol=1e-4, rtol=1e-4)
 
 
 @pytest.mark.parametrize('kernel_cls', [HMC, NUTS])
@@ -603,7 +610,7 @@ def model():
     mcmc._compile(rng_key)
     # no delay after compiling
     mcmc.warmup(rng_key)
-    mcmc.run(mcmc._warmup_state.rng_key)
+    mcmc.run(mcmc.last_state.rng_key)
     actual_samples = mcmc.get_samples()["x"]
 
     assert_allclose(actual_samples, expected_samples)