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

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

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

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

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

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

MoE由兩個主要組件定義：

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

在采用MoE的LLM的每個層中，我們都會發(fā)現(xiàn)（某種程度上專業(yè)化的）專家：

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

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

路由器（Router）（門控網(wǎng)絡(luò)）選擇最適合給定輸入的專家：

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

專家（The Experts）

為了探究專家代表什么以及它們?nèi)绾喂ぷ鳎覀兪紫刃枰私釳oE旨在替代什么：稠密層（Dense Layers）。

稠密層（Dense Layers）

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

回想一下，標(biāo)準(zhǔn)的僅解碼器Transformer架構(gòu)在層歸一化后應(yīng)用FFNN：

FFNN支持模型利用注意力機(jī)制創(chuàng)建的上下文信息，對其進(jìn)行進(jìn)一步轉(zhuǎn)換以捕捉數(shù)據(jù)中更復(fù)雜的關(guān)系。

然而，F(xiàn)FNN的規(guī)模增長迅速。為了學(xué)習(xí)這些復(fù)雜關(guān)系，它通常會擴(kuò)展其接收到的輸入：

稀疏層（Sparse Layers）

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

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

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

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

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

當(dāng)被問及一個問題時，我們可以選擇最適合給定任務(wù)的專家：

專家學(xué)習(xí)什么？

正如我們之前所見，專家學(xué)習(xí)的信息比整個領(lǐng)域更細(xì)粒度1。因此，稱它們?yōu)椤皩＜摇庇袝r被認(rèn)為是具有誤導(dǎo)性的。

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

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

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

此圖也表明，專家傾向于關(guān)注語法而非特定領(lǐng)域。

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

專家的系統(tǒng)架構(gòu)

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

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

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

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

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

路由機(jī)制

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

在專家層之前，會添加一個路由器（Router）（也稱為門控網(wǎng)絡(luò)），它經(jīng)過訓(xùn)練，可以為給定的token選擇專家。

路由器（The Router）

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

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

路由器與專家（其中只有少數(shù)被選中）共同構(gòu)成了MoE層：

給定的MoE層有兩種大小：稀疏（sparse）或稠密（dense）混合專家模型。

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

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

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

專家選擇

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

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

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

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

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

讓我們把所有內(nèi)容放在一起，探究輸入如何流經(jīng)路由器和專家：

路由的復(fù)雜性

然而，這個簡單的函數(shù)常常導(dǎo)致路由器選擇相同的專家，因為某些專家可能比其他專家學(xué)習(xí)得更快：

這不僅會導(dǎo)致專家選擇分布不均，還會導(dǎo)致一些專家?guī)缀醯貌坏接?xùn)練。這在訓(xùn)練和推理過程中都會導(dǎo)致問題。

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

負(fù)載均衡

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

KeepTopK

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

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

通過將這些權(quán)重設(shè)置為-∞，SoftMax對這些權(quán)重輸出的概率將為0：

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

Token選擇（Token Choice）

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

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

一個主要的好處是它允許對專家各自的貢獻(xiàn)進(jìn)行加權(quán)和整合。

輔助損失（Auxiliary Loss）

為了在訓(xùn)練期間獲得更均勻的專家分布，網(wǎng)絡(luò)常規(guī)損失中增加了輔助損失（auxiliary loss）（也稱為負(fù)載均衡損失）。

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

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

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

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

例如，如果重要性分?jǐn)?shù)差異很大，CV將很高：

相反，如果所有專家具有相似的重要性分?jǐn)?shù)，CV將很低（這是我們的目標(biāo)）：

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

最后，輔助損失作為單獨的損失在訓(xùn)練期間進(jìn)行優(yōu)化。

專家容量（Expert Capacity）

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

例如，如果輸入token不成比例地發(fā)送給一個專家而不是另一個專家，那么這也可能導(dǎo)致訓(xùn)練不足：

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

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

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

使用Switch Transformer簡化MoE

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

切換層（The Switching Layer）

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

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

這種架構(gòu)（top-1路由）假設(shè)路由器只需要1個專家就可以學(xué)習(xí)如何路由輸入。這與我們之前看到的假設(shè)相反，我們之前假設(shè)token應(yīng)該被路由到多個專家（top-k路由）以學(xué)習(xí)路由行為。

容量因子（Capacity Factor）

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

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

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

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

輔助損失（Auxiliary Loss）

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

這種簡化的損失不再計算變異系數(shù)，而是權(quán)衡了調(diào)度token的比例與每個專家路由器概率的比例：

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

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

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

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

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

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

當(dāng)這些小塊進(jìn)入編碼器時，它們的處理方式與token類似，這使得這種架構(gòu)非常適合MoE。

Vision-MoE

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

這使得ViT模型（通常比語言模型小）能夠通過添加專家實現(xiàn)大規(guī)模擴(kuò)展。

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

為了保持低容量，網(wǎng)絡(luò)會給小塊分配重要性分?jǐn)?shù)，并優(yōu)先處理那些分?jǐn)?shù)較高的小塊，從而使溢出的小塊通常不那么重要。這被稱為批量優(yōu)先級路由（Batch Priority Routing）。

因此，即使token的百分比降低，我們?nèi)匀粦?yīng)該看到重要的補丁被路由。

優(yōu)先級路由通過關(guān)注最重要的補丁來處理更少的補丁。

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

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

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

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

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

更新后的補丁嵌入本質(zhì)上是所有補丁嵌入的加權(quán)平均值。

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

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

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

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

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

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

換句話說，我們?nèi)匀恍枰獙⒄麄€模型（包括所有專家）加載到設(shè)備上（稀疏參數(shù)），但當(dāng)我們運行推理時，我們只需要使用一個子集（激活參數(shù)）。MoE模型需要更多的VRAM來加載所有專家，但在推理期間運行速度更快。

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

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

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

結(jié)論

至此，我們的混合專家模型之旅告一段落！希望這篇文章能讓您更好地理解這項有趣技術(shù)的潛力。現(xiàn)在幾乎所有模型系列都至少包含一個MoE變體，感覺它將持續(xù)存在。

參考文獻(xiàn)

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).

本文轉(zhuǎn)載自?????Andy730?????，作者：常華?