MLNLP社群是國內外知名的機器學習與自然語言處理社群，受眾覆蓋國內外NLP碩博生、高校老師以及企業研究人員。

社群的願景是促進國內外自然語言處理，機器學習學術界、產業界和廣大愛好者之間的交流和進步，特別是初學者同學們的進步。

來源 | 知乎

作者 | IlikeShuhuaMilk

最近倒騰了一波RLHF，從ColossalAI到TRLX以及DeepSpeed-Chat，最後基於DeepSpeed-Chat成功訓練上了自己的模型，最後效果也是肉眼可見的提升。對這一部分進行下總結，包括原理，程式碼以及踩坑與解決方案。

基本概念

首先還是解釋一下一些概念，從NLP的角度舉一些例子。

首先是RL中的Policy，State，Action。

• Policy π：這是我們需要學習的策略函式，我們使用語言模型來近似這個函式。
• State S：模型接受的輸入句子，這裡的話狀態是無限的，因為輸入可能為任意句子。
• Action A：根據所處的狀態，策略函式做出動作。這裡的話就是模型接受上文，預測下一個字元。這裡的動作是有限集合，大小為vocab size。

接下來介紹Reward，Return，Q，V。

• Reward（獎勵）：在t時刻，給定上文，預測下一個字元，所獲得的獎勵。
• return（回報）：在t時刻，給定上文，持續預測下一個字元，獲得獎勵，直到生成結束符，到達最終狀態，獲得獎勵，所獲得的獎勵和。
• Q（動作的價值）：在t時刻，給定上文，預測下一個字元，可以獲得所有獎勵的期望。。
• V（狀態的價值）：在t時刻，基於給定上文，可以獲得所有獎勵的期望。。這裡就是在上文狀態下，求和所有下一個字元出現機率與字元對應價值的乘積。

PS：這裡要注意區分價值和獎勵：價值是未來累計獎勵的期望。獎勵是我們做出該動作後立即獲取的收益。

RLHF過程

整個過程主要是分為三步：SFT， Training Reward Model，RLHF。

這裡主要介紹一下Training Reward Model和RLHF。

Step2：Training Reward Model

資料

首先是這一部分資料的構成，每一條資料由query，chosen response， rejected response構成，chosen response相較於rejected response質量更高。我使用的資料來自Cohere/miracl-zh-queries-22-12^[1]，裡面對於每條query包含了positive_passages和negative_passages兩個部分，positive部分可以視為chosen，negative部分視為rejected，就可以取樣構建我們訓練Reward Model的資料了：

{"query": "聯合國總部在哪裡？","chosen": "聯合國總部大樓（亦稱聯合國大廈）是聯合國總部的所在地，位於美國紐約市曼哈頓東側，屬於國際領土，因此只要是會員國國民持有護照就可以進入，包括與美國無邦交的聯合國會員國。聯合國總部大樓位於紐約市，其西側邊界為第一大道、南側邊界為東42街、北側邊界為東48街、東側邊界為東河，從聯合國總部大樓可以俯瞰東河。此大樓於1949年和1950年間興建，土地購自於當時的紐約房地產家，面積闊達17英畝（約6.87973公頃）。在此之前，洛克斐勒家族有意提供其在紐約州威斯特徹斯特郡洛克菲勒莊園的土地，但因距離曼哈頓遙遠而作罷；之後，納爾遜·洛克菲勒便協助新的土地的購買，其父小約翰·戴維森·洛克菲勒則捐助了850萬美元協助興建大樓。", "rejected": "聯合國的15個專門機構（如教科文組織）都沒有設在總部。然而，有一些“自治附屬機構”（如聯合國兒童基金會）的總部設在聯合國總部。"}

模型結構

Reward Model相較於原始的SFT Model，在後面加上了一個value head，value head是一個Linear，輸入維度為模型的hidden_dim，輸出維度為1，輸出表示模型預測每一字元獲取的得分。DeepSpeed-Chat中使用最後一個字元的得分作為整個response的得分（當然也可以使用整個句子中每個字元的平均分作為整體的得分）。

訓練目標

訓練Reward Model是一個排序任務，針對query，輸入chosen和rejected response，訓練目標儘可能的使得chosen和rejected的差值更大，損失函式為：

以上就是第二步Training Reward Model的全部過程，基於rank loss訓練了一個打分模型。在第三步強化學習中，reward模型將扮演環境的角色，針對模型預測的字元給出獎勵分數。

Step3：RLHF

整體結構

首先來從整體上看一下這部分（這裡就只介紹RL部分，PTX就是加上了預訓練任務）：

RLHF基於A2C方法，這一步包含了四個模型:

• Actor Model：由SFT之後的模型初始化而來。作為策略（policy）模型，用於接收上文，做出動作，預測下一個字元。學習完畢之後，我們最終使用的就是這個模型。
• Reference Model：和Actor Model同樣初始化自SFT Model，訓練過程中凍結引數，用於和Actor Model做對比，保證模型不要偏離原始SFT Model太多。
• Reward Model：作為環境（env），訓練過程中凍結引數，針對每一個狀態，給出獎勵分數。
• Critic Model：由Reward Model初始化而來，用於近似價值函式，輸入為狀態s，估計當前狀態的價值V。

訓練過程

接下來梳理一遍訓練過程。訓練過程整體分為兩步：maker experience和learn。

首先是make_experience，首先在訓練資料中抽取一部分query，然後Actor Model生成答案。然後我們依據這條答案獲取我們所需要的經驗：

• actor_logits:由Actor Model產生，包含對答案所有詞的機率分佈。
• reference_logits：由Reference Model產生，包含對答案所有詞語的機率分佈，用於和actor logits進行對比，防止actor model偏離SFT Model太遠。
• reward_score: 由Reward Model產生，為當前句子狀態下，立即獲取的收益分數。
• values：由Critic Model產生，估計當前句子狀態下，到完成生成，可以獲取的回報。

整體流程如下：

然後在learn的時候，透過所產生的經驗進行學習。我們透過Actor Model與Critic Model近似策略函式和價值函式，整體流程如下：

關於learn這部分，詳細介紹下Critic Model訓練和Actor Model訓練過程。

這裡對下標new和old做一下說明（最近自己複習一下，看自己的文章，想了半天也沒想起，最後被迫又去看程式碼了，可惡）

問：這裡獲得的actor prob 和之前 make experience中獲得的有什麼區別，為什麼要區分old和new，包括critic也是？

答：因為模型會有dropout 所以會有一些區別。此外如果ppo_epoch大於1，那麼actor產生經驗之後，在forward， backward之後，也會有不同^[2]

Critic Loss

Critic Model估計當前狀態可以獲取的價值，也就是我們前面所說的V值。模型的輸入為狀態s，也就是當前模型生成的句子，輸出為狀態價值 V(s) 。

這裡有一個重要的近似（基於t時刻的狀態s的價值=t時刻獲取的獎勵 +折扣率時刻狀態的價值）：

因此，我們在make experience的時候產生 , 使用其估計Critic Model的目標值。為TD target，是Critic Model在t+1時刻，對做出的估計。

在learn的時候，我們計算當前Critic Model的最新預測值 ,用其估計。由於中包含真實觀測獎勵，因此更為接近真實目標，我們需要將進行對齊，使得Critic Model預測的value更接近真實值。可以使用MSE損失來最佳化Critic Model，這部分損失表示為：

這裡舉一個例子：

早上起床去公司，地圖預計通行時間是50min，然後走了10分鐘到馬連窪，地圖預計用時還需要三十分鐘。那麼一開始預估的總體時間是50min，而基於一部分真實觀測的總體時間是10+30=40min，那麼我們最終的誤差就是(10+30 – 50) **2啦。

Actor Loss

Actor model基於當前輸入上文s，預測下一個字元的機率分佈。

基於策略梯度演算法，我們期望最大化狀態價值，因此目標函式

但是基於這種方法，我們很多時候無論採取何種動作，我們得到的預期回報只要大於0，我們模型就會向這個動作方向更新。

舉個例子，如果我們模型有“上，下，左，右”四個動作，分別有累計獎勵“10，20，30，40”，我們做出任意動作，都會獲取正向的累計獎勵，因此模型也會向這個動作更新引數。而實際上，我們累計獎勵的平均期望為25，對於動作“上，下”我們都應該受到懲罰。

因此我們需要引入一個基線來作為參考，在A2C方法中，引入基線為狀態價值，為狀態s的平均預期回報。因此我們的目標函式變為：

其中為優勢值，

最後，我們在透過PPO演算法來進一步最佳化。關於PPO詳細講解，可以看我的文章PPO詳解。PPO的想法是希望限制訓練過程中對policy進行的更改，避免policy進行過大的更新，來提升policy訓練的穩定性。基於PPO演算法，我們的目標函式變為：

其中。ratio用於比較當前策略和之前舊策略之間的變化，同時將其限制在區間之中。

以上就是第三步的核心內容，RL過程整體分為兩步，make experience和learn。我們首先採樣資料，然後生成結果，Reward Model給出環境的獎勵，以及Critic Model對結果進行評判。之後我們依據所獲取的經驗來對模型進行更新。

DeepSpeed-Chat實踐與踩坑

Step2 Training Reward Model

這個步驟基本就是全部按照DeepSpeed-Chat程式碼來了，使用cohere-zh的資料，大約構造了2w條chosen-rejected pair用於訓練。最後訓練訓練了一個epoch，在驗證集上準確率到了0.79左右。

***** Evaluating reward, Epoch 1/1 *****step: 499/2287, chosen_last_scores (higher is better) : 5.074454307556152,reject_last_scores (lower is better) : 0.5599770545959473, acc (higher is better) : 0.812000036239624step: 999/2287, chosen_last_scores (higher is better) : 5.084388732910156,reject_last_scores (lower is better) : 0.7938708662986755, acc (higher is better) : 0.7940000295639038step: 1499/2287, chosen_last_scores (higher is better) : 5.106724262237549,reject_last_scores (lower is better) : 0.7971451878547668, acc (higher is better) : 0.7986666560173035step: 1999/2287, chosen_last_scores (higher is better) : 5.0183587074279785,reject_last_scores (lower is better) : 0.672178328037262, acc (higher is better) : 0.7955000400543213chosen_last_scores (higher is better) : 5.028912544250488,reject_last_scores (lower is better) : 0.7077188491821289, acc (higher is better) : 0.7936161160469055

Step3 RLHF

踩坑&解決方案

在這個步驟中，從跑通到收斂還是有不少麻煩，分享一些比較重要的點：

訓練過程中發現make experience的時候，model.generate()產生的答案全是重複的無意義字。這裡就很奇怪了，最後發現是開啟了DeepSpeed Hybrid Engine所導致的，查看了issue之後，也發現了有類似的問題^[3]。

不過目前DeepSpeed-Chat也沒有解決，需要關閉Hybrid Engine進行訓練。

DeepSpeed-Chat還有一個很嚴重的問題就是，在make experience的時候，強制Actor Model生成到最大長度（設定max_length=min_length=max_min_length），這樣子導致模型生成偏差很大。對於一個簡單的問題，模型可能本來生成簡單的一句話就可以完美回答了，但是卻必須強制生成到最大長度，這樣訓練的模型和我們實際用起來的模型是有區別的。

對於這個問題，可以透過修改generate中的引數eos_token_id來解決。設定一個虛假的結束符，然後模型就可能生成正常的結束符。然後我們在構建attention_mask來遮蔽掉結束符後面的回答。例如：

seq = [0,0,0,0, prompt, answer, eos_token, other_word]mask = [0,0,0,0,1(prompt),1(answer),1(eos_token),0(other_word)]

透過以上兩步基本就可以跑通流程了，但是訓練過程中還遇到了一個比較大的問題，就是Critic Loss並不收斂，越來越大。

具體的原因是：訓練後期，隨著模型輸出答案越來越好，我們的reward值也會越來越高，導致我們最終累計回報return的區間也會越來越大。而前面說過，我們Critic Loss是一個MSE損失，因此訓練後期，隨著return的估計範圍越來越大，Critic Model就越難估計。在我訓練的過程中，一開始return的範圍是在3-4左右，訓練後期漲到了18-20，因此我們需要想點辦法來約束一下我們的return。

首先的超參的調節，γ引數為折扣回報率，DeepSpeed-Chat中初始設定為1，可以將其調小一些來緩解。其次的話，使用reward scale這一trick幫助也非常大。

透過以上這些步驟，基本上我們正常訓練了，模型最後也能看到一些效果，但是需要取得更好的效果，我們就需要引入一些trick了。

Trick

trick方面主要參考了The 37 Implementation Details of Proximal Policy Optimization^[4]以及影響PPO演算法效能的10個關鍵技巧（附PPO演算法簡潔Pytorch實現）^[5]。對於部分trick，進行了嘗試。

Normalization of Advantages

• 將Advantage進行歸一化：adv=(adv-mean)/std
• 在mini-batch上進行
• 沒有指標的量化，從我個人看結果而言，感覺提升不大

Overall Loss and Entropy Bonus

• 為了提高演算法的探索能力，在actor的loss中增加一項策略熵，並乘以一個係數entropy_coef，使得在最佳化actor_loss的同時，讓策略的熵儘可能大。一般我們設定entropy_coef=0.01
• 總體損失變為：loss = policy_loss – entropy * entropy_coefficient + value_loss * value_coefficient
• 沒有指標的量化，從我個人看結果而言，感覺沒有太多提升

Reward Scale

• 對reward進行縮放，將reward除以標準差
• 從訓練log來看，對穩定critic loss效果很好，畢竟將reward 進行縮放之後，降低了return的估計區間。
• 沒有指標的量化，從我個人看結果而言，提升很大，推薦使用。

關於其他的一些trick，如學習率衰減、梯度裁剪、Value Clipping等，本身框架就包含了，就不進行特別說明了。當然，這些trick在不同的資料或者模型上都會有不同的效果，需要自己進行探索。

透過以上這些方法，最後也順利的完成了訓練，希望對大家能有所幫助，也希望大家能分享自己有用的經驗與理解，關於RL自己瞭解的也很少，還需要更多的學習。

引用連結

[1] Cohere/miracl-zh-queries-22-12: https://huggingface.co/datasets/Cohere/miracl-zh-queries-22-12/viewer/Cohere–miracl-zh-queries-22-12/train?row=0[2] 不同: https://github.com/deepspeedai/DeepSpeedExamples/issues/533[3] 類似的問題: https://github.com/deepspeedai/DeepSpeedExamples/issues/503[4] The 37 Implementation Details of Proximal Policy Optimization: https://iclr-blog-track.github.io/2022/03/25/ppo-implementation-details/[5] 影響PPO演算法效能的10個關鍵技巧（附PPO演算法簡潔Pytorch實現）: https://zhuanlan.zhihu.com/p/512327050