【Edward】MCMCの数学的基礎からStochastic Gradient Langevin Dynamicsの実装まで

こんにちは。初めまして。
データ分析部新入りのmathetake(@mathetake)と申します。

先日個人ブログでこんなエントリを書いた人です:

mathetake.hatenablog.com

そんなこんなでTwitter就活芸人(?)として活動(?)してましたが、これからは真面目に頑張っていこうと思います。

今日はみんな大好きベイズモデリングおいて、事後分布推定に欠かせないアルゴリズム(群)の一つである*1

通称MCMCに関するエントリです。より具体的に、

MCMCの意義(§1.)から始め、マルコフ連鎖の数学的な基礎(§2.,3.,4.)、MCMCの代表的なアルゴリズムであるMetropolis-Hastings法(§5.)、その例の１つである*2Langevin Dynamics(§6.)、そして(僕の中で)絶賛大流行中のライブラリEdwardを使ってより発展的(?)なアルゴリズムであるStochastic Gradient Langevin Dynamicsの説明＆実装(§7.,8.)していきたいと思います。

今までのデータ分析ブログのエントリと少しテイストが違うかもしれませんが、お楽しみいただけたら幸いです。

細心の注意を払ってはいますが、数学的正しさを保証する記事ではありませんので、詳細が気になる方は§Referenceにある資料を御覧ください。

はじめに
離散値マルコフ連鎖
マルコフ連鎖収束定理とMCMC
連続値マルコフ連鎖の場合
Metropolis-Hastings法
Langevin Dynamics
Stochastic Gradient Langevin Dynamics(SGLD)
EdwardでのSGLDの実装
Reference
Appendix

はじめに

本題に入る前にまず、MCMCの意義について軽く触れておきます。

ベイズ統計におけるモデルパラメータの推定には事後分布からのサンプリングがかかせません。即ち、モデルのパラメータを ${\theta \in \Theta}$ ,データの集合を ${\mathcal{D}}$ として、ベイズの定理から得られる次のような確率分布

${ p \left( \theta \ | \ \mathcal{D} \right) = \dfrac{p \left( \mathcal{D} \ | \ \theta \right)p(\theta) }{ \displaystyle \int_\Theta p \left( \mathcal{D} \ | \ \theta \right) p(\theta) d\theta } }$

からサンプリングする必要がありますが、一般に分母( 正規化定数 )は解析的に求められず、サンプリングは困難です。

そこを、マルコフ連鎖の収束定理等の数学的基礎理論を使って上手くサンプリングする手法がMCMCで、広く応用されています。*3

ベイズ統計の数学的背景については、僕の個人ブログのエントリを御覧ください

mathetake.hatenablog.com

離散値マルコフ連鎖

${ \mathcal{X}:=\{X^0,X^1,X^2,\dots,X^t,\dots \} }$ を離散な集合 ${ \{1,2,3,4,..,k\}(k \in \mathbb{N}) }$ に値を取る、離散な確率変数の集合とします。確率的な値の時間発展だと考え、確率過程と呼ぶことにします。

(Definition 1.) ${\mathcal{X} }$ がマルコフ連鎖であるとは、全ての ${ \{x^{s}\}_{s=0}^{t+1} \subset \{1,2,\dots,k \}}$ に対して
${ \mathbb{P} \left(X^{t+1} = x^{t+1} \ | \ X^t=x^t ,\dots X^0=x^0 \right) = \mathbb{P}\left(X^{t+1} =x \ | \ X^t=x^t \right)\ \ \ \cdots (1) }$
が成立する事。

感覚的には、次の時刻における値の分布は現在の値のみで決まり、それ以前の値には影響されない確率過程の事です。

数学的便宜上、次の定義も用意しておきます。

(Definition 2.) マルコフ連鎖 ${\mathcal{X} }$ が斉時的であるとは、
${p_{i,j}^t:= \mathbb{P}\left(X^{t+1} =j \ | \ X^{t} = i \right) }$
が時間に依存しない事。またこの時行列 ${ p_{i,j} := p_{i,j}^t }$ を 遷移行列 と呼ぶ。

つまり状態遷移の確率が時間に依らず一定なマルコフ連鎖の事で、状態遷移を表す行列を遷移行列 としています。

ずっと数学の話になってしまいましたが、最後に一つだけMCMCに関わる重要な命題とその系を述べておきます。

(Proposition 3.) ${\mathcal{X} }$ を斉時的なマルコフ連鎖とする。この時 ${X^t}$ の分布は遷移行列 ${ p_{i,j} }$ と ${X^0}$ の分布(初期分布)により完全に決定される。

(Proof) ${X^0,X^1,\dots}$ の分布を ${\pi^0, \pi^1,\dots,}$ とし、 ${\pi^t_j}$ で確率変数 ${X^t}$ が値 ${j}$ を取る確率とする。マルコフ連鎖の性質(1)より
${ \begin{eqnarray} \pi^1_j &=& \sum_{i=1}^k \mathbb{P} \left( X^1=j \ | \ X^0=i \right) \mathbb{P} \left( X^0 =i \right) \\ &=& \sum_{i=1}^k \pi^0_i p_{i,j} \ \ \ \cdots (2) \end{eqnarray}}$
また、
${ \begin{eqnarray} \pi^t_j &=& \sum_{i=1}^k \mathbb{P} \left( X^t=j \ | \ X^{t-1}=i \right) \mathbb{P} \left( X^{t-1} =i \right) \\ &=& \sum_{i=1}^k \pi^{t-1}_i p_{i,j} \ \ \ \cdots (3) \end{eqnarray} }$
が成立するので、帰納的に示される □

この命題の系として次が得られます

(Corollary 4.) 初期分布 ${\pi^{0}}$ と遷移行列 ${T:= (p_{i,j})}$ を与えることで、式(2),(3)により斉時的マルコフ連鎖が得られる。

マルコフ連鎖収束定理とMCMC

あと少しMCMCの説明まで辿りつきます。もう少々数学にお付き合いください。

§1. で紹介したように、サンプリングしたい確率分布 ${ \pi =(\pi_i) }$ が手元(?)にあるとします。

(Definition 5.)遷移行列が ${T:= (p_{i,j})}$ により与えられる、斉時的マルコフ連鎖 ${ \mathcal{X}}$ が ${\pi}$ を不変分布に持つとは
$\sum_{i=1}^k \pi_i p_{i,j} = \pi_j \ \ \ ( \ \forall j \in \{ 1,\dots,k\})$
が成立する事。行列の式で書けば
${ \pi T = \pi }$
が成立すること。この時 ${\pi}$ を ${ \mathcal{X}}$ の不変分布と言う。

さて、MCMCの肝となる定理は次のものです

(Theorem 6.(離散値マルコフ連鎖の収束定理))
${\pi }$ を不変分布に持つ斉時的マルコフ連鎖 ${ \mathcal{X}}$ が (i) 非周期的 かつ (ii) 既約*4 である時、マルコフ連鎖は不変分布 ${\pi }$ に収束する。即ち

$\lim_{t \rightarrow \infty} \mathbb{P}\left( X^t = i \right) = \pi_i$
が成立する。

ここまで来てやっと、MCMCの定義を与えることができます：

(Definition 7.) MCMC(Markov chain Monte Carlo)とは、サンプリングしたい確率分布 ${\pi}$ を不変分布とするような既約で非周期的なマルコフ連鎖を構築＆サンプリングするアルゴリズムの事。

Theorem 6.により、MCMCにより生成されるサンプルの列 ${ \{x^t\}^\infty_{t=0} }$ は確率分布 ${\pi}$ からのサンプルに収束し、目的を達成することができます。

既約性と非周期性を満たすようなマルコフ連鎖を構築するのはそんなに難しくはありません、が、サンプリングしたい確率分布 ${\pi}$ を不変分布とするようなマルコフ連鎖を構築するのは一般に困難です。

そこでよく用いられるのが*5、詳細釣り合い条件と呼ばれる不変分布を持つための十分条件です:

(Definition 8.)遷移行列が ${T:= (p_{i,j})}$ により与えられる、斉時的マルコフ連鎖 ${ \mathcal{X}}$ が不変分布 ${\pi}$ に対して詳細釣り合い条件を満たすとは

$\pi_i \ p_{i,j} = \pi_j \ p_{j,i} \ \ \ \ ( \forall \ i, j \in \{ 1,\dots,k\})$

が成立する事。またこの時、マルコフ連鎖 ${ \mathcal{X}}$ は ${\pi}$ を不変分布に持つ。

連続値マルコフ連鎖の場合

今まで簡単のため、離散な値を持つマルコフ連鎖について話をしてきましたが、連続な確率変数の話に一般化する事ができます。

連続値( 便宜上 ${\mathbb{R}^d}$ 値とする )なマルコフ連鎖 ${\mathcal{X} = \{ X^0,X^1,\dots,X^t,\dots \}}$ が斉時的である時、離散値の場合の推移行列に対応する推移核 ${ T: \mathbb{R}^d \times \mathbb{R}^d \rightarrow \mathbb{R} }$ を

${ \mathbb{P} \left( X^{t} \in A \ | \ X^{t-1} =x \right) = \displaystyle \int_A T(x,y) dy \ \ \ \ \ \ (A \subset \mathbb{R}^d ) }$

を満たすものとして定めます。感覚的には現時刻で値 ${x}$ を取る時、次の時刻の分布を表す密度関数です。

また、 ${ \mathcal{X}}$ が分布 ${\pi}$ を不変分布に持つとは

$\int_{\mathbb{R}^d} \pi(x)T(x,y)dx = \pi(y)$

を満たすことであり、詳細釣り合い条件は

${ \pi(x) \ T(x,y) = \pi(y) \ T(y,x) \ \ \ \cdots (4)}$

で与えられます。

(Theorem 9.(連続値マルコフ連鎖の収束定理))
${\pi }$ を不変分布に持つ斉時的マルコフ連鎖 ${ \mathcal{X}}$ が (i) 非周期的 かつ (ii) 既約である時、 ${ \mathcal{X}}$ はtotal variation distanceの意味で ${\pi}$ に収束する。即ち、

$\lim_{t \rightarrow \infty} \sup_{A \subset \mathbb{R}^d } \left| \pi^t(A) - \pi(A) \right| \ = \ 0$

が成立する。

Metropolis-Hastings法

MCMCの代表的なアルゴリズム(群)である、Metropolis-Hastings法(以下M-H法)について説明します。

M-H法では、各 ${x}$ に対して提案分布と呼ばれる確率分布 ${q(x,y)dy}$ を用意し、採択確率と呼ばれる確率

${\alpha(x,y):= \min \left\{ 1, \dfrac{\pi(y)q(y,x) }{\pi(x)q(x,y)} \right\} }$

を準備します。そして推移核 ${T(x,y)}$ を

${T(x,y) = \alpha(x,y)q(x,y) + A(x)\delta_{x,y} }$

として定義＆適当な初期分布 ${\pi^0 }$ を与えることで斉時的マルコフ連鎖 ${\mathcal{X}_{MH} }$ を考えます。(ここで ${\delta_{x,y} }$ は ${x=y }$ の時に1でそれ以外の時に0を取る関数、 ${A(x)}$ は正規化定数を与える関数。)

(Theorem 10.) 斉時的マルコフ連鎖 ${\mathcal{X}_{MH} }$ は ${\pi}$ を定常分布に持つ

(Proof) 定義から
${ \begin{eqnarray} \pi(x)q(x,y)\alpha(x,y) &=& \pi(x)q(x,y) \min \left\{ 1, \dfrac{\pi(y)q(y,x) }{\pi(x)q(x,y)} \right\}\\ &=& \min \left\{ \pi(x)q(x,y), \pi(y)q(y,x) \right\}\\ &=& \pi(y)q(y,x) \min \left\{ \dfrac{\pi(x)q(x,y)}{\pi(y)q(y,x)}, 1 \right\}\\ &=& \pi(y)q(y,x)\alpha(y,x) \end{eqnarray}}$
が従うので、クロネッカーのデルタの定義から(4)式が成立し詳細釣り合い条件が満たされる。□

Theorem 10.だけでは保証されない既約性や非周期性を満たすような推移核の具体的な設計は重要な課題ですが、ここではそのような性質が満たされ Theorem 6.(or 7.) が成立すると仮定しましょう。*6

この時マルコフ連鎖 ${\mathcal{X}_{MH}}$ からサンプリングする次のようなアルゴリズムを Metropolis-Hastings法と呼びます。

(Metropolis-Hastings法)
(1) 初期分布 ${ \pi^0}$ から ${x^0}$ をサンプリングする:
${x^0 \sim \pi^0 }$
(2) ${t =0,1,2,3, \dots}$ に対して、以下を実行する
　(i) 標準一様分布から乱数 ${u}$ を生成する:
${u \sim U(0,1) }$
　(ii) ${y}$ を提案分布 ${q(x^t, y)}$ からサンプリングする:
${y \sim q(x^t, y) }$
　(ii) 次の式により"次の点" ${x^{t+1} }$ を決める:
${x^{t+1} \ := \ \displaystyle \begin{cases} y \ \ \ \ \ if \ \ u < \alpha(x^t, y) \\ x^t \ \ \ \ otherwise \end{cases} }$

このアルゴリズムが、実際に上で与えた ${\mathcal{X}_{MH} }$ からサンプリングしている事は明らかでしょう。

§0. で述べたように採択確率の計算には目標の分布 ${\pi(x))}$ の正規化定数が必要ない事に注目して下さい。

Langevin Dynamics

M-H法を実際に実行するためには、提案分布 ${q(x,y)dy}$ を定義する必要があります。

ここではその例として、Langevin Dynamics法(Metropolis-adjusted Langevin Algorithm)(以下LD法) を紹介します。

LD法では提案分布を次のように定義します:

${q(x,y) := \mathcal{N} \left( x + \dfrac{\epsilon}{2} \dfrac{\partial \log \pi }{\partial x}(x), \ \ \epsilon I \right) }$

ここで ${\mathcal{N}\left( \mu, \Sigma \right)}$ は平均 ${\mu}$ , 分散 ${\Sigma}$ に従う正規分布で、 ${\epsilon}$ はstep size(またはlearning rate)と呼ばれるハイパーパラメータです。

ベイズモデルの事後分布に適応する場合において、決定論的に眺めると、LD法は事前分布で正規化した、ノイズ入り勾配降下法のようなものであると解釈することができます。

実際、事後分布推定において、 ${\pi}$ は観測データを ${\mathcal{D = \{d_i \}}_{i=1}^N }$ として

${ \pi(\theta) = p \left( \theta \ | \ \mathcal{D} \right) = \dfrac{p(\theta) \prod_{i=1}^N p\left( d_i \ | \ \theta \right) }{ \displaystyle \int_\Theta p \left( \mathcal{D} \ | \ \theta \right) p(\theta) d\theta } }$

で与えられ、その対数微分は

$\dfrac{\partial}{\partial \theta} \log \pi \left( \theta \right) =\dfrac{\partial}{\partial \theta} \log p \left( \theta \right) + \dfrac{\partial}{\partial \theta} \sum_{i=1}^N \log p\left(d_i \ | \ \theta \right) \ \ \ \cdots (5)$

のように計算できるので、上のような解釈ができます。

実はLG法は、Stanで有名になったHamiltonian Monte Carlo法の特別なケースと等価になっているので、気になる方は [1]や[2] を御覧ください。

Stochastic Gradient Langevin Dynamics(SGLD)

ここまでMCMCの数学的基礎からM-H法、そしてその具体例としてLG法を紹介しました。

LG法の問題点として、(5) 式の計算量がサンプル数が増えるほど膨大になっていく点があります。

近年はビッグデータと呼ばれるバズワードもあるように、サンプル数＆パラメータ数が巨大なセッティングでモデリングする事が多いのでこのままLG法を適用する事はできません。*7

その問題点を克服するサンプリング手法として、ここで紹介するのが Stochastic Gradient Langevin Dynamics法 [8](以下SGLD法)です。

SGLD法では次のように初期分布からサンプルしたパラメータを更新していきます:

まず、次の２つの条件

$\sum_{t=0}^\infty \epsilon_t^2 < \infty \ , \ \ \ \ \ \ \sum_{t=0}^\infty \epsilon_t = \infty$

を満たす数列 ${ \{\epsilon_t \}_{t=0}^\infty }$ を用意し、パラメータサンプル ${ \{ \theta_t \}_{t=0}^\infty }$ を次の式によって取得して行きます

${\begin{eqnarray} \displaystyle & \theta_{t+1} &\leftarrow \ \ \ \theta_t + \dfrac{\epsilon_t}{2} \dfrac{\partial L}{\partial \theta} (\theta_t) + \eta_t \ , \ \ \ \eta_t \sim \mathcal{N}\left( 0 \ ,\epsilon_t \right) \\ \ \\ &\dfrac{\partial L}{\partial \theta} (\theta_t)& = \dfrac{\partial \log p}{\partial \theta_t} \left( \theta_t \right) + \dfrac{N}{\left| \mathcal{S_t} \right|}\sum_{d \in \mathcal{S_t} } \dfrac{\partial \log p \left(d \ | \ \theta \right) }{\partial \theta}(\theta_t) \end{eqnarray} }$

ここで、 ${\mathcal{S_t}}$ はデータ ${ \{d_i\}_{i=1}^N }$ からランダムに抽出された ${N}$ より十分小さいミニバッチとします。

注意として、このアルゴリズムに対応するマルコフ連鎖は斉時的ではないので、上述の収束定理は適用できません。

ですが、例えば [9]で収束性に関する解析がされています。

EdwardでのSGLDの実装

最後に確率モデリング用ライブラリEdwardを用いて、SGLDをベイズ的線形回帰に適用してみようと思います。

Edwardの詳しい使い方は公式チュートリアルまたは次の論文

[1701.03757] Deep Probabilistic Programming

[1610.09787] Edward: A library for probabilistic modeling, inference, and criticism

をご覧ください。また質問等ありましたら@mathetakeまで気軽にリプライorDMください。

まず各種ライブラリをimportします。

import numpy as np
import tensorflow as tf
import edward as ed
from edward.models import Normal, Empirical
import time

次にデータセットを用意します。

N = 20000  # サンプル数
D = 50  # 特徴量の次元
N_ITER = 10000  # MCMCのiteration
MINI_BATCH_SIZE = 2500  #ミニバッチのサイズ 

# toy dataset. 切片=0はなし.
def build_toy_dataset(N, D, noise_std=0.1):
    w = np.random.randn(D).astype(np.float32)
    X = np.random.randn(N, D).astype(np.float32)
    Y = np.dot(X, w) + np.random.normal(0, noise_std, size=N)
    return w, X, Y

# データ生成。観測値のノイズの分散は既知とする。
w_true, X_data, Y_data = build_toy_dataset(N, D)

# ミニバッチを返す関数 
def next_batch(mini_batch_size=128):
    indexes = np.random.randint(N, size=mini_batch_size)
    return X_data[indexes], Y_data[indexes]

モデルを構築し, 推論のためのインスタンスを作ります。

# 観測データを挿入するためのデータを収めるplaceholder
x = tf.placeholder(tf.float32, [MINI_BATCH_SIZE, D])
y_ph = tf.placeholder(tf.float32, [MINI_BATCH_SIZE])

w = Normal(mu=tf.zeros(D), sigma=tf.ones(D))
b = Normal(mu=tf.zeros(1), sigma=tf.ones(1))
y = Normal(mu=ed.dot(x, w) + b, sigma=tf.ones(MINI_BATCH_SIZE)*0.1)

# 経験分布をposteriorの近似に使う
qw = Empirical(params=tf.Variable(tf.random_normal([N_ITER, D])))
qb = Empirical(params=tf.Variable(tf.random_normal([N_ITER, 1])))

# SGLD法用インスタンス
SGLD = ed.SGLD(latent_vars={w: qw, b: qb}, data={y: y_ph})

最後に推論を実行します。

# 推論GO 
# data辞書にはobservedな確率変数の観測データを送る。
# xの値は確率変数ではないので、updateの際feed_dictで送る。
SGLD = ed.SGLD(latent_vars={w: qw, b: qb}, data={y: y_ph})
SGLD.initialize(scale={y: float(N) / MINI_BATCH_SIZE}, step_size=0.00001, n_iter=N_ITER)

start = time.time()
init = tf.global_variables_initializer()
init.run()
for _ in tqdm(range(N_ITER)):
    X_batch, Y_batch = next_batch(MINI_BATCH_SIZE)
    _ = SGLD.update(feed_dict={x: X_batch, y_ph: Y_batch})
elapsed_time = time.time() - start
print("elapsed_time:{}".format(elapsed_time))

実行結果ですが、ミニバッチで勾配計算をしない通常のLD法と比較してみました:

アルゴリズム	実行時間(s)	MSE of W	Estimated b
SGLD法	26.4	${4.9 \times 10^{-5}}$	0.0020
LD法	55.4	${5.3 \times 10^{-7}}$	0.0001

といった感じです。精度はLD法には劣りますが、実行時間は約２倍短いと言った感じです。

より大きなデータセット＆複雑なモデルの場合、この差はより顕著になるでしょう。

Reference

◯ハミルトニアンモンテカルロ法について

[1] [1206.1901] MCMC using Hamiltonian dynamics

[2] ハミルトニアンＭＣＭＣの解説（PDF)

◯M-H法について

[3] The Metropolitan-Hastings Algorithm and Extensions

[4] マルコフ連鎖モンテカルロ法入門

[5] Simple conditions for the convergence of the Gibbs sampler and Metropolis-Hastings algorithms

[6] Markov Charin For Exploring Posterior Distribution

[7] Markov Chains on Continuous State Space

◯Stochastic Gradient Langevin Dynamicsについて

[8] Bayesian Learning via Stochastic Gradient Langevin Dynamics

[9] Approximation Analysis of Stochastic Gradient Langevin Dynamics by using Fokker-Planck Equation and Ito Process

Appendix

(Definition) 斉時的マルコフ連鎖 ${ \mathcal{X}}$ が非周期的であるとは、任意の ${i \in \{1,\dots,k\} }$ に対して
${gcd \left\{ n \ | \ \mathbb{P} \left( X^n=i \ | \ X^0 =i \right) > 0 \right\} = 1 }$
が成立する事。

(Definition) 斉時的マルコフ連鎖 ${ \mathcal{X}}$ が既約であるとは, 任意の ${i, j \in \{1,\dots,k \}}$ に対してある ${t_{i,j} \in \mathbb{N} }$ が存在して
${ \mathbb{P}\left(X^{t_{i,j}}=j \ | \ X^0 =i \right) >0 }$
が成立する事。