演算法、系統和應用,三個視角全面讀懂混合專家(MoE)

机器之心發表於2024-07-26
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 相關的關鍵性挑戰,並指出了有希望獲得重要成果的未來研究方向。下面簡要列出了這些挑戰和研究方向,更多詳情請檢視原論文。

  • 訓練穩定性和負載平衡

  • 可擴充套件性和通訊開銷

  • 專家的專業化和協作

  • 稀疏啟用和計算效率

  • 泛化和穩健性

  • 可解釋性和透明性

  • 最優的專家架構

  • 與現有框架整合

擴充套件閱讀:MoE 相關報導

基礎:

  • 30 年曆史回顧,Jeff Dean:我們整理了一份「稀疏專家模型」研究綜述

  • 為什麼基於 MoE 的大模型更值得關注?

  • 被 OpenAI、Mistral AI 帶火的 MoE 是怎麼回事?一文貫通專家混合架構部署

  • 吸引機器學習圈眼球的 MoE,會成為 NLP 與 CV 的未來嗎?

  • 手把手教你,從零開始實現一個稀疏混合專家架構語言模型(MoE)

前沿:

  • 單一作者論文,谷歌提出百萬專家 Mixture,超越密集前饋、稀疏 MoE

  • 微軟讓 MoE 長出多個頭,大幅提升專家啟用率

  • 將多模態大模型稀疏化,3B 模型 MoE-LLaVA 媲美 LLaVA-1.5-7B

  • MoE 與 Mamba 強強聯合,將狀態空間模型擴充套件到數百億引數

  • 開源大模型王座再易主,1320 億引數 DBRX 上線,基礎、微調模型都有

  • CVPR 2024 | 基於 MoE 的通用影像融合模型,新增 2.8% 引數完成多項任務

  • CVPR 2023 | 模組化 MoE 將成為視覺多工學習基礎模型

  • 谷歌 Gemini1.5 火速上線:MoE 架構,100 萬上下文

  • 蘋果大模型 MM1 殺入場:300 億引數、多模態、MoE 架構,超半數作者是華人

  • 8x7B MoE 與 Flash Attention 2 結合,不到 10 行程式碼實現快速推理

  • 打破 MoE 訓練效率與效能瓶頸,華為盤古稀疏大模型全新架構 LocMoE 出爐

  • 單個 4090 可推理,2000 億稀疏大模型「天工 MoE」開源

  • Mistral 開源 8X22B 大模型,OpenAI 更新 GPT-4 Turbo 視覺,都在欺負谷歌

  • 一條磁力連結席捲 AI 圈,87GB 種子直接開源 8x7B MoE 模型

  • 比 MoE 更有潛力?進化演算法融合模型的新路徑是否值得一試?

  • 清華髮布 SmartMoE:一鍵實現高效能 MoE 稀疏大模型分散式訓練

  • 一塊錢 100 萬 token,超強 MoE 模型開源,效能直逼 GPT-4-Turbo

相關文章