當查看大型語言模型（LLMs）的最新發布時，你經常會看到標題中帶有“MoE”。這個“MoE”代表什么？為什么這么多LLMs都在使用它？

在這份視覺指南中，我們將花時間探索這個重要組件——專家混合（MoE），通過超過 50 個可視化進行深入探討！

在此視覺指南中，我們將介紹 MoE 的兩個主要組件，即專家和路由器，以及它們在典型的基于LLM架構中的應用。

查看與LLMs相關的更多可視化內容，并支持此通訊，請參閱我寫的關于大型語言模型的書籍！

什么是混合專家（MoE）？

混合專家（MoE）是一種使用許多不同的子模型（或“專家”）來提高LLMs質量的技術。

兩個主要組件定義了 MoE：

• 專家 - 每個 FFNN 層現在都有一個“專家”集合，其中可以選擇子集。這些“專家”通常是 FFNN 本身。
• 路由器或網關網絡 - 確定哪些令牌發送給哪些專家。

在具有 MoE 的LLM的每一層中，我們發現（某種程度上是專業的）專家：

“專家”并非在特定領域如“心理學”或“生物學”中專業化。最多，它只是詞法方面的專家：

更具體地說，他們的專長在于處理特定語境中的特定標記。

路由器（網關網絡）選擇最適合給定輸入的專家（們）：

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

專家Expert, MoE中的E

為了探索專家所代表的內容以及他們的工作方式，讓我們首先考察 MoE 應該取代的內容；密集層。

密集層Dence layer

混合專家（MoE）都始于相對基本的功能LLMs，即前饋神經網絡（FFNN）。

請記住，標準解碼器 Transformer 架構在層歸一化之后應用 FFNN：

一個 FFNN 允許模型使用由注意力機制創建的上下文信息，進一步將其轉化為捕捉數據中更復雜關系的能力。

FFNN 然而，其規模卻迅速增長。為了學習這些復雜關系，它通常會在接收到的輸入基礎上進行擴展：

稀疏層

在傳統的 Transformer 中，FFNN 被稱為密集模型，因為所有參數（其權重和偏差）都被激活。沒有任何東西被留下，所有內容都用于計算輸出。

如果我們仔細觀察密集模型，注意輸入如何在一定程度上激活所有參數：

相比之下，稀疏模型僅激活其總參數的一部分，并且與專家混合模型密切相關。

為了說明，我們可以將我們的密集模型切割成多個部分（所謂專家），重新訓練它，并在給定時間只激活一組專家：

底層理念是每個專家在訓練期間學習不同的信息。然后，在運行推理時，僅使用與特定任務最相關的特定專家。

當被問及問題時，我們可以選擇最適合執行特定任務的專家：

專家學到了什么？

如我們所見，專家比整個領域學習到更細粒度的信息。1 因此，有時將他們稱為“專家”被視為誤導。

編碼器模型在 ST-MoE 論文中的專家專業化。

解碼器模型專家似乎沒有相同的類型專業化。但這并不意味著所有專家都是平等的。

一個很好的例子可以在 Mixtral 8x7B 論文中找到，其中每個標記都涂上了第一個專家的選擇。

此視覺圖還表明，專家往往關注語法而非特定領域。

因此，盡管解碼器專家似乎沒有特定的專長，但他們似乎一直被用于某些類型的標記。

專家架構

盡管將專家視為密集模型中切割成塊的隱藏層聽起來很酷，但他們本身往往是完整的全連接神經網絡（FFNN）：

由于大多數LLMs具有多個解碼器塊，給定文本在生成之前將通過多個專家：

所選專家可能在標記之間有所不同，從而導致采取不同的“路徑”：

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

解碼器塊現在具有多個 FFNN（每個都是一個“專家”），它在推理期間可以使用。

路由機制

現在我們有一批專家，模型是如何知道使用哪些專家的？

在專家之前，添加了一個路由器（也稱為網關網絡），該路由器經過訓練，能夠選擇為給定令牌選擇哪個專家。

路由器

路由器（或網關網絡）也是一種前饋神經網絡（FFNN），用于根據特定輸入選擇專家。它輸出概率，并使用這些概率來選擇最佳匹配的專家：

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

路由器與專家（其中只有少數被選中）共同構成 MoE 層：

給定 MoE 層有兩種大小，要么是稀疏的專家混合，要么是密集的專家混合。

兩者都使用路由器來選擇專家，但稀疏 MoE 只選擇少數，而密集 MoE 則選擇所有，但可能在不同的分布中。

例如，給定一組標記，MoE 會將標記分配給所有專家，而稀疏 MoE 只會選擇少數專家。

在當前狀態LLMs下，當你看到“MoE”時，它通常是一個稀疏 MoE，因為它允許你使用專家子集。這在計算上更便宜，這是LLMs的一個重要特性。

專家選擇

門控網絡可以說是任何 MoE 最重要的組成部分，因為它不僅決定了推理期間選擇哪些專家，還包括訓練階段。

在其最基本的形式中，我們將輸入（x）乘以路由權重矩陣（W）：

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

路由器使用這種概率分布來選擇給定輸入的最佳匹配專家。

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

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

路由的復雜性

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

不僅選定的專家分布不均，而且一些專家幾乎都沒有經過培訓。這導致在訓練和推理過程中都存在問題。

相反，我們希望在訓練和推理過程中專家之間保持同等重要性，我們稱之為負載均衡。從某種意義上講，這是為了防止對同一專家的過度擬合。

負載均衡

為了平衡專家的重要性，我們需要將路由器視為主要組件，以決定在特定時間選擇哪些專家。

保持頂部 K

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

然后，除了您想要激活的前 k 個專家（例如 2 個）之外，其余專家的權重將被設置為-∞：

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

The KeepTopK 策略是許多 LLMs 盡管有許多有希望的替代方案仍在使用的策略。請注意，KeepTopK 也可以在不添加額外噪聲的情況下使用。

Token 選擇

The KeepTopK 策略將每個令牌路由到幾位選定的專家。這種方法稱為 Token Choice 3，允許將特定令牌發送給一位專家（top-1 路由）：

或超過一個專家（top-k 路由）：

一個主要好處是它允許權衡和整合專家各自的貢獻。

輔助損失

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

它增加了一個約束，迫使專家們具有同等重要性。

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

這為我們提供了每位專家的重要性得分，它代表了無論輸入如何，某位專家被選中的可能性。

我們可以用這個來計算變異系數（CV），它告訴我們專家之間的重要性分數差異有多大。

例如，如果重要性得分差異很大，CV 將會很高：

相比之下，如果所有專家的重要性評分都相似，則 CV 值將較低（這正是我們追求的目標）：

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

最后，將輔助損失作為一個單獨的損失項用于訓練優化。

專家能力

不平衡不僅存在于被選中的專家中，也存在于發送給專家的代幣分布中。

例如，如果輸入標記不成比例地發送給某個專家而不是另一個，這也可能導致欠訓練：

這里，不僅僅是關于使用哪些專家，而是他們被使用了多少。

該問題的解決方案是限制特定專家可以處理的令牌數量，即專家容量 4。當專家達到容量時，產生的令牌將被發送到下一個專家：

如果兩位專家都已達到其容量，則令牌不會被任何專家處理，而是發送到下一層。這被稱為令牌溢出。

簡化 MoE 與 Switch Transformer

之一是處理 MoE（如負載均衡）訓練不穩定問題的第一個基于 transformer 的 MoE 模型——Switch Transformer。5 它簡化了大部分架構和訓練過程，同時提高了訓練穩定性。

交換層

開關變換器是一種 T5 模型（編碼器-解碼器），它用切換層替換了傳統的 FFNN 層。切換層是一個稀疏 MoE 層，為每個標記選擇一個專家（Top-1 路由）。

路由器在計算選擇哪個專家時沒有特殊技巧，只是將輸入乘以專家的權重后取 softmax（與我們之前做的一樣）。

這種架構（top-1 路由）假設路由器只需要 1 位專家就能學習如何路由輸入。這與我們之前所見的假設不同，我們之前假設應該將令牌路由到多個專家（top-k 路由）以學習路由行為。

容量系數

容量系數是一個重要值，因為它決定了專家可以處理多少個標記。Switch Transformer 通過引入直接影響專家容量的容量系數來擴展這一點。

專家能力組成部分簡單明了：

如果我們提高容量系數，每位專家將能夠處理更多的標記。

然而，如果容量系數過大，我們會浪費計算資源。相比之下，如果容量系數過小，由于標記溢出，模型性能將下降。

輔助損失

為進一步防止丟失標記，引入了輔助損失的簡化版本。

而不是計算變異系數，這種簡化的損失衡量了分配的標記分數與每個專家的路由概率分數之比：

由于目標是讓 N 個專家之間具有統一的令牌路由，我們希望向量 P 和 f 的值為 1/N。

α 是一個超參數，我們可以用它來微調訓練過程中這個損失函數的重要性。值過高會超過主損失函數，值過低對負載均衡的幫助很小。

專家混合模型

MoE 并非僅限于語言模型的技巧。視覺模型（如 ViT）利用基于 transformer 的架構，因此有潛力使用 MoE。

作為快速回顧，ViT（視覺 Transformer）是一種將圖像分割成類似標記處理的補丁的架構。6

這些補丁（或令牌）隨后被投影到嵌入（帶有額外的位置嵌入）中，然后再輸入到常規編碼器中：

這些補丁進入編碼器的那一刻，它們就像標記一樣被處理，這使得這種架構非常適合 MoE。

Vision-MoE

Vision-MoE（V-MoE）是圖像模型中 MoE 的第一個實現之一。7 它采用我們之前看到的 ViT，并將編碼器中的密集 FFNN 替換為稀疏 MoE。

這允許 ViT 模型，通常比語言模型尺寸更小，通過添加專家進行大規模擴展。

每個專家使用了一個小的預定義專家容量，以減少硬件限制，因為圖像通常有許多塊。然而，低容量往往會造成塊被丟棄（類似于令牌溢出）。

為了保持容量低，網絡將重要性分數分配給補丁，并首先處理這些補丁，以便溢出的補丁通常不太重要。這被稱為批量優先路由。

因此，即使令牌百分比下降，我們仍然應該看到重要的補丁被路由。

優先路由通過關注最重要的部分，允許處理更少的補丁。

在 V-MoE 中，優先級評分器幫助區分重要程度不同的補丁。然而，補丁被分配給每個專家，未處理補丁中的信息丟失。

軟 MoE 旨在通過混合補丁從離散補丁（標記）分配過渡到軟補丁（標記）分配。8

在第一步中，我們將輸入 x（補丁嵌入）與可學習的矩陣Φ相乘。這為我們提供了路由信息，告訴我們某個標記與特定專家的相關程度。

通過隨后對路由信息矩陣（按列）進行 softmax 操作，我們更新每個補丁的嵌入。

更新的補丁嵌入實際上是所有補丁嵌入的加權平均值。

視覺上，就像所有的補丁都被混合在一起了。這些組合補丁隨后被發送給每位專家。在生成輸出后，它們再次與路由矩陣相乘。

路由矩陣影響令牌級別的輸入和專家級別的輸出。

因此，我們得到“軟”補丁/令牌，而不是處理離散輸入。

活躍參數與稀疏參數在 Mixtral 8x7B 中

莫埃（MoE）引人入勝的一大原因是其計算需求。由于在特定時間只使用專家子集，我們能夠訪問比實際使用的更多參數。

盡管給定的 MoE 有更多參數需要加載（稀疏參數），但由于我們只在推理時使用一些專家（激活參數），因此激活的參數較少。

換句話說，我們仍然需要在您的設備上加載整個模型（包括所有專家）（稀疏參數），但在進行推理時，我們只需要使用子集（活動參數）。MoE 模型需要更多的 VRAM 來加載所有專家，但在推理時運行更快。

讓我們通過一個例子來探討稀疏參數與活躍參數的數量，Mixtral 8x7B。9

這里，我們可以看到每個專家的大小是 5.6B，而不是 7B（盡管有 8 個專家）。

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

結論

這標志著我們與專家混合模型的旅程結束！希望這篇帖子能讓你更好地理解這種有趣技術的潛力。現在，幾乎所有模型集都至少包含一個 MoE 變體，感覺它將在這里停留。

本文轉載自??AI大模型世界??，作者： Maarten ????