Population Annealing with Emcee

Introduction

ある確率分布からサンプリングをしたいとき, 確率分布が解析的に得られない場合にはモンテカルロ法をもちいることが多くあります. 広く活用されているマルコフ連鎖モンテカルロ法では粒子の座標を時間発展させ, 滞在時間が確率分布に従うようにコントロールすることでサンプリングをおこないます. 一方で, たくさんの粒子を用意して, 粒子の分布がターゲットとなる確率分布に従うように座標を時間発展させるというアプローチもあります. こういった手法は粒子フィルタ (particle filter) や逐次モンテカルロ法 (sequential Monte Carlo) と呼ばれています. ここでは粒子法と表現します. 粒子法にもいくつか種類がありますが, ここでは Hukushima & Iba (2003) によって提案された Population Annealing を簡略化してデモンストレーションしてみます.

Target distribution

粒子法の効果を見るためにターゲットとなる確率分布は強い多峰性を持つ分布とします. ここでは 2 次元平面で定義された 2 つのコンパクトな Gaussian 分布の混合をターゲットとしました. 混合比率は 2:1 としています. 対数尤度は以下のように定義できます.

def log_prob(xy):
    ''' target log-likelihood '''
    gaussian = lambda x, μ, σ: np.exp(- 0.5 * ((x - μ) / σ)**2) / σ

    return np.log(
        2 * gaussian(xy[:, 0], -3.0, 0.10) * gaussian(xy[:, 1], 0.0, 0.10)
        + 2 * gaussian(xy[:, 0], +3.0, 0.20) * gaussian(xy[:, 1], 0.0, 0.40)
    )

対数尤度を可視化してみましょう. \(x,y \in [-6, 6]\) の範囲に 10,000 個のデータ点をランダムに分布させて, 各点における対数尤度の値をカラーマップで示しました. 2 つのコンパクトな分布が \(x\) 方向に並んでおり, 2 つの分布の間を行き来するためには確率分布関数が非常に小さい領域を通過しなければいけません.

Sampling by emcee

試しに天文学で広く使われている emcee というパッケージでサンプリングをしてみます. emcee はマルコフ連鎖をつなげていく一般的な MCMC ですが, 粒子の位置を発展させるときに他の粒子の情報を使用する Ensemble Sampler という独特な手法を採用しています.

以下では粒子 (walker) を 50 個用意して領域 \(x,y \in [-10, 10]\) 内に一様に分布させました. 100 ステップは burn-in として 5000 ステップ分サンプリングをしました. 最終的に 50×5000 個のデータがサンプルされたことになります.

ndim = 2
rng = np.random.default_rng(42)

nwalkers = 50
p0 = rng.uniform(-10, 10, size=(nwalkers, 2))

sampler = emcee.EnsembleSampler(
    nwalkers=nwalkers, ndim=ndim, log_prob_fn=log_prob)

state = sampler.run_mcmc(p0, 100)
sampler.reset()
state = sampler.run_mcmc(state, 5000)

さて, 得られた分布を確認してみます. 左側にはサンプルされたデータを \(xy\)-平面に散布図で示しています. 先程確認したように 2 つの山ができていることが確認できます. それでは \(y\) 方向を周辺化して導出した \(x\) 方向の確率分布 (ヒストグラム) を確認してみます. 青色の実線は emcee でサンプリングした結果, オレンジ色の破線は今回仮定した確率分布です. それぞれの分布の形状は適切に近似できていますが, 2 つの分布の相対的な大きさは真の値から大きくずれています.

これはターゲットとなる確率分布が強い多峰性を持っていることに起因しています. 片方の分布からもう片方の分布に移動する確率が極端に低いため, 2 つの分布のどちらが尤度が高いかを比較できていないと考えられます.

実際に emcee でサンプルされた粒子の軌跡 (トレース) を確認してみます. 12 個の walker の軌跡を以下に示しました. 1 つの walker を除いてすべて \(x>0\) の分布をうろついていました. また, 5000 ステップで 2 つの分布の間を移動したサンプルはひとつもありませんでした.

emcee は多数の粒子をもちいて次のステップを決める特徴的なアプローチを採用していますが, 確率分布の強い多峰性を乗り越えることは難しいようです. こうした分布から MCMC でサンプルするためには, パラレルテンパリングやマルチカノニカルマルコフ連鎖モンテカルロ法といった手法によって確率が非常に低い領域をまたいで効率よく状態を混合する必要があります.

Population annealing

それでは粒子法を試してみましょう. Population annealing では最適化手法のひとつアニーリング法 (simulated annealing) と同様に温度をパラメタとして確率分布の系列を定義します. 温度パラメタ \(T\) を取り入れた対数尤度は以下のように定義できます.

def log_prob_t(xy, T):
    ''' target log-likelihood with temperature '''
    return log_prob(xy) / T

Population annealing によるサンプリング手順は以下のとおりです.¹

サンプルしたい数 \((N)\) だけ粒子を用意して初期状態を与えます.²
温度の系列 \(\{T_k\}\) を用意します. 高温からスタートして最終的に \(T_k \to 1\) となるように定義します.
ある温度 \(T_k\) で定義した対数尤度 \(\mathcal{L}(x | T_{k})\) をターゲットとしてマルコフ連鎖を進めます.
粒子の分布が温度 \(T_k\) に対して平衡状態に達したら温度を次の系列に進めます. このとき
1. 粒子 \(i\) の温度 \(T_k\) における対数尤度を \(\mathcal{L}(x_i | T_k)\) とします.
2. この座標のまま温度 \(T_{k+1}\) に下げたときの粒子 \(i\) の対数尤度を \(\mathcal{L}(x_i | T_{k+1})\) とします.
3. 粒子 \(i\) の重みを \(w_i = \exp\left(-(\mathcal{L}(x_i | T_{k}) - \mathcal{L}(x_i | T_{k+1}))\right)\) とします.
4. \(N\) 個の粒子を重み \(\{w_i\}\) で再サンプリングします.
3 に戻って \(T_k = 1\) になるまで繰り返します.

この手順 (3–5) を 1 ステップ実行するための関数を定義します. マルコフ連鎖モンテカルロ法を計算する必要があるため, その部分は先ほどと同様 emcee を採用することにしました.

def resample(state, log_prob_fn, T, nstep=10, pool=None):
    ''' resample population

    Parameters:
        state: emcee Monte Carlo state
        log_prob_fn: log likelihood function (with temperature)
        T: temperature
        nstep: number of inner MCMC steps
        pool: Pool instance for multiprocessing
    '''
    if not isinstance(state, emcee.state.State):
        nwalker, ndim = state.shape
        sampler = emcee.EnsembleSampler(
            nwalkers=nwalker, ndim=ndim,
            log_prob_fn=log_prob_fn,
            kwargs={'T': T}, pool=pool)
        state = sampler.run_mcmc(state, nstep)
        sampler.reset()

    loc, log_p, seed = state
    log_q = log_prob_fn(loc, T=T) - log_p
    nwalker, ndim = loc.shape

    p = np.nan_to_num(np.exp(log_q), 0)
    p = p / p.sum()
    idx = np.random.choice(nwalker, nwalker, p=p)
    loc = loc[idx, :]

    sampler = emcee.EnsembleSampler(
        nwalkers=nwalker, ndim=ndim,
        log_prob_fn=log_prob_fn, kwargs={'T': T}, pool=pool)
    return sampler.run_mcmc(loc, nstep)

それでは Population Annealing を実行してみます. 10,000 個サンプルするために walker を 10,000 だけ用意しました. 温度は \(T_0 = 1000\) から始めてゆっくりと低下させていきます. 各温度では MCMC ステップを 10 回だけ進めています.

with Pool() as pool:
    state = rng.uniform(-10, 10, size=(10000, 2))
    state = resample(state, log_prob_t, 2000, nstep=100, pool=pool)
    states = []
    for T in tqdm(1000 / np.logspace(1, 3, 50)):
        state = resample(state, log_prob_t, T, pool=pool)
        states.append(state)

以下にPopulation annealing によってサンプルされた結果を図示します. 左側には \(xy\)-平面での散布図を示しました. emcee のときと同様に 2 つのコンパクトな分布に分かれていることがわかります. 右側には \(y\) 方向に周辺化した \(x\) 方向の分布を示しています. Population annealing によってサンプルされた分布 (緑色実線) は真の分布 (オレンジ色破線) をよく再現しています.

Population annealing によって粒子の分布がどのように更新されていったのかを可視化してみます. 各温度での MCMC 更新が完了したあとの \(x\) 方向のヒストグラムを動画として以下に示しました. 高温 (\(T \simeq 1000\)) では確率分布は薄く広がっており, 2 つの分布のあいだで粒子のやり取りが十分に可能になっています. その後, 緩やかに温度が低下して \(T \simeq 300\) 程度では 2 つの分布はほぼ分断されています. 最終的に \(T = 1\) まで温度が下がると, もともとターゲットとしていた分布に収束していきます.³

このように単純な MCMC ではサンプルすることが難しい確率分布を Population annealing では比較的シンプルな仕組みで再現することができます. また, Population annealing では (リサンプリングの時を除いて) 各粒子は独立に計算できるため高い並列計算効率が期待できます. 今回は MCMC の更新に emcee を使用したためにあまり恩恵がありませんでしたが, 巨大な問題に対するスケーラビリティという点でもおもしろい手法だと思います.

ここではかなり簡略化した手順を示します. 正確な手続きについては論文 Hukushima & Iba (2003) を参照してください. ↩
このステップは emcee における walker の初期化と似ていますが, 用意する粒子の数がかなり多くなります. ↩
サンプリングというよりはアニーリング法によって確率分布関数を最適化しているというイメージに近いのかもしれません. ↩