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

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

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

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

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

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

MoE由兩個(gè)主要組件定義：

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

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

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

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

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

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

專家（The Experts）

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

稠密層（Dense Layers）

混合專家模型（MoE）都源于LLM的一個(gè)相對基本的功能，即前饋神經(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被稱為稠密模型，因?yàn)樗乃袇?shù)（權(quán)重和偏置）都被激活。沒有任何部分被遺漏，所有參數(shù)都用于計(jì)算輸出。

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

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

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

其基本思想是，每個(gè)專家在訓(xùn)練過程中學(xué)習(xí)不同的信息。然后，在推理時(shí)，只使用特定的專家，因?yàn)樗鼈兣c給定任務(wù)最相關(guān)。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

路由機(jī)制

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

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

路由器（The Router）

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

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

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

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

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

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

對于當(dāng)前的LLM，當(dāng)您看到“MoE”時(shí)，它通常指的是稀疏MoE，因?yàn)樗С帜褂脤＜易蛹＿@在計(jì)算上成本較低，而這對于LLM來說是一個(gè)重要特性。

專家選擇

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

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

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

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

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

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

路由的復(fù)雜性

然而，這個(gè)簡單的函數(shù)常常導(dǎo)致路由器選擇相同的專家，因?yàn)槟承＜铱赡鼙绕渌麑＜覍W(xué)習(xí)得更快：

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

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

負(fù)載均衡

為了平衡專家的重要性，我們需要關(guān)注路由器，因?yàn)樗菦Q定在給定時(shí)間選擇哪些專家的主要組件。

KeepTopK

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

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

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

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

Token選擇（Token Choice）

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

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

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

輔助損失（Auxiliary Loss）

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

它增加了一個(gè)約束，強(qiáng)制專家具有同等重要性。

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

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

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

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

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

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

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

專家容量（Expert Capacity）

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

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

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

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

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

使用Switch Transformer簡化MoE

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

切換層（The Switching Layer）

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

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

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

容量因子（Capacity Factor）

容量因子是一個(gè)重要值，因?yàn)樗鼪Q定了專家可以處理多少token。Switch Transformer通過引入直接影響專家容量的容量因子來擴(kuò)展這一點(diǎn)。

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

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

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

輔助損失（Auxiliary Loss）

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

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

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

α 是一個(gè)超參數(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)入編碼器時(shí)，它們的處理方式與token類似，這使得這種架構(gòu)非常適合MoE。

Vision-MoE

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

盡管一個(gè)MoE需要加載更多的參數(shù)（稀疏參數(shù)），但激活的參數(shù)較少，因?yàn)槲覀冊谕评頃r(shí)只使用部分專家（激活參數(shù)）。

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

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

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

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

結(jié)論

至此，我們的混合專家模型之旅告一段落！希望這篇文章能讓您更好地理解這項(xiàng)有趣技術(shù)的潛力。現(xiàn)在幾乎所有模型系列都至少包含一個(gè)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?????，作者：常華?