專家混合 (MoE) 是 LLM 中常用的一種技術，旨在提高其效率和準確性。這種方法的工作原理是將復雜的任務劃分為更小、更易于管理的子任務，每個子任務都由專門的迷你模型或「專家」處理。

早些時候，有人爆料 GPT-4 是采用了由 8 個專家模型組成的集成系統(tǒng)。近日，Mistral AI 發(fā)布的 Mixtral 8x7B 同樣采用這種架構，實現(xiàn)了非常不錯的性能（傳送門：一條磁力鏈接席卷 AI 圈，87GB 種子直接開源 8x7B MoE 模型）。

OpenAI 和 Mistral AI 的兩波推力，讓 MoE 一時間成為開放人工智能社區(qū)最熱門的話題 。

本文將介紹 MoE 的構建模塊、訓練方法以及在使用它們進行推理時需要考慮的權衡因素。

混合專家架構簡稱 MoE，它的特點如下：

• 與稠密模型相比，預訓練速度更快；
• 與參數(shù)數(shù)量相同的模型相比，推理速度更快；
• 因為要把所有專家模型都加載在內存中，所以需要大量顯存；
• 在微調方面面臨許多挑戰(zhàn)，但最近在 MoE 指令微調方面的工作很有希望能解決這些問題。

什么是混合專家（MoE）？

模型的規(guī)模是決定模型質量的最重要因素之一。在預算固定的情況下，用較少的步驟訓練較大的模型要優(yōu)于用較多的步驟訓練較小的模型。

MoE 可以用較少的計算開銷對模型進行預訓練，這意味著可以用與稠密模型相同的計算開銷，大幅擴大模型或數(shù)據(jù)集的規(guī)模。特別是，在預訓練過程中，MoE 模型能更快地達到與稠密模型相同的性能。

那么，究竟什么是 MoE？從 Transformer 模型的角度來說，MoE 包含兩個主要元素：

• 使用稀疏的 MoE 層代替稠密的前饋網絡（FFN）層。MoE 層中有一定數(shù)量（如 8 個）的 「專家」，每個專家都是一個神經網絡。實際上，專家可以是 FFN，也可以是更復雜的網絡，甚至是 MoE 本身，這樣就會形成有多層 MoE 的 MoE。
• 使用門控網絡或者路由來決定將哪個 token 發(fā)送給哪個專家。例如，在下圖中，「More」被發(fā)送給第二個專家，而 「Parameters」被發(fā)送到第一個專家。如何將 token 通過路由發(fā)送給專家是使用 MoE 時需要計劃的重點之一，這是因為路由同樣由學習到的參數(shù)組成，并與網絡的其他部分同時進行預訓練。

Switch Transformers 中的 MoE 層 (https://arxiv.org/abs/2101.03961)

簡而言之，在 MoE 中，一個 MoE 層取代了 transformer 中的每個 FFN 層，MoE 層由一個門控網絡和一定數(shù)量的專家網絡組成。

雖然與稠密模型相比，MoE 具有高效預訓練和快速推理等優(yōu)點，但也面臨著一些挑戰(zhàn)：

• 訓練：MoE 能夠大大提高預訓練的計算效率，但在微調過程中難以實現(xiàn)泛化，從而導致過擬合。
• 推理：雖然 MoE 可能有很多參數(shù)，但在推理過程中只使用其中的一部分。與參數(shù)數(shù)量相同的稠密模型相比，推理速度要快得多。然而，所有參數(shù)都需要加載到 RAM 中，因此對內存的要求很高。例如，給定一個像 Mixtral 8x7B 這樣的 MoE，需要有足夠的 顯存 來容納一個 47B 參數(shù)的稠密模型。為什么是 47B 參數(shù)而不是 8 x 7B = 56B？這是因為在 MoE 模型中，只有 FFN 層被視為獨立的專家網絡，其余模型的參數(shù)都是共享的。同時，假設每個 token 只傳入兩個專家網絡，那么推理速度（FLOPs）就像使用 12B 模型（而不是 14B 模型），因為它進行的是 2x7B 的矩陣乘法運算，同時有些層是共享的（后文將詳細介紹）。

在對 MoE 有了一個大致的介紹后，一起來了解一下 MoE 的發(fā)展軌跡。

MoE 簡史

MoE 起源于 1991 年的論文《Adaptive Mixture of Local Experts》（https://www.cs.toronto.edu/~hinton/absps/jjnh91.pdf）。該論文的理念與集合方法類似，都是為由不同網絡組成的系統(tǒng)提供監(jiān)督程序，每個網絡處理不同的訓練集子集。每個獨立的網絡或者說專家擅長于輸入空間的不同區(qū)域。至于如何選擇專家這個問題，是由門控網絡來決定每個專家網絡的權重。在訓練過程中，專家網絡和門控網絡都要接受訓練。

2010-2015 年間，兩個研究領域的發(fā)展共同促成了 MoE 后來的進步：

• 專家網絡組件化：在傳統(tǒng)的 MoE 中，整個系統(tǒng)由一個門控網絡和多個專家網絡組成。在 SVM、高斯過程和其他方法中，研究者們對作為整體模型的 MoE 進行了探索。Eigen、Ranzato 和 Ilya 的研究，將 MoE 作為更深層網絡的組成部分進行了探索。MoE 能夠作為多層網絡中的組成部分，使得模型既龐大又高效成為可能。
• 條件計算：傳統(tǒng)網絡每一層都會對所有輸入數(shù)據(jù)進行處理。接著，Yoshua Bengio 研究了根據(jù)輸入的 token 動態(tài)激活或停用組件網絡的方法。

這些工作促使研究者們在 NLP 的背景下探索混合專家模型。具體來說，Shazeer 及 Geoffrey Hinton 、Jeff Dean，谷歌的 Chuck Norris 通過對引入稀疏網絡，將這一想法擴展到了 137B LSTM（https://arxiv.org/abs/1701.06538），從而即使在大規(guī)模下也能保持極快的推理速度。這項工作的重點目標是機器翻譯，同時也存在一些缺點，如通信成本高和訓練不穩(wěn)定。

論文 Outrageously Large Neural Network 中的 MoE 層

MoE 可以訓練數(shù)萬億級參數(shù)的模型，例如開源的 1.6T 參數(shù)的 Switch Transformer。計算機視覺領域也在探索 MoE，但這里還是先側重講解 NLP 領域。

什么是稀疏化？

稀疏化一詞來源于條件計算理念。在稠密模型中，所有參數(shù)都發(fā)揮作用，而稀疏化可以只運行整個系統(tǒng)的某些部分。

前文提到 Shazeer 對機器翻譯中的 MoE 進行了探索。條件計算（網絡中只有某些部分處于活動狀態(tài)）使得在不增加計算量的情況下能夠擴大模型的規(guī)模，因此，每層 MoE 都可以包含成千上萬的專家網絡。

但是這種設計帶來了一些挑戰(zhàn)。例如，雖然擴大 batch size 通常更有利于提高模型性能，但 MOE 中的 batch size 會隨著數(shù)據(jù)在激活狀態(tài)的專家網絡中的流動而縮小。例如，如果 batch size 為 10 個 token，其中 5 個 token 可能在一個專家網絡中結束，而另外 5 個 token 可能在 5 個不同的專家網絡中結束，從而導致 batch size 大小不均和利用率不足的情況。

如何解決這個問題？方法之一是讓學習后的門控網絡（G）決定向哪些專家網絡（E）傳達輸入信息：

在這種情況下，全部的專家網絡會對所有輸入進行運算 — 用一種加權乘法的方式。但是，如果 G 為 0 時會怎樣呢？這種情況下，就不需要經過相應的專家網絡運算，節(jié)省了計算開銷。那么典型的門控函數(shù)又是怎么樣的呢？在最傳統(tǒng)的設置中，只需使用一個帶有 softmax 函數(shù)的簡單網絡。該網絡將學習向哪個專家傳遞輸入數(shù)據(jù)。

Shazeer 的研究還探索了其他門控機制，如 Top-K 噪聲門控。這種門控方法會引入一些（可調整的）噪聲，然后保留前 K 個。也就是說：

1. 添加一些噪音

2. 只選擇前 k 個

3. 用 softmax 激活。

這種稀疏化操作帶來了一些有趣的特性。通過使用足夠小的 k（例如一個或兩個），研究者發(fā)現(xiàn)訓練和推理速度比激活許多專家網絡的設置更快。那為什么不直接保留 top=1 時的專家網絡呢？研究者最初的猜想是，需要路由到一個以上的專家，才能讓門控網絡學習如何路由到不同的專家，因此至少要選擇兩個專家。Switch Transformer 的部分將重新討論了這一決策。

為什么要添加噪聲？這是為了負載平衡。

MoE 的 token 負載平衡

如前所述，如果所有的 token 都只發(fā)送給少數(shù)幾個受歡迎的專家網絡，訓練效率將變得低下。通常 MoE 訓練中，門控網絡會收斂到頻繁激活同樣的幾個專家網絡。這種情況會隨著訓練的進行趨于顯著，因為受青睞的專家網絡會更快地得到訓練，從而更容易被選中。為了緩解這種情況，可以添加一個輔助損失，以鼓勵給予所有專家同等的重要性。該損失可確保所有專家獲得大致相同數(shù)量的訓練樣本。后文將探討專家能力的概念，即專家可以處理 token 的數(shù)量閾值。在 transformer 中，輔助損失通過 aux_loss 參數(shù)顯示。

MoE 和 Transformer

Transformer 是一個增加參數(shù)數(shù)量可以提高性能的非常明顯的例子，因此，谷歌順理成章地在 GShard 上沿用這種思維，將其中 Transformers 的參數(shù)量擴展到 6000 億以上。

GShard 在編碼器和解碼器中都使用了 top-2 門控技術，用 MoE 層取代了 FFN 層。下圖顯示了編碼器的部分情況。

GShard 中的 MoE Transformer 編碼器

這種設置對大規(guī)模計算非常有利：當擴展到多個設備時，MoE 層將在設備間共享，而所有其他層都將被復制。具體細節(jié)在下文的 「高效訓練 MoE」中進一步討論。

為了保持均衡負載和規(guī)模效率，GShard 的作者除了采用與上一節(jié)討論的類似的輔助損耗外，還引入了一些改進：

• 隨機路由：在 top-2 的設置中，研究者總是選擇排名第一的專家，同時第二位專家以與其權重成正比的概率被選中。
• 專家容量：還可以設定一個閾值，即一個專家可以處理多少 token。如果兩位專家的處理能力都已達到閾值，則 token 會被視為溢出，并通過剩余連接發(fā)送到下一層（或在其他項目中完全丟棄）。這一概念將成為 MoEs 最重要的概念之一。為什么需要設置專家容量呢？因為所有張量形狀都是在編譯時靜態(tài)確定的，但無法提前知道每個專家將獲得多少 token 作為輸入，所以需要確定容量。

GShard 的貢獻在于為 MoEs 確定了并行計算模式。需要注意的是：推理時，只有部分專家網絡會被觸發(fā)。與此同時，還有一些需要共享數(shù)據(jù)的步驟，比如自注意力機制，適用于所有 token。這就是為什么對于一個由 8 個專家組成的 47B 模型時，GShard 可以用一個 12B 的稠密模型進行計算。如果使用 top-2，則需要 14B 個參數(shù)。但考慮到注意力操作是共享的，實際使用的參數(shù)數(shù)量為 12B。

Switch Transformer

盡管 MoE 前景廣闊，但它們在訓練和微調不穩(wěn)定性方面仍有弊端。Switch Transformer （https://arxiv.org/abs/2101.03961）的出現(xiàn)意義非凡，作者甚至在 Hugging Face 發(fā)布了一個擁有 2048 個專家的 1.6 萬億個參數(shù)的模型（https://huggingface.co/google/switch-c-2048）。與 T5-XXL 相比，Switch Transformer 的預訓練速度提高了 4 倍。

Switch Transformer 論文中的 Switch Transformer 層

正如在 GShard 中作者用 MoE 層取代了 FFN 層一樣，Switch Transformer 論文提出了一個 Switch Transformer 層，該層接收兩個輸入（兩個不同的 token），有四個專家網絡。

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

• 減少了路由的計算量
• 每個專家的 batch size 至少減半
• 通信成本降低
• 模型的質量仍能得到保證

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

上面推薦的容量計算方法是將 batch size 中的 token 數(shù)量平均分配給專家。如果使用大于 1 的容量系數(shù)，就能在 token 不完全平衡時提供緩沖。增加容量會導致設備間的通信費用增加，因此需要注意 trade-off。尤其是在容量系數(shù)較低（1-1.25）的情況下，Switch Transformer 的表現(xiàn)尤為出色。

Switch Transformer 的作者還重新審視并簡化了章節(jié)中提到的負載平衡損耗。對于每個交換層，輔助損耗會在訓練過程中添加到模型總損耗中。這種損耗會鼓勵模型傾向于統(tǒng)一路由，并可使用超參數(shù)加權。

Switch Transformer 的作者還嘗試了選擇性精度，例如使用 bfloat16 的參數(shù)精度來訓練專家，而在其他計算中使用全精度。較低的精度可以降低處理器之間的通信成本、計算成本和用于存儲張量的內存。在最初的實驗中，專家和門控網絡都使用 bfloat16 進行了訓練，但訓練結果并不穩(wěn)定。這主要是由于路由網絡參與計算造成：由于路由網絡具有指數(shù)函數(shù)，因此更高的精度非常重要。為了減少不穩(wěn)定性，路由也使用了全精度。

使用選擇性精度不會降低質量，并能更快地訓練模型。

微調部分中，Switch Transformer 使用了編碼器 - 解碼器設置，并將 T5 與 MoE 對應。GLaM （https://arxiv.org/abs/2112.06905）探討了如何用三分之一的計算開銷訓練出與 GPT-3 質量相匹配的模型，從而擴大模型的規(guī)模。GLaM 作者的研究重點是純解碼器模型以及小樣本及零樣本的評估結果，并不是微調。他們使用了 Top-2 路由和更大的容量系數(shù)。此外，他們還將容量因子作為一個指標進行了調研，在訓練和驗證過程中，可以根據(jù)想要使用的計算量進行改變。

穩(wěn)定訓練基于 Z-loss 損失函數(shù)的路由網絡

上文討論的平衡損失可能會導致不穩(wěn)定問題。但是可以使用許多方法以犧牲質量為代價來穩(wěn)定稀疏模型。例如，引入 dropout 可以提高穩(wěn)定性，但有損于模型的性能質量。另一方面，添加更多的乘法分量可以提高模型的性能質量，但會降低穩(wěn)定性。

ST-MoE（https://arxiv.org/abs/2202.08906）中引入的路由器 z 損失，通過懲罰進入門控網絡的較大的對數(shù)值，在不降低質量的情況下顯著提高了訓練的穩(wěn)定性。由于這種損失會鼓勵數(shù)值變小，因此舍入誤差會減少，這對門控中的指數(shù)函數(shù)會產生很大影響。

專家網絡能學到什么？

根據(jù) ST-MoE 作者的觀察，編碼專家網絡專注于組別 token 或淺層概念。例如，標點符號專家、專有名詞專家等。另一方面，解碼專家網絡的專業(yè)化程度較低。其作者還在多語言環(huán)境中進行了訓練。不同于想象之中，每個專家網絡都精通一種語言：由于 token 被路由分發(fā)和負載平衡的原因，沒有任何一種語言的專家網絡是專業(yè)的。

ST-MoE 論文中的表格，顯示不同的 token 組分別被發(fā)送給了哪個專家。

擴大專家數(shù)量對預訓練的影響

專家數(shù)量越多，采樣效率越高，速度越快，但收益也會遞減（尤其是在 256 或 512 的量級之后），推理需要更多的顯存。Switch Transformers 中研究的大規(guī)模特性在小規(guī)模中也是一致的，即使每層有 2、4 或 8 個專家也是如此。

微調 MoE

Mixtral 支持 4.36.0 版本的 transformers 。使用 pip install "transformers==4.36.0 --upgrade 進行更新。

稠密模型和稀疏模型的過擬合動態(tài)截然不同。稀疏模型更容易出現(xiàn)過擬合，因此可以在專家本身內部探索更高的正則化（例如，可以為稠密層設定一個 dropout，為稀疏層設定另一個更高的 dropout）。

還有一個需要決策的問題是：是否使用輔助損失（auxiliary loss ）進行微調。ST-MoE 的作者曾嘗試關閉輔助損耗，結果發(fā)現(xiàn)即使有高達 11% 的 token 被丟棄，質量也沒有受到明顯影響。token dropping 可能是正則化的一種形式，有助于防止過擬合。

Switch Transformer 的作者觀察到，在固定的預訓練困惑度下，稀疏模型在下游任務中的表現(xiàn)不如稠密模型，尤其是在推理任務繁重的任務中，如 SuperGLUE。另一方面，在 TriviaQA 等知識密集型語料集中，稀疏模型的表現(xiàn)卻好得出乎意料。作者還觀察到，專家較少有助于微調。對泛化問題的另一個觀察結果是，模型在較小的任務中表現(xiàn)較差，但在較大的任務中表現(xiàn)良好。

在小型任務（左圖）中，可以看到明顯的過擬合，稀疏模型在驗證集中的表現(xiàn)要差得多。而在大型任務中（右圖），MoE 的表現(xiàn)很好。圖片來自 ST-MoE 論文。

嘗試凍結所有非專家權重導致了性能的大幅下降，不過這在意料之中，因為 MoE 層占據(jù)了大部分網絡。嘗試相反的方法：只凍結 MoE 層的參數(shù)后，結果發(fā)現(xiàn)效果幾乎與更新所有參數(shù)一樣好。這個發(fā)現(xiàn)有助于加快微調速度并減少內存占用。

通過只凍結 MoE 層，可以在保證質量的同時加快訓練速度。本圖來自 ST-MoE 論文。

在微調稀疏化 MoE 時需要考慮的最后一個問題是，不同的 MoE 有不同的微調超參數(shù) — 例如，稀疏模型往往更受益于較小的批大小和較高的學習率。

稀疏模型的微調質量隨著學習率的增大和 batch size 的降低而提高。本圖來自 ST-MoE 論文。

研究者們一直在努力對 MoE 進行微調，這個過程是充滿曲折的。最近的一篇論文《MoEs Meets Instruction Tuning》行了這樣的實驗：

• 單任務微調
• 多任務指令微調
• 多任務指令調整后進行單任務微調

論文中，當作者微調 MoE 和 T5 后，T5 等效模擬輸出的效果更好。當作者對 Flan T5 和 MoE 進行微調后，MoE 的表現(xiàn)明顯更好。不僅如此，F(xiàn)lan-MoE 相對于 MoE 的改進幅度大于 Flan T5 相對于 T5 的改進幅度，這表明 MoE 可能比密集模型更受益于指令調優(yōu)。任務變多會讓 MoE 受益更多。與之前建議關閉輔助損失函數(shù)的結論不同，損失函數(shù)實際上可以防止過擬合。

與稠密模型相比，稀疏模型更受益于 instruct-tuning。圖片來自《MoEs Meets Instruction Tuning》論文

何時使用稀疏 MoE，何時使用稠密 MoE？

專家模型適用于使用多臺機器的高吞吐量場景。在預訓練運算預算固定的情況下，稀疏模型將更為理想。對于顯存較少的低吞吐量場景，稠密模型會更好。

注意：不能直接比較稀疏模型和稠密模型的參數(shù)量，因為二者所代表的意義明顯不同。

高效訓練 MoE

最初的 MoE 工作將 MoE 層作為一個分支設置，導致模型計算速度緩慢，因為 GPU 并非為此而設計，而且由于設備需要向其他設備發(fā)送信息，網絡帶寬成為瓶頸。本節(jié)將討論現(xiàn)有的一些工作，以使這些模型的預訓練和推理更加實用。

并行計算

并行計算的種類：

• 數(shù)據(jù)并行：相同的權重在所有內核上復制，數(shù)據(jù)在內核上分割。
• 模型并行：模型在不同內核間分區(qū)，數(shù)據(jù)在不同內核間復制。
• 模型和數(shù)據(jù)并行：在不同內核間劃分模型和數(shù)據(jù)。請注意，不同的內核處理不同批次的數(shù)據(jù)。
• 專家并行：將專家放在不同的工作站上。如果與數(shù)據(jù)并行相結合，則每個內核都有不同的專家，數(shù)據(jù)在所有內核之間進行分割

在專家并行模式下，專家被置于不同的工作站上，每個工作站采集不同批次的訓練樣本。對于非 MoE 層，專家并行的行為與數(shù)據(jù)并行相同。對于 MoE 層，序列中的 token 會被發(fā)送到所需的專家所在的工作站。

Switch Transformers 論文中的插圖，顯示了數(shù)據(jù)和模型是如何通過不同的并行技術在內核上分割的。

容量系數(shù)和通信成本

提高容量因子 (CF) 可以提高模型質量，但會增加通信成本和內存開銷。如果全連接通信（all-to-all communications）速度較慢，使用較小的容量因子會更好。這里提供一個可以參考的配置：使用容量系數(shù)為 1.25 的 top-2 路由機制，每個內核留有一名專家。在驗證過程中，可以改變容量系數(shù)以減少計算量。

服務器

mistralai/Mixtral-8x7B-Instruct-v0.1 可以被部署到推理終端上。

MoE 的一大缺點是參數(shù)較多。本地用例可能希望使用更小的模型。以下是幾種有助于本地部署的技術：

• Switch Transformer 的作者進行過蒸餾實驗。將 MoE 蒸餾到稠密模型中，可以保留 30-40% 的稀疏化增益。因此，蒸餾可以帶來更快的預處理速度以及模型更小的好處。
• 創(chuàng)新的路由算法：將完整的句子或任務路由至專家網絡，使用用于提取的子網絡來提供服務。
• MoE 聚合：這種技術可以合并專家權重，從而減少推理時使用的參數(shù)。

其他高效訓練的方法

FasterMoE（2022 年 3 月提出）分析了高效分布式系統(tǒng)中 MoE 的性能，并分析了不同并行策略的理論極限，還分析了傾斜專家受歡迎程度的技術、減少延遲的細粒度通信調度，以及根據(jù)最低延遲挑選專家的拓撲感知門控，從而使速度提高了 17 倍。

Megablocks（https://arxiv.org/abs/2211.15841）推出了一款全新的 GPU 內核，能處理 MoE 中存在的動態(tài)問題，探索了高效的稀疏化預訓練。論文中建議不要丟棄任何一個 token，并實現(xiàn)了高效地 token 映射技術，從而顯著提高了速度。訣竅在于，傳統(tǒng)的 MoE 使用分批矩陣乘法，假定所有專家都具有相同的形狀和相同數(shù)量的 token。相比之下，Megablocks 將 MoE 層表示為塊稀疏運算，可以適應不平衡分配。

針對不同大小的專家和 token 數(shù)量的塊稀疏矩陣乘法。圖片摘自 MegaBlocks 論文。

開源的 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

已發(fā)布的 MoE：

• Switch Transformer (Google)：從 8 到 2048 個專家的基于 T5 的模擬引擎集合。最大的模型有 1.6 萬億個參數(shù)。
• NLLB MoE (Meta)：NLLB 翻譯模型的 MoE 變體。
• OpenMoE：基于 Llama 的 MoE 。
• Mixtral 8x7B (Mistral)：一種高質量的翻譯模型，其性能優(yōu)于 Llama 2 70B，推理速度也更快。

未來展望

以下是一些值得探索的有趣的領域：

• 將 Mixtral 提煉為稠密模型
• 探索專家模型的合并技術，分析其對推理時間的影響
• Mixtral 的量化技術