Add adaptation on HMC trajectory length to increase the expected Change in the Estimator of the Expected Square (ChEES-HMC)

[albcab]

Closes #382

[codecov]

Codecov Report

Merging #421 (56eb275) into main (d056670) will decrease coverage by 0.01%.
The diff coverage is 99.04%.

@@            Coverage Diff             @@
##             main     #421      +/-   ##
==========================================
- Coverage   99.14%   99.14%   -0.01%     
==========================================
  Files          49       50       +1     
  Lines        2117     2218     +101     
==========================================
+ Hits         2099     2199     +100     
- Misses         18       19       +1     

Files	Coverage Δ
blackjax/__init__.py	100.00% <100.00%> (ø)
blackjax/adaptation/__init__.py	100.00% <100.00%> (ø)
blackjax/adaptation/chees_adaptation.py	100.00% <100.00%> (ø)
blackjax/mcmc/hmc.py	98.92% <90.90%> (-1.08%)	⬇️

[rlouf]

Isn't it called ChEES? 😁

[albcab]

Oh no 🤦

[junpenglao]

@albcab time to pick this up again (no pressure lol)?

[twiecki]

What's missing here? Would be great to get this sampler in!

[albcab]

Need to rebase and resolve conflicts for all the remaining, restructuring, etc. that have been done during the past year. Other than that, the algorithm is done.

[twiecki]

@albcab Great, let us know if you need any help!

[twiecki]

Nice, seems like conflicts are solved, just some mypy errors:

blackjax/adaptation/chees_adaptation.py:257: error: Incompatible default for argument "jitter_seed" (default has type "None", argument has type "int")  [assignment]
blackjax/adaptation/chees_adaptation.py:257: note: PEP 484 prohibits implicit Optional. Accordingly, mypy has changed its default to no_implicit_optional=True
blackjax/adaptation/chees_adaptation.py:257: note: Use https://github.com/hauntsaninja/no_implicit_optional to automatically upgrade your codebase
tests/smc/test_tempered_smc.py:69: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]

[junpenglao]

Sending out first round of review, I will need a bit more time on chees_adaptation.py

[junpenglao]

Could you add a test to check that it learns a good step size and trajectory_length_moving_average?
Similar to https://github.com/tensorflow/probability/blob/bc2e0e0d0969d2b5e7cb0a8fbaaf93ba1e0adfb2/tensorflow_probability/python/experimental/mcmc/gradient_based_trajectory_length_adaptation_test.py#L334

[albcab]

Could you add a test to check that it learns a good step size and trajectory_length_moving_average? Similar to https://github.com/tensorflow/probability/blob/bc2e0e0d0969d2b5e7cb0a8fbaaf93ba1e0adfb2/tensorflow_probability/python/experimental/mcmc/gradient_based_trajectory_length_adaptation_test.py#L334

I don't get it; how do we know that an avg step size of 1.5 and an avg number of steps of 15 are the "correct" amounts? Also, they target an acceptance rate of .75 while we do .651, which would change the target step size and number of steps.

[junpenglao]

I don't get it; how do we know that an avg step size of 1.5 and an avg number of steps of 15 are the "correct" amounts? Also, they target an acceptance rate of .75 while we do .651, which would change the target step size and number of steps.

I think those values comes from empirical, we should get similar number with the similar log_density set up. Could you turn .651 into an kwarg?

[albcab]

I think those values comes from empirical, we should get similar number with the similar log_density set up. Could you turn .651 into an kwarg?

I've added the test, although the number of integration steps depends on the optimizer/learning rate used.

[junpenglao]

Great job, thanks for working on this!

[junpenglao]

I've added the test, although the number of integration steps depends on the optimizer/learning rate used.

Interesting. Makes sense.

[junpenglao]

Should be

Suggested change

	Value of the num_integration_steps / step_size parameter of
	Value of the num_integration_steps * step_size parameter of

right?

[twiecki]

Amazing! 🥳