PPO 小註解

Proximal Policy Optimization 是一個 fine-tune Model 的方式, 另一個主流方式是 DPO (Direct Preference Optimization). 這邊主要是整理課程影片中 PPO 的公式 (只用大約 2 分鐘超高速閃過…), 以便有空時可以查詢.

[基本架構]

有一個 LLM model, 和可調整的參數: θ, query: X, response: Y, 所有輸出入組合 rollout 表示為 (X, Y).

因為我們要微調這個 model, 所以要給一個獎勵函數 (reward function): r(X,Y), 用來調整模型 πθ (X,Y).

獎勵是對所有輸出的 Y 取期望值

EY~πθ (X,Y) [r(X,Y)], 其中 Y~πθ (X,Y), 表示 Y 在 πθ (X,Y) 的輸出組合.

也對每個 query X 計算

EX~D[EY~πθ (X,Y) [r(X,Y)]] ——————— (1)

其中 X~D, D 表示 X 的散度 (Divergence), 依此類推.

我們找出一個得到最多獎勵的那組參數, 就算是 fine-tune 完成.

π* (X,Y) = arg maxπ {EX~D[EY~πθ (X,Y) [r(X,Y)]]}, {} 裡面就是上面那包.

[Reference model]

做得好不好, 我們要有一個 reference model 來當評審.

πref (X,Y)

Model 的輸出愈接近 reference model 就表示學得愈好.

[懲罰值]

KL 是 Kullback-Leibler 兩個人名的縮寫. KL penality coefficient 是它們衡量懲罰的函數. 基於 reference model πref (X,Y) 下, 對 πθ (X,Y) 產生的懲罰 DKL 表示為

DKLθ (X,Y) || πref (X,Y)]

[綜合期望值和懲罰值]

π* (X,Y) = arg maxπ {EX~D[EY~πθ (X,Y) [r(X,Y)]] – β * 懲罰值} =>

π* (X,Y) = arg maxπ {EX~D[EY~πθ (X,Y) [r(X,Y)]] – β DKLθ (X,Y) || πref (X,Y)]}


為了要知道往哪個方向調整, 我們要找 θ 的梯度.

[只看 reward function]

這是 reward function 在特定 model 參數 θ 的期望值

E[rY ] = EY~πθ (X,Y) [r(X,Y)]

將它表示為解析函數 [1], 期望值等於所有 Y 可能的值乘上它的機率的和

E[rY ] = ∑Y r(X,Y)πθ (Y|X) —— (2)

最優化的參數 θ⌃ 就是所有 θ 中 reward 最好的. (我打不出 hat theta, 所以把帽子放右邊)

θ⌃ = arg maxθ [E[rY ]] = arg maxθ [∑Y r(X,Y)πθ (Y|X)]

對式子 (2) 取 θ 的梯度,

θ E[rY ] = ∑Y r(X,Y)θ πθ (Y|X)] ——– (3)

按照微積分的 chain rule [3]

θ log (πθ (Y|X)) = θπθ (Y|X) / πθ (Y|X) =>

θπθ (Y|X) = θ log (πθ (Y|X)) * πθ (Y|X) —– (4)

把這個變長的 (4) 帶回去 (3),

θ E[rY ] = ∑Y r(X,Y)θ πθ (Y|X)] =>

θ E[rY ] = ∑Y r(X,Y)θ log (πθ (Y|X)) * πθ (Y|X) =>

θ E[rY ] = EY~πθ (Y|X) r(X,Y)θ log (πθ (Y|X))

再把 X 考慮進去, 就得到

Ex~D [θ E[rY ] ] = θ Ex~D [ E[rY ] ]

意思是說, 當 β = 0, 忽略掉懲罰的話, 對式子 (1) 取梯度, 等效只對 reward function 取梯度. 反之亦然.

這有什麼意義呢? 我請 AI 幫我解釋. 我試了幾次都沒有它講得好.

θ E[rY ] = EY~πθ (Y|X) r(X,Y)θ log (πθ (Y|X)) 這個結果是 策略梯度定理(Policy Gradient Theorem)[2] 的核心公式。其意義在於:

  • 將梯度轉換為期望形式:原始梯度需遍歷所有可能的輸出 θY(計算量巨大),但此公式表明:梯度可通過從策略 πθ   中採樣 Y 來近似計算。
  • 避開解析計算:無需知道所有 Y 的機率分佈,只需對採樣的 Y 計算  r(X,Y)θ log (πθ (Y|X)) 的平均值。

它與 Monte Carlo Method 的相似處在於: 透過採樣直接估計期望值,無需精確計算所有可能的 Y.

蒙特卡羅方法收斂速度為 O(1/√N​), N 夠大就好. 而窮舉法不可行.

最後再筆記一下懲罰項.

Kullback-Leibler 方法其實是要避免訓練後新的參數和舊的參數偏差太遠. 以至於學會新的東西就忘了舊的 – 災難性遺忘 (Catastrophic Forgetting), 和 reward 無關. 因此討論原理時可以忽略它.

但它對於穩定系統非常重要, 其主要角色是作為正則化項 (Regularizer), 將策略更新限制在信賴域 (Trust Region) 內. 在PPO算法中,KL 散度還會動態調整更新步長. 當 DKL​ 超過閾值時自動縮小學習率,此設計本質上將 “避免遺忘" 和 “促進學習" 綁定為同一過程。

[REF]

  1. 解析函數 – 維基百科,自由的百科全書
  2. https://en.wikipedia.org/wiki/Policy_gradient_method
  3. Slide 1

關於 LLM 的自我審查

先前用過 ChatGPT 3 的人應該都記得這些 Model 動不動就說自己不知道, 不能說, … 但變個花樣問, 不能說的又全說了!

那麼 LLM 怎麼做到自我審查呢? 基本的做法是 RLHF (Reinforcement Learning from Human Feedback) [1], 也就是靠人類的意見來約束 LLM. 但我們不知道 LLM 會被問什麼問題, 總不能一筆一筆叫網路警察來核准吧! 所以我們先根據真人的反饋, 來建立一個獎勵模型 (reward model), 以後就叫 reward model 來代替人工.

當大量的問題, 都有被人類標記為優選的答案, 他們就是一個標準答案 reward model. 它其實也是一個 LLM, 但我們先忘記這件事. 下面講到 LLM 都是指那未經自我審查 / fine tune 的 LLM.

顯然地, 我們拿 LLM 的標準問題 prompt 以及輸出 completion對, 去和 reward model 的標準 prompt + completion 比較, 計算兩者 logits 層 (未經 softmax 之類 activation function 的原始輸出) 的差異, 就知道這個 LLM 回答得像不像人? 然後把 LLM 的 logits 往人類標記的方向調整 [註 A], 就會讓 LLM 的答案愈來愈像人的回答.

我們也可以叫 reward model 標記 LLM 的輸出是否有害? 例如做個 remodel model 專門學 LLM 的輸出是否拉高仇恨值? LLM 答題主要依據 helpful, honest, 和 harmless 三原則. 身為一個機器人, 誠實是基本的. 不能出現言語傷害 (harmless 原則), 也不能因為怕講錯話而淨說沒用的廢話 (helpful 原則) [註 B].

本圖取材自 https://primo.ai/index.php?title=Reinforcement_Learning_(RL)from_Human_Feedback(RLHF

上面的論述略過了好幾段段情節, 在此補充:

[註 A] Reward model 怎麼微調 LLM 參數. 目前常用的一個演算法是 PPO (proximal policy optimization). 我們將它理解為可以調整參數, 但一次不要調太多, 免得把辛苦 train 了半天的 model 調壞. 因此就算要調整, 也是利用 PEFT (Parameter-Efficient Fine-Tuning) 的做法, 包括有名的 Lora.

[註 B] 為避免 Reward model 把 LLM 帶偏, 變成只求不出錯就好, 我們同樣限制調整過的參數和原本的參數 (reference model or frozen model) 不能差太多. 可用的演算法包括 Kullback-Leibler (KL) divergence, 它可以用 softmax 的輸出來比較像不像, 所以根本不用管輸出的 token 是啥以減少計算量. 它可以跟 reward model 共存.

最後, 不免有人會問, 如果一定要有 Human 才能幫 LLM 做思想審查, 是不是太不 AI 了. 沒錯! 其實 Human 也可以換成 AI. 但我們不會叫小學生去教小學生, 我們先叫 LLM 產生一堆負面教材, 然後 train LLM 避免生成這些仇恨、色情、暴力的言論即可. 於是乎, 當我們問 LLM, 不可以上那些色情網站時, 這些剛好都是 RLAIF (Reinforcement Learning from AI Feedback) 的紅隊演練考題, 因此他們侃侃而談, 不知道中計! 當然, 這已經是過去式了~~

[REF]

  1. https://huggingface.co/blog/rlhf
  2. https://huggingface.co/docs/peft/index