將混合專家推向極限:只需更新0.32%的參數就能完成模型微調
眾所周知,大模型的訓練成本很高,但其實對預訓練后的模型進行微調也需要一定的成本,還好我們已經有了 (IA)3 或 LORA 等一些參數高效型微調(PEFT)方法。
近日,AI 創業公司 Cohere 更進一步,將混合專家方法與 PEFT 組合,實現了極其參數高效的微調 —— 即使是在未曾見過的任務上,這種新方法只需更新不到 1% 的參數,就能實現與完全微調方法相當的性能。
Cohere 公司在攻克這些挑戰上邁出了重要一步,提出了一個新框架 —— 可在大幅受限的計算環境中利用 MoE 的優勢。
- 論文鏈接:https://arxiv.org/abs/2309.05444
- 代碼鏈接:https://github.com/for-ai/parameter-efficient-moe
傳統的訓練范式是將模型的權重應用于每個輸入。有理由相信這種方法效率不高,因為單個輸入可能無法讓模型發揮全部力量。
相比之下,混合專家(MoE/Mixture of Experts)則是基于多個子模塊組件(即所謂的專家),這些組件可以專門處理不同類型的輸入。這種對于條件計算的強調對效率具有重要的影響,比如會有恒定的推理成本。這使得 MoE 成了大規模 Transformer 時代一個重要的研究領域并得到了廣泛的采用;而在這些應用中,規模的擴張也會造成部署成本和延遲上升。
盡管截至目前的大部分研究關注的是將 MoE 用作一種預訓練策略,但 MoE 的內在動機并不僅局限于預訓練。事實上,MoE 具有一些非常適合指令微調的特性;在指令微調設置中,數據會被刻意地構建,以代表多樣化的任務,這通常也被稱為多任務微調。
基于這一觀察,AI 創業公司 Cohere 提出了一個問題:能否使用 MoE 來進行指令微調?MoE 范式的一大主要缺點是會引入大量參數。盡管是基于條件執行計算,但完全微調 MoE 架構需要更新所有參數,這需要非常大量的計算。對于大多數實踐者來說,考慮到現代 LLM 的規模,這樣的計算成本是無力支撐的。
于是,Cohere 公司關注的是更為現實的實踐場景,即將 MoE 用于參數高效型微調(PEFT)方法,如 (IA)3 或 LORA,它們微調的參數數量會少很多。這是一個重大挑戰,不僅因為目標是僅更新所有參數的一小部分,而且因為還要在更受限的環境中解決 MoE 固有的優化挑戰。
他們引入了 Mixture of Vectors(MoV)和 Mixture of LORA(MoLORA),一種參數高效型混合專家適應方法。不同于標準的 MoE,這種新框架很輕量,可用于對參數有限制的場景。
值得注意的是,在未曾見過的任務上,這種新方法只需更新不到 1% 的參數,就能實現與完全微調方法相當的性能。其表現也能輕松勝過 (IA)3 或 LORA 等基礎的參數高效型技術。
該團隊基于 55 個數據集,在 12 個不同任務上,用 770M 到 11B 的不同大小 T5 模型進行了實驗,得到了相當一致的結果。
本文的主要貢獻包括:
1. 提出了極其參數高效的 MoE。該架構利用了模塊化和輕量級的專家,可在更加現實的設置中使用 MoE。使用這個新提出的 MoE 來訓練一個密集模型時,僅需更新其不到 1% 的參數。
2. 使用新提出的方法執行指令微調時,在未曾見過的任務上能取得持續優于傳統參數高效方法的性能,同時還能為不同大小的模型維持較高的參數效率。相較于 3B 和 11B 模型大小的標準 (IA)3,混合 (IA)3 向量(MoV)分別實現了 14.57% 和 8.39% 的提升。這種優越性在各種大小的模型、各種類型的專家和可訓練參數負載上都有體現。
3. 研究表明新提出的方法能取得與大規模完全微調方法相當的性能表現,同時只需更新一小部分模型參數。研究者在 8 個未曾見過的任務上進行了實驗,結果表明 MoV 僅會分別更新 3B 和 11B 模型中 0.32% 和 0.86% 的參數,卻能以顯著更低的計算成本實現與完全微調方法相當的表現。
4. 他們也進行了一系列廣泛的消融研究,系統性地評估了各種 MoE 架構和 PEFT 策略的效能,涉及多種模型大小、不同的適應器類型、專家的數量、路由機制和優化超參數的重要性,尤其是考慮到 MoE 的敏感性。
方法
指令微調可以形式化地表述為:存在一個任務集合 T,將其分為訓練任務集 T_train 和留存評估集 T_eval。
首先在 T_train 上對基礎預訓練模型進行微調,然后使用零樣本方式在 T_eval 中的每個任務上評估微調后的模型。標準的微調方法是微調所有模型參數,這樣會有很高的計算和內存成本。新提出的方法卻有很高的參數效率,其使用了參數高效型混合專家,下面詳細描述一下該框架。
(IA)3 和 LORA 參數高效型微調方法
參數高效型微調(PEFT)方法只會更新少量參數的權重。為了展示新方法適用于不同的 PEFT 技術,研究者實驗了 (IA)3 和 LORA。這些方法會向已有的預訓練模型模型添加少量參數。這里就不對 (IA)3 和 LORA 做更詳細的介紹了,感興趣的讀者可參閱原論文以及機器之心的報道《調教 LLaMA 類模型沒那么難,LoRA 將模型微調縮減到幾小時》。
極其參數高效的混合專家
新提出的極其參數高效的 MoE 框架使用了輕量級的「適應器」作為預訓練密集模型之上的專家。
具體來說,MoE 是一系列神經網絡架構,可以通過基于門控機制(路由器)激活的多個專家來進行條件計算。MoE 層由一個路由器網絡 R 和一個專家集合 E 構成,其中 E 包含 n 個專家,每個專家 E_i 都是一個已參數化的函數。至于路由器網絡 R 的設計,該團隊參照了論文《Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity》,通常由一個密集層后帶一個 softmax 函數構成;其中這個密集層有可訓練的權重 W_g,softmax 函數則是以中間 token 表征 x 為輸入并基于門控分數 s_1, ..., s_n 將每個專家的輸出組合到一起:
對于 Transformer 模型,其中的密集型前向層被 MoE 層取代,該層中每個專家 E_i 都對應于一個獨立的密集型前向網絡。隨著每個專家的規模和專家數量的增加,模型的參數總數會成倍增大。但是,在新提出的參數高效型 MoE 架構中,每一個專家都被一個輕量級 PEFT 適應器(比如 (IA)3 向量或 LORA 適應器)替代。
在微調過程中,密集層的預訓練權重保持固定,而專家層和路由器層則從頭開始訓練。不同于標準 MoE,這種新提出的輕量級專家可以在微調時間學習適應預訓練 Transformer 層。這樣一來,新的 MoE 框架只需要少量參數更新,而不是整體對大模型進行更新。
除了參數效率之外,研究者選擇的 PEFT 適應器可通過 soft merging 實現路由計算。具體來說,由于 (IA)3 向量和 LORA 適應器都是線性函數,所以他們首先會計算專家的加權平均,然后使用組合后的專家 E_mix 應用一個 PEFT 變換:
對于把 (IA)3 向量或 LORA 適應器分別用作專家的方法,研究者將它們分別命名為 Mixture of Vectors (MoV) 和 Mixture of LORA (MoLORA),這兩種方法都能持續穩定地優于對應的 PEFT 方法。下圖 2 展示了 MoV 層的架構以及相應的偽代碼。
通過 MoV 和 MoLORA 只更新少部分參數有一些實際的好處 —— 不僅有益于訓練,也有助于推理;對于后者的優勢是 MoE 架構所獨有的。下面簡要描述一些這些好處:
訓練效率
這種極其參數高效的 MoE 形式可以顯著減少內存消耗。通過在訓練期間凍結大多數參數,不僅能減少計算模型參數的梯度的計算開銷,而且還可降低存儲模型的優化器狀態的內存需求。根據優化器的選擇,后者可能非常重要,舉個例子,AdamW 等 Adam 優化器變體為了存儲優化器狀態,每個參數需要兩倍的內存(根據一階矩和二階矩的估計),而 Adafactor 則能通過分解二階參數矩的估計來將這種開銷減半。
推理效率
MoV 和 MoLORA 方法的固有結構模塊化特性能為推理帶來顯著的內存效益。對于傳統的 MoE 模型,很多完全用于前向的塊的副本(甚至基于特定架構的模型的完整副本)需要在推理時間被存儲在內存中,這種做法的成本很高。
使用這種方法,不管確切的類型如何,都只需要將模型主干的一個副本保存在內存中,外加上輕量級參數高效型專家。這能在推理時間顯著降低對內存的需求。
實驗
參數高效型 MoE 與 PEFT 方法:新的 MoE 方法與單專家型 PEFT 方法相比如何?下表 1 比較了 PEFT 方法((IA)3 和 LORA)與新提出的參數高效型 MoE 方法(MoV 和 MoLORA)的零樣本性能,其中基礎模型是 T5-3B。
可以觀察到,相比于標準的 (IA)3 向量和 LORA 適應器,MoE 變體(MoV 和 MoLORA)的性能有顯著優勢。
比如,使用 30 個專家的 MoV 的表現優于密集型 (IA)3 方法 14.57%,使用 15 個專家的 MoLORA 則在平均中值分數上提升了 5.70%。
在同樣的參數負載下,MoV 的表現勝過 MoLORA。研究者也比較了新提出的兩種方法:使用 3B 參數的基礎模型時,MoV 在性能與參數成本的權衡方面表現更好。
參數高效型 MoE 與完全微調方法的對比。如表 1 所示,與完全微調的 T0-3B 相比,使用 10 個專家的 MoV 和 MoLORA 的性能相當。這是一個非常亮眼的結果,因為 MoV-10 僅更新了全體模型參數中的 0.32%。而如果將 MoV 的專家數增至 15,將 MoLORA 的專家數增至 30,新方法的效果甚至能小幅勝過完全微調方法。
此外,研究者還通過實驗探索了參數高效型 MoE 隨基礎模型大小增長的性能變化情況、專家數量對下游任務性能的影響、最優的路由策略等方面,詳見原論文。