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

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

在查看最新發布的LLMs時，你可能會在標題中看到“MoE”這個詞。那么，這個“MoE”到底代表什么？為什么現在有這么多 LLM 都在使用它呢？

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

640.png

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

目錄

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

第二部分：專家的角色

• 密集層（Dense Layers）
• 稀疏層（Sparse Layers）
• 專家能學到什么？
• 專家的架構

第三部分：路由機制

• 路由器（Router）
• 專家的選擇
• 路由的復雜性

第四部分：負載平衡與優化

• KeepTopK 策略

a.Token 選擇策略

b.輔助損失函數

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

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

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

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

什么是專家混合模型？

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

MoE 的兩個主要組成部分為：

專家（Experts）：每個前饋神經網絡層（FFNN）現在都有一組可以選擇的“專家”。這些“專家”通常本身也是前饋神經網絡（FFNN）。

路由器或門控網絡（Router 或 Gate Network）：決定哪些 token 被發送到哪些專家。

在具有 MoE 的 LLM 的每一層中，我們可以發現一些（相對專業化的）專家：

請注意，這些“專家”并不像人類領域中的“心理學”或“生物學”專家那樣在特定學科上表現出高度專業化。實際上，它們更多是從詞匯層次上學習句法信息：

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

路由器（或門控網絡）負責為每個輸入選擇最合適的專家：

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

專家

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

密集層（Dense Layers）

專家混合模型（MoE）最初是從大型語言模型（LLMs）中相對基礎的功能開始的，即前饋神經網絡（Feedforward Neural Network, FFNN）。

請記住，在標準的Decode-only Transformer 架構中，FFNN 通常應用于層歸一化（Layer Normalization）之后：

FFNN 允許模型利用注意力機制生成的上下文信息，并進一步轉化這些信息，從而捕捉數據中更復雜的關系。

然而，FFNN 的規模增長速度很快。為了學習這些復雜的關系，它通常需要對接收到的輸入進行擴展：

稀疏層（Sparse Layers）

在傳統的 Transformer 中，FFNN 被稱為密集模型（Dense Model），因為它的所有參數（包括權重和偏置項）都會被激活。所有參數都被用于計算輸出，沒有任何部分被遺棄。

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

640_6.png

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

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

其核心思想是，每個專家在訓練過程中學習不同的信息。而在推理時，僅使用與當前任務最相關的特定專家。

當面對一個問題時，我們可以選擇最適合該任務的專家：

專家學習了什么？

正如我們之前所見，專家學習的信息比整個領域的信息更加細粒度。因此，有時將它們稱為“專家”可能會產生誤導。

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

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

https://arxiv.org/pdf/2401.04088

這一可視化結果也表明，專家往往更關注句法，而不是某個特定領域的內容。

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

專家的架構

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

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

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

如果我們更新解碼器塊的可視化，現在它將包含多個 FFNN（每個 FFNN 對應一個“專家”）：

640_14.png

解碼器塊現在包含多個可以在推理時使用的 FFNN（即“專家”）。

路由機制（The Routing Mechanism）

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

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

路由器（The Router）

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

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

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

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

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

例如，面對一組 tokens，MoE 會將這些 tokens 分布到所有專家，而稀疏 MoE 則只會選擇少數幾個專家。

在目前的 LLM 中，當你看到“MoE”時，它通常指的是稀疏 MoE，因為稀疏 MoE 允許使用部分專家，從而減少計算開銷，這對于 LLM 來說是一個重要的特性。

專家的選擇

門控網絡可以說是 MoE 中最重要的組件，因為它不僅決定了推理時要選擇哪些專家，還決定了訓練時的選擇。

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

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

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

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

讓我們將所有內容整合起來，探索輸入如何在路由器和專家中流動：

路由的復雜性

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

這不僅會導致專家選擇的不均勻分布，還會導致某些專家幾乎沒有被訓練過。這會在訓練和推理過程中引發問題。

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

負載平衡與優化

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

KeepTopK

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

640_23.png

https://arxiv.org/pdf/1701.06538

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

640_24.png

將這些權重設為 -∞ 時，SoftMax 操作后的輸出概率將變為 0：

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

Token 選擇策略

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

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

這種策略的主要優點在于，它可以對各個專家的貢獻進行加權，并將其整合起來。

輔助損失（Auxiliary Loss）

為了在訓練期間實現專家的均勻分布，網絡的常規損失中加入了輔助損失（也稱為負載平衡損失）。

輔助損失增加了一個約束，強制專家在訓練過程中具有相同的重要性。

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

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

我們可以使用這些重要性得分計算變異系數（Coefficient of 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。它通過簡化架構和訓練過程，提高了訓練穩定性。

切換層（Switching Layer）

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

路由器在選擇專家時并沒有使用特殊的方法，只是對輸入與專家權重相乘的結果取 softmax（與之前的方法相同）。

這種架構（top-1 路由）假設每個 token 只需要一個專家來學習如何進行路由。這與我們之前討論的 top-k 路由（將 token 分配給多個專家）有所不同。

容量因子（Capacity Factor）

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

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

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

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

輔助損失（Auxiliary Loss）

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

在簡化版的輔助損失中，不再計算變異系數，而是將分配的 token 數量與每個專家的路由概率進行加權比較：

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

α 是一個超參數，用于在訓練過程中微調此損失的重要性。值過高會影響主要損失函數，而值過低則無法有效進行負載平衡。

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

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

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

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

當這些圖像塊進入編碼器時，它們會像 tokens 一樣被處理，這使得這種架構非常適合 MoE。

Vision-MoE

Vision-MoE（V-MoE）是圖像模型中第一個實現 MoE 的例子之一。它將 ViT 中的密集 FFNN 層替換為稀疏 MoE。

640_44.png

這種改進使得 ViT 模型（通常比語言模型?。┠軌蛲ㄟ^增加專家的數量來大幅擴展。

為了降低硬件限制，每個專家都設置了一個較小的預定義容量，因為圖像通常包含大量的圖像塊。然而，低容量往往會導致圖像塊被丟棄（類似于 token 溢出）。

為了保持容量較低，網絡會為每個圖像塊分配重要性得分，并優先處理這些得分較高的圖像塊，從而避免溢出圖像塊的丟失。這種方法被稱為批量優先路由（Batch Priority Routing）。

因此，即使 token 數量減少，我們仍然能夠看到重要的圖像塊被成功路由。

優先路由使得在處理較少的圖像塊時，仍能聚焦于最重要的圖像塊。

從稀疏 MoE 到軟 MoE

在 V-MoE 中，優先評分機制能夠區分出重要和不重要的圖像塊。然而，圖像塊被分配給每個專家后，未被處理的圖像塊中的信息就會丟失。

軟 MoE（Soft-MoE）旨在通過混合圖像塊，從離散的圖像塊（token）分配轉變為軟分配。

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

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

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

從視覺上看，這就像是所有圖像塊被混合。這些組合后的圖像塊被發送到每個專家。生成輸出后，它們再次與路由矩陣相乘。

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

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

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

MoE 的一個重要特點是其計算需求。由于在同一時刻只會使用部分專家，我們可以擁有比實際使用的更多的參數。

盡管給定的 MoE 擁有更多的參數（稀疏參數），但由于我們在推理時只使用部分專家（活躍參數），因此激活的參數較少。

換句話說，我們仍然需要將整個模型（包括所有專家）加載到設備中（稀疏參數），但在實際運行推理時，我們只需要使用部分參數（活躍參數）。MoE 模型需要更多的顯存來加載所有專家，但推理時運行速度更快。

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

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

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

結論

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