リパラメータ化トリック (Reparameterization Trick)

深層学習モデルにおいて、確率分布からサンプリングされた変数を扱う場合、そのサンプリング操作は通常微分不可能です。そのため、誤差逆伝播法（バックプロパゲーション）を用いて勾配を計算し、モデルのパラメータを更新することができません。

この問題を解決するために、D. P. KingmaとM. Wellingが変分オートエンコーダ（VAE）の論文で導入した手法が、リパラメータ化トリック (Reparameterization Trick) です。

原理

リパラメータ化トリックは、確率分布からのサンプリング操作を、パラメータに依存しないランダムノイズと、パラメータに依存する決定的な変換に分離することで、微分可能にする手法です。

例えば、平均 \(\mu\) 、分散 \(\sigma^2\) の正規分布 \(\mathcal{N}(\mu, \sigma^2)\) から確率変数 \(z\) をサンプリングする場合を考えます。

このサンプリング操作は微分不可能です。しかし、標準正規分布 \(\mathcal{N}(0, 1)\) からサンプリングされたランダムノイズ \(\epsilon\) を用いると、\(z\) は以下のように表現できます。

（\(\odot\) は要素ごとの積を表します。多次元の場合、\(\sigma\) は標準偏差のベクトルまたは共分散行列の平方根に対応します。）

この変換により、\(z\) は \(\mu\) と \(\sigma\) に対して決定的な関数となり、\(\mu\) と \(\sigma\) はニューラルネットワークの出力として微分可能になります。ランダム性は \(\epsilon\) にカプセル化され、勾配は \(\mu\) と \(\sigma\) を通じて逆伝播できるようになります。

なぜ再パラメータ化が必要か：勾配とサンプリングの交換問題

VAEの学習では、エンコーダのパラメータ \(\phi\) に依存する分布 \(q_\phi(z)\) からサンプリングした \(z\) を使って、目的関数（ELBOの一部）

を最大化したいとします（\(f\) はデコーダの対数尤度など）。勾配法で \(\phi\) を更新するには \(\nabla_\phi L(\phi)\) が必要ですが、これは一見思うほど単純ではありません。\(q_\phi\) 自体が \(\phi\) に依存しているため、

という点に注意が必要です。右辺は「\(z\) を固定した上で \(f\) を \(\phi\) で微分する」操作ですが、\(z\) の分布そのものが \(\phi\) とともに動くため、これだけでは勾配の情報が欠落します。実際、積分と勾配を形式的に交換すると

となり、被積分関数は \(f(z) \nabla_\phi q_\phi(z)\) であって \(\nabla_\phi f(z)\) ではありません。さらに \(\nabla_\phi q_\phi(z)\) はそれ自体が確率密度ではない（積分して1になる保証がない）ため、式(1)を素朴にモンテカルロ近似する「\(q_\phi\) からサンプリングして平均を取る」という手続きがそのままでは使えません。これが「サンプリング操作を含む期待値の勾配をどう推定するか」という問題の核心です。

再パラメータ化トリックは、この問題を次のように回避します。\(z\) を、\(\phi\) に依存しない基底分布 \(p(\epsilon)\) からのサンプル \(\epsilon\) と、\(\phi\) に依存する決定的な変換 \(g_\phi\) の合成

として表現できれば（ガウス分布の場合は \(g_\phi(\epsilon) = \mu + \sigma \odot \epsilon\) ）、期待値は

と書き換えられます。ここでは積分（期待値）を取る対象の分布 \(p(\epsilon)\) が \(\phi\) に依存しない定数分布になっているため、緩やかな正則条件（\(f \circ g_\phi\) が \(\phi\) について可積分な優関数で抑えられることなど）のもとで勾配と積分を素直に交換でき、

が成立します。式(2)は \(\epsilon \sim p(\epsilon)\) から1つ以上サンプリングし、\(f(g_\phi(\epsilon))\) を通常の連鎖律（誤差逆伝播）で \(\phi\) について微分するだけで不偏推定できます。ランダム性の源である \(\epsilon\) の分布が固定されているため、「サンプリング」と「勾配計算」の順序を気にする必要がなくなる、というのが再パラメータ化トリックの本質です。

REINFORCE（スコア関数推定量）との比較

式(1)を直接モンテカルロ推定する別の方法として、**対数微分トリック（log-derivative trick）**があります。確率密度の勾配は

と書けます（\(\nabla_\phi \log q_\phi(z) = \nabla_\phi q_\phi(z) / q_\phi(z)\) の両辺に \(q_\phi(z)\) を掛けるだけ）。これを式(1)に代入すると、

となり、今度は期待値が（\(\phi\) に依存する）\(q_\phi(z)\) のもとで取られているにもかかわらず、被積分関数が通常の関数 \(f(z)\nabla_\phi \log q_\phi(z)\) になっているため、\(z \sim q_\phi\) を通常通りサンプリングして平均するだけで不偏推定量が得られます。これが REINFORCE推定量（スコア関数推定量とも呼ばれる）で、強化学習の 方策勾配法 における方策勾配定理の対数微分トリックと全く同じ導出です。REINFORCE推定量の最大の利点は、\(z\) が \(\phi\) の微分可能な関数として表現できる必要がない点です。つまり \(z\) が離散変数であっても適用できます。一方で再パラメータ化トリックは \(z = g_\phi(\epsilon)\) という微分可能な変換の存在を要求するため、原理的に連続変数にしか使えません。

その代わり、REINFORCE推定量は一般に分散が大きいことが知られています。これを具体的に検証するために、\(q_\phi = \mathcal{N}(\mu, \sigma^2)\) （\(\phi = \mu\) 、\(\sigma\) は定数）、\(f(z) = z^2\) という単純な設定で、2つの推定量を導出し比較します。

• REINFORCE推定量: \(\log q_\phi(z) = -\frac{1}{2}\log(2\pi\sigma^2) - \frac{(z-\mu)^2}{2\sigma^2}\) より \(\nabla_\mu \log q_\phi(z) = (z-\mu)/\sigma^2\) 。よって式(3)の被積分関数（1サンプルあたりの推定量）は

• 再パラメータ化推定量: \(z = \mu + \sigma\epsilon\) 、\(f(z)=z^2\) なので \(\nabla_\mu f(z) = 2z \cdot \frac{\partial z}{\partial \mu} = 2z\) 。よって

真の勾配は \(\mathbb{E}[z^2] = \mu^2+\sigma^2\) より \(\nabla_\mu \mathbb{E}[z^2] = 2\mu\) です。\(u = z-\mu \sim \mathcal{N}(0,\sigma^2)\) とおいてガウス分布の中心モーメント（\(\mathbb{E}[u]=\mathbb{E}[u^3]=\mathbb{E}[u^5]=0\) 、\(\mathbb{E}[u^2]=\sigma^2\) 、\(\mathbb{E}[u^4]=3\sigma^4\) 、\(\mathbb{E}[u^6]=15\sigma^6\) ）を使うと、両推定量の分散を解析的に計算できます。

（後者は \(\hat g_{\text{REINFORCE}}=(u^3+2\mu u^2+\mu^2 u)/\sigma^2\) を展開し、奇数次中心モーメントが0になることを使って \(\mathbb{E}[\hat g_{\text{REINFORCE}}^2]\) を計算し、\((2\mu)^2\) を引くことで得られます。）\(\mu=1,\sigma=1\) を代入すると、分散比は

と厳密に求まります。この理論値を実際にPythonでモンテカルロ検証しました。

import numpy as np

rng = np.random.default_rng(42)
mu, sigma = 1.0, 1.0
n = 5_000_000

z = rng.normal(mu, sigma, size=n)
g_reinforce = z**2 * (z - mu) / sigma**2
g_reparam = 2 * z

print("Var(REINFORCE):", g_reinforce.var())   # 理論値 30
print("Var(reparam):  ", g_reparam.var())     # 理論値 4
print("分散比:", g_reinforce.var() / g_reparam.var())  # 理論値 7.5

実行結果は Var(REINFORCE) = 30.037、Var(reparam) = 3.998、分散比 7.513 となり、理論値（30, 4, 7.5）と極めて良く一致しました。同様に \(\mu=0\) とすると理論分散比は \(15/4=3.75\) ですが、これも実測で Var(REINFORCE)=15.023、Var(reparam)=3.998、比 3.758 と理論値に一致することを確認しています（\(\mu=0\) でもREINFORCE推定量の分散はゼロにならない点に注意が必要です）。

サンプル数 \(N\) を増やして平均をとった場合の分散（\(\mu=1,\sigma=1\) 、各設定20,000試行でモンテカルロ評価）も確認しました。

どちらの推定量も分散は \(1/N\) で減少し（中心極限定理の帰結）、両者とも真の勾配 \(2\mu=2.0\) に不偏収束しますが、REINFORCE推定量の分散は再パラメータ化推定量よりおよそ7.5倍大きいまま、比はサンプル数によらずほぼ一定です。\(N=1\) から \(N=2000\) まで対数間隔で細かく評価した結果を可視化すると、両曲線が並行に減衰していく（＝比が一定である）様子がはっきり見えます。

この分散比の差は次元数 \(d\) が増えるほど深刻になることが知られており（REINFORCE推定量の分散は一般に次元に対して不利にスケールする）、これが深層生成モデルの学習で再パラメータ化トリックが（使える場面では）標準的に選ばれる実務上の理由です。逆に言えば、\(z\) が離散変数などで再パラメータ化できない場合は、分散削減のための工夫（ベースライン減算など、 Actor-Criticのアドバンテージ関数 と同種の技法）を併用したREINFORCE推定量に頼らざるを得ません。

多変量・非ガウス分布への拡張

対角共分散の多変量ガウス

VAEのエンコーダが出力する \(\mu, \sigma \in \mathbb{R}^d\) が独立な対角共分散 \(\Sigma = \mathrm{diag}(\sigma_1^2,\dots,\sigma_d^2)\) を持つ場合、既に示した式 \(z = \mu + \sigma \odot \epsilon\) （\(\epsilon \sim \mathcal{N}(0, I_d)\) 、\(\odot\) は要素積）がそのまま多変量版の再パラメータ化になります。各次元が独立なので、この変換で得られる \(z\) の共分散が正しく \(\Sigma\) になることは各次元ごとの1次元の場合に帰着して自明です。

フル共分散（コレスキー分解によるサンプリング）

次元間に相関がある一般の共分散行列 \(\Sigma\) （半正定値対称行列）を扱いたい場合は、コレスキー分解 \(\Sigma = LL^\top\) （\(L\) は下三角行列）を用います。\(\epsilon \sim \mathcal{N}(0, I_d)\) に対して

とすると、\(z\) は平均 \(\mu\) 、共分散 \(\mathrm{Cov}(z) = L\,\mathrm{Cov}(\epsilon)\,L^\top = L I_d L^\top = LL^\top = \Sigma\) の多変量正規分布に従います（線形変換の共分散が \(\mathrm{Cov}(Az)=A\,\mathrm{Cov}(z)\,A^\top\) となることを使っています）。これも \(L\) が \(\phi\) の微分可能な関数である限り、勾配は \(L\) を通じて逆伝播できます。実際にランダムな3次元のSPD（対称正定値）行列 \(\Sigma\) と平均ベクトル \(\mu\) に対してこの手続きを検証しました。

import numpy as np

rng = np.random.default_rng(1)
d = 3
A = rng.normal(size=(d, d))
Sigma = A @ A.T + d * np.eye(d)   # ランダムなSPD行列
mu = np.array([1.0, -2.0, 0.5])

L = np.linalg.cholesky(Sigma)
eps = rng.normal(size=(500_000, d))
Z = mu + eps @ L.T

print("sample mean:", Z.mean(axis=0))
print("sample cov:\n", np.cov(Z.T))

50万サンプルでの実行結果は、標本平均が [1.0014, -2.0034, 0.5067]（目標 [1.0, -2.0, 0.5] との最大絶対誤差 0.0067）、標本共分散が目標 \(\Sigma\) との最大絶対誤差 0.0130 となり、サンプルサイズに応じた統計誤差の範囲内で理論通りに一致することを確認しました。

離散変数への拡張の困難とGumbel-Softmax

再パラメータ化トリックは、離散変数（カテゴリカル変数）には単純には適用できません。カテゴリ確率 \(\pi_1,\dots,\pi_K\) （\(\phi\) に依存）からのサンプリングは、例えば一様乱数 \(u\sim U(0,1)\) を累積分布関数で閾値処理する、あるいは後述のGumbel-Max法のように実装できますが、いずれの場合も「基底ノイズ → 出力」の写像がステップ関数（区分定数）になります。ステップ関数は各点でほとんど至るところ勾配がゼロであり、\(\phi\) を連続的に動かしても出力がジャンプ的にしか変化しないため、式(2)の \(\nabla_\phi f(g_\phi(\epsilon))\) が意味のある情報を持ちません。これが「離散変数には単純な再パラメータ化が使えない」理由です。

この問題への対処として、Gumbel-Maxトリックと**Gumbel-Softmax（Concrete分布）**があります。まずGumbel-Maxトリックは、カテゴリカル分布からのサンプリングを

という形で厳密に（分布として同値に）再パラメータ化します（\(g_k = -\log(-\log u_k)\) 、\(u_k \sim U(0,1)\) で生成）。ランダム性は \(\phi\) に依存しない \(g_k\) に押し込められていますが、\(\arg\max\) 自体が微分不可能なため、これだけではまだ勾配は流れません。そこでGumbel-Softmaxは \(\arg\max\) を温度パラメータ \(\tau>0\) 付きのsoftmaxで連続緩和します。

\(\tau \to 0\) の極限で \(\tilde z\) はGumbel-Maxのone-hotサンプルに確率収束し、\(\tau\) が大きいほど各カテゴリに確率が分散した「柔らかい」サンプルになります。この温度依存の振る舞いを、\(\pi=(0.5,0.3,0.15,0.05)\) の4値カテゴリカル分布で実際に検証しました。

import numpy as np

rng = np.random.default_rng(7)
pi = np.array([0.5, 0.3, 0.15, 0.05])
logits = np.log(pi)
n = 200_000

g = -np.log(-np.log(rng.uniform(size=(n, 4))))
y = logits + g

# Gumbel-Max: 各カテゴリの選択頻度
onehot = np.eye(4)[y.argmax(axis=1)]
print("Gumbel-Max 実測頻度:", onehot.mean(axis=0))

for tau in [0.1, 1.0]:
    soft = np.exp(y / tau)
    soft /= soft.sum(axis=1, keepdims=True)
    print(f"tau={tau}: 最大成分の平均値 = {soft.max(axis=1).mean():.4f}")

Gumbel-Maxの実測カテゴリ頻度は [0.5006, 0.3005, 0.1494, 0.0495] となり、目標の \(\pi=(0.5,0.3,0.15,0.05)\) とほぼ一致しました（式として厳密に同値であることの数値的裏付け）。また温度依存性については、\(\tau=0.1\) （低温）では出力ベクトルの最大成分の平均値が 0.9554 と1（one-hot）に近く離散に近い挙動を示すのに対し、\(\tau=1.0\) （高温）では 0.6376 にとどまり、カテゴリ間に確率が分散した「柔らかい」サンプルになることが確認できました。これは低温にするほど離散分布への近似は良くなる一方、\(\tau\to0\) で \(\nabla_\phi \tilde z\) の分散が発散して勾配推定が不安定になるというトレードオフに対応しています（実務では \(\tau\) をアニーリングする、あるいはStraight-Through推定量と組み合わせるといった工夫が使われます）。

実務上の注意点

分散パラメータを log-variance で扱う理由

VAEの実装では、エンコーダのネットワークに \(\sigma\) や \(\sigma^2\) を直接出力させるのではなく、対数分散 \(\log\sigma^2\) （あるいは \(\log\sigma\) ）を出力させ、\(\sigma = \exp(0.5 \log\sigma^2)\) として使うのが標準的です。理由は主に3つあります。

1. 制約なし最適化にできる: \(\sigma^2>0\) という制約を、\(\log\sigma^2 \in \mathbb{R}\) という無制約なパラメータ化に変換できます。ニューラルネットワークの出力層は通常、値域を制限しないため、制約付き最適化や出力への追加の非負化処理（ReLUやsoftplusでの丸め）を回避できます。
2. 勾配の数値的な健全性: \(\sigma\to 0\) の付近で \(\sigma\) を直接パラメータ化すると、勾配やKLダイバージェンス項の \(\log \sigma^2\) の計算が不安定になりがちです。\(\log\sigma^2\) を変数として扱えば \(\sigma^2 \to 0\) は \(\log\sigma^2 \to -\infty\) という滑らかな極限になり、勾配爆発・消失を起こしにくくなります。
3. 正値性の保証: \(\sigma = \exp(0.5\log\sigma^2)\) は入力がどんな実数であっても常に正の値を返すため、学習の初期段階でパラメータが大きく動いても \(\sigma\) が負や複素数になる心配がありません。

KLダイバージェンス項の解析解

VAEの損失関数には、事後分布 \(q_\phi(z|x)=\mathcal{N}(\mu,\mathrm{diag}(\sigma^2))\) と事前分布 \(p(z)=\mathcal{N}(0,I)\) の間のKLダイバージェンス \(D_{KL}(q_\phi\|p)\) が含まれます。事前分布が標準正規分布の場合、このKL項は積分計算なしに閉形式で求まります。1次元で導出すると、

なので、

この期待値 \(\mathbb{E}_q[\cdot]\) を取ると、\(\mathbb{E}_q[(z-\mu)^2]=\sigma^2\) より第2項は \(-\frac12\) になり、\(\mathbb{E}_q[z^2] = \mu^2+\sigma^2\) （分散の定義 \(\mathrm{Var}(z)=\mathbb{E}[z^2]-\mathbb{E}[z]^2\) より）を使うと、

が得られます。\(d\) 次元の対角共分散の場合は各次元が独立なのでKLダイバージェンスは加法的になり、

という、VAEの実装で頻繁に見かける式が導けます。この閉形式をモンテカルロ積分と突き合わせて検証しました。

import numpy as np

rng = np.random.default_rng(1)
mu = np.array([0.5, -1.2, 2.0])
log_var = np.array([-0.3, 0.8, 0.1])

# 解析解（式4）
kl_analytic = 0.5 * np.sum(np.exp(log_var) + mu**2 - 1.0 - log_var)

# モンテカルロ積分
sigma = np.exp(0.5 * log_var)
n = 2_000_000
z = mu + sigma * rng.normal(size=(n, 3))
log_q = -0.5*np.sum(np.log(2*np.pi*sigma**2) + ((z-mu)/sigma)**2, axis=1)
log_p = -0.5*np.sum(np.log(2*np.pi) + z**2, axis=1)
kl_mc = np.mean(log_q - log_p)

print(f"解析解: {kl_analytic:.6f}")
print(f"モンテカルロ推定(n=2e6): {kl_mc:.6f}")

実行結果は解析解 3.080765、モンテカルロ推定（200万サンプル）3.081169 で、差はわずか 0.000404 でした。これは式(4)の導出が正しいことの数値的な裏付けです。この閉形式のおかげで、KL項の計算に追加のサンプリングやモンテカルロ近似が不要になり、ELBOの再構成誤差項（再パラメータ化トリックで推定）とKL項（解析解）を組み合わせるのがVAEの標準的な実装になっています。

PyTorchでのリパラメータ化トリックを用いたサンプリング

PyTorchの torch.distributions モジュールは、リパラメータ化トリックをサポートしている分布を提供しています。Normal クラスなどの分布オブジェクトは、rsample() メソッドを通じてリパラメータ化トリックを用いたサンプリングをサポートしています。

import torch
from torch.distributions import Normal

# ニューラルネットワークの出力として得られた平均と標準偏差
# 例として、mu=0, sigma=1 の正規分布を考える
mu = torch.tensor(0.0, requires_grad=True)
sigma = torch.tensor(1.0, requires_grad=True)

# 正規分布オブジェクトを作成
m = Normal(mu, sigma)

# rsample() メソッドを使ってサンプリング
# このサンプリングはリパラメータ化トリックが適用されるため、勾配が計算可能
z = m.rsample()

print(f"サンプリングされたz: {z}")
print(f"zの勾配計算可能フラグ: {z.requires_grad}") # Trueになる

m.has_rsample プロパティを確認することで、その分布がリパラメータ化トリックをサポートしているかどうかを判別できます。

print(f"Normal分布はrsampleをサポートしているか: {m.has_rsample}") # True

先に導出したKLダイバージェンスの閉形式（式4）は、PyTorchの torch.distributions.kl_divergence 関数の結果とも一致することを確認できます。

import torch
import torch.distributions as D

mu = torch.tensor([0.5, -1.2, 2.0])
log_var = torch.tensor([-0.3, 0.8, 0.1])

# 自前実装（式4）
kl_manual = 0.5 * torch.sum(torch.exp(log_var) + mu**2 - 1.0 - log_var)

# PyTorch組み込み関数
q = D.Normal(mu, torch.exp(0.5 * log_var))
p = D.Normal(torch.zeros(3), torch.ones(3))
kl_builtin = D.kl_divergence(q, p).sum()

print(f"自前実装: {kl_manual.item():.7f}")
print(f"torch.distributions.kl_divergence: {kl_builtin.item():.7f}")

実行結果は自前実装が 3.0807652、kl_divergence が 3.0807648 となり、float32の丸め誤差レベル（\(10^{-7}\) オーダー）で完全に一致しました。これは式(4)の導出が正しいことのもう一つの裏付けです。

リパラメータ化トリックは、VAEだけでなく、強化学習の確率的方策勾配法など、確率的な要素を含む深層学習モデルの学習において非常に重要な技術です。\(z\) が微分可能な変換で表現できる限り低分散な勾配推定が可能ですが、離散変数や一部の複雑な分布ではREINFORCE推定量やGumbel-Softmaxのような代替・緩和手法との使い分けが必要になる、という点を押さえておくと実装上の判断がしやすくなります。

参考

• Kingma, D. P., & Welling, M. (2013). “Auto-Encoding Variational Bayes”. arXiv preprint arXiv:1312.6114.
• Williams, R. J. (1992). “Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning”. Machine Learning, 8, 229–256.（REINFORCE推定量の原論文）
• Jang, E., Gu, S., & Poole, B. (2016). “Categorical Reparameterization with Gumbel-Softmax”. arXiv preprint arXiv:1611.01144.
• Maddison, C. J., Mnih, A., & Teh, Y. W. (2016). “The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables”. arXiv preprint arXiv:1611.00712.
• 方策勾配法 — 対数微分トリックによる期待報酬の勾配導出。REINFORCE推定量と本質的に同じ数学的技巧を扱っています。
• ReNom, 変分オートエンコーダ