什么是混合專家模型MoE呢？

混合專家模型MoE是一種能讓大型語言模型（LLM）質量變好的技術。

它主要由兩部分組成：

專家——現在每個前饋神經網絡（FFNN）層都有好多個“專家”，每次可以用其中一部分。“專家”一般也是前饋神經網絡。

門控網絡——它能決定把哪些tokens送到哪些專家那兒。

在帶混合專家模型MoE的大型語言模型的每一層，都能看到這些有點專門作用的專家:

要知道，這里的‘專家’可不是像‘心理學’‘生物學’這種特定領域的專家哦。最多就是學學詞語層面的語法信息：

說得更細點兒，它們擅長處理特定情境下的特定詞。

門控網絡會挑出最適合給定輸入的專家：

每個專家也不是一個完整的大型語言模型，只是大型語言模型架構里的一部分小模型。

專家

要搞清楚專家代表啥、咋工作的，咱們先瞅瞅混合專家模型MoE是想替代啥；就是密集層。

密集層

混合專家模型MoE都是從大型語言模型（LLM）一個比較基礎的功能開始的，就是前饋神經網絡（FFNN）。

得記得，標準的只解碼器的Transformer架構里，前饋神經網絡是在層歸一化之后用的。

前饋神經網絡（FFNN）能讓模型利用注意力機制產生的上下文信息，進一步轉換它，從而捕捉數據里更復雜的關系。

不過，前饋神經網絡的大小可是會迅速變大的。為了學習這些復雜關系，它通常會擴充它接收到的輸入：

稀疏層

傳統Transformer里的前饋神經網絡（FFNN）被稱為密集模型，因為它所有的參數（權重和偏置）都被激活了。啥都不落下，全用來算輸出。

要是仔細瞅瞅密集模型，你會發現輸入會不同程度地激活所有參數：

相比之下，稀疏模型只激活一部分總參數，和混合專家模型MoE關系很近。

舉個例子，我們可以把密集模型切塊（這些塊就是所謂的專家），重新訓練它，然后在特定時間只激活一小部分專家：

背后的想法是，每個專家在訓練時學不同的信息。然后在推理時，只用特定的專家，因為它們對特定任務最相關。

當被問問題時，我們可以挑出最適合特定任務的專家：

專家都學了些啥呢？

之前說過，專家學的不是整個領域的知識，而是更細致的信息，所以有時候把它們叫‘專家’還挺容易讓人誤解的。

在ST-MoE論文里，編碼器模型的專家有專門的分工。不過，解碼器模型的專家好像就沒這種專門分工了。但這并不意味著所有專家都一樣。

Mixtral 8x7B論文里有個很好的例子，每個詞都用第一個選中的專家來上色。

這個圖也說明了專家更關注語法，而不是特定領域。所以，雖然解碼器專家好像沒有專門的領域，但它們在處理某些類型的詞時確實會一直被用到。

專家的架構

雖然把專家想象成密集模型里切塊的隱藏層挺形象的，但它們通常本身就是完整的前饋神經網絡（FFNN）：

因為大多數大型語言模型（LLM）有好幾個解碼器模塊，一段文本在生成之前要經過多個專家處理：

選中的專家在不同詞之間可能不一樣，這就導致走不同的“路徑”：

要是我們更新一下解碼器模塊的可視化，那現在里面就會有更多前饋神經網絡（每個專家一個）：

解碼器模塊現在有多個前饋神經網絡（每個都是一個“專家”），在推理時可以用它們。

路由機制

既然我們有了這么一幫專家，那模型咋知道該用哪個專家呢？在專家前面加了個門控網絡，它會學著挑出給定詞用哪個專家。

路由器

路由器（門控網絡）也是前饋神經網絡，根據特定輸入來選專家。它會輸出概率，用這些概率來挑出最匹配的專家：

專家層會返回選中的專家的輸出乘以門控值（選擇概率）。

路由器加上專家（只選幾個）就組成了專家混合（MoE）層：

一個MoE層有兩種尺寸，要么是稀疏的，要么是密集的專家混合。

兩者都用路由器選專家，但稀疏MoE只挑幾個，密集MoE雖然全選，但可能按不同分布來選。

比如，給定一串詞，MoE會把詞分給所有專家，而稀疏MoE就只挑幾個專家。

在目前大型語言模型（LLM）的狀況下，看到‘MoE’一般指稀疏MoE，因為它能用專家的一個子集。這樣計算成本低，對大型語言模型來說很重要。

專家的選擇

門控網絡可以說是專家混合（MoE）里最重要的部分，因為它不僅決定了推理時用哪些專家，訓練時也一樣。

最基礎的形式下，我們把輸入（x）和路由器的權重矩陣（W）相乘：

然后，對輸出應用 SoftMax 函數，為每個專家生成一個概率分布

路由器用這個概率分布來挑出給定輸入最適合的專家。

最后，我們把每個路由器的輸出和每個選中的專家的輸出相乘，再把結果加起來。

咱們把所有步驟串起來，看看輸入是怎么通過路由器和專家的：

路由的復雜性

不過，這個簡單的函數常常會導致路由器總是選同一個專家，因為有些專家可能學得比其他專家快：

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

為了解決這個問題，我們希望在訓練和推理過程中，所有專家都能得到平等的重視，這被稱為負載均衡。這樣可以防止模型過度依賴少數幾個專家，從而避免過擬合。

負載均衡

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

KeepTopK

一種負載均衡路由器的方法是通過一個簡單的擴展，稱為 KeepTopK。通過引入可訓練的高斯噪聲，有助于打破對特定專家的偏好，使選擇更加隨機化，這樣我們可以防止總是選擇相同的專家：

然后，除了你想要激活的頂部 k 個專家（例如 2 個）之外，其他專家的權重會被設置為 ?∞：

通過將這些權重設置為 ?∞，SoftMax 函數在這些權重上的輸出將導致概率為 0：

SoftMax 函數的性質： SoftMax 函數將一個向量轉換為概率分布，每個元素的值在 0 和 1 之間，且所有元素的和為 1。 SoftMax 函數的公式為：

其中，是輸入向量，是向量的長度，是向量的第 個元素。 設置權重為 ：

當某個元素的值為 時，的值為 0。

因此，對于任何包含的向量，SoftMax 函數將這些元素的輸出概率設置為 0。 具體步驟

假設我們有中間向量 ，其中某些元素被設置為：

應用 SoftMax 函數：

具體計算：由于e^{-\infty} = 0，上述表達式簡化為：\text{SoftMax}(\math

計算具體值：\text{SoftMax}(\mathbf{z}_{\text{top-k}}) = \left, \frac{1.6487}{1.6487 + $2.22$55}, 0, \frac{2.2255}{1.6487 + 2.2255}, 0 \right] = \left[ 0, 0.4256, 0,0.5744, 0 \right]

結果

通過將權重設置為 ，SoftMax 函數確保了這些專家的激活概率為 0，從而只激活頂部 個專家。

優點

負載均衡：確保每個專家都有機會被激活，避免某些專家過度使用而其他專家閑置。

防止過擬合：通過引入噪聲，防止模型過度依賴少數幾個專家，從而提高模型的泛化能力。

靈活性：可以根據具體需求調整的值，靈活控制專家的選擇數量。

KeepTopK 策略是許多大型語言模型（LLM）仍在使用的一種方法，盡管有許多有前景的替代方案。需要注意的是，KeepTopK 也可以在不添加額外噪聲的情況下使用。

Token Choice

KeepTopK 策略將每個詞路由到幾個選定的專家。這種方法稱為 Token Choice，允許給定的詞被發送到一個專家（top-1 路由）或多個專家（top-k 路由）。

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

輔助損失

一個主要的好處是，它允許對各個專家的貢獻進行加權和整合。

為了在訓練過程中實現專家的更均勻分布，輔助損失（也稱為負載均衡損失）被添加到了網絡的常規損失中。

它增加了一個約束，迫使專家具有相等的重要性。 這個輔助損失的第一個組成部分是在整個批次中對每個專家的路由器值進行求和：

這為我們提供了每個專家的重要性分數，這些分數表示無論輸入是什么，給定專家被選中的可能性。

我們可以利用這些分數來計算變異系數（CV），它告訴我們專家之間的重要性分數的差異程度。

例如，如果重要性分數差異很大，變異系數（CV）將會很高：

相反，如果所有專家的重要性分數相似，變異系數（CV）將會很低（這是我們所期望的）：

利用這個 CV 分數，我們可以在訓練過程中更新輔助損失，使其盡可能降低 CV 分數（從而給予每個專家相等的重要性）：

最后，輔助損失作為單獨的損失在訓練過程中進行優化。

專家容量

不平衡不僅存在于被選擇的專家中，還存在于發送到專家的令牌分布中。

例如，如果輸入令牌不成比例地發送到一個專家而不是另一個專家，這可能會導致訓練不足。

這里，問題不僅僅在于使用了哪些專家，還在于對它們的使用程度。

一個解決方案是限制給定專家可以處理的令牌數量，即專家容量（Expert Capacity）。當一個專家達到其容量時，后續的令牌將被發送給下一位專家。

如果兩個專家都達到容量了，這個詞就不會被任何專家處理，而是直接送到下一層。這就叫做詞溢出。

簡單來說，就好比有兩個專家在處理任務，但他們都忙不過來了，手上的活兒已經堆得滿滿的。這時候再來的新任務，他們就沒辦法處理了，只能把這個任務往后傳，讓下一層的專家或者別的機制來處理。這就叫詞溢出。

用 Switch Transformer 簡化 MoE

最早解決 MoE 訓練不穩定性問題（如負載均衡）的基于 Transformer 的 MoE 模型之一是 Switch Transformer。它極大地簡化了架構和訓練過程，同時提高了訓練的穩定性。

Switch Transformer

Switch Transformer 是 Google 在 2022 年發表的一篇論文中提出的一種模型，它簡化了 MoE 的路由算法，減少了計算量和通信量，同時支持 bfloat16 精度進行訓練。基于 T5-Base 和 T5-Large 設計的模型在相同的算力下訓練速度提升了 7 倍；同時發布了 1.6 萬億參數的 MoE 模型，相比 T5-XXL 模型訓練速度提升了 4 倍。

模型介紹

Switch Transformer 是一個 T5 模型（編碼器-解碼器），它用切換層替換了傳統的前饋神經網絡層。切換層是一個稀疏的 MoE 層，它為每個詞元選擇一個專家（Top-1 路由）。

簡化稀疏路由

在傳統的 MoE 層中，網絡接收一個輸入詞元 ，輸出會路由到 top-K 個專家進行計算。Switch Transformer 采用了一種簡化的路由策略，即每次只發給一個專家，這樣可以顯著減少路由的計算量，同時保證模型的性能。

高效稀疏路由

Switch Transformer 通過以下方式實現高效稀疏路由：

路由計算量減少：只有一個專家被激活，減少了計算量。

專家容量減半：專家中的 batch_size 至少減半，進一步減少了計算量。

簡化路由實現：減少了傳統 MoE 方法中通信的代價。

分布式 Switch 實現

Switch Transformer 的分布式實現中，所有核心被分配給數據分區維度 ，這也對應于模型中的專家數量。對于每個核心的每個詞元，路由器會本地計算分配給專家的分配。輸出是一個大小為 的二進制矩陣，該矩陣在第一維度上進行分區，并確定專家的分配。

負載均衡損失

為了促使每個專家都可以拿到近似均勻分布的樣本，Switch Transformer 引入了負載均衡損失。 當 時，損失是最小的。

切換層

Switch Transformer 是一個 T5 模型（編碼器-解碼器），它用切換層替換了傳統的前饋神經網絡層。切換層是一個稀疏的 MoE 層，它為每個詞元選擇一個專家（Top-1 路由）。

簡單來說，就是路由器在挑專家的時候，就是按照老規矩來，把輸入數據和專家的權重相乘，然后用 SoftMax 函數處理一下，看看哪個專家最合適。沒有啥復雜的操作，就是常規操作。

Switch Transformer 用了一種很簡單的辦法 來挑專家。以前的方法是讓每個詞元找好幾個專家來處理，這種方法叫 top-k 路由。但 Switch Transformer 說，其實每個詞元找一個專家就足夠了，這種方法叫 top-1 路由。這樣做的好處是，計算量變小了，通信成本也降低了，而且訓練起來更穩定。

容量因子

容量因子就像是給每個專家分配任務的“額度”。這個值很重要，因為它決定了每個專家能處理多少個詞元。Switch Transformer 在這方面做了改進，通過直接調整容量因子來控制每個專家的處理能力。這樣做的好處是，可以更好地平衡每個專家的負載，避免有些專家忙不過來，而有些專家卻很閑。

tokens per batch：每個批次中的令牌總數。

number of experts：專家的數量。

capacity factor：容量因子，通常大于 1，以提供額外的緩沖空間。

專家容量的組成部分其實挺簡單的。專家容量主要由兩個因素決定：每個批次中的令牌數量和專家的數量。

如果我們把容量因子調大，每個專家就能處理更多的詞元。這聽起來好像挺好的，因為專家們能干更多的活兒。

但是，如果把這個因子調得太大，就會浪費很多計算資源，因為有些專家可能會閑著沒事干。反過來，如果容量因子太小，專家們就會忙不過來，處理不了所有的詞元，這就叫令牌溢出。結果就是，模型的性能會變差，因為有些詞元沒被處理好。所以，容量因子得找個合適的值，這樣才能讓專家們既不過載，也不閑置，模型性能也能保持在最佳狀態。

輔助損失

輔助損失這個東西，就是為了防止模型在處理數據時丟掉一些重要的信息。想象一下，每個專家都有自己的工作量，我們希望這些工作量分配得均勻一些。為了達到這個目的，我們引入了一個簡化的輔助損失函數。

這個輔助損失函數不是去計算那些復雜的變異系數，而是直接看每個專家被分配到的令牌比例和路由器預測的概率。具體來說，我們希望每個專家處理的令牌數量和它們被選中的概率都差不多，也就是每個專家都能均勻地分到活兒干。

α 這個超參數就像是一個調節旋鈕，用來控制這個輔助損失在訓練過程中的重要性。如果 α 值設得太高，輔助損失就會太強勢，把主要的損失函數都蓋過去了；如果 α 值太低，那這個輔助損失就起不到什么作用，專家們的工作量還是不平衡。所以，α 值需要好好調整，才能讓模型既不丟信息，又能讓每個專家都忙得過來。

混合專家模型MoE在視覺模型中的應用

混合專家模型（MoE）這種技術不僅在語言模型里有用，在視覺模型里也能大顯身手。比如 ViT（Vision-Transformer）這種模型，它用的是 Transformer 的架構，所以完全可以把 MoE 用起來。

簡單來說，ViT 這個模型是把圖像切成很多小塊，這些小塊就像是語言模型里的詞元一樣被處理。具體來說，比如一張 224x224 像素的圖像，可以切成 16x16 像素的小塊，這樣就能得到 196 個小塊。每個小塊再被轉換成一維向量，這樣就能用 Transformer 來處理了。

這種方法的好處是，可以利用 Transformer 在處理序列數據上的強大能力，來處理圖像這種二維數據。這樣，視覺模型也能像語言模型一樣，通過 MoE 來提高效率和可擴展性。

在 Vision Transformer（ViT）模型中，我們先把圖像切成很多小塊，這些小塊就像是語言模型里的詞元一樣。然后，我們把這些小塊轉換成嵌入向量，這個過程有點像給每個小塊加上一個標簽，幫助模型更好地理解它們。

這個轉換過程包括兩個步驟：

線性變換：每個小塊被展平成一維向量，然后通過一個線性變換（全連接層）映射到一個固定的維度 D，形成每個小塊的嵌入向量。這個嵌入向量的維度就是 Transformer 的輸入維度。

位置嵌入：為了幫助模型理解每個小塊在圖像中的位置，我們還會加上位置嵌入。位置嵌入是一個額外的向量，它包含了每個小塊的位置信息。這樣，模型不僅知道每個小塊的內容，還知道它們在圖像中的位置。

最后，這些帶有位置嵌入的嵌入向量被輸入到常規的 Transformer 編碼器中，編碼器會通過自注意力機制來處理這些向量，提取圖像的全局特征。這個過程有點像語言模型處理詞元序列一樣，只不過這里是處理圖像的小塊序列。

當這些小塊（也就是圖像被切成的小方塊）進入編碼器的時候，它們的處理方式和語言模型里的詞元一樣。這就意味著，每個小塊都會被當作一個獨立的單元來處理，就像處理句子中的每個詞一樣。這種處理方式讓 Vision Transformer（ViT）這種架構非常適合用混合專家（MoE）模型來優化。因為 MoE 模型本來就是設計來處理這種“分而治之”的任務的，每個專家可以專注于處理一部分數據，這樣既能提高效率，又能保持高性能。

Vision-MoE

Vision-MoE（V-MoE）是圖像模型中第一個實現 MoE 的例子之一。它把 ViT 中的密集前饋神經網絡（FFNN）層換成了稀疏 MoE。這樣，ViT 模型（通常比語言模型小）就能通過增加專家數量來大規模擴展。

因為圖像通常包含很多小塊，為了減少硬件限制，每個專家使用了一個較小的預定義專家容量。但是，較低的容量會導致一些小塊被丟棄，就像令牌溢出一樣。

為了應對這個問題，網絡會為小塊分配重要性分數，并優先處理這些小塊，這樣溢出的小塊通常就不那么重要了。這個方法叫做批量優先級路由（Batch Priority Routing）。

所以，即使處理的小塊數量減少了，重要的小塊仍然會被路由。

優先級路由允許模型通過專注于最重要的小塊來減少處理的小塊數量。這樣，模型在保持高性能的同時，還能減少計算資源的消耗。

活躍與稀疏參數：以 Mixtral 8x7B 為例

MoE（混合專家）模型的一個關鍵優勢在于它的計算效率。

以 Mixtral 8x7B 為例，這個模型雖然總參數量很大，但在實際運行時，每次只激活一部分專家。這意味著，雖然模型有 47B 參數，但實際計算量相當于 12B 參數的密集模型。這種設計讓模型在保持高性能的同時，大大減少了計算資源的消耗。

具體來說，MoE 模型在訓練和推理時有以下優勢：

訓練速度更快：與相同參數規模的密集模型相比，MoE 可以用更少的計算資源和時間完成預訓練。例如，Google 的 Switch Transformer 模型大小是 T5-XXL 的 15 倍，但在相同計算資源下，Switch Transformer 模型在達到固定困惑度 PPL 時，比 T5-XXL 模型快 4 倍。

推理速度快：雖然模型總參數量大，但推理時只激活部分專家網絡。對于同等參數規模的密集模型，MoE 推理速度更快。例如，Mixtral 8x7B 雖然有 47B 參數，但推理時計算量相當于 12B 參數的密集模型。

擴展性強：MoE 允許模型在保持計算成本不變的情況下增加參數數量，這使得它能夠擴展到非常大的模型規模，如萬億參數模型。

多任務學習能力：MoE 在多任務學習中表現出色，能夠同時處理多種不同的任務，每個專家可以專注于特定的任務或數據類型。

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

Mixtral 8x7B 就是這樣一個模型，它有很多參數（稀疏參數），但實際推理時只用到其中的一部分（活躍參數）。

具體來說，Mixtral 8x7B 模型有 8 個專家，每個專家有 5.6B 參數。雖然模型總參數量是 46.7B（8 個專家的參數加上共享參數），但在推理時，每個輸入只用到 2 個專家，所以實際激活的參數量是 12.8B。這意味著，雖然加載整個模型需要很多顯存，但實際運行時效率很高。

這種設計讓 Mixtral 8x7B 在處理大規模數據時非常高效，既能保持高性能，又減少了計算資源的消耗。這就是為什么 MoE 模型在大型語言模型和視覺模型中越來越受歡迎的原因。

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

本文轉載自 ??智駐未來??，作者： 小智