圖解 MoE 模型

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

下面我們將通過 50 多個圖示來詳細探討這一重要組件：專家混合模型（MoE）！本指南將圍繞 MoE 的兩個主要組成部分——專家（Experts）和路由器（Router）——在典型的基於 LLM 架構中的應用展開討論。

目錄

第一部分：什麼是專家混合模型？
第二部分：專家的角色
第三部分：路由機制
第四部分：負載平衡與優化
第五部分：視覺模型中的專家混合模型
第六部分：帶有 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），因爲它的所有參數（包括權重和偏置項）都會被激活。所有參數都被用於計算輸出，沒有任何部分被遺棄。

如果我們仔細觀察密集模型，可以發現輸入在某種程度上激活了所有參數：相比之下，稀疏模型（Sparse Models）僅激活總參數中的一部分，這與專家混合模型（MoE）密切相關。

爲了說明這一點，我們可以將密集模型分解爲多個部分（稱爲專家），並重新訓練它。隨後，在某一時刻只激活部分專家：其核心思想是，每個專家在訓練過程中學習不同的信息。而在推理時，僅使用與當前任務最相關的特定專家。

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

專家學習了什麼？

正如我們之前所見，專家學習的信息比整個領域的信息更加細粒度。因此，有時將它們稱爲 “專家” 可能會產生誤導。然而，在解碼器模型中的專家似乎並沒有表現出相同類型的專業化。這並不意味着所有專家都是等同的。

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

https://arxiv.org/pdf/2401.04088

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

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

專家的架構

雖然將專家可視化爲密集模型的隱藏層被切分成若干部分很有幫助，但實際上它們通常本身就是完整的 FFNN。由於大多數 LLM 具有多個解碼器塊（Decoder Blocks），因此給定的文本會在生成之前通過多個專家：不同 tokens 被選中的專家可能各不相同，這導致了不同的 “路徑” 被選擇：如果我們更新解碼器塊的可視化，現在它將包含多個 FFNN（每個 FFNN 對應一個 “專家”）：解碼器塊現在包含多個可以在推理時使用的 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。通過引入可訓練的（高斯）噪聲，我們可以防止總是選擇相同的專家：

https://arxiv.org/pdf/1701.06538

然後，除希望激活的前 k 個專家（例如 2 個）以外的所有專家權重都將被設爲 -∞：將這些權重設爲 -∞ 時，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。這種改進使得 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 變體，這也預示着它可能會長期存在下去。

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/W8gWtjzGJBi9-t_--o_RHQ

猜你喜歡