【LLM】在線合并優(yōu)化器以提升獎勵并減輕對齊開銷

sbf_2000

發(fā)布于 2024-6-7 12:59

瀏覽

0收藏

【LLM】在線合并優(yōu)化器以提升獎勵并減輕對齊開銷-AI.x社區(qū)

一、結(jié)論寫在前面

在強化學(xué)習(xí)人類反饋（RLHF）中，有效對齊大型語言模型（LLMs）與以人為中心的價值，同時防止通過預(yù)訓(xùn)練和監(jiān)督微調(diào)（SFT）獲得的能力退化，是一個核心挑戰(zhàn)。插值RLIF和SFT模型參數(shù)可以調(diào)整人類偏好與基本能力之間的權(quán)衡，從而以犧牲對齊獎勵為代價減少對齊開銷(Alignment Tax)。

受此啟發(fā)，論文提出在RLHF的每個優(yōu)化步驟中整合RL策略和SFT模型，以持續(xù)調(diào)節(jié)訓(xùn)練方向，引入在線合并優(yōu)化器。具體而言，論文通過合并SFT與預(yù)訓(xùn)練模型之間的參數(shù)差異來融合梯度，有效地引導(dǎo)梯度向SFT優(yōu)化的方向最大化獎勵。

對各種骨干LLMs進行的大量實驗表明，與正則化和離線合并基線相比，在線合并優(yōu)化器能更好地緩解對齊成本并實現(xiàn)更優(yōu)的對齊性能。此外，論文提出了step-K在線合并優(yōu)化器，彌合了在線和離線合并之間的差距，并深入分析了超參數(shù)和消融的效果。論文證明了論文的優(yōu)化器適用于不同的LLM家族，如Qwen和LLaMA，跨越從1.8B到8B的各種模型大小，以及DPO和KTO等不同的RLHF算法，以及現(xiàn)有的模型合并方法。它顯著提升了對齊獎勵，同時減輕了對齊成本開銷，在14個基準(zhǔn)測試中實現(xiàn)了更高的整體性能。

限制：在線合并優(yōu)化器的主要限制與參數(shù)效率相關(guān)。在線合并優(yōu)化器增加了內(nèi)存需求，因為它們需要緩存參考模型的額外增量參數(shù)，以對應(yīng)訓(xùn)練中的增量更新權(quán)重。同時，除非參考模型也使用LoRA適配器進行訓(xùn)練，否則它們不能應(yīng)用于LoRA訓(xùn)練。然而，通過將GaLore與在線合并優(yōu)化器結(jié)合使用，可以進一步消除這一限制。

二、論文的簡單介紹

2.1 論文的背景

理想情況下，一個最佳的RLHF策略應(yīng)在保持對齊性的獎勵同時避免相關(guān)的代價，力求在最大化獎勵的同時最小化遺忘。依賴于神經(jīng)網(wǎng)絡(luò)的線性模式連通性，模型能力之間的權(quán)衡可簡潔地描述為模型參數(shù)的插值。研究表明，通過權(quán)重插值從同一預(yù)訓(xùn)練模型微調(diào)的不同模型結(jié)合，往往能在原始模型間實現(xiàn)更平衡的性能。

受此啟發(fā)，論文初步探索了將RLHF模型與其訓(xùn)練自的參考SFT模型進行合并。論文的觀察表明，這種離線模型合并有效地緩解了對齊成本。如表4所示，離線合并模型在語言基準(zhǔn)和語言熟練度上恢復(fù)了與SFT模型相當(dāng)?shù)男阅堋Ｈ欢@種改進是以相對于RLHF模型偏好評分的減少為代價的。

鑒于單次參數(shù)插值僅允許在固定能力的模型之間進行權(quán)衡，離線合并帶來的適度性能提升并不令人意外。在RLHF訓(xùn)練過程中，每個優(yōu)化步驟都提升了模型的能力。因此，論文有機會確保這些變化的方向與參考SFT模型保持一致。本文中，論文將模型合并整合到每個RLIIF優(yōu)化步驟中，并引入了在線合并優(yōu)化器。這一創(chuàng)新的優(yōu)化器在提升獎勵方面比傳統(tǒng)的優(yōu)化器如AdamW更有效，同時也在減少對齊成本方面，類似于離線合并。

【LLM】在線合并優(yōu)化器以提升獎勵并減輕對齊開銷-AI.x社區(qū)

圖1：RLHF與在線合并優(yōu)化器的示意圖。在每個RLHF迭代中，論文首先獲取更新權(quán)重A((0，然后對其進行稀疏化處理，并與參考模型的delta參數(shù)達成共識。論文使用這種合并的delta作為本次迭代中策略模型的更新。論文還對比了在線合并與離線合并，如圖下方所示

對齊開銷(Alignment Tax)。通常，使大型語言模型（LLMs）與人類偏好對齊涉及兩個階段：首先進行監(jiān)督微調(diào)（SFT）以建立遵循指令的模型，隨后通過人類反饋強化學(xué)習(xí)（RLHF）來增強人類偏好。當(dāng)前主流的RLHF方法，如PPO和DPO，指導(dǎo)模型在優(yōu)化獎勵的同時，引入Kullback-Leibler（KL）散度懲罰項，該懲罰項存在于學(xué)習(xí)到的RL策略輸出與參考的SFT模型之間。這一懲罰機制防止策略在追求偏好獎勵時偏離其原始目標(biāo)，即保持已獲得的語言能力。

論文在DPO設(shè)置中嘗試了不同的KL散度權(quán)重β，如圖4所示。β的增加與基準(zhǔn)性能的平均提升相關(guān)，但代價是MT-Bench和AlpacaEval上的性能下降。相反，降低β會導(dǎo)致模型失去其基本能力。在獎勵優(yōu)化與維持語言分類學(xué)之間尋求平衡已成為RLHF訓(xùn)練中的首要挑戰(zhàn)。

2.2 論文的方法--在線合并優(yōu)化器

受到離線合并的啟發(fā)，本節(jié)論文探討將模型合并融入到RLHF優(yōu)化步驟中。論文首先審視常用的基于梯度的優(yōu)化器。

2.2.1 從基于梯度的優(yōu)化器到在線合并優(yōu)化器

離線任務(wù)運算技術(shù)通過將LLMs的增量參數(shù)聚合來合并LLMs。相應(yīng)地，論文旨在合并參考SFT模型和第t次訓(xùn)練步驟的政策模型。

然而，論文實證發(fā)現(xiàn)直接優(yōu)化等式是不穩(wěn)定的且難以收斂，并且等式需要額外的緩存來存儲預(yù)訓(xùn)練模型的參數(shù)。因此，論文對等式進行了松弛處理，這種松弛仍然與論文在優(yōu)化步驟中應(yīng)用離線合并的動機很好地一致。這種松弛的另一個重要好處是避免了緩存額外的參數(shù)，增強了內(nèi)存效率。通過這種松弛，論文表明在每個優(yōu)化步驟中的在線合并可以通過基于梯度的增量權(quán)重與參考模型的增量參數(shù)之間的整合來近似。

2.2.2 實現(xiàn)方法

需要注意的是，論文的優(yōu)化器框架具有高度的靈活性，并與現(xiàn)有的模型合并方法兼容。論文基于廣泛使用的模型合并方法DARE和TIES開發(fā)了兩種在線合并優(yōu)化器：

OnDARE優(yōu)化器。DARE采用隨機稀疏化方法和線性組合作為共識方法。

OnTIES優(yōu)化器。TIES使用top-k百分比稀疏化和基于符號的共識方法。具體來說，它從每個合并候選中保留關(guān)于絕對值的top-p百分比的參數(shù)。它根據(jù)符號和范數(shù)計算元素級多數(shù)符號，丟棄與多數(shù)符號不同的參數(shù)，并對剩余的參數(shù)進行加權(quán)求和。

2.3 論文的效果

2.3.1 實驗設(shè)置

數(shù)據(jù)集。論文在廣泛使用的偏好數(shù)據(jù)集ULTRAFEEDBACK上進行實驗。與原始版本相比，ULTRAFEEDBACK的訓(xùn)練和評估部分分別包含約61K和2K個偏好對，由GPT-4進行排序，并通過人工努力進行清理，同時使用流行的基準(zhǔn)測試(如TruthfulQA)進行了無害化處理。ULTRAFEEDBACK中的提示規(guī)模龐大、細(xì)粒度且來源多樣。

訓(xùn)練。論文主要探索了在ULTRAFEEDBACK數(shù)據(jù)集上直接偏好優(yōu)化(DPO)中的在線合并優(yōu)化器，因為與近端策略優(yōu)化(PPO)相比，DPO由于其較低的訓(xùn)練成本而在當(dāng)今大規(guī)模LLM對齊中得到廣泛應(yīng)用。一般的DPO包括從策略模型中采樣和標(biāo)注響應(yīng)。在這項工作中，論文使用DPO的離線策略設(shè)置，直接在ULTRAFEEDBACK數(shù)據(jù)集上訓(xùn)練論文的策略模型，這也被證明在提高有益性和無害性方面是有效的。

評估。評估對齊的大型語言模型是一項具有挑戰(zhàn)性的任務(wù)。論文遵循全面評估產(chǎn)生可信評估的直接原則。因此，論文的評估包括7個類別中的12個公共基準(zhǔn)，評估對齊LLM的綜合能力:

數(shù)學(xué)：(1)GSM8K(2)Math401(3)Math23K;

編碼:(1)HumanEval(2)MBPP(3)DS1000(4)CodeApex;

指令遵循(IF):(1)IFEval;

閱讀理解(RC):(1)COQA(2)DROP;

知識：(1)MMLU;

agent：(1)NousResearch;

代碼混合;

論文使用類別內(nèi)的平均分?jǐn)?shù)作為最終類別分?jǐn)?shù)，使用所有類別中所有基準(zhǔn)的平均分?jǐn)?shù)作為整體評估。論文還使用基于GPT-4的評估器，通過長度控制評分，引入了MT-Bench 4和AlpacaEval 2.0，這是兩個領(lǐng)先且流行的基準(zhǔn)，用于評估LLM與人類偏好的一致性。

基線。論文方法的一個簡單基線是普通的AdamW。論文進一步考慮離線合并方法，如線性合并、DARE和TIES作為論文的強基線，因為Lin et al. (2024)表明簡單合并可以減輕對齊成本。由于對齊成本與遺忘有關(guān)，論文還將傳統(tǒng)的正則化方法，如KL懲罰、EMA和ChildTuning作為論文的基線。

具體來說，DPO算法使用超參數(shù)β調(diào)整KL懲罰。論文還將LoRA作為論文的基線之一，因為參數(shù)高效方法在訓(xùn)練中應(yīng)用了關(guān)于權(quán)重空間的正則化。

配置。論文使用三種LLM尺寸進行實驗，即Qwen1.5-1.8B、Qwen1.5-7B和LLaMa3-8B系列。

具體而言，論文使用Qwen-1.8B-Base、Qwen-7B-Base和LLaMa-3-8B作為在線優(yōu)化器的基礎(chǔ)模型。論文在Qwen1.5-1.8B-SFT、Qwen1.5-7B-SFT和LLaMa-3-8B-it作為參考模型上對ULTRAFEEDBACK進行直接偏好優(yōu)化。兩個Qwen1.5監(jiān)督微調(diào)模型在多語言指令數(shù)據(jù)上進行了訓(xùn)練，但與ULTRAFEEDBACK數(shù)據(jù)集沒有重疊。

【LLM】在線合并優(yōu)化器以提升獎勵并減輕對齊開銷-AI.x社區(qū)

2.3.2 主要結(jié)果

論文在表4中展示了論文的主要結(jié)果，其中展示了基線方法和論文提出的在線合并優(yōu)化器在ULTRAFEEDBACK上的性能，分別在Qwen1.5-1.8BChat、Qwen1.5-7B-Chat和LLaMa-3-8B-Chat上訓(xùn)練。總的來說，與原始的AdamW相比，正則化和離線模型合并方法在大多數(shù)設(shè)置下并沒有顯著提高RLHF模型在基準(zhǔn)測試上的平均性能，反而導(dǎo)致MT-Bench和AlpacaEval 2.0的偏好分?jǐn)?shù)下降。這表明，僅僅依靠梯度dropout(ChildTuning)、抑制模型梯度更新的變化(EMA)或基于SFT模型對RLHF模型參數(shù)進行一次性調(diào)整(Merging)等技術(shù)，都不能有效解決對齊獎勵-稅收的權(quán)衡問題。

正則化基線在LLama-3-8B-It上表現(xiàn)特別好，所有正則化方法在平均基準(zhǔn)測試分?jǐn)?shù)以及MT-Bench和AlpacaEval分?jǐn)?shù)上都取得了一致的改進。相比之下，論文提出的在線合并優(yōu)化器，特別是OnDARE變體，在所有測試集上都取得了最顯著的改進。OnDARE在基準(zhǔn)測試平均分上取得了最高的改進，并在所有三個主干LLM上持續(xù)增強了MT-Bench和AlpacaEval 2.0，顯著超過其他基線，尤其是在LLaMa-3-8B-Instruct實驗中，分別在基準(zhǔn)測試、MT-Bench和AlpacaEval上取得了1.3、0.19和1.57的改進。

盡管OnTIES和OnDARE在提高獎勵和減輕稅收方面都顯示出了有效性，但在大多數(shù)情況下，OnDARE在平均基準(zhǔn)測試分?jǐn)?shù)方面略優(yōu)于OnTIES，而OnTIES在AlpacaEval 2.0上有時具有更高的LC勝率。每個基準(zhǔn)測試的詳細(xì)分?jǐn)?shù)在表5、表6、表7中報告。

2.3.3 超參數(shù)效應(yīng)

本節(jié)分析了兩個主要超參數(shù)，即參數(shù)保留率和合并權(quán)重，如何影響在線合并優(yōu)化器的整體性能。

參數(shù)保留率是在線合并期間的參數(shù)保留率。論文在Qwen1.5-1.8B-Chat上探索了從1到1e-5的保留率，以在有限的計算資源內(nèi)最大化搜索空間。如圖2所示，即使在低至5e-4的低參數(shù)保留率下，在線合并優(yōu)化器仍然保持穩(wěn)健。這表明，在每個RLHF步驟中丟棄99.95%的基于梯度的參數(shù)修改仍然可以得到穩(wěn)定的訓(xùn)練。與OnDARE相比，OnTIES對極低的參數(shù)保留率更敏感。這種敏感性是由于OnDARE采用了無偏隨機稀疏化方法，而OnTIES使用的top-k稀疏化在訓(xùn)練過程中引入了顯著的偏差。

【LLM】在線合并優(yōu)化器以提升獎勵并減輕對齊開銷-AI.x社區(qū)

合并權(quán)重是參考(SFT)模型的梯度和增量參數(shù)的聚合權(quán)重。較大的合并權(quán)重在在線合并優(yōu)化器中引入更強的正則化。論文嘗試了從10?4到10?7的各種合并權(quán)重，并在表2中報告了結(jié)果。

隨著合并權(quán)重的增加，由于訓(xùn)練過程中添加的正則化減少，MT-Bench分?jǐn)?shù)上升，而平均基準(zhǔn)測試分?jǐn)?shù)先增加后減少，在α=5e?7時達到峰值。與EMA訓(xùn)練中的指數(shù)系數(shù)類似，較大的合并權(quán)重會導(dǎo)致訓(xùn)練不穩(wěn)定。因此，論文建議從10?7這樣的小數(shù)字開始合并權(quán)重的超參數(shù)搜索。值得注意的是，特殊情況合并權(quán)重=0使OnDARE等同于梯度dropout正則化方法，如論文基線中的ChildTuning。

2.3.4 RLHF算法的影響

論文進一步研究了它們在其他RLHF算法中的應(yīng)用。具體而言，論文在IPO和KTO中實現(xiàn)了OnDARE和OnTIES。論文在這些設(shè)置下在ULTRAFEEDBACK數(shù)據(jù)集上訓(xùn)練Qwen1.5-7B-Chat，并在表3中展示了結(jié)果。

除了IPO算法的平均基準(zhǔn)測試分?jǐn)?shù)外，在線合并優(yōu)化器OnDARE和OnTIES在MT-Bench上的表現(xiàn)都優(yōu)于AdamW。這表明它們的有效性可以擴展到多個RLHF算法變體。

【LLM】在線合并優(yōu)化器以提升獎勵并減輕對齊開銷-AI.x社區(qū)