Detailed API for AdvancedHMC.jl
An important design goal of AdvancedHMC.jl is modularity; we would like to support algorithmic research on HMC. This modularity means that different HMC variants can be easily constructed by composing various components, such as preconditioning metric (i.e., mass matrix), leapfrog integrators, trajectories (static or dynamic), adaption schemes, etc. In this section, we will explain the detailed usage of different modules in AdancedHMC.jl to provide a comprehensive udnerstanding of how AdvancedHMC.jl can achieve both modularity and efficiency. The section highlights the key components of AdvancedHMC.jl, with a complete documentation provided at the end.
Hamiltonian mass matrix (metric
)
- Unit metric:
UnitEuclideanMetric(dim)
- Diagonal metric:
DiagEuclideanMetric(dim)
- Dense metric:
DenseEuclideanMetric(dim)
where dim
is the dimensionality of the sampling space.
Integrator (integrator
)
- Ordinary leapfrog integrator:
Leapfrog(ϵ)
- Jittered leapfrog integrator with jitter rate
n
:JitteredLeapfrog(ϵ, n)
- Tempered leapfrog integrator with tempering rate
a
:TemperedLeapfrog(ϵ, a)
where ϵ
is the step size of leapfrog integration.
Kernel (kernel
)
- Static HMC with a fixed number of steps (
n_steps
) from Neal [1]:HMCKernel(Trajectory{EndPointTS}(integrator, FixedNSteps(integrator)))
- HMC with a fixed total trajectory length (
trajectory_length
) from Neal [1]:HMCKernel(Trajectory{EndPointTS}(integrator, FixedIntegrationTime(trajectory_length)))
- Original NUTS with slice sampling from Hoffman et al. [2]:
HMCKernel(Trajectory{SliceTS}(integrator, ClassicNoUTurn()))
- Generalised NUTS with slice sampling from Betancourt [3]:
HMCKernel(Trajectory{SliceTS}(integrator, GeneralisedNoUTurn()))
- Original NUTS with multinomial sampling from Betancourt [3]:
HMCKernel(Trajectory{MultinomialTS}(integrator, ClassicNoUTurn()))
- Generalised NUTS with multinomial sampling from Betancourt [3]:
HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
Adaptor (adaptor
)
Adapt the mass matrix
metric
of the Hamiltonian dynamics:mma = MassMatrixAdaptor(metric)
- This is lowered to
UnitMassMatrix
,WelfordVar
orWelfordCov
based on the type of the mass matrixmetric
- This is lowered to
Adapt the step size of the leapfrog integrator
integrator
:ssa = StepSizeAdaptor(δ, integrator)
- It uses Nesterov's dual averaging with
δ
as the target acceptance rate.
- It uses Nesterov's dual averaging with
Combine the two above naively:
NaiveHMCAdaptor(mma, ssa)
Combine the first two using Stan's windowed adaptation:
StanHMCAdaptor(mma, ssa)
The sample
functions
sample(
rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
h::Hamiltonian,
κ::HMCKernel,
θ::AbstractVector{<:AbstractFloat},
n_samples::Int;
adaptor::AbstractAdaptor=NoAdaptation(),
n_adapts::Int=min(div(n_samples, 10), 1_000),
drop_warmup=false,
verbose::Bool=true,
progress::Bool=false,
)
Draw n_samples
samples using the kernel κ
under the Hamiltonian system h
The randomness is controlled by
rng
.- If
rng
is not provided, the default random number generator (Random.default_rng()
) will be used.
- If
The initial point is given by
θ
.The adaptor is set by
adaptor
, for which the default is no adaptation.- It will perform
n_adapts
steps of adaptation, for which the default is1_000
or 10% ofn_samples
, whichever is lower.
- It will perform
drop_warmup
specifies whether to drop samples.verbose
controls the verbosity.progress
controls whether to show the progress meter or not.
Note that the function signature of the sample
function exported by AdvancedHMC.jl
differs from the sample
function used by Turing.jl
. We refer to the documentation of Turing.jl
for more details on the latter.
Full documentation of APIs in AdvancedHMC.jl
AdvancedHMC.AbstractIntegrator
— Typeabstract type AbstractIntegrator
Represents an integrator used to simulate the Hamiltonian system.
Implementation
A AbstractIntegrator
is expected to have the following implementations:
stat
(@ref)nom_step_size
(@ref)step_size
(@ref)
AdvancedHMC.AbstractMCMCKernel
— Typeabstract type AbstractMCMCKernel
Abstract type for HMC kernels.
AdvancedHMC.AbstractMetric
— Typeabstract type AbstractMetric
Abstract type for preconditioning metrics.
AdvancedHMC.AbstractTerminationCriterion
— Typeabstract type AbstractTerminationCriterion
Abstract type for termination criteria for Hamiltonian trajectories, e.g. no-U-turn and fixed number of leapfrog integration steps.
AdvancedHMC.AbstractTrajectorySampler
— TypeHow to sample a phase-point from the simulated trajectory.
AdvancedHMC.BinaryTree
— TypeA full binary tree trajectory with only necessary leaves and information stored.
AdvancedHMC.ClassicNoUTurn
— Typestruct ClassicNoUTurn{F<:AbstractFloat} <: AdvancedHMC.DynamicTerminationCriterion
Classic No-U-Turn criterion as described in Eq. (9) in [1].
Informally, this will terminate the trajectory expansion if continuing the simulation either forwards or backwards in time will decrease the distance between the left-most and right-most positions.
Fields
max_depth::Int64
Δ_max::AbstractFloat
References
- Hoffman, M. D., & Gelman, A. (2014). The No-U-Turn Sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(1), 1593-1623. (arXiv)
AdvancedHMC.DynamicTerminationCriterion
— Typeabstract type DynamicTerminationCriterion <: AdvancedHMC.AbstractTerminationCriterion
Abstract type for dynamic Hamiltonian trajectory termination criteria.
AdvancedHMC.EndPointTS
— TypeSamples the end-point of the trajectory.
AdvancedHMC.FixedIntegrationTime
— Typestruct FixedIntegrationTime{F<:AbstractFloat} <: AdvancedHMC.StaticTerminationCriterion
Standard HMC implementation with a fixed integration time.
Fields
λ::AbstractFloat
: Total length of the trajectory, i.e. takefloor(λ / integrator_step_size)
number of leapfrog steps.
References
- Neal, R. M. (2011). MCMC using Hamiltonian dynamics. Handbook of Markov chain Monte Carlo, 2(11), 2. (arXiv)
AdvancedHMC.FixedNSteps
— Typestruct FixedNSteps <: AdvancedHMC.StaticTerminationCriterion
Static HMC with a fixed number of leapfrog steps.
Fields
L::Int64
: Number of steps to simulate, i.e. length of trajectory will beL + 1
.
References
- Neal, R. M. (2011). MCMC using Hamiltonian dynamics. Handbook of Markov chain Monte Carlo, 2(11), 2. (arXiv)
AdvancedHMC.FullMomentumRefreshment
— TypeCompletly resample new momentum.
AdvancedHMC.GeneralisedNoUTurn
— Typestruct GeneralisedNoUTurn{F<:AbstractFloat} <: AdvancedHMC.DynamicTerminationCriterion
Generalised No-U-Turn criterion as described in Section A.4.2 in [1].
Fields
max_depth::Int64
Δ_max::AbstractFloat
References
- Betancourt, M. (2017). A Conceptual Introduction to Hamiltonian Monte Carlo. arXiv preprint arXiv:1701.02434.
AdvancedHMC.GeneralizedLeapfrog
— Typestruct GeneralizedLeapfrog{T<:(Union{AbstractVector{var"#s20"}, var"#s20"} where var"#s20"<:AbstractFloat)} <: AdvancedHMC.AbstractLeapfrog{T<:(Union{AbstractVector{var"#s20"}, var"#s20"} where var"#s20"<:AbstractFloat)}
Generalized leapfrog integrator with fixed step size ϵ
.
Fields
ϵ::Union{AbstractVector{var"#s20"}, var"#s20"} where var"#s20"<:AbstractFloat
: Step size.n::Int64
References
- Girolami, Mark, and Ben Calderhead. "Riemann manifold Langevin and Hamiltonian Monte Carlo methods." Journal of the Royal Statistical Society Series B: Statistical Methodology 73, no. 2 (2011): 123-214.
AdvancedHMC.HMC
— TypeHMC(ϵ::Real, n_leapfrog::Int)
Hamiltonian Monte Carlo sampler with static trajectory.
Fields
n_leapfrog
: Number of leapfrog steps.integrator
: Choice of integrator, specified either using aSymbol
orAbstractIntegrator
metric
: Choice of initial metric;Symbol
means it is automatically initialised. The metric type will be preserved during automatic initialisation and adaption.
AdvancedHMC.HMCDA
— TypeHMCDA(δ::Real, λ::Real, integrator = :leapfrog, metric = :diagonal)
Hamiltonian Monte Carlo sampler with Dual Averaging algorithm.
Fields
δ
: Target acceptance rate for dual averaging.λ
: Target leapfrog length.integrator
: Choice of integrator, specified either using aSymbol
orAbstractIntegrator
metric
: Choice of initial metric;Symbol
means it is automatically initialised. The metric type will be preserved during automatic initialisation and adaption.
Notes
For more information, please view the following paper (arXiv link):
- Hoffman, Matthew D., and Andrew Gelman. "The No-U-turn sampler: adaptively setting path lengths in Hamiltonian Monte Carlo." Journal of Machine Learning Research 15, no. 1 (2014): 1593-1623.
AdvancedHMC.HMCProgressCallback
— TypeHMCProgressCallback
A callback to be used with AbstractMCMC.jl's interface, replicating the logging behavior of the non-AbstractMCMC sample
.
Fields
pm
:Progress
meter from ProgressMeters.jl, ornothing
.verbose
: Ifpm === nothing
and this istrue
some information will be logged upon completion of adaptation.num_divergent_transitions
: Number of divergent transitions.num_divergent_transitions_during_adaption
AdvancedHMC.HMCSampler
— TypeHMCSampler
An AbstractMCMC.AbstractSampler
for kernels in AdvancedHMC.jl.
Fields
metric
: Choice of initial metricAbstractMetric
. The metric type will be preserved during adaption.adaptor
:AdvancedHMC.Adaptation.AbstractAdaptor
.
Notes
Note that all the fields have the prefix initial_
to indicate that these will not necessarily correspond to the kernel
, metric
, and adaptor
after sampling.
To access the updated fields, use the resulting HMCState
.
AdvancedHMC.HMCState
— TypeHMCState
Represents the state of a HMCSampler
.
Fields
i
: Index of current iteration.transition
: CurrentTransition
.metric
: CurrentAbstractMetric
, possibly adapted.κ
: CurrentAbstractMCMCKernel
.adaptor
: CurrentAbstractAdaptor
.
AdvancedHMC.JitteredLeapfrog
— Typestruct JitteredLeapfrog{FT<:AbstractFloat, T<:Union{AbstractArray{FT<:AbstractFloat, 1}, FT<:AbstractFloat}} <: AdvancedHMC.AbstractLeapfrog{T<:Union{AbstractArray{FT<:AbstractFloat, 1}, FT<:AbstractFloat}}
Leapfrog integrator with randomly "jittered" step size ϵ
for every trajectory.
Fields
ϵ0::Union{AbstractVector{FT}, FT} where FT<:AbstractFloat
: Nominal (non-jittered) step size.jitter::AbstractFloat
: The proportion of the nominal step sizeϵ0
that may be added or subtracted.ϵ::Union{AbstractVector{FT}, FT} where FT<:AbstractFloat
: Current (jittered) step size.
Description
This is the same as LeapFrog
(@ref) but with a "jittered" step size. This means that at the beginning of each trajectory we sample a step size ϵ
by adding or subtracting from the nominal/base step size ϵ0
some random proportion of ϵ0
, with the proportion specified by jitter
, i.e. ϵ = ϵ0 - jitter * ϵ0 * rand()
. p Jittering might help alleviate issues related to poor interactions with a fixed step size:
- In regions with high "curvature" the current choice of step size might mean over-shoot leading to almost all steps being rejected. Randomly sampling the step size at the beginning of the trajectories can therefore increase the probability of escaping such high-curvature regions.
- Exact periodicity of the simulated trajectories might occur, i.e. you might be so unlucky as to simulate the trajectory forwards in time
L ϵ
and ending up at the same point (which results in non-ergodicity; see Section 3.2 in [1]). If momentum is refreshed before each trajectory, then this should not happen exactly but it can still be an issue in practice. Randomly choosing the step-sizeϵ
might help alleviate such problems.
References
- Neal, R. M. (2011). MCMC using Hamiltonian dynamics. Handbook of Markov chain Monte Carlo, 2(11), 2. (arXiv)
AdvancedHMC.Leapfrog
— Typestruct Leapfrog{T<:(Union{AbstractVector{var"#s29"}, var"#s29"} where var"#s29"<:AbstractFloat)} <: AdvancedHMC.AbstractLeapfrog{T<:(Union{AbstractVector{var"#s29"}, var"#s29"} where var"#s29"<:AbstractFloat)}
Leapfrog integrator with fixed step size ϵ
.
Fields
ϵ::Union{AbstractVector{var"#s29"}, var"#s29"} where var"#s29"<:AbstractFloat
: Step size.
AdvancedHMC.MultinomialTS
— Methodstruct MultinomialTS{F<:AbstractFloat, P<:AdvancedHMC.PhasePoint} <: AdvancedHMC.AbstractTrajectorySampler
Multinomial sampler for a trajectory consisting only a leaf node.
- tree weight is the (unnormalised) energy of the leaf.
AdvancedHMC.MultinomialTS
— Methodstruct MultinomialTS{F<:AbstractFloat, P<:AdvancedHMC.PhasePoint} <: AdvancedHMC.AbstractTrajectorySampler
Multinomial sampler for the starting single leaf tree. (Log) weights for leaf nodes are their (unnormalised) Hamiltonian energies.
Ref: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/hmc/nuts/base_nuts.hpp#L226
AdvancedHMC.MultinomialTS
— Typestruct MultinomialTS{F<:AbstractFloat, P<:AdvancedHMC.PhasePoint} <: AdvancedHMC.AbstractTrajectorySampler
Multinomial trajectory sampler carried during the building of the tree. It contains the weight of the tree, defined as the total probabilities of the leaves.
Fields
zcand::AdvancedHMC.PhasePoint
: Sampled candidatePhasePoint
.ℓw::AbstractFloat
: Total energy for the given tree, i.e. the sum of energies of all leaves.
AdvancedHMC.NUTS
— TypeNUTS(δ::Real; max_depth::Int=10, Δ_max::Real=1000, integrator = :leapfrog, metric = :diagonal)
No-U-Turn Sampler (NUTS) sampler.
Fields
δ
: Target acceptance rate for dual averaging.max_depth
: Maximum doubling tree depth.Δ_max
: Maximum divergence during doubling tree.integrator
: Choice of integrator, specified either using aSymbol
orAbstractIntegrator
metric
: Choice of initial metric;Symbol
means it is automatically initialised. The metric type will be preserved during automatic initialisation and adaption.
AdvancedHMC.PartialMomentumRefreshment
— Typestruct PartialMomentumRefreshment{F<:AbstractFloat} <: AdvancedHMC.AbstractMomentumRefreshment
Partial momentum refreshment with refresh rate α
.
Fields
α::AbstractFloat
See equation (5.19) [1]
r' = α⋅r + sqrt(1-α²)⋅G
where r is the momentum and G is a Gaussian random variable.
References
- Neal, Radford M. "MCMC using Hamiltonian dynamics." Handbook of markov chain monte carlo 2.11 (2011): 2.
AdvancedHMC.SliceTS
— Methodstruct SliceTS{F<:AbstractFloat, P<:AdvancedHMC.PhasePoint} <: AdvancedHMC.AbstractTrajectorySampler
Slice sampler for the starting single leaf tree. Slice variable is initialized.
AdvancedHMC.SliceTS
— Methodstruct SliceTS{F<:AbstractFloat, P<:AdvancedHMC.PhasePoint} <: AdvancedHMC.AbstractTrajectorySampler
Create a slice sampler for a single leaf tree:
- the slice variable is copied from the passed-in sampler
s
and - the number of acceptable candicates is computed by comparing the slice variable against the current energy.
AdvancedHMC.SliceTS
— Typestruct SliceTS{F<:AbstractFloat, P<:AdvancedHMC.PhasePoint} <: AdvancedHMC.AbstractTrajectorySampler
Trajectory slice sampler carried during the building of the tree. It contains the slice variable and the number of acceptable condidates in the tree.
Fields
zcand::AdvancedHMC.PhasePoint
: Sampled candidatePhasePoint
.ℓu::AbstractFloat
: Slice variable in log-space.n::Int64
: Number of acceptable candidates, i.e. those with probability larger than slice variableu
.
AdvancedHMC.StaticTerminationCriterion
— Typeabstract type StaticTerminationCriterion <: AdvancedHMC.AbstractTerminationCriterion
Abstract type for a fixed number of leapfrog integration steps.
AdvancedHMC.StrictGeneralisedNoUTurn
— Typestruct StrictGeneralisedNoUTurn{F<:AbstractFloat} <: AdvancedHMC.DynamicTerminationCriterion
Generalised No-U-Turn criterion as described in Section A.4.2 in [1] with added U-turn check as described in [2].
Fields
max_depth::Int64
Δ_max::AbstractFloat
References
- Betancourt, M. (2017). A Conceptual Introduction to Hamiltonian Monte Carlo. arXiv preprint arXiv:1701.02434.
- https://github.com/stan-dev/stan/pull/2800
AdvancedHMC.TemperedLeapfrog
— Typestruct TemperedLeapfrog{FT<:AbstractFloat, T<:Union{AbstractArray{FT<:AbstractFloat, 1}, FT<:AbstractFloat}} <: AdvancedHMC.AbstractLeapfrog{T<:Union{AbstractArray{FT<:AbstractFloat, 1}, FT<:AbstractFloat}}
Tempered leapfrog integrator with fixed step size ϵ
and "temperature" α
.
Fields
ϵ::Union{AbstractVector{FT}, FT} where FT<:AbstractFloat
: Step size.α::AbstractFloat
: Temperature parameter.
Description
Tempering can potentially allow greater exploration of the posterior, e.g. in a multi-modal posterior jumps between the modes can be more likely to occur.
AdvancedHMC.Termination
— MethodTermination(s, nt, H0, H′)
Check termination of a Hamiltonian trajectory.
AdvancedHMC.Termination
— MethodTermination(s, nt, H0, H′)
Check termination of a Hamiltonian trajectory.
AdvancedHMC.Termination
— TypeTermination
Termination reasons
dynamic
: due to stoping criterianumerical
: due to large energy deviation from starting (possibly numerical errors)
AdvancedHMC.Trajectory
— Typestruct Trajectory{TS<:AdvancedHMC.AbstractTrajectorySampler, I<:AdvancedHMC.AbstractIntegrator, TC<:AdvancedHMC.AbstractTerminationCriterion}
Numerically simulated Hamiltonian trajectories.
AdvancedHMC.Transition
— Typestruct Transition{P<:AdvancedHMC.PhasePoint, NT<:NamedTuple}
A transition that contains the phase point and other statistics of the transition.
Fields
z::AdvancedHMC.PhasePoint
: Phase-point for the transition.stat::NamedTuple
: Statistics related to the transition, e.g. energy.
AdvancedHMC.Adaptation.AbstractAdaptor
— Typeabstract type AbstractAdaptor
Abstract type for HMC adaptors.
AdvancedHMC.Adaptation.NesterovDualAveraging
— TypeAn implementation of the Nesterov dual averaging algorithm to tune step size.
References
Hoffman, M. D., & Gelman, A. (2014). The No-U-Turn Sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(1), 1593-1623. Nesterov, Y. (2009). Primal-dual subgradient methods for convex problems. Mathematical programming, 120(1), 221-259.