編者按： 對大語言模型進行擴展的過程中，如何在保持高性能的同時有效控制計算資源消耗？混合專家模型（MoE）作為一種新興的架構設計方案正在得到越來越多的關注，但它究竟是如何工作的？為什么越來越多的大語言模型都在采用這種架構？

這篇文章通過50多張精心繪制的示意圖，系統拆解了MoE的核心原理。從最基礎的專家網絡、路由機制，到負載均衡、容量控制等進階話題，每個概念都配有直觀的可視化講解。作者還介紹了如何將MoE應用于視覺模型，例如Vision-MoE和Soft-MoE。

作者 | Maarten Grootendorst

編譯 | 岳揚

在查看最新推出的大語言模型（LLMs）時，你可能會注意到標題中頻繁出現的“MoE”字樣。這個“MoE”究竟是什么，它為何會成為眾多 LLMs 的標配？

在這份可視化指南中，我們將借助超過 50 張圖表來詳細了解這一關鍵組件，逐步深入混合專家模型（Mixture of Experts, MoE）的工作原理！

在這份可視化指南里，我們將逐一剖析 MoE 的核心要素：專家網絡（Experts）和路由器（Router），并展示它們在常見的大語言模型（LLM）架構中是如何發揮作用的。

01 什么是混合專家模型？

混合專家模型（MoE）技術通過使用眾多不同的子模型（即“專家網絡”）來提升大語言模型（LLMs）的性能。

MoE 的兩大核心組成部分如下：

• 專家網絡- 當前，每個前饋神經網絡層（FFNN）都配備有一組“專家網絡”，可以根據需要選擇其中的子集。這些“專家網絡”通常就是 FFNN 本身。
• 路由器或門控網絡- 它負責決定哪些 tokens 應該被傳遞給哪些專家網絡。

在集成了 MoE 的 LLM 的每一個模型層中，我們都能找到（具有一定專業性的）專家網絡：

需要注意的是，“專家網絡”并非專注于某個特定領域（如“心理學”或“生物學”）的專家，它們更多的是在單詞層面上學習句法信息：

更準確地說，它們的專長在于在特定的上下文中處理特定的詞匯單元。

路由器（門控網絡）會為每個輸入挑選出最合適的專家網絡：

每個專家網絡并非一個完整的 LLM，而是 LLM 架構中的一個組成部分。

02 專家網絡

為了深入了解專家網絡的含義及其運作機制，我們先來考察混合專家模型（Mixture of Experts，MoE）試圖優化的結構 —— 密集層（dense layers）。

2.1 密集層（Dense Layers）

混合專家模型（Mixture of Experts，MoE）的概念源自大語言模型（LLMs）的一個基本組件，也就是前饋神經網絡（Feedforward Neural Network，FFNN）。

我們需要知道，在標準的僅包含解碼器的 Transformer 架構中，前饋神經網絡通常應用于層歸一化（layer normalization）之后：

前饋神經網絡利用注意力機制生成的上下文信息，進一步處理以捕捉數據中的復雜關系。

但是，前饋神經網絡的規模增長速度很快。為了掌握這些復雜的關系，它往往需要對輸入信息進行擴展處理：

2.2 稀疏層（Sparse Layers）

在傳統的 Transformer 中，FFNN 被稱作密集模型（dense model），因為它的所有參數（包括權重和偏置項）都會被激活。沒有任何參數被閑置，每一個都參與到輸出的計算中。

如果我們仔細觀察密集模型，會發現輸入數據在某種程度上激活了所有參數：

與此相對的是，稀疏模型（sparse models）只激活了其參數總數的一部分，這與混合專家模型（Mixture of Experts）密切相關。

打個比方，我們可以把密集模型分解成若干小塊（即所謂的“專家網絡”），重新訓練它們，并且在任何特定時刻只激活一部分“專家網絡”：

這種方法的核心理念是，每個“專家網絡”在訓練過程中會學習到不同的信息。而在推理階段，只有那些與當前任務最相關的“專家網絡”會被調用。

這樣在面對問題時，我們可以挑選出最適合處理該任務的“專家網絡”：

2.3 專家網絡學習了什么？

正如我們之前所見，專家網絡所學習的信息比泛泛地學習整個領域的所有信息更為精細1。因此，將它們稱為“專家”有時會被認為是一種誤導。

在 ST-MoE 論文中提到的編碼器模型的專業化

相比之下，解碼器模型中的專家網絡并沒有展現出同樣的專業化。但這并不意味著所有專家網絡都是等同的。

在 Mixtral 8x7B 論文[1]中提供了一個很好的案例，其中每個 token 都根據首次選擇的專家網絡進行了著色。

這張圖片也表明，專家網絡通常更關注句法結構，而不是特定領域的內容。

因此，盡管解碼器中的專家網絡沒有顯示出特定的專業領域，但它們似乎確實會一致地用于處理某些特定類型的 tokens 。

2.4 專家網絡架構

雖然將專家網絡可視化成被分解的密集模型中的隱藏層確實很有幫助，但實際上，他們本身往往就是一個完整的前饋神經網絡（FFNN）：

大語言模型（LLM）通常包含多個解碼器塊，這意味著一段文本在最終生成之前，會經過一系列專家網絡的處理：

每個 tokens 可能由不同的專家網絡處理，這就形成了多樣化的處理“路徑”：

如果我們重新構想解碼器塊的構造，會發現它現在集成了更多的前饋神經網絡（每個專家網絡對應一個）：

因此，解碼器塊在推理階段能夠調用多個前饋神經網絡（每個都被視為一個“專家網絡”）來完成任務。

03 路由機制

現在已經有了一組專家網絡，模型如何知道使用哪些專家網絡呢？

就在專家網絡之前，會添加一個路由器（也稱為門控網絡），它會被訓練來為給定的 token 選擇專家網絡。

3.1 路由器

路由器（或稱門控網絡）本身也是一個前饋神經網絡（FFNN），其作用是基于特定的輸入來挑選合適的專家網絡。它會輸出一系列概率值，用以確定最匹配的專家網絡：

專家層（expert layer）將所選專家網絡的輸出結果與門控值（即選擇概率）相乘后返回。

路由器與被選中的少數專家網絡共同組成了 MoE 層：

MoE 層有兩種類型，一種是稀疏混合專家（sparse mixture of experts），另一種是密集混合專家（dense mixture of experts）。

這兩種模式都通過路由器來挑選專家網絡，不同的是，稀疏 MoE 只挑選部分專家網絡，而密集 MoE 則會選擇所有專家網絡，但每個專家網絡的參與程度可能不同。

打個比方，面對一組 tokens ，MoE 會將其分配給所有專家網絡處理，而稀疏 MoE 則只挑選少數幾個專家網絡來處理。

在現有的 LLMs 中，提到的 “MoE” 通常指的是稀疏 MoE，因為它只使用部分專家網絡，大大降低了計算成本，對于 LLMs 來說，這一點至關重要。

3.2 專家網絡的選擇

門控網絡可以說是任何 MoE 最重要的組成部分，因為它不僅決定推理過程中選擇哪些專家網絡，還決定訓練過程中的選擇。

最基本的形式是：我們將輸入（x）乘以路由器權重矩陣（W）：

然后，我們對輸出應用 SoftMax 操作，為每個專家網絡創建一個概率分布G(x)：

路由器利用該概率分布來為給定輸入選擇最匹配的專家網絡。

最后，我們將每個路由器的輸出與每個選定專家網絡的輸出相乘，并將結果相加。

現在將整個流程組合起來，探索輸入是如何流經路由器和專家網絡的：

3.3 路由過程的復雜性

然而，這種簡單的功能往往使得路由器傾向于選擇相同的專家網絡，因為某些專家網絡的學習速度可能比其他專家網絡要快：

這樣不僅會造成專家網絡選擇的不均衡，還可能導致部分專家網絡幾乎得不到訓練。這種現象在訓練和推理階段都會引發問題。

為了解決這個問題，我們希望在訓練和推理過程中，所有專家網絡都能得到平等的重視，我們稱之為“專家網絡負載均衡”。這樣做一定程度上是為了避免模型對某些專家網絡產生過度的依賴，從而避免過擬合。

04 負載均衡

為了平衡專家網絡的重要性，我們需要重點關注路由器，因為它是決定在特定時間選擇哪些專家網絡的主要組成部分。

4.1 KeepTopK

對路由器進行負載平衡的一種方法是一個名為 KeepTopK[2]2 的簡易擴展策略。通過引入可訓練的高斯噪聲，避免了總是重復選擇同一批專家網絡的情況：

接下來，除了我們希望激活的前 k 個專家網絡（比如 2 個）之外，其他專家網絡的權重將被調整至負無窮大：

這樣做之后，這些權重在 SoftMax 函數作用下的輸出概率將會是 0：

盡管還存在許多有潛力的替代策略，但許多大語言模型（LLMs）仍然采用 KeepTopK 策略。需要注意的是，即使不添加額外的噪聲，KeepTopK 策略同樣可以使用。

4.1.1 Token 選擇策略

KeepTopK 策略能夠將每個 token 分配給若干選定的專家網絡。這種技術被稱為 Token Choice3，它允許每個 token 可以被路由至一個專家網絡（單專家網絡路由，即 top-1）：

或者同時發送至多個專家網絡（多專家網絡路由，即 top-k）：

這種方法的一個主要優點在于，它能夠對各個專家的貢獻進行評估權衡，并將其整合起來。

4.1.2 輔助損失（Auxiliary Loss）

為了在訓練過程中實現專家網絡的均衡分配，我們在網絡的常規損失基礎上增加了輔助損失（也稱作負載平衡損失）。

這一措施增加了一個約束條件，旨在強制確保每個專家網絡都擁有同等的重要性。

輔助損失的第一部分是對整個批次中每個專家網絡的路由器值（router values）進行求和：

這樣我們就得到了每個專家網絡的重要性得分（importance scores），它代表了在不考慮輸入的情況下，給定專家網絡被選中的可能性。

通過這些重要性得分（importance scores），我們可以計算出變異系數（CV），它表示不同專家網絡之間重要性得分的差異程度。

比如，如果專家網絡之間的重要性得分差異顯著，CV 值就會偏高：

反之，如果所有專家網絡的重要性得分相近，CV 值就會偏低（這正是我們希望達到的效果）：

利用 CV 值，我們可以在訓練過程中調整輔助損失，力求最大限度地降低 CV 值（以此確保每個專家網絡都具有相同的重要性）：

最終，這個輔助損失將作為一個獨立的 loss，參與訓練優化。

4.2 專家網絡容量

不平衡不僅體現在被選中的專家網絡上，還體現在分配給專家網絡的 token 分布上。

例如，如果輸入的 token 過多地被分配給某一個專家網絡，而不是均勻分配，可能會導致某些專家網絡得不到充分的訓練：

這里的關鍵不僅僅在于哪些專家網絡被使用，更在于它們被使用的程度。

為了解決這個問題，我們可以采用一種方法，即限制每個專家網絡能夠處理的 token 數量，我們稱之為專家網絡容量（Expert Capacity4）。當一個專家網絡的處理能力達到上限時，新來的 token 將會被轉交給下一個專家網絡：

如果所有相關專家網絡都已達到其處理容量，那么這些 tokens 將不會由任何專家網絡處理，而是直接傳遞到網絡的下一層。這種現象我們稱之為 token 溢出。

4.3 使用 Switch Transformer 簡化 MoE

在處理 MoE 模型訓練不穩定問題（例如負載平衡）的早期基于 Transformer 的模型中，Switch Transformer 是其中的一個。5它通過簡化架構和訓練流程，有效提升了訓練的穩定性。

4.3.1 Switching Layer

Switch Transformer 是一個 T5 模型（編碼器-解碼器結構），它用 Switching Layer 替換了傳統的 FFNN 層。Switching Layer 是一個稀疏的 MoE 層，它為每個 token 選擇一個專家網絡（Top-1 路由）。

路由器在計算選擇哪個專家網絡時沒有使用特殊技巧，只是將輸入乘以專家網絡的權重后取softmax（與我們之前做的一樣）。

這種架構（Top-1路由）假定路由器只需要一個專家網絡就能學會如何對輸入進行路由。這與我們之前看到的情況截然不同，之前我們認為 token 應該被路由到多個專家網絡（Top-k 路由）以學習路由行為。

4.3.2 容量因子（Capacity Factor）

容量因子是一個重要的參數，它決定了專家網絡可以處理多少個 token。Switch Transformer 在此基礎上進行了擴展，引入了一個直接影響專家網絡容量的容量因子。

專家網絡容量的組成部分非常簡單：

如果我們增加容量因子，則每個專家網絡將能夠處理更多的 token。

然而，如果容量因子過大，就會浪費計算資源。相反，如果容量因子太小，由于 token 溢出，模型性能將會下降。

4.3.3 輔助損失（Auxiliary Loss）

為了進一步防止 token 被丟棄，Switch Transformer 引入了一個簡化版的輔助損失（Auxiliary Loss）。

這個簡化版的輔助損失不再計算系數變化，而是根據每個專家網絡的路由器概率與分發的 token 數量進行加權比較：

由于我們的目標是讓 token 在 N 個專家網絡之間均勻路由，因此我們希望向量 P 和 f 的值都是 1/N。

α 是一個超參數，我們可以在訓練過程中用它來微調這個輔助損失的重要性。如果 α 值過高，會影響主要損失函數；如果 α 值過低，則無法有效進行負載平衡。

05 視覺模型中的 MoE

MoE不僅僅是一種僅適用于語言模型的技術。視覺模型（如 ViT）使用基于 transformer 的架構，因此也有潛力使用 MoE。

簡單回顧一下，ViT（Vision-Transformer）是一種將圖像分割成 patches 的架構，這些 patch 的處理方式與 token 類似。6

這些 patchs（或 tokens）隨后會被投影到嵌入向量中（并加上位置嵌入向量），然后被送入常規的編碼器中：

一旦這些 patches 進入編碼器，它們就像 tokens 一樣被處理，這使得這種架構非常適合用于MoE。

5.1 Vision-MoE

Vision-MoE（V-MoE）是第一次在圖像模型中實現 MoE 的嘗試。7它采用了我們之前看到的ViT，并將編碼器中的密集 FFNN 層替換為稀疏 MoE。

這使得 ViT 模型（通常模型規模小于語言模型），可以通過添加專家網絡來進行大規模擴展。

為了減少硬件限制，每個專家網絡都設置了較小的預定義容量，因為圖像通常包含許多 patch 。然而，低容量往往導致 patch 被丟棄（類似于 token 溢出）。

為了保持較低的專家網絡容量，網絡會為每個 patch 分配重要性分數，并優先處理這些得分較高的 patches，這樣使得溢出的 patches 一般就不那么重要了。這種方法被稱為批量優先級路由（Batch Priority Routing）。

因此，即使 token 的數量減少，我們仍然應該看到重要的 patch 被成功路由。

優先級路由（priority routing）可以集中處理最重要的 patches，從而減少處理的 patch 數量。

5.2 從稀疏 MoE 到軟 MoE

在 V-MoE 中，優先級評分器（priority scorer）能夠幫助區分重要和不重要的 patch。然而，patch 被分配給每個專家網絡后，未被處理的 patch 中的信息就會丟失。

Soft-MoE 的目標是通過混合 patch 來從離散的 patch（token）分配轉變為軟分配。8

在第一步中，我們將輸入 x（patch 嵌入）與一個可學習的矩陣 Φ 相乘。這將生成路由器信息（router information），告訴我們某個 token 與給定專家網絡的相關程度。

然后，通過對路由器信息矩陣（按列）進行 softmax 操作，從而更新每個 patch 的嵌入向量。

更新后的 patch 嵌入本質上是所有 patch 嵌入的加權平均。

從視覺上看，就好像所有的 patch 都被混合在一起。然后，這些合并后的 patch 會被發送給每個專家網絡。生成輸出后，它們再次與路由器矩陣相乘。

路由器矩陣在 token 層面上影響輸入，在專家網絡層面上影響輸出。

因此，我們得到的是經過處理的“軟” patch / token，而不是離散的輸入。

06 Mixtral 8x7B 的激活與稀疏參數對比

MoE 的一大亮點在于其計算需求。由于在給定時間內只會使用部分專家網絡，我們可以訪問的參數比我們正在使用的參數要多。

盡管給定的 MoE 有更多的參數需要加載（稀疏參數），但由于在推理過程中只使用部分專家網絡（活躍參數），因此實際激活的參數數量較少。

換句話說，我們仍然需要將整個模型（包括所有專家網絡）加載到設備上（稀疏參數），但是當我們進行推理時，只需要使用部分參數（活躍參數）。MoE 模型在加載所有專家網絡時需要更多的 VRAM，但在推理過程中運行得更快。

讓我們以 Mixtral 8x7B 為例，來探討稀疏參數與活躍參數的數量差異。9

在這里，我們可以看到每個專家網絡的參數量是 5.6B，而不是7B（有 8 個專家網絡）。

我們將不得不加載 8x5.6B（46.7B）的參數（以及所有共享參數），但在推理過程中我們只需要使用 2x5.6B（12.8B）的參數。

07 Conclusion

我們的混合專家模型（MoE）學習之旅到此結束！希望這篇文章能讓你更好地理解這種有趣技術的潛力。現在，幾乎所有的模型系列都至少包含一個 MoE 變體，感覺它已經成為了主流。

想要查看更多與大語言模型（LLMs）相關的可視化內容，請查看作者編寫的大語言模型書籍！

書籍相關網站[3]。你可以在亞馬遜[4]上訂購這本書。所有代碼都已上傳至GitHub[5]。

Resources

如果你想要更深入地了解 MoE，我建議你學習以下內容：

• 這兩篇論文[6][7]是關于 MoE最新創新技術的概述。
• 這篇關于 expert choice routing[8] 的論文已經引起了業界的一些關注。
• 一篇很棒的博客文章[9]，梳理了一些關于 MoE 的主要論文（及其發現）。
• 這篇博客文章[10] 講述了 MoE 的發展時間線。

1.Zoph, Barret, et al. "St-moe: Designing stable and transferable sparse expert models. arXiv 2022." arXiv preprint arXiv:2202.08906.

2.Shazeer, Noam, et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." arXiv preprint arXiv:1701.06538 (2017).

3.Shazeer, Noam, et al. "Outrageously large neural networks: The sparsely-gated mixture-of-experts layer." arXiv preprint arXiv:1701.06538 (2017).

4.Lepikhin, Dmitry, et al. "Gshard: Scaling giant models with conditional computation and automatic sharding." arXiv preprint arXiv:2006.16668 (2020).

5.Fedus, William, Barret Zoph, and Noam Shazeer. "Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity." Journal of Machine Learning Research 23.120 (2022): 1-39.

6.Dosovitskiy, Alexey. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).

7.Riquelme, Carlos, et al. "Scaling vision with sparse mixture of experts." Advances in Neural Information Processing Systems 34 (2021): 8583-8595.

8.Puigcerver, Joan, et al. "From sparse to soft mixtures of experts." arXiv preprint arXiv:2308.00951 (2023).

9.Jiang, Albert Q., et al. "Mixtral of experts." arXiv preprint arXiv:2401.04088 (2024).

Thanks for reading!

Hope you have enjoyed and learned new things from this blog!

About the authors

Maarten Grootendorst

Data Scientist | Psychologist | Writer | Open Source Developer (BERTopic, PolyFuzz, KeyBERT) | At the intersection of Artificial Intelligence and Psychology

END

本期互動內容 ??

?如果讓你用一個生活中的比喻來解釋 MoE 的工作原理，你會怎么比喻？歡迎在評論區分享你的想法！

??文中鏈接??

[1]??https://arxiv.org/pdf/2401.04088??

[2]??https://arxiv.org/pdf/1701.06538??

[3]??https://www.llm-book.com/??

[4]??https://www.amazon.com/Hands-Large-Language-Models-Understanding/dp/1098150961??

[5]??https://github.com/handsOnLLM/Hands-On-Large-Language-Models??

[6]??https://arxiv.org/pdf/2209.01667??

[7]??https://arxiv.org/pdf/2407.06204??

[8]??https://arxiv.org/pdf/2202.09368??

[9]??https://cameronrwolfe.substack.com/p/conditional-computation-the-birth??

[10]??https://brunomaga.github.io/Mixture-of-Experts??

原文鏈接：

??https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mixture-of-experts??