Auto-Encoding Variational Bayes

論文解説

問題設定と動機

連続潜在変数 $\mathbf{z}$ を持つ有向確率モデルは、生成モデリング、教師なし表現学習、周辺推論に有用です。しかし、最尤法やベイズ手法による学習は、以下の状況において極めて大きな計算上の困難に直面します：

1. 事後分布の難解性（Posterior Intractability）： 非線形な関係性（深層ニューラルネットワークなど）が存在する場合、周辺尤度の積分 $p_{\mathbf{\theta }}(\mathbf{x})=\int p_{\mathbf{\theta }}(\mathbf{z})p_{\mathbf{\theta }}(\mathbf{x}|\mathbf{z})d\mathbf{z}$ や事後分布 $p_{\mathbf{\theta }}(\mathbf{z}|\mathbf{x})$ は解析解を持たず、標準的なExpectation-Maximization (EM) や平均場変分ベイズ法が適用できません。
2. 大規模データセット： Monte Carlo EMなどの伝統的なサンプリング手法は、データセット全体の各データ点に対して低速な反復MCMCサンプリングを必要とするため、大規模なミニバッチ最適化には計算コストが高すぎます。

標準的かつ素朴なスコア関数勾配推定量 ${\nabla }_{\mathbf{\varphi }}{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z})}[f(\mathbf{z})]$ は非常に分散が大きく、確率的勾配上昇法による安定した最適化が困難でした。

主要なアイデア

• Reparameterization Trick（再パラメータ化トリック）： 潜在変数 $\mathbf{z}$ をパラメータ化された分布 $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ から直接サンプリングする代わりに、独立したノイズ変数 $\mathbf{ϵ}\sim p(\mathbf{ϵ})$ の決定論的変換 $\mathbf{z}=g_{\mathbf{\varphi }}(\mathbf{ϵ},\mathbf{x})$ として書き直します。これにより、標準的な誤差逆伝播法を用いて $\mathbf{\varphi }$ に関する微分が可能になり、分散を劇的に抑えることができます。
• Amortized Inference（認識モデル）： データ点ごとに変分パラメータを個別に最適化するのではなく、エンコーダとなるニューラルネットワーク $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ が $\mathbf{x}$ を条件として変分パラメータ（潜在ガウス分布の平均や分散など）を直接予測します。
• Auto-Encoding Variational Bayes (AEVB)： 再パラメータ化された変分下界に対して確率的勾配上昇法を適用することで、生成パラメータ $\mathbf{\theta }$（確率的デコーダ）と変分パラメータ $\mathbf{\varphi }$（確率的エンコーダ）を同時に最適化します。

手法

変分目的関数は、証拠下界（ELBO: Evidence Lower Bound）$\mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})$ を最大化します： $$\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)})\ge \mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})=-D_{KL}(q_{\mathbf{\varphi }}(\mathbf{z}|{\mathbf{x}}^{(i)})\parallel p_{\mathbf{\theta }}(\mathbf{z}))+{\mathbb{E}}_{q_{\mathbf{\varphi }}(\mathbf{z}|{\mathbf{x}}^{(i)})}\left[\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)}|\mathbf{z})\right]$$

${\mathbf{ϵ}}^{(i,l)}\sim p(\mathbf{ϵ})$ を用いて再パラメータ化トリック ${\mathbf{z}}^{(i,l)}=g_{\mathbf{\varphi }}({\mathbf{ϵ}}^{(i,l)},{\mathbf{x}}^{(i)})$ を適用すると、一般的なSGVB推定量（Estimator A）は以下のようになります： $${\tilde{\mathcal{L}}}^A(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})=\frac{1}{L}\sum _{l=1}^L\left[\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)},{\mathbf{z}}^{(i,l)})-\log q_{\mathbf{\varphi }}({\mathbf{z}}^{(i,l)}|{\mathbf{x}}^{(i)})\right]$$

事前分布 $p_{\mathbf{\theta }}(\mathbf{z})=\mathcal{N}(0,\mathbf{I})$ および近似事後分布 $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})=\mathcal{N}(\mathbf{\mu },{\mathbf{\sigma }}^2\mathbf{I})$ がガウス分布の場合、KLダイバージェンスは解析的に計算でき、推定量Bが得られます： $$\mathcal{L}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})\simeq \frac{1}{2}\sum _{j=1}^J\left(1+\log (({\sigma }_j^{(i)}{)}^2)-({\mu }_j^{(i)}{)}^2-({\sigma }_j^{(i)}{)}^2\right)+\frac{1}{L}\sum _{l=1}^L\log p_{\mathbf{\theta }}({\mathbf{x}}^{(i)}|{\mathbf{z}}^{(i,l)})$$ ここで ${\mathbf{z}}^{(i,l)}={\mathbf{\mu }}^{(i)}+{\mathbf{\sigma }}^{(i)}\odot {\mathbf{ϵ}}^{(l)}$ かつ ${\mathbf{ϵ}}^{(l)}\sim \mathcal{N}(0,\mathbf{I})$ です。

サイズ $M$ のミニバッチに対する全体の下界推定量は次の通りです： $${\tilde{\mathcal{L}}}^M(\mathbf{\theta },\mathbf{\varphi };{\mathbf{X}}^M)=\frac{N}{M}\sum _{i=1}^M\tilde{\mathcal{L}}(\mathbf{\theta },\mathbf{\varphi };{\mathbf{x}}^{(i)})$$ 実際には、ミニバッチサイズを $M=100$ 程度に設定すれば、データ点あたりのサンプル数は $L=1$ で十分機能します。

結果

• モデル設定: MNIST（ベルヌーイMLPデコーダ、隠れユニット500）およびFrey Face（ガウスMLPデコーダ、隠れユニット200）で生成モデルを評価しました。最適化にはAdagradを用い、学習率は $\{0.01,0.02,0.1\}$、重み減衰を適用しました。
• 下界の最適化: Frey FaceおよびMNISTにおける様々な潜在次元（$N_{\mathbf{z}}\in \{2,5,10,20,400\}$）において、AEVBはwake-sleepアルゴリズムよりも大幅に速く収束し、より高い下界を達成しました。
• 過学習耐性: 事前分布に向かう解析的KLダイバージェンス項の正則化効果により、$\mathbf{z}$ の次元数を増やしても過学習が発生しませんでした。
• 周辺尤度の推定: HMCに基づく周辺尤度推定量を用いた3次元潜在空間のMNIST実験において、AEVBはwake-sleepよりも速く収束し、Monte Carlo EM（MNIST全体への適用には遅すぎる）の現実性を大きく凌駕しました。
• 実行時間: Intel Xeon CPU（約40 GFLOPS）上で、100万サンプルあたり約20〜40分で訓練が完了しました。

限界と未解決の課題

• 連続潜在変数への依存: 再パラメータ化トリックはパラメータに関する微分可能性に直接依存しているため、離散潜在変数の分布へ直接適用することは容易ではありません。
• 近似のギャップ: $q_{\mathbf{\varphi }}(\mathbf{z}|\mathbf{x})$ に対角ガウス分布を仮定しているため、真の事後分布が多峰性であったり、強い相関を持っていたりする場合に表現能力が制限されます。
• 完全な変分ベイズへの拡張: 本実験の評価は $\mathbf{\theta }$ のMAP/MLパラメータ推定に焦点を当てていました。グローバルパラメータ $\mathbf{\theta }$ に対する変分推論は付録で定式化されているものの、実験的な検証は今後の課題とされました。