畳庵〜tatamiya practice〜

機械学習・統計モデリング練習帳。読書・source code読みの記録や，実装など。

Leap-frog法はなぜ誤差が蓄積しないのか？を視覚的に理解する〜Hamiltonian Monte Carlo法の補足の補足〜

ベイズ推定 MCMC

久々の更新、お久しぶりです。

最近このようなベイズ統計モデリングの本を共著で書きました：

PyMC3、Pyroで学ぶ統計モデリング発展

PyMC3、Pyroで学ぶ統計モデリング発展

自分の担当箇所は第二章のHamiltonian Monte Carlo法（以下HMC）の節と、AppendixのHMCの詳細です。

HMCは任意の確率分布からサンプリングを行う際に用いられる手法で、 StanやPyMCといった確率的プログラミングの裏の処理にも使われています。

この手法は物理学の考え方をベースとしておりとっつきにくいうえに、 PRMLなどの書籍でも簡単にしか触れられておらず読んでいてピンとこなかったので、自分なりに間を埋めて解説してみました。

kindle unlimitedなら無料で読めますので、もしご興味があれば手にっていただけると嬉しい限りです。

さて、この記事では、HMCで使われるLeap-frog法について、上掲書の補足を行います。

HMC法では、求めたい分布に従うような運動を表現するための力学モデルを作り、 運動方程式を差分法で解くことによってサンプル候補を作り出しています。

そして、運動方程式を数値的に解く際には、leap-frog法という特殊な差分法を用いています。

このleap-frog法は、大体どの書籍でもいきなり出てくるので、面食らったか狐に摘まれたような気分になった方も多数いるかと思います。

また、微分方程式の数値計算手法を勉強したことのある人からすると、なぜEuler法やRunge-Kutta法といった一般的な手法を使わないのだろう？ という疑問が湧いたかと思います。

上掲書の中でも解説しましたが、leap-frog法を用いる理由は、

  • 数値計算誤差が蓄積していかないから
  • 時間反転対象性がMetropolis法を適用するうえで不可欠だから

ということが挙げられます。

このうち前者の数値計算誤差について、ハミルトニアンとは別に「影のハミルトニアン」という別の保存量が存在するから真の軌道に限りなく近い数値軌道が得られる、との補足を入れました。

詳細に突っ込みすぎるため本の中ではそれ以上の言及は控えましたが、 これだけ読んでもあまりにも抽象的でよくわからないぞ、という感想が大半だろうと思いますので、 この場にて「補足の補足」をさせていただきます。

この記事の内容

目的

Leap-frog法ではなぜ誤差が蓄積していないかを、一次元単振動を例に、他の一般的なアルゴリズムとの比較を通して視覚的に理解する

伝えたいこと

  • Leap-frog法で誤差が蓄積しない理由は、真の軌道に近い閉じた軌道を描くからである
    • 刻み幅を小さくすればするだけ厳密な軌道に連続的に近づく
    • 一方で、他の手法ではどんなに刻み幅を小さくしようが軌道が発散or収束してしまう

話さないこと

  • 物理の基本（運動方程式など）
    • ハミルトン方程式については冒頭で紹介した本にも簡単な導入があります。
  • HMCのアルゴリズムの解説
    • 気になる方は、是非PRMLや、冒頭の本を手にお取りください。
  • 一般に影のハミルトニアンが存在することの証明
    • 単振動における影のハミルトニアンの求め方については、最後に補足という形で解説します。

Leap-frog法

復習

まずはleap-frog法についておさらいします。

以下のように、時刻をtとして、ポテンシャルU(x)のもとで動く質量mの質点の運動方程式が与えられたとします。

{ m\ddot{x} = -\frac{dU(x)}{dx} }

これを運動量p=m\dot{x}を用いてハミルトン方程式の形で書くと、 以下のようになります：

{ \begin{aligned} \frac{dp}{dt} &= -\frac{\partial H}{\partial x} =-\frac{dU(x)}{dx} \\ \frac{dx}{dt} &= \frac{\partial H}{\partial p} = \frac{p}{m}. \end{aligned} }

ハミルトニアンは

{ \begin{aligned} H(p, x) &= \frac{1}{2m}p^2 + U(x) \end{aligned} }

のように表され、時間によらず一定です。

一部の例外を別として、一般的にはこれを厳密に解く術は存在しないので、差分化して数値的に解きます。

このときleap-frog法では以下のように差分化します。

{ \begin{aligned} p(t+\varepsilon/2) &= p(t) - \frac{\varepsilon}{2}\left. \frac{d U(x)}{dx}\right|_{x=x(t)} \\ x(t+\varepsilon) &= x(t) + \varepsilon \frac{p(t+\varepsilon/2)}{m} \\ p(t+\varepsilon) &= p(t+\varepsilon/2) - \frac{\varepsilon}{2}\left. \frac{d U(x)}{dx}\right|_{x=x(t+\varepsilon)}. \end{aligned} }

\varepsilonは刻み幅です。運動量と位置を交互に更新していくのが特徴です。

単振動におけるleap-frog法の数値軌道

冒頭で述べたように、leap-frog法では誤差が蓄積していかない点でEuler法やRunge-Kutta法と比べて都合が良いのですが、このことを単純なモデルを使って実際に数値軌道を見ながらお話ししていこうと思います。

単振動のハミルトニアンとハミルトン方程式は以下のように書けます。

{ \begin{aligned} H(p, x) &= \frac{1}{2m}p^2 + \frac{1}{2}kx^2 \\ \frac{dp}{dt} &= -kx \\ \frac{dx}{dt} &= \frac{p}{m} \end{aligned} \tag{1} \label{1} }

この場合の解は簡単に求めることができて、ハミルトニアンが一定値（Eとおく）をとることから、(p, x)空間上では以下のような楕円曲線上を動いていきます：

{ \left(\frac{p}{\sqrt{2mE}}\right)^2 + \left(\frac{x}{\sqrt{2E/k}}\right)^2 = 1 }

Eの値は、初期値によって決まります。

一方でこれをleap-frog法で解く場合、以下のような差分方程式で時間発展させていくことになります。

{ \begin{aligned} p(t+\varepsilon/2) &= p(t) - \frac{\varepsilon}{2}kx(t) \\ x(t+\varepsilon) &= x(t) + \varepsilon \frac{p(t+\varepsilon/2)}{m} \\ p(t+\varepsilon) &= p(t+\varepsilon/2) - \frac{\varepsilon}{2}kx(t+\varepsilon) \end{aligned} \tag{2} \label{2} }

Leap-frog法で差分化して得られた軌道を以下の図に青丸でプロットしました。 この解軌道もまた、灰色の楕円曲線上に乗るのですが、赤線で描いた厳密に解いた場合の解曲線と比べると、やや横にひしゃげているのが分かります。

f:id:tatamiyatamiyatatatamiya:20201218084013p:plain
leap-frog法で計算した単振動の軌道

なお、このleap-frog法の解曲線（灰色）は、刻み幅\varepsilonを小さくしていくことで限りなく厳密な解曲線（赤）に近づいていきます。

このように、leap-frog法の数値解は、真の解曲線を連続的に変形させた曲線の上を、真の軌道から近づきつつ離れつつを繰り返しながら動いていくのです。

このことを式の上からも確認してみます。

今回は単振動でモデルがシンプルなため、青色の数値解が乗っている楕円曲線（灰色）を求めることができます（導出方法は記事の最後に補足します）。

{ \begin{aligned} \tilde{H}(p, x) = \frac{1}{2m}p^2 + \frac{1}{2}k\left(1 - \frac{k\varepsilon^2}{4m}\right)x^2 = const \end{aligned} }

これが「影のハミルトニアン」と呼ばれるものの正体で、Leap-frog法の場合は式\eqref{1}で表した真のハミルトニアンに代わってこちらの量が時間によらず一定となっています。[^1] \varepsilon=0を代入してみると、真のハミルトニアンと一致していることが確認できるかと思います。

Euler法、Runge-Kutta法との比較

これに対して、常微分方程式の数値計算法として一般的なEuler法とRunge-Kutta法でも単振動を計算してみて、その時の軌道の違いを比較してみます。

Euler法

{ \begin{aligned} p(t+\varepsilon) &= p(t) - \varepsilon \left. \frac{d U(x)}{dx}\right|_{x=x(t)} \\ x(t+\varepsilon) &= x(t) + \varepsilon \frac{p(t)}{m} . \end{aligned} }

f:id:tatamiyatamiyatatatamiya:20201218084239p:plain
Euler法で計算した単振動の軌道

Euler法では真の軌道から螺旋を描くように遠ざかっていき、最終的に無限遠まで飛んでいってしまいます。 このことは、刻み幅を小さくしようが変わりません。

Runge-Kutta法

Runge-Kutta法では、少々複雑ですが以下のようにして差分化します:

{ \begin{aligned} p(t+\varepsilon) &= p(t) + \frac{\varepsilon}{6}(k_1^{(p)}+2k_2^{(p)}+2k_3^{(p)}+k_4^{(p)})\\ x(t+\varepsilon) &= x(t) + \frac{\varepsilon}{6}(k_1^{(x)}+2k_2^{(x)}+2k_3^{(x)}+k_4^{(x)})\\ \end{aligned} }

ただし、

{ \begin{aligned} k_1^{(p)} &= - \left. \frac{dU(x)}{dx}\right|_{x=x(t)} \quad &k_1^{(x)} &= \frac{p(t)}{m}\\ k_2^{(p)} &= - \left. \frac{dU(x)}{dx}\right|_{x=x(t)+\frac{\varepsilon}{2}k_1^{(x)}} \quad &k_2^{(x)} &= \frac{p(t) + \frac{\varepsilon}{2}k_1^{(p)}}{m}\\ k_3^{(p)} &= - \left. \frac{dU(x)}{dx}\right|_{x=x(t)+\frac{\varepsilon}{2}k_2^{(x)}} \quad &k_3^{(x)} &= \frac{p(t) + \frac{\varepsilon}{2}k_2^{(p)}}{m}\\ k_4^{(p)} &= - \left. \frac{dU(x)}{dx}\right|_{x=x(t)+\varepsilon k_3^{(x)}} \quad &k_4^{(x)} &= \frac{p(t) + \varepsilon k_3^{(p)}}{m}\\ \end{aligned} }

f:id:tatamiyatamiyatatatamiya:20201218084152p:plain
Runge-Kutta法で計算した単振動の軌道

こちらでは、刻み幅\varepsilonが一定値より大きいと発散、小さい場合は逆に一点(0, 0)へと収束していってしまいます。 刻み幅\varepsilonを上手く調整すれば厳密解と一致するのですが、一般にそのようなパラメータを探すのは困難です。

まとめ

以上からLeap-frog法では、真のハミルトニアンに近い保存量＝影のハミルトニアンが存在するため、真の軌道を連続的に変形した軌道を描くことが分かりました。

このような特別な性質により誤差が蓄積しないため、外力のない運動方程式を解く際には、Euler法やRunge-Kutta法よりLeap-frog法の方が好まれるのです。

このことは、運動方程式を拠り所とするHMCにおいてLeap-frog方が用いられる理由の1つでもあります（主たる理由は時間反転対称性ですが、こちらについては議論を割愛します）。

ちなみに、上記のようにハミルトニアンとは別の保存量を作る差分手法は一般にSymplectic数値積分法と呼ばれ、その中でLeap-frog法は2次の差分法として位置付けられます。

HMCでLeap-frog法より高次のSymplectic法が使われない理由は、単純に計算コストが見合わないためと思われます。

Symplectic数値積分法について詳しく知りたい方は、参考文献をご参照ください。

参考文献

補足：単振動の影のハミルトニアンの導出

差分式の行列表記

単振動をLeap-frog法で解く際の差分式（式\eqref{2}）を、少し書き直します。

{ \begin{aligned} x(t+\varepsilon) &= x(t) + \varepsilon \frac{p(t+\varepsilon/2)}{m} \\ &= x(t) + \varepsilon \frac{p(t)}{m} - \frac{1}{m} \frac{\varepsilon^2}{2}kx(t)\\ &=\left(1 - \frac{k\varepsilon^2}{2m}\right)x(t) +\varepsilon \frac{p(t)}{m} \\ p(t+\varepsilon) &= p(t+\varepsilon/2) - \frac{\varepsilon}{2}kx(t+\varepsilon)\\ &= p(t) - \frac{\varepsilon}{2}kx(t)- \frac{\varepsilon}{2}k \left(1 - \frac{k\varepsilon^2}{2m}\right)x(t) - \frac{k\varepsilon^2}{2} \frac{p(t)}{m}\\ &= \left(-k\varepsilon + \frac{k^2\varepsilon^3}{4m} \right)x(t) +\left(1 - \frac{k\varepsilon^2}{2m}\right)p(t) \end{aligned} }

さらにこれをベクトルと行列を用いて表現すると、以下のようになります：

{ \begin{aligned} \begin{pmatrix} x(t+\varepsilon)\\ p(t+\varepsilon) \end{pmatrix} = \begin{pmatrix} 1 - \frac{k\varepsilon^2}{2m} & \frac{\varepsilon}{m} \\ -k\varepsilon + \frac{k^2\varepsilon^3}{4m} & 1 - \frac{k\varepsilon^2}{2m} \end{pmatrix} \begin{pmatrix} x(t) \\ p(t) \end{pmatrix} \end{aligned}. }

固有値分解による保存量の表現

このとき出てきた行列を、

{ \begin{aligned} A &= \begin{pmatrix} 1 - \frac{k\varepsilon^2}{2m} & \frac{\varepsilon}{m} \\ -k\varepsilon + \frac{k^2\varepsilon^3}{4m} & 1 - \frac{k\varepsilon^2}{2m} \end{pmatrix}\\ &= U^{-1} \begin{pmatrix} \lambda_+ & 0\\ 0 & \lambda_- \end{pmatrix} U \end{aligned} }

のように固有値分解します。

なお、この時の固有値の積\lambda _ {+} \lambda _ {-}が1になることは、行列Aの行列式が|A|=1となることから分かります。

式(??)に左から行列Uをかけ、さらに新しい座標系

{ \begin{aligned} \begin{pmatrix} \tilde{x}(t)\\ \tilde{p}(t) \end{pmatrix} = U \begin{pmatrix} x(t)\\ p(t) \end{pmatrix} \end{aligned} }

を用いることで以下のように表すことができます：

{ \begin{aligned} \begin{pmatrix} \tilde{x}(t+\varepsilon)\\ \tilde{p}(t+\varepsilon) \end{pmatrix} = \begin{pmatrix} \lambda_+ & 0\\ 0 & \lambda_- \end{pmatrix} \begin{pmatrix} \tilde{x}(t)\\ \tilde{p}(t) \end{pmatrix} \end{aligned}. }

これにより(\tilde{x}, \tilde{p})の時間発展が独立な2つの差分式

{ \begin{aligned} \begin{cases} \tilde{x}(t+\varepsilon)&=\lambda_+\tilde{x}(t)\\ \tilde{p}(t+\varepsilon)&=\lambda_-\tilde{p}(t) \end{cases} \end{aligned} }

で表現できることがわかりました。

ここで固有値の積が \lambda _ {+} \lambda _ {-} = 1であったことを思い出すと、以下のような関係が成り立つことがわかります：

{ \begin{aligned} \tilde{x}(t+\varepsilon)\tilde{p}(t+\varepsilon) = \tilde{x}(t)\tilde{p}(t) \end{aligned} }

つまり、\tilde{x}(t)\tilde{p}(t)は時間発展に対して一定の値をとる、保存量となっているのです。

保存量の座標逆変換

それでは、この保存量を元の座標系(x, p)で表し直します。

固有値\lambda_\pmに対応する行列Aの固有ベクトルを(s _ x ^ {\pm}, s _ p ^ {\pm}) ^ Tのように表すと、 固有値分解の際に現れた行列Uとの間には、以下のような関係が成り立ちます：

{ \begin{aligned} U^{-1} &= \begin{pmatrix} s_x^+ & s_x^-\\ s_p^+ & s_p^- \end{pmatrix}\\ \therefore U &= \frac{1}{s_x^+s_p^- - s_p^+ s_x^-} \begin{pmatrix} s_p^- & -s_x^-\\ -s_p^+ & s_x^+ \end{pmatrix}. \end{aligned} }

これを使うと、先ほどの保存量\tilde{x}\tilde{p}は、元の座標系(x, p)を用いて以下のように表せます：

{ \begin{aligned} \tilde{x}\tilde{p} &= \frac{1}{(s_x^+ s_p^- - s_p^+s_x^-)^2} \left \lbrace -s_x^+s_x^- p^2 + (s_x^+ s_p^- + s_p^+ s_x^-)xp - s_p^+ s_p^- x^2\right \rbrace \\ &= - s_x^+s_x^- \left \lbrace p^2 - \left( \frac{s_p^+}{s_x^+} + \frac{s_p^-}{s_x^-}\right)xp + \frac{s_p^+}{s_x^+}\frac{s_p^-}{s_x^-} x^2\right \rbrace. \end{aligned} }

ここで、行列Aと固有値・固有ベクトルの間には

{ \begin{aligned} A \begin{pmatrix} s_x^\pm\\ s_p^\pm \end{pmatrix} = \lambda_\pm \begin{pmatrix} s_x^\pm\\ s_p^\pm \end{pmatrix} \end{aligned} }

の関係が成り立っていなくてはいけないので、このうちx成分だけ取り出すと、

{ \begin{aligned} \left( 1 - \frac{k\varepsilon^2}{2m} \right)s_x^\pm + \frac{\varepsilon}{m} s_p^\pm = \lambda_\pm s_x^\pm\\ \therefore \frac{s_p^\pm}{s_x^\pm} = \frac{m}{\varepsilon}\left(\lambda_\pm -1 + \frac{k\varepsilon^2}{2m} \right). \end{aligned} }

これを用いることで、

{ \begin{aligned} \frac{s_p^+}{s_x^+} + \frac{s_p^-}{s_x^-} &= \frac{m}{\varepsilon}\left( \lambda_+ + \lambda_- - 2 + \frac{k\varepsilon^2}{m} \right)\\ &= \frac{m}{\varepsilon}\left \lbrace 2 \left( 1 - \frac{k\varepsilon^2}{2m} \right) - 2 + \frac{k\varepsilon^2}{m} \right \rbrace\\ &=0\\ \frac{s_p^+}{s_x^+} \frac{s_p^-}{s_x^-} &= \frac{m^2}{\varepsilon^2} \left \lbrace \lambda_+ \lambda_- - \left( 1 - \frac{k\varepsilon^2}{2m}\right)(\lambda_+ + \lambda_-) + \left( 1 - \frac{k\varepsilon^2}{2m}\right)^2 \right \rbrace \\ &= \frac{m^2}{\varepsilon^2} \left \lbrace 1 - \left( 1 - \frac{k\varepsilon^2}{2m}\right)^2 \right \rbrace\\ &=mk \left( 1 - \frac{k\varepsilon^2}{4m}\right) \end{aligned} }

が得られます。 このとき、

{ \begin{aligned} \lambda_+ \lambda_- &= 1\\ \lambda_+ + \lambda_- &= 2 \left( 1 - \frac{k\varepsilon^2}{2m}\right)\\ \end{aligned} }

という関係を用いました（後者は行列Aの対角成分の和）。

以上を代入すると、

{ \begin{aligned} \tilde{x}\tilde{p} &= - s_x^+s_x^- \left \lbrace p^2 + mk \left( 1 - \frac{k\varepsilon^2}{4m}\right)x^2\right \rbrace\\ &\propto \frac{1}{2m}p^2 + \frac{1}{2} k \left( 1 - \frac{k\varepsilon^2}{4m}\right) x^2 = const \end{aligned} }

が得られます。

係数の部分は本質的でないので、最終的に\varepsilon=0で真のハミルトニアンに一致するように調整しました。