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

如您關注最新發布的大型語言模型（LLM），可能會頻繁在標題中看到“MoE”這個詞。那么，“MoE”究竟是什么？為何眾多LLM都開始采用它呢？

在這份圖解中，我們將通過50多張可視化圖，詳細探索這個重要組件——混合專家模型（Mixture of Experts, MoE）！

在這份圖解中，我們將深入探討MoE的兩個主要組成部分，即專家（Experts）和路由器（Router），及其在典型LLM架構中的應用。

什么是混合專家模型（MoE）？

混合專家模型（MoE）是一種利用多個不同子模型（或“專家”）來提升LLM質量的技術。

MoE由兩個主要組件定義：

• 專家（Experts）- 每個前饋神經網絡（FFNN）層現在擁有一組“專家”，可以從中選擇一個子集。這些“專家”通常本身就是FFNN。
• 路由器（Router）或門控網絡（Gate Network）- 決定哪些token被發送到哪些專家。

在采用MoE的LLM的每個層中，我們都會發現（某種程度上專業化的）專家：

請注意，“專家”并非專攻“心理學”或“生物學”等特定領域。它們至多學習詞匯層面的句法信息：

更具體地說，它們的專長在于處理特定上下文中的特定token。

路由器（Router）（門控網絡）選擇最適合給定輸入的專家：

每個專家都不是一個完整的LLM，而是LLM架構中的一個子模型。

專家（The Experts）

為了探究專家代表什么以及它們如何工作，我們首先需要了解MoE旨在替代什么：稠密層（Dense Layers）。

稠密層（Dense Layers）

混合專家模型（MoE）都源于LLM的一個相對基本的功能，即前饋神經網絡（Feedforward Neural Network, FFNN）。

回想一下，標準的僅解碼器Transformer架構在層歸一化后應用FFNN：

FFNN支持模型利用注意力機制創建的上下文信息，對其進行進一步轉換以捕捉數據中更復雜的關系。

然而，FFNN的規模增長迅速。為了學習這些復雜關系，它通常會擴展其接收到的輸入：

稀疏層（Sparse Layers）

傳統Transformer中的FFNN被稱為稠密模型，因為它的所有參數（權重和偏置）都被激活。沒有任何部分被遺漏，所有參數都用于計算輸出。

如果我們仔細觀察稠密模型，會注意到輸入在某種程度上激活了所有參數：

相比之下，稀疏模型只激活其總參數的一部分，并且與混合專家模型密切相關。

舉例來說，我們可以將稠密模型切分成多個部分（即所謂的專家），重新訓練它，并在給定時間僅激活這些專家中的一個子集：

其基本思想是，每個專家在訓練過程中學習不同的信息。然后，在推理時，只使用特定的專家，因為它們與給定任務最相關。

當被問及一個問題時，我們可以選擇最適合給定任務的專家：

專家學習什么？

正如我們之前所見，專家學習的信息比整個領域更細粒度1。因此，稱它們為“專家”有時被認為是具有誤導性的。

ST-MoE論文中編碼器模型的專家專業化示例

然而，解碼器模型中的專家似乎沒有相同的專業化類型。但這并不意味著所有專家都是平等的。

Mixtral 8x7B論文中有一個很好的例子，其中每個token都用其第一個專家選擇進行著色。

此圖也表明，專家傾向于關注語法而非特定領域。

因此，盡管解碼器專家似乎沒有特定專長，但它們確實似乎持續用于某些類型的token。

專家的系統架構

雖然將專家可視化為稠密模型的隱藏層切片很方便，但它們本身通常是完整的FFNN：

由于大多數LLM具有多個解碼器塊，給定文本在生成之前會通過多個專家：

所選專家可能因token而異，從而導致采取不同的“路徑”：

如果我們更新解碼器塊的可視圖，它現在將包含更多的FFNN（每個專家一個），而不是一個：

解碼器塊現在擁有多個FFNN（每個都是一個“專家”），可以在推理期間使用。

路由機制

現在我們有了一組專家，模型如何知道該使用哪些專家呢？

在專家層之前，會添加一個路由器（Router）（也稱為門控網絡），它經過訓練，可以為給定的token選擇專家。

路由器（The Router）

路由器（Router）（或門控網絡）也是一個FFNN，用于根據特定輸入選擇專家。它輸出概率，并利用這些概率選擇最匹配的專家：

專家層返回所選專家的輸出乘以門控值（選擇概率）的結果。

路由器與專家（其中只有少數被選中）共同構成了MoE層：

給定的MoE層有兩種大?。合∈瑁╯parse）或稠密（dense）混合專家模型。

兩者都使用路由器來選擇專家，但稀疏MoE只選擇少數，而稠密MoE選擇所有專家，但可能采用不同的分布。

例如，給定一組token，MoE會將其token分配給所有專家，而稀疏MoE只會選擇少數專家。

對于當前的LLM，當您看到“MoE”時，它通常指的是稀疏MoE，因為它支持您使用專家子集。這在計算上成本較低，而這對于LLM來說是一個重要特性。

專家選擇

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

在其最基本的形式中，我們將輸入 (x) 與路由器權重矩陣 (W) 相乘：

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

路由器使用此概率分布為給定輸入選擇最匹配的專家。

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

讓我們把所有內容放在一起，探究輸入如何流經路由器和專家：

路由的復雜性

然而，這個簡單的函數常常導致路由器選擇相同的專家，因為某些專家可能比其他專家學習得更快：

這不僅會導致專家選擇分布不均，還會導致一些專家幾乎得不到訓練。這在訓練和推理過程中都會導致問題。

相反，我們希望在訓練和推理過程中專家之間具有同等的重要性，這被稱為負載均衡（load balancing）。從某種意義上說，這是為了防止對相同的專家過度擬合。

負載均衡

為了平衡專家的重要性，我們需要關注路由器，因為它是決定在給定時間選擇哪些專家的主要組件。

KeepTopK

一種對路由器進行負載均衡的方法是通過一個直接的擴展，稱為KeepTopK2。通過引入可訓練的（高斯）噪聲，我們可以防止總是選擇相同的專家：

然后，除了你想要激活的K個頂部專家（例如2個）之外，所有其他專家的權重都將被設置為-∞：

通過將這些權重設置為-∞，SoftMax對這些權重輸出的概率將為0：

KeepTopK策略是許多LLM仍在使用的策略，盡管存在許多有前景的替代方案。請注意，KeepTopK也可以在不添加額外噪聲的情況下使用。

Token選擇（Token Choice）

KeepTopK策略將每個token路由到幾個選定的專家。這種方法被稱為Token Choice3，它支持將給定token發送到一個專家（top-1路由）：

或發送到多個專家（top-k路由）：

一個主要的好處是它允許對專家各自的貢獻進行加權和整合。

輔助損失（Auxiliary Loss）

為了在訓練期間獲得更均勻的專家分布，網絡常規損失中增加了輔助損失（auxiliary loss）（也稱為負載均衡損失）。

它增加了一個約束，強制專家具有同等重要性。

這個輔助損失的第一個組成部分是將整個批次中每個專家的路由器值求和：

這為我們提供了每個專家的重要性分數（importance scores），它代表了無論輸入如何，給定專家被選中的可能性。

我們可以用它來計算變異系數（coefficient variation, CV），它告訴我們專家之間重要性分數的差異程度。

例如，如果重要性分數差異很大，CV將很高：

相反，如果所有專家具有相似的重要性分數，CV將很低（這是我們的目標）：

利用這個CV分數，我們可以在訓練期間更新輔助損失，使其旨在盡可能降低CV分數（從而給予每個專家同等的重要性）：

最后，輔助損失作為單獨的損失在訓練期間進行優化。

專家容量（Expert Capacity）

不平衡不僅存在于被選擇的專家中，還存在于發送給專家的token分布中。

例如，如果輸入token不成比例地發送給一個專家而不是另一個專家，那么這也可能導致訓練不足：

在這里，重要的不僅是使用了哪些專家，還在于使用了多少。

解決這個問題的一個方案是限制給定專家可以處理的token數量，即專家容量（Expert Capacity）?。當一個專家達到容量時，剩余的token將被發送到下一個專家：

如果兩個專家都已達到容量，則token將不會被任何專家處理，而是發送到下一個層。這被稱為token溢出（token overflow）。

使用Switch Transformer簡化MoE

首批處理MoE訓練不穩定性問題（如負載均衡）的基于Transformer的MoE模型之一是Switch Transformer?。它極大地簡化了架構和訓練過程，同時提高了訓練穩定性。

切換層（The Switching Layer）

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

路由器在計算選擇哪個專家時沒有特殊技巧，它將輸入與專家權重相乘后進行softmax（與我們之前做的相同）。

這種架構（top-1路由）假設路由器只需要1個專家就可以學習如何路由輸入。這與我們之前看到的假設相反，我們之前假設token應該被路由到多個專家（top-k路由）以學習路由行為。

容量因子（Capacity Factor）

容量因子是一個重要值，因為它決定了專家可以處理多少token。Switch Transformer通過引入直接影響專家容量的容量因子來擴展這一點。

專家容量的組成部分很簡單：

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

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

輔助損失（Auxiliary Loss）

為了進一步防止token丟失，引入了簡化版的輔助損失。

這種簡化的損失不再計算變異系數，而是權衡了調度token的比例與每個專家路由器概率的比例：

由于目標是實現token在$N$個專家間的均勻路由，我們希望向量$P$和$f$的值均為1/N。

α 是一個超參數，我們可以用它來微調訓練期間此損失的重要性。值過高會壓倒主損失函數，值過低則對負載均衡幾乎沒有作用。

視覺模型中的混合專家模型

MoE并非僅適用于語言模型的技術。視覺模型（如ViT）利用基于Transformer的架構，因此有潛力使用MoE。

快速回顧一下，ViT（Vision-Transformer）是一種將圖像分割成小塊（patches）的架構，這些小塊的處理方式類似于token?。

這些小塊（或token）隨后被投影成嵌入（附加位置嵌入），然后被送入常規編碼器：

當這些小塊進入編碼器時，它們的處理方式與token類似，這使得這種架構非常適合MoE。

Vision-MoE

Vision-MoE（V-MoE）是MoE在圖像模型中的首批實現之一?。它采用我們之前看到的ViT，并將編碼器中的稠密FFNN替換為稀疏MoE。

這使得ViT模型（通常比語言模型?。┠軌蛲ㄟ^添加專家實現大規模擴展。

每個專家都使用了預定義的小專家容量，以減少硬件限制，因為圖像通常有許多小塊。然而，低容量往往會導致小塊被丟棄（類似于token溢出）。

為了保持低容量，網絡會給小塊分配重要性分數，并優先處理那些分數較高的小塊，從而使溢出的小塊通常不那么重要。這被稱為批量優先級路由（Batch Priority Routing）。

因此，即使token的百分比降低，我們仍然應該看到重要的補丁被路由。

優先級路由通過關注最重要的補丁來處理更少的補丁。

從稀疏到軟MoE（From Sparse to Soft MoE）

在V-MoE中，優先級評分器幫助區分重要性更高和更低的補丁。然而，補丁被分配給每個專家，未處理補丁中的信息會丟失。

軟MoE（Soft-MoE）旨在通過混合補丁，從離散到軟補丁（token）分配?。

第一步，我們將輸入$x$（補丁嵌入）與一個可學習矩陣$\Phi$相乘。這為我們提供了路由器信息，它告訴我們某個token與給定專家的相關程度。

然后，通過對路由器信息矩陣（按列）進行softmax運算，我們更新每個補丁的嵌入。

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

從視覺上看，這就像所有的補丁都被混合在一起了。這些組合后的補丁然后被發送到每個專家。生成輸出后，它們再次與路由器矩陣相乘。

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

因此，我們得到的是“軟”補丁/token，而不是離散輸入。

Mixtral 8x7B中的激活參數與稀疏參數（Active vs. Sparse Parameters with Mixtral 8x7B）

MoE之所以引人注目，很大程度上在于其計算需求。由于在給定時間只使用專家子集，我們能夠訪問比實際使用更多的參數。

盡管一個MoE需要加載更多的參數（稀疏參數），但激活的參數較少，因為我們在推理時只使用部分專家（激活參數）。

換句話說，我們仍然需要將整個模型（包括所有專家）加載到設備上（稀疏參數），但當我們運行推理時，我們只需要使用一個子集（激活參數）。MoE模型需要更多的VRAM來加載所有專家，但在推理期間運行速度更快。

讓我們以Mixtral 8x7B?為例，探究稀疏參數與激活參數的數量。

這里，我們可以看到每個專家的大小是5.6B，而不是7B（盡管有8個專家）。

我們將需要加載8x5.6B (46.7B) 參數（以及所有共享參數），但我們只需要使用2x5.6B (12.8B) 參數進行推理。

結論

至此，我們的混合專家模型之旅告一段落！希望這篇文章能讓您更好地理解這項有趣技術的潛力?，F在幾乎所有模型系列都至少包含一個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).

本文轉載自?????Andy730?????，作者：常華?