勾配推定の3つの方法:reparameterization・REINFORCE・Gumbel-softmaxをわかりやすく

ニューラルネットの学習は「損失を $\theta$ で微分して、勾配の逆向きに少し動かす」の繰り返しです。ところが、現代の機械学習では、この「微分」が素直にできない場面が頻繁に現れます。途中にサンプリング(乱数を引く操作)が挟まるときです。

たとえば次のような量を最小化(あるいは最大化)したいとします。

$$ J(\theta) = \mathbb{E}_{z \sim p_\theta(z)}\big[\, f(z) \,\big] $$

これは「分布 $p_\theta(z)$ から $z$ をサンプリングし、$f(z)$ を測り、その平均をとる」量です。やっかいなのは、期待値をとる分布そのものが $\theta$ に依存していることです。$\theta$ を少し変えると、$f$ の値が変わるのではなく、$z$ が引かれてくる「確率」が変わります。サンプリングという確率的な操作の真ん中を、どうやって微分すればよいのでしょうか。

この問題は具体的な場面で次々と顔を出します。

  • VAE(変分オートエンコーダ):潜在変数 $z$ をエンコーダの分布 $q_\theta(z\mid x)$ からサンプリングし、デコーダで再構成します。エンコーダのパラメータ $\theta$ を学習するには、まさにこの「サンプリングを挟んだ期待値」を微分する必要があります。
  • 強化学習の方策勾配:方策 $\pi_\theta(a\mid s)$ に従って行動 $a$ をサンプリングし、得られた報酬の期待値を最大化します。環境(報酬関数)は微分できない黒箱なのに、方策のパラメータを勾配で更新したいのです。

この記事では、この「サンプリングを挟んだ期待値の勾配」を推定する代表的な3つの方法を、直感・導出・分散の比較・適用条件まで通して解説します。スコア関数推定量(REINFORCE)reparameterizationトリック、そして離散変数のためのGumbel-softmaxです。3つの使い分けが腹に落ちると、生成モデルも強化学習も同じ土俵で見えるようになります。

本記事の内容

  • サンプリングを挟んだ期待値の勾配が難しい理由
  • スコア関数推定量(REINFORCE)の導出と分散低減(baseline)
  • reparameterizationトリックの仕組みと低分散性
  • Gumbel-softmaxによる離散変数のreparameterization
  • Pythonによる不偏性・分散の実測と3手法の使い分け

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

なぜサンプリングを挟むと微分できないのか

まず、難しさの核心を絵で押さえましょう。私たちが計算したいのは「分布から引いた値の平均」です。$\theta$ を動かすと何が動くのか、図で見ます。

分布p_theta(z)自体がthetaで動くため期待値の勾配が難しいことを示す図

3本の山は、$\theta$(ここでは平均 $\mu$)を変えたときの分布 $p_\theta(z)$ です。$\theta$ を大きくすると、山が右へ動いていきます。つまり $\theta$ は「$f$ の値」を直接いじっているのではなく、「どの $z$ がよく出るか」という確率の重みを動かしているのです。勾配 $\nabla_\theta J$ は、この重みの移動が期待値をどう変えるかを表します。

もう少し丁寧に式で見ます。期待値は積分で書けます。

$$ J(\theta) = \mathbb{E}_{z \sim p_\theta(z)}[f(z)] = \int f(z)\, p_\theta(z)\, dz $$

ここで $f(z)$ は $\theta$ を含みません。$\theta$ が入っているのは確率密度 $p_\theta(z)$ の方だけです。だから、積分と微分の順序が入れ替えられるなら、勾配は次のように書けます。

$$ \nabla_\theta J(\theta) = \int f(z)\, \nabla_\theta p_\theta(z)\, dz $$

式の上ではきれいですが、これはそのままでは期待値の形になっていません。$\nabla_\theta p_\theta(z)$ は確率密度ではないので、「$p_\theta$ からサンプリングして平均する」というモンテカルロ近似が使えないのです。実装上は、サンプリングという乱数操作のところで計算グラフが切れてしまいます。

計算グラフでサンプリングが勾配の逆伝播を遮断することを示す図

この図が問題のすべてを語っています。$\theta \to$ サンプリング $\to f(z) \to$ 損失、と順方向には流れます。しかし損失からの勾配を逆向きに流そうとすると、真ん中の「サンプリング(乱数を引く)」のところで止まります。乱数を引く操作は微分できないからです。勾配を $\theta$ まで届かせるための工夫、それが本記事の3手法です。

ここからは、この壁をどう乗り越えるかを1つずつ見ていきます。まずは、最も一般的に使える「スコア関数推定量」からです。

手法1:スコア関数推定量(REINFORCE)

最初の手法は、$f$ の中身を一切のぞき込まずに勾配を作る方法です。$f$ が微分できなくても、離散でも、黒箱のシミュレータでも使えます。強化学習では REINFORCE、統計では尤度比推定量(likelihood ratio estimator) とも呼ばれます。

対数微分トリック

鍵になるのは、対数の微分公式を逆向きに使う「対数微分トリック」です。まず、対数の微分から次が成り立ちます。

$$ \nabla_\theta \log p_\theta(z) = \frac{\nabla_\theta p_\theta(z)}{p_\theta(z)} $$

両辺に $p_\theta(z)$ を掛けると、欲しかった $\nabla_\theta p_\theta(z)$ を「密度 $\times$ 対数の勾配」に書き換えられます。

$$ \nabla_\theta p_\theta(z) = p_\theta(z)\, \nabla_\theta \log p_\theta(z) $$

この一手が効きます。先ほどの勾配の式にこれを代入してみましょう。

$$ \nabla_\theta J(\theta) = \int f(z)\, \nabla_\theta p_\theta(z)\, dz = \int f(z)\, p_\theta(z)\, \nabla_\theta \log p_\theta(z)\, dz $$

右辺をよく見ると、$p_\theta(z)$ が密度として復活しています。つまり、これは $p_\theta$ に関する期待値の形に戻ったのです。

$$ \boxed{\ \nabla_\theta J(\theta) = \mathbb{E}_{z \sim p_\theta(z)}\big[\, f(z)\, \nabla_\theta \log p_\theta(z) \,\big]\ } $$

期待値になったということは、モンテカルロ近似ができます。$z_1,\dots,z_N$ を $p_\theta$ からサンプリングして、平均をとればよいのです。

$$ \nabla_\theta J(\theta) \approx \frac{1}{N}\sum_{i=1}^{N} f(z_i)\, \nabla_\theta \log p_\theta(z_i) $$

REINFORCEがスコア関数で勾配を構成する仕組みの図

この図がREINFORCEの本質です。サンプル $z$ そのものは微分しません。代わりに、各サンプルに対して スコア $\nabla_\theta \log p_\theta(z)$ を計算し、その値を $f(z)$ で重み付けして平均します。直感的には「報酬 $f$ が高かったサンプルが出やすくなる向きに $\theta$ を動かす」という更新です。報酬が高い行動の確率を上げ、低い行動の確率を下げる――方策勾配定理の正体は、まさにこの式なのです。

一般性の裏にある弱点:高分散

スコア関数推定量のすごいところは、$f$ について何も仮定しないことです。$f$ が微分できなくても、$z$ が離散カテゴリでも、$f$ がゲームの勝敗のような0/1でも、そのまま使えます。$\nabla_\theta \log p_\theta(z)$ さえ計算できれば成立します。これが強化学習で広く使われる理由です。

しかし代償があります。分散が非常に大きいのです。$f(z)$ の値が丸ごと $\nabla_\theta \log p_\theta(z)$ に掛かるため、$f$ の絶対値が大きいとスコアが大きく振れます。サンプルごとの推定量がばらつき、学習が不安定になります。この分散をどう抑えるかが、スコア関数推定量を実用にするための最大の課題です。

baseline(制御変量)による分散低減

分散を下げる定番が baseline の導入です。アイデアはシンプルで、$f(z)$ から定数 $b$ を引きます。

$$ \nabla_\theta J(\theta) = \mathbb{E}\big[\, (f(z) – b)\, \nabla_\theta \log p_\theta(z) \,\big] $$

ここで重要なのは、$b$ を引いても推定量の期待値(平均)は変わらないことです。なぜなら、スコアの期待値は0だからです。これを確認しましょう。$b$ を定数とすると、引いた項の期待値は

$$ \mathbb{E}\big[\, b\, \nabla_\theta \log p_\theta(z) \,\big] = b \int p_\theta(z)\, \nabla_\theta \log p_\theta(z)\, dz = b \int \nabla_\theta p_\theta(z)\, dz $$

となります。ここで対数微分トリックを逆向きに使いました。さらに積分と微分を入れ替えると、

$$ b \int \nabla_\theta p_\theta(z)\, dz = b\, \nabla_\theta \int p_\theta(z)\, dz = b\, \nabla_\theta 1 = 0 $$

です。確率密度の積分は常に1なので、その勾配は0になります。つまり baseline を引いても勾配の期待値は不変(不偏性を保つ)のまま、分散だけを下げられます。これは統計で「制御変量法」と呼ばれる分散低減テクニックそのものです。

実用上は $b$ として $f(z)$ の平均(移動平均)を使ったり、強化学習では状態価値 $V(s)$ を使ったりします。$f(z) – b$ が0付近に集まるほど、スコアに掛かる重みが小さくなり、分散が下がります。

スコア関数推定量は「何にでも使えるが分散が大きい」万能選手でした。では、$f$ が微分できるとき、もっと分散の小さい推定量は作れないでしょうか。次の手法がそれに答えます。

手法2:reparameterizationトリック

$f(z)$ が $z$ について微分できて、しかも分布が「滑らかに作れる」とき、スコア関数推定量よりはるかに低分散な推定量が手に入ります。それが reparameterizationトリック です。VAEを実用にした立役者でもあります。

乱数を経路の外に追い出す

核心のアイデアは、「サンプリング」を「決定的な変換 $+$ パラメータに依存しない乱数」に分解することです。たとえば $z \sim \mathcal{N}(\mu, \sigma^2)$ は、次のように書き換えられます。

$$ z = \mu + \sigma\, \varepsilon, \qquad \varepsilon \sim \mathcal{N}(0, 1) $$

ここで $\varepsilon$ は標準正規分布からの乱数で、$\theta = (\mu, \sigma)$ に一切依存しません。$z$ をいきなり $\mathcal{N}(\mu,\sigma^2)$ から引く代わりに、$\theta$ に依存しない $\varepsilon$ を引いてから、決定的な式 $z = \mu + \sigma\varepsilon$ で変換するのです。一般には $z = g(\theta, \varepsilon)$ と書きます。

reparameterizationトリックで乱数を入力に追い出す計算グラフの図

図のポイントは、乱数 $\varepsilon$ が計算の「入口」に移動したことです。$\theta$ から $z$、$z$ から $f$ への道筋は、すべて決定的な(微分できる)関数だけでつながっています。乱数はもう道の途中にいないので、勾配を $\theta$ まで一直線に流せます。先ほど「サンプリングで切れていた」計算グラフが、つながったのです。

勾配が期待値の中にそのまま入る

数式で確かめます。$z = g(\theta, \varepsilon)$ と置くと、期待値は $\varepsilon$ についての期待値に書き換わります。

$$ J(\theta) = \mathbb{E}_{z \sim p_\theta(z)}[f(z)] = \mathbb{E}_{\varepsilon \sim p(\varepsilon)}\big[\, f(g(\theta, \varepsilon)) \,\big] $$

ここが決定的に違う点です。右辺で期待値をとる分布 $p(\varepsilon)$ は $\theta$ に依存しません。 だから、微分を期待値の中にそのまま入れられます。

$$ \boxed{\ \nabla_\theta J(\theta) = \mathbb{E}_{\varepsilon \sim p(\varepsilon)}\big[\, \nabla_\theta f(g(\theta, \varepsilon)) \,\big]\ } $$

あとは連鎖律で $\nabla_\theta f(g(\theta,\varepsilon)) = f'(g)\cdot \nabla_\theta g(\theta,\varepsilon)$ を計算するだけです。モンテカルロ近似は

$$ \nabla_\theta J(\theta) \approx \frac{1}{N}\sum_{i=1}^{N} \nabla_\theta f\big(g(\theta, \varepsilon_i)\big), \qquad \varepsilon_i \sim p(\varepsilon) $$

となります。スコア関数推定量が「$f(z)$ をスコアで重み付け」だったのに対し、こちらは「$f$ を直接 $\theta$ で微分」しています。$f$ の勾配情報を使えるぶん、推定量のばらつきが小さくなるのが直感的に分かります。

なぜ低分散なのか、そして適用条件

reparameterizationが低分散になるのは、$f$ の傾き($f’$)という構造的な情報を勾配にそのまま使うからです。スコア関数推定量は $f$ の値だけを使い、傾きは使いません。情報量の差が分散の差になって現れます。

ただし、使える条件は限定的です。

  • $f$ が $z$ について微分可能であること(黒箱の報酬関数では使えない)
  • 分布が $z = g(\theta, \varepsilon)$ の形で微分可能に書き換えられること

連続分布の多く(正規分布、対数正規分布など)はこの形に書けますが、離散分布(カテゴリ分布)はこのままでは書けません。カテゴリを選ぶ操作は階段状で、微分が0かつ不連続だからです。この壁を越えるのが3つ目の手法です。

その前に、ここまでの2手法の「不偏性」と「分散」を、実際に数値で確かめてみましょう。

Pythonで不偏性と分散を実測する

理論で「reparameterizationは低分散」「baselineは不偏なまま分散を下げる」と述べました。これを、解析的な答えが分かるトイ問題で実測します。題材は次の期待値です。

$$ J(\mu, \sigma) = \mathbb{E}_{z \sim \mathcal{N}(\mu, \sigma^2)}[\, z^2 \,] $$

$z^2$ の期待値は $\mathbb{E}[z^2] = \mu^2 + \sigma^2$ なので、勾配は手計算できます。

$$ \frac{\partial J}{\partial \mu} = 2\mu, \qquad \frac{\partial J}{\partial \sigma} = 2\sigma $$

真の値が分かっているので、各推定量が「平均すると真の勾配に一致するか(不偏性)」「どれだけばらつくか(分散)」を直接測れます。

import numpy as np

mu, sigma = 1.0, 1.5
true_grad_mu = 2 * mu        # = 2.0
true_grad_sigma = 2 * sigma  # = 3.0

def estimators(n, seed):
    g = np.random.default_rng(seed)
    eps = g.standard_normal(n)
    z = mu + sigma * eps
    f = z ** 2

    # reparam: f(z)=z^2, z=mu+sigma*eps を直接 theta で微分
    rep_mu  = 2 * z            # df/dmu = 2z * dz/dmu = 2z
    rep_sig = 2 * z * eps      # df/dsigma = 2z * dz/dsigma = 2z*eps

    # score function (REINFORCE): f * d/dtheta log N(z;mu,sigma)
    score_mu  = (z - mu) / sigma ** 2
    score_sig = ((z - mu) ** 2 - sigma ** 2) / sigma ** 3
    rf_mu  = f * score_mu
    rf_sig = f * score_sig

    # baseline b = サンプル平均 (制御変量)
    b = f.mean()
    rfb_mu  = (f - b) * score_mu
    rfb_sig = (f - b) * score_sig

    return dict(rep_mu=rep_mu, rf_mu=rf_mu, rfb_mu=rfb_mu)

N = 20000
e = estimators(N, 1)
for name in ["rep_mu", "rf_mu", "rfb_mu"]:
    print(f"{name}: 平均={e[name].mean():.3f}  分散={e[name].var():.2f}")

実行すると、おおよそ次の結果になります。

rep_mu: 平均=1.966  分散=8.89
rf_mu:  平均=1.945  分散=47.53
rfb_mu: 平均=1.969  分散=30.52

この出力から3つのことが読み取れます。第一に、3手法とも平均が真の勾配 $2.0$ に近い――どれも不偏です。第二に、reparameterの分散(8.9)はREINFORCEの分散(47.5)の約5分の1で、理論どおり大幅に低分散です。第三に、baselineを引くとREINFORCEの分散が47.5から30.5へ下がりつつ、平均は1.97のままで、不偏性を保ったまま分散だけ低減できています。

この「ばらつきの差」を可視化したのが次のヒストグラムです。

REINFORCEとreparameterizationの単一サンプル勾配推定量の分布を比較するヒストグラム

横軸は単一サンプルから計算した勾配推定量の値、縦軸は頻度です。オレンジ(REINFORCE)は0付近に鋭いピークがある一方、左右に長く裾を引いています。多くのサンプルでは値が小さいのに、ときどき極端に大きな値が出るため、平均すると分散が膨らむのです。これに対して水色(reparam)は真の勾配 $2.0$ のまわりにきれいに集中しています。分布の「広がり」そのものが分散の差です。

次に、サンプル数 $N$ を増やすと推定誤差がどう減るかを見ます。

import numpy as np

ns = [1, 2, 5, 10, 20, 50, 100, 200, 500]
reps = 400
for label, key in [("REINFORCE","rf_mu"), ("+baseline","rfb_mu"), ("reparam","rep_mu")]:
    stds = []
    for n in ns:
        means = [estimators(n, 1000 + r)[key].mean() for r in range(reps)]
        stds.append(np.std(means))
    print(label, [f"{s:.2f}" for s in stds])

サンプル数と勾配推定の標準偏差の関係を示す両対数グラフ

両対数グラフで、3手法とも傾き $-1/2$ の直線($\propto N^{-1/2}$)に乗っています。これはモンテカルロ推定の基本性質で、サンプルを4倍にすると誤差が半分になることを意味します。注目すべきは直線の高さです。reparam(水色)が最も下、baseline付きREINFORCE(紫)が中間、素のREINFORCE(オレンジ)が最も上にあります。同じ誤差に到達するのに必要なサンプル数が、手法によって何倍も違うことが読み取れます。reparameterが使える場面では、圧倒的に少ないサンプルで済むのです。

ここまでは連続変数の話でした。では、$z$ がカテゴリ(離散)のときはどうすればよいのでしょうか。reparameterizationは使えそうにありません。次の手法がこの難所を切り開きます。

手法3:Gumbel-softmax(Concrete分布)

潜在変数が「カテゴリA・B・Cのどれか」のような離散値のとき、reparameterizationは直接は使えません。カテゴリを選ぶ操作($\arg\max$)が微分不可能だからです。Gumbel-softmax(別名 Concrete分布)は、この離散サンプリングを「微分可能な近似」に置き換えるトリックです。

Gumbel-maxトリック

出発点は、カテゴリ分布からのサンプリングを「ノイズ $+$ $\arg\max$」で書けるという事実です。カテゴリ $k$ の確率を $\pi_k$ とすると、次の手順がカテゴリ分布からのサンプリングと完全に等価になります。

$$ z = \arg\max_k \big(\, \log \pi_k + g_k \,\big), \qquad g_k \sim \text{Gumbel}(0, 1) $$

ここで $g_k$ は各カテゴリに独立に足す Gumbelノイズ で、$g_k = -\log(-\log u_k),\ u_k \sim \text{Uniform}(0,1)$ で生成します。「対数確率にGumbelノイズを足して、最大のものを選ぶ」だけで、ちゃんと $\pi$ に従ってサンプリングできるのです。これが Gumbel-maxトリック です。

Gumbel-maxトリックが真の確率からカテゴリを正しくサンプリングすることを示す棒グラフ

水色が真の確率 $\pi_k$、オレンジがGumbel-maxトリックで4000回サンプリングした経験頻度です。両者がほぼ重なっており、トリックが正しくカテゴリ分布を再現していることが分かります。reparameterizationと同じく、乱数 $g_k$ が $\theta$(ここでは $\pi$)に依存しない形に追い出せている点が重要です。

softmaxによる緩和

Gumbel-maxトリックで乱数は外に出せましたが、肝心の $\arg\max$ がまだ微分できません。そこで $\arg\max$ を、温度パラメータ $\tau$ つきの softmax で「緩和(relax)」します。

$$ y_k = \frac{\exp\big((\log \pi_k + g_k)/\tau\big)}{\sum_{j} \exp\big((\log \pi_j + g_j)/\tau\big)} $$

出力 $y = (y_1,\dots,y_K)$ は、各要素が0以上で和が1になる連続なベクトルです。これは1-hotベクトルの「やわらかい版」で、しかも全体が微分可能です。温度 $\tau$ がこの「やわらかさ」を制御します。

  • $\tau$ が大きいと、$y$ は一様分布に近づき、なめらかになります(勾配は安定だが離散から遠い)。
  • $\tau \to 0$ にすると、$y$ は1-hotベクトル($\arg\max$)に近づきます(離散に忠実だが勾配が暴れる)。

温度tauを下げるとGumbel-softmaxが1-hotに近づく様子の温度スイープ図

同じGumbelノイズで $\tau$ だけを変えた4つの出力です。$\tau=2.0$ では4カテゴリにそれなりに分散していますが、$\tau$ を下げるにつれて1つのカテゴリに質量が集中し、$\tau=0.1$ ではほぼ1-hot(離散サンプル)になっています。温度を下げると離散に近づくが、勾配のばらつきは増える――この緩和(離散への忠実さ)となめらかさ(勾配の質)のトレードオフこそ、Gumbel-softmaxの実装で $\tau$ を調整する理由です。実用では、学習初期は $\tau$ を大きく、徐々に下げる「アニーリング」がよく使われます。

straight-through推定量

「順伝播では完全に離散な値を使いたいが、逆伝播では勾配を流したい」――この欲張りを叶えるのが straight-through(ST)推定量 です。

straight-through推定量で順伝播は離散・逆伝播は連続を使う仕組みの図

アイデアは図のとおりです。順伝播(forward) では $y$ に $\arg\max$ をかけた完全な1-hotベクトル $y_{\text{hard}}$ を使います。一方、逆伝播(backward) では $\arg\max$ を「なかったこと」にして、連続な $y$ の勾配をそのまま下流へ流します。実装では y_hard - y.detach() + y のような式で、値は $y_{\text{hard}}$、勾配は $y$、という二枚舌を実現します。これにより、ネットワークの入力には真に離散なベクトルが入りつつ、$\theta$ には勾配が届きます。

STは「順伝播の離散性」が重要なタスク(離散表現の学習、ハードな注意機構など)で効果を発揮します。勾配はあくまで近似なので不偏ではありませんが、実用上はうまく機能することが多い手法です。

ここまでで3つの道具が揃いました。最後に、どの場面でどれを選ぶかを整理します。

3手法の使い分け

3つの推定量は競合するというより、適用できる条件が異なるので、状況に応じて選びます。

3つの勾配推定法の適用条件マップ

表を縦に読むと、それぞれの素性がはっきりします。

  • REINFORCE(スコア関数):対象が連続でも離散でもよく、$f$ が微分不可能な黒箱でも使えます。最も一般的ですが分散が高く、baselineなどの分散低減が事実上必須です。強化学習・方策勾配の中心的な道具です。
  • reparameterization:$f$ が微分可能で、分布が $z=g(\theta,\varepsilon)$ の形に書けるときに限り使えます。そのぶん圧倒的に低分散です。連続潜在変数を扱う VAE の標準テクニックです。
  • Gumbel-softmax:離散カテゴリ変数のためのreparameterizationです。Gumbelノイズで乱数を外に出し、温度 $\tau$ つきsoftmaxで微分可能にします。離散潜在変数や構造予測で使われます。分散は $\tau$ に依存します。

選び方の指針はシンプルです。まず「$f$ を微分できるか」「$z$ が連続か離散か」を問うこと。連続で微分可能ならreparameterization、離散ならGumbel-softmax、どちらも無理(黒箱・離散で緩和も難しい)ならREINFORCE+baseline、という順で検討すれば、ほとんどの場面をカバーできます。

なお、近年はこれらを組み合わせたり、より高度な制御変量(RELAX、REBARなど)でREINFORCEの分散をreparameterization並みに下げる研究も進んでいます。基礎となる3手法を押さえておけば、こうした発展手法も「どの弱点をどう補っているか」という視点で読み解けます。

ここまでは、直感と要点を優先して導出を「飛ばし気味」に進めてきました。本文だけでも使い分けは分かりますが、なぜそれぞれが正しいのか――推定量が本当に不偏なのか、Gumbel-maxが本当にカテゴリ分布を再現するのか――を腹の底から納得するには、証明を1行ずつ追うのが一番です。最後の節は、その「飛ばした証明」をすべて埋める補遺です。数学はできるがこの分野は初見、という読者が、外部の教科書を開かずに本記事だけで閉じられることを目指します。

補遺:飛ばした証明をすべて埋める

この補遺は読み飛ばしても大丈夫です。 ここから先は、本文で結果だけ述べた等式(不偏性・等価性・Gumbel-maxの正しさなど)を前提を明示して一行ずつ証明する部分です。前半までで「3手法の使い分け」のイメージは押さえられているので、結論を使えれば十分という方は次の節へ進んでください。証明を自分の手で追いたい方、論文を読めるようになりたい方のための部分です。

ここからは、本文で結果だけ述べた等式や主張を、前提を明示しながら全行で証明します。表記を最初にそろえておきます。

  • $\theta$ は分布のパラメータ(ベクトル)。$\nabla_\theta$ は $\theta$ に関する勾配(各成分の偏微分を並べたベクトル)を表す微分演算子です。
  • $\mathbb{E}_{z\sim p_\theta}[\,\cdot\,] = \int (\cdot)\, p_\theta(z)\, dz$ は分布 $p_\theta$ に関する期待値(積分)です。連続変数を例に書きますが、離散変数なら積分を総和 $\sum_z$ に読み替えればそのまま成り立ちます。
  • $\arg\max_i a_i$ は $a_i$ を最大にする添字 $i$ を返す操作、$\mathrm{softmax}(a)_i = e^{a_i}/\sum_j e^{a_j}$ は実ベクトルを「和が1の正のベクトル」に押し込む写像です。

証明全体で何度も使う正則性条件(積分と微分を交換してよいための条件)を先にまとめます。以降の「交換」はすべてこの条件のもとで正当化されます。

正則性条件(微分と積分の交換)。 関数 $h(\theta, z)$ について、(i) ほとんどすべての $z$ で $\theta \mapsto h(\theta,z)$ が微分可能で、(ii) ある可積分な関数 $M(z)$($\int M(z)\,dz < \infty$)が存在して $\|\nabla_\theta h(\theta,z)\| \le M(z)$ が $\theta$ の近傍で一様に成り立つとき、 $$ > \nabla_\theta \int h(\theta,z)\, dz = \int \nabla_\theta h(\theta,z)\, dz > $$ が成り立つ(ルベーグの優収束定理による微分積分交換定理)。加えて、$p_\theta(z)$ の台(support、$p_\theta(z)>0$ となる $z$ の範囲)が $\theta$ に依存しないことも仮定します。台が $\theta$ で動くと境界項が出て、以下の議論が崩れるためです。

実用上、正規分布など滑らかな指数型分布族はこの条件を満たします。一方、一様分布 $\mathrm{Uniform}(0,\theta)$ のように台の端が $\theta$ で動く分布では成り立たないことに注意してください。

証明1:スコア関数推定量(REINFORCE)の不偏性

主張。 正則性条件のもとで、 $$ \nabla_\theta\, \mathbb{E}_{z\sim p_\theta}[f(z)] = \mathbb{E}_{z\sim p_\theta}\big[\, f(z)\, \nabla_\theta \log p_\theta(z) \,\big]. $$

証明。 期待値を積分で書き、勾配を施します。$f(z)$ は $\theta$ を含まない($z$ だけの関数)ことを使います。 $$ \nabla_\theta\, \mathbb{E}_{z\sim p_\theta}[f(z)] = \nabla_\theta \int f(z)\, p_\theta(z)\, dz. $$ ここで正則性条件により微分と積分を交換します($h(\theta,z)=f(z)p_\theta(z)$ に上の定理を適用)。$\theta$ が掛かるのは $p_\theta(z)$ だけなので、勾配は $p_\theta$ にのみ作用します。 $$ = \int f(z)\, \nabla_\theta p_\theta(z)\, dz. $$ 次に対数微分トリックを使います。これは合成関数の微分 $\nabla_\theta \log p_\theta(z) = \dfrac{\nabla_\theta p_\theta(z)}{p_\theta(z)}$ を変形した $$ \nabla_\theta p_\theta(z) = p_\theta(z)\, \nabla_\theta \log p_\theta(z) $$ という恒等式です($p_\theta(z)>0$ の台の上で成立)。これを上の被積分関数に代入すると、 $$ = \int f(z)\, p_\theta(z)\, \nabla_\theta \log p_\theta(z)\, dz. $$ 最後に、$p_\theta(z)\,dz$ が期待値の測度であることに注目すると、この積分はそのまま $p_\theta$ に関する期待値です。 $$ = \mathbb{E}_{z\sim p_\theta}\big[\, f(z)\, \nabla_\theta \log p_\theta(z) \,\big]. \qquad \blacksquare $$ これで、右辺の期待値を $p_\theta$ からのサンプル平均で置き換えたモンテカルロ推定量が、左辺の真の勾配の不偏推定量(平均すれば真の値に一致する推定量)であることが示せました。

証明2:baseline(制御変量)が不偏性を保つこと

主張。 任意の定数 $b$(より一般に $z$ に依存しない量)について、 $$ \mathbb{E}_{z\sim p_\theta}\big[\, (f(z)-b)\, \nabla_\theta \log p_\theta(z) \,\big] = \nabla_\theta\, \mathbb{E}_{z\sim p_\theta}[f(z)]. $$ すなわち $f$ から $b$ を引いても勾配推定量の期待値は変わらない。

証明。 期待値の線形性で2項に分けます。 $$ \mathbb{E}\big[(f(z)-b)\nabla_\theta \log p_\theta(z)\big] = \underbrace{\mathbb{E}\big[f(z)\nabla_\theta \log p_\theta(z)\big]}_{(\mathrm{A})} – \underbrace{b\,\mathbb{E}\big[\nabla_\theta \log p_\theta(z)\big]}_{(\mathrm{B})}. $$ 第1項 (A) は証明1により $\nabla_\theta \mathbb{E}[f(z)]$ そのものです。よって第2項 (B) が $0$ であることを示せば証明が終わります。鍵はスコアの期待値が常に $0$ であること、すなわち $\mathbb{E}_{z\sim p_\theta}[\nabla_\theta \log p_\theta(z)] = 0$ です。これを示します。定義どおり書き、対数微分トリックを使うと、 $$ \mathbb{E}_{z\sim p_\theta}\big[\nabla_\theta \log p_\theta(z)\big] = \int p_\theta(z)\, \nabla_\theta \log p_\theta(z)\, dz = \int \nabla_\theta p_\theta(z)\, dz. $$ ここで正則性条件により微分と積分を交換します。 $$ = \nabla_\theta \int p_\theta(z)\, dz. $$ $p_\theta$ は確率密度なので積分は恒等的に $1$、その勾配は $0$ です。 $$ = \nabla_\theta 1 = 0. $$ したがって (B) $= b\cdot 0 = 0$ となり、 $$ \mathbb{E}\big[(f(z)-b)\nabla_\theta \log p_\theta(z)\big] = (\mathrm{A}) – 0 = \nabla_\theta\, \mathbb{E}[f(z)]. \qquad \blacksquare $$

分散を最小化する最適 baseline。 $b$ は期待値(不偏性)には影響しませんが、分散には影響します。推定量を $G(b) = (f(z)-b)\,s(z)$(ただし $s(z)=\nabla_\theta \log p_\theta(z)$ はスコア)と書くと、$\mathbb{E}[G(b)]$ は $b$ によらず一定なので、分散 $\mathrm{Var}[G(b)] = \mathbb{E}[G(b)^2] – (\mathbb{E}[G(b)])^2$ を最小化することは $\mathbb{E}[G(b)^2]$ を最小化することと同値です。1次元成分で $$ \mathbb{E}[G(b)^2] = \mathbb{E}\big[(f(z)-b)^2 s(z)^2\big] = \mathbb{E}[f^2 s^2] – 2b\,\mathbb{E}[f s^2] + b^2\,\mathbb{E}[s^2] $$ は $b$ の下に凸な2次関数です。$b$ で微分して $0$ と置くと、 $$ \frac{d}{db}\mathbb{E}[G(b)^2] = -2\,\mathbb{E}[f s^2] + 2b\,\mathbb{E}[s^2] = 0 \;\;\Longrightarrow\;\; b^\star = \frac{\mathbb{E}\big[f(z)\, s(z)^2\big]}{\mathbb{E}\big[s(z)^2\big]}. $$ つまり最適 baseline は「スコアの2乗 $s^2$ で重み付けした $f$ の加重平均」です。スコアの大きさを無視した単純平均 $\mathbb{E}[f]$ はその近似であり、実装で移動平均がよく使われる理由がここにあります。多次元の $\theta$ では各成分ごとにこの式が成り立ちます。$\blacksquare$

証明3:reparameterization トリックの等価性と経路微分

主張。 $z = g(\theta, \varepsilon)$ が $\varepsilon$ について(ほとんどいたるところ)可逆で滑らかな変換であり、$\varepsilon \sim p(\varepsilon)$ が $\theta$ に依存しないとき、(可逆性は下の等価性を変数変換で示すために使う仮定です。後半の勾配を期待値の中に入れる操作には、$g$ が $\theta$ について微分可能であることと微分・積分交換の正則性条件だけがあれば十分で、可逆性は本質的に不要です。) $$ \mathbb{E}_{z\sim p_\theta(z)}[f(z)] = \mathbb{E}_{\varepsilon\sim p(\varepsilon)}\big[f(g(\theta,\varepsilon))\big], $$ かつ正則性条件のもとで $$ \nabla_\theta\, \mathbb{E}_{z\sim p_\theta(z)}[f(z)] = \mathbb{E}_{\varepsilon\sim p(\varepsilon)}\big[\nabla_\theta f(g(\theta,\varepsilon))\big]. $$

証明(等価性)。 固定した $\theta$ のもとで、$z = g(\theta,\varepsilon)$ を $\varepsilon$ から $z$ への変数変換とみなします。変数変換の公式(ヤコビアンを伴う密度変換)により、$\varepsilon\sim p(\varepsilon)$ を変換して得られる $z$ の密度はちょうど $p_\theta(z)$ になります。具体的に1次元なら、$z=g(\theta,\varepsilon)$ の密度は $$ p_\theta(z) = p(\varepsilon)\left|\frac{d\varepsilon}{dz}\right|_{\varepsilon = g^{-1}(\theta,z)} $$ で与えられます(これが reparameterization が「同じ分布を作る」ことの定義です)。よって任意の可測関数 $f$ について、無意識の統計家の法則(変換した確率変数の期待値は、元の変数で測度を変換して計算してよい、という変数変換の定理)により $$ \mathbb{E}_{z\sim p_\theta(z)}[f(z)] = \int f(z)\, p_\theta(z)\, dz = \int f(g(\theta,\varepsilon))\, p(\varepsilon)\, d\varepsilon = \mathbb{E}_{\varepsilon\sim p(\varepsilon)}\big[f(g(\theta,\varepsilon))\big]. $$ 2つ目の等号で $z=g(\theta,\varepsilon)$ を代入し、$p_\theta(z)\,dz = p(\varepsilon)\,d\varepsilon$(ヤコビアンが密度変換とちょうど打ち消し合う)を使いました。これで2つの期待値が等しいことが示せました。

証明(勾配を期待値の中へ)。 等価性で得た右辺に勾配を施します。ここで決定的に効くのは、期待値をとる測度 $p(\varepsilon)$ が $\theta$ を含まないことです。したがって正則性条件のもとで微分と積分(期待値)を交換できます。 $$ \nabla_\theta\, \mathbb{E}_{\varepsilon\sim p(\varepsilon)}\big[f(g(\theta,\varepsilon))\big] = \nabla_\theta \int f(g(\theta,\varepsilon))\, p(\varepsilon)\, d\varepsilon = \int \nabla_\theta f(g(\theta,\varepsilon))\, p(\varepsilon)\, d\varepsilon = \mathbb{E}_{\varepsilon\sim p(\varepsilon)}\big[\nabla_\theta f(g(\theta,\varepsilon))\big]. $$ 証明1ではスコア $\nabla_\theta \log p_\theta$ を経由して密度を微分しましたが、ここでは被積分関数 $f(g(\theta,\varepsilon))$ を直接 $\theta$ で微分できる点が本質的な違いです。最後に連鎖律で展開すると(1次元 $z$ の場合)、 $$ \nabla_\theta f(g(\theta,\varepsilon)) = f'(g(\theta,\varepsilon))\, \nabla_\theta g(\theta,\varepsilon) $$ となり、$f$ の傾き $f’$ が推定量に明示的に現れます。これが経路微分(pathwise derivative) と呼ばれる所以です。$\blacksquare$

なぜ低分散か(直感)。 reparameterization 推定量は $f’$($f$ の傾き)という構造情報を勾配にそのまま使います。$\theta$ を少し動かすと $z=g(\theta,\varepsilon)$ がどちらへ動き、それで $f$ がどれだけ増減するかを、各サンプルが「局所的な傾き」として教えてくれます。同じ $\varepsilon$ なら $\theta$ の摂動に対する応答が滑らかにつながるため、サンプル間のばらつきが抑えられます。一方スコア関数推定量は $f$ のしか使わず、傾きの情報を捨てています。次の証明4で、この差が分散の差として定量的に効くことを見ます。

証明4:REINFORCE と reparameterization の分散差

設定。 直感を壊さずに最小の例で比較します。$z\sim\mathcal{N}(\theta,\sigma^2)$($\sigma$ は固定、平均 $\theta$ で微分)とし、$f(z)$ を考えます。正規分布のスコアは $$ \log p_\theta(z) = -\frac{(z-\theta)^2}{2\sigma^2} + \text{const}, \qquad \nabla_\theta \log p_\theta(z) = \frac{z-\theta}{\sigma^2} $$ です。reparameterization は $z = \theta + \sigma\varepsilon,\ \varepsilon\sim\mathcal{N}(0,1)$ とします。

スコア関数推定量(単一サンプル)。 $$ G_{\mathrm{SF}} = f(z)\,\frac{z-\theta}{\sigma^2}, \qquad z\sim\mathcal{N}(\theta,\sigma^2). $$ ここで $f(z)$ の値そのものが因子 $\frac{z-\theta}{\sigma^2}$ に掛かっています。$f$ を定数だけ持ち上げて $f \to f + C$ としても勾配の真値は変わらない(証明2)のに、$G_{\mathrm{SF}}$ には $C\cdot\frac{z-\theta}{\sigma^2}$ という余分な揺らぎが乗り、その分散は $$ \mathrm{Var}\!\left[C\,\frac{z-\theta}{\sigma^2}\right] = \frac{C^2}{\sigma^4}\,\mathrm{Var}[z-\theta] = \frac{C^2}{\sigma^2} $$ と、定数 $C$(=$f$ の絶対的なスケール)が大きいほど際限なく増えます。$f$ のスケールがそのまま分散に乗るのがスコア関数推定量の弱点です(baseline はこの $C$ を引き戻して打ち消す装置にほかなりません)。

reparameterization 推定量(単一サンプル)。 $$ G_{\mathrm{RP}} = \nabla_\theta f(\theta + \sigma\varepsilon) = f'(\theta + \sigma\varepsilon), \qquad \varepsilon\sim\mathcal{N}(0,1). $$ こちらは $f$ の傾き $f’$ だけが現れ、$f$ への定数加算 $f\to f+C$ は $f’$ を一切変えません($\frac{d}{dz}C = 0$)。したがって $f$ のスケール由来の揺らぎは構造的にキャンセルされ、分散は $f’$ の変動だけで決まります。具体的に $f(z)=z^2$ なら $f'(z)=2z$ なので $G_{\mathrm{RP}} = 2(\theta+\sigma\varepsilon)$、その分散は $$ \mathrm{Var}[G_{\mathrm{RP}}] = 4\sigma^2\,\mathrm{Var}[\varepsilon] = 4\sigma^2. $$ 一方スコア関数側は $G_{\mathrm{SF}} = z^2\frac{z-\theta}{\sigma^2}$ で、$z$ の3次モーメントまで効くため分散はずっと大きくなります(後述の実測がそれを裏づけます)。

実測との接続。 本文の Python 実測では、同じトイ問題 $J=\mathbb{E}[z^2]$、$\mu=1,\ \sigma=1.5$ で $$ \mathrm{Var}[G_{\mathrm{RP}}] \approx 8.89, \qquad \mathrm{Var}[G_{\mathrm{SF}}] \approx 47.53 $$ でした。reparameterization の理論値は上式で $\sigma=1.5$ を入れると $4\sigma^2 = 4\times 2.25 = 9.0$ となり、実測 $8.89$ とよく一致します。スコア関数側が約 $5$ 倍大きいのは、まさに「$f(z)=z^2$ のスケールがスコア $\frac{z-\theta}{\sigma^2}$ に乗って増幅される」ためで、baseline を引くと $47.53\to30.52$ へ下がるのも、上で見た定数項 $C$ の打ち消しとして理解できます。$\blacksquare$

証明5:Gumbel-max トリックの正しさ

主張。 カテゴリ確率を $\pi_1,\dots,\pi_K$($\pi_i>0,\ \sum_i\pi_i=1$)とし、独立な $G_i\sim\mathrm{Gumbel}(0,1)$ を引くとき、 $$ k = \arg\max_i\big(\log\pi_i + G_i\big) $$ の分布は $P(k=i)=\pi_i$ に一致する。

準備(Gumbel 分布の CDF/PDF)。 $\mathrm{Gumbel}(0,1)$ の累積分布関数(CDF)と確率密度(PDF)は $$ F(g) = \exp(-e^{-g}), \qquad f(g) = F'(g) = e^{-g}\exp(-e^{-g}) $$ です。$G_i = -\log(-\log U_i),\ U_i\sim\mathrm{Uniform}(0,1)$ で生成できることも、$P(G_i\le g) = P(U_i \le e^{-e^{-g}}) = e^{-e^{-g}}$ から確認できます。

証明。 $a_i := \log\pi_i$ と置きます。カテゴリ $i$ が最大になる確率を、$G_i = g$ で条件付けてから周辺化して計算します。$i$ が最大であるとは「すべての $j\ne i$ で $a_j + G_j < a_i + g$」、すなわち $G_j < a_i - a_j + g$ が同時に成り立つことです。$G_j$ たちは独立なので、その同時確率は各 CDF の積になります。 $$ P\big(k=i \,\big|\, G_i = g\big) = \prod_{j\ne i} F\big(a_i - a_j + g\big) = \prod_{j\ne i}\exp\!\big(-e^{-(a_i - a_j + g)}\big). $$ 指数の積は指数の和にまとめられます。 $$ = \exp\!\Big(-\sum_{j\ne i} e^{-(a_i - a_j + g)}\Big) = \exp\!\Big(-e^{-g}\,e^{-a_i}\sum_{j\ne i} e^{a_j}\Big). $$ ここで $G_i$ について周辺化します。$G_i$ の密度 $f(g)=e^{-g}\exp(-e^{-g})$ を掛けて積分します。 $$ P(k=i) = \int_{-\infty}^{\infty} f(g)\, P(k=i\mid G_i=g)\, dg = \int_{-\infty}^{\infty} e^{-g}\exp(-e^{-g})\, \exp\!\Big(-e^{-g}\,e^{-a_i}\sum_{j\ne i} e^{a_j}\Big)\, dg. $$ 2つの指数をまとめます。$i$ 自身の項を加えて $\sum_{j\ne i}e^{a_j} + e^{a_i} = \sum_{j} e^{a_j}$ とできるよう、$\exp(-e^{-g})$ を $\exp(-e^{-g}e^{-a_i}e^{a_i})$ と書き直してから合体させます。 $$ = \int_{-\infty}^{\infty} e^{-g}\, \exp\!\Big(-e^{-g}\,e^{-a_i}\big(\textstyle\sum_{j}e^{a_j}\big)\Big)\, dg. $$ ここで $S := \sum_j e^{a_j}$ と置き、変数変換 $t = e^{-g}$ を行います。$dt = -e^{-g}\,dg$ すなわち $e^{-g}\,dg = -dt$ で、$g:-\infty\to\infty$ のとき $t:\infty\to 0$ です。符号と積分範囲の反転が打ち消し合って、 $$ = \int_{0}^{\infty} \exp\!\big(-t\, e^{-a_i} S\big)\, dt = \left[\frac{-1}{e^{-a_i}S}\exp\!\big(-t\,e^{-a_i}S\big)\right]_{0}^{\infty} = \frac{1}{e^{-a_i}S} = \frac{e^{a_i}}{S}. $$ 最後に $a_i = \log\pi_i$ を戻すと $e^{a_i} = \pi_i$、$S = \sum_j e^{a_j} = \sum_j \pi_j = 1$ なので、 $$ P(k=i) = \frac{\pi_i}{\sum_j \pi_j} = \frac{\pi_i}{1} = \pi_i. \qquad \blacksquare $$ これで、Gumbel ノイズを足して $\arg\max$ を取る操作が、厳密にカテゴリ分布 $\mathrm{Cat}(\pi)$ からのサンプリングと一致することが示せました。本文の棒グラフ(経験頻度と $\pi$ の一致)は、この定理の数値的な確認になっています。なお途中で正規化していない $a_i = \log w_i$($w_i$ は非正規化重み)を使っても、最後の式が $w_i / \sum_j w_j$ となり、自動的に正規化されたカテゴリ確率を返す点も読み取れます。

証明6:Gumbel-softmax 緩和と straight-through

緩和。 $\arg\max$ は微分できないので、温度 $\tau>0$ つきの softmax で連続化します。$\arg\max$ の入力を $h_i := \log\pi_i + G_i$ として、 $$ y_i = \mathrm{softmax}\big(h/\tau\big)_i = \frac{\exp(h_i/\tau)}{\sum_j \exp(h_j/\tau)}. $$ 出力 $y=(y_1,\dots,y_K)$ は $y_i>0,\ \sum_i y_i = 1$ を満たす連続ベクトルで、$h$ について(したがって $\theta=\log\pi$ について)滑らかに微分できます。

$\tau\to 0$ で 1-hot($\arg\max$)に収束すること。 最大成分の添字を $m = \arg\max_i h_i$ とし、唯一の最大(同点なし、確率1で成立)と仮定します。$i\ne m$ の成分について、softmax の比を $y_m$ で割って整理すると $$ \frac{y_i}{y_m} = \exp\!\Big(\frac{h_i – h_m}{\tau}\Big). $$ $h_i – h_m < 0$($m$ が最大)なので、$\tau\to 0^+$ のとき指数の中身は $-\infty$ に発散し、$y_i/y_m \to 0$。すなわち $i\ne m$ の全成分が $y_m$ に対して無視できるほど小さくなります。$\sum_i y_i = 1$ と合わせると $y_m\to 1$、$y_i\to 0\ (i\ne m)$、つまり $y$ は $m$ 番目だけが $1$ の 1-hot ベクトル、これは $\arg\max$ の出力そのものです。逆に $\tau\to\infty$ では $h_i/\tau\to 0$ で全成分が等しくなり、$y\to(1/K,\dots,1/K)$ の一様ベクトルに近づきます。本文の温度スイープ図はこの両極限の補間を可視化したものです。

微分可能性(なぜ勾配が流せるか)。 softmax は初等関数の合成なので $y_i$ は $h_j$ で偏微分でき、よく知られた形 $$ \frac{\partial y_i}{\partial h_j} = \frac{1}{\tau}\, y_i\big(\delta_{ij} – y_j\big) \qquad(\delta_{ij}\text{ はクロネッカーのデルタ}) $$ を持ちます。さらに $h_i = \log\pi_i + G_i$ で $G_i$ は $\theta$ に依存しないため、$\partial h_i/\partial(\log\pi_i) = 1$ を通じて勾配が $\theta$ まで連鎖律で届きます。これは reparameterization と同じ構図――乱数 $G_i$ を外に出し、残りを微分可能関数でつなぐ――の離散版です。$\tau$ が小さいほど $\partial y_i/\partial h_j$ に $1/\tau$ が掛かって勾配が大きく振れる(高分散になる)ことも、この式から読み取れます。

straight-through(ST)推定量。 順伝播では真に離散な 1-hot を使いたいが、逆伝播では連続な勾配を流したい、という要求に応えます。順伝播の出力を $y_{\text{hard}} = \mathrm{onehot}(\arg\max_i y_i)$ とし、実装では $$ y_{\text{ST}} = \underbrace{(y_{\text{hard}} – \mathrm{sg}[y])}_{\text{値は }y_{\text{hard}},\ \text{勾配は }0} + \underbrace{y}_{\text{値も勾配もあり}} $$ と書きます($\mathrm{sg}[\cdot]$ は stop-gradient = 順伝播では恒等、逆伝播では勾配 $0$ とみなす演算で、PyTorch では .detach())。順伝播では $\mathrm{sg}[y]=y$ なので値は $y_{\text{hard}} – y + y = y_{\text{hard}}$ となり離散ベクトルが流れます。逆伝播では $y_{\text{hard}}$ と $\mathrm{sg}[y]$ の勾配が共に $0$ なので、$\partial y_{\text{ST}}/\partial\theta = \partial y/\partial\theta$ となり、連続な $y$ の勾配がそのまま下流へ伝わります。この勾配は「$\arg\max$ の勾配を $\mathrm{softmax}$ の勾配で代用する」近似なので不偏ではありませんが、順伝播の離散性が重要なタスクで実用的に機能します。$\blacksquare$

以上で、本文で結果だけ述べた6つの命題――REINFORCE の不偏性、baseline の不偏性と最適形、reparameterization の等価性と経路微分、両者の分散差、Gumbel-max の正しさ、Gumbel-softmax 緩和と ST――をすべて全行で埋めました。証明を一度たどっておくと、本文の実測値や図が「なぜそうなるか」まで含めて腑に落ちるはずです。

まとめ

本記事では、サンプリングを挟んだ期待値 $\mathbb{E}_{z\sim p_\theta(z)}[f(z)]$ の勾配を推定する3つの方法を解説しました。

  • 問題の核心:期待値をとる分布そのものが $\theta$ に依存するため、サンプリングの箇所で計算グラフが切れ、素直に微分できない。
  • スコア関数推定量(REINFORCE):対数微分トリックで勾配を $\mathbb{E}[f(z)\nabla_\theta \log p_\theta(z)]$ と期待値の形に戻す。$f$ が黒箱・離散でも使える万能選手だが高分散。baselineで不偏性を保ったまま分散低減。
  • reparameterizationトリック:$z=g(\theta,\varepsilon)$ と乱数を経路の外に追い出し、$\nabla_\theta f$ を直接計算。$f$ の傾き情報を使うため低分散。VAEの要。
  • Gumbel-softmax:Gumbel-maxトリックでカテゴリサンプリングを「ノイズ+$\arg\max$」に分解し、温度 $\tau$ つきsoftmaxで微分可能に緩和。離散変数のreparameterization。
  • 実測:トイ問題で3手法とも不偏、reparameterの分散はREINFORCEの約1/5、baselineで分散低減を確認した。

この3手法は、生成モデルと強化学習をつなぐ共通言語です。VAEのreparameterizationも、方策勾配のREINFORCEも、根は同じ「確率的勾配の推定」問題でした。次のステップとして、以下の記事も参考にしてください。