はじめに
今回は期待値 $\mathbb E[f _ \theta (z)]$ という形式で表される式の勾配を求める方法について説明します。
モチベーション
通常の機械学習においては、例えば二乗損失関数
$$ l(\theta) = \frac{1}{N} \sum _ i (y _ i - f _ \theta( x _ i )) ^ 2 $$
に関する $\theta$ の勾配を計算し、勾配法によってモデル $f _ \theta (\cdot)$ を学習させるなどのことを行います。これは実際のところ、データ分布からのサンプリング
$$ X, Y \sim p(x, y) $$
によって得られたであろうデータ $\{ (x _ i , y _ i) \} _ {i = 1, \cdots, N} $ を用いて
$$ \mathbb E _ {p(x, y)} [ (y - f _ \theta( x )) ^ 2] \simeq \frac{1}{N} \sum _ i (y _ i - f _ \theta( x _ i )) ^ 2 $$
と有限データで近似しているというものです。通常、データ分布での期待値の方を予測損失と呼び、データを使って計算している右辺の方を経験損失と呼びます。
さて、変分推論や強化学習では、本来期待値で表されている式を、手持ちのデータで近似するという方法が取れない場合があります。例えば、変分推論では、変分近似分布 $q(z)$ による様々な関数の期待値を要求されます（例えばVAEでは"再構成誤差項"と"近似分布と事前分布とのKLダイバージェンス項"に関して、それぞれ変分近似分布での期待値を計算しなければならない）。すなわち、一般的に $q (z)$ に従う確率変数 $Z$ と何らかのパラメトリックな関数 $f _ \theta(z)$ があり、
$$ \mathbb E _ {q(z)} [f _ \theta (z)] $$
という式が、目的関数の中にそのまま含まれているようなケースがあるということです。そのような場合でも勾配法を実施したい場合は
$$ \frac{\partial }{\partial \theta} \mathbb E _ {q (z)} [ f _\theta (z)] $$
を評価しなければ行けないということになります。そのような場合にどうすれば良いのかというのが今回の題材です。特に $f _ \theta(z)$ という関数のパラメータ $\theta$ を期待値を取る分布の方も $q (z) = q _ \theta (z)$ と取っているケースを想定します。
なぜなら、$q(z)$ は何らかのパラメトリックなモデルであることが多く、かつ、$f(z)$ はその $q (z)$ を評価するための何がしかの式であるケースが多々あるからです。すなわち、期待値を取る分布の方にも期待値を取られる関数の方にもパラメータ $\theta$ があり、その $\theta$ に関する勾配
$$ \frac{\partial }{\partial \theta} \mathbb E _ {q _ \theta (z)} [ f _ \theta (z)] $$
を計算したいという問題を考えていきます。
期待値の勾配の計算方法
モンテカルロ近似と再パラメータ化
最も簡単に思いつくのは、兎にも角にも $\mathbb E _ {p(x)} [f(x)] $ という形式に持っていってしまい、予測損失を経験損失で代替しているのと同じ発想で、分布 $p (x)$ から実際に乱数を生成して近似してしまえという方法です。
今、
$$ \frac{\partial }{\partial \theta} \mathbb E _ {q _ \theta (z)} [ f _ \theta (z)] $$
という実際に評価したい"期待値の勾配"を"○○の期待値"というものにすげ替える事ができれば上記の発想が使えます。まずは、ちゃんと式を見てみましょう。
$$ \frac{\partial }{\partial \theta} \mathbb E _ {q _ \theta(z)} [ f _ \theta( z )) ] = \frac{\partial }{\partial \theta} \int _ z f _ \theta (z) q _ \theta (z) dz = \int _ z \frac{\partial }{\partial \theta} \left [ f _ \theta (z) q _ \theta (z) \right] dz $$
となっています（ここで積分と微分が交換できるようなものを扱っているということを前提としました）。さて、これは期待値の形をしていませんのでそのままではモンテカルロ近似を実行できません。そこで再パラメータ化というテクニックを用います。
$Z \sim q _ \theta (z)$ という確率分布からのサンプリングは下記のような "$\theta$ に依らない分布から生ずる確率変数" を$\theta$ に依存する関数 $g$ で変換するという形式に置き換えます。
$$ \begin{align} \alpha &\sim q(\alpha) \\ z &= g _ \theta (\alpha) \end{align} $$
このようなことができる確率分布 $q _ \theta (z) $ を扱っているのだと思えば $f _ \theta (z) = f (g _ \theta (\alpha))$ の期待値を計算する話にしてしまうことで
$$ \begin{align} \frac{\partial }{\partial \theta} \mathbb E _ {q _ \theta (z)} [ f _ \theta( z )) ] &= \mathbb E _ {\alpha \sim q(\alpha)} [ \frac{\partial }{\partial \theta} f _ \theta (g _ \theta (\alpha) )] \\ &\simeq \frac{1}{M} \sum _ i \frac{\partial }{\partial \theta} f _ \theta(g _ \theta (\alpha _ i) ) \end{align} $$
と書くことができるということです。仮に上記のようにできれば、和の中身は合成関数 $f _ \theta (g _ \theta(\cdot))$ の勾配の計算式になっており、バックプロパゲーション等が利用でき、また、モンテカルロ近似によって期待値の勾配の不偏推定量が得られます。最も単純な例は $q _ \theta (z) = \mathcal N (z \mid \mu, \sigma)$ のような場合です。この場合は $g _ \theta (\alpha) = \mu + \sigma \alpha$ という変換を作れば、$\alpha \sim \mathcal N(0, 1)$ というパラメータに依らない固定された分布からのサンプリングで乱数生成を実施することができます。
再パラメータ化できない場合
さて、再パラメータ化できない場合はどうしましょう。まずは評価したい勾配を、数式ベースで真面目に追っていくことをしてみます。
$$ \frac{\partial }{\partial \theta} \mathbb E _ {q _ \theta(z)} [ f _ \theta( z )) ] = \frac{\partial }{\partial \theta} \int _ z f _ \theta (z) q _ \theta (z) dz $$
という積分をちゃんと書いた式になります。さて、ここで、$z$ に関する積分と $\theta$ に関する微分が交換できるようなケースを考えることにしましょう。そうすると微分が積分の中に入り、
$$ \begin{align} \frac{\partial }{\partial \theta} \int _ z f _ \theta (z) q _ \theta (z) dz &= \int _ z \frac{\partial }{\partial \theta}\left [ f _ \theta (z) q _ \theta(z) \right] dz\\ & = \int _ z \left[ \frac{\partial }{\partial \theta} f _ \theta (z) q _ \theta (z) + f _ \theta (z) \frac{\partial }{\partial \theta} q _ \theta (z) \right] dz \end{align} $$
と、一筋縄ではいかなさそうな形になってしまいました。ところで少々飛躍的ではありますが、 期待値を計算するためには $p(x)dx$ みたいな形が欲しいわけなので、この形をひねり出すために対数関数の微分を応用してみることにします。
$$ \frac{\partial }{\partial \theta} \log (h) = \frac{1}{h} \frac{\partial }{\partial \theta} h $$
となる一般的な性質を用いれば、
$$ \frac{\partial }{\partial \theta} q _ \theta (x) = q _ \theta (z) \frac{\partial }{\partial \theta} \log \{ q _ \theta (x) \} $$
という形を得ることができます。これを先程の積分の中にねじ込んでやると
\begin{align} \frac{\partial }{\partial \theta} \int _ z f _ \theta (z) q _ \theta (z) dz & = \int _ z \left[ \frac{\partial }{\partial \theta} f _ \theta (z) q _ \theta (z) + f _ \theta (z) \frac{\partial }{\partial \theta} q _ \theta (z) \right] dz \\ & = \int _ z \left[ \frac{\partial }{\partial \theta} f _ \theta (z) q _ \theta (z) + f _ \theta (z) q _ \theta (z) \frac{\partial }{\partial \theta} \log \{ q _ \theta (x) \} \right] dz \\ & = \int _ z \left[ \frac{\partial }{\partial \theta} f _ \theta (z) + f _ \theta (z) \frac{\partial }{\partial \theta} \log \{q _ \theta (x)\} \right] q _ \theta (z) dz \\ & = \mathbb E _ {q _ \theta(z)}\left[\frac{\partial }{\partial \theta} f _ \theta (z) + f _ \theta (z) \frac{\partial }{\partial \theta} \log \{q _ \theta (x)\}\right] \end{align}
と期待値の形でなんとか表してやることに成功しました。後は $z \sim q _ \theta(z)$ からサンプリングをしてみてモンテカルロ近似を実施すればよいということになります。これは不偏推定量ではあるのですが、実を言うと分散がかなり大きくなり、サンプルサイズが小さいと勾配の推定量として信用がならん値を使うことになります（無論 $q _ \theta$ からサンプリングを実施するのは時間さえあれば可能かもしれないが、通常、学習を実施する1ステップにどれ程の時間を掛けることが許されるのだろうか…）。
一先ずココまでで、期待値の勾配をモンテカルロ近似する術を手に入れたので、数式上で導出された期待値で表される何らかの目的関数に対する勾配が、具体的にどのようにコンピュータで計算されているかの大枠を知ったことになるでしょう。実際には推定量の分散を抑えるための工夫が施されます（ベースラインと呼ばれる手法が、変分推論でも強化学習の方策勾配法でも知られている）。
