『直感的に理解するTransformerの仕組み』の続編である『仕組みから理解するChatGPT』の印刷版のサポートページです。主に追加コンテンツや誤植が見つかった場合の正誤表の作成、カラー画像の確認が行えるように作成を行いました。誤植につきましては見つかり次第都度追加いたしますので、お気づきの方は気軽にご指摘ください。
・サポートページ：直感的に理解するTransformerの仕組み
・仕組みから理解するChatGPT
Contents
追加コンテンツ
$(\mathrm{o.xx})$形式の式番号は『仕組みから理解するChatGPT』の式番号に対応しますのでご注意ください。
Transformer decoder
『仕組みから理解するChatGPT』ではGPT$3$に用いられるTransformer decoderの詳細は取り扱わなかったので、以下では詳しく確認を行う。
ChatGPTのベースに用いられるGPT$3$の主要な処理はTransformerに基づく一方で、Encoder-Decoder形式のTransformerではなくDecoder-onlyのTransformerが用いられる。このようなDecoder-onlyのTransformerはTransformer decoder論文の処理に基づく。
上記に出てくる$(m^{1}, \cdots , m^{n}) \mapsto (y^{1}, \cdots , y^{\eta})$は入力の$(m^{1}, \cdots , m^{n})$を出力の$(y^{1}, \cdots , y^{\eta})$に変換するタスクを表す。このタスクをそのまま取り扱うにはオーソドックスなTransformerのようにEncoder-Decoder形式のDeepLearningを用いる必要があるので、Transformer decoderの論文では$(m^{1}, \cdots , m^{n}) \mapsto (y^{1}, \cdots , y^{\eta})$を$(w^{1}, \cdots , w^{n+\eta+1}) = (m^{1}, \cdots , m^{n}, \delta , y^{1}, \cdots , y^{\eta})$のように置き換え、下記のように定義する同時確率の最大化のタスクの最大化に基づくパラメータの学習を行う。
$$
\large
\begin{align}
p(w^{1}, \cdots , w^{n+\eta+1}) = \prod_{j=1}^{n+\eta} p(w^{j+1}|w^{1}, \cdots , w^{j})
\end{align}
$$
上記に基づいてタスクの定義を行うことで、Encoderを省略し、Decoderのみを用いた処理が可能になる。このような処理を行うことで、必要なパラメータを半分にすることができる。
Decoder onlyのTransformerについては下記でも詳しく取り扱った。
モンテカルロ法と確率分布
InstructGPTで用いる方策勾配法では勾配の期待値の近似にあたってサンプリングに基づくモンテカルロ法を用いる。モンテカルロ法による期待値の近似については$(2.27)$式で簡単に取り扱った。
$$
\large
\begin{align}
\mu = \mathbb{E}[X] \simeq \frac{1}{n} \sum_{i=1}^{n} x_i \quad (2.27)
\end{align}
$$
上記のモンテカルロ法による近似計算は『仕組みから理解するChatGPT』では下記のようにサンプリングを行った結果に基づいて定義した。
$$
\large
\begin{align}
x_1, x_2, \cdots , x_n \sim \mathcal{N}(\mu, \sigma^{2}), \quad \mathrm{i.i.d.}
\end{align}
$$
ここで注意が必要なのが、「期待値計算で用いられる確率分布とサンプリングに用いる確率分布が一致する必要がある」ということである。$(2.27)$式の期待値$\mathbb{E}[X]$は正規分布$\mathcal{N}(\mu, \sigma^{2})$の確率密度関数$f(x)$を用いて下記のように定義される。
$$
\large
\begin{align}
\mathbb{E}[X] = \int_{-\infty}^{\infty} x f(x) dx \quad (1)
\end{align}
$$
上記の$(1)$式に用いる確率密度関数$f(x)$と$(2.27)$式に用いるサンプルの$x_1, x_2, \cdots , x_n$は同一の確率分布$\mathcal{N}(\mu, \sigma^{2})$に対応する。
「期待値計算で用いられる確率分布とサンプリングに用いる確率分布が一致する必要がある」ことを逆に考えると、「モンテカルロ法に基づく近似を行うにあたってはサンプリングを行う確率分布に基づく期待値の式を得る必要がある」と解釈できる。$3$-$2$-$3$節の「方策勾配法における勾配の計算」や$5$-$2$-$3$節の「Reinforcement learning」でLog-Derivative Trickを逆に用いて期待値の式を得るのは「サンプリングを行う分布の期待値の式を得ることでモンテカルロ法による近似を行う」というのが主な狙いである。
ここまでは「サンプリングに用いる確率分布」と「期待値の定義式に出てくる確率関数・確率密度関数に対応する確率分布」が一致する場合を確認したが、一致しない場合に用いられるのが重点サンプリング(Importance Sampling)である。重点サンプリングでは下記のような式が出てくる。
$$
\large
\begin{align}
\mathbb{E}_{\pi}[x] &= \sum x \pi(x) \\
&= \sum x \frac{b(x)}{b(x)} \pi(x) \\
&= \sum x \frac{\pi(x)}{b(x)} b(x) = \mathbb{E}_{b} \left[ \frac{\pi(x)}{b(x)} x \right]
\end{align}
$$
重点サンプリングは上記の式変形に基づいて「特定の確率分布の期待値を別の確率分布からサンプリングした値に基づいて計算する手法」である。式変形の解釈については詳しくは下記で取りまとめたので当項では省略する。
InstructGPTでは重点サンプリングは出てこないが、TRPOやPPOの論文では重点サンプリングが出てくる。たとえば『仕組みから理解するChatGPT』の$(3.28)$式には重点サンプリングの式が用いられる。
$$
\large
\begin{align}
L^{KLPEN}(\theta) = \mathbb{E}_{t} \left[ \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t|s_t)} A_t – \mathrm{KL_term} \right] \quad (3.28)
\end{align}
$$
上記に出てくる$\displaystyle \frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta_{\mathrm{old}}}(a_t|s_t)}$は重点サンプリングに基づいて用いられる。「InstructGPTでは重点サンプリングが出てこないのに対してPPOでは重点サンプリングが用いられる」のは、InstructGPTでは$\pi_{\phi}^{RL}$を用いてサンプリングを行っているのに対し、PPOでは$\pi_{\theta_{\mathrm{old}}}$を用いてサンプリングを行うからである。
InstructGPTの強化学習における目的関数の式をPPOの式の形式を元に解釈するにあたっては、当項で取り扱ったように「サンプリングをどの確率分布に基づいて行ったか」の視点が重要である。
a per-token KL penaltyの勾配
$$
\large
\begin{align}
\mathbb{E}_{(x,y) \sim D_{\pi_{\phi}^{RL}}} \left[ \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] = \sum_{y_i \in \mathcal{Y}} \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i}) \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \quad (5.12)’
\end{align}
$$
上記は概ね参照元の$(5.12)$式に対応するが、$\displaystyle \sum$がわかりやすくなるように出力層に対応する語彙の集合$\mathcal{Y}$を用いて$\displaystyle \sum_{y_i \in \mathcal{Y}}$のように表した。
この$\displaystyle \sum_{y_i \in \mathcal{Y}}$の中の項の$\phi$に関する勾配は下記のように得られる。
$$
\begin{align}
\nabla_{\phi} \left[ \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i}) \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] = \nabla_{\phi} \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i}) \left[ 1 + \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] \quad (5.13)
\end{align}
$$
ここで$(5.13)$式にLog-Derivative Trickを用いると、$(5.12)$式の勾配は下記のように変形できる。
$$
\large
\begin{align}
& \nabla_{\phi} \mathbb{E}_{(x,y) \sim D_{\pi_{\phi}^{RL}}} \left[ \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] \\
&= \nabla_{\phi} \sum_{y_i \in \mathcal{Y}} \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i}) \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \quad (5.12)’ \\
&= \sum_{y_i \in \mathcal{Y}} \nabla_{\phi} \left[ \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i}) \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] \\
&= \sum_{y_i \in \mathcal{Y}} \nabla_{\phi} \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i}) \left[ 1 + \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] \quad (5.13) \\
&= \sum_{y_i \in \mathcal{Y}} \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i}) \frac{\nabla_{\phi} \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})} \left[ 1 + \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] \\
&= \mathbb{E}_{(x,y) \sim D_{\pi_{\phi}^{RL}}} \left[ \frac{\nabla_{\phi} \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})} \left( 1 + \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right) \right] \\
&= \mathbb{E}_{(x,y) \sim D_{\pi_{\phi}^{RL}}} \left[ \nabla_{\phi} \log{[\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})]} \left( 1 + \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right) \right] \quad (2)
\end{align}
$$
ここでさらに$(5.10)$式を用いることで、$(5.8)$式のKL penaltyに関する項の勾配は下記のように得ることができる。
$$
\large
\begin{align}
& \nabla_{\phi} \mathbb{E}_{(x,y) \sim D_{\pi_{\phi}^{RL}}} \left[ – \beta \log{ \frac{\pi_{\phi}^{RL}(y|x)}{\pi_{\phi}^{SFT}(y|x)} } \right] = -\beta \nabla_{\phi} \mathbb{E}_{(x,y) \sim D_{\pi_{\phi}^{RL}}} \left[ \log{ \frac{\pi_{\phi}^{RL}(y|x)}{\pi_{\phi}^{SFT}(y|x)} } \right] \\
&= -\beta \nabla_{\phi} \mathbb{E}_{(x,y) \sim D_{\pi_{\phi}^{RL}}} \left[ \sum_{i=1}^{n} \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] \quad (5.10) \\
&= -\beta \nabla_{\phi} \sum_{i=1}^{n} \mathbb{E}_{(x,y_i) \sim D_{\pi_{\phi}^{RL}}} \left[ \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] \\
&= -\beta \nabla_{\phi} \sum_{i=1}^{n} \sum_{y_i \in \mathcal{Y}} \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i}) \left[ \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] \\
&= -\beta \sum_{i=1}^{n} \sum_{y_i \in \mathcal{Y}} \nabla_{\phi} \left[ \pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i}) \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right] \\
&= -\beta \sum_{i=1}^{n} \mathbb{E}_{(x,y_i) \sim D_{\pi_{\phi}^{RL}}} \left[ \nabla_{\phi} \log{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})} \left( 1 + \log{ \frac{\pi_{\phi}^{RL}(y_i|x,\mathbf{y}_{:i})}{\pi_{\phi}^{SFT}(y_i|x,\mathbf{y}_{:i})} } \right) \right] \quad (3)
\end{align}
$$
上記の$(3)$式がKL penaltyに関する項の勾配に対応する。
一方、「DeepLearningによって方策が計算できるのだからサンプリングを行わずに期待値の勾配を$\displaystyle \sum$で展開した式から計算できるのではないか」という見方も可能であるように一見見える。しかしながら、この計算にあたっては$y$の全パターンについて確率を計算する必要があり、$y$の構成にあたっては組合せ爆発が生じることから現実的ではない。たとえば単語が$30,000$種類かつ$y$の長さが$10$であるだけで$30,000^{10}$通りであり、サンプリングを用いる意義について理解できる。