Numpyro: Best method for complicated multimodal posteriors

[Joshuaalbert]

Are there any plans to implement something like pyMC3's SMC which works quite well for multimodal posteriors. I perform inference in angular coordinates where phase wrapping occurs and this results in periodic structure in the posterior. NUTS often fails. I'll provide data simulation and model for completion. Note the effective sample size is low and the Rhat is high even after 1M draws.

import argparse
import time
import matplotlib.pyplot as plt
import numpy as np
from jax import random
import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS

def simulate_data(rng_key):
    """Simulate data"""
    freqs = jnp.linspace(121e6,168e6,24)
    tec = 55.
    const = 1.
    tec_conv = -8.4479745e6 / freqs
    phase = tec * tec_conv + const
    Y = jnp.concatenate([jnp.cos(phase), jnp.sin(phase)], axis=-1)
    Y_obs = Y + 0.5*random.normal(rng_key, shape=Y.shape)
    Sigma = 0.5*jnp.eye(48)
    return Y_obs, Sigma, tec_conv

def uniform_const_uniform_tec(Y_obs, Sigma, tec_conv):
    """Model"""
    sigma = jnp.sqrt(jnp.diag(Sigma))
    const = numpyro.sample('const', dist.Uniform(-jnp.pi, jnp.pi))
    tec = numpyro.sample('tec', dist.Uniform(-100., 100.))
    phase = tec * tec_conv + const
    Y = jnp.concatenate([jnp.cos(phase), jnp.sin(phase)], axis=-1)
    numpyro.sample('Y', dist.Normal(Y, sigma), obs=Y_obs)

def print_results(posterior, model):
    """Plot results"""
    print("Results for {}".format(model.__name__))
    plt.hist(posterior['tec'], bins=100)
    plt.title('tec maginal')
    plt.show()
    plt.hist(posterior['const'], bins=100)
    plt.title('const marginal')
    plt.show()
    plt.hist2d(posterior['const'], posterior['tec'],bins=100)
    plt.xlabel('const')
    plt.ylabel('tec')
    plt.title(model.__name__)
    plt.show()

def main(args):
    print('Simulating data...')
    (Y_obs, Sigma, tec_conv) = simulate_data(random.PRNGKey(1))
    print('Starting inference...')
    rng_key = random.PRNGKey(3)
    for model in [uniform_const_uniform_tec]:
        start = time.time()
        kernel = NUTS(model, target_accept_prob=args.target_prob, init_strategy=numpyro.infer.init_to_prior())
        mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains, progress_bar=False)
        mcmc.run(rng_key, Y_obs, Sigma, tec_conv)
        samples = mcmc.get_samples(group_by_chain=True)
        numpyro.diagnostics.print_summary(samples)
        print("Ground truth tec = 55. and const = 1.")
        print('\nMCMC elapsed time:', time.time() - start)
        samples = {k:v.reshape((v.shape[0]*v.shape[1],)+v.shape[2:]) for k,v in samples.items()}
        print_results(samples, model)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Get copula of const.')
    parser.add_argument('-n', '--num-samples', nargs='?', default=1000000, type=int)
    parser.add_argument('--num-warmup', nargs='?', default=10000, type=int)
    parser.add_argument("--num-chains", nargs='?', default=100, type=int)
    parser.add_argument("--target-prob", nargs='?', default=0.8, type=float)
    parser.add_argument('--device', default='cpu', type=str, help='use "cpu" or "gpu".')
    args = parser.parse_args()
    numpyro.set_platform(args.device)
    numpyro.set_host_device_count(args.num_chains)
    main(args)

It prints this:

                mean       std    median      5.0%     95.0%     n_eff     r_hat
     const      1.66      1.45      2.09     -0.19      3.14     20.96      1.53
       tec     40.29     43.83     56.57    -27.34     93.16     13.89      2.70

MCMC elapsed time: 101.53514671325684

[Joshuaalbert]

This would be useful example for https://github.com/pyro-ppl/numpyro/issues/189 once solved.

[fehiepsi]

Hi @Joshuaalbert, currently, we don't have a plan to implement SMC. However, we welcome contributions on new inference methods because NumPyro is community-oriented (like other Pyro projects).

I am not familiar with SMC. For multi-modal problems, I tend to use normalization flows to neutralize posterior geometry. For example,

from jax import lax
from numpyro.infer import ELBO, MCMC, NUTS, SVI
from numpyro.infer.autoguide import AutoBNAFNormal
from numpyro.infer.reparam import NeuTraReparam
import numpyro.optim as optim

guide = AutoBNAFNormal(model, num_flows=1, hidden_factors=[8, 8])
svi = SVI(model, guide, optim.Adam(0.003), ELBO(), Y_obs=Y_obs, Sigma=Sigma, tec_conv=tec_conv)
svi_state = svi.init(random.PRNGKey(2))
last_state, losses = lax.scan(lambda state, i: svi.update(state), svi_state, jnp.zeros(10000))
neutra = NeuTraReparam(guide, svi.get_params(last_state))
kernel = NUTS(neutra.reparam(model), target_accept_prob=args.target_prob)
mcmc = MCMC(kernel, args.num_warmup, args.num_samples, num_chains=args.num_chains, progress_bar=False)
mcmc.run(rng_key, Y_obs, Sigma, tec_conv)
zs = mcmc.get_samples(group_by_chain=True)["auto_shared_latent"]
samples = neutra.transform_sample(zs)

with num_samples=10000, I get

                mean       std    median      5.0%     95.0%     n_eff     r_hat
     const      1.35      1.06      1.53     -0.15      2.87 373331.97      1.00
       tec     61.66     18.08     64.75     35.40     88.28 372451.50      1.00

Ground truth tec = 55. and const = 1.

MCMC elapsed time: 26.009264945983887

Is it expected?

This would be useful example for #189 once solved.

Sure!!! Do you want to contribute one? ;)

[Joshuaalbert]

Thanks for this thoughtful reply. This result looks ok considering the priors were vague.
I would would be happy to contribute my use case once I am happy with it.
As for tempered sampling, I would love to have it, and would consider contributing it, but I'll need some guidance. I'll start another issue for that.

[fehiepsi]

Sure, please ping us for any questions that you have! The forum is also a good place to discuss because it will catch the attention of other users and it is easier to search for the old topics there.

[Joshuaalbert]

Hmm, I'm getting

ModuleNotFoundError: No module named 'numpyro.infer.autoguide'

Both on pip 0.2.4 and on master HEAD.

[Joshuaalbert]

Ah, no worries had to first delete package and reinstall

[Joshuaalbert]

@fehiepsi I get strange results with NeuTra and it seems the "good" result with the simulation parameters I posted was a coincidence. Firstly, the posteriors that I get out sometimes have shrunken or clipped domains (even if I use priors with real-line support I often see sharp cutoffs in the posterior, which I can only imagine is due to the nonlinearity of the NN). The posteriors are just wrong for most simulations parameters. However, as a sanity check, if I make the observational uncertainty very small then it works as expected. It could simply be that this inference problem is very difficult for any appreciable uncertainty. Might be interesting to expose more of the guide internals, e.g. the definition of the ARN model. Do you have any other tips for multi modal complex behaviour posteriors?

[Joshuaalbert]

NeuTra seems to be aimed at making HMC efficient and making better geometry when the curvature changes a lot throughout the domain. Can it actually handle multi-modality? Is there a way to visually see this?

[fehiepsi]

Opps, I don't have much experience with multi-modal problems. Using block neural autoregressive is by far the best approach I know. It can cover multi-modal pretty well as you can see from this image https://images.app.goo.gl/2ntJEW8kX33vJ6RK6

I did try NeuTra in large dataset such as covtype and it also helps mixing MCMC a lot. But that is the farthest I go. About some issues that you raised, I have such experience with AutoIAFNormal but not with AutoBNAFNormal. In the last comment, I used the default configs for BNAF. You might want to try other settings to see if the loss converges before applying MCMC.

One way to see if the guide covers multi-modal domains is to generate samples from it, then plot those samples as in https://github.com/pyro-ppl/numpyro/blob/2c9a68d5a8c1db75167bca4586d092835fcb6224/examples/neutra.py#L76

Could you update your model so I can see why BNAF failed for you? I also want to learn more from your experience on this stuff.

(BNAF uses tanh for nonlinearity. It is theoretically smooth but not in practice (its domain is about -20,20 in float32 mode). If using other invertable nonlinearity helps, we will expose it in the contructor of AutoBNAF)

[Joshuaalbert]

@fehiepsi I'll use this issue to post some of my findings. For a small introduction to problematic posteriors people might first read this. I think I'll focus in this thread on NeuTra. I think we'll need a few modifications to it for a truely robust implementation that "just works".

[fehiepsi]

Sure, I am looking forward to seeing your results. I just read that Stan article and tried to use NeuTra for such examples

y = random.normal(random.PRNGKey(0), (100,))

def model(y):
    lambda1 = numpyro.sample('lambda1', dist.ImproperUniform(dist.constraints.real, (), ()))
    lambda2 = numpyro.sample('lambda2', dist.ImproperUniform(dist.constraints.real, (), ()))
    sigma = numpyro.sample('sigma', dist.ImproperUniform(dist.constraints.positive, (), ()))
    mu = numpyro.deterministic('mu', lambda1 + lambda2)
    numpyro.sample('y', dist.Normal(mu, sigma), obs=y)

and got

                mean       std    median      5.0%     95.0%     n_eff     r_hat
   lambda1     -0.33      2.29     -0.34     -3.59      3.49    540.24      1.00
   lambda2      0.42      2.30      0.45     -3.37      3.72    538.99      1.00
        mu      0.09      0.10      0.09     -0.08      0.24    747.59      1.00
     sigma      0.96      0.07      0.95      0.83      1.06    490.47      1.00

which is pretty good comparing to vanilla NUTS without NeuTra in that Stan article. I guess NeuTra plays some regularization job here...

(#647 is needed to run SVI for models with improper distributions)

[Joshuaalbert]

@fehiepsi I did not find a nice way to fit nested sampling into numpyro for several reasons, so I wrote my own package: https://github.com/Joshuaalbert/jaxns. It does nested sampling using JAX and is pretty fast. I'll continue to support it there. If you feel it could come into numpyro (or some port of it) let me know.

[fehiepsi]

@Joshuaalbert It is cool! I think that your implementation will be very helpful for many tricky models that NUTS can't do well or requires further transformation steps.

I'll continue to support it there. If you feel it could come into numpyro (or some port of it) let me know.

Please just follow your gut. :) I am not sure I can assess the complication for an integration to NumPyro. FYI, having new inference algorithms is one of the main focus for the future of Pyro so we are very welcome you to contribute. Personally, I am happy to review and share my thoughts on your PR. IIUC, nested sampling is different from MCMC or SVI inferences so you have freedom to choose an interface for it (with a support for NumPyro model, rather than the pair log_likelihood and prior_chain). I'm looking forward to trying nested sampling on real problems whether it is ported to NumPyro or not. :D

[Joshuaalbert]

@fehiepsi Thanks, I'll improve it and at some point make a PR for numpyro.

[fehiepsi]

Hi @Joshuaalbert, we have an interesting thread on forum to discuss about this multimodal issue. Would you like to join and give us some insight? :)

[fehiepsi]

@joshuaalbert I just go over your examples. It seems to me that we can make a wrapper for NestedSampler to make it work with pyro models. I hope users can use NestedSamplerWrapper with numpyro model without having to specify anything new (e.g. PriorChain, UniformPrior) - like how we wrapped TFP mcmc kernels to make them work with numpyro models. Under the hood, we will

• extract a likelihood function (similar to how we convert a model to a joint density function in hmc) from a pyro model
• convert pyro priors to a PriorChain

What do you think? Can we setup a call to discuss things in more details if needed? (e.g. best place for this wrapper, how we support the related issues, testing,...) cc @fbartolic

I'm curious about priors in PriorChain. Is there anything special about it? If I understand correctly, then those are transformed distributions with base_dist are Uniform(0,1) - something like reparam can do. If so, it would be even simpler if we create Reparam classes and reparam the model. Then all priors in that reparam-ed model will be Uniform(0, 1). I might misunderstand the role of PriorChain here so please also correct me if I am wrong.

[Joshuaalbert]

I support this direction, @fehiepsi. I can meet after Nov 14. Perhaps Nov 16?

PriorChain is quite simple. Given a model M, and joint distribution factored into prior and likelihood like p(y, theta | M) = p(y|theta,M) p(theta|M) then the PriorChain represents a decomposition of the p(theta|M) into a particular form like p(theta|M) = p(theta_0|M) p(theta_1|theta_0,M) ... p(theta_n|theta_0,...,theta_n-1,M). The user specifies this decomposition in the usual way of probabilistic programming, by making instances of objects that represent those RVs and using them potentially in a hierarchical manner. The user should then "push" these prior objects onto a chain (in any order). If you push a prior object that depends on another prior object, then it will also push the dependencies.

A = DistrbutionA("A")
B = DistributionB("B", A)
prior_chain = PriorChain().push(A).push(B)
# same thing as
prior_chain = PriorChain().push(B)

The prior chain then has a __call__ method that transforms some base RV vector into a dictionary of the prior RVs. We use the uniform distribution as the base distribution. Therefore, each PriorTransform object requires a method of transforming uniform RVs to the target prior. This reparametrisation is an essential aspect of nested sampling. The prior RVs are dimensionfull, whereas the base uniform RVs are dimensionless. Doing operations like clustering and measuring distance between samples in the dimensionless spaces makes more sense.

Does this make sense?

[fehiepsi]

Thank you, @Joshuaalbert! Your explanation makes sense to me. I will try to have something working before we discuss. :+1:

[denijane]

Are there any plans to implement something like pyMC3's SMC which works quite well for multimodal posteriors. I perform inference in angular coordinates where phase wrapping occurs and this results in periodic structure in the posterior. NUTS often fails. I'll provide data simulation and model for completion. Note the effective sample size is low and the Rhat is high even after 1M draws.

def simulate_data(rng_key):
"""Simulate data"""
freqs = jnp.linspace(121e6,168e6,24)
tec = 55.
const = 1.
tec_conv = -8.4479745e6 / freqs
phase = tec * tec_conv + const
Y = jnp.concatenate([jnp.cos(phase), jnp.sin(phase)], axis=-1)
Y_obs = Y + 0.5random.normal(rng_key, shape=Y.shape)
Sigma = 0.5jnp.eye(48)
return Y_obs, Sigma, tec_conv

def uniform_const_uniform_tec(Y_obs, Sigma, tec_conv):
"""Model"""
sigma = jnp.sqrt(jnp.diag(Sigma))
const = numpyro.sample('const', dist.Uniform(-jnp.pi, jnp.pi))
tec = numpyro.sample('tec', dist.Uniform(-100., 100.))
phase = tec * tec_conv + const
Y = jnp.concatenate([jnp.cos(phase), jnp.sin(phase)], axis=-1)
numpyro.sample('Y', dist.Normal(Y, sigma), obs=Y_obs)

Hi, I'm sorry for the noob question, but could you explain to me what are tec and const? I've seen the tec and tec_conv in many example codes and I couldn't figure what they correspond to.

[Joshuaalbert]

tec means total electron content in units if [1/m^2], and tec_conv is a physical factor with units of [m^2] that converts tec to phase, representing path length distortions in radio astronomy calibration.

[denijane]

tec means total electron content in units if [1/m^2], and tec_conv is a physical factor with units of [m^2] that converts tec to phase, representing path length distortions in radio astronomy calibration.

Thanks a lot! It helped me understand how exactly you introduced the priors and now it's running.

[fehiepsi]

This thread has been a nice discussion thread so far. It seems that the new package jaxns is a good candidate for multimodal posteriors. For those who are interested in it, please try it out. We can move the further discussions to the forum if needed. :)