Repository: pyro-ppl/pyro

Issue:
[Bug?] HMC and NUTS are too sensitive to step-size for simplle Gamma/Beta model
### Issue Description

 A very simple generative model---a Beta likelihood whose two parameters have a Gamma prior on them:

$$\alpha\sim Gamma(1, 1)$$
$$\beta\sim Gamma(1, 1)$$
$$x\sim Beta(\alpha, \beta)$$

and then running `NUTS` to perform MCMC inference on it, leads to
the acceptance probability to be 0 poor and hence garbage estimates for $\alpha$ and $\beta$.

Changing the stepsize in `HMC`fixes the issue, but `NUTS` should be more robust to this scenario.

Plus, without enabling debugging log messages, the fact that 0 uncorrelated samples were actually provided was totally hidden. I guess I should open a new issue for a feature request for better output?

This issue has been referenced also in this forum post https://forum.pyro.ai/t/simple-mcmc-on-gamma-beta-providing-wrong-estimates/259

### Environment
 - Ubuntu `17.04`, python `3.6`
 - PyTorch `0.4.0`
 - Pyro `0.2.1`

### Code Snippet
```
import numpy as np
import matplotlib.pyplot as plt
import scipy.stats as ss
import torch
from torch.autograd import Variable
import pyro
import pyro.distributions as dist
import pyro.infer.mcmc as mcmc
from pyro.infer.abstract_infer import EmpiricalMarginal


a = 5
b = 1
n_samples = 1000
beta_data = ss.beta(a=a, b=b).rvs(n_samples).astype(np.float32)
print(beta_data.shape)
plt.hist(beta_data, bins=100)
plt.show()

def beta_model(data):
    alpha_0 = torch.ones(1)
    beta_0 = torch.ones(1)
    alpha_prior = pyro.sample('alpha', dist.Gamma(concentration=alpha_0, rate=beta_0))
    beta_prior = pyro.sample('beta', dist.Gamma(concentration=alpha_0, rate=beta_0))
    x = pyro.sample('x', dist.Beta(concentration1=alpha_prior, concentration0=beta_prior), obs=data)
    return x

# hmc_kernel = mcmc.HMC(beta_model, step_size=0.0855, num_steps=4)
nuts_kernel = mcmc.NUTS(beta_model, adapt_step_size=True)
mcmc_run = mcmc.MCMC(nuts_kernel, num_samples=1000, warmup_steps=10000).run(torch.from_numpy(beta_data))
beta_posterior = EmpiricalMarginal(mcmc_run, 'beta')
print('beta', beta_posterior.mean) 
alpha_posterior = EmpiricalMarginal(mcmc_run, 'alpha')
print('alpha', alpha_posterior.mean)
```



Discussion:
> Plus, without enabling debugging log messages, the fact that 0 uncorrelated samples were actually provided was totally hidden.

Yes, I think we should have a better default handler, as this is crucial for debugging how HMC/NUTS is doing. Yes, please create an issue for this, and I'll take care of that! 

Thanks a lot @neerajprad 

Resolving patch:
diff --git a/pyro/infer/mcmc/hmc.py b/pyro/infer/mcmc/hmc.py
--- a/pyro/infer/mcmc/hmc.py
+++ b/pyro/infer/mcmc/hmc.py
@@ -200,7 +200,8 @@ def setup(self, *args, **kwargs):
             z = {name: node["value"] for name, node in trace.iter_stochastic_nodes()}
             for name, transform in self.transforms.items():
                 z[name] = transform(z[name])
-            self.step_size = self._find_reasonable_step_size(z)
+            with pyro.validation_enabled(False):
+                self.step_size = self._find_reasonable_step_size(z)
             self.num_steps = max(1, int(self.trajectory_length / self.step_size))
             # make prox-center for Dual Averaging scheme
             loc = math.log(10 * self.step_size)
diff --git a/pyro/infer/mcmc/nuts.py b/pyro/infer/mcmc/nuts.py
--- a/pyro/infer/mcmc/nuts.py
+++ b/pyro/infer/mcmc/nuts.py
@@ -37,7 +37,9 @@ class NUTS(HMC):
     **References**
 
     [1] `The No-U-turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo`,
-    Matthew D. Hoffman, and Andrew Gelman
+    Matthew D. Hoffman, and Andrew Gelman.
+    [2] `A Conceptual Introduction to Hamiltonian Monte Carlo`, Michael Betancourt
+    [3] `Slice Sampling`, Radford M. Neal
 
     :param model: Python callable containing Pyro primitives.
     :param float step_size: Determines the size of a single step taken by the
@@ -211,17 +213,17 @@ def sample(self, trace):
         #     first sampling u from initial state (z_0, r_0) according to u ~ Uniform(0, p(z_0, r_0)),
         #     then sampling state (z, r) from the integrator trajectory according to
         #         (z, r) ~ Uniform({(z', r') in trajectory | p(z', r') >= u}).
-        # For more information about slice sampling method, see
-        #     `Slice sampling` by Radford M. Neal.
+        #
+        # For more information about slice sampling method, see [3].
         # For another version of NUTS which uses multinomial sampling instead of slice sampling, see
-        #     `A Conceptual Introduction to Hamiltonian Monte Carlo` by Michael Betancourt.
-        joint_prob = torch.exp(-energy_current)
-        if joint_prob == 0:
-            slice_var = energy_current.new_tensor(0.0)
-        else:
-            slice_var = pyro.sample("slicevar_t={}".format(self._t),
-                                    dist.Uniform(torch.zeros(1), joint_prob))
-        log_slice = slice_var.log()
+        # [2].
+
+        # Rather than sampling the slice variable from `Uniform(0, exp(-energy))`, we can
+        # sample log_slice directly using `energy`, so as to avoid potential underflow or
+        # overflow issues ([2]).
+        slice_exp_term = pyro.sample("slicevar_exp_t={}".format(self._t),
+                                     dist.Exponential(energy_current.new_tensor(1.)))
+        log_slice = -energy_current - slice_exp_term
 
         z_left = z_right = z
         r_left = r_right = r


Test patch:
diff --git a/tests/infer/mcmc/test_hmc.py b/tests/infer/mcmc/test_hmc.py
--- a/tests/infer/mcmc/test_hmc.py
+++ b/tests/infer/mcmc/test_hmc.py
@@ -171,7 +171,7 @@ def model(data):
     assert_equal(rmse(true_coefs, beta_posterior.mean).item(), 0.0, prec=0.1)
 
 
-def test_bernoulli_beta():
+def test_beta_bernoulli():
     def model(data):
         alpha = torch.tensor([1.1, 1.1])
         beta = torch.tensor([1.1, 1.1])
@@ -187,7 +187,7 @@ def model(data):
     assert_equal(posterior.mean, true_probs, prec=0.05)
 
 
-def test_normal_gamma():
+def test_gamma_normal():
     def model(data):
         rate = torch.tensor([1.0, 1.0])
         concentration = torch.tensor([1.0, 1.0])
@@ -203,7 +203,7 @@ def model(data):
     assert_equal(posterior.mean, true_std, prec=0.05)
 
 
-def test_categorical_dirichlet():
+def test_dirichlet_categorical():
     def model(data):
         concentration = torch.tensor([1.0, 1.0, 1.0])
         p_latent = pyro.sample('p_latent', dist.Dirichlet(concentration))
@@ -236,7 +236,7 @@ def model(data):
     assert_equal(rmse(posterior.mean, true_coefs).item(), 0.0, prec=0.1)
 
 
-def test_bernoulli_beta_with_dual_averaging():
+def test_beta_bernoulli_with_dual_averaging():
     def model(data):
         alpha = torch.tensor([1.1, 1.1])
         beta = torch.tensor([1.1, 1.1])
@@ -252,7 +252,7 @@ def model(data):
     assert_equal(posterior.mean, true_probs, prec=0.05)
 
 
-def test_normal_gamma_with_dual_averaging():
+def test_gamma_normal_with_dual_averaging():
     def model(data):
         rate = torch.tensor([1.0, 1.0])
         concentration = torch.tensor([1.0, 1.0])
diff --git a/tests/infer/mcmc/test_nuts.py b/tests/infer/mcmc/test_nuts.py
--- a/tests/infer/mcmc/test_nuts.py
+++ b/tests/infer/mcmc/test_nuts.py
@@ -20,8 +20,11 @@
 T2 = T(*TEST_CASES[2].values)._replace(num_samples=800, warmup_steps=200)
 TEST_CASES[2] = pytest.param(*T2, marks=pytest.mark.skipif(
     'CI' in os.environ and os.environ['CI'] == 'true', reason='Slow test - skip on CI'))
-T3 = T(*TEST_CASES[3].values)._replace(num_samples=700, warmup_steps=100)
-TEST_CASES[3] = T3
+T3 = T(*TEST_CASES[3].values)._replace(num_samples=1000, warmup_steps=200)
+TEST_CASES[3] = pytest.param(*T3, marks=[
+    pytest.mark.skipif('CI' in os.environ and os.environ['CI'] == 'true',
+                       reason='Slow test - skip on CI')]
+)
 
 
 @pytest.mark.parametrize(
@@ -82,7 +85,7 @@ def model(data):
     assert_equal(rmse(true_coefs, posterior.mean).item(), 0.0, prec=0.1)
 
 
-def test_bernoulli_beta():
+def test_beta_bernoulli():
     def model(data):
         alpha = torch.tensor([1.1, 1.1])
         beta = torch.tensor([1.1, 1.1])
@@ -98,7 +101,7 @@ def model(data):
     assert_equal(posterior.mean, true_probs, prec=0.02)
 
 
-def test_normal_gamma():
+def test_gamma_normal():
     def model(data):
         rate = torch.tensor([1.0, 1.0])
         concentration = torch.tensor([1.0, 1.0])
@@ -132,8 +135,7 @@ def model(data):
     assert_equal(rmse(true_coefs, posterior.mean).item(), 0.0, prec=0.1)
 
 
-@pytest.mark.filterwarnings("ignore:Encountered NAN")
-def test_bernoulli_beta_with_dual_averaging():
+def test_beta_bernoulli_with_dual_averaging():
     def model(data):
         alpha = torch.tensor([1.1, 1.1])
         beta = torch.tensor([1.1, 1.1])
@@ -149,7 +151,7 @@ def model(data):
     assert_equal(posterior.mean, true_probs, prec=0.03)
 
 
-def test_categorical_dirichlet():
+def test_dirichlet_categorical():
     def model(data):
         concentration = torch.tensor([1.0, 1.0, 1.0])
         p_latent = pyro.sample('p_latent', dist.Dirichlet(concentration))
@@ -162,3 +164,18 @@ def model(data):
     mcmc_run = MCMC(nuts_kernel, num_samples=200, warmup_steps=100).run(data)
     posterior = EmpiricalMarginal(mcmc_run, sites='p_latent')
     assert_equal(posterior.mean, true_probs, prec=0.02)
+
+
+def test_gamma_beta():
+    def model(data):
+        alpha_prior = pyro.sample('alpha', dist.Gamma(concentration=1., rate=1.))
+        beta_prior = pyro.sample('beta', dist.Gamma(concentration=1., rate=1.))
+        pyro.sample('x', dist.Beta(concentration1=alpha_prior, concentration0=beta_prior), obs=data)
+
+    true_alpha = torch.tensor(5.)
+    true_beta = torch.tensor(1.)
+    data = dist.Beta(concentration1=true_alpha, concentration0=true_beta).sample(torch.Size((5000,)))
+    nuts_kernel = NUTS(model, adapt_step_size=True)
+    mcmc_run = MCMC(nuts_kernel, num_samples=500, warmup_steps=200).run(data)
+    posterior = EmpiricalMarginal(mcmc_run, sites=['alpha', 'beta'])
+    assert_equal(posterior.mean, torch.stack([true_alpha, true_beta]), prec=0.05)