機器之心報道
編輯:Panda W
LLM 很強,而為了實作 LLM 的可持續擴充,有必要找到并實作能提升其效率的方法,混合專家(MoE)就是這類方法的一大重要成員。
最近,各家科技公司提出的新一代大模型不約而同地正在使用混合專家(Mixture of Experts:MoE)方法。
混合專家這一概念最早誕生于 1991 年的論文《Adaptive mixtures of local experts》,三十多年來已經得到了廣泛的探索和發展。近年來,随着稀疏門控 MoE 的出現和發展,尤其是與基于 Transformer 的大型語言模型相結合,這種已有三十多年曆史的技術煥發出了新的生機。
MoE 架構基于一個簡單卻又強大思想:模型的不同部分(稱為專家)專注于不同的任務或資料的不同方面。
使用這一範式時,對于一個輸入,僅有與之相關的專家(Expert)才會參與處理,這樣一來便能控制計算成本,同時仍能受益于大量專業知識。是以,MoE 可在不大幅提升計算需求的前提下提升大語言模型的能力。
如圖 1 所示,MoE 相關研究增長強勁,尤其是在 2024 年 Mixtral-8x7B 以及 Grok-1、DBRX、Arctic、DeepSeek-V2 等各種産業級 LLM 出現之後。
這張圖來自香港科技大學(廣州)的一個研究團隊近日釋出的一篇 MoE 綜述報告,其中清晰且全面地總結了 MoE 相關研究,并提出了一種全新的分類法,将這些研究歸類到了算法、系統和應用三大類。
論文标題:A Survey on Mixture of Experts
論文位址:https://arxiv.org/pdf/2407.06204
機器之心整理了這篇綜述報告的内容主幹,以幫助讀者了解目前 MoE 的發展概況,更多詳情請閱讀原論文。此外,我們也在文末整理了一些與 MoE 相關的報道。
混合專家的背景知識
在基于 Transformer 的大型語言模型(LLM)中,每個混合專家(MoE)層的組成形式通常是 個「專家網絡」{_1, ... , _} 搭配一個「門控網絡」G。
這個門控網絡的形式通常是一個使用 softmax 激活函數的線性網絡,其作用是将輸入引導至合适的專家網絡。MoE 層的放置位置是在 Transformer 子產品内,作用是選取前向網絡(FFN),通常位于自注意力(SA)子層之後。這種放置方式很關鍵,因為随着模型增大,FFN 的計算需求也會增加。舉個例子,在參數量達到 5400 億的 PaLM 模型中,90% 的參數都位于其 FFN 層中。
用數學形式描述的話:每個專家網絡 _ (通常是一個線性 - ReLU - 線性網絡)都由 W_ 進行參數化,其接收同一輸入 x 并生成輸出 _ (x; W_)。同時,參數為 Θ 的門控網絡 G(通常由一個線性 - ReLU - 線性 - softmax 網絡構成)得到輸出 G (x; Θ)。再根據門控函數的設計方式,可以将 MoE 層大緻分為以下兩類。
密集 MoE
密集混合專家層是在每次疊代過程中激活所有專家網絡 {_1, ... , _}。早期的 MoE 研究普遍采用了這一政策。近段時間也有一些研究采用了密集 MoE,比如 EvoMoE、MoLE 、LoRAMoE 和 DS-MoE。圖 2a 給出了密集 MoE 層的結構。是以,密集 MoE 層的輸出可以表示成:
其中,(x; Θ) 是 softmax 運算之前的門控值。
稀疏 MoE
盡管密集混合專家的預測準确度通常更高,但其計算負載也非常高。
為了解決這個問題,Shazeer et al. 的論文《Outrageously large neural networks: The sparsely-gated mixture-of-experts layer》引入了稀疏門控 MoE 層,其能在每次前向通過時僅激活標明的專家子集。該政策實作稀疏性的方式是計算 top-k 個專家的輸出的權重和,而非将所有專家的輸出聚合到一起。圖 2b 展示了這種稀疏 MoE 層的結構。
根據上述論文提出的架構,可對 2.2 式進行修改以反映稀疏門控機制:
這裡解釋一下:TopK (・, ) 函數是僅保留向量原始值的前 k 項,同時将其它項設定為 −∞。之後是 softmax 運算,所有 −∞ 項都會變成近似于零。超參數 k 要根據具體應用選取,常見選項是 = 1 或 = 2。加入噪聲項 R_noise 是訓練稀疏門控 MoE 層的一種常用政策,可促進專家之間的探索并提升 MoE 訓練的穩定性。
盡管稀疏門控 G (x; Θ) 可在不增加相應計算成本的前提下顯著擴充模型的參數空間,但也會導緻負載平衡問題。負載平衡問題是指各專家的負載分布不均 —— 某些專家被頻繁使用,而另一些專家則很少被使用甚至完全不上場。
為了解決這個問題,每個 MoE 層都要內建一個輔助損失函數,其作用是敦促每批次的 token 被均勻配置設定給各個專家。從數學形式描述來看,首先定義一個包含 T 個 token 的查詢批次 B = {x_1 , x_2, ... , x_ } 以及 N 個專家。則對于其的輔助負載平衡損失定義為:
其中 D_i 是配置設定給專家 i 的 token 比例,P_i 是配置設定給專家 i 的門控機率比例。為了確定該批次在 N 個專家之間均勻分布,應當最小化負載平衡損失函數 L_{load-balancing}。當每個專家都被配置設定了同等數量的 token D_ = 1/ 和同等的門控機率 P_ = 1/ 時,即達到最優條件:
此時各專家的負載達到平衡。
在後文中,除非另有明确說明,則「MoE」這一術語單指「稀疏 MoE」。
混合專家的分類
為了幫助研究者在大量采用 MoE 的 LLM 研究中找到目标,該團隊開發了一套分類方法,根據三個方面對這些模型進行了分類:算法設計、系統設計和應用。
圖 3 展示了這種分類法以及一些代表性研究成果。
下面将全面深入地介紹各類别的情況。
混合專家的算法設計
門控函數
門控函數(也被稱為路由函數或路由器)是所有 MoE 架構的基礎元件,其作用是協調使用專家計算以及組合各專家的輸出。
根據對每個輸入的處理方法,該門控可分為三種類型:稀疏式、密集式和 soft 式。其中稀疏式門控機制是激活部分專家,而密集式是激活所有專家,soft 式則包括完全可微方法,包括輸入 token 融合和專家融合。圖 4 展示了 MoE 模型中使用的各種門控函數。
稀疏式
稀疏門控函數在處理各個輸入 token 時會激活被選中的部分專家,這可被視為一種形式的條件計算。
門控函數可以實作多種形式的門控決策,比如二進制決策、稀疏或連續決策、随機或确定性決策;其已經得到了深入的研究,可使用各種形式的強化學習和反向傳播來訓練。
Shazeer et al. 的研究《Outrageously large neural networks: The sparsely-gated mixture-of-experts layer》開創性地提出了一種使用輔助負載平衡損失的可微分啟發式方法,其中可根據選取機率對專家計算的輸出進行權重。這為門控過程引入了可微性,由此可通過梯度來引導門控函數的優化。
後來,這一範式便成了 MoE 研究領域的主導範式。由于這種方法會針對每個輸入 token 選擇專家,是以可将其看作是 token 選擇式門控函數。
以下為這一小節的要點,詳見原論文:
token 選擇式門控
用于 token 選擇式門控的輔助損失
token 選擇式門控的專家容量
token 選擇式門控的其它進展
不可訓練的 token 選擇式門控
專家選擇式門控
密集式
密集 MoE 是指處理每個輸入時都激活所有專家。
雖然稀疏 MoE 有效率方面的優勢,但密集 MoE 方向依然在不斷迎來創新。尤其值得一提的是,密集激活在 LoRA-MoE 微調方面表現很好,并且 LoRA 專家的計算開銷相對較低。這種方法能夠有效靈活地內建多個 LoRA 以完成各種下遊任務。這能保留原始預訓練模型的生成能力,同時保留各個 LoRA 針對各個任務的獨有特性。
soft 式
對稀疏 MoE 來說,一大基本離散優化難題是如何決定為每個 token 配置設定哪些合适的專家。為了確定專家能平衡地參與并盡可能減少無配置設定 token,這通常必須啟發式的輔助損失。在涉及分布外資料的場景(比如推理批次小、有全新輸入或遷移學習)中,這個問題尤其顯著。
類似于密集 MoE,soft MoE 方法在處理每個輸入時也會使用所有專家,進而維持完全可微性,進而避免離散專家選擇方法的固有問題。soft MoE 與密集 MoE 的不同在于前者會通過對輸入 token 或專家進行門控權重的融合來緩解計算需求。
專家
這一節會介紹 MoE 架構内專家網絡的架構,并會讨論協調這些專家的激活的門控函數。
網絡類型
自從 MoE 被整合到 Transformer 架構中以來,其通常會替代這些模型中的前向網絡(FFN)子產品。通常來說,MoE 層中的每個專家都會複制其替換的 FFN 的架構。
這種将 FFN 用作專家的範式到現在依然是主流,但人們也對此做了不少改進。
超參數
稀疏 MoE 模型的規模由幾個關鍵超參數控制,包括:
每個 MoE 層的專家數量
每個專家的大小
MoE 層在整個模型中的放置頻率
這些超參數的選擇至關重要,因為它會深刻影響模型在各種任務中的性能和計算效率。是以,要根據特定的應用要求和計算基礎設施來選擇最佳超參數。表 2 給出了一些使用 MoE 的模型的配置情況。
另外,表 3 列舉了一些近期的開源模型的參數數量和基準性能。
激活函數
基于密集 Transformer 架構建構的稀疏 MoE 模型采用了與 BERT、T5、GPT 和 LLAMA 等領先的密集 LLM 類似的激活函數。激活函數已經從 ReLU 發展出了 GeLU、GeGLU、SwiGLU 等更先進的選擇。
這一趨勢也擴充到了 MoE 模型的其它元件,它們經常整合均方根層歸一化(RMSNorm)、分組查詢注意力(GQA)和旋轉位置嵌入(RoPE)等技術。
共享專家
DeepSpeed-MoE 創新性地引入了殘差 MoE(Residual-MoE)架構,其中每個 token 都由一個固定專家外加一個門控選擇的專家進行處理,實作了每一層都有兩個專家參與處理,同時也不會讓通信成本超過 top-1 門控方法。這種方法是把門控選擇的 MoE 專家當作是固定密集 FFN 的糾錯輔助。
NLLB 中使用的條件式 MoE 路由(CMR/Conditional MoE Routing)也采用了類似的方法,将密集 FFN 和 MoE 層的輸出組合起來使用。
将固定 FFN 和稀疏 MoE 整合起來的範式通常被稱為共享專家,如圖 5b 所示。
近期有 DeepSeekMoE、OpenMoE、Qwen1.5-MoE 和 MoCLE 等模型采用這一範式,表明其正在成為一種主流配置。不過 DeepSeekMoE 和 Qwen1.5-MoE 采用了多個共享專家,而不是單個。
混合參數高效型專家
參數高效型微調(PEFT)是一種提升微調效率的方法。簡單來說,PEFT 就是在微調時僅更新基礎模型的一小部分參數。
PEFT 很成功,但由于其可訓練的參數有限以及可能存在的災難性遺忘問題,該方法難以用于需要泛化到多個任務的情況。
為了緩解這些局限,混合參數高效型專家(MoPE)誕生了,其将 MoE 架構與 PEFT 整合到了一起。MoPE 內建了 MoE 的門控機制與多專家架構,同時每個專家都使用了 PEFT 技術進行建構。這種巧妙的組合能極大提升 PEFT 在多任務場景中的性能。此外,由于使用了 PEFT 來建構專家,是以 MoPE 使用的參數也更少,資源效率比傳統 MoE 模型高得多。
MoPE 融合了 MoE 的多任務特性與 PEFT 的資源效率,是一個極具前景的研究方向。圖 6 根據在 Transformer 模型架構中的位置對 MoPE 進行了分類。至于 MoPE 方面更詳細的研究成果介紹,請參看原論文。
訓練和推理方案
混合專家在進步發展,相關的訓練和推理方案也在進步發展。
初始的訓練和推理方案需要從頭開始訓練 MoE 模型,直接采用訓練的模型配置來執行推理。
但現在,MoE 模型的訓練和推理方面已經出現了許多新範式,包括組合密集模型和稀疏模型的優勢實作取長補短。
圖 7 展示了與 MoE 相關的訓練和推理方案,可以看到新出現的方案可分為三類:
密集到稀疏:從密集模型訓練開始,逐漸過渡到稀疏 MoE 配置;
稀疏到密集:涉及到将稀疏 MoE 模型降格為密集形式,這有利于将推理實作為硬體形式;
專家模型融合:将多個預訓練密集專家模型整合成一個統一的 MoE 模型。
MoE 的衍生技術
混合專家(MoE)啟發了許多不同的變體技術。舉個例子,Xue et al. 的論文《Go wider instead of deeper》提出了模型寬度增大的 WideNet,其做法是将前向網絡(FFN)替換成 MoE 層,同時維持 Transformer 層上的共享可訓練參數,但歸一化層除外。
另外還有 Tan et al. 提出的 SYT(稀疏通用 Transformer)、Antoniak et al. 提出的 MoT(混合 token)、Choi et al. 提出的 SMoP(稀疏混合提詞)、Chen et al. 提出的 Lifelong-MoE、Raposo et al. 提出的 MoD(混合深度)等。
總結一下,MoE 衍生技術的發展揭示了一個趨勢:MoE 的功能越來越多,越來越能适應不同的領域。
混合專家的系統設計
混合專家(MoE)雖然能增強大型語言模型的能力,但也帶來了新的技術挑戰,因為其具有稀疏且動态的計算負載。
GShard 引入了專家并行化(expert parallelism),可根據專家能力的負載平衡限制來排程切分後的局部 token,進而實作并行的門控和專家計算。該範式已經成為促進 MoE 模型高效擴充的基礎政策。我們可以将該方法看作是增強版的資料并行化 ——MoE 層中的每個專家都被配置設定到一台不同裝置,同時所有裝置上都重複配備所有非專家層。
如圖 8a 所示,專家并行化的工作流程是按順序執行以下操作:門路由、輸入編碼、All-to-All 排程、專家計算、All-to-All 組合、輸出解碼。
一般來說,GEMM 的輸入大小需要足夠大,以便充分利用計算裝置。是以,要使用輸入編碼将同一個專家的輸入 token 聚合到連續的記憶體空間中,這由門路由中的「token - 專家映射」決定。之後,All-to-All 排程的作用是将輸入 token 分發給各裝置上對應的專家。之後是專家的本地化計算。計算完成後再通過 All-to-All 組合彙總,然後解碼輸出,根據門控索引恢複原始資料的布局。
此外,也有研究者探索專家并行化與其它已有并行政策(比如張量、管道化、序列并行化)的協同,以提升 MoE 模型在大規模分布式環境中的可擴充性和效率。
圖 8 中給出了一些混合并行化示例,包括 (b) 資料 + 專家 + 張量并行化、(c) 資料 + 專家 + 管道并行化、(d) 專家 + 張量并行。
需要認識到,計算效率、通信負載、記憶體占用之間存在複雜的互相作用,分布式并行化政策的選擇會對其産生影響,并且也會被不同的硬體配置影響。是以,在部署用于實際應用的政策時,必須細緻地權衡考慮并針對具體場景進行調整。
之後,該團隊分計算、通信和存儲三大闆塊介紹了 MoE 模型開發所面臨的系統設計難題以及解決這些難題的研究成果,詳見原論文。表 4 給出了開源 MoE 架構的概況。
混合專家的應用
在目前 Transformer 主導的大型語言模型(LLM)領域,混合專家(MoE)範式頗具吸引力,因為其能在不給訓練和推理階段引入過大計算需求的前提下大幅提升模型能力。這類技術能顯著 LLM 在多種下遊任務上的性能,甚至造就了一些超越人類水準的 AI 應用。
有傳言說強大如斯的 GPT-4 可能也采用了某種 MoE 架構 —— 由 8 個 2200 億參數的專家構成,在多樣化的資料集和任務上完成了訓練,并使用了一種 16 次疊代的推理過程。有關該傳言的更多詳情可參閱機器之心報道《終極「揭秘」:GPT-4 模型架構、訓練成本、資料集資訊都被扒出來了》。
是以,毫不奇怪 MoE 在自然語言處理、計算機視覺、推薦系統和多模态應用中遍地開花了。
這些應用本質上就需要使用條件計算來大幅提升模型的參數量,以此增強模型在固定計算成本下的性能,或通過門控機制實作動态專家選擇來實作高效多任務學習。
該團隊也介紹了這些不同領域的代表性 MoE 應用,可幫助讀者了解如何将 MoE 用于具體任務。詳見原論文。
挑戰與機遇
混合專家,功能強大,降低成本,提升性能。前景雖好,仍有挑戰。
這一節中,該團隊梳理了 MoE 相關的關鍵性挑戰,并指出了有希望獲得重要成果的未來研究方向。下面簡要列出了這些挑戰和研究方向,更多詳情請檢視原論文。
- 訓練穩定性和負載平衡
- 可擴充性和通信開銷
- 專家的專業化和協作
- 稀疏激活和計算效率
- 泛化和穩健性
- 可解釋性和透明性
- 最優的專家架構
- 與現有架構整合