本文僅做記錄，圖挺形象的。

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

在查看最新發(fā)布的LLMs時(shí)，你可能會(huì)在標(biāo)題中看到“MoE”這個(gè)詞。那么，這個(gè)“MoE”到底代表什么？為什么現(xiàn)在有這么多 LLM 都在使用它呢？

在本視覺指南中，我們將通過 50 多個(gè)圖示來詳細(xì)探討這一重要組件：專家混合模型（MoE）！

640.png

本指南將圍繞 MoE 的兩個(gè)主要組成部分——專家（Experts）和路由器（Router）——在典型的基于 LLM 架構(gòu)中的應(yīng)用展開討論。

目錄

第一部分：什么是專家混合模型？

第二部分：專家的角色

• 密集層（Dense Layers）
• 稀疏層（Sparse Layers）
• 專家能學(xué)到什么？
• 專家的架構(gòu)

第三部分：路由機(jī)制

• 路由器（Router）
• 專家的選擇
• 路由的復(fù)雜性

第四部分：負(fù)載平衡與優(yōu)化

• KeepTopK 策略

a.Token 選擇策略

b.輔助損失函數(shù)

• 專家容量
• 使用 Switch Transformer 簡(jiǎn)化 MoE
• 切換層（Switching Layer）
• 容量因子（Capacity Factor）
• 簡(jiǎn)化的輔助損失函數(shù)

第五部分：視覺模型中的專家混合模型

• Vision-MoE
• 從稀疏 MoE 到軟 MoE

第六部分：帶有 Mixtral 8x7B 的激活與稀疏參數(shù)

什么是專家混合模型？

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

MoE 的兩個(gè)主要組成部分為：

專家（Experts）：每個(gè)前饋神經(jīng)網(wǎng)絡(luò)層（FFNN）現(xiàn)在都有一組可以選擇的“專家”。這些“專家”通常本身也是前饋神經(jīng)網(wǎng)絡(luò)（FFNN）。

路由器或門控網(wǎng)絡(luò)（Router 或 Gate Network）：決定哪些 token 被發(fā)送到哪些專家。

在具有 MoE 的 LLM 的每一層中，我們可以發(fā)現(xiàn)一些（相對(duì)專業(yè)化的）專家：

請(qǐng)注意，這些“專家”并不像人類領(lǐng)域中的“心理學(xué)”或“生物學(xué)”專家那樣在特定學(xué)科上表現(xiàn)出高度專業(yè)化。實(shí)際上，它們更多是從詞匯層次上學(xué)習(xí)句法信息：

更具體地說，它們擅長(zhǎng)于在特定上下文中處理特定的 tokens。

路由器（或門控網(wǎng)絡(luò)）負(fù)責(zé)為每個(gè)輸入選擇最合適的專家：

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

專家

要理解專家的含義及其工作方式，首先需要了解 MoE 用來替代的內(nèi)容：密集層（Dense Layers）。

密集層（Dense Layers）

專家混合模型（MoE）最初是從大型語言模型（LLMs）中相對(duì)基礎(chǔ)的功能開始的，即前饋神經(jīng)網(wǎng)絡(luò)（Feedforward Neural Network, FFNN）。

請(qǐng)記住，在標(biāo)準(zhǔn)的Decode-only Transformer 架構(gòu)中，F(xiàn)FNN 通常應(yīng)用于層歸一化（Layer Normalization）之后：

FFNN 允許模型利用注意力機(jī)制生成的上下文信息，并進(jìn)一步轉(zhuǎn)化這些信息，從而捕捉數(shù)據(jù)中更復(fù)雜的關(guān)系。

然而，F(xiàn)FNN 的規(guī)模增長(zhǎng)速度很快。為了學(xué)習(xí)這些復(fù)雜的關(guān)系，它通常需要對(duì)接收到的輸入進(jìn)行擴(kuò)展：

稀疏層（Sparse Layers）

在傳統(tǒng)的 Transformer 中，F(xiàn)FNN 被稱為密集模型（Dense Model），因?yàn)樗乃袇?shù)（包括權(quán)重和偏置項(xiàng)）都會(huì)被激活。所有參數(shù)都被用于計(jì)算輸出，沒有任何部分被遺棄。

如果我們仔細(xì)觀察密集模型，可以發(fā)現(xiàn)輸入在某種程度上激活了所有參數(shù)：

640_6.png

相比之下，稀疏模型（Sparse Models）僅激活總參數(shù)中的一部分，這與專家混合模型（MoE）密切相關(guān)。

為了說明這一點(diǎn)，我們可以將密集模型分解為多個(gè)部分（稱為專家），并重新訓(xùn)練它。隨后，在某一時(shí)刻只激活部分專家：

其核心思想是，每個(gè)專家在訓(xùn)練過程中學(xué)習(xí)不同的信息。而在推理時(shí)，僅使用與當(dāng)前任務(wù)最相關(guān)的特定專家。

當(dāng)面對(duì)一個(gè)問題時(shí)，我們可以選擇最適合該任務(wù)的專家：

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

正如我們之前所見，專家學(xué)習(xí)的信息比整個(gè)領(lǐng)域的信息更加細(xì)粒度。因此，有時(shí)將它們稱為“專家”可能會(huì)產(chǎn)生誤導(dǎo)。

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

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

https://arxiv.org/pdf/2401.04088

這一可視化結(jié)果也表明，專家往往更關(guān)注句法，而不是某個(gè)特定領(lǐng)域的內(nèi)容。

因此，盡管解碼器中的專家似乎沒有特定的專業(yè)化，但它們確實(shí)在某些類型的 tokens 上表現(xiàn)得更加一致。

專家的架構(gòu)

雖然將專家可視化為密集模型的隱藏層被切分成若干部分很有幫助，但實(shí)際上它們通常本身就是完整的 FFNN。

由于大多數(shù) LLM 具有多個(gè)解碼器塊（Decoder Blocks），因此給定的文本會(huì)在生成之前通過多個(gè)專家：

不同 tokens 被選中的專家可能各不相同，這導(dǎo)致了不同的“路徑”被選擇：

如果我們更新解碼器塊的可視化，現(xiàn)在它將包含多個(gè) FFNN（每個(gè) FFNN 對(duì)應(yīng)一個(gè)“專家”）：

640_14.png

解碼器塊現(xiàn)在包含多個(gè)可以在推理時(shí)使用的 FFNN（即“專家”）。

路由機(jī)制（The Routing Mechanism）

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

在專家之前會(huì)加入一個(gè)路由器（也稱為門控網(wǎng)絡(luò)），它會(huì)被訓(xùn)練來選擇每個(gè) token 應(yīng)選擇的專家。

路由器（The Router）

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

專家層返回被選定專家的輸出，并乘以門控值（選擇概率）。

路由器和專家（其中僅選擇少部分）共同構(gòu)成了 MoE 層：

給定的 MoE 層有兩種類型：稀疏專家混合模型（Sparse Mixture of Experts）和密集專家混合模型（Dense Mixture of Experts）。

兩者都使用路由器來選擇專家，但稀疏 MoE 只選擇少數(shù)幾個(gè)專家，而密集 MoE 則選擇全部專家，但可能會(huì)以不同的分布進(jìn)行選擇。

例如，面對(duì)一組 tokens，MoE 會(huì)將這些 tokens 分布到所有專家，而稀疏 MoE 則只會(huì)選擇少數(shù)幾個(gè)專家。

在目前的 LLM 中，當(dāng)你看到“MoE”時(shí)，它通常指的是稀疏 MoE，因?yàn)橄∈?MoE 允許使用部分專家，從而減少計(jì)算開銷，這對(duì)于 LLM 來說是一個(gè)重要的特性。

專家的選擇

門控網(wǎng)絡(luò)可以說是 MoE 中最重要的組件，因?yàn)樗粌H決定了推理時(shí)要選擇哪些專家，還決定了訓(xùn)練時(shí)的選擇。

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

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

路由器利用這個(gè)概率分布來為給定的輸入選擇最匹配的專家。

最后，我們將每個(gè)路由器的輸出與各自選擇的專家輸出相乘，并將結(jié)果相加：

讓我們將所有內(nèi)容整合起來，探索輸入如何在路由器和專家中流動(dòng)：

路由的復(fù)雜性

然而，這個(gè)簡(jiǎn)單的功能往往會(huì)導(dǎo)致路由器總是選擇相同的專家，因?yàn)槟承＜铱赡鼙绕渌麑＜覍W(xué)習(xí)得更快：

這不僅會(huì)導(dǎo)致專家選擇的不均勻分布，還會(huì)導(dǎo)致某些專家?guī)缀鯖]有被訓(xùn)練過。這會(huì)在訓(xùn)練和推理過程中引發(fā)問題。

因此，我們希望在訓(xùn)練和推理期間，各個(gè)專家的使用具有同等的重要性，這就是所謂的負(fù)載平衡。某種程度上，這是為了防止模型在同一組專家上過擬合。

負(fù)載平衡與優(yōu)化

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

KeepTopK

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

640_23.png

https://arxiv.org/pdf/1701.06538

然后，除希望激活的前 k 個(gè)專家（例如 2 個(gè)）以外的所有專家權(quán)重都將被設(shè)為 -∞：

640_24.png

將這些權(quán)重設(shè)為 -∞ 時(shí)，SoftMax 操作后的輸出概率將變?yōu)?0：

盡管存在許多有前景的替代方案，許多 LLM 仍然使用 KeepTopK 策略。需要注意的是，KeepTopK 也可以不使用額外的噪聲。

Token 選擇策略

KeepTopK 策略會(huì)將每個(gè) token 路由到若干選定的專家。這種方法被稱為 Token 選擇策略（Token Choice），它允許一個(gè)給定的 token 被路由到一個(gè)專家（top-1 路由）：

或者被路由到多個(gè)專家（top-k 路由）：

這種策略的主要優(yōu)點(diǎn)在于，它可以對(duì)各個(gè)專家的貢獻(xiàn)進(jìn)行加權(quán)，并將其整合起來。

輔助損失（Auxiliary Loss）

為了在訓(xùn)練期間實(shí)現(xiàn)專家的均勻分布，網(wǎng)絡(luò)的常規(guī)損失中加入了輔助損失（也稱為負(fù)載平衡損失）。

輔助損失增加了一個(gè)約束，強(qiáng)制專家在訓(xùn)練過程中具有相同的重要性。

輔助損失的第一個(gè)組成部分是對(duì)整個(gè)批次中每個(gè)專家的路由值進(jìn)行求和：

這為我們提供了每個(gè)專家的重要性得分，即在不考慮輸入的情況下，給定專家被選中的概率。

我們可以使用這些重要性得分計(jì)算變異系數(shù)（Coefficient of Variation, CV），它表示各個(gè)專家的重要性得分之間的差異程度。

例如，如果重要性得分之間的差異較大，那么 CV 值就會(huì)較高：

相反，如果所有專家的得分都相似，則 CV 值較低（這是我們期望的情況）：

通過使用這個(gè) CV 得分，我們可以在訓(xùn)練過程中更新輔助損失，使其盡可能降低 CV 得分（從而使每個(gè)專家具有相同的重要性）：

最后，輔助損失將作為一個(gè)獨(dú)立的損失項(xiàng)，參與訓(xùn)練優(yōu)化。

專家容量（Expert Capacity）

專家的不平衡不僅體現(xiàn)在被選中的專家上，還體現(xiàn)在分配給這些專家的 token 分布上。

例如，如果輸入 token 被不成比例地分配到某些專家上，而不是平均分配，這可能導(dǎo)致某些專家的訓(xùn)練不足：

這里不僅要考慮使用了哪些專家，還需要關(guān)注這些專家被使用的頻率。

解決這個(gè)問題的方法是限制每個(gè)專家能夠處理的 token 數(shù)量，即專家容量（Expert Capacity）。當(dāng)一個(gè)專家達(dá)到其容量時(shí)，多余的 token 將被分配到下一個(gè)專家：

如果兩個(gè)專家都達(dá)到了其容量，token 將不會(huì)被任何專家處理，而是直接傳遞到下一層。這種情況被稱為 token 溢出（token overflow）。

使用 Switch Transformer 簡(jiǎn)化 MoE

第一個(gè)解決 MoE 訓(xùn)練不穩(wěn)定性問題（如負(fù)載平衡）的基于 Transformer 的 MoE 模型是 Switch Transformer。它通過簡(jiǎn)化架構(gòu)和訓(xùn)練過程，提高了訓(xùn)練穩(wěn)定性。

切換層（Switching Layer）

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

路由器在選擇專家時(shí)并沒有使用特殊的方法，只是對(duì)輸入與專家權(quán)重相乘的結(jié)果取 softmax（與之前的方法相同）。

這種架構(gòu)（top-1 路由）假設(shè)每個(gè) token 只需要一個(gè)專家來學(xué)習(xí)如何進(jìn)行路由。這與我們之前討論的 top-k 路由（將 token 分配給多個(gè)專家）有所不同。

容量因子（Capacity Factor）

容量因子是一個(gè)重要參數(shù)，它決定了每個(gè)專家可以處理的 token 數(shù)量。Switch Transformer 通過引入直接影響專家容量的容量因子擴(kuò)展了這一概念。

專家容量的組成部分非常簡(jiǎn)單：

如果我們?cè)黾尤萘恳蜃樱瑒t每個(gè)專家能夠處理更多的 token。

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

輔助損失（Auxiliary Loss）

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

在簡(jiǎn)化版的輔助損失中，不再計(jì)算變異系數(shù)，而是將分配的 token 數(shù)量與每個(gè)專家的路由概率進(jìn)行加權(quán)比較：

由于目標(biāo)是希望將 token 在 N 個(gè)專家中均勻分配，因此我們希望向量 P 和 f 的值為 1/N。

α 是一個(gè)超參數(shù)，用于在訓(xùn)練過程中微調(diào)此損失的重要性。值過高會(huì)影響主要損失函數(shù)，而值過低則無法有效進(jìn)行負(fù)載平衡。

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

MoE 并不僅限于語言模型。視覺模型（如 ViT）使用 Transformer 架構(gòu)，因此也有潛力使用 MoE。

快速回顧一下，ViT（Vision Transformer）是一種將圖像分割為若干塊，并將其作為 tokens 處理的架構(gòu)。

這些圖像塊（或 tokens）會(huì)被投射到嵌入向量（加上額外的位置嵌入向量）中，然后輸入到常規(guī)編碼器中：

當(dāng)這些圖像塊進(jìn)入編碼器時(shí)，它們會(huì)像 tokens 一樣被處理，這使得這種架構(gòu)非常適合 MoE。

Vision-MoE

Vision-MoE（V-MoE）是圖像模型中第一個(gè)實(shí)現(xiàn) MoE 的例子之一。它將 ViT 中的密集 FFNN 層替換為稀疏 MoE。

640_44.png

這種改進(jìn)使得 ViT 模型（通常比語言模型?。┠軌蛲ㄟ^增加專家的數(shù)量來大幅擴(kuò)展。

為了降低硬件限制，每個(gè)專家都設(shè)置了一個(gè)較小的預(yù)定義容量，因?yàn)閳D像通常包含大量的圖像塊。然而，低容量往往會(huì)導(dǎo)致圖像塊被丟棄（類似于 token 溢出）。

為了保持容量較低，網(wǎng)絡(luò)會(huì)為每個(gè)圖像塊分配重要性得分，并優(yōu)先處理這些得分較高的圖像塊，從而避免溢出圖像塊的丟失。這種方法被稱為批量?jī)?yōu)先路由（Batch Priority Routing）。

因此，即使 token 數(shù)量減少，我們?nèi)匀荒軌蚩吹街匾膱D像塊被成功路由。

優(yōu)先路由使得在處理較少的圖像塊時(shí)，仍能聚焦于最重要的圖像塊。

從稀疏 MoE 到軟 MoE

在 V-MoE 中，優(yōu)先評(píng)分機(jī)制能夠區(qū)分出重要和不重要的圖像塊。然而，圖像塊被分配給每個(gè)專家后，未被處理的圖像塊中的信息就會(huì)丟失。

軟 MoE（Soft-MoE）旨在通過混合圖像塊，從離散的圖像塊（token）分配轉(zhuǎn)變?yōu)檐浄峙洹?/span>

第一步，我們將輸入 x（圖像塊嵌入）與一個(gè)可學(xué)習(xí)矩陣 Φ 相乘。這將生成路由信息，它告訴我們某個(gè) token 與某個(gè)專家的相關(guān)程度。

然后對(duì)路由信息矩陣進(jìn)行 softmax 操作（在列上），從而更新每個(gè)圖像塊的嵌入向量。

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

從視覺上看，這就像是所有圖像塊被混合。這些組合后的圖像塊被發(fā)送到每個(gè)專家。生成輸出后，它們?cè)俅闻c路由矩陣相乘。

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

因此，我們獲得了“軟”圖像塊/token，這些 token 被處理，而不是離散輸入。

Mixtral 8x7B 的激活與稀疏參數(shù)對(duì)比

MoE 的一個(gè)重要特點(diǎn)是其計(jì)算需求。由于在同一時(shí)刻只會(huì)使用部分專家，我們可以擁有比實(shí)際使用的更多的參數(shù)。

盡管給定的 MoE 擁有更多的參數(shù)（稀疏參數(shù)），但由于我們?cè)谕评頃r(shí)只使用部分專家（活躍參數(shù)），因此激活的參數(shù)較少。

換句話說，我們?nèi)匀恍枰獙⒄麄€(gè)模型（包括所有專家）加載到設(shè)備中（稀疏參數(shù)），但在實(shí)際運(yùn)行推理時(shí)，我們只需要使用部分參數(shù)（活躍參數(shù)）。MoE 模型需要更多的顯存來加載所有專家，但推理時(shí)運(yùn)行速度更快。

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

在此例中，我們可以看到每個(gè)專家的參數(shù)量為 5.6B，而不是 7B（盡管一共有 8 個(gè)專家）。

我們需要加載 8x5.6B（46.7B）的參數(shù)（加上所有共享參數(shù)），但推理時(shí)只需要使用 2x5.6B（12.8B）的參數(shù)。

結(jié)論

以上就是我們對(duì)專家混合模型（MoE）的探索之旅！希望這篇文章能幫助你更好地理解這一有趣技術(shù)的潛力。如今幾乎所有的模型架構(gòu)中都有 MoE 變體，這也預(yù)示著它可能會(huì)長(zhǎng)期存在下去。