Numpyro: HMC fails to sample posterior properly while PyMC3 works fine

[fbartolic]

I'm having trouble with sampling from a posterior using NUTS for a relatively simple linear model. The chains fail to converge. When I implement the same model with the same priors in PyMC3, the sampler converges easily. Here is a notebook which reproduces the issue for me.

I can't figure out if there's a problem with the model or the sampler. I've used pretty much the same settings for the sampler in PyMC3 and numpyro and the two models look identical to me. I would print out the initial log density for both models to check that they're returning the same numerical value but I'm not sure how to do that in numpyro. I've also tried different versions of numpyro and that didn't change anything.

Can anyone else reproduce this issue?

[martinjankowiak]

what are the mean/std of p_scale for both samplers to more digits of precision?

[fbartolic]

Here's a higher precision output for a few parameters from the numpyro chains:

and the pymc3 chains:

I'm interested in the x-s.

I initialized the numpyro sampler at a sample from the pymc3 model but the numpyro sampler fails to converge independent of where I initialize the chains or how many tuning steps I choose.

[fehiepsi]

To get potential energy, you can use

mcmc.run(rng_key, extra_fields=("potential_energy",))
mcmc.get_extra_fields()["potential_energy"]

or you can compute it by hand to confirm. In my system, numpyro returns

array([-63874.68430024, -63874.68750484, -63872.97797588, -63872.96589239, ...

with the same init values init_vals = {"p": 1e-06 * np.random.rand(n), "fs": 0.99999, "p_scale": 2e-05}, while pymc3 potential energy stays at -74k. So it seems to me that numpyro returns better results...

[fbartolic]

Potential energy is minus the total logp (right?) and 74k >> 63k so pymc3 gives much higher logp. Anyway, if you plot the model in data space (see the updated notebook) you can clearly see that if you initialize the numpyro model at the same point as the pymc3 model, it fails to find the typical set. If instead you initialize it at a sample from the pymc3 chains then the chains get stuck and fail to explore the posterior properly (look at the R-hats and ESS). So something is definitely not right.

[fehiepsi]

Oops, you are right -74k < -63k... (usually -logp is a positive number, so I was confused). There are some source of this issue that I can think of for now:

• logp or its grad numerical error
• mass matrix is not adapted well

[fehiepsi]

@fbartolic It seems to me that setting regularize=False solves the issue. Could you confirm?

If that's the case, then we need to take some action here:

• setting it to False by default
• introduce a flag in HMC/NUTS that governs this behavior ~(to avoid API change, probably we can introduce adapt_mass_matrix=(True, {"regularize": False}))~

[fbartolic]

Haha yeah, I often get confused about the logp values.

Setting regularize=False does seem to fix the issue, the chains are converging now. Thanks a lot!

I see that regularize=True rescales the diagonal of the the covariance matrix to speed up sampling. Do you have an idea why it doesn't work for this model? I guess that the large dynamic range in the parameters has something to do with it.

[fehiepsi]

I think regularize is useful to deal with numerical issues when taking cholesky of the dense matrix when the number of parameters is large. Not sure about its role when dense_mass=False though.

@martinjankowiak Which action do you prefer in this situation?

I found a similar situation in this blog post

    # Stan uses a regularized estimator for the covariance matrix to
    # be less sensitive to numerical issues for large parameter spaces.
    # In the test case for this blog post, this isn't necessary and it
    # actually makes the performance worse so I'll disable it, but I
    # wanted to include the implementation here for completeness

[martinjankowiak]

it could be nice to expose regularize but i think the default should remain the same. this adaptation scheme is battle-tested in stan and i don't think a single model is enough evidence to justify changing defaults. the reason this particular model leads to problems is because it exhibits pretty extreme curvature. so e.g. if p were reparameterized/scaled appropriately numpyro nuts should work fine with default settings.

[fehiepsi]

Sounds reasonable to me. Let's introduce a kind of regularize_mass_matrix=True/False to HMC for convenience (for users, it is hard to debug this). But you're right, we should encourage users to rescale the model if someone faces this issue again.

@fbartolic the issue happens because the scale of the inverse mass matrix is so small that regularize term negatively affects the estimation. rescaling the parameters as Martin mentioned will help.