跳到主要內容
AI News HubLIVE
站內改寫3 分鐘閱讀

分開訓練,合併使用:基於混合專家模型的分模塊後訓練

文章摘要

BAR(分支-適配-路由)是一種逐步訓練語言模型能力的方案:獨立訓練領域專家,將它們合併為一個混合專家模型,並能在不影響其他模塊的情況下升級任一專家。

來源Ai2 Blog
分開訓練,合併使用:基於混合專家模型的分模塊後訓練
報告錯誤

更正渠道尚未開通,可先複製下方文章資訊留存。

查看更正說明
直接讀正文

在預訓練之後,語言模型需要經過一系列中期和後訓練階段才能變得實用——學習遵循指令、推理問題、可靠地調用工具等。然而,在這些階段之後更新或擴展模型往往具有挑戰性。最可靠的選擇是從頭開始重新訓練,但成本高昂且需要完全訪問原始訓練設置。進一步訓練新數據更便宜,但可能導致模型丟失已有能力。而且,由於後訓練通常涉及多個階段(每個階段有自己的數據和目標),添加新技能意味着需要重新運行或調整每個階段以避免破壞已有成果。

為了解決這些問題,艾倫人工智能研究所(Ai2)的研究團隊提出了BAR(Branch-Adapt-Route,分支-適配-路由)方法,一種模塊化後訓練的方案。與一次訓練一個處理所有數據的單一模型不同,BAR獨立訓練領域專家——每個專家通過自己的完整訓練流程——並通過混合專家(MoE)架構將它們組合成一個統一模型。每個專家都可以獨立開發、升級或替換,而無需影響其他專家。該團隊發佈了配方、技術報告以及用於驗證該方法的模型檢查點。

BAR的動機源於之前的FlexOlmo工作。FlexOlmo表明,基於MoE的模塊化訓練在預訓練階段效果良好:可以從共享基礎分支,訓練特定領域的FFN專家同時凍結所有共享層,然後合併回。然而,研究人員發現這種方案不適用於後訓練。原因在於:預訓練主要更新知識表示,這些知識主要存在於FFN層;而後訓練則引入行為變化(如新的輸出格式、推理模式和安全約束),這些變化需要改變共享參數(如注意力層、嵌入層和語言模型頭)。例如,當研究人員在強化學習與驗證獎勵(RLVR)中直接使用FlexOlmo方法時,獎勵曲線完全平坦;模型在所有共享參數凍結的情況下根本學不會。這促使他們專門為後訓練開發了新的配方。

BAR包含三個階段。階段一:獨立專家訓練。每個領域專家實例化為包含兩個專家的MoE:一個凍結的“錨定”專家(保留基礎模型的FFN權重)和一個可訓練的專家。專家根據領域需求經歷相應的訓練階段。在實驗中,數學和代碼專家經歷中期訓練、有監督微調(SFT)和RLVR;工具使用和安全專家僅使用SFT。關鍵的技術貢獻是跨階段的漸進式解凍計劃:中期訓練時所有共享層凍結;SFT時解凍嵌入層和語言模型頭,這對於引入新特殊令牌(如工具使用的函數調用格式)是必要的;RLVR時解凍所有共享參數,包括注意力層。此外,每個專家還訓練混合了領域特定和通用SFT數據的數據集,因為僅領域SFT會嚴重損害通用能力。

階段二:專家合併。訓練後,所有專家合併為一個單一的MoE模型。跨專家運行中分化的共享參數(由於在SFT或RLVR期間解凍)被簡單地平均。研究發現,與任何單個專家相比,這種平均化在領域特定評估中幾乎沒有引入可測量的性能損失。階段三:路由器訓練。最後,在MoE內部訓練路由器,同時凍結所有其他專家和共享權重。研究人員發現,僅用SFT數據的分層5%樣本就足以實現有效路由,使得這一階段快速且廉價。

實驗表明,BAR模型在19個基準測試中超越了所有不需要從零重新訓練中期訓練的基線方法。平均而言,BAR優於僅進行後訓練重訓練的模型(49.1 vs 47.8),尤其在數學(+7.8)和代碼(+4.7)方面提升顯著。這歸因於模塊化訓練的結構性優勢:在單一流水線中,後期對數學和代碼的RL可能會削弱早期SFT階段學到的安全能力,而模塊化訓練完全避免了這一點。值得注意的是,中期訓練後的稠密模型合併效果極差(整體6.5),而BAR在無需中期訓練的情況下仍然大幅領先。

BAR最實用的特性之一是專家可以獨立升級。研究展示了兩種升級方式:用更高質量數據和RL訓練的代碼專家替換原有專家,使代碼性能提升16.5點,而其他領域幾乎不變;在已有數學專家基礎上添加RL訓練,使數學性能提升13點。在這兩種情況下,只有受影響的專家和輕量級路由器需要重新訓練,而傳統單一流水線則需要重新訓練整個模型。這使得BAR的領域更新成本呈線性增長,而單一模型更新成本則幾乎是二次方增長。

研究還得出一些實用經驗:後訓練需要比預訓練更大的靈活性,漸進式解凍至關重要;僅使用領域SFT數據是不夠的,必須混合通用SFT數據;解凍後的權重平均效果出奇地好;並非所有專家都需要在推理時激活,激活4/5專家即可達到幾乎相同的性能。未來工作包括從原生稀疏架構開始,而非將稠密模型轉換為稀疏模型,這可能提高模塊化方法的效率和可擴展性。

展開要點與分析

文章情報

工程師進階

要點

  • BAR通過獨立訓練領域專家並使用混合專家架構組合,實現了模塊化後訓練。
  • 漸進式解凍共享參數至關重要:SFT階段解凍嵌入和語言模型頭,RL階段解凍注意力層。
  • BAR在不重新訓練整個模型的情況下,優於多數基線方法,並支持獨立升級專家。
  • 將通用SFT數據與領域數據混合訓練是保持模型通用能力的關鍵。

要點與分析由自動化流程生成,可能有誤,請結合原始來源核實。