変分推論(Variational Inference; VI)は、解析的に扱いにくい事後分布 $p(\mathbf{z} \mid \mathbf{x})$ を扱いやすい近似分布 $q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})$ で近似する枠組みである。古典的な変分推論では近似分布の族を手動で設計する必要があったが、ニューラルネットワーク(NN)を用いることで、複雑な事後分布を柔軟かつスケーラブルに近似できるようになった。本稿では、この枠組みの数理的基礎から、変分オートエンコーダ(VAE)を中心とした実装・拡張まで詳述する。
$N$ 個の観測データ $\mathbf{X} = \{\mathbf{x}_i\}_{i=1}^N$ が与えられたとする。各 $\mathbf{x}_i \in \mathbb{R}^D$ に対し、連続潜在変数 $\mathbf{z}_i \in \mathbb{R}^L$($L \ll D$)が存在すると仮定する生成モデルを考える:\[p_{\boldsymbol{\theta}}(\mathbf{x}、 \mathbf{z}) = p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})\, p(\mathbf{z})\]ここで $p(\mathbf{z})$ は潜在変数の事前分布(典型的には標準正規分布 $\mathcal{N}(\mathbf{0}、 \mathbf{I})$)、$p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})$ はパラメータ $\boldsymbol{\theta}$ を持つデコーダ(生成ネットワーク)である。
目標はデータの対数周辺尤度(エビデンス)\[\log p_{\boldsymbol{\theta}}(\mathbf{x})= \log \int p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})\, p(\mathbf{z})\, d\mathbf{z}\]を最大化する生成パラメータ $\boldsymbol{\theta}$ を学習することである。この積分は一般に解析的に求まらず($p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})$ が NN で表現される場合は特に)、真の事後分布\[p_{\boldsymbol{\theta}}(\mathbf{z} \mid \mathbf{x})= \frac{p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})\, p(\mathbf{z})}{p_{\boldsymbol{\theta}}(\mathbf{x})}\]も計算困難(intractable)となる。
変分推論では、近似分布族 $\{q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})\}$ を導入し、真の事後分布との KL ダイバージェンスを最小化する:\[\boldsymbol{\phi}^* = \arg\min_{\boldsymbol{\phi}}\, \mathrm{KL}\!\left(q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}) \, \|\, p_{\boldsymbol{\theta}}(\mathbf{z} \mid \mathbf{x})\right)\]KL ダイバージェンスは直接計算できないが、対数周辺尤度との関係式\[\log p_{\boldsymbol{\theta}}(\mathbf{x})= \mathcal{L}(\boldsymbol{\theta}, \boldsymbol{\phi}; \mathbf{x}) + \mathrm{KL}\!\left(q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}) \, \|\, p_{\boldsymbol{\theta}}(\mathbf{z} \mid \mathbf{x})\right)\]から、KL の最小化はELBO(Evidence Lower BOund)の最大化と等価であることがわかる:\[\mathcal{L}(\boldsymbol{\theta}、 \boldsymbol{\phi}; \mathbf{x})= \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})}\!\left[\log p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})\right] - \mathrm{KL}\!\left(q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}) \, \|\, p(\mathbf{z})\right)\]第1項は再構成誤差(reconstruction term)、第2項は近似事後分布を事前分布へ引き寄せる正則化項(regularization term)と解釈できる。$\mathrm{KL} \geq 0$ より $\mathcal{L} \leq \log p_{\boldsymbol{\theta}}(\mathbf{x})$ が常に成立する。
ELBO を展開すると、完全データ対数尤度の期待値としても表現できる:\[\mathcal{L}(\boldsymbol{\theta}, \boldsymbol{\phi}; \mathbf{x})= \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})}\!\left[\log p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z}) - \log q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})\right]\]$N$ 個のデータ全体に対する目的関数は\[\mathcal{L}(\boldsymbol{\theta}, \boldsymbol{\phi})= \sum_{i=1}^{N} \mathcal{L}(\boldsymbol{\theta}, \boldsymbol{\phi}; \mathbf{x}_i)\]であり、$\boldsymbol{\theta}$ と $\boldsymbol{\phi}$ を同時に勾配上昇法で最適化する。
古典的な変分推論では、各データ点 $\mathbf{x}_i$ ごとに変分パラメータ $\boldsymbol{\phi}_i$ を個別に最適化する(局所変分パラメータ)。これは $N$ が大きい場合に計算コストが膨大となり、また新データへの汎化も困難である。
ニューラルネットワークによる変分推論(償却推論; Amortized Inference)では、この局所最適化を NN による関数近似で置き換える:\[q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}) = \mathcal{N}\!\left(\mathbf{z};\, \boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x})\, \mathrm{diag}(\boldsymbol{\sigma}^2_{\boldsymbol{\phi}}(\mathbf{x}))\right)\]ここで $\boldsymbol{\mu}_{\boldsymbol{\phi}} : \mathbb{R}^D \to \mathbb{R}^L$ および $\boldsymbol{\sigma}^2_{\boldsymbol{\phi}} : \mathbb{R}^D \to \mathbb{R}^L_{>0}$はパラメータ $\boldsymbol{\phi}$ を共有するニューラルネットワーク(エンコーダまたは推論ネットワーク)の出力である。
エンコーダの構造は以下のように表される。入力 $\mathbf{x}$ を受け取り、中間表現 $\mathbf{h}$ を経て平均と対数分散を出力する:\[\mathbf{h} = f_{\boldsymbol{\phi}}(\mathbf{x}), \quad\boldsymbol{\mu} = W_\mu \mathbf{h} + \mathbf{b}_\mu, \quad\log \boldsymbol{\sigma}^2 = W_\sigma \mathbf{h} + \mathbf{b}_\sigma\]$f_{\boldsymbol{\phi}}$ は多層パーセプトロン(MLP)、畳み込みネットワーク(CNN)など、データの構造に応じた任意のアーキテクチャを採用できる。償却推論の利点は、推論を単一の順伝播に帰着させることで計算コストを $O(1)$ とし、未観測の新データに対しても即座に事後分布を評価できる汎化性能を持つことである。
ELBO の $\boldsymbol{\phi}$ に関する勾配\[\nabla_{\boldsymbol{\phi}}\, \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})}\!\left[f(\mathbf{z})\right]\]は、期待値の分布 $q_{\boldsymbol{\phi}}$ 自体が $\boldsymbol{\phi}$ に依存するため、直接モンテカルロ推定ができない。スコア関数推定量(REINFORCE)を用いる方法もあるが、分散が非常に大きく学習が不安定になる。
再パラメータ化トリック(Reparameterization Trick)は、確率変数のサンプリングをパラメータから切り離す変数変換である。$q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}) = \mathcal{N}(\boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x})、 \mathrm{diag}(\boldsymbol{\sigma}^2_{\boldsymbol{\phi}}(\mathbf{x})))$ の場合、\[\mathbf{z} = \boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x}) + \boldsymbol{\sigma}_{\boldsymbol{\phi}}(\mathbf{x}) \odot \boldsymbol{\varepsilon}, \quad \boldsymbol{\varepsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})\]と変換すると、期待値の確率変数は $\boldsymbol{\varepsilon}$ に関するものとなり、$\boldsymbol{\phi}$ の依存性がすべて決定論的な写像に移る:\[\mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})}\!\left[f(\mathbf{z})\right]= \mathbb{E}_{\boldsymbol{\varepsilon} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})}\!\left[f\!\left(\boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x}) + \boldsymbol{\sigma}_{\boldsymbol{\phi}}(\mathbf{x}) \odot \boldsymbol{\varepsilon}\right)\right]\]これにより、$\boldsymbol{\phi}$ に関する勾配を通常の誤差逆伝播法で計算できる:\[\nabla_{\boldsymbol{\phi}}\, \mathbb{E}_{q_{\boldsymbol{\phi}}}\!\left[f(\mathbf{z})\right]= \mathbb{E}_{\boldsymbol{\varepsilon}}\!\left[\nabla_{\boldsymbol{\phi}}\, f\!\left(\boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x}) + \boldsymbol{\sigma}_{\boldsymbol{\phi}}(\mathbf{x}) \odot \boldsymbol{\varepsilon}\right)\right]\]右辺はモンテカルロサンプル $\boldsymbol{\varepsilon}^{(s)} \sim \mathcal{N}(\mathbf{0}、 \mathbf{I})$ により低分散に推定できる。実用上は $S = 1$ のサンプルで十分な精度が得られることが多い。
再パラメータ化が可能な分布族は広く、ガンマ分布、ディリクレ分布、フォン・ミーゼス分布なども適切な変換(例:Implicit Reparameterization Gradients)により扱える。
変分オートエンコーダ(Variational Autoencoder; VAE)(Kingma & Welling、 2013)は、上述のエンコーダ・デコーダ・再パラメータ化トリックを統合した生成モデルである。
モデルの構成は以下の通りである:
VAE の ELBO は、近似事後分布が正規分布、事前分布が標準正規分布の場合、KL 項が解析的に計算できる:\[\mathrm{KL}\!\left(\mathcal{N}(\boldsymbol{\mu}、 \mathrm{diag}(\boldsymbol{\sigma}^2)) \, \|\, \mathcal{N}(\mathbf{0}、 \mathbf{I})\right)= \frac{1}{2} \sum_{l=1}^{L}\!\left(\mu_l^2 + \sigma_l^2 - \log \sigma_l^2 - 1\right)\]再構成誤差項はモンテカルロ推定により計算される。データが連続値の場合はガウス尤度\[\log p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z}) = -\frac{1}{2}\left\|\mathbf{x} - \boldsymbol{\mu}_{\boldsymbol{\theta}}(\mathbf{z})\right\|^2 + \text{const}\]すなわち平均二乗誤差(MSE)に対応し、2値データの場合はベルヌーイ尤度\[\log p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z}) = \sum_{d=1}^{D}\!\left[x_d \log f_d(\mathbf{z}) + (1-x_d)\log(1-f_d(\mathbf{z}))\right]\]すなわちバイナリクロスエントロピーに対応する。
学習アルゴリズムは、ミニバッチ $\mathcal{M} \subset \{1, \ldots, N\}$ を用いた確率的勾配上昇法(SGVB 推定量)となる:\[\tilde{\mathcal{L}}(\boldsymbol{\theta}、 \boldsymbol{\phi})= \frac{N}{|\mathcal{M}|} \sum_{i \in \mathcal{M}} \left(-\mathrm{KL}\!\left(q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}_i) \, \|\, p(\mathbf{z})\right) + \frac{1}{S}\sum_{s=1}^{S} \log p_{\boldsymbol{\theta}}\!\left(\mathbf{x}_i \mid \mathbf{z}_i^{(s)}\right)\right)\]ここで $\mathbf{z}_i^{(s)} = \boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x}_i) + \boldsymbol{\sigma}_{\boldsymbol{\phi}}(\mathbf{x}_i) \odot \boldsymbol{\varepsilon}^{(s)}$、$\boldsymbol{\varepsilon}^{(s)} \sim \mathcal{N}(\mathbf{0}、 \mathbf{I})$ である。
平均場正規近似(対角共分散)は柔軟性に限界がある。より表現力の高い近似族として以下の手法が提案されている。
フルコバリアンス正規分布:\[q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}) = \mathcal{N}(\boldsymbol{\mu}_{\boldsymbol{\phi}}(\mathbf{x}), \Sigma_{\boldsymbol{\phi}}(\mathbf{x}))\]共分散行列 $\Sigma_{\boldsymbol{\phi}}$ をコレスキー因子 $L_{\boldsymbol{\phi}}$($\Sigma = L L^\top$)で表現し、$O(L^2)$ パラメータで潜在変数間の相関を捉える。
正規化流(Normalizing Flows):単純な基底分布 $q_0(\mathbf{z}_0)$(例:標準正規分布)に対し、可逆変換列 $f_1, f_2, \ldots, f_T$ を合成して表現力の高い分布を構成する:\[\mathbf{z}_T = f_T \circ \cdots \circ f_1(\mathbf{z}_0), \quad \mathbf{z}_0 \sim q_0\]変数変換公式(change-of-variables formula)により、\[\log q_T(\mathbf{z}_T) = \log q_0(\mathbf{z}_0) - \sum_{t=1}^{T} \log \left|\det \frac{\partial f_t}{\partial \mathbf{z}_{t-1}}\right|\]ヤコビアン行列式が効率的に計算可能な変換族(RealNVP、Glow、Masked Autoregressive Flow など)が多数提案されており、任意の複雑な分布を近似できる。
IAF(Inverse Autoregressive Flow):自己回帰変換を逆方向に適用することでサンプリングを並列化する。$\mathbf{z}_0 \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ から出発し、\[z_{t,l} = \mu_{t,l}(\mathbf{z}_{t-1, < l}) + \sigma_{t,l}(\mathbf{z}_{t-1, < l}) \cdot z_{t-1, l}\]の変換を $T$ ステップ適用する。サンプリングは並列だが密度評価は逐次となる。
潜在変数を複数層に積み重ねた階層的変分オートエンコーダでは、生成モデルと推論モデルをそれぞれ\[p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z}_1, \ldots, \mathbf{z}_L)= p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z}_1)\prod_{l=1}^{L-1} p_{\boldsymbol{\theta}}(\mathbf{z}_l \mid \mathbf{z}_{l+1}) \cdot p(\mathbf{z}_L)\]\[q_{\boldsymbol{\phi}}(\mathbf{z}_1, \ldots, \mathbf{z}_L \mid \mathbf{x})= q_{\boldsymbol{\phi}}(\mathbf{z}_L \mid \mathbf{x}) \prod_{l=1}^{L-1} q_{\boldsymbol{\phi}}(\mathbf{z}_l \mid \mathbf{z}_{l+1}, \mathbf{x})\]と定義する(ボトムアップ推論)。ELBO は\[\mathcal{L} = \mathbb{E}_{q}\!\left[\log p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z}_1)\right]- \sum_{l=1}^{L} \mathrm{KL}\!\left(q_{\boldsymbol{\phi}}(\mathbf{z}_l \mid \cdot) \, \|\, p_{\boldsymbol{\theta}}(\mathbf{z}_l \mid \cdot)\right)\]と分解される。NVAE や Very Deep VAE などの手法では、残差接続や重み付き KL 等の工夫により非常に深い階層構造を学習している。
$\beta$-VAE(Higgins et al。、 2017)は、KL 正則化項に重み $\beta > 1$ を付加した目的関数\[\mathcal{L}_\beta(\boldsymbol{\theta}, \boldsymbol{\phi}; \mathbf{x})= \mathbb{E}_{q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})}\!\left[\log p_{\boldsymbol{\theta}}(\mathbf{x} \mid \mathbf{z})\right] - \beta\, \mathrm{KL}\!\left(q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}) \, \|\, p(\mathbf{z})\right)\]を最大化する。$\beta$ を大きくするほど、潜在空間に強い独立性の圧力がかかり、解釈可能な分離表現(disentangled representation)の学習が促進される。これは情報ボトルネック原理の観点からも理解でき、潜在変数 $\mathbf{z}$ と入力 $\mathbf{x}$ の間の相互情報量 $I(\mathbf{z}; \mathbf{x})$ を圧縮しながら、ラベル $\mathbf{y}$ に関する情報 $I(\mathbf{z}; \mathbf{y})$ を保持する:\[\min_{\boldsymbol{\phi}}\, I(\mathbf{z}; \mathbf{x}) - \gamma\, I(\mathbf{z}; \mathbf{y})\]
IWAE(Importance Weighted Autoencoder)(Burda et al., 2015)は、複数サンプルによる重要度重み付き推定量を用いて、ELBO より厳密な対数尤度の下界を構成する:\[\mathcal{L}_K(\boldsymbol{\theta}, \boldsymbol{\phi}; \mathbf{x})= \mathbb{E}_{\mathbf{z}^{(1)}, \ldots, \mathbf{z}^{(K)} \sim q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x})} \!\left[\log \frac{1}{K} \sum_{k=1}^{K} \frac{p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z}^{(k)})}{q_{\boldsymbol{\phi}}(\mathbf{z}^{(k)} \mid \mathbf{x})}\right]\]$K = 1$ のとき通常の ELBO に一致し、$K \to \infty$ のとき真の対数尤度 $\log p_{\boldsymbol{\theta}}(\mathbf{x})$ に収束する:\[\log p_{\boldsymbol{\theta}}(\mathbf{x}) \geq \mathcal{L}_{K+1} \geq \mathcal{L}_K \geq \mathcal{L}_1 = \mathcal{L}\]重要度重みは\[\tilde{w}_k = \frac{p_{\boldsymbol{\theta}}(\mathbf{x}, \mathbf{z}^{(k)})}{q_{\boldsymbol{\phi}}(\mathbf{z}^{(k)} \mid \mathbf{x})}, \quad w_k = \frac{\tilde{w}_k}{\sum_{j=1}^{K} \tilde{w}_j}\]と定義され、$\boldsymbol{\phi}$ の勾配は $\sum_k w_k \nabla_{\boldsymbol{\phi}} \log q_{\boldsymbol{\phi}}(\mathbf{z}^{(k)} \mid \mathbf{x})$ の形で計算される。これはより正確な推論を可能にするが、$K$ が大きいと重要度重みの分散が増大し、勾配推定量の質が低下する(weight degeneracy)ことが知られている。
VAE の学習において重要な概念として、集約事後分布(aggregated posterior)がある:\[q_{\boldsymbol{\phi}}(\mathbf{z}) = \frac{1}{N} \sum_{i=1}^{N} q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}_i)\]生成モデルの質の観点から、集約事後分布が事前分布 $p(\mathbf{z})$ と一致することが望ましい。$q_{\boldsymbol{\phi}}(\mathbf{z}) \approx p(\mathbf{z})$ が成立しないと、事前分布からサンプリングした $\mathbf{z}$ に対してデコーダが意味のある $\mathbf{x}$ を生成できない(ホールの問題)。
もう一つの重要な問題として事後崩壊(posterior collapse)がある。デコーダの表現力が高い場合(特に自己回帰型デコーダを用いる場合)、モデルは潜在変数 $\mathbf{z}$ を無視して $q_{\boldsymbol{\phi}}(\mathbf{z} \mid \mathbf{x}) \to p(\mathbf{z})$ となる局所解に落ちやすい。このとき KL 項はゼロになるが、$\mathbf{z}$ が意味のある情報を持たなくなる。対処法として、KL アニーリング($\beta$ を 0 から 1 へ徐々に増加)、自由ビット(Free Bits; $\mathrm{KL}_l \geq \lambda$)、$\delta$-VAE などが提案されている。
ニューラルネットワークによる変分推論は、古典的変分推論の3つの課題——スケーラビリティ、柔軟性、汎化——を同時に解決する枠組みである。その中核をなす要素を整理すると以下の通りである。
VAE はこれらを統合した代表的な生成モデルであり、潜在空間での補間・生成・半教師あり学習など多様なタスクに応用されている。さらに正規化流、階層的 VAE、IWAE、$\beta$-VAE などの拡張により、表現力・推論精度・解釈可能性の各側面で継続的に発展を遂げており、拡散モデルや大規模生成モデルの基盤理論としても重要な位置を占めている。
Mathematics is the language with which God has written the universe.