VAE(変分オートエンコーダ)の学習において、変分下界(ELBO)の勾配を計算する際、近似事後分布 $q_{\phi}(z|x)$ からのサンプリング操作が微分不可能であるという問題に直面する。この課題を解決し、バックプロパゲーションを可能にする手法が再パラメータ化トリック(Reparameterization Trick)である。この手法は、確率的ノードを含む計算グラフを連続的かつ微分可能な形へと再構成することで、確率的最適化と深層学習を整合的に接続する中核的技術である。
VAE の目的関数である ELBO を最大化するには、デコーダのパラメータ $\theta$ とエンコーダのパラメータ $\phi$ に関する勾配が必要である。ELBO は通常、再構成項と正則化項の和として表されるが、とりわけ問題となるのは再構成項に含まれる期待値である。
\[\mathbb{E}_{q_{\phi}(z|x)}[\log p_{\theta}(x|z)]\]この項に対して $\theta$ に関する勾配は、サンプル $z$ を固定した条件付き尤度の微分として容易に計算できる。一方で $\phi$ に関する勾配は、分布 $q_{\phi}(z|x)$ 自体が $\phi$ に依存するため、単純な微分では扱えない。一般に、この問題は以下の形で抽象化される。
\[\nabla_{\phi} \mathbb{E}_{q_{\phi}(z)}[f(z)]\]ここでの本質的困難は、「サンプリング操作が非連続であり、計算グラフ上で微分可能な経路を持たない」点にある。古典的にはスコア関数推定量(score function estimator)、すなわち
\[\nabla_{\phi} \mathbb{E}_{q_{\phi}(z)}[f(z)]=\mathbb{E}_{q_{\phi}(z)}[f(z)\nabla_{\phi} \log q_{\phi}(z)]\]が知られているが、この推定量は分散が極めて大きく、深層ニューラルネットワークと組み合わせた場合には実用的な収束を得ることが困難である。このため、より低分散な勾配推定手法が必要となる。
再パラメータ化トリックの核心は、確率変数 $z$ を「パラメータに依存しないノイズ」と「決定論的変換」に分解することである。すなわち、
\[z = g_{\phi}(x, \epsilon), \quad \epsilon \sim p(\epsilon)\]と表現する。このとき、$p(\epsilon)$ は $\phi$ に依存しない固定分布である。これにより、確率性はすべて $\epsilon$ に押し込められ、$z$ は $\phi$ に関して微分可能な関数として扱えるようになる。
VAE において最も一般的な設定は、近似事後分布を対角共分散を持つ多変量正規分布と仮定する場合である:
\[q_{\phi}(z|x) = \mathcal{N}(z; \mu_{\phi}(x), \mathrm{diag}(\sigma_{\phi}^2(x)))\]このとき、サンプリングは以下のように再表現される:
\[z = \mu_{\phi}(x) + \sigma_{\phi}(x) \odot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)\]ここで $\odot$ はアダマール積である。この変換により、確率的ノードは「外生ノイズ $\epsilon$ のサンプリング」と「決定論的写像」の合成として分離される。
上記の変換を用いることで、期待値は次のように書き換えられる:
\[\mathbb{E}_{q_{\phi}(z|x)}[f(z)]=\mathbb{E}_{p(\epsilon)}[f(g_{\phi}(x, \epsilon))]\]この形においては、分布 $p(\epsilon)$ が $\phi$ に依存しないため、勾配演算子を期待値の内部へ移動させることが可能となる:
\[\begin{aligned}\nabla_{\phi} \mathbb{E}_{q_{\phi}(z|x)}[f(z)]&= \nabla_{\phi} \mathbb{E}_{p(\epsilon)}[f(g_{\phi}(x, \epsilon))] \\&= \mathbb{E}_{p(\epsilon)}[\nabla_{\phi} f(g_{\phi}(x, \epsilon))]\end{aligned}\]ここでチェインルールを適用すると、
\[\nabla_{\phi} f(g_{\phi}(x,\epsilon))=\frac{\partial f}{\partial z}\frac{\partial z}{\partial \phi}\]となり、勾配はエンコーダネットワークの出力を通じて安定に伝播する。この形式は pathwise derivative estimator とも呼ばれる。
実際の計算では、期待値は有限個のサンプルによって近似される:
\[\mathbb{E}_{p(\epsilon)}[f(g_{\phi}(x, \epsilon))]\approx\frac{1}{L} \sum_{l=1}^{L} f(g_{\phi}(x, \epsilon^{(l)}))\]ここで $\epsilon^{(l)} \sim p(\epsilon)$ である。実務上は $L=1$ であっても十分に低分散であり、ミニバッチ学習と組み合わせることで安定した最適化が可能である。
VAE の ELBO は再構成項に加えて KL ダイバージェンス項を含む:
\[D_{\mathrm{KL}}(q_{\phi}(z|x) \,\|\, p(z))\]ガウス分布同士の場合、この項は閉形式で計算可能である:
\[D_{\mathrm{KL}} =\frac{1}{2}\sum_i\left(\mu_i^2 + \sigma_i^2 - \log \sigma_i^2 - 1\right)\]このため、再パラメータ化トリックは主として再構成項の勾配計算において本質的役割を果たす。
再パラメータ化トリックの本質的利点は、勾配推定量の分散が低い点にある。スコア関数法が分布全体の形状変化に依存するのに対し、本手法はサンプルの位置変化を通じて局所的に勾配を評価する。この違いにより、高次元空間においても安定した学習が可能となる。また、この手法は確率的計算グラフにおける微分の一般理論(stochastic backpropagation)の具体例としても理解される。
再パラメータ化トリックは連続分布に対しては自然に適用できるが、離散分布には直接適用できない。この問題に対しては、Gumbel-Softmax(Concrete 分布)などの連続緩和手法が提案されている。また、一般の分布に対しては、正規化フロー(normalizing flows)を用いて可逆変換を導入することで、より表現力の高い再パラメータ化が可能となる。
再パラメータ化トリックは、確率的サンプリングという非微分的操作を、外生ノイズと決定論的変換の組として再構成することで、勾配ベース最適化を可能にする手法である。この枠組みにより、VAE はエンドツーエンドで学習可能な確率モデルとして実現され、深層生成モデルの発展に決定的な影響を与えた。すなわち本手法は、変分推論・確率モデリング・深層学習を統一する基盤的技術であると言える。
Mathematics is the language with which God has written the universe.