turinglang / coupledhmc.jl

CoupledHMC.jl

This package implemented the coupled HMC kernels used in the AISTATS 2021 paper Couplings for Multinomial Hamiltonian Monte Carlo by Kai Xu, Tor Erlend Fjelde, Charles Sutton and Hong Ge.

See the research folder for scripts to reproduce the results.

You can cite this paper using the BibTeX entry below.

@inproceedings{xu2021couplings,
  title={{Couplings for Multinomial {H}amiltonian Monte Carlo}},
  author={Kai Xu and Fjelde, Tor Erlend and Sutton, Charles and Ge, Hong},
  booktitle={International Conference on Artificial Intelligence and Statistics},
  year={2021}
}

Using this package directly

If you would like to use CoupledHMC without going through our research folder, you will have to install VecTargets(https://github.com/xukai92/VecTargets.jl) and the master branch of AdvancedHMC manually. For example, running the command below will install them globally

julia -e 'using Pkg; Pkg.add(url="https://github.com/xukai92/VecTargets.jl"); Pkg.add(name="AdvancedHMC", rev="master")'

Or below to install inside your local environment

julia --project=@. -e 'using Pkg; Pkg.add(url="https://github.com/xukai92/VecTargets.jl"); Pkg.add(name="AdvancedHMC", rev="master")'

coupledhmc.jl's Issues

Useful helper

I found it's helpful to have a sample call that does: run untill meeting with a fixed number of step followed.
Something like below.

using CoupledHMC: get_ahmc_primitives, sample

function sample_after_meeting(target, alg::CoupledHMCSampler, n_after_samples::Int; theta0=nothing, progress=false)
    rng, hamiltonian, proposal, theta0 = get_ahmc_primitives(target, alg, theta0)
    samples = sample_until_meeting(rng, hamiltonian, proposal, theta0)
    samples_after_meeting = sample(
        target, HMCSampler(rinit=alg.rinit, TS=alg.TS, ϵ=alg.ϵ, L=alg.L), n_after_samples; 
        theta0=samples[end][:,1], progress=progress
    )
    # Duplicate the `n_samples_left` part to make the dimension consistent
    samples_after_meeting = map(s -> cat(s, s; dims=2), samples_after_meeting)
    return cat(samples, samples_after_meeting; dims=1)
end

This way of simulation kind of ensures unbiaseness with some level of confidence on variance (as we ask to always run for centain steps after the meeting).

Clean-up

IMO we should do the following:

1. Add examples in README.md
2. Add testing for the things that can be tested.
3. Allow arbitrary metrics to be used in the optimal transport coupling, i.e. just let the user pass in a function that computes the matrix?
4. Move stuff that is required for experiments only to research/, e.g. get_ahmc_primitives, the dependency on VecTargets.jl.
5. Add some more documentation to parts of the code.

Things that would be nice to have:

1. Actually make this an installable package.
2. Make a new release of AdvancedHMC.jl with the recent changes so that we don't need to depend on #master.
3. Actual documentation using Documenter.jl. Here we could also have some examples.

AISTATS camera-ready preparation

To-Dos

Code

The repo is currently organised as a Julia package together with scripts located in research that produces the results in our paper. We will need to do some clean-ups before making things public.

• Move research/src/VecTarget.jl into a separate repo
  • Can we merge it with ChallengingPosteriorBenchmarks.jl or so?
• Update the kernel interface in AdvancedHMC.jl
• Remove hacks (https://github.com/xukai92/CoupledHMC.jl/blob/master/src/CoupledHMC.jl#L14-L45) and adapt to the new kernel interface here
• Remove deprecated files (e.g. CoupledHMC2.jl)

Some extra experiments

• Gaussians with increasing dimensionality
• Confirm meeting time with adapted NUTS parameters

Paper

The draft is almost there but we still need to finish the followings.

• Trim down to the allowed number of pages~~
  • @yebai We are currently at 8.5 pages. Do you think it's possible to ask if this length is allowed as camera-ready?
  • Now 9-page is allowed 🎉
  • Slight over 9-page ATM
• Proofread the paper, especially proofs
• Check all bib entries
• Double check the new notation for coupled kernel from Tor
• Include a link to this repo
• Also submit to arXiv

Please leave comments and feel free to add/edit bullet points here directly @yebai @torfjelde

Make coupled HMC avaiable at AdvancedHMC by default

Options

• Porting codes to AdvancedHMC.jl
  • At least the maximal coupling part has no extra depedency
• Loading CoupledHMC from AdvancedHMC
  • We can make those needs JuMP depedency optional only when JuMP is loaded.