turinglang / advancedhmc.jl Goto Github PK

@yebai It seems that I don't have access to edit the description of the repo after the transfer. Maybe we want something like "Efficient building blocks for modern HMC samplers"?

As this package mainly aim to provide robust building blocks for HMC (instead of providing a robust sample interface directly), may be we want to re-name it? @yebai @trappmartin

Currently initialised M inverse need to be passed in. We probably want one without them and initialise them as identity.

TuringLang/Turing.jl#691

Neal, R. M. (2011). MCMC using Hamiltonian dynamics. Handbook of Markov chain Monte Carlo, 2(11), 2. (pdf)

Betancourt, M. (2017). A Conceptual Introduction to Hamiltonian Monte Carlo. arXiv preprint arXiv:1701.02434.

Girolami, M., & Calderhead, B. (2011). Riemann manifold Langevin and Hamiltonian Monte Carlo methods. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 73(2), 123-214. (link)

Betancourt, M. J., Byrne, S., & Girolami, M. (2014). Optimizing the integrator step size for Hamiltonian Monte Carlo. arXiv preprint arXiv:1411.6669.

Betancourt, M. (2016). Identifying the optimal integration time in Hamiltonian Monte Carlo. arXiv preprint arXiv:1601.00225.

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.

Currently during test the verbosity is disable for clean test output thus show() functions are not tests at all. We may want to include tests for the explicitly.

We currently have a method for each AbstractMetric which enables us to make another instance of the metric of the following form:

metric = ...
new_metric_of_same_type = metric(inverse_mass_matrix)

Why do we need such a method?

The tag name "0.1.1" is not of the appropriate SemVer form (vX.Y.Z).
cc: @xukai92

Hi @xukai92, I'd like to start a repository heavily based on your code, but with some modifications that lead to incompatibility. I don't see a license file. Is there one you prefer? I'd like to use MIT or BSD if possible. Thanks!

• Consider introduce the PhasePoint type #17 #50
• Update gradient function to return both value and gradient #56
• RFC find_good_eps #27
• More robust and modular way of detecting divergence #16
• Return more information for each step #59
• Implement Base.show for all user side types #36
• RFC Adapter types #28
• Instantiation of new AbstractMetrics #45
• Simplifying Metric types by using AbstractPDMat #34
• Adding noise to step size? #10

We should try to have a modular interface for different momentum refreshment methods (e.g. partial refreshments in Neal 2011)

The NUTS sampler computes Cholesky decomposition on the fly for each Hamiltonian step. This is likely one of the performance bottlenecks for NUTS since Cholesky decomposition is a costly operation and should be computed only once adaption is done.

https://github.com/TuringLang/Turing.jl/blob/5c94eebd5a8a0fe2471ac3fa08341d3e2d054186/src/inference/support/hmc_core.jl#L107

Related: https://github.com/TuringLang/Turing.jl/issues/559

Copied from @yebai comment in #23 (comment)

    r = rand_momentum(rng, h)
    H = hamiltonian_energy(h, θ, r)

out of find_good_eps and wraping

        θ′, r′, _is_valid = step(Leapfrog(ϵ′), h, θ′, r′)
        H_new = _is_valid ? hamiltonian_energy(h, θ′, r′) : Inf

into a function

function A(ϵ, h, θ, r) 
        θ′, r′, _is_valid = step(Leapfrog(ϵ′), h, θ′, r′)
        return H_new = _is_valid ? hamiltonian_energy(h, θ′, r′) : Inf
end

Then, we can drop the dependency on AdvancedHMC from this function: find_good_eps(h::Hamiltonian, θ::AbstractVector{T}, A::Function; max_n_iters::Int=100), and move it into adaption/stepsize.jl.

I just came across the error message at:

This error message should be more informative so that the user gets an idea why he gets it, i.e. in my case I suppose the model I'm looking at is misspecified. It would be good to know what is Inf.

We should consider adding the following: https://arxiv.org/abs/1905.11916

Maybe we should make the sample function not return those samples by default (and support an optional keyword for it)?

Discussions in #5 (comment)

@mohamed82008 I reverted some functions about pre-allocating memory so that only intermediate variables are pre-allocated and all functions remains immutable. This is to avoid un-expected behaviour due to mutability. We probably want to implement mutable version of functions to pre-allocate memory for return variables explicitly later. Does it sound good to you?

Sorry, I missed this question. I think it would be nice to support mutable and immutable array types, e.g. StaticArrays.SArray and StaticArrays.MArray or Vector. I saw Chris mention somewhere that DiffEq does this by defining 2 versions of the function, one for mutable arrays and one for immutable. We might want to explore something like this in a future PR. The mutating version takes a first argument the pre-allocated output vector whereas the non-mutating version doesn't. Supporting StaticArrays would be an interesting direction to pursue to gain some speedup for parameter vectors of size <1000 or so. Small arrays is the sweet spot of StaticArrays putting ordinary arrays to shame in many cases.

At the moment, we are using the following notations:

• logπ - log target distribution
• ∂logπ∂θ - gradient
• prop - transitioning kernel
• θ_init - initial value for sampling states / parameters

I think this can be improved by

• ∂logπ∂θ ==> ∇logπ
• prop ==> κ (kappa) or τ (tau) - MH transitioning kernel
• θ - parameters, r - momentum
• θ_init ==> θ₀ (theta_{0}), r_init ==> r₀ (r_{0})
• z := (θ, r) - (position, momentum) pair

Any other notations missing?

I was speaking with Kai about this and agreed it would be good to be able to identify divergence of a particular transition (Sec. 6.2 in here ) allowing us to identify pathologies which may bias the posterior sample.

Basically those appear in the READM.md for now.

• metric
• Hamiltonian
• Leapfrog
• Trajectory
• Adaptor

Approximation error of leapfrog integration (i.e. accumulated Hamiltonian energy error) can sometimes explode, for example when the curvature of the current region is very high. This type of approximation error is sometimes called divergence [1] since it shifts a leapfrog simulation away from the correct solution.

In Turing, this type of errors is currently caught a relatively ad-hoc function called is_valid ,

is_valid can catch cases where one or more elements of the parameter vector is either nan or inf. This has several drawbacks

• it's not general enough for catching all leapfrog approximation errors, e.g. when the parameter vector is valid but Hamiltonian energy is invalid.
• it may be delayed because numerical errors can appear in Hamiltonian energy earlier than in parameter vectors
• it's hard to propagate the exception around, i.e. at the moment we use a heuristic to find the previous valid point before approximation/numerical error happens

  • AdvancedHMC.jl/src/integrator.jl

    Line 26 in 734c0fa

    !_is_valid && return θ, r, false

  • AdvancedHMC.jl/src/integrator.jl

    Line 45 in 734c0fa

    r, _ = lf_momentum(-lf.ϵ / 2, h, θ, r)

Therefore, we might want to refactor the current code a bit for a more robust mechanism for handling leapfrog approximation errors. Perhaps we can learn from the DynamicHMC implementation:

https://github.com/tpapp/DynamicHMC.jl/blob/master/src/buildingblocks.jl#L168

[1]: Betancourt, M. (2017). A Conceptual Introduction to Hamiltonian Monte Carlo. arXiv preprint arXiv:1701.02434.

1. avoiding running unnecessary pass of the pdf function of the target distribution

• the current design separates the pdf and its gradient and there are some duplicated evaluation of the pdf

2. update is doing unnecessary decomposition of matrix (#24)

• need to have a flag to indicate if M inverse is update by adaptor

Both this package and Distributions.jl define a sample function, as opposed to this package extending Distributions.jl's sample function with new methods. This causes a name clash which, given that these packages will likely be used to together a lot, is annoying. Should we maybe extend Distributions.jl's sample rather than define our own?

Moved from #69

#################################
### Hong's abstraction starts ###
#################################

###
### Create a `Termination` type for each `Trajectory` type, e.g. HMC, NUTS etc.
### Merge all `Trajectory` types, and make `transition` dispatch on `Termination`,
### such that we can overload `transition` for different HMC samplers.
### NOTE:  stopping creteria, max_depth::Int, Δ_max::AbstractFloat, n_steps, λ
###

"""
Abstract type for termination.
"""
abstract type AbstractTermination end

# Termination type for HMC and HMCDA
struct StaticTermination{D<:AbstractFloat} <: AbstractTermination
    n_steps :: Int
    Δ_max :: D
end

# NoUTurnTermination
struct NoUTurnTermination{D<:AbstractFloat} <: AbstractTermination
    max_depth :: Int
    Δ_max :: D
    # TODO: add other necessary fields for No-U-Turn stopping creteria.
end

struct Trajectory{I<:AbstractIntegrator} <: AbstractTrajectory{I}
    integrator :: I
    n_steps :: Int # Counter for total leapfrog steps already applied.
    Δ :: AbstractFloat # Current hamiltonian energy minus starting hamiltonian energy
    # TODO: replace all ``*Trajectory` types with `Trajectory`.
    # TODO: add turn statistic, divergent statistic, proposal statistic
end

isterminated(
    x::StaticTermination,
    τ::Trajectory
) = τ.n_steps >= x.n_steps || τ.Δ >= x.Δ_max

# Combine trajectories, e.g. those created by the build_tree algorithm.
#  NOTE: combine proposal (via slice/multinomial sampling), combine turn statistic,
#       and combine divergent statistic.
combine_trajectory(τ′::Trajectory, τ′′::Trajectory) = nothing # To-be-implemented.

## TODO: move slice variable `logu` into `Trajectory`?
combine_proposal(τ′::Trajectory, τ′′::Trajectory) = nothing # To-be-implemented.
combine_turn(τ′::Trajectory, τ′′::Trajectory) = nothing # To-be-implemented.
combine_divergence(τ′::Trajectory, τ′′::Trajectory) = nothing # To-be-implemented.

transition(
    τ::Trajectory{I},
    h::Hamiltonian,
    z::PhasePoint,
    t::T
) where {I<:AbstractIntegrator,T<:AbstractTermination} = nothing

###############################
### Hong's abstraction ends ###
###############################

Time used for collecting 2 million samples

• Turing.NUTS: 3318.293540 seconds (33.35 G allocations: 1.299 TiB, 29.35% gc time)
• Turing.DynamicNUTS: 848.741643 seconds (7.86 G allocations: 251.076 GiB, 32.13% gc time)

using the following

using DynamicHMC, Turing, Test

@model gdemo(x, y) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt(s))
  x ~ Normal(m, sqrt(s))
  y ~ Normal(m, sqrt(s))
  return s, m
end

mf = gdemo(1.5, 2.0)

@time chn1 = sample(mf, DynamicNUTS(2000000));

@time chn2 = sample(mf, Turing.NUTS(2000000, 0.65));

• long simulation and check average Hamiltonian energy
• Eq 2.11 from https://arxiv.org/abs/1206.1901

The REQUIRE file could not be found.
cc: @yebai

migrated from TuringLang/Turing.jl#571

As mentioned on page 59 in [1], multinomial sampling instead of slice sampling may provide an improvement in overall performance.

[1] https://arxiv.org/abs/1701.02434

@trappmartin pointed that we should aim to support GPU via CuArrays. I believe this is possible and a good idea.

• \alpha
• H
• Leapfrog step size
• Number of leapfrog steps
• Number of value evals
• Number of grad evals

migrated from TuringLang/Turing.jl#268

As requested by @jonasmac16, it would be nice to have this implemented; see Section 5.7 of Neal, 2012.

After the RF this might be an interesting extension of the current code, especially RMHMC with SoftAbs Metric.

https://arxiv.org/pdf/1212.4693.pdf

Return value doesn't introduce any cost and it would save some computation for place where value is needed right after gradient evaluation.

The following 3 types in src/metric.jl:

• AdvancedHMC.jl/src/metric.jl

  Lines 16 to 24 in a0af7cc

  	struct DiagEuclideanMetric{T<:Real,A<:AbstractVector{T}} <: AbstractMetric{T}
  	dim :: Int
  	# Diagnal of the inverse of the mass matrix
  	M⁻¹ :: A
  	# Sqare root of the inverse of the mass matrix
  	sqrtM⁻¹ :: A
  	# Pre-allocation for intermediate variables
  	_temp :: A
  	end

• AdvancedHMC.jl/src/metric.jl

  Lines 41 to 49 in a0af7cc

  	struct DenseEuclideanMetric{T<:Real,AV<:AbstractVector{T},AM<:AbstractMatrix{T}} <: AbstractMetric{T}
  	dim :: Int
  	# Inverse of the mass matrix
  	M⁻¹ :: AM
  	# U of the Cholesky decomposition of the mass matrix
  	cholM⁻¹ :: UpperTriangular{T,AM}
  	# Pre-allocation for intermediate variables
  	_temp :: AV
  	end

• AdvancedHMC.jl/src/metric.jl

  Lines 3 to 5 in a0af7cc

  	struct UnitEuclideanMetric{T<:Real} <: AbstractMetric{T}
  	dim :: Int
  	end

can be merged into one type using AbstractPDMat

struct EuclideanMetric{T<:Real,A<:AbstractPDMat{T}} <: AbstractMetric{T}
    M⁻¹     ::  A # contains dim, M⁻¹ and its Cholesky decomposition. 
    # Pre-allocation for intermediate variables
    _temp   ::  A
end

Reference: https://github.com/JuliaStats/PDMats.jl

This might be interesting to add at some point.

https://arxiv.org/pdf/1705.08510.pdf

https://arxiv.org/pdf/1711.09268.pdf

The REQUIRE file could not be found.
cc: @xukai92

Will do after #31 gets merged.

Copied from @yebai's comments in #23 (comment)

can we

• move this file's content into adaptation/Adaptation.jl
• remove the file src/adaptation.jl
• and include adaptation/Adaptation.jl from AdvancedHMC directly?

Copied from @yebai's comments in #23 (comment)

consider merging the following three functions

• getss(fss::FixedStepSize)
• getss(da::DualAveraging)
• getss(mssa::ManualSSAdapter)

into

• getϵ(a:: StepSizeAdapter) = a.ϵ, or more generally
• getϵ(a:: AbstractAdapter) = a.ϵ

Copied from @yebai's comments in #23 (comment)

consider removing ManualSSAdapter and related code, since we can always use FixedStepSize

Copied from @yebai's comments in #23 (comment)

can we make DualAveraging an immutable type, and let

• adapt_stepsize!, adapt!
• finalize!

always return a new DualAveraging instance. This can improve both readability and performance since immutable types are allocated on stacks. If so, we can also remove

• reset!

and merge

• adapt!(da::DualAveraging, θ::AbstractVector{<:AbstractFloat}, α::AbstractFloat)
• adapt_stepsize!(da::DualAveraging, α::AbstractFloat)

into

-adapt(d::AbstractAdaptor, α::AbstractFloat=1.0, θ::AbstractVector{<:AbstractFloat}=zeros(0))

We can probably do similar things for Metric adaptors.

Copied from @yebai's comments in #23 (comment)

Consider

• removing `AbstractCompositeAdapter, and
• making NaiveCompAdaptor and ThreePhaseAdaptor direct subtypes of AbstractAdapter

Since AbstractCompositeAdapter is only used to define

• getM⁻¹(ca::AbstractCompositeAdapter)
• getss(ca::AbstractCompositeAdapter)
• update(h::Hamiltonian, p::AbstractProposal, a::Adaptation.AbstractCompositeAdapter)

These functions can be defined for AbstractAdaptor directly:

• getM⁻¹(a::AbstractAdapter) = a.var
• getss(a::AbstractAdapter) = a.ϵ
• update(h::Hamiltonian, p::AbstractProposal, a::Adaptation.AbstractCompositeAdapter)

Then we can overload these functions for subtypes of AbstractAdaptor

• getM⁻¹(a:: ThreePhaseAdapter) = getM⁻¹(a.pc)
• getss(a:: ThreePhaseAdapter) = getss(a.ssa)
• getM⁻¹(dpc::UnitPreConditioner) = nothing (perhaps return 1.0 instead of nothing)
• getM⁻¹(dpc::DensePreConditioner) = dpc.covar

Related: TuringLang/Turing.jl#641