怎樣讓PPO訓練更穩(wěn)定？早期人類征服RLHF的馴化經(jīng)驗

做一個專門面向年輕NLPer的每周在線論文分享平臺

寫在前面

今天給大家?guī)硪黄?strong>“如何穩(wěn)定且有效地訓練 PPO”的論文解讀，來自知乎@何枝（已授權）。在這篇文章中我們將學習：哪些技巧能夠穩(wěn)定訓練過程、哪些指標能夠代表著訓練的順利進行等內(nèi)容。

作為 Reinforcement Learning 中的頂流算法，PPO 已經(jīng)統(tǒng)領這個領域多年。直到InstructGPT的爆火，PPO 開始進軍 LLM 領域，憑借其 label-free 的特性不斷拔高基座的性能，在 Llama2 、 Baichuan 的工作中都能看到 RLHF 的身影。

于是你開始摩拳擦掌，躍躍欲試，準備利用這項強大的技術來進化自己的基座。但當你信心滿滿地跑通訓練任務時，你看到的情況很有可能是這樣的：

各種形形色色的失敗案例

盡管魯迅先生曾言：真的強化敢于直面慘淡的結果，敢于正視崩壞的曲線。但日復一復地開盲盒難免會讓人心臟承受不了，好在前人們留下了寶貴的馴化經(jīng)驗，今天讓我們一起看看“如何穩(wěn)定且有效地訓練PPO”。

知乎：https://zhuanlan.zhihu.com/p/666455333
Paper：https://arxiv.org/pdf/2307.04964.pdf
Code：https://github.com/OpenLMLab/MOSS-RLHF/tree/main

1. Reward Model 訓練

RL 的整個訓練目標都是圍繞著 reward 來進行，傳統(tǒng) RM 的訓練公式為拉開好/壞樣本之間的得分差：

y_w 為 selected 樣本，y_l 為 rejecte 樣本

但是，僅僅是「拉開得分差」這一個目標很有可能讓 RM 陷入到「鉆牛角尖」的困境中。

為了保持住 RM 的本質還是一個「語言模型」，文章在原本的 loss 中加入了對「好樣本」的 LM loss：

在原來的 loss 基礎上順便學習寫出「優(yōu)秀樣本」，保持住模型能寫句子的能力

值得一提的是：文章中的 r' 是用了另外一個 RM' 來算 loss 的，RM' 的結構和 RM 一樣，只不過輸出的維度不是 1，而是 vocab_size。但其實我認為也可以使用一個帶有 ValueHead 的模型來既訓練打分又訓練寫句子，畢竟這 2 個任務都需要模型知道什么的句子是「好句子」—— 還能省顯存。

以下是論文訓練 RM 的詳細參數(shù)：

一般的，我們會使用 prefered sample - disprefered sample 的分差來衡量 RM 的效果：

圖左為在中文標注數(shù)據(jù)集上的分差分布，圖右為在英文數(shù)據(jù)集上的分差分布（區(qū)分度不如中文）

完全理想的狀況下，prefered - disprefered 應該都在 0 的右邊（好樣本的分數(shù)更高），但考慮到標注中的噪聲、模型的擬合能力等，存在少小部分 0 左邊的樣本是合理的，拉出來人工評估下即可。

此外，文中還提到：只看 Acc 并不能夠很好的衡量 RM 的性能，但尚未給出其他可以衡量的指標。

2. PPO 的穩(wěn)定訓練方法

2.1 及時發(fā)現(xiàn)訓練過程中的異常

PPO 訓練中很常見的一個問題是「模式崩潰」，其典型特征為：長度很長且無意義的文字。

而對于這種「崩潰的輸出」Reward Model 往往還容易打出一個很高的分數(shù)，這將導致我們無法在訓練過程中及時的發(fā)現(xiàn)問題，等訓完對著一個滿意分數(shù)的 checkpoint 看生成結果的時候才發(fā)現(xiàn)空歡喜一場。

對于上述這種問題，我們可以通過 3 個指標來監(jiān)控：KL、Response Length、Perplexity。

訓練過程中的各種指標，從約第 420 step 開始: 1. reward 出現(xiàn)驟增。2. KL 出現(xiàn)驟增。3. Perplexity 出現(xiàn)驟降。4. Response Length 出現(xiàn)驟增。5. 訓練效果出現(xiàn)驟降（圖左上紅線）

因此我們可以總結出幾種指標異常的情況：

• Reward 出現(xiàn)驟增：很可能 Policy Model 找到了某條 shortcut，比如通過模式崩潰來獲得更高的分數(shù)。
• KL 出現(xiàn)驟增：同上，很可能此時的輸出模式已經(jīng)完全崩潰。
• Perplexity 驟降：由于 PPL 是指代模式對當前生成結果的「確定性」，一般來講，句子的生成都會帶有一定的不確定性，當 Policy Model 對某一個生成結果突然「非常確定」的時候（無論是什么樣的上文都很確定接下來應該輸出什么），那么它大概率是已經(jīng)擬合到了一個確定的「崩壞模式」上了。
• Response Length 驟增：這個對應我們之前給的 bad case，回復長度的驟增也可能代表當前輸出已經(jīng)崩潰。

2.2 Score Normalization & Clipping

PPO 的整個訓練都是圍繞優(yōu)化 Score 作為目標來進行的，和 Score 相關的變量有 2 個：

• Reward：由 RM（≈ Human） 直接給出的反饋。
• Advantages：由 Reward 和 Critic Model 共同決定的優(yōu)勢值，最終用于 loss 計算。

對于這 2 個值，我們都可以對其進行「歸一化」和「裁剪」。

Reward 的處理公式如下：

Reward Normalization & Clipping

上述式子將 reward 化成了一個均值為 0 的標準分布，均值為 0 是為了保證在訓練過程中得到的正負獎勵能夠盡可能的均勻，如果一段時間內(nèi)全為負或全為正從直覺上來講不太利于模型學習。

文中提到，使用 clipping 可以限制模型進化的「最終分數(shù)」沒有那么高，鑒于之前「分數(shù)越高，并不一定有更好的效果」的結論，作者認為使用 clipping 對最終的效果是有益的。

至于 Advantages，在 PPO 的標準流程里已經(jīng)會對其進行 Normalization，而 advantage clipping 和 reward clipping 在本質上其實很相似，則只用在 reward 階段進行截斷即可，所以對于 Advantage 來講不需要做太多其他額外的操作。

2.3 Policy Loss 設計

在傳統(tǒng)的 PPO 流程中，我們通常會對 Policy Molde 的 Loss 上做以下 2 種操作：

• Importance Sampling：這是 PPO 中加快 On-Policy 算法訓練速度的關鍵步驟，即一次采樣的數(shù)據(jù)可以進行多次更新（通過系數(shù)補償）。這種方法通常和 KL 懲罰一起使用，實驗表明這樣能夠更加穩(wěn)定 PPO 的訓練，但對最終的效果會存在一定折損（所以最好的還是 1 輪 sample 只做一次 update，退化為原始的 PG 流程）。
• Entropy Loss：一般為了鼓勵 Policy 在進化的同時保留「探索」的能力，我們會在 loss 中加入 entropy（確定性）loss，但在 RLHF 中這項設置對超參非常敏感，很容易就崩掉。鑒于 KL 和 Entropy 有著相似的效果，因此作者更推薦使用 KL 來代替 Entropy Loss。

除了上述 2 個傳統(tǒng)設置外，RLHF 中加入一個新的指標：Token Level KL-Penalty。

在傳統(tǒng)的 RL 流程中，agent 每采取一個 action 后都會得到一個 action reward，對比到文本生成任務中，每新生成一個 token 就等于做出了一次 action，但實際上我們無法每生成一個 token 就打出一個分數(shù)，我們只能在一個完整句子（Trajectory）生成完成之后打出一個 Total Reward。

這就比較痛苦了，當我們只有一個長序列的最后得分時，前面的每一個 step 的得分估計就變得非常困難。因此，為了避免「sparse reward」的同時限制 Policy Model 朝著「相對合理的方向」進化，我們會通過計算每個生成 token 與參考模型之間的 KL 來作為單個 token 的 reward 分數(shù)。

式子的前半部分為 reward（discounted）后半部分為 KL 懲罰分數(shù)

2.4 模型初始化

PPO 繼承自 Actor-Critic 框架，因此算法中一共包含 2 個模型：Actor 和 Critic。

• Actor Model（Policy Model）

Policy Model 是指我們最終訓練后使用的生成模型，Policy Model 需要具備一定基本的能力才能保證訓練的穩(wěn)定性，通常會選用 SFT 之后的模型。這個比較好理解，如果我們選用 Pretrained Model 為起點的話，探索空間會非常大，同時也更加的不穩(wěn)定（對 Reward Model 要求更高）。

• Critic Model

一種很直覺的想法是：同樣是「評判任務」，我們直接使用 Reward Model 來當作 Critic Model 就好了。

但其實這種想法不完全正確，從本質上來講 Critic 需要對每一個 token 的狀態(tài)進行打分，而 RM 是對整個句子進行綜合得分評估，這兩個任務還是存在一定的區(qū)別。

因此，一種更好的方式是：先訓練 Critic Model一段時間，直到 Critic Loss 降的相對較低為止。預先訓練能夠幫助在正式訓練的初期 Critic 能夠進行較為正確的 value 預估，從而穩(wěn)定訓練過程，至于使用 SFT 還是 RM 作為 Critic 的結構，實驗結果顯示并沒有非常明顯的區(qū)別。

2.5 最優(yōu)策略集合（PPO-max）

文章的末尾給出了作者匯聚了各種實驗結果給出的一套推薦的策略：

• reward normalize：使用歷史獲得過的所有 reward 的均值和方差進行標準化。
• token KL penalty：限制模型更新方向。
• Critic Model：使用 RM 初始化 Critic，并在 PPO 正式訓練之前先進行 Critic 預訓練。
• Global Gradient Clipping
• 使用相對較小的 Experience Buffer。
• Pretrain Loss：在 PPO 訓練 loss 中加入 Pretrain Language Model Loss，和 [InstructGPT] 中保持一致。

PPO-max 所使用的方法合集（標星的方法）

---