JavaScriptを有効にしてください

Stochastic Gradient Langevin Dynamicsを理解する

 ·  ☕ 4 min read

はじめに

  • MCMCの一種
  • 目標: ある分布 $\pi(x)$からのサンプリングを行いたい
    1. Metropolis-Hastingsアルゴリズム (MH)
    2. Hamiltonian Monte Carlo (HMC)
    3. Langevin Dynamics (Metropolis-adjusted Langevin Algorithm)
    4. Stochastic Gradient Langevin Dynamics (SGLD)
  • の順に見ていくと理解しやすい

Metropolis-Hastings

  • Metropolis-Hastingsについては既知のもとする
    • 提案分布 $q(z)$を元に判定関数を用いて受容・棄却を行うMCMC
    • cf. GibbsSamplingとHM
    • 提案分布が対称, すなわち $q(z|z_\ast) = q(z_\ast|z)$ならば, 判定関数はただの $\min(1,\frac{\pi(z_*)}{\pi(z)})$となる
      • これを単にメトロポリス法と呼ぶ

Hamiltonian Monte Carlo

  • Hamiltonian Monte Carlo (HMC)
    • ここでハミルトニアンが登場する
      $$H := H(q,p;t) =\mathcal{K}(p)+\mathcal{U}(q)$$

    • HMCのお気持ち

      • 解析力学に則り, ある系における座標と運動量の時間に関する軌道を元に, サンプリングを行う
      • 例えば, $(z,p)$があるとして, 一定時間が経ったときの終点 $(z_\ast,p_\ast)$が得られればサンプル取れるよね
      • このとき, 判定関数がハミルトニアンの差で表せるので嬉しい
        • ハミルトニアン自体は系全体のエネルギーそのものなので, エネルギー保存則が成り立つ.
        • ゆえにハミルトニアンの差分は本来0なので嬉しい.
    • 簡単のため, 質量を $1$としておく ($\mathcal{K}$が下記のように表せるため)
      $$\mathcal{K}(p) = \frac{1}{2}p^\top \mathrm{I} p$$

    • サンプリングしたい分布 $\pi(z) \propto \tilde{\pi}(z)$について, 補助変数 $p$を導入して, $\pi(z,p) := \pi(z) = \pi(z)\pi(p)$と拡張する.

    • また $\pi(p) = \mathcal{N}(p|0,I)$とし, $\ln{\tilde{\pi}(z)} = - \mathcal{U}(z)$とすると,
      $$\pi(p) = \frac{1}{{(2\pi)}^{k/2}}\exp(-\frac{\mathbf{p}^{\top}\mathrm{I}\mathbf{p}}{2})$$

    • から,
      $$\begin{align}\pi(z,p) &= \pi(z)\pi(p) \\ &= \exp(\ln{\pi(z)} + \ln{\pi(p)}) \\
      &\propto \exp(\ln{\tilde{\pi}(z)}-\frac{{p}^{\top}\mathrm{I}{p}}{2}) \\
      &= \exp(\ln{\tilde{\pi}(z)}- \mathcal{U}(p)) \\
      &= \exp(-\mathcal{K}(p)- \mathcal{U}(p)) \\
      &= \exp(-H(z,p))\end{align}$$

    • となる

    • ここで, 提案分布を $\mathcal{N}(0,I)$とすると, Metropolis-Hastings→メトロポリス法となるので,

        1. $p $を $\mathcal{N}(0,I)$からサンプリング
        1. Leapfrog法で $(z,p)$から候補点 $(z_\ast,p_\ast)$を取得 (ステップサイズ $\epsilon$・ステップ数 $L$)
    • とすれば, 判定関数 $A(z_\ast,p_\ast)$は
      $$A(z_\ast,p_\ast) = \min(1,\frac{\pi(z_\ast,p_\ast)}{\pi(z,p)}) = \exp(-H(z_\ast,p_\ast) + H(z,p))$$

    • となる. だが, そもそもハミルトニアンは系全体のエネルギーそのものなので時間変化せず, 常に $H(z_\ast,p_\ast)$と $H(z,p)$は一致するので, 判定関数 $A(z_\ast,p_\ast)$は常に $1$になる

      • めでたしめでたし
      • (実際はLeapfrogの誤差の影響で棄却されることもある)
    • ここからはハイパラのお話

      • Leapfrogではステップサイズ $\epsilon$・ステップ数 $L$というパラメタが存在する
      • $\epsilon$が小さければ探索空間が小さく非効率だし, 大きすぎると棄却される
      • なので, 基本は
        • ① $L$固定で $\epsilon$動かす
        • ② $\epsilon$固定で $L$動かす
      • のどちらかだが, 後者は後者で計算コストが高すぎる
      • したがって, ここらへんは慎重に選択する必要がある

Leapfrog

  • 一旦ここでLeapfrogのおさらいをする
    • ハミルトニアンの正準方程式をLeapfrogで解こうとすると,
      $$\begin{align}p_i (t+\frac{\epsilon}{2}) &= p_i(t) - \frac{\epsilon}{2}\frac{d\mathcal{U}}{dz_i} \\
      z_i(t + \epsilon) &= z_i(t) + \epsilon p_i(t + \frac{\epsilon}{2})\end{align}$$
    • 一つの式にまとめると
      $$z_i(t + \epsilon) = z_i(t) - \frac{\epsilon^2}{2}\frac{d\mathcal{U}}{dz_i} + \epsilon p_i(t)$$

Langevin Dynamics

  • Langevin Dynamics (LD)
    • LDは先程のHMCの特殊例である. 具体的には $L = 1$としたときのHMCである

    • LDにおいて, あるモデルのパラメタ $\theta$に関して, 確率分布 $\pi(\theta)$からのサンプリングを考えてみる

    • 上のLeapfrogの式より
      $$\begin{align}\Delta \theta := \theta_{t+1} - \theta_{t} &= - \frac{\epsilon^2}{2} \nabla \mathcal{U} + \epsilon p_i(t)\\
      &= \frac{\epsilon^2}{2} \nabla \ln{{\pi}(\theta)} + \epsilon p_i(t) (\therefore \ln{\tilde{\pi}(\theta)} = - \mathcal{U}(\theta))\end{align}$$

    • ここで, データ $(X,Y)$について, $\pi(\theta) \leftarrow \pi(\theta|Y,X)$であるとき, ベイズより
      $$\begin{align}\Delta \theta &= \frac{\epsilon^2}{2} \nabla \ln{\pi(\theta|Y,X)} + \epsilon p_t\\
      &= \frac{\epsilon^2}{2}( \nabla \ln{\pi(Y|\theta,X)} + \nabla \ln{\pi(\theta)}) + \epsilon p_t\\
      &= \frac{\epsilon^2}{2}( \sum_n\nabla \ln{\pi(y_n|\theta,x_n)} + \nabla \ln{\pi(\theta)}) + \epsilon p_t\end{align}$$

    • ただし, $p_t$は先程の説明通り $\mathcal{N}(0,I)$からサンプリング

    • なので, これはノイズを加えた勾配法と同じようになる

Stochastic Gradient Langevin Dynamics

  • Stochastic Gradient Langevin Dynamics (SGLD)
    • お待ちかねのSGLDだが, これは単純に上の式の第一項をサブサンプリングするだけ
    • なので更新幅は
      $$\Delta \theta = \frac{\epsilon^2}{2}( \frac{N}{M} \sum_{n=1}^M \nabla \ln{\pi(y_n|\theta,x_n)} + \nabla \ln{\pi(\theta)}) + \epsilon p$$
    • これで, 棄却率がほぼ0 ∧ 計算量激減のMCMCが出来上がったことになる
共有

YuWd (Yuiga Wada)
著者
YuWd (Yuiga Wada)
機械学習・競プロ・iOS・Web