Linear regression tutorial mixing

[ndmlny-qs]

Issue Description

The linear regression tutorial has been updated in branch ot-linear-reg-tutorial to have interactive
plotting. The notebook shows poor mixing of model chains except when using the BMGInference
method.

Steps to Reproduce

Execute the notebook to in the branch to see the mixing behavior.

Expected Behavior

Similar mixing behavior between the models.

System Info

Please provide information about your setup

• PyTorch Version (run print(torch.__version__)
  • 1.11.0
• Python version
  • 3.10.4

Additional Context

[horizon-blue]

Thanks for surfacing this, Andy. I can reproduce the same issue on my machine even without your modifications, so there is definitely some regression on BM side.

I am suspecting this is caused by improper initialization, because when I rerun the inference (without resetting the seed), sometimes NMC can still mix well. The same also applies when I try initialize from prior (by setting initialize_fn=bm.world. init_from_prior on infer). Could you try setting a different seed or changing the initialization function to see if either of these approach helps resolving the issue?

[ndmlny-qs]

• I changed the seed three times, all with the same issue surfacing.
• Reverting back to the original seed (12) and setting initialize_fn=bm.world.init_from_prior in the infer method still had poor mixing.
• I tried a few different initialization functions using a gamma, uniform, and halfnormal distribution. All ended up with poor mixing. The only time I got better mixing was when I used the gamma distribution in the initialize_fn with a concentration greater than the rate parameter, but this was not always the case, and I was not able to figure out a pattern.