選考チューニングのRLHFを数式ベースで理解する（強化学習編）

大規模言語モデル入門および大規模言語モデル入門II出てくるRLHFについて、少し数式で詰まった部分があったので、メモがてら数式の意味を記載していこうと思います。
章番号や数式の番号は書籍のものなので、下記書籍を持ってないとなんのこっちゃだと思います。

https://gihyo.jp/book/2023/978-4-297-13633-8

https://gihyo.jp/book/2024/978-4-297-14393-0

対象は、4.5.2項および12.1.1項です。

強化学習の数式

eq(4.4)

$$P(y∣x,ϕ)=∏i=1Nπϕ(wi∣x,w<i) P(y|x, \phi) = \prod_{i=1}^{N} \pi_\phi(w_i|x, w_{<i})$$

この式は、プロンプト$xx$とパラメータ$ϕ\phi$において、テキスト$y=w1,w2,...,wNy = w_1, w_2, ..., w_N$が生成される確率を表している。
$πϕ(wi∣x,w<i)\pi_\phi(w_i|x, w_{<i})$は、プロンプト$xx$と生成途中のテキスト$w<i=W1,w2,...,wi−1w_{<i} = W_1, w_2, ..., w_{i-1}$がLLMに与えられたときに、$wiw_i$が出力される確率となる(softmaxを適用した、語彙に対する確率分布のうちの、$wiw_i$の確率)。

総乗記号$∏\prod$の慣れてない方のために、書き下しておきます。

$$πϕ(w1∣x)×πϕ(w2∣x,w1)×πϕ(w3∣x,w1,w2)×...×πϕ(wN∣x,w1,w2,...,wN−1) \pi_\phi(w_1|x) \times \pi_\phi(w_2|x, w_1) \times \pi_\phi(w_3|x, w_1, w_2) \times ... \times \pi_\phi(w_N|x, w_1, w_2, ..., w_{N-1})$$

eq(4.5)

$$ϕ^=argmaxϕEx∼DrlEy∼P(y∣x,ϕ)[R(x,y)] \hat{\phi} = \text{argmax}_{\phi} \mathbb{E}_{x \sim D_\text{rl}} \mathbb{E}_{y \sim P(y|x,\phi)} [R(x, y)]$$

報酬$R(x,y)R(x, y)$を最大化する方策のパラメータ$ϕ^\hat{\phi}$を求める式です。

• $ϕ^\hat{\phi}$: 最適なパラメータ値
• $argmaxϕ\text{argmax}_{\phi}$: $Ex∼DrEy∼P(y∣x,ϕ)[R(x,y)]\mathbb{E}_{x \sim D_r} \mathbb{E}_{y \sim P(y|x,\phi)} [R(x, y)]$を最大にするパラメータ$ϕ\phi$を求める、ということ
• $Ex∼Drl\mathbb{E}_{x \sim D_\text{rl}}$: データセット$DrlD_\text{rl}$からサンプリングされた$xx$に関する期待値
• $Ey∼P(y∣x,ϕ)\mathbb{E}_{y \sim P(y|x,\phi)}$: テキスト$yy$が生成される確率$P(y∣x,ϕ)P(y|x,\phi)$に関する期待値
• $R(x,y)R(x, y)$：eq(4.6)で説明されてる報酬

期待値が2つあるため混乱するかもしれませんが、これは具体的には：

1. 外側の期待値$Ex∼Drl\mathbb{E}_{x \sim D_\text{rl}}$は、データ分布からサンプリングされた様々な入力$xx$に対する平均を計算します
2. 内側の期待値$Ey∼P(y∣x,ϕ)\mathbb{E}_{y \sim P(y|x,\phi)}$は、特定の入力$xx$とパラメータ$ϕ\phi$が与えられたときに、モデルが生成する可能性のある様々な出力$yy$に対する平均報酬を計算

を表します。
つまり、この式は「データ分布から得られる様々な入力に対して、モデルが生成する出力の期待報酬を最大化するようなパラメータ$ϕ\phi$を見つける」という最適化問題を表しています。

eq(4.6)

$$R(x,y)=rθ(x,y)−βlog⁡P(y∣x,ϕ)P(y∣x,ϕinst) R(x, y) = r_{\theta}(x, y) - \beta \log \frac{P(y|x, \phi)}{P(y|x, \phi_{\text{inst}})}$$

上記にも出てきた$R(x,y)R(x, y)$の式です。
第1項の$rθ(x,y)r_{\theta}(x, y)$は報酬モデルの出力(報酬)です。
第2項の分母$P(y∣x,ϕinst)P(y|x, \phi_{\text{inst}})$は、選考チューニングをする前のLLMがテキスト$yy$を出力する確率となっており、パラメータ$ϕinst\phi_{\text{inst}}$は学習させずに保存しておきます。このモデルは、参照モデルとよばれます。
第2項の分子$P(y∣x,ϕ)P(y|x, \phi)$は、今回の選考チューニングで学習対象のLLMがテキスト$yy$を出力する確率です。このモデルは方策モデルと呼ばれます。
この第2項は正則化項です。

さて、この報酬$R(x,y)R(x, y)$は、どんな時に高くなり、どんな時に低くなるのでしょうか。

第1項の$rθ(x,y)r_{\theta}(x, y)$はそのまま存在するので、報酬モデルが出力する報酬$rθ(x,y)r_{\theta}(x, y)$が高ければ全体の報酬$R(x,y)R(x, y)$が高くなり、報酬$rθ(x,y)r_{\theta}(x, y)$が低ければ$R(x,y)R(x, y)$は低くなります(当たり前)。

第2項は少しややこしいですが、$β>0\beta > 0$であることを考慮するとこんな風に解釈できるのではないでしょうか？こういう時は場合分けするとわかりやすいですね。

まずは$P(y∣x,ϕ)>P(y∣x,ϕinst)P(y|x, \phi) > P(y|x, \phi_{\text{inst}})$の時、つまり方策モデルが$yy$を出力する確率が、参照モデルが$yy$を出力する確率より大きい場合です。
$log⁡(⋅)\log(\cdot)$が正の値をとるので、全体報酬$R(x,y)R(x, y)$は低くなります。
これは、参照モデルが出力しにくいテキスト$yy$を無理やり出力した場合に報酬を低くすることで、指示チューニングで学習した内容を忘れることを防ぎたい、というわけです。

その逆$P(y∣x,ϕ)<P(y∣x,ϕinst)P(y|x, \phi) < P(y|x, \phi_{\text{inst}})$だと全体報酬$R(x,y)R(x, y)$は高くなり、両方の確率が等しい場合だと正則化項の寄与は0となります。

eq(4.7)、eq(4.8)およびeq(12.3)は、既に説明済み事項を別の書きかたにしただけなので省略します。

eq(4.6)からわかるように、RLHFの学習には3つのモデルを使用します。そのため多量の計算リソースを必要となり、時間もかかります。
そのため、改善された手法としてDPO(Direct Preference Optimization)があります。

eq(4.9)およびeq(4.10)

$$ϕ^=argmaxϕ(Ex∼DrlEy∼P(y∣x,ϕ)[R(x,y)]+γEx∼Dpt[log⁡P(x∣ϕ)]) \hat{\phi} = \text{argmax}_{\phi} \left( \mathbb{E}_{x \sim D_\text{rl}} \mathbb{E}_{y \sim P(y|x,\phi)} [R(x, y)] + \gamma \mathbb{E}_{x \sim D_\text{pt}} [ \log P(x|\phi) ] \right)$$

ここでは、使われてるデータセットの違いに注意しましょう。
$DrlD_\text{rl}$は、選考チューニングで使うデータセットです。また、$DptD_\text{pt}$は、事前学習に使用したデータセットと同様です。

第1項はeq(4.5)と同じです。

第2項は、

$$Ex∼Dpt[log⁡P(x∣ϕ)]=1∣Dpt∣∑x∑ilog⁡πϕ(ui∣ui−K,…,ui−1) \mathbb{E}_{x \sim D_{\text{pt}}} [\log P(x|\phi)] = \frac{1}{|D_{\text{pt}}|} \sum_{x}\sum_{i} \log \pi_{\phi}(u_i|u_{i-K}, \ldots, u_{i-1})$$

となっていますね。
$Ex∼Dpt[⋅]\mathbb{E}_{x \sim D_{\text{pt}}} [ \cdot ]$の部分は、$xx$について期待値を取るので右辺の$1∣Dpt∣∑x\frac{1}{|D_{\text{pt}}|} \sum_{x}$に対応しています。
残りの$∑ilog⁡πϕ(ui∣ui−K,…,ui−1)\sum_{i} \log \pi_{\phi}(u_i|u_{i-K}, \ldots, u_{i-1})$はeq(3.2: P35)と文字が変わってるだけで同じです。
$ui−K,…,ui−1u_{i-K}, \ldots, u_{i-1}$となっているのは、モデルの入力トークンの最大長が$KK$となっているので、出力トークン$uiu_i$からさかのぼって$KK$個分しかLLMに入力できないためです。