Hamiltonian Monte Carlo 法

Julia と Turing エコシステムを用いて

Process
Sampling
Julia
MCMC
Author

司馬 博文

Published

7/03/2024

概要
Julia に存在する HMC 関連のパッケージの実装と,その使い方をまとめる.

Julia による MCMC パッケージの概観は次の稿も参照:

1 AdvancedHMC.jl

Cauchy 分布に HMC を適用するとぶっ壊れる!?

using AdvancedHMC, ForwardDiff
using LogDensityProblems
using LinearAlgebra
using Plots

struct LogTargetDensityCauchy
    loc::Float64
    scale::Float64
end

# Define the target distribution (1D Cauchy) using the `LogDensityProblem` interface

LogDensityProblems.logdensity(p::LogTargetDensityCauchy, θ) = -log(π) - log(p.scale) - log(1 + ((θ[1] - p.loc)/p.scale)^2)
LogDensityProblems.dimension(p::LogTargetDensityCauchy) = 1
LogDensityProblems.capabilities(::Type{LogTargetDensityCauchy}) = LogDensityProblems.LogDensityOrder{0}()

function HMC_sample(initial_θ)

    # Choose initial parameter value for 1D
    initial_θ = [initial_θ]

    # Define the Cauchy distribution with location and scale
    loc, scale = 0.0, 1.0
    ℓπ = LogTargetDensityCauchy(loc, scale)

    # Set the number of samples to draw and warmup iterations
    n_samples, n_adapts = 2_000, 1

    # Define a Hamiltonian system
    metric = DiagEuclideanMetric(1)
    hamiltonian = Hamiltonian(metric, ℓπ, ForwardDiff)

    # Define a leapfrog solver, with the initial step size chosen heuristically
    initial_ϵ = find_good_stepsize(hamiltonian, initial_θ)
    integrator = Leapfrog(initial_ϵ)

    # Define an HMC sampler with the following components
    #   - multinomial sampling scheme,
    #   - generalised No-U-Turn criteria, and
    #   - windowed adaption for step-size and diagonal mass matrix
    kernel = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
    adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))

    # Run the sampler to draw samples from the specified Cauchy distribution, where
    #   - `samples` will store the samples
    #   - `stats` will store diagnostic statistics for each sample
    samples, stats = sample(hamiltonian, kernel, initial_θ, n_samples, adaptor, n_adapts; progress=true)

    # Print the results
    sample_values = [s[1] for s in samples]

    p = plot(1:length(samples), sample_values,
                label="HMC trajectory",
                title="1D HMC Sampler (Cauchy distribution)",
                xlabel="t",
                ylabel="X",
                linewidth=2,
                marker=:circle,
                markersize=2,
                markeralpha=0.6,
                color="#78C2AD")
end

HMC_sample(50.0)
┌ Warning: ProgressMeter by default refresh meters with additional information in IJulia via `IJulia.clear_output`, which clears all outputs in the cell. 
│  - To prevent this behaviour, do `ProgressMeter.ijulia_behavior(:append)`. 
│  - To disable this warning message, do `ProgressMeter.ijulia_behavior(:clear)`.
└ @ ProgressMeter ~/.julia/packages/ProgressMeter/gA66f/src/ProgressMeter.jl:590
Sampling 100%|███████████████████████████████| Time: 0:00:00
  iterations:                                   2000
  ratio_divergent_transitions:                  0.02
  ratio_divergent_transitions_during_adaption:  0.0
  n_steps:                                      3
  is_accept:                                    true
  acceptance_rate:                              4.073791219423035e-7
  log_density:                                  -10.758856435114664
  hamiltonian_energy:                           13.458401323991495
  hamiltonian_energy_error:                     0.0
  max_hamiltonian_energy_error:                 811.5344241268535
  tree_depth:                                   2
  numerical_error:                              false
  step_size:                                    324.39138009045865
  nom_step_size:                                324.39138009045865
  is_adapt:                                     false
  mass_matrix:                                  DiagEuclideanMetric([1.0])
┌ Info: Finished 2000 sampling steps for 1 chains in 1.282884958 (s)
│   h = Hamiltonian(metric=DiagEuclideanMetric([1.0]), kinetic=GaussianKinetic())
│   κ = HMCKernel{AdvancedHMC.FullMomentumRefreshment, Trajectory{MultinomialTS, Leapfrog{Float64}, GeneralisedNoUTurn{Float64}}}(AdvancedHMC.FullMomentumRefreshment(), Trajectory{MultinomialTS}(integrator=Leapfrog(ϵ=324.0), tc=GeneralisedNoUTurn{Float64}(10, 1000.0)))
│   EBFMI_est = 0.30445073109268245
└   average_acceptance_rate = 0.22613549485908152

スタート地点を 50 にした場合

HMC_sample(0.0)
┌ Info: Finished 2000 sampling steps for 1 chains in 0.044786417 (s)
│   h = Hamiltonian(metric=DiagEuclideanMetric([1.0]), kinetic=GaussianKinetic())
│   κ = HMCKernel{AdvancedHMC.FullMomentumRefreshment, Trajectory{MultinomialTS, Leapfrog{Float64}, GeneralisedNoUTurn{Float64}}}(AdvancedHMC.FullMomentumRefreshment(), Trajectory{MultinomialTS}(integrator=Leapfrog(ϵ=14.7), tc=GeneralisedNoUTurn{Float64}(10, 1000.0)))
│   EBFMI_est = 2.018632634147994
└   average_acceptance_rate = 0.00031706816850295914

スタート地点を 0 にした場合

HMC_sample(500.0)
┌ Info: Finished 2000 sampling steps for 1 chains in 0.056288833 (s)
│   h = Hamiltonian(metric=DiagEuclideanMetric([1.0]), kinetic=GaussianKinetic())
│   κ = HMCKernel{AdvancedHMC.FullMomentumRefreshment, Trajectory{MultinomialTS, Leapfrog{Float64}, GeneralisedNoUTurn{Float64}}}(AdvancedHMC.FullMomentumRefreshment(), Trajectory{MultinomialTS}(integrator=Leapfrog(ϵ=1670.0), tc=GeneralisedNoUTurn{Float64}(10, 1000.0)))
│   EBFMI_est = 0.2517834551566722
└   average_acceptance_rate = 0.20632568153387063

スタート地点を 500 にした場合

HMC_sample(100.0)
┌ Info: Finished 2000 sampling steps for 1 chains in 0.043402167 (s)
│   h = Hamiltonian(metric=DiagEuclideanMetric([1.0]), kinetic=GaussianKinetic())
│   κ = HMCKernel{AdvancedHMC.FullMomentumRefreshment, Trajectory{MultinomialTS, Leapfrog{Float64}, GeneralisedNoUTurn{Float64}}}(AdvancedHMC.FullMomentumRefreshment(), Trajectory{MultinomialTS}(integrator=Leapfrog(ϵ=956.0), tc=GeneralisedNoUTurn{Float64}(10, 1000.0)))
│   EBFMI_est = 1.998017004885298
└   average_acceptance_rate = 6.0828891145860786e-9

スタート地点を 100 にした場合