Simple model inference problem

[andreaRapuzzi]

Issue Description

I'm building a very simple model to generate a signal and a corresponding BeanMachine model to infer its parameters. But in some specific configurations, the inference does not work.

Steps to Reproduce

import beanmachine.ppl as bm
import torch.distributions as dist
from torch import tensor

# an average signal property that our model shall infer
signal_avg = 116
# our signal is modeled as a Poisson function with lambda = signal_avg
s1 = dist.Poisson(signal_avg)

# let's generate some observations
s1_obs = s1.sample((100, )).tolist()

# BeanMachine Model
@bm.random_variable
def signal_avg_rv():
    # our priori knowledge is deliberately foresighted
    return dist.Normal(116, 0.1)

@bm.random_variable
def s1_rv(i):
    return dist.Poisson(signal_avg_rv())

observations = {s1_rv(i): tensor(s1_obs[i]) for i in range(len(s1_obs))}

samples = bm.CompositionalInference().infer(
    queries=[signal_avg_rv()],
    observations=observations,
    num_samples=700,
    num_adaptive_samples=300,
    num_chains=3,
)

Expected Behavior

I expect that all the traces provide sensible values for the posterior distribution of signal_avg_rv. Instead, a lot of the traces are stuck at small values (like 0.2, -0.2, etc) instead of being around 116.

Additional Info

1. If s1 and s1_rv follow a Normal distribution (instead of a Poisson) everything works fine and a good posterior for signal_avg_rv is always inferred.
2. If the prior on signal_avg_rv is bounded (like dist.Uniform(100, 150)) everything works fine. The prior we have set has, yes, an unbounded domain, but is, in practice super-informative (e.g. we have a very strong and precise belief about it).
3. When we use other samplers (like SingleSiteNoUTurnSampler or GlobalNoUTurnSampler) the inference sometimes get stuck (does not progress). Bonus track :-) : is it possible to verify which sampler has been used (e.g. querying the samples)?
4. The same simple model seems to run with other engines (e.g. pymc3)
5. Adding samples (up to a quarter of a million) does not help.

System Info

Python: 3.8.12
PyTorch: 1.10.1+cpu
BeanMachine: 0.1.0.post1

[jpchen]

Thanks for the detailed issue @andreaRapuzzi! There are a few possible issues I see:
Bean Machine initializes from Uniform(-2, 2), which is the same initialization used in Stan . Since your prior has very low probability at 0 , after transforming the initialized value, the sampler does not do a good job moving in the right direction. You can initialize from the prior instead:

from beanmachine.ppl.world import init_from_prior
bm.GlobalNoUTurnSampler(..., initialize_fn=init_from_prior)

Or reparameterize your model:

@bm.random_variable
def signal_avg_shifted():
    # our priori knowledge is deliberately foresighted
    return dist.Normal(0., 0.1)

@bm.functional
def signal_avg_rv():
    # shift the mean
    return signal_avg_shifted() + 116.0

@bm.random_variable
def s1_rv(i):
    return dist.Poisson(signal_avg_rv())

is it possible to verify which sampler has been used

Compositional uses Newtonian Monte Carlo for continuous RVs and uniform for discrete RVs by default, which we will change to use GlobalNoUTurnSampler in future releases which may be more robust. Since proposers (and their transforms) are assigned per random variable, you can call the get_proposers method after specifying which random variable and world you are looking at (You can construct a random world to see what Compositional will assign them).

[andreaRapuzzi]

Fantastic @jpchen !!!
Both the proposed solutions work and the get_proposers method is what I needed for diagnostics.
Thanks a lot!!!