NUTS Posterior Sampling

numpyro
1

Hello,

I’m not sure why sampling from the posterior after running MCMC, seems to be doing nothing. A toy example:

key = random.PRNGKey(42)
my_data = dist.Normal(loc=1.3, scale=0.7).sample(key, sample_shape=(100,))

def model(data):
    mu = numpyro.sample('mu', dist.Uniform(low=-2, high=2))
    std = numpyro.sample('std', dist.Uniform(low=0, high=1.5))
    with numpyro.plate("data", data.shape[0]):
        numpyro.sample("obs", dist.Normal(loc=mu, scale=std), obs=data)

# Inference
kernel = NUTS(model)
key, subkey = random.split(key)
posterior = MCMC(kernel, num_samples=1100, num_warmup=100)
posterior.run(subkey, my_data)

# Prior Sampling
key, subkey = random.split(key)
prior_samples = Predictive(model, {}, num_samples=10)(subkey, my_data)
x = jnp.linspace(-3, 3, 201)
for mu, s in zip(prior_samples["mu"], prior_samples["std"]):
    y = dist.Normal(loc=mu, scale=s).log_prob(x)
    plt.plot(x, jnp.exp(y), c="k", alpha=0.4)

Sampling from the prior works as expected, but sampling from the posterior seems to be doing nothing

# Posterior Sampling
key, subkey = random.split(key)
posterior_pred = Predictive(model, posterior.get_samples(), return_sites=['mu', 'std'])
post = posterior_pred(subkey, jnp.array([1]))
posterior.get_samples()['mu'] == post['mu']

the above code returns all Trues. On the other hand, the following code

key, subkey = random.split(key)
posterior_pred = Predictive(model, posterior.get_samples())
posterior_pred(subkey, jnp.array([1]))

Returns all ones.

Any guidance on how to properly use Preictive, will be highly appreciated

Thanks,
Fouad

2

You should set data=None in the model if you want to predict its value. E.g. post_pred(key, data=None). In the model, to specify plate size when data is None, you can use size=data.shape[0] if data is not None else 10.

3

Thank you!
Now I’m able to sample from the model in the posterior albeit not with the shape I want, but still can’t sample from the posterior of the parameters (which I could do with AutoLaplaceApproximation#sample_posterior

Not specifying return_sties and then invoking the predictive post_obs = posterior_pred(subkey, data=None)['obs'] post_obs has shape [number_of_mcmc_samples, plate_size_when_None] in my case [1100, 10].

Specifying return site, I still get what was passed in from the posterior.get_samples() which I don’t think I understand why it happens.

Thanks

4

Oh, you are using posterior samples from MCMC and want to get new posterior samples for them? I guess you want to use resample or something, not predictive. Or you can run MCMC again with new seeds. B

5

Thank you, I think I get it now. To sample from the the posterior parameters, we run MCMC more.

I’ve seen examples in the docs where Predictive is used to sample from the model given the samples from the MCMC run (e.g. here Bayesian Regression Using NumPyro — NumPyro documentation) Is it possible to do this in a manner independent of the shape of the posterior samples.

Really appreciate all the support

6

What do you mean by sampling from the model? If you need to get posterior predictive, you can run Predictive and get obs samples. mu and std samples are your MCMC posterior samples. We don’t have the notion of posterior predictive for latent variables mu and std.

To sample from the the posterior parameters, we run MCMC more

It should be “to get more posterior samples, we run MCMC more” I think. I guess I misunderstand what you meant by “sample”.

7

Sorry for the sloppy writing, I was just agreeing with you that to sample parameters one can run MCMC again.

Regarding the obs samples, my question is that I can’t seem to control the shape of the output of Predictive. I don’t understand why it follows the shape of the MCMC samples fed in.

8

why it follows the shape of the MCMC samples fed in

In Predictive, for each posterior sample, we will draw a corresponding obs sample. If you have 1000 posterior samples, you will get 1000 obs. If you want to get 200 x 1000 obs, one way is to broadcasting those 1000 posterior samples into (200 x 1000) posterior samples, then using Predictive with those 200 x 1000 posterior samples. If you want to get 100 obs, you can draw 100 random samples from those 1000 posterior samples so that you can provide 100 posterior samples to Predictive and get 100 obs.

1 Like