專家混合（MoE）詳解

隨著 Mixtral 8x7B（公告，模型卡）的釋出，一類 Transformer 模型成為開放 AI 社群最熱門的話題：專家混合（Mixture of Experts），簡稱 MoE。在這篇部落格文章中，我們將深入探討 MoE 的基本組成部分、它們的訓練方式以及在推理時需要考慮的權衡。

讓我們深入瞭解！

目錄

• 什麼是專家混合？
• MoE 的簡要歷史
• 什麼是稀疏性？
• MoE 的令牌負載均衡
• MoE 與 Transformers
• Switch Transformers
• 透過路由器 Z-損失穩定訓練
• 專家學到了什麼？
• 擴充套件專家數量如何影響預訓練？
• MoE 的微調
• 何時使用稀疏 MoE 與密集模型？
• 讓 MoE 加速執行
  • 專家並行
  • 容量因子和通訊成本
  • 服務技術
  • 高效訓練
• 開源 MoE
• 令人興奮的工作方向
• 一些資源

摘要

MoE 模型

• 比密集模型預訓練速度快得多
• 與相同引數數量的模型相比，推理速度更快
• 需要高 VRAM，因為所有專家都載入到記憶體中
• 在微調方面面臨諸多挑戰，但近期關於 MoE 指令微調的研究很有前景

讓我們深入瞭解！

什麼是專家混合（MoE）？

模型規模是提高模型質量最重要的因素之一。在固定的計算預算下，訓練一個更大的模型，步數更少，比訓練一個更小的模型，步數更多，效果更好。

專家混合（MoE）使得模型能夠以更少的計算量進行預訓練，這意味著在相同的計算預算下，你可以大幅擴充套件模型或資料集的大小，而使用密集模型則無法做到。具體而言，一個 MoE 模型在預訓練期間應該比其密集模型更快地達到相同的質量。

那麼，MoE 到底是什麼？在 Transformer 模型的背景下，MoE 由兩個主要元素組成：

• 稀疏 MoE 層用於代替密集的全連線網路（FFN）層。MoE 層包含一定數量的“專家”（例如 8 個），每個專家都是一個神經網路。實際上，這些專家是 FFN，但它們也可以是更復雜的網路，甚至是 MoE 本身，從而形成分層 MoE！
• 一個門控網路或路由器，它決定哪些令牌被髮送到哪個專家。例如，在下圖中，“More”令牌被髮送到第二個專家，而“Parameters”令牌被髮送到第一個網路。正如我們稍後將探討的，我們可以將一個令牌傳送到多個專家。如何將令牌路由到專家是使用 MoE 時的一個重要決策——路由器由可學習引數組成，並與網路的其餘部分同時進行預訓練。

總而言之，在 MoE 中，我們用 MoE 層取代了 Transformer 模型的每個 FFN 層，該 MoE 層由一個門控網路和一定數量的專家組成。

儘管 MoE 提供了高效預訓練和比密集模型更快的推理等優勢，但它們也帶來了挑戰：

• 訓練： MoE 顯著提高了計算效率，從而實現了高效預訓練，但它們在微調過程中歷史上一直難以泛化，導致過擬合。
• 推理： 儘管 MoE 可能擁有許多引數，但推理時只使用了其中一部分。這使得與相同引數數量的密集模型相比，推理速度快得多。然而，所有引數都需要載入到 RAM 中，因此記憶體需求很高。例如，對於像 Mixtral 8x7B 這樣的 MoE，我們需要足夠的 VRAM 來容納一個 47B 引數的密集模型。為什麼是 47B 引數而不是 8 x 7B = 56B？那是因為在 MoE 模型中，只有 FFN 層被視為獨立的專家，模型的其餘引數是共享的。同時，假設每個令牌只使用兩個專家，推理速度（FLOPs）就像使用一個 12B 模型（而不是 14B 模型），因為它計算 2x7B 矩陣乘法，但有一些層是共享的（稍後會詳細介紹）。

現在我們對 MoE 有了一個大致的瞭解，接下來讓我們看看導致其發明的研究進展。

MoE 的簡要歷史

MoE 的根源可以追溯到 1991 年的論文《區域性專家的自適應混合》。其思想類似於整合方法，即為由獨立網路組成的系統提供一個監督程式，每個網路處理訓練案例的不同子集。每個獨立的網路或專家專注於輸入空間的不同區域。如何選擇專家？一個門控網路決定每個專家的權重。在訓練期間，專家和門控都被訓練。

在 2010-2015 年間，兩個不同的研究領域為 MoE 的後續發展做出了貢獻：

• 專家作為元件：在傳統的 MoE 設定中，整個系統由一個門控網路和多個專家組成。MoE 作為整體模型已被應用於 SVM、高斯過程和其他方法。Eigen、Ranzato 和 Ilya 的工作探討了將 MoE 作為更深層網路的元件。這使得在多層網路中擁有 MoE 層成為可能，從而使模型能夠同時兼具大型和高效的特點。
• 條件計算：傳統網路透過每一層處理所有輸入資料。在此期間，Yoshua Bengio 研究了根據輸入令牌動態啟用或停用元件的方法。

這些工作促使人們在 NLP 領域探索專家混合。具體來說，Shazeer 等人（2017 年，其中的“等人”包括 Geoffrey Hinton 和 Jeff Dean，Google 的 Chuck Norris）透過引入稀疏性，將這一思想擴充套件到 137B 的 LSTM（當時事實上的 NLP 架構，由 Schmidhuber 建立），從而即使在大型模型中也能保持極快的推理速度。這項工作主要集中在翻譯領域，但也面臨許多挑戰，例如高通訊成本和訓練不穩定性。

MoE 使得訓練萬億引數模型成為可能，例如開源的 1.6T 引數 Switch Transformers 等。MoE 也已在計算機視覺領域進行了探索，但本部落格文章將重點關注 NLP 領域。

什麼是稀疏性？

稀疏性利用了條件計算的思想。在密集模型中，所有引數都用於所有輸入，而稀疏性允許我們只執行整個系統的一部分。

讓我們深入探討 Shazeer 對翻譯 MoE 的探索。條件計算的思想（網路的一部分根據每個示例動態啟用或停用）允許在不增加計算量的情況下擴充套件模型大小，因此，這使得每個 MoE 層中使用數千個專家成為可能。

這種設定帶來了一些挑戰。例如，儘管通常大批次大小對效能更好，但 MoE 中的批次大小會隨著資料流經活躍專家而有效減少。例如，如果我們的批處理輸入包含 10 個令牌，其中五個令牌可能進入一個專家，而另外五個令牌可能進入五個不同的專家，導致批次大小不均勻和利用率不足。下面的“讓 MoE 加速執行”部分將討論其他挑戰和解決方案。

我們如何解決這個問題？一個學習到的門控網路（G）決定將輸入的一部分發送到哪個專家（E）：

$$y = \sum_{i=1}^{n} G(x)_i E_i(x)$$

在這種設定中，所有專家都針對所有輸入執行——這是一種加權乘法。但是，如果 G 為 0 呢？如果是這種情況，就不需要計算相應的專家操作，從而節省了計算量。典型的門控函式是什麼？在最傳統的設定中，我們只使用一個帶有 Softmax 函式的簡單網路。該網路將學習將輸入傳送到哪個專家。

$$G_\sigma(x) = \text{Softmax}(x \cdot W_g)$$

Shazeer 的工作還探討了其他門控機制，例如 Noisy Top-k Gating。這種門控方法引入了一些（可調的）噪聲，然後保留了 Top-k 值。也就是說：

1. 我們添加了一些噪聲：

$$H(x)_i = (x \cdot W_{\text{g}})_i + \text{StandardNormal()} \cdot \text{Softplus}((x \cdot W_{\text{noise}})_i)$$

2. 我們只選擇 Top k：

$$\text{KeepTopK}(v, k)_i = \begin{cases} v_i & \text{if } v_i \text{ is in the top } k \text{ elements of } v, \\ -\infty & \text{otherwise.} \end{cases}$$

3. 我們應用 Softmax。

$$G(x) = \text{Softmax}(\text{KeepTopK}(H(x), k))$$

這種稀疏性引入了一些有趣的特性。透過使用足夠低的 k 值（例如一個或兩個），我們可以比啟用許多專家時更快地進行訓練和推理。為什麼不只選擇頂級專家？最初的推測是，路由到多個專家是必要的，以使門控網路學習如何路由到不同的專家，因此至少必須選擇兩個專家。Switch Transformers 部分重新審視了這一決策。

為什麼要新增噪聲？這是為了負載均衡！

MoE 的令牌負載均衡

如前所述，如果所有令牌都只發送給少數幾個熱門專家，那會使訓練效率低下。在正常的 MoE 訓練中，門控網路會收斂到主要啟用相同的幾個專家。這種現象會自我強化，因為受青睞的專家訓練得更快，因此被選擇得更多。為了緩解這個問題，添加了一個輔助損失來鼓勵所有專家獲得同等的重要性。這個損失確保所有專家接收到大致相同數量的訓練樣本。以下章節還將探討專家容量的概念，該概念引入了專家可以處理多少令牌的閾值。在 transformers 中，輔助損失透過 aux_loss 引數暴露。

MoE 與 Transformers

Transformer 模型很明顯地證明了引數數量的擴充套件能夠提高效能，因此 Google 在 GShard 中探索了這一點，它將 Transformer 模型擴充套件到超過 6000 億個引數，這並不令人驚訝。

GShard 在編碼器和解碼器中都用 MoE 層（使用 top-2 門控）替換了隔層 FFN 層。下圖展示了編碼器部分的這種設定。這種設定對於大規模計算非常有利：當我們擴充套件到多個裝置時，MoE 層在裝置之間共享，而所有其他層都進行復制。這將在“讓 MoE 加速執行”部分進一步討論。

為了在大規模下保持負載均衡和效率，GShard 作者除了引入類似於上一節中討論的輔助損失之外，還引入了一些更改：

• 隨機路由：在 top-2 設定中，我們總是選擇排名第一的專家，而第二個專家則以與其權重成比例的機率進行選擇。
• 專家容量：我們可以設定一個閾值，限制一個專家可以處理多少個令牌。如果兩個專家都達到了容量，則該令牌被視為溢位，並透過殘差連線傳送到下一層（或者在其他專案中完全丟棄）。這個概念將成為 MoE 最重要的概念之一。為什麼需要專家容量？因為所有張量形狀在編譯時都是靜態確定的，但我們無法提前知道有多少令牌會發送到每個專家，所以我們需要固定容量因子。

GShard 論文透過表達適用於 MoE 的平行計算模式做出了貢獻，但討論這一點超出了本部落格文章的範圍。

注意：在進行推理時，只有一部分專家會被觸發。同時，還存在一些共享計算，例如自注意力，它適用於所有令牌。這就是為什麼當我們談論一個 8 個專家的 47B 模型時，我們可以以 12B 密集模型的計算量執行。如果使用 top-2，將使用 14B 引數。但考慮到注意力操作（以及其他操作）是共享的，實際使用的引數數量是 12B。

Switch Transformers

儘管 MoE 展現出巨大的潛力，但它們在訓練和微調穩定性方面面臨挑戰。Switch Transformers 是一項非常激動人心的工作，深入探討了這些主題。作者甚至在 Hugging Face 上釋出了一個擁有 2048 個專家、1.6 萬億引數的 MoE 模型，你可以使用 transformers 庫執行它。Switch Transformers 在預訓練速度上比 T5-XXL 提升了 4 倍。

正如 GShard 中一樣，作者用 MoE 層替換了 FFN 層。Switch Transformers 論文提出了一種 Switch Transformer 層，它接收兩個輸入（兩個不同的令牌）並擁有四個專家。

與最初使用至少兩個專家的想法相反，Switch Transformers 使用了簡化的單專家策略。這種方法的效果是：

• 路由器計算量減少
• 每個專家的批次大小至少可以減半
• 通訊成本降低
• 質量得以保持

Switch Transformers 還探索了專家容量的概念。

$$\text{Expert Capacity} = \left(\frac{\text{tokens per batch}}{\text{number of experts}}\right) \times \text{capacity factor}$$

上述建議的容量將批處理中的令牌數量平均分配給專家數量。如果使用大於 1 的容量因子，則在令牌未完美平衡時提供了一個緩衝區。增加容量將導致更昂貴的裝置間通訊，因此需要權衡。特別是，Switch Transformers 在低容量因子（1-1.25）下表現良好。

Switch Transformer 的作者還重新審視並簡化了前面章節中提到的負載均衡損失。對於每個 Switch 層，輔助損失被新增到訓練過程中的總模型損失中。這種損失鼓勵均勻路由，並且可以使用超引數進行加權。

作者還嘗試了選擇性精度，例如用 bfloat16 訓練專家，而其餘計算則使用全精度。較低精度可以降低處理器之間的通訊成本、計算成本以及儲存張量的記憶體。最初的實驗中，專家和門控網路都用 bfloat16 訓練，但導致訓練更不穩定。這尤其歸因於路由器計算：由於路由器具有指數函式，因此更高的精度很重要。為了緩解不穩定性，路由器也使用了全精度。

這個筆記本展示瞭如何微調 Switch Transformers 進行摘要，但我們建議首先回顧微調部分。

Switch Transformers 使用編碼器-解碼器設定，其中他們製作了 T5 的 MoE 版本。GLaM 論文探索了透過訓練一個匹配 GPT-3 質量的模型來提高這些模型的規模，該模型僅使用 1/3 的能量（是的，由於訓練 MoE 所需的計算量更少，它們可以將碳足跡減少一個數量級）。作者專注於僅解碼器模型以及少量樣本和單樣本評估，而不是微調。他們使用了 Top-2 路由和更大的容量因子。此外，他們探索了容量因子作為一個可以在訓練和評估期間根據所需的計算量進行更改的指標。

透過路由器 Z-損失穩定訓練

前面討論的平衡損失可能導致不穩定性問題。我們可以使用許多方法來穩定稀疏模型，但代價是質量下降。例如，引入 dropout 可以提高穩定性，但會導致模型質量損失。另一方面，新增更多的乘性元件可以提高質量，但會降低穩定性。

路由器 Z-損失（由 ST-MoE 引入）透過懲罰進入門控網路的大 logits，顯著提高了訓練穩定性，而不會降低質量。由於這種損失鼓勵值保持較小的絕對幅度，因此舍入誤差會減少，這對於指數函式（如門控）可能產生相當大的影響。我們建議查閱論文以獲取詳細資訊。

專家學到了什麼？

ST-MoE 的作者觀察到，編碼器專家專注於一組令牌或淺層概念。例如，我們可能最終得到一個標點專家、一個專有名詞專家等。另一方面，解碼器專家的專業化程度較低。作者還在多語言設定中進行了訓練。儘管人們可能會想象每個專家都專注於一種語言，但情況恰恰相反：由於令牌路由和負載均衡，沒有任何一個專家專門針對任何給定語言。

擴充套件專家數量如何影響預訓練？

更多的專家可以提高樣本效率和加速，但這些增益是遞減的（尤其是在 256 或 512 之後），並且推理需要更多的 VRAM。在 Switch Transformers 中大規模研究的特性在小規模下也是一致的，即使每層只有 2、4 或 8 個專家。

MoE 的微調

Mixtral 支援 transformers 4.36.0 版本。您可以透過 pip install transformers==4.36.0 --upgrade 進行安裝。

密集模型和稀疏模型之間的過擬合動態非常不同。稀疏模型更容易過擬合，因此我們可以在專家內部探索更高的正則化（例如 dropout）（例如，我們可以為密集層設定一個 dropout 率，為稀疏層設定另一個更高的 dropout 率）。

一個問題是是否在微調時使用輔助損失。ST-MoE 的作者嘗試關閉輔助損失，即使有高達 11% 的令牌被丟棄，質量也沒有受到顯著影響。令牌丟棄可能是一種有助於防止過擬合的正則化形式。

Switch Transformers 觀察到，在固定的預訓練困惑度下，稀疏模型在下游任務中表現不如密集模型，尤其是在 SuperGLUE 等推理密集型任務上。另一方面，對於 TriviaQA 等知識密集型任務，稀疏模型表現得異常出色。作者還觀察到，少量專家有助於微調。另一個證實泛化問題的觀察是，模型在較小任務中表現較差，但在較大任務中表現良好。

可以嘗試凍結所有非專家權重。也就是說，我們只更新 MoE 層。這會導致效能大幅下降。我們可以嘗試相反的做法：只凍結 MoE 層中的引數，這幾乎與更新所有引數的效果一樣好。這有助於加快微調速度並減少記憶體。這可能有些反直覺，因為 80% 的引數都在 MoE 層中（在 ST-MoE 專案中）。他們對該架構的假設是，由於專家層只出現在每 1/4 層，並且每個令牌每層最多看到兩個專家，因此更新 MoE 引數影響的層數遠少於更新其他引數。

微調稀疏 MoE 時要考慮的最後一點是，它們具有不同的微調超引數設定——例如，稀疏模型往往更能受益於較小的批次大小和較高的學習率。

到目前為止，你可能對人們在微調 MoE 方面遇到的困難感到有些沮喪。令人興奮的是，一篇最近的論文《MoEs 遇見指令微調》（2023 年 7 月）進行了以下實驗：

• 單任務微調
• 多工指令微調
• 多工指令微調後進行單任務微調

當作者微調 MoE 和 T5 等效模型時，T5 等效模型表現更好。當作者微調 Flan T5 (T5 instruct 等效) MoE 時，MoE 模型的表現顯著更好。不僅如此，Flan-MoE 相對於 MoE 的提升比 Flan T5 相對於 T5 的提升更大，這表明 MoE 可能比密集模型從指令微調中受益更多。MoE 從更多工中受益更多。與之前建議關閉輔助損失函式的討論不同，該損失實際上可以防止過擬合。

何時使用稀疏 MoE 與密集模型？

專家在具有多臺機器的高吞吐量場景中非常有用。給定預訓練的固定計算預算，稀疏模型將更最佳化。對於 VRAM 很少的低吞吐量場景，密集模型將更好。

注意：不能直接比較稀疏模型和密集模型的引數數量，因為兩者代表著截然不同的事物。

讓 MoE 加速執行

最初的 MoE 工作將 MoE 層呈現為分支結構，導致計算速度緩慢，因為 GPU 不是為此設計的，並且導致網路頻寬成為瓶頸，因為裝置需要向其他裝置傳送資訊。本節將討論一些現有工作，以使這些模型的預訓練和推理更加實用。MoE 加速執行。

並行性

我們來簡要回顧一下並行性

• 資料並行： 相同的權重在所有核心上覆制，資料在核心之間劃分。
• 模型並行： 模型在核心之間劃分，資料在核心之間複製。
• 模型和資料並行： 我們可以將模型和資料劃分到不同的核心上。請注意，不同的核心處理不同的資料批次。
• 專家並行：專家被放置在不同的工作節點上。如果與資料並行結合，每個核心都有一個不同的專家，並且資料在所有核心上進行分割槽。

在專家並行中，專家被放置在不同的工作節點上，每個工作節點處理不同批次的訓練樣本。對於非 MoE 層，專家並行與資料並行行為相同。對於 MoE 層，序列中的令牌被髮送到所需專家所在的工作節點。

容量因子和通訊成本

增加容量因子（CF）可以提高質量，但會增加通訊成本和啟用記憶體。如果全對全通訊速度較慢，使用較小的容量因子會更好。一個好的起點是使用 top-2 路由，容量因子為 1.25，並且每個核心有一個專家。在評估期間，可以更改容量因子以減少計算量。

服務技術

您可以將 mistralai/Mixtral-8x7B-Instruct-v0.1 部署到推理端點。

MoE 的一個主要缺點是引數數量龐大。對於本地用例，可能需要使用較小的模型。讓我們快速討論一些有助於服務部署的技術：

• Switch Transformers 的作者進行了早期的蒸餾實驗。透過將 MoE 蒸餾回其密集對應物，他們可以保留 30-40% 的稀疏性收益。因此，蒸餾提供了更快預訓練和在生產中使用更小模型的好處。
• 最近的方法修改了路由，將整個句子或任務路由到一個專家，從而允許提取子網路進行服務。
• 專家聚合（MoE）：該技術合併了專家的權重，從而減少了推理時的引數數量。

更多高效訓練

FasterMoE (2022 年 3 月) 分析了 MoE 在高效分散式系統中的效能，並分析了不同並行策略的理論極限，以及傾斜專家受歡迎度的技術、減少延遲的精細通訊排程，以及根據最低延遲選擇專家的調整後的拓撲感知門控，從而實現了 17 倍的加速。

Megablocks (2022 年 11 月) 透過提供新的 GPU 核心來探索高效稀疏預訓練，這些核心可以處理 MoE 中存在的動態性。他們的方案從不丟棄令牌並能高效地對映到現代硬體，從而實現顯著的加速。訣竅是什麼？傳統的 MoE 使用批次矩陣乘法，這假設所有專家都具有相同的形狀和相同的令牌數量。相比之下，Megablocks 將 MoE 層表示為塊稀疏操作，可以適應不平衡的分配。

開源 MoE

現在有幾個開源專案可以訓練 MoE：

• Megablocks: https://github.com/stanford-futuredata/megablocks
• Fairseq: https://github.com/facebookresearch/fairseq/tree/main/examples/moe_lm
• OpenMoE: https://github.com/XueFuzhao/OpenMoE

在已釋出的開源 MoE 領域，您可以檢視：

• Switch Transformers (Google)：T5 基礎的 MoE 集合，從 8 個專家到 2048 個專家。最大的模型擁有 1.6 萬億引數。
• NLLB MoE (Meta)：NLLB 翻譯模型的 MoE 變體。
• OpenMoE：一項社群努力，釋出了基於 Llama 的 MoE。
• Mixtral 8x7B (Mistral)：一種高質量的 MoE，效能優於 Llama 2 70B，並且推理速度更快。還發布了指令微調模型。欲瞭解更多資訊，請閱讀釋出部落格文章。

令人興奮的工作方向

進一步實驗將稀疏 MoE 蒸餾回引數更少但質量相似的密集模型。

另一個領域將是 MoE 的量化。QMoE（2023 年 10 月）是朝著這個方向邁出的良好一步，它將 MoE 量化到每個引數少於 1 位元，從而將使用 3.2TB 加速器的 1.6T Switch Transformer 壓縮到僅 160GB。

所以，簡而言之，一些值得探索的有趣領域包括：

• 將 Mixtral 蒸餾成密集模型
• 探索專家模型合併技術及其對推理時間的影響
• 對 Mixtral 進行極端量化技術

一些資源

• 區域性專家的自適應混合（1991）
• 在深度專家混合中學習因子化表示（2013）
• 驚人龐大的神經網路：稀疏門控專家混合層（2017）
• GShard：使用條件計算和自動分片擴充套件巨型模型（2020 年 6 月）
• GLaM：使用專家混合高效擴充套件語言模型（2021 年 12 月）
• Switch Transformers：透過簡單高效的稀疏性擴充套件到萬億引數模型（2022 年 1 月）
• ST-MoE：設計穩定且可遷移的稀疏專家模型（2022 年 2 月）
• FasterMoE：大型動態預訓練模型的建模與訓練最佳化（2022 年 4 月）
• MegaBlocks：專家混合實現高效稀疏訓練（2022 年 11 月）
• 專家混合遇見指令微調：大型語言模型的致勝組合（2023 年 5 月）
• Mixtral-8x7B-v0.1, Mixtral-8x7B-Instruct-v0.1。

引用

@misc {sanseviero2023moe,
    author       = { Omar Sanseviero and
                     Lewis Tunstall and
                     Philipp Schmid and
                     Sourab Mangrulkar and
                     Younes Belkada and
                     Pedro Cuenca
                   },
    title        = { Mixture of Experts Explained },
    year         = 2023,
    url          = { https://huggingface.co/blog/moe },
    publisher    = { Hugging Face Blog }
}

Sanseviero, et al., "Mixture of Experts Explained", Hugging Face Blog, 2023.