A TensorFlow implementation of Hamiltonian Annealed Importance Sampling (HAIS). We implement the method described by Sohl-Dickstein and Culpepper in their paper "Hamiltonian Annealed Importance Sampling for partition function estimation".
- Partial momentum refresh (from HAIS paper). This preserves some fraction of the Hamiltonian Monte Carlo momentum across annealing distributions resulting in more accurate estimation.
- Adaptive step size for Hamiltonian Monte Carlo. This is a simple scheme to adjust the step size for each chain in order to push the smoothed acceptance rate towards a theoretical optimum.
We have used ideas and built upon the code from some of the following repositories:
- BayesFlow TensorFlow 1.4 - 1.6 contribution. This is now integrated into TensorFlow Probability.
- Sohl-Dickstein's Matlab implementation
- Xuechen Li's PyTorch (0.2.0) implementation of Bi-Directional Monte Carlo from "Sandwiching the marginal likelihood using bidirectional Monte Carlo"
- Tony Wu's Theano/Lasagne implementation of the methods described in "On the Quantitative Analysis of Decoder-Based Generative Models"
- jiamings's (unfinished?) TensorFlow implementation based on Tony Wu's Theano code.
- Stefan Webb's HMC AIS in tensorflow repository.
The tests that appear to be working include:
test-hmc: a simple test of the HMC implementationtest-hmc-mvn: a test of the HMC implementation that samples from a multivariate normaltest-hais-log-gamma: a simple test to sample from and calculate the log normaliser of an unnormalised log-Gamma density.test-hais-model1a-gaussian: a test that estimates the log marginal likelihood for model 1a with a Gaussian prior from Sohl-Dickstein and Culpepper (2011).
Install either the GPU version of TensorFlow (I don't know why but tensorflow-gpu==1.8 and
tensorflow-gpu==1.9 are >10x slower than 1.7 on my machine)
pip install tensorflow-gpu==1.7or the CPU version
pip install tensorflowthen install the project
pip install git+https://github.com/JohnReid/HAISThe implementation contains some documentation generated from the docstrings that may be useful. However it is probably easier to examine the test scripts and adapt them to your needs.