Modify turning condition for nuts #1466

fehiepsi · 2018-10-17T23:25:12Z

With the introduction of mass matrix, we have to modify the turning condition. It seems not straightforward, as discussed in section A.4.2. of A Conceptual Introduction to
Hamiltonian Monte Carlo. I will follow that section for the implementation (and compare the current implementation in Stan).

Add a flag to eliminate starting point from candidates of a trajectory
Implement new turning condition
Test
Debug the slowness (or why stepsize is so small when adapting)
Debug: gaussian conjugate tests fail

It would be fun btw. After this, I will add the multinomial sampling option for nuts. Hope that it would not be so complicated. ^^

neerajprad · 2018-10-18T08:06:18Z

pyro/infer/mcmc/nuts.py

+        self._eliminate_starting_point = True
+
+    def _is_turning(self, r_left, r_right, r_sum):
+        # We follow the strategy in Section A.4.2 of [2] for this implementation.


Would also suggest including this reference which derives the termination criterion explicitly.

In [2], the author also derives the criterion too. I find it easier to understand than differential geometry style (though it is a nice language).

Oh..its the same derivation. I was looking at 4.2 rather than the appendix.

fehiepsi · 2018-10-18T21:04:19Z

pyro/infer/mcmc/nuts.py

                    accepted = True
                    z = new_tree.z_proposal

-                if self._is_turning(z_left, r_left, z_right, r_right):  # stop doubling
+                r_sum += new_tree.r_sum


what a bug!!! it took me a lot of time to detect it. We should never do shortcut assignment for tensor.

Yeah, in place tensor ops can give non deterministic results, in many cases. Just curious - what was the cause of failure here?

@neerajprad When tree_depth = 0, we create a new_tree with 1 element. Hence r_sum = r_left + r_right, but the shortcut assignment gives r_sum = r_left. Because of this, is_turning will return True (r_sum - r_left = 0), which makes NUTS not run at all or just run with 1 velocity verlet step.

fehiepsi · 2018-10-18T21:47:58Z

tests/infer/mcmc/test_nuts.py

-        (0.02, False, False, False),
-        (0.02, False, True, False),
+        (0.1, False, False, False),
+        (0.1, False, True, False),


the old step_size works too but it is slow

neerajprad · 2018-10-18T21:48:07Z

tests/infer/mcmc/test_nuts.py

-    pytest.mark.skipif('CI' in os.environ or 'CUDA_TEST' in os.environ,
-                       reason='Slow test - skip on CI/CUDA')]
-)
+TEST_CASES = [


+1, this has been bugging me too for a while. 😄

yeah, me too ^^

fehiepsi · 2018-10-18T21:56:40Z

@neerajprad from the gaussian test, I can see that using mass matrix makes things easier to sample. Now we don't need many samples to pass the test.

neerajprad · 2018-10-18T21:57:01Z

Makefile

@@ -5,8 +5,8 @@ all: docs test
 install: FORCE
 	pip install -e .[dev,profile]

-uninstall: FORCE
-	pip uninstall pyro-ppl
+reinstall: FORCE


When do you need this? If you do an editable install via -e, it should pick up any local changes.

whoa, I have uninstalled and installed again and again to test examples, so I made this command. Thx @neerajprad!

neerajprad

Just some minor comments. Looks great otherwise!

neerajprad · 2018-10-18T22:26:21Z

Makefile

@@ -5,9 +5,6 @@ all: docs test
 install: FORCE
 	pip install -e .[dev,profile]



let us put the uninstall option back in.

neerajprad · 2018-10-18T22:53:54Z

pyro/infer/mcmc/nuts.py

+        # TODO: change to torch.dot for pytorch 1.0
+        if self.full_mass:
+            if ((r_sum - r_left_flat) * (self._inverse_mass_matrix.matmul(r_left_flat))).sum() > 0:
+                if ((r_sum - r_right_flat) * (self._inverse_mass_matrix.matmul(r_right_flat))) \


nit: why not just and this condition?

fehiepsi · 2018-10-19T00:19:07Z

@neerajprad Do you think that it is better to keep the old terminate condition (it is not a bug indeed) by using a hidden flag _riemannian_turning_condition? It might be good to have different versions to compare. If you think that it is good to have that flag, I will add it in the multinomial pull request. :)

neerajprad · 2018-10-19T00:46:20Z

Do you think that it is better to keep the old terminate condition (it is not a bug indeed) by using a hidden flag _riemannian_turning_condition?

The old terminating condition with the identity mass matrix? Do you mean that it is a valid terminating condition in terms of preserving detailed balance, even though it might not generate the longest trajectories? Unless you are already observing cases where the old terminating condition is yielding better results, I would be more inclined to just remove that option. In any case, this should be equivalent to our old code with mass matrix adaptation disabled, so we can always compare against that.

fehiepsi · 2018-10-19T01:36:18Z

Never mind, ignoring it is totally fine to me. :)

Summary: Pull Request resolved: #853 The original U-turn condition is well-deﬁned only for Euclidean manifolds (as discussed in [Appendix 4.2 of this paper](https://arxiv.org/pdf/1701.02434.pdf) or equivalently, [this paper](https://arxiv.org/pdf/1304.1920.pdf)), before introducing more complicated kinetic energy functions, it'd be better to first replace the original U-turn condition with the generalized version, i.e. terminate when either of these conditions is true: {F619168889} where {F619168973} (M^{-1} would be an identity matrix in this diff as we haven't implemented mass matrix adaptation scheme yet) and {F619169021} This this diff is analogous to pyro-ppl/pyro#1466. I noticed that in [Pyro and Numpyro's implementation](https://github.com/pyro-ppl/pyro/blob/dev/pyro/infer/mcmc/nuts.py#L163-L172), `rho` is defined slightly differently by excluding momentums at the boundary -- this was briefly discussed in [Stan's forum](https://discourse.mc-stan.org/t/nuts-misses-u-turns-runs-in-circles-until-max-treedepth/9727/44). The post mentioned another issue with the U turn condition which I will address in D28735950 to make the reviewing process easier :). Reviewed By: jpchen, neerajprad Differential Revision: D28424431 fbshipit-source-id: 4aa477c263f2902891f4cc39a6f6d820b3692f9f

fehiepsi added 3 commits October 17, 2018 06:47

temp save

9e33250

Merge remote-tracking branch 'upstream/dev' into multinomial

a88e5ac

temp save

27f0535

fehiepsi added the WIP label Oct 17, 2018

fehiepsi added 3 commits October 17, 2018 21:42

implement turning cond. with mass matrix

59a6b84

fix bug

0dcab88

modify makefile to fasten testing

ec9ea88

neerajprad reviewed Oct 18, 2018

View reviewed changes

fehiepsi commented Oct 18, 2018

View reviewed changes

fehiepsi added 2 commits October 18, 2018 17:43

fix the bug and move def of HMC conjugate gaussian test to nuts

d098434

revert num_samples in gamma test

db009bf

fehiepsi commented Oct 18, 2018

View reviewed changes

neerajprad reviewed Oct 18, 2018

View reviewed changes

fehiepsi added awaiting review and removed WIP labels Oct 18, 2018

neerajprad reviewed Oct 18, 2018

View reviewed changes

remove makefile reinstall, change torch.dot

5e71225

neerajprad previously approved these changes Oct 18, 2018

View reviewed changes

address comment

ff1ca4e

fehiepsi dismissed neerajprad’s stale review via ff1ca4e October 18, 2018 23:07

neerajprad approved these changes Oct 19, 2018

View reviewed changes

neerajprad merged commit 0f00b77 into pyro-ppl:dev Oct 19, 2018

horizon-blue mentioned this pull request May 27, 2021

Generalized U-turn condition facebookresearch/beanmachine#853

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Modify turning condition for nuts #1466

Modify turning condition for nuts #1466

fehiepsi commented Oct 17, 2018 •

edited

Loading

neerajprad Oct 18, 2018

fehiepsi Oct 18, 2018

neerajprad Oct 18, 2018

fehiepsi Oct 18, 2018

neerajprad Oct 18, 2018

fehiepsi Oct 18, 2018

fehiepsi Oct 18, 2018

neerajprad Oct 18, 2018

fehiepsi Oct 18, 2018

fehiepsi commented Oct 18, 2018

neerajprad Oct 18, 2018

fehiepsi Oct 18, 2018

neerajprad left a comment

neerajprad Oct 18, 2018

neerajprad Oct 18, 2018

fehiepsi commented Oct 19, 2018

neerajprad commented Oct 19, 2018

fehiepsi commented Oct 19, 2018

		@@ -5,9 +5,6 @@ all: docs test
		install: FORCE
		pip install -e .[dev,profile]

Modify turning condition for nuts #1466

Modify turning condition for nuts #1466

Conversation

fehiepsi commented Oct 17, 2018 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fehiepsi commented Oct 18, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

neerajprad left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fehiepsi commented Oct 19, 2018

neerajprad commented Oct 19, 2018

fehiepsi commented Oct 19, 2018

fehiepsi commented Oct 17, 2018 •

edited

Loading