Juliaを使ってみる(18/22)：ハミルトニアン・モンテカルロ法 - 「大人の教養・知識・気付き」を伸ばすブログ

今回やりたいこと

　ハミルトニアン・モンテカルロ法をJuliaでやってみる。

Juliaで作って学ぶベイズ統計学 (KS情報科学専門書)

作者:須山敦志
講談社

Amazon

を参照する。

前回

power-of-awareness.com

今回やりたいこと
前回
1.　ハミルトニアン・モンテカルロ法
次回

1.　ハミルトニアン・モンテカルロ法

　ハミルトニアン・モンテカルロ法はマルコフ連鎖モンテカルロ法の1つで確率的プログラミング言語でもデファクトスタンダードとされている近似推論手法である。近似のための分布を仮定する必要が無く、理論上、厳密な事後分布からの標本を得ることができる。以下ではプログラミングするに当たっての最低限の概要を説明する。

1.1　マルコフ連鎖モンテカルロ法

　サンプル候補の生成に当たり、与えられたデータの情報を加味し、1つ1つのシミュレーションの情報を引き継ぐことができれば、シミュレーションの効率を向上し得ると考えられる。これらを加味できるのがマルコフ連鎖モンテカルロ法( $\mathrm{MCMC}$ )である。

マルコフ連鎖モンテカルロ法( $\mathrm{MCMC}$ )のアイディア

初期サンプル $\boldsymbol{\mu}_0$ を定める

初期サンプル $\boldsymbol{\mu}_0$ から遷移カーネルに基づき2つ目のサンプル $\boldsymbol{\mu}_1$ を確率的に選定する。

以降、既定の回数まで同様にサンプリングを続ける。

1.2　メトロポリス・ヘイスティング法

　上述したマルコフ連鎖モンテカルロ法のアイディアに基づくとき、遷移カーネルを適切に設定することで、各サンプルが事後分布に従うようにする。この遷移カーネルの設計を容易にする方法がメトロポリス・ヘイスティング法である。
　メトロポリス・ヘイスティング法では、まず何らかの分布(提案分布)を定めて事前のサンプル $\boldsymbol{\mu}_i$ から $q(\boldsymbol{\mu}_{\mathrm{tmp}}|\boldsymbol{\mu}_i)$ に従うサンプル候補 $\boldsymbol{\mu}_{\mathrm{tmp}}$ を生成し、これを新しいサンプルとして受容して $\boldsymbol{\mu}_{i+1}=\boldsymbol{\mu}_{\mathrm{tmp}}$ とするか棄却して $\boldsymbol{\mu}_{i+1}=\boldsymbol{\mu}_i$ とするかを決定する。

提案分布 $q$ を定める。

非正規化後事後分布 $\tilde{p}$ を関数として記述する。

初期サンプル $\boldsymbol{\mu}_0$ およびサンプリング回数 $n$ を定める。

$i=0$ とする。

$\boldsymbol{\mu}_i$ に基づき $q(\boldsymbol{\mu}_{\mathrm{tmp}}|\boldsymbol{\mu}_i)$ に従うサンプル候補 $\boldsymbol{\mu}_{\mathrm{tmp}}$ を生成する。

比率
$\begin{aligned}r=\displaystyle{\frac{\tilde{p}(\boldsymbol{\mu}_{\mathrm{tmp}})q(\boldsymbol{\mu}_i|\boldsymbol{\mu}_{\mathrm{tmp}})}{\tilde{p(\boldsymbol{\mu}_i)q(\boldsymbol{\mu}_{\mathrm{tmp}}|\boldsymbol{\mu}_i)}}}\end{aligned}$
を計算する。

確率 $\min(1,r)$ で $\boldsymbol{\mu}_{i+1}=\boldsymbol{\mu}_{\mathrm{tmp}}$ とし、そうでない場合には $\boldsymbol{\mu}_{i+1}=\boldsymbol{\mu}_{i}$ とする。

$i\lt n$ ならば $i=i+1$ とした後に5.に戻る。そうでなければサンプリングを終了する。

1.3　ハミルトニアン・モンテカルロ法

　ハミルトニアン・モンテカルロ法は、メトロポリス・ヘイスティング法に基づいた手法で、提案分布 $q(\boldsymbol{\mu}_{\mathrm{tmp}}|\boldsymbol{\mu}_{i})$ に小球運動に着想を得た計算手法を取り入れたものである。具体的には位置 $\boldsymbol{\mu}_{i}$ にある小球に対して、ランダムに運動量 $\boldsymbol{p}$ を与えて小球をはじき、一定自国 $\varepsilon L$ 後に小球を止めてその位置を $\boldsymbol{\mu}_{\mathrm{tmp}}$ として記録する。ここで $\varepsilon\gt0$ は単位時間を、 $L\gt0$ は計算を行う回数に対応する。

運動量 $\boldsymbol{p}$ として多変量正規分布から $\boldsymbol{p}^{(0)}\sim\mathcal{N}(\boldsymbol{p}^{(0)}|\boldsymbol{0},\boldsymbol{I})$ する。

$\boldsymbol{\mu}^{(0)}=\boldsymbol{\mu}_{i-1}$ とし、運動量を
$\begin{aligned}\boldsymbol{p}^{(1)}=\boldsymbol{p}^{(0)}+\displaystyle{\frac{1}{2}\varepsilon\nabla\log \tilde{p}\left(\boldsymbol{\mu}^{(0)}\right)}\end{aligned}$
で更新する。

$j=1,2,\cdots,L-1$ に関して
$\begin{aligned}\boldsymbol{\mu}^{(j)}&=\boldsymbol{\mu}^{(j-1)}+\varepsilon\boldsymbol{p}^{(j)},\\\boldsymbol{p}^{(j+1)}&=\boldsymbol{p}^{(j)}+\varepsilon\nabla\log\tilde{p}\left(\boldsymbol{\mu}^{(j)}\right)\end{aligned}$

$\boldsymbol{\mu}_{\mathrm{tmp}},\boldsymbol{p}_{\mathrm{tmp}}$ を
$\begin{aligned}\boldsymbol{\mu}_{\mathrm{tmp}}&=\boldsymbol{\mu}^{(L-1)}+\varepsilon\boldsymbol{p}^{(L)},\\\boldsymbol{p}_{\mathrm{tmp}}&=\boldsymbol{p}^{(L)}+\varepsilon\nabla\log\tilde{p}\left(\boldsymbol{\mu}_{\mathrm{tmp}}\right)\end{aligned}$
で設定する。

1.3.1　実装

using Distributions, PyPlot, ForwardDiff, LinearAlgebra

# n次元単位行列
eye(n) = Diagonal{Float64}(I, n)

# グラフの諸設定を行う関数
function set_options(ax, xlabel, ylabel, title;
                      grid = true, gridy = false, legend = false)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    if grid
        if gridy
            ax.grid(axis = "y")
        else
            ax.grid()
        end
    end
    legend && ax.legend()
end

# ガウス分布を提案分布としたMetropolis-Hasting法
function GaussianMH(log_p_tilde, μ₀;
        maxiter::Int64, σ::Float64)
    # サンプルを格納する配列
    D = length(μ₀)
    μ_samples = Array{typeof(μ₀[1]), 2}(undef, D, maxiter)
    
    μ_samples[:,1] = μ₀
    
    # 受容サンプル数
    num_accepted = 1
    
    for i in 2:maxiter
        #
        μ_tmp = rand(MvNormal(μ_samples[:, i-1], σ*eye(D)))
        
        # 比率 r
        log_r = (log_p_tilde(μ_tmp) +
                 logpdf(MvNormal(μ_tmp, σ), μ_samples[:, i-1])) -
                (log_p_tilde(μ_samples[:, i-1]) + 
                 logpdf(MvNormal(μ_samples[:,i-1], σ), μ_tmp))
        #
        is_accepted = min(1, exp(log_r)) > rand()
        new_sample = is_accepted ? μ_tmp : μ_samples[:, i-1]
        
        #
        μ_samples[:, i] = new_sample
        
        #
        num_accepted += is_accepted
    end
    
    μ_samples, num_accepted
end

# Hamiltonian Monte Carlo method
function HMC(log_p_tilde, μ₀; maxiter::Int64, L::Int64, ε::Float64)
    # leapfrog による値の更新
    function leapflog(grad, p_in,μ_in, L, ε)
        μ=μ_in
        p = p_in + 0.5 *ε*grad(μ)
        for l in 1: L-1
            μ +=ε*p
            p += ε*grad(μ)
        end
        μ += ε*p
        p += 0.5 *ε*grad(μ)
        p, μ
    end
    
    # 非正規化対数事後分布の勾配関数
    grad(μ) = ForwardDiff.gradient(log_p_tilde, μ)
    
    # サンプルを格納する配列
    D = length(μ₀)
    μ_samples = Array{typeof(μ₀[1]), 2}(undef, D, maxiter)
    
    # 初期サンプル
    μ_samples[:, 1] = μ₀
    
    # 受容されたサンプル数
    num_accepted = 1
    
    for i in 2:maxiter
        # 運動量の生成
        p_in = randn(size(μ₀))
        
        # リープフロッグ
        p_out, μ_out = leapflog(grad, p_in, μ_samples[:, i-1], L, ε)
        
        # 比率rの対数を計算
        μ_in = μ_samples[:, i-1]
        log_r = (log_p_tilde(μ_out) + 
                 logpdf(MvNormal(zeros(D),eye(D)), vec(p_out))) -
                (log_p_tilde(μ_in) + 
                 logpdf(MvNormal(zeros(D),eye(D)),vec(p_in)))
                
        # 確率rでサンプル受容
        is_accepted = min(1, exp(log_r)) > rand()
        new_sample = is_accepted ? μ_out : μ_in
        
        # 新サンプルを格納
        μ_samples[:,i] = new_sample
        
        # 受容された場合、合計を加算する
        num_accepted += is_accepted
    end
    
    μ_samples, num_accepted
end


### ラッパー関数 ###
function inference_wrapper_GMH(log_joint, params, w_init; maxiter::Int64 = 100_000, σ::Float64 = 1.0)
    ulp(w) = log_joint(w, params...)
    GaussianMH(ulp, w_init; maxiter = maxiter, σ = σ)
end

function inference_wrapper_HMC(log_joint, params, w_init; maxiter::Int64 = 100_000, L::Int64 = 100, ε::Float64 = 1e-1)
    ulp(w) = log_joint(w, params...)
    HMC(ulp, w_init; maxiter = maxiter, L=L, ε=ε)
end

###
X_obs = [-2, 1, 5]
Y_obs = [-2.2, -1.0, 1.5]

σ = 1.0
μ₁ = 0.0
μ₂ = 0.0
σ₁ = 10.0
σ₂ = 10.0

log_joint(w, X, Y, σ, μ₁,σ₁,μ₂,σ₂) =
          sum(logpdf.(Normal.(w[1]*X .+ w[2],σ), Y)) +
          logpdf(Normal(μ₁,σ₁), w[1]) +
          logpdf(Normal(μ₂,σ₂), w[2])
params = (X_obs, Y_obs,σ ,μ₁, σ₁,μ₂,σ₂)

ulp(w) = log_joint(w, params...)


###
# 初期値
w_init = randn(2)
maxiter = 1000
param_posterior_GMH, num_accepted_GMH = 
    inference_wrapper_GMH(log_joint, params, w_init, maxiter = maxiter, σ = 1.0)
param_posterior_HMC, num_accepted_HMC = 
    inference_wrapper_HMC(log_joint, params, w_init, maxiter = maxiter, L = 10, ε = 1e-1)

# 
fig, axes = subplots(2,1,figsize=(8,4))

axes[1].plot(param_posterior_GMH[1,:])
set_options(axes[1], "iteration", "w₁", "w₁ sequence (GMH)")

axes[2].plot(param_posterior_GMH[2,:])
set_options(axes[2], "iteration", "w₂", "w₂ sequence (GMH)")

tight_layout()
println("acceptance rate (GMH) = $(num_accepted_GMH/maxiter)")

# 
fig, axes = subplots(2,1,figsize=(8,4))

axes[1].plot(param_posterior_HMC[1,:])
set_options(axes[1], "iteration", "w₁", "w₁ sequence (HMC)")

axes[2].plot(param_posterior_HMC[2,:])
set_options(axes[2], "iteration", "w₂", "w₂ sequence (HMC)")

tight_layout()
println("acceptance rate (HMC) = $(num_accepted_HMC/maxiter)")

次回

power-of-awareness.com

今回やりたいこと

前回

1. ハミルトニアン・モンテカルロ法

1.1 マルコフ連鎖モンテカルロ法

1.2 メトロポリス・ヘイスティング法

1.3 ハミルトニアン・モンテカルロ法

1.3.1 実装

次回

1.　ハミルトニアン・モンテカルロ法

1.1　マルコフ連鎖モンテカルロ法

1.2　メトロポリス・ヘイスティング法

1.3　ハミルトニアン・モンテカルロ法

1.3.1　実装