作者:LLM-Finder,某廠研究大語言模型和多模態學習
寫這篇文章的動機
1. 在筆者看來RLHF是LLMs智能的關鍵之一;
2. 國內廠商在這方面投入比較少,目前看起來并沒有很重視;
3. 大家偏向于認為ChatGPT的RLHF做法最多的線索來源于InstructGPT,但是InstructGPT原文的描述也挺含糊的,很多東西只能靠猜和結合開源的實現來解讀;
4. 通常學習強化學習所依賴鏈路比較長,筆者希望以最直觀的方式幫助大家通關。
筆者會分兩篇文章來介紹,第一篇是理論篇,第二篇是實踐篇。讀者會在第一篇學習到PPO的原理和instrcutGPT中的RLHF做法;在第二篇中學習到目前影響比較大的開源RLHF實現。
據公開可獲得的信息來看,ChatGPT需要有大致三個階段的訓練過程,如上圖所示:
1.Pretraining: 在大規模“無監督”的語料上訓練,訓練任務是預測下一個詞。
2.Supervised Fine-Tuning(SFT):在人類標注上進行微調,所謂人類標注就是人類寫Prompt,人類寫答案。然后語言模型學習模仿人類是如何作答的。這部分通常要求數據集多樣性很好,也因為標注成本很高,通常量級很小。
3.Reinforcement Learning with human feedback(RLHF):對于同一個Prompt把模型的多個輸出給人類排序,獲取人類偏好標注。用人類的偏好標注,訓練一個reward model。訓練得到的reward model會作為PPO算法中的reawrd function,來繼續優化SFT得到的模型。
通常來說,第一步最有資源門檻,第三步最有技術門檻(同時也需要大量的資源),第二步最簡單。所以目前很多廠商是直接拿了開源的第一步的模型,做SFT,或者continue-pretrain(比較小規模的無監督訓練)再SFT。他們PR的時候可能會嘴一句,無需復雜的RLHF,只需做細致的微調也能達到很好的效果。
后面兩個步驟,通常被視作是人類偏好對齊(alignment),讓模型更好地跟隨人類的指令作回復。而一些研究發現,對齊后的模型是會有對齊稅的現象的(alignment tax),即在通用能力上會有所下降。
因此,不少人是這樣認為的:第一步預訓練得到的模型就已經決定了后續模型的能力上限;后面兩步要做的事情僅僅是在盡可能減少對齊稅的情況下,對齊人類偏好。
這里可以分兩種情況分析:
? SFT過數據太多遍了,導致大模型出現遺忘;
?安全性對齊很多模型能回答的問題,強制不讓回答肯定會對模型能力有所牽制。
在筆者看來,某種意義下RL提供了對LLM的response的Global-level的監督,在一些需要答案非常精確的場景上,RL可能可以發揮出更大的威力。這個看法的依據也很樸素:
1. 比如在coding、數學推導等場景,只要response在關鍵的地方犯了一點點錯給人的感覺就是模型不會,但是SFT的loss可能區分不出來是犯錯了還是只是寫法風格的差異。
2. SFT給定了標準答案,LLM的上限可能會被標注者的水平所限制;RLHF則只給定了人類偏好,得到了一定(有可能是很大)程度的解放,模型有可能探索出更高程度的智能。這一點并不是無中生有的想法,在游戲AI領域有太多的驗證,即在模仿人類玩法(imitation learning)之后,再用RL訓練出來的模型,就是能獲得更高的智能。這里語言模型跟游戲又有多少本質的區別呢。
InstructGPT中的RLHF
這里簡要帶過具體數據構造和訓練細節,后面會專門有一篇對InstructGPT像素級的解讀。
如前文所述,InstructGPT也是包含3階段的訓練,同時我們應該注意到,RLHF這一步訓練,實則包含兩步訓練:
1. 訓練Reward Model(RM);
2. 用Reward Model和SFT Model構造Reward Function,基于PPO算法來訓練LLM。
數據集
SFT、RM和PPO用到的數據集數據量如下表所示:
注意,上表統計的是prompts數量,在RM數據中每個prompt,對應會有4~9個responses。
在構造RM數據的時候,作者采集了用戶的prompts,每個prompts包含4~9個模型的輸出,模型的輸出會給標注員進行排序。
訓練Reward Model(RM)
目標:給pormpt-response pair打分,擬合人類的偏好。
模型:這InstructGPT的paper中,雖然用了1.3B、6B和175B的GPT-3來做實驗,但是綜合考慮下,只用6B的模型來訓練Reward Model,因為作者發現用175B的模型會不穩定。把最后的unembedding層換成一個輸出為scalar的線性層。這里讀者可能會有點混亂,眾所周知,GPT的模型結構是sequence-in,sequence-out的,怎么變成scalar呢?這里文章似乎也沒提到,根據筆者的判斷和開源實現,推測是直接用最后一個token的輸出接一個linear。
Reward Model的初始化:6B的GPT-3模型在多個公開數據((ARC, BoolQ, CoQA, DROP, MultiNLI, OpenBookQA, QuAC, RACE, and Winogrande)上fintune。不過Paper中提到其實從預訓練模型或者SFT模型開始訓練結果也差不多。
訓練:以前的做法是,RM每次比較兩個模型輸出的好壞,做法很簡單類似對比學習,兩個樣本對應兩個類別,RM對這兩個樣本分別輸出兩個得分,拼成一個logits向量;人類標注比較好的那個輸出作為label,比如第一個比較好那么label為0,第二個比較好label為1;用cross entropy約束之。
但是作者發現這么做很容易過擬合;也不高效,因為每比較一次都要重新過一下reward model。
因此作者的做法是,在一個batch里面,把每個Prompt對應的所有的模型輸出,都過一遍Reward model,并把所有兩兩組合都比較一遍。比如一個Prompt有K個模型輸出,那么模型則只需要處理K個樣本就可以一氣兒做(2K)次比較。loss的設計如下:
很直觀,其中,x是prompt,yw和yl分別是較好和較差的模型response,rθ(x,y)是Reward Model的輸出。σ在文中似乎沒有解釋,不過根據公式推斷和開源實現,應該是sigmod函數。
這里要注意一個細節:在RM訓練完之后,會讓RM的輸出減去一個bias,使得reward score在人類寫的答案上(labeler demonstrations)的平均分為0。這里筆者沒找到具體在什么數據上統計的,猜測是在SFT數據上做的,如果有讀者知道是怎么做的歡迎指出。
Reinforcement Learning(RL)
直接看需要最大化的目標函數
其中,πΦRL和πSFT分別是正在用RL訓練的語言模型和SFT訓練得到的模型。
上式中,
第一項期望式是在最大化reward的同時,最小化和SFT模型的per-token KL penalty,可以理解為是一種正則手段,兩者組合成關于prompt-Responce pair最終的Reward:R(x,y)=rθ(x,y)?βlog(πΦRL(y∣x)/πSFT(y∣x))。per-token KL penalty的好處如下:
1. 充當熵紅利(Entropy bonus),鼓勵policy探索并阻止其坍塌為單一模式。
2. 確保策略模型產生的輸出 與 Reward Model在訓練期間看到的輸出 不會相差太大,保證Reward的可靠性。僅含這一項就是單純使用了PPO。這里也可以看出來,Reward model的能力可能會成為RLHF的瓶頸。
第二項期望式是可選項,注意到它其實是使用預訓練的數據來做跟預訓練同樣的任務(predict next word),因為這一項的數據不是模型生成的其實跟RL是并行的目標。包含這一項的算法稱之為PPO-ptx。
PPO算法
本小節以最小知識補充為前提,快速介紹PPO,不用犯怵,很簡單而直觀。
通常來說,對于一個強化學習模型,會有一個做動作的策略網絡π,它根據自己觀測的狀態(si)做出動作(ai)跟環境交互,然后會拿到一個即刻的reward(ri), 同時進入到下一個狀態(si+1);策略網絡再繼續觀測狀態si+1做下一個動作ai+1...直到達到最終狀態。這樣,策略網絡和環境的一系列互動后最終會得到一個軌跡(trajectory):τ=s1,a1,r1,s2,a2,r2,...,sT,aT,rT。
那么,在語言模型的場景下,策略網絡就是待微調的LLM,它所能做的動作就是預測下一個token,它觀測的轉狀態就是預測下一個token時所能觀測到的context(Prompt+這個token前所生成的所有tokens)。
reward除了最后一個rT等于上文提到的R(x,y)=rθ(x,y)?βlog(πRLΦ(yT∣x,y1,...,yT?1)/πSFT(yT∣x,y1,...,yT?1))
其他的ri=?βlog(πRLΦ(yi∣x,y1,...,yi?1)/πSFT(yi∣x,y1,...,yi?1))。
好,在LLM的場景中,現在可以統一一下符號:s1=x,ai=yi,si=cat(x,y1,y2,...,yi?1),其中x是prompt,yi是第i步蹦的token。看到這,了解PPO的同學基本上就清晰了RLHF具體是怎么做優化的了,可以直接跳過下面的科普部分。
因為PPO原文是基于Actor-Critic算法做的,Actor-Critic算法是進階版的Policy Gradient算法。下面我們從policy gradient到Actor-Critic,再到PPO,幫助RL背景比較弱的讀者串一遍。
Policy Gradient(PG)算法
核心要義:用“Reward”作為權重,最大化策略網絡所做出的動作的概率。
偽代碼核心部分一句話的事:
用策略網絡πθ采樣出一個軌跡,然后根據即刻得到的rewardrt計算 discounted rewardRt=∑i=tTγi?tri;用Rt作為權重,最大化這個軌跡下所采取的動作的概率log(π(at∣st))?Rt,用梯度上升優化之。
雖然在強化學習算法中對每一步都有一個即時的“reward”,但是每一步對后面的可能狀態都是有影響的。
即,后面的動作獲取的即時“reward”都能累計到前面的動作的貢獻。但是直接加上去可能不好,畢竟不是前面的動作直接獲取的reward,但是可以打個折扣再加上去,即乘個小于1的γ。
這里面讀者可能會有個問題:可是不好的動作也要最大化概率嗎?
這里有必要稍微展開一下:
1.Rt也可以是負的,對負的Rt那就是最小化動作at的概率,這也是為什么前面提到要對RM的輸出做歸一化的其中一個原因之一。
2.即便Rt都是正的,但只要充分采樣,同一個狀態下相對的Rt較小的動作也是會被抑制的,因為同一個狀態下的動作概率求和等于1,此消彼長,只有權重最大的動作才會得到獎勵。
可是,比如同一個狀態下,有兩個動作的Rt是正的,但是因為動作采樣本來就很稀疏的,我們很可能不幸運采樣到了相對較小的Rt對應的動作,而沒有采樣到相對較大的。但因為它是正的,這時候當前的機制下,還是會鼓勵這個動作,這樣的話網絡很容易一直沿著不太好的策略去優化。為了解決這個問題,我們引入Actor-Critic算法。
Actor-Critic (AC)算法
核心要義:再增加一個Critic網絡來構造一個Reward baseline,只有獲得的reward比baseline要好才獎勵這個動作,否則抑制它。
Actor指的是策略網絡πθ;Criticb?目的就是給定一個策略網絡,預估每個狀態st,策略網絡所能拿到期望rewardb?(st)是多少。什么是期望reward,無非就是在狀態st,對πθ采樣不同的動作at所能獲取的Rt的平均值嘛。我們要選擇的動作當然是獲取的reward比平均reward要好的動作,不比baseline好的動作就得抑制它。
觀測上面算法2,其實對比PG算法就加了兩行:
1. 原來用Reward function來加權,現在用Advantage function來加權。現在我們把b?(st)當作一個baseline方法所能拿到的reward, 用采樣出來的at所拿到的rewardRt減去b?(st)作為最大化當前動作概率的權重:At=Rt?b?(st)。其中 A_t 通常被稱作是Advantage function(或Advantage estimator),即優勢函數。
2.拉近b?(st)和Rt的距離,初學者對這個可能會費解。實則很好理解,記住b?在做什么,要預估當前策略下Rt的期望,我只要不管三七二十一,每來一個動作的Rt都拉近一下距離,其實就是在預估平均值。更一般地:
其實上面用到的b?,它無非是換了皮的Vπθ?(st)(簡寫成V?(st)),即RL中的重要概念V function:給定策略πθ在st上的期望reward。那么最后一步 T 到達的state sT通常來講是沒有隨機性的(比如下棋,最后一個state決定贏輸就是固定的reward;LLM,最后一個token生成完,response確定了,reward也就確定了),因此rT應該和V?(sT)相等。
所以我們可以重寫上面的優勢函數:
A^t=?V?(st)+rt+γrt+1+?+γT?tV?(sT)
寫成Generalized Advantage Estimation,當λ=1 下式等于上式:
A^t=δt+(γλ)δt+1+?+(γλ)T?t+1δT?1
其中,δt=rt+γV?(st+1)?V?(st)是時序差分式(TD error)。
記住這個結論:這樣我們可以用A^t優化πθ,現在我們可以用▽θlog(πθ(at∣st))?A^t來更新策略網絡了。
PPO Finetuning
上面提到的算法,有一個最嚴重的弊端是,一個軌跡只用一次就丟掉了。可是,采樣軌跡通常是很耗時的,對應到在LLM場景則需要做推理,眾所周知LLM的推理是比訓練費勁很多的,它需要一個一個地蹦詞。可是直接用之前的策略采樣出來的樣本來優化現在的策略網絡肯定不行,如何合理復用樣本則是PPO要做的事情。
做法巨簡單,大致可以用這個思想來更新:
定義 動作概率比rt(θ)=πθold(at∣st)πθ(at∣st),用▽θrt(θ)?A^t去梯度上升更新策略網絡,注意這里stat和A^t都是只之前的策略網絡πθold采樣得到的。這個公式,在筆者看來沒有直觀的解釋,需要一丟丟推導,因為是科普向這里讀者先承認就好了,后面筆者會單開一篇文章再重新梳理一遍。
本質上是最大化這個目標函數:
但是如果πθ和πθold如果差別太大,就不能用這個式子優化了,PPO給出的做法是給rt(θ)卡閾值,太大或太小就不用這一步的樣本更新了:
上面的目標函數可以分類討論進行分析,對優勢函數A^t大于0和小于0兩種情況分析,這個目標函數的圖像長這樣:
觀測圖像:
當A^t大于0,要提高動作的概率,但是如果概率比之前大比較多了(πθ是πθold的1+?倍),就不提高了
當A^t小于0,要減少動作的概率,但是如果概率比之前小比較多了(πθ是πθold的1??倍),就不減少了
偽代碼如下:
科普到此結束,看到這讀者就可以看懂RLHF的代碼。值得注意的是為了減少讀者負擔做了大量的敘述上的簡化,方法上是比較完備的,但是說法上不夠嚴謹。Again,更詳細的強化學習科普會單開一篇文章。
大語言模型的PPO
稍微整理一下,符號和上面的科普部分不一致,不過應該不影響理解
1.現在我們的actor是SFT初始化的LLMπΦRL;
2.為了計算reward,我們需要兩個凍住參數網絡,一個RM,一個是凍住的SFT模型πSFT用來計算KL散度,參考下面兩式子:rT=R(x,y)=rθ(x,y)?βlog(πRLΦ(yT∣x,y1,...,yT?1)/πSFT(yT∣x,y1,...,yT?1))其他步的ri=?βlog(πRLΦ(yi∣x,y1,...,yi?1)/πSFT(yi∣x,y1,...,yi?1));
3.為了執行PPO算法,我們需要引入一個估計V值的網絡Vη,它初始化來自RM。所以統共,有4個網絡,兩個訓練的actorπΦRL和criticVη;兩個用來計算reward的SFT模型πSFT和RM模型。然后actor初始化來自SFT,critic初始化來自RM。
把這四個網絡,結合reward的構造,帶入到上面提到的PPO算法中,整個過程就比較清晰了。
盜一下DeepSpeed-Chat的圖,圖解如下:
看到這,相信讀者已經可以輕易看懂的DeepSpeed-Chat代碼了。??
審核編輯:黃飛
-
ChatGPT
+關注
關注
29文章
1568瀏覽量
8060
原文標題:PPO算法
文章出處:【微信號:zenRRan,微信公眾號:深度學習自然語言處理】歡迎添加關注!文章轉載請注明出處。
發布評論請先 登錄
相關推薦
評論