Langevin Monte Carlo法には棄却が必要か

機械学習系のベイズっぽい論文を読んでいると、SGLD [WT11] やSGRLD [PT13] といった文字列を見ることがあります。これらが何かというとマルコフ連鎖モンテカルロ法 (MCMC) の一種で、正規化定数がわからない高次元の確率分布からのサンプリングを得たい場合などに使われます。

アルゴリズムの位置づけとしては、

  • Langevin Monte Carlo (LMC) とか Langevin Dynamics などという名前で呼ばれている既存アルゴリズムがまずあり、
  • それに伴う勾配計算をサブサンプリングを利用して簡略化したもの
    という感じです。勾配降下法 (GD) を確率的勾配降下法 (SGD) に拡張することにインスパイアされているのだったと思います。

LMCのモチベーションとしてよく言われるのは、LMCでは (Metropolis–Hastings法に見られるような) 棄却ステップが必要ないということです。MHは、提案分布がうまく選べていなかったりすると棄却がたくさん発生してしまうことが知られていて、サンプリングに無駄な時間がかかってしまいます。MHの採択確率を（ほとんど）1にするというのは全ベイジアンの夢ですが、ある意味それを達成しているのがLMC系のアルゴリズムです。実際に、SGLDの元論文 [WT11] などでは、「棄却率を低くしたい」という動機がかなり強調して書いてあったりします。

動機だけを聞くとLMCは棄却が要らないという誤解が生じる可能性があるのですが、残念ながらそういうわけではない です。実際には、見た目上は同じアルゴリズムでも、サンプルを得たい分布の性質によって棄却が必要な場合とそうでない場合が存在します（ただし、その理論的な境界線をひとことで説明するのは難しい模様です）。

上のような知識を何度か論文で見かけたことはあったのですが、棄却せずに強行突破すると何が起きてしまうか正直よくわかってなかったです。
なので実際にやってみたというのが本記事の趣旨です。

LMC

LMCの概要をラフに説明します。

$U(\theta)$ を $\mathbb{R}^d$ 上のある関数として、$\exp(-U)$ に比例した密度をもつ分布からのサンプルが欲しいとします。
このとき、
$$
\theta_{k + 1} = \theta_{k} - h \nabla U(\theta_k) + \sqrt{2h} \xi_{k+1}
\tag{1}
$$
という更新式によって $\theta_k$ を得るアルゴリズムをLangevin Monte Carlo法といいます。ただし $h > 0$ はステップ幅で、簡単のため固定とします。$\xi_{k+1} \sim N(0, I)$ は $d$ 次元の標準正規分布に独立に従う確率変数の実現とします。「$U$ の勾配方向に $h$ だけ降りてから分散 $2h$ の正規ノイズで摂動する」という感じです。

例えば、ベイズの事後分布からのサンプリングを行いたいときは、
$$
U(\theta) = - \log p(\theta) - \sum_{i=1}^n \log p(x_i | \theta)
$$
のようなポテンシャルを考えればよいです。

理論的な背景としては、拡散過程
$$
\dd X_t = - \nabla U(X_t) \dd t + \sqrt{2} \dd W_t
\tag{2}
$$
の定常分布が $\pi(\theta) \propto \exp(-U(\theta))$ だということが知られているので、これを離散化 (Euler–Maruyama近似) したものがLMC法に相当します。

ただし、$U$ がどういう関数かによって

  • (2) が定常分布を持たない
  • (2) は定常分布を持つが、その離散化 (1) が定常分布を持たない

というケースが発生し、アルゴリズムがうまく動かない場合があります。それを実験で確かめたいです。

うまく動く場合 (t分布)

1次元で実験します。自由度1のStudentのt分布 の密度関数を $f$ とすると
$$
U(x) = - \log f(x) = \log (1 + t^2) + const.
$$
と書けます。これを (1) 式に突っ込んでシミュレーションを回すとt分布からのサンプルが得られるはずです。

やってみるとこんな感じになります。左が生成されたパス、右がサンプルのヒストグラムです。うまくいっている気がします。

うまく動かない場合1

$$
\pi_{1/2} (x) = \frac{1}{4} \exp(- \sqrt{|x|})
$$
という分布を考えます。Laplace分布よりもっと裾が重い分布になります。

この分布に対して何も考えずにLMCを動かしてみると、見事にパスが発散します。

実は、この分布に形式的に対応する $U = -\log \pi_{1/2}$ を考えると、実は (2) の連続時間の確率過程すら定常分布を持たないことがわかります。こういうときはLMC系の手法を諦める事案だと思われます。

うまく動かない場合2

さらに、もうちょっと微妙な、できそうなのにできないケースがあります。

$$
\pi_4 (x) = \frac{\exp(- x^4)}{2 \Gamma(5/4)}
$$
という分布を考えます。正規分布よりも裾が軽い分布になります。そのまま (1) 式に突っ込むのであれば、 $\nabla \log U(x) = -4 x^3$ とすればよいはずです。

この場合は (2) の確率過程は定常分布を持ち、それは $\pi_4$ に一致します。ですので、LMCも動いて欲しい気がします。

ところが実際にやってみると、どのようにステップ幅を設定しても、最初はうまくいっているかと思いきや途中でいきなり発散して終わります。

実はこのケースでは、Euler–Maruyama法がエルゴード性を持たない (=収束しない) ということが理論的に示唆されています [RT96]。 数学的な条件としては、$\nabla U$ がLipschitzでないことがおそらく問題になっています。

棄却ステップの導入

Metropolis Adjusted Langevin Algorithm (MALA) というアルゴリズムがあります。提案分布を (1) で定義した上で、Metropolis–Hastings と同じような棄却ステップを導入します。詳細はWikipediaや元論文 [RT96]、本 [RC10] などに載っています。

こちらのアルゴリズムを上の $\pi_4 (x)$ に対して導入してみると、パスが途中で発散することはなく、さらに目的の分布に近いヒストグラムがちゃんと得られているように見えます。ただしステップ幅を大きく取りすぎると採択率が下がり、計算が全く終わらなくなります。

PDF版

もうちょっと詳しく書いたPDFを置いときます。

参考文献

[WT11] M. Welling and Y. W. Teh. Bayesian Learning via Stochastic Gradient Langevin Dynamics. In ICML, 2011.
[PT13] S. Patterson and Y. W. Teh. Stochastic Gradient Riemannian Langevin Dynamics on the Probability Simplex. In NIPS, 2013.
[RT96] G. O. Robert and R. L. Tweedie. Exponential convergence of Langevin distributions and their discrete
approximations. Bernoulli, 2(4):341–363, 1996.
[RC10] C. P. Robert and G. Casella. Monte Carlo Statistical Methods. Springer, 2010.