No-U-Turn Sampler (NUTS) を理解したい

概要

ベイズモデルを構築し，MCMCを用いてパラメータの推定ができるPythonライブラリとしてPyMCが有名です．

PyMCのMCMCでは，デフォルトでNo-U-Turn Sampler (NUTS) [Matthew D. Homan, 2014]が利用されます．

例

import pymc3 as pm

# 分布の定義
with pm.Model() as model:
    mu = pm.Normal('mu', mu=0, sigma=1)
    pm.Normal('theta', mu=mu, sigma=1)

# MCMCでサンプリング
with model:
    trace = pm.sample(random_seed=0)

NUTSはハミルトニアンモンテカルロ (HMC)をベースにした手法で，人手によるパラメータの調整が必要ない優れたアルゴリズムです．

TensorFlow Probabilityでも実装されていますね．
tfp.mcmc.NoUTurnSampler

今回はNUTSのアルゴリズムを実装してみたいと思います．

HMC

No-U-Turn SamplerのベースとなるHMCについておさらいです．

エネルギー保存則

運動量$r$，位置(高さ)$\theta$とした時，運動エネルギーと位置エネルギーをそれぞれ$K(r),U(\theta)$とすると，ハミルトニアン$H(\theta, r)$は以下で定義されます．

\begin{aligned}
H(\theta, r) &= K(r) + U(\theta)
\\
\\
K(r) &= \frac{1}{2}m r^2
\\
U(\theta) &= mg\theta
\end{aligned}

$m, g$はそれぞれ物体の質量と重力加速度です．

ハミルトニアンは時間によらず一定

リープフロッグ法による時間発展の近似

ハミルトンの運動方程式
ハミルトニアンが一定に保たれることから，時刻$\tau$について物体の運動を記述する運動方程式が導けます．これをハミルトンの運動方程式と呼びます．

\begin{aligned}
\frac{\partial r}{\partial \tau} &= - \frac{\partial U}{\partial \theta}
\\
\frac{\partial \theta}{\partial \tau} &= \frac{\partial K}{\partial r}
\end{aligned}

ハミルトンの運動方程式による時間発展を離散化して得られるのがリープフロッグ法です．

リープフロッグ法による時間発展

\begin{aligned}
r_{\tau+\frac{\epsilon}{2}} &= r_{\tau} - \frac{\epsilon}{2}\frac{\partial U}{\partial \theta_{\tau}} 
\\
\theta_{\tau+\epsilon} &= \theta_{\tau} + \epsilon r_{\tau+\frac{\epsilon}{2}}
\\
r_{\tau+\epsilon} &= r_{\tau+\frac{\epsilon}{2}} - \frac{\epsilon}{2}\frac{\partial U}{\partial \theta_{\tau+\epsilon}} 
\end{aligned}

$\epsilon$は任意のステップサイズです．

HMCのアルゴリズム

MH法では以下のことが課題でした．

• 受容率が低く，サンプル効率が悪い（特に高次元の場合）
• ステップ幅を小さくとると，受容率は上がるが自己相関が高くなる．
  → 膨大なサンプルがないと，パラメータ空間を網羅できない．

そこで，HMCでは自己相関を減らし，効率よくサンプルを取得することを目指します．

積分に重要な配位は確率密度が高い領域です
確率密度関数$P(\theta)$について

P(\theta) = \frac{e^{-S(\theta)}}{Z}

「確率密度が大きい領域」 ＝ 「作用$S(\theta)$が小さい領域」が重要な配位になります．

HMCでは，作用$S(\theta)$を位置エネルギーとみなす

U(\theta) = S(\theta)

→ 重要な配位は標高が低い配位

つまり，標高が低いサンプルを効率的に集めることが重要になります．

そこで，HMCでは、たとえ斜面を登る側に動いたとしても，ある程度時間が経てば谷底に戻ってくるような力学的な時間発展を取り入れることでサンプル効率を上げる工夫をしています．

以下は，イテレーション$t$でのサンプル$\theta^{t}$に対し，次のサンプル$\theta^{t+1}$を決める手順です．
1. 初期運動量を生成

r_0 \sim \mathcal{N}(0, I)

動き出す方向と強さをランダムに生成しています．

2. 初期状態でのハミルトニアンを計算

H_{0} = \underbrace{U(\theta^{t})}_{位置エネルギーに対応} + \underbrace{\frac{1}{2}r_0\cdot r_0}_{運動エネルギーに対応}

3. リープフロッグによる時間発展
任意のステップサイズ$\epsilon$で$L$回リープフロッグによる時間発展を行います．

時間発展後の配位と運動量をそれぞれ$\theta^{t}_{\rm fin}$, $r_{\rm fin}$ とします．

※ 力学に基づいた時間発展であるため，斜面を登る側に動き始めてもある程度時間が経てば戻ってきます．

4. 終状態でのハミルトニアンを計算

H_{\rm fin} = U(\theta^{t}_{\rm fin}) + \frac{1}{2}r_{\rm fin}\cdot r_{\rm fin}

5. メトロポリステスト
ハミルトニアンの変化量を用いてメトロポリステストを行います．

• 確率$\min(1, e^{H_0 - H_{\rm fin}})$でサンプルを受容．$\theta^{t+1} = \theta^{t}_{\rm fin}$
• 棄却する場合は$\theta^{t+1} = \theta^{t}_{0}$

※ 詳細釣り合い条件の証明などはここでは省略します．

リープフロッグの$\epsilon$が小さいほど，ハミルトニアンの変化量は小さくなるので受容確率が高くなります．また，$L$が大きいほど配位の変化が大きくなるので自己相関長が短くなります．

つまり，パタメータを適切に選択すれば，自己相関が短く，重要な配位のサンプルを効率的に集めることができます.

スライスサンプリング

スライスサンプリング [R. Neal, 2003]もMCMCの一種です．
NUTSではスライスサンプリングの技術を応用しているので，簡単にここでも整理しておきます．

MH法などでは、分布の山が複数ある多峰分布などに局所的に多くのサンプリングをしてしまう場合がある．（ステップサイズの調整が課題）

→ 近傍に限らず，山が高いところから新しいサンプルを得る

現在のサンプル$\theta^t$の時，補助変数$u$をサンプリングします．

u \sim {\rm Uniform}(0, e^{-S(\theta^t)})

これよりも確率密度の高いサンプル集合からの一様サンプリングで次のサンプル$\theta^{t+1}$を決定します．

\theta^{t+1} \sim {\rm Uniform}\{\theta'|e^{-S(\theta')} \ge u\}

このとき，同時分布$p(\theta, u)$は以下のようになります．

p(\theta, u) = 
\left\{
\begin{array}{ll}
\frac{1}{Z} & 0 \le u \le e^{-S(\theta)} \\
0 & {\rm others}
\end{array}
\right.

$u$について周辺化すると

\begin{aligned}
\int p(\theta, u) du &= \int_0^{e^{-S(\theta)}} \frac{1}{Z}du
\\
&= \frac{e^{-S(\theta)}}{Z}
\\
&= p(\theta)
\end{aligned}

つまり，$\theta$と$u$の同時分布からサンプリングを行い，$u$を無視することで$p(\theta)$からのサンプリングになっています．

実際のアルゴリズム
実際には，必ずしも連結でない集合$\{ \theta'|e^{-S(\theta')} \ge u \}$から一様サンプリングすることは難しいです．
そこで，$\theta^t$を含む領域$(\theta^{l}, \theta^r)$を見つけ，そこから一様にサンプリングしています．
領域の決め方としては，stepping-out法やdoubling法などがあるようです．

No-U-Turn Sampler

いよいよ本題のNUTSです．
HMCでは、ステップ数$L$とステップサイズ$\epsilon$の調整が課題でした．

ステップ数について

まずはリープフロッグのステップ数 $L$ に焦点を当てます．

リープフロッグによる時間発展が「十分長く」行われることを知るための基準が必要

時間発展後の配位と運動量を$\theta', r'$としたとき

「十分長く」 \Longrightarrow \theta_0と\theta'の距離が増加しなくなる.

つまり，Uターンを始めたら十分長く時間発展が行われたと考えることができます．

簡単に考えると，以下の基準値が$0$よりも小さくなるまでリープフロッグによる時間発展を続ければいいように思えます．

(\theta_0 - \theta')\cdot r'

ただし，この時間発展は可逆性を保証しなくなってしまいます．

可逆性を保証しない　$\Longrightarrow$　詳細釣り合い条件を満たさない

詳細釣り合い条件を満たさないと正しい分布に収束しない.
(MCMCの原則に違反)

Unfortunately this algorithm does not guarantee time reversibility, and is therefore not guaranteed to converge to the correct distribution. NUTS overcomes this issue by means of a recursive algorithm that preserves reversibility by running the Hamiltonian simulation both forward and backward in time. 　([Matthew D. Homan, 2014]から引用)

試しに簡単な分布に対して，上記の基準で$L$を動的に決めてHMCを実行してみましたが悲惨な結果でした．

NUTSではスライスサンプリングと同じ考え方を取り入れ，詳細釣り合い条件を満たすように次のサンプルを決定します．

初期の配位を$\theta^t_0$とした時，初期運動量をサンプリングします．

r_0 \sim \mathcal{N}(0, I)

ここで，補助変数$u$を導入します．

u \sim {\rm Uniform}(0, e^{-U(\theta^t_0)-\frac{1}{2}r_0\cdot r_0})

次のサンプル$\theta^{t+1}$は以下のスライスサンプリング同様以下の集合からの一様サンプリングで決定します．

\theta^{t+1} \sim {\rm Uniform}\{\theta', r'|e^{-U(\theta')-\frac{1}{2}r'\cdot r'} \ge u\}

実際には，リープフロッグによる時間発展を前進方向と後退方向両側に再帰的に行いながら，$\theta^t_0$を中心とした候補集合を構築し，そこから一様サンプリングを行います．

時間方向はステップごとにランダムに決定しています．

簡単な例で実装して確認
平均$[0, 0]$，分散共分散行列$\begin{pmatrix}5 & 2 \\ 2 & 1\end{pmatrix}$の2変量ガウス分布で候補集合の構築を試してみます．
$\theta=[x_1, x_2]$，作用(位置エネルギー)は$U(\theta) = \frac{x_1^2 + 5x_2^2 + 4x_1x_2}{2}$です．
リープフロッグのために勾配を計算しておきます．

\nabla_{\theta}U = [x_1 - 2x_2, 5x_2 - 2x_1]

$\theta_0 = [0.2, -0.3]$から始め，$j=5$までの経路生成をシミュレートしてみます．
※ ステップサイズ$\epsilon$はまだ適当に仮決めしてます．

import numpy as np

def hamiltonian(theta: np.array, r: np.array):
    """
    ハミルトニアンの計算
    位置エネルギー＋運動エネルギー
    """
    x = theta[0]
    y = theta[1]
    U = (1/2)*(x**2 + 5*y**2 - 4*x*y)
    P = 0.5 * np.dot(r, r)
    return U + P


def grad_U(theta: np.array):
    """
    作用Uのtheta(配位)微分
    """
    x = theta[0]
    y = theta[1]
    return np.array([x-2*y, 5*y-2*x])

    
def leapfrog(theta: np.array, r: np.array, eps):
    """
    リープフロッグによる時間発展(1ステップ)
    Args:
        theta: 配位
        r: 運動量
        eps: ステップサイズ
    """
    r = r - (eps/2)*grad_U(theta)
    theta = theta + eps*r
    r = r - (eps/2)*grad_U(theta)
    return theta, r
    
    
def build_route(theta, r, v, u, j, eps: float=0.3):
    """
    再帰的２進木構築による経路生成
    Args:
        tehta: 配位
        r: 運動量
        v: 時間方向{-1, 1}
        u: 棄却サンプリング用の補助変数
        j: 現在ステップ
    """
    if j == 0:
        theta, r = leapfrog(theta, r, v*eps)
        # 棄却サンプリング
        if u <= np.exp(- hamiltonian(theta, r)):
            C_ = {tuple(theta.tolist())}
        else:
            C_ = set()
        return theta, r, theta, r, C_

    theta_m, r_m, theta_p, r_p, C_ = build_route(theta, r, v, u, j-1)
    if v == -1:
        theta_m, r_m, _, _, C__ = build_route(theta_m, r_m, v, u, j-1)
    else:
        _, _, theta_p, r_p, C__ = build_route(theta_p, r_p, v, u, j-1)
    
    C_ = (C_|C__)
    return theta_m, r_m, theta_p, r_p, C_

#############
# ここからメイン
#############
theta = np.array([0.2, -0.3])  # 初期配位
r = np.random.randn(2)  # 初期運動量を生成

# 初期化
theta_m = theta.copy()
theta_p = theta.copy()
r_m = r.copy()
r_p = r.copy()

# スライス変数をサンプリング
u = random.uniform(0, np.exp(- hamiltonian(theta, r)))

J = 5
C = {tuple(theta.tolist())}  # 候補集合
for j in range(J):
    # 時間方向をランダムに選択
    v = -1 if random.randint(0, 1) == 0 else 1
    # 経路構築
    if v == -1:
        theta_m, r_m, _, _, C_ = build_route(theta_m, r_m, v, u, j)
    else:
        _, _, theta_p, r_p, C_ = build_route(theta_p, r_p, v, u, j)

    C = (C|C_)

これを各$j$ごとに色を変えてプロットしてみると以下のようになりました．

時間発展の停止基準
永遠に候補集合を拡大し続けるのは無駄が多いので，Uターンを始めたら経路生成をストップします． (上図でもUターンしてますね.)

前進方向の先頭の配位と運動量を$\theta^+, r^+$，後退方向の先頭の配位と運動量を$\theta^-, r^-$した時，以下の条件に合致した場合はそこで終了します．

(\theta^+ - \theta^-)\cdot r^- < 0 \quad {\rm or} \quad (\theta^+ - \theta^-)\cdot r^+ < 0 

また，新たに候補集合に追加される確率が極端に低い状態になった場合も経路生成を終了します．つまり，以下の条件を満たす場合も終了します．

u \ge e^{\Delta_{\rm max} - U(\theta') -\frac{1}{2}r'\cdot r'}

ここで，$\Delta_{\rm max}$は経路生成に支障をきたさないように1000くらいの大きな値に設定することが推奨されています．

ステップサイズについて

• $\epsilon$が小さすぎる → 計算量を浪費
• $\epsilon$が大きすぎる → 棄却率が高くなる(MCMCの精度低下)

考え方
$\epsilon$を変数として，期待するメトロポリステストの受容確率を$\delta$，イテレーション$t\ge1$での実際の受容確率を$\alpha_t$とした時，アルゴリズムの振る舞いに関する統計量$h(\epsilon)$を以下のように定義します．

\begin{aligned}
h(\epsilon)&\equiv\mathbb{E}_t[H_t|\epsilon]
\\
H_t &= \delta - \alpha_t
\end{aligned}

$h(\epsilon) = 0$となるパラメータ$\epsilon \in \mathbb{R}$を求めることが目的になります．
$h(\epsilon)$は$\epsilon$の非減少関数と仮定できるので凸最適化問題になり，これをNesterovの双対平均化法を用いて最適化します．

※ 強い仮定のもとでは，標準的なHMCでは、メトロポリステストの平均受容確率の最適値は0.65らしいです．[Beskos et al., 2010], [Neal, 2011]

Nesterovの双対平均化法(dual averaging) [Xiao, L, 2009]
※ここについては理解が浅く，ほぼ写経になります．

バッチ最適化におけるNesterovの主双対法をオンライン確率的最適化に拡張したもの

$h(x) \equiv\mathbb{E}_t[H_t|x]=0$となるパラメータ$x\in\mathbb{R}$を求める際，以下の更新を行います．

\begin{aligned}
x_{t+1} &\leftarrow \mu -\frac{\sqrt{t}}{\gamma}\frac{1}{t+t_0}\sum_{i=1}^t H_{i;}
\\
\bar{x}_{t+1} &\leftarrow \eta_{t}x_{t+1} + (1-\eta_t)\bar{x}_t
\end{aligned}
\tag{1}

• $\mu$ : 収縮目標
• $\gamma > 0$ : 収縮量を制御するパラメータ
• $t_0 \ge 0$ : 初期反復を安定させるためのパラメータ
• $\eta \equiv t^{-\kappa}$ : ステップサイズスケジュール

$H_t$が有界である限り、$h(\bar{x}_t)$が0に収束することが保障されます。

NUTSの場合

NUTSでは単一アクセプト/リジェクトのステップはないので、メトロポリステストの受容率に代わる統計量を定義しなければならない

受容確率に変わる統計量$H_t^{\rm NUTS}$を以下のように定義します

H_t^{\rm NUTS} \equiv \frac{1}{|\mathcal{B}_t^{\rm fin}|}
\sum_{\theta, r \in \mathcal{B}_t^{\rm fin}} 
\min
\left\{
1, \frac{p(\theta, r)}{p(\theta^{t-1}, r^{t}_0)}
\right\}

• $\mathcal{B}_t^{\rm fin}$ : イテレーション$t$で探索された全ての状態の集合(棄却されたものも含む)
• $(\theta^{t-1}, r^{t}_0)$ イテレーション$t$での初期状態

ここで、$H^{\rm NUTS}$は$\epsilon$について非減少と仮定できます。
よって、式(1)において、$H_t = \delta - H^{\rm NUTS}_t$, $x = \log \epsilon$として双対平均化法を適用することで$\mathbb{E}_t[H^{\rm NUTS}] = \delta$を強制できます。

実際の更新式は以下のようになります.

\begin{aligned}
\bar{H}_{t} 
&\leftarrow \left(1 - \frac{1}{t + t_0}\right)\bar{H}_{t-1} + \frac{1}{t + t_0}(\delta - \frac{\alpha}{n_{\alpha}})
\\
\log\epsilon_t 
&\leftarrow
\mu - \frac{\sqrt{t}}{\gamma}\bar{H}_t
\\
\log\bar{\epsilon}_t 
&\leftarrow
t^{-\kappa}\log\epsilon_t + (1 - t^{-\kappa})\log\bar{\epsilon}_{t-1}
\end{aligned}

$\alpha, n_{\alpha}$は逐次的に$H_t$を計算するための補助変数で、$n_{\alpha}$はイテレーション$t$での候補数、$\alpha$は各候補の受容率の合計です。

各パラメータについては、実験的に以下の値が推奨されています。
$\mu = \log(10\epsilon_0), \gamma=0.05, t_0=10, \kappa=0.75, \bar{\epsilon}_0 = 1$

この更新を指定したイテレーションまで行い(warmup phase)、以降は$\epsilon=\bar{\epsilon}$で固定します.

NUTSの実装・シミュレーション

以上のことを踏まえ、NUTSを実装・シミュレーションしてみます.
ここでも、動作確認のために平均$[0, 0]$, 分散共分散行列$\begin{pmatrix}5 & 2 \\ 2 & 1\end{pmatrix}$の2変量ガウス分布で試してみます.

NUTS

import numpy as np
import random

def hamiltonian(theta: np.array, r: np.array):
    """
    ハミルトニアンの計算
    Args:
        theta: 配位
        r: 運動量
    """
    x = theta[0]
    y = theta[1]
    U = (1/2)*(x**2 + 5*y**2 - 4*x*y)
    P = 0.5 * np.dot(r, r)
    return U + P


def grad_U(theta: np.array):
    """
    作用Uのtheta(配位)微分
    """
    x = theta[0]
    y = theta[1]
    return np.array([x-2*y, 5*y-2*x])


def leapfrog(theta: np.array, r: np.array, eps):
    """
    リープフロッグによる時間発展(1ステップ)
    Args:
        theta: 配位
        r: 運動量
        eps: ステップサイズ
    """
    r = r - (eps/2)*grad_U(theta)
    theta = theta + eps*r
    r = r - (eps/2)*grad_U(theta)
    return theta, r
    

def build_route(theta, r, v, u, j, ham_0, eps):
    """
    再帰的２進木構築による経路生成
    Args:
        tehta: 配位
        r: 運動量
        v: 時間方向{-1, 1}
        u: 棄却サンプリング用の補助変数
        j: 現在ステップ
        ham_0: 初期状態でのハミルトニアン
        eps: リープフロッグのステップサイズ
    """
    if j == 0:
        theta, r = leapfrog(theta, r, v*eps)
        n_ = 1 if u <= np.exp(- hamiltonian(theta, r)) else 0
        s_ = 1 if u < np.exp(1000 - hamiltonian(theta, r)) else 0
        alpha_ = min(1, np.exp(- hamiltonian(theta, r) + ham_0))
        n_alpha_ = 1
        return theta, r, theta, r, theta, n_, s_, alpha_, n_alpha_

    theta_m, r_m, theta_p, r_p, theta_, n_, s_, alpha_, n_alpha_ = build_route(theta, r, v, u, j-1, ham_0, eps)
    if s_ == 1:  
        if v == -1:
            theta_m, r_m, _, _, theta__, n__, s__, alpha__, n_alpha__ = build_route(theta_m, r_m, v, u, j-1, ham_0, eps)
        else:
            _, _, theta_p, r_p, theta__, n__, s__, alpha__, n_alpha__ = build_route(theta_p, r_p, v, u, j-1, ham_0, eps)
        
        n_rate = 0 if n__ == 0 else (n__ / (n_ + n__))
        if np.random.uniform(0,1) < n_rate:
            theta_ = theta__
        alpha_ = alpha_ + alpha__
        n_alpha_ = n_alpha_ + n_alpha__
        n_ = n_ + n__
        s_ = s__ * (1 if np.dot((theta_p - theta_m), r_m) >= 0 else 0) * (1 if np.dot((theta_p - theta_m), r_p) >= 0 else 0)
        
    return theta_m, r_m, theta_p, r_p, theta_, n_, s_, alpha_, n_alpha_


def init_eps(theta):
    """
    収束を早めるためのepsの初期値
    ※ ここの説明は今回は省略(必須ではない)
    """
    eps = 1
    r = np.random.randn(2) 
    theta_old, r_old = theta, r
    theta, r = leapfrog(theta_old, r_old, eps)
    a = 2 * (1 if (np.exp(-hamiltonian(theta, r)) / np.exp(-hamiltonian(theta_old, r_old))) > 0.5 else 0) - 1
    
    while (np.exp(-hamiltonian(theta, r)) / np.exp(-hamiltonian(theta_old, r_old)))**a > 2**(-a):
        eps = (2**a) * eps
        theta_old, r_old = theta, r
        theta, r = leapfrog(theta_old, r_old, eps)
    return eps


#####################
# ここからNUTSのメイン処理
#####################
import tqdm

# 初期点
theta = np.array([0.0, 0.0])  
# epsの初期値
eps = init_eps(theta)

# パラメータ
M = 50000  # イテレーション数
M_adapt = 1000  # warmup
delta = 0.8
mu = np.log(10*eps)
eps_ = 1
log_eps_ = np.log(eps_)
H_ = 0
t_0 = 10
gamma = 0.05
kappa = 0.75

samples = [] # MCMCのサンプル格納用
#for m in range(1, M):
for m in tqdm.tqdm(range(1, M)):
    # 運動量をサンプリング
    r = np.random.randn(2) 
    # 初期状態でのハミルトニアン
    hami_0 = hamiltonian(theta, r)
    # スライス変数をサンプリング
    u = random.uniform(0, np.exp(- hami_0))

    # 初期化
    theta_m = theta
    theta_p = theta
    r_m = r
    r_p = r
    s = 1
    n = 1
    j = 0
    
    # Uターンするまでリープフロッグ&棄却サンプリング
    while s == 1:
        v = -1 if random.randint(0, 1) == 0 else 1 # 時間発展方向
        if v == -1:
            theta_m, r_m, _, _, theta_, n_, s_, alpha, n_alpha = build_route(theta_m, r_m, v, u, j, hami_0, eps)
        else:
            _, _, theta_p, r_p, theta_, n_, s_, alpha, n_alpha = build_route(theta_p, r_p, v, u, j, hami_0, eps)
            
        if s_ == 1:
            # 逐次評価で一様サンプリングを表現（全ての状態を保持するのはメモリ効率が悪いため）
            if np.random.uniform(0,1) < (n_/n):
                theta = theta_
                
        n = n + n_
        s = s_ * (1 if np.dot((theta_p - theta_m), r_m) >= 0 else 0) * (1 if np.dot((theta_p - theta_m), r_p) >= 0 else 0)
        j += 1
    
    # 双対平均化法による更新
    if m <= M_adapt:
        # warmup phaseの間はepsを更新
        H_ = (1 - (1 / (m+t_0))) * H_ + (1 / (m+t_0)) * (delta - (alpha / n_alpha))
        log_eps = mu - (np.sqrt(m) / gamma)*H_
        log_eps_ = m**(-kappa) * log_eps + (1 - m**(-kappa)) * log_eps_
        eps = np.exp(log_eps)
        eps_ = np.exp(log_eps_)
    elif m == (M_adapt+1):
        # epsを固定
        eps = eps_
    
    # サンプルを格納
    samples.append(theta)
            
# warmup phaseのサンプルを捨てる
samples = samples[M_adapt:]


#####################
# プロット
#####################
import matplotlib.pyplot as plt
import japanize_matplotlib

x1s = []
x2s = []
for x1, x2 in samples:
    x1s.append(x1)
    x2s.append(x2)
    
fig = plt.figure(figsize=(20, 5))
ax1 = fig.add_subplot(1, 3, 1)
ax2 = fig.add_subplot(1, 3, 2)
ax3 = fig.add_subplot(1, 3, 3)

ax1.scatter(x1s, x2s, alpha=0.01)
ax1.set_title('散布図')
ax1.set_xlabel('x1')
ax1.set_ylabel('x2')
ax2.hist(x1s, bins=100)
ax2.set_title('x1ヒストグラム')
ax2.set_xlabel('x1')
ax3.hist(x2s, bins=100)
ax3.set_title('x2ヒストグラム')
ax3.set_xlabel('x2')

$\delta = 0.9$の時の結果が以下です。

warmup終了時点での$\epsilon$は$0.4911279110929324$でした

np.mean(x1s), np.mean(x2s)

(0.0055021458937409414, 0.0013560857774292641)

$\delta = 0.1$の時の結果が以下です。

warmup終了時点での$\epsilon$は$1.4020990054443532$でした

np.mean(x1s), np.mean(x2s)

(0.058092148366324134, 0.025138085185976995)

上記の結果からも、$\delta$を大きく設定することで$\epsilon$が小さくなり、高精度なMCMCが実行できてることが確認できました。

今回の実行時間は5秒程度でした.

100%|██████████| 49999/49999 [00:05<00:00, 8478.97it/s]

参考文献

• Slice sampling, [R. Neal, 2003]
• ゼロからできるMCMC
  • HMCを始め、MCMC全般の勉強にとても役立ちました
• はじめてのMCMC (スライス・サンプリング)
  • スライスサンプリングの理解に役立ちました
• スライスサンプリング
  • スライスサンプリングの理解に役立ちました
• Stanを使ってNUTSを実装する
  • NUTSの実装の参考になりました
• 機械学習における確率的最適化