Zig-Zag 過程によるサンプリング

ジャンプと確定的な動きによる新たな MCMC 手法

Process
Sampling
Julia
MCMC
Author

司馬博文

Published

7/03/2024

Modified

7/18/2024

概要
Zig-Zag サンプラー定義とエルゴード性を解説する．続いて，Zig-Zag サンプラーは非対称なダイナミクスを持つために，従来の MCMC よりも速い収束が期待されることを，MALA との比較でみる．最後に，Zig-Zag サンプラーの実装に用いたパッケージとその利用方法を示す．

1 Zig-Zag 過程

1.1 はじめに

１次元の Zig-Zag 過程は元々，Curie-Weiss 模型 における Glauber 動力学を lifting により非可逆化して得る Markov 連鎖の，スケーリング極限として特定された Feller-Dynkin 過程である (Bierkens and Roberts, 2017)．

区分確定的 Markov 過程（PDMP）といい，ランダムな時刻にランダムな動きをする以外は，決定論的な動きをする過程である．

\(\mathbb{R}^2\) 上の Gauss 分布に収束する Zig-Zag 過程の軌跡

PDMP の一般論については次の記事も参照：

ただし Zig-Zag 過程は，(Goldstein, 1951) で電信方程式と関連して，同様の過程が扱われた歴史もある．

1.2 設定

Zig-Zag 過程 \(Z=(X,\Theta)\) の状態空間は，力学的な立場に立って \(E=\mathbb{R}^d\times\{\pm1\}^d\) と見ることが多い．

力学における相空間と同様，\(X\in\mathbb{R}^d\) が位置を，\(\theta\in\{\pm1\}^d\) が速度を表すと解する．すなわち，Zig-Zag 過程は全座標系と \(45\) 度をなす方向に，常に一定の単位速さで \(\mathbb{R}^d\) 上を運動する粒子とみなせる．

すなわち，\((x,\theta)\in E\) から出発する Zig-Zag 過程は，次の微分方程式系で定まる決定論的なフロー \(\phi_{(x,\theta)}:\mathbb{R}\to\mathbb{R}^d\) に従って運動する粒子とみなせる： \[ \frac{d \phi_{(x,\theta)}(t)}{d t}=\theta,\qquad \frac{d \Theta_t}{d t}=0. \]

1.3 アルゴリズム

1.3.1 全体像

Zig-Zag 過程 \(Z\) は次のようにしてシミュレーションできる：

Zig-Zag 過程のシミュレーション
  1. レート関数 \(\lambda_1,\cdots,\lambda_d\) から定まる強度関数 \[ m_i(t):=\lambda_i(x+\theta t,\theta),\qquad i\in[d], \] を持つ，\(d\) 個の独立な \(\mathbb{R}_+\) 上の非一様 Poisson 点過程の，最初の到着時刻 \(T_1,\cdots,T_d\) をシミュレーションする．
  2. 最初に到着した座標番号 \(j:=\operatorname*{argmin}_{i\in[d]}T_i\) について，時刻 \(T_j\) に速度成分 \(\theta_j\) の符号を反転させる．すなわち，関数 \[ F_j(\theta)_i:=\begin{cases}-\theta_i&i=j\\\theta_i&i\ne j\end{cases} \] に従ってジャンプする．
  3. １に \(t=T_j\) として戻って，くり返す．

もう一つ，MCMC の文脈で自然な見方は，状態空間を \[ E=\bigcup_{\theta\in\{\pm1\}^d}\mathbb{R}^d\times\{\theta\} \] と取る見方である．これは (Davis, 1993) による 一般の PDMP の設定 と対応する．

この \(E\) 上で，レート関数 \[ \lambda(x,\theta):=\sum_{i=1}^d\lambda_i(x,\theta) \] が定める強度 \[ M(t):=\lambda(x+t\theta,\theta) \] を持った \(\mathbb{R}_+\) 上の非一様 Poisson 点過程に従ってジャンプが訪れる．

この点過程に対して，確率核 \[ Q((x,\theta),-):=\sum_{i=1}^d\frac{\lambda_i(x,\theta)}{\lambda(x,\theta)}\delta_{(x,F_i(\theta))}(-) \] に 印付けられた点過程 が，\(Z\) の跳躍測度である．

Zig-Zag 過程に対する２つの定義を与えた．これら２つが同分布の過程を定めることは (Corbella et al., 2022), (Hardcastle et al., 2024) などさまざまなところで触れられているが，証明が必要である．

まず，\(\min_{i\in[d]}T_i\) が，強度関数 \(M\) が定める到着時刻に同分布であることを示す．

各 \(T_i\) の密度は \[ p_i(t)=m_i(t)e^{-M_i(t)}1_{(0,\infty)}(t) \] で与えられ，\(T_i\) は互いに独立だから，\((T_1,\cdots,T_d)\) の結合密度もわかる．

\(T_1,\cdots,T_d\) を昇順に並べた順序統計量を \[ T_{(1)}\le\cdots\le T_{(d)} \] で表すとする．この \(d\) 次元確率ベクトルの密度 \(p\) は， \[ p(t_1,\cdots,t_d)=1_{\left\{t_1\le\cdots\le t_d\right\}}(t_1,\cdots,t_d)\left(\sum_{\sigma\in\mathfrak{S}_d}\prod_{i=1}^dm_i(t_{\sigma(i)})e^{-M_i(t_{\sigma(i)})}\right) \] と計算できる．

この \(p\) を \(t_2,\cdots,t_d\) に関して積分することで，\(T_1\) の密度が得られる：1 \[\begin{align*} p_{(1)}(t)&=\int_{(0,\infty)^{d-1}}p(t_1,\cdots,t_d)\,dt_2\cdots dt_d\\ &=\biggr(\sum_{i=1}^dm_i(t_1)\biggl)\exp\left(-\sum_{i=1}^dM_i(t_1)\right)=m(t_1)e^{-M(t_1)}. \end{align*}\]

これは確かに，強度関数 \(m\) が定める到着時刻の密度である．

続いて，\(j=\operatorname*{argmin}_{i\in[d]}T_i\) の，\(\min_{i\in[d]}T_i\) に関する条件付き確率質量関数が \[ q(i|t)=\frac{m_i(t)}{\sum_{i=1}^dm_i(t)} \] であることを示す．

そのためには，任意の \(i\in[d]\) と \(A\in\mathcal{B}(\mathbb{R}^+)\) とに関して \(\left\{T_{(1)}\in A,T_{(1)}=T_i\right\}\) という形の事象を計算し，密度が積の形で与えられることを見れば良い．

\[\begin{align*} &\qquad\operatorname{P}[T_{(1)}\in A,T_{(1)}=T_i]\\ &=\operatorname{P}[T_i\in A,\forall_{j\ne i}\;T_i\le T_j]\\ &=\int_Ap_i(t_i)\,dt_i\left(\sum_{\sigma\in\mathrm{Aut}([d]\setminus\{i\})}\int^\infty_{t_i}p_{\sigma(1)}(t_{\sigma(1)})\,dt_{\sigma(1)}\int^\infty_{t_{\sigma(1)}}p_{\sigma(2)}(t_{\sigma(2)})\,dt_{\sigma(2)}\cdots\int^\infty_{t_{\sigma(d-1)}}p_{\sigma(d)}(t_{\sigma(d)})\,dt_{\sigma(d)}\right)\\ &=\int_Am_i(t_i)\exp\left(-\sum_{i=1}^dm_i(t_i)\right)\,dt_i\\ &=\int_A\frac{m_i(t_i)}{m(t_i)}m(t_i)e^{-M(t_i)}\,dt_i. \end{align*}\]

よって，\(\min_{i\in[d]}T_i\) と \(\operatorname*{argmin}_{i\in[d]}T_i\) とに関する結合密度は，2 \[ q(i|t)p_{(1)}(t) \] という積の形で与えられることがわかった．

まとめ
  1. 前述の定義は，\(\min_{i\in[d]}T_i\) の形で密度 \(p_{(1)}\) からシミュレーションし，\(\operatorname*{argmin}_{i\in[d]}T_i\) の形で \(q\) からシミュレーションしている．
  2. 後述の定義は，\(p_{(1)}(t)\) から直接シミュレーションし，再び \(q(i|t)\) から直接シミュレーションをする．

１が２に等価であることがわかった．

1.3.2 到着時刻 \(T_i\) のシミュレーション方法

Zig-Zag 過程のシミュレーションは，ほとんど強度

\[ M_i(t):=\int^t_0m_i(s)\,ds \] を持つ非一様 Poisson 点過程のシミュレーションに帰着される．

実はこれは，指数分布確率変数 \(E_i\overset{\text{iid}}{\sim}\operatorname{Exp}(1)\) について \[ T_i\overset{\text{d}}{=}M_i^{-1}(E_i) \] と求まる．

仮にこの逆関数 \(M_i^{-1}\) が得られない場合でも，剪定 (Lewis and Shedler, 1979) によって \(T_i\) は正確なシミュレーションが可能である．

1.4 レート関数の条件

Zig-Zag 過程 \(Z\) がどのような分布に従うかは，全てレート関数 \(\lambda_1,\cdots,\lambda_d\) に委ねられている．

Zig-Zag 過程のレート関数 \(\lambda_1,\cdots,\lambda_d:E\to\mathbb{R}_+\) は，負の対数密度 \(U\in C^1(\mathbb{R}^d)\) に対して， \[ \lambda_i(x,\theta):=(\theta_i\partial_iU(x))_++\gamma_i(x,\theta_{-i})\quad(i\in[d]) \] と定める．

ただし，次を仮定する：

  • \(\gamma_i:E\to\mathbb{R}_+\) は，\(\theta_i\) のみには依らない任意の非負連続関数3 \[ \gamma_i(x,\theta)=\gamma_i(x,F_i(\theta)). \]
  • \(e^{-U}\in\mathcal{L}^1(\mathbb{R}^d)\) が成り立ち，\(\pi(dx)\,\propto\,e^{-U(x)}dx\) が確率測度を定める．
  • \(M_i\) は \(t\to\infty\) の極限で発散する： \[ M_i(t):=\int^t_0\lambda_i(x+t\theta,\theta)\,dt \]

また， \[ M_i(t):=\int^t_0\lambda_i(x+t\theta,\theta)\,dt \] は \(t\to\infty\) の極限で発散する必要がある．

さもなくば，\(M_i:(0,L)\to(0,\infty)\;(L\in(0,\infty])\) の形で定まらず，\(M_i\) がこのような可微分同相を与えない場合は \[ T_i:=M_i^{-1}(E_i),\qquad E_i\overset{\text{iid}}{\sim}\operatorname{Exp}(1), \] によるシミュレーションも不正確になる．

(Bierkens et al., 2019, pp. 1294 定理2.2)

上述のリフレッシュレート \(\lambda_1,\cdots,\lambda_d\) に対して，定義 1.3 で定まる Zig-Zag 過程 \(Z\) は次の分布 \(\widetilde{\pi}=\pi\otimes\mathrm{U}(\{\pm1\}^{d})\in\mathcal{P}(E)\) を不変にする： \[ \widetilde{\pi}(dxd\theta)=\frac{1}{2^d}\frac{e^{-U(x)}}{\mathcal{Z}}\,dxd\theta \]

\(\{\pm1\}^d\) 上の周辺分布が一様分布になっていること，勾配ベクトル \(DU\) の情報のみを使っており，座標に沿った方向しか見ていないため \(U\) の異方性に大きく左右されること，これらが「必ずしもそうある必要はない」拡張可能な点である．

1.5 エルゴード性の条件

\(\pi\) が不変分布になるための十分条件 1.4 は極めて緩かったが，MCMC として使えるためにはエルゴード性が成り立つ必要がある．

(Bierkens and Roberts, 2017 定理５)

\(d=1\) で，レート関数 \(\lambda:E\to\mathbb{R}_+\) はある \(x_0>0\) が存在して次を満たすとする： \[ \inf_{x\ge x_0}\lambda(x,1)>\sup_{x\ge x_0}\lambda(x,-1), \] \[ \inf_{x\le-x_0}\lambda(x,-1)>\sup_{x\le-x_0}\lambda(x,1). \]

このとき，ある関数 \(f:E\to[1,\infty)\) が存在して \(f(x,\theta)\to\infty\;(\lvert x\rvert\to\infty)\) が成り立ち，かつ \[ \left\|P^t\left((x,\theta),-\right)-\pi\right\|_\mathrm{TV}\le\kappa f(x,\theta)\rho^t,\qquad(x,\theta)\in E,t\ge0,\rho\in(0,1). \]

1.6 Subsampling Technology

Zig-Zag 過程はレート関数 \(\lambda\) の設計に大きな自由度があった（第 1.4 節）．

これにより，Zig-Zag 過程ではバイアスを導入しない subsampling が可能であり，これを通じて データサイズに依らない一定効率での事後分布サンプリングが可能になる という super-efficiency (Bierkens et al., 2019) と呼ばれる性質を持つ．

この性質が実用上は最も重要である．詳しくは，次の記事を参照：

2 シミュレーション

2.1 １次元での例

ZigZag サンプラーは非対称なダイナミクスを持っており，その点が MALA (Metropolis-adjusted Langevin Algorithm) や HMC (Hamiltonian Monte Carlo) などの従来手法と異なる．

１次元でその違いを確認するために，Cauchy 分布という裾の重い分布を用いる．Cauchy 分布 \(\mathrm{C}(\mu,\sigma)\) は次のような密度を持つ： \[ f(x)=\frac{1}{\pi\sigma}\frac{1}{1+\left(\frac{x-\mu}{\sigma}\right)^2}. \]

その裾の重さ故，平均も分散も存在しない（発散する）．

このとき，次のような観察が得られる：

  • ZigZag サンプラーは Cauchy 分布に対して，最頻値から十分遠くから開始しても，高速に最頻値に戻ってくる．
  • MALA は diffusive behaviour が見られ，最頻値に戻るまでに時間がかかる．

C(0,1) に対する ZigZag サンプラーの軌道

MALA と比較すると，その再帰の速さが歴然としている：4

Code 
using AdvancedHMC, AdvancedMH, ForwardDiff
using LinearAlgebra
using LogDensityProblems
using LogDensityProblemsAD
using StructArrays
using Distributions
using Random

Random.seed!(2)

# Define the target distribution (1D Cauchy) using the `LogDensityProblem` interface
struct LogTargetDensityCauchy
    loc::Float64
    scale::Float64
end

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}()

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

# Use automatic differentiation to compute gradients
model_with_ad = LogDensityProblemsAD.ADgradient(Val(:ForwardDiff), LogTargetDensityCauchy(0.0, 1.0))

# Set up the sampler with a multivariate Gaussian proposal.
σ² = 100
spl = MALA(x -> MvNormal((σ² / 2) .* x, σ² * I))

# Sample from the posterior.
chain = sample(model_with_ad, spl, 2000; initial_params=initial_θ, chain_type=StructArray, param_names=["θ"])

# plot
θ_vector = chain.θ
sample_values = zip(1:length(θ_vector), θ_vector)
p2 = plot_1dtraj(sample_values, title="1D MALA Sampler (Cauchy Distribution)", markersize=0, ylim=(-30, 750), label="MALA_1D")

C(0,1) に対する MALA サンプラーの軌道

２つを並べて比較すると，Langevin ダイナミクスの方は，少し diffusive な動き（random walk property と呼ばれる）が見られることがわかる．

NUTS サンプラーはステップサイズを極めて大きくするため，プロットによるダイナミクスの比較があまり意味を持たなくなってくる．

実際見てみると恐ろしいものである：

C(0,1) に対する NUTS サンプラーの軌道

2.2 ２次元での例

  1. 勾配 \(-\nabla\log\pi\) を計算し，∇ϕ(x, i, Γ) の形で定義する．

\[ \Sigma^{-1}=\begin{pmatrix}2&-1\\-1&2\end{pmatrix} \]

で定まる分散共分散行列 \(\Sigma\) を持った中心化された正規分布 \(\pi(x)dx=\mathrm{N}_2(0,\Sigma)(dx)\) に対しては，対数尤度は

\[\begin{align*} \log\pi(x)&=-\log\left((2\pi)^{d/2}(\det\Sigma)^{1/2}\right)\\ &\qquad\quad-\frac{1}{2}x^\top\Sigma^{-1}x \end{align*}\]

であるから，\(\phi:=-\log\pi\) の第 \(i\) 成分に関する微分は

\[\begin{align*} \partial_i\phi(x)&=\frac{\partial }{\partial x_i}\biggr(\frac{1}{2}x^\top\Sigma^{-1}x\biggl)\\ &=\Sigma^{-1}x. \end{align*}\]

using ZigZagBoomerang
using SparseArrays

d = 2

# 対数尤度関数 ϕ の第 i 成分に関する微分を計算
Γ = sparse([1,1,2,2], [1,2,1,2], [2.0,-1.0,-1.0,2.0])
∇ϕ(x, i, Γ) = ZigZagBoomerang.idot(Γ, i, x)

# 初期値
t0 = 0.0
x0 = randn(d)
θ0 = rand([-1.0,1.0], d)

# Rejection bounds
c = 1.0 * ones(length(x0))

# ZigZag 過程をインスタンス化
Z = ZigZag(Γ, x0*0)

# シミュレーション実行
T = 20.0
zigzag_trace, (tT, xT, θT), (acc, num) = spdmp(∇ϕ, t0, x0, θ0, T, c, Z, Γ; adapt=true)

# 軌跡を離散化
traj = collect(zigzag_trace)
1
勾配関数∇ϕの計算のためには，共分散行列の逆（精度行列ともいう）をSparseMatrixCSC型で指定する必要があることに注意．idotの実装 も参照．
2
idotは，疎行列Γの第i列と，疎ベクトルxとの内積を高速に計算する関数．

idot(A,j,u)は，疎行列Aの第j列と，疎ベクトルuとの内積を高速に計算する関数である．

function idot(A::SparseMatrixCSC, j, x)
    rows = rowvals(A)
    vals = nonzeros(A)
    s = zero(eltype(x))
    @inbounds for i in nzrange(A, j)
        s += vals[i]'*x[rows[i]][2]
    end
    s
end
1
パッケージ内部で，位置 \(x\in\mathbb{R}^d\) は全て SparseSate 型に統一されている？
2
疎行列 A の行インデックスを取得．rowvals(A)はベクトルであり，第１列から順番に，非零要素のある行番号が格納されている．
3
非零要素の値が格納されている．
4
@inbounds は，範囲外アクセスを許容するマクロ．高速化のためだろう．nzrange は，A の第 j 列に非零要素がある範囲を，第 \(1\) 列から累積して何番目かで返す．すなわち，rows[i]で正確に第j列の非零要素の行番号を狙い撃ちしてイテレーションできる．
5
xの非零要素がある行番号 rows[i] における成分の値 u[rows[i]][2] はこのような表記になる．これと，A の非零要素 vals[i] との内積を計算．

なお，通常の行列に対しては，次のように実装されている：

idot(A, j, x) = dot((@view A[:, j]), x)

3 Zig-Zag サンプラーの実装

ZigZagBoomerang の実装を紹介する．

Julia の MCMC パッケージ一般については次の稿を参照：

Article Image

Julia による MCMC サンプリング

新時代の確率的プログラミング環境の構築に向けて

3.1 ZigZag サンプラーを提供しているパッケージ一覧

3.1.1 Rパッケージ

Joris Bierkens ら開発の R パッケージ RZigZag (GitHub / CRAN) が最も手軽に実行できる．

3.1.2 Juliaパッケージ

一方で，(Bierkens et al., 2022) では Julia によるパッケージ ZigZagBoomerang（GitHub / ANN / docs） も提供している．名前によらず，2022.1 月リリースの v.0.11 以降は BPS もサポートしている．

using Pkg
Pkg.add("ZigZagBoomerang")

3.2 ZigZagオブジェクト

ZigZag <: ContinuousDynamics <: Any は次のフィールドを持つ：5

  • Γ::SparseMatrixCSC：ポテンシャル関数
  • μ::Vector：平均
struct ZigZag{T,S,S2,R} <: ContinuousDynamics
    Γ::T
    μ::S
    σ::S2
    λref::R
    ρ::R
    ρ̄::R
end
ZigZag(Γ, μ, σ=(Vector(diag(Γ))).^(-0.5); λref=0.0, ρ=0.0) = ZigZag(Γ, μ, σ, λref, ρ, sqrt(1-ρ^2))

dynamics.jlにて，現在時刻にランダム時刻τを加算し，位置xを更新するが，θにはまだ触れない．

function move_forward!(τ, t, x, θ, Z::Union{BouncyParticle, ZigZag})
    t += τ
    x .+= θ .* τ
    t, x, θ
end

内積∇ϕxを用いて，位置xに対する第i成分の反射を計算する．

function reflect!(i, ∇ϕx::Number, x, θ, F::Union{ZigZag, FactBoomerang})
    θ[i] = -θ[i]
    θ
end
function reflect!(i, ∇ϕx, x, θ, F::Union{ZigZag, FactBoomerang})
    θ[i] = θ[i] - (2*dot(∇ϕx, θ[i])/normsq(∇ϕx))*∇ϕx
    θ
end

References

Andrieu, C., and Livingstone, S. (2021). Peskun–Tierney ordering for Markovian Monte Carlo: Beyond the reversible scenario. The Annals of Statistics, 49(4), 1958–1981.
Bierkens, J., Fearnhead, P., and Roberts, G. (2019). The Zig-Zag Process and Super-Efficient Sampling for Bayesian Analysis of Big Data. The Annals of Statistics, 47(3), 1288–1320.
Bierkens, J., Grazzi, S., Meulen, F. van der, and Schauer, M. (2022). Sticky PDMP Samplers for Sparse and Local Inference Problems. Statistics and Computing, 33(1), 8.
Bierkens, J., Nyquist, P., and Schlottke, M. C. (2021). Large Deviations for the Empirical Measure of the Zig-Zag Process. The Annals of Applied Probability, 31(6), 2811–2843.
Bierkens, J., and Roberts, G. (2017). A Piecewise Deterministic Scaling Limit of Lifted Metropolis-Hastings in the Curie-Weiss Model. The Annals of Applied Probability, 27(2), 846–882.
Corbella, A., Spencer, S. E. F., and Roberts, G. O. (2022). Automatic Zig-Zag Sampling in Practice. Statistics and Computing, 32(6), 107.
Davis, M. H. A. (1993). Markov models and optimization,Vol. 49. Chapman & Hall.
Goldstein, S. (1951). On Diffusion by Discontinuous Movements, and on the Telegraph Equation. The Quarterly Journal of Mechanics and Applied Mathematics, 4(2), 129–156.
Hardcastle, L., Livingstone, S., and Baio, G. (2024). Averaging polyhazard models using piecewise deterministic monte carlo with applications to data with long-term survivors.
Lewis, P. A. W., and Shedler, G. S. (1979). Simulation of nonhomogeneous poisson processes by thinning. Naval Research Logistics Quarterly, 26(3), 403–413.

Footnotes

  1. 計算過程は省略したが，\(d=2\) の場合と，\(d=3\) の場合を少しやってみると良い．↩︎

  2. 参照測度は，\([d]\) 上のものは計数測度 \(\#\) をとっている．↩︎

  3. 従って，レート関数 \(\lambda\) は連続とする．この関数 \(\gamma_i\) は，\(U\) の情報には依らない追加のリフレッシュ動作を仮定に加える．実際，\(\lambda_i(x,\theta)-\lambda_i(x,F_i(\theta))=\theta_i\partial_iU(x)\) である限り，\(\theta\) と \(F_i(\theta)\) の往来には影響を与えず釣り合っているため，どのような \(\gamma_i\) をとっても，平衡分布には影響を与えない．しかし，高くするごとにアルゴリズムの対称性が上がるため，\(\gamma\equiv0\) とすることが Monte Carlo 推定量の漸近分散を最小にするという (Andrieu and Livingstone, 2021)．(Bierkens et al., 2021) でも同様の洞察がなされている．↩︎

  4. (Bierkens et al., 2019) にある提示の仕方である．Zig-Zag の 2000 単位時間を単純に MALA と比較はできないと筆者も考えるが，ダイナミクスに注目していただきたい．実際，自分で実装してみると，シード値をいじらないと，Zig-Zag は必ずしも 500 単位時間前後でモード \(0\) に戻るわけではない．↩︎

  5. 実は６つ持つ．他の初期値は σ=(Vector(diag(Γ))).^(-0.5); λref=0.0, ρ=0.0↩︎