Implement Riemannian HMC

[xukai92]

This is a draft to implement Riemannian HMC. There are many things to discuss. I put the high-level points here while leaving more specific ones in the code.

• Do we want to unify all function signatures for Hamiltonian to a non-separable one (i.e. position dependent)?
• Do we need to update our metric abstraction? Looks like we should decouple unit/diagonal/dense vs Euclidean/Riemannian.
• Where should SoftAbs best locate? I current implemented it as an external function to provide G to a generic Riemannian HMC implementation.
• The efficiency is pretty bad and there is fewer optimization opportunities due to our decoupled abstraction (e.g. compared to what derived in [1]).

To-dos

• The current SoftAbs implementation has numerical issues when running on Neal's funnel. I need to double check the if the manual Jacobian implementation is correct, or we could register softabs with ReverseDiff to see if this solves the issue. I could use someone else help on this.

How to play with this PR

I provided a notebook (which contains the same content as the test/experimental/riemannian_hmc.jl file) to play with the code. The notebook has some simple validation on the implementation and also shows the current numerical issue of SoftAbs. I highly recommend you to try this.

---

[1] Betancourt, M., 2013, August. A general metric for Riemannian manifold Hamiltonian Monte Carlo. In International Conference on Geometric Science of Information (pp. 327-334). Springer, Berlin, Heidelberg.

[torfjelde]

One thing that is worth discussing wrt. this PR is how much of this functionality we should put into the metric types themselves vs. just implementing it all as a change of variables.

If we instead consider these position-dependent metrics as a change of variables for the potential instead, we don't have to deal with non-separable Hamiltonians internally and hence should hopefully make hte internals simpler.

Are there any obvious downsides to this?

[xukai92]

One thing that is worth discussing wrt. this PR is how much of this functionality we should put into the metric types themselves vs. just implementing it all as a change of variables. If we instead consider these position-dependent metrics as a change of variables for the potential instead, we don't have to deal with non-separable Hamiltonians internally and hence should hopefully make hte internals simpler. Are there any obvious downsides to this?

I also think so as I put a note here: https://github.com/TuringLang/AdvancedHMC.jl/blob/master/src/metric.jl#L116. Ideally we could interface Bijectors.jl here and do more than just affine transformation.

With this being said, the best route to implement this in AHMC is to do an internal refactoring PR first and then implement here. I feel that would require quite a bit of discussion.

Besides, I think most of the callback change to handle non-separable Hamiltonians that introduced in this PR is still needed there.

[torfjelde]

I also think so as I put a note here

Ah didn't catch that; cool! Btw, are we aware of any downsides to just doing a reparameterization for something like RHMC, given that the original paper didn't do this but instead went the route of Genearlized Leapfrog, which requires a root-solve in to determine the implictly defined step?

[xukai92]

Btw, are we aware of any downsides to just doing a reparameterization for something like RHMC, given that the original paper didn't do this but instead went the route of Genearlized Leapfrog, which requires a root-solve in to determine the implictly defined step?

I don't quite get the question. What the original paper does can also be seen as a form of reparameterization and we are doing the same thing together with genearlized Leapfrog, for which the implict updates are done by fixed-point iterations.

[xukai92]

Updates as of ed15539

1. Introduced types to control how Hessian is mapped (https://github.com/TuringLang/AdvancedHMC.jl/blob/kx/rhmc-draft/research/src/riemannian_hmc.jl#L136-L155). This gives an easy way to access internal parameters of the SoftAbs map and specialize the internals based on different maps (IdentityMap or SoftAbsMap).
   • PS: I made this choice mostly because of the easiness of implementation and I'm happy to consider other designs.
2. Fixed a bug in constructing J (https://github.com/TuringLang/AdvancedHMC.jl/blob/kx/rhmc-draft/research/src/riemannian_hmc.jl#L258) by adding a missing coefficient to csch
   • With 1 and 2, I refactored ∂H∂θ for the SoftAbs metric (https://github.com/TuringLang/AdvancedHMC.jl/blob/kx/rhmc-draft/research/src/riemannian_hmc.jl#L271-L300) and the sampling results visually look OK on Neal's funnel now:
3. The incorrect way of handling numerical issue is removed and any handling is done by users (e.g. https://github.com/TuringLang/AdvancedHMC.jl/blob/kx/rhmc-draft/research/tests/riemannian_hmc.jl#L40)
4. Performed Geweke test for RHMC (), which indicates that there are still bugs in the RHMC sampler with SoftAbs metric: The test is done by the model below

   @model function TuringFunnel(θ=missing, x=missing)
   if ismissing(θ)
       θ = Vector(undef, 2)
   end
   θ[1] ~ Normal(0, 3)
   s = exp(θ[1] / 2)
   θ[2] ~ Normal(0, s)
   x ~ Normal(0, s)
   return θ, x
   end

   with MCMCDebugging.jl.

PS: For reference, the test results of a normal HMC sampler under the same model is below:

I'm looking into the codes to see where the bug might be but if anyone has some ideas for debugging or intuition of where the bug might be, please let me know.

[xukai92]

Note of meeting with Hong on 23 Feb

Potential optimization

• G(θ) is unchanged in the fixed-point iteration for (16). We don't need to compute it for every call of ∂H∂θ.
• The first call of ∂H∂r is unchanged during (17). We can put it outside of the loop.

Tests

• Use finite differentiation to test our Hessian/metric computation and partial derivatives
• Test neg_energy using MvNormal
  • Double check the usage

[yebai]

@xukai92 I did some minor refactoring. I am happy to merge this PR as-is and revisit some design issues (e.g. ∂H∂θ_cache) when the research is complete.