TuringLang / AdvancedHMC.jl

Robust, modular and efficient implementation of advanced Hamiltonian Monte Carlo algorithms
https://turinglang.org/AdvancedHMC.jl/
MIT License
237 stars 41 forks source link

Add low-rank metric adaptation #277

Open sethaxen opened 3 years ago

sethaxen commented 3 years ago

The following paper presented a novel approach for adapting the metric:

Ben Bales, Arya Pourzanjani, Aki Vehtari, Linda Petzold, Selecting the Metric in Hamiltonian Monte Carlo, 2019. https://arxiv.org/abs/1905.11916

The idea is to adapt a low-rank (plus diagonal) approximation to the covariance matrix using a diagonal estimate and a low-rank approximation of the Hessian matrix. The main advantage of the estimated covariance being low-rank is that it can be stably estimated with many fewer draws than the dense covariance matrix. The paper also included a selection criterion for determining which rank to use.

They demonstrated in several benchmarks that the method often (but not always) outperformed diagonal and dense metric adaptation in terms of ESS/s.

sethaxen commented 2 years ago

See also #70

sethaxen commented 2 years ago

As noted by @nsiccha on Slack, the reference code for the paper is available at https://github.com/stan-dev/stan/compare/develop...bbbales2:stan:experimental/warmup-wishart.

A good first step would be to implement a low-rank update Euclidean metric type. This is already implemented in Pathfinder (https://github.com/mlcolab/Pathfinder.jl/blob/96f76ef4471c98964693c6afad92c5c19f209138/src/integration/advancedhmc.jl) and could be ported here, avoiding a PDMats dependency by simply storing the component matrices and factorizations in the metric object.

sethaxen commented 2 years ago

As noted by @nsiccha on Slack, the reference code for the paper is available at stan-dev/stan@develop...bbbales2:stan:experimental/warmup-wishart.

Specifically, these files seem to contain all code for the metric adaptation, while the rest is mostly glue code so it can be used with NUTS:

nsiccha commented 2 years ago

Yeah, after having a really short look, I also arrived at those two files. Haven't had the opportunity to port i anything yet though.