當(dāng)查看大型語(yǔ)言模型（LLMs）的最新發(fā)布時(shí)，你經(jīng)常會(huì)看到標(biāo)題中帶有“MoE”。這個(gè)“MoE”代表什么？為什么這么多LLMs都在使用它？

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

在此視覺指南中，我們將介紹 MoE 的兩個(gè)主要組件，即專家和路由器，以及它們?cè)诘湫偷幕贚LM架構(gòu)中的應(yīng)用。

查看與LLMs相關(guān)的更多可視化內(nèi)容，并支持此通訊，請(qǐng)參閱我寫的關(guān)于大型語(yǔ)言模型的書籍！

什么是混合專家（MoE）？

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

兩個(gè)主要組件定義了 MoE：

• 專家 - 每個(gè) FFNN 層現(xiàn)在都有一個(gè)“專家”集合，其中可以選擇子集。這些“專家”通常是 FFNN 本身。
• 路由器或網(wǎng)關(guān)網(wǎng)絡(luò) - 確定哪些令牌發(fā)送給哪些專家。

在具有 MoE 的LLM的每一層中，我們發(fā)現(xiàn)（某種程度上是專業(yè)的）專家：

“專家”并非在特定領(lǐng)域如“心理學(xué)”或“生物學(xué)”中專業(yè)化。最多，它只是詞法方面的專家：

更具體地說，他們的專長(zhǎng)在于處理特定語(yǔ)境中的特定標(biāo)記。

路由器（網(wǎng)關(guān)網(wǎng)絡(luò)）選擇最適合給定輸入的專家（們）：

每位專家不是一個(gè)完整的LLM，而是LLM架構(gòu)的子模型部分。

專家Expert, MoE中的E

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

密集層Dence layer

混合專家（MoE）都始于相對(duì)基本的功能LLMs，即前饋神經(jīng)網(wǎng)絡(luò)（FFNN）。

請(qǐng)記住，標(biāo)準(zhǔn)解碼器 Transformer 架構(gòu)在層歸一化之后應(yīng)用 FFNN：

一個(gè) FFNN 允許模型使用由注意力機(jī)制創(chuàng)建的上下文信息，進(jìn)一步將其轉(zhuǎn)化為捕捉數(shù)據(jù)中更復(fù)雜關(guān)系的能力。

FFNN 然而，其規(guī)模卻迅速增長(zhǎng)。為了學(xué)習(xí)這些復(fù)雜關(guān)系，它通常會(huì)在接收到的輸入基礎(chǔ)上進(jìn)行擴(kuò)展：

稀疏層

在傳統(tǒng)的 Transformer 中，F(xiàn)FNN 被稱為密集模型，因?yàn)樗袇?shù)（其權(quán)重和偏差）都被激活。沒有任何東西被留下，所有內(nèi)容都用于計(jì)算輸出。

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

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

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

底層理念是每個(gè)專家在訓(xùn)練期間學(xué)習(xí)不同的信息。然后，在運(yùn)行推理時(shí)，僅使用與特定任務(wù)最相關(guān)的特定專家。

當(dāng)被問及問題時(shí)，我們可以選擇最適合執(zhí)行特定任務(wù)的專家：

專家學(xué)到了什么？

如我們所見，專家比整個(gè)領(lǐng)域?qū)W習(xí)到更細(xì)粒度的信息。1 因此，有時(shí)將他們稱為“專家”被視為誤導(dǎo)。

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

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

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

此視覺圖還表明，專家往往關(guān)注語(yǔ)法而非特定領(lǐng)域。

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

專家架構(gòu)

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

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

所選專家可能在標(biāo)記之間有所不同，從而導(dǎo)致采取不同的“路徑”：

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

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

路由機(jī)制

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

在專家之前，添加了一個(gè)路由器（也稱為網(wǎng)關(guān)網(wǎng)絡(luò)），該路由器經(jīng)過訓(xùn)練，能夠選擇為給定令牌選擇哪個(gè)專家。

路由器

路由器（或網(wǎng)關(guān)網(wǎng)絡(luò)）也是一種前饋神經(jīng)網(wǎng)絡(luò)（FFNN），用于根據(jù)特定輸入選擇專家。它輸出概率，并使用這些概率來選擇最佳匹配的專家：

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

路由器與專家（其中只有少數(shù)被選中）共同構(gòu)成 MoE 層：

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

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

例如，給定一組標(biāo)記，MoE 會(huì)將標(biāo)記分配給所有專家，而稀疏 MoE 只會(huì)選擇少數(shù)專家。

在當(dāng)前狀態(tài)LLMs下，當(dāng)你看到“MoE”時(shí)，它通常是一個(gè)稀疏 MoE，因?yàn)樗试S你使用專家子集。這在計(jì)算上更便宜，這是LLMs的一個(gè)重要特性。

專家選擇

門控網(wǎng)絡(luò)可以說是任何 MoE 最重要的組成部分，因?yàn)樗粌H決定了推理期間選擇哪些專家，還包括訓(xùn)練階段。

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

然后，我們對(duì)輸出應(yīng)用 SoftMax，為每位專家創(chuàng)建概率分布 G(x)：

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

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

讓我們把所有內(nèi)容整合起來，探索輸入如何通過路由器和專家流動(dòng)：

路由的復(fù)雜性

然而，這個(gè)簡(jiǎn)單的函數(shù)往往導(dǎo)致路由器選擇相同的專家，因?yàn)槟承＜铱赡鼙绕渌麑＜覍W(xué)習(xí)得更快：

不僅選定的專家分布不均，而且一些專家?guī)缀醵紱]有經(jīng)過培訓(xùn)。這導(dǎo)致在訓(xùn)練和推理過程中都存在問題。

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

負(fù)載均衡

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

保持頂部 K

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

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

通過將這些權(quán)重設(shè)置為 -∞，這些權(quán)重上的 SoftMax 輸出將導(dǎo)致概率為 0：

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

Token 選擇

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

或超過一個(gè)專家（top-k 路由）：

一個(gè)主要好處是它允許權(quán)衡和整合專家各自的貢獻(xiàn)。

輔助損失

為了在訓(xùn)練期間獲得更均勻的專家分布，網(wǎng)絡(luò)常規(guī)損失中添加了輔助損失（也稱為負(fù)載均衡損失）。

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

該輔助損失的第一個(gè)組成部分是對(duì)整個(gè)批次中每個(gè)專家的路由器值進(jìn)行求和：

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

我們可以用這個(gè)來計(jì)算變異系數(shù)（CV），它告訴我們專家之間的重要性分?jǐn)?shù)差異有多大。

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

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

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

最后，將輔助損失作為一個(gè)單獨(dú)的損失項(xiàng)用于訓(xùn)練優(yōu)化。

專家能力

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

例如，如果輸入標(biāo)記不成比例地發(fā)送給某個(gè)專家而不是另一個(gè)，這也可能導(dǎo)致欠訓(xùn)練：

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

該問題的解決方案是限制特定專家可以處理的令牌數(shù)量，即專家容量 4。當(dāng)專家達(dá)到容量時(shí)，產(chǎn)生的令牌將被發(fā)送到下一個(gè)專家：

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

簡(jiǎn)化 MoE 與 Switch Transformer

之一是處理 MoE（如負(fù)載均衡）訓(xùn)練不穩(wěn)定問題的第一個(gè)基于 transformer 的 MoE 模型——Switch Transformer。5 它簡(jiǎn)化了大部分架構(gòu)和訓(xùn)練過程，同時(shí)提高了訓(xùn)練穩(wěn)定性。

交換層

開關(guān)變換器是一種 T5 模型（編碼器-解碼器），它用切換層替換了傳統(tǒng)的 FFNN 層。切換層是一個(gè)稀疏 MoE 層，為每個(gè)標(biāo)記選擇一個(gè)專家（Top-1 路由）。

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

這種架構(gòu)（top-1 路由）假設(shè)路由器只需要 1 位專家就能學(xué)習(xí)如何路由輸入。這與我們之前所見的假設(shè)不同，我們之前假設(shè)應(yīng)該將令牌路由到多個(gè)專家（top-k 路由）以學(xué)習(xí)路由行為。

容量系數(shù)

容量系數(shù)是一個(gè)重要值，因?yàn)樗鼪Q定了專家可以處理多少個(gè)標(biāo)記。Switch Transformer 通過引入直接影響專家容量的容量系數(shù)來擴(kuò)展這一點(diǎn)。

專家能力組成部分簡(jiǎn)單明了：

如果我們提高容量系數(shù)，每位專家將能夠處理更多的標(biāo)記。

然而，如果容量系數(shù)過大，我們會(huì)浪費(fèi)計(jì)算資源。相比之下，如果容量系數(shù)過小，由于標(biāo)記溢出，模型性能將下降。

輔助損失

為進(jìn)一步防止丟失標(biāo)記，引入了輔助損失的簡(jiǎn)化版本。

而不是計(jì)算變異系數(shù)，這種簡(jiǎn)化的損失衡量了分配的標(biāo)記分?jǐn)?shù)與每個(gè)專家的路由概率分?jǐn)?shù)之比：

由于目標(biāo)是讓 N 個(gè)專家之間具有統(tǒng)一的令牌路由，我們希望向量 P 和 f 的值為 1/N。

α 是一個(gè)超參數(shù)，我們可以用它來微調(diào)訓(xùn)練過程中這個(gè)損失函數(shù)的重要性。值過高會(huì)超過主損失函數(shù)，值過低對(duì)負(fù)載均衡的幫助很小。

專家混合模型

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

作為快速回顧，ViT（視覺 Transformer）是一種將圖像分割成類似標(biāo)記處理的補(bǔ)丁的架構(gòu)。6

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

這些補(bǔ)丁進(jìn)入編碼器的那一刻，它們就像標(biāo)記一樣被處理，這使得這種架構(gòu)非常適合 MoE。

Vision-MoE

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

這允許 ViT 模型，通常比語(yǔ)言模型尺寸更小，通過添加專家進(jìn)行大規(guī)模擴(kuò)展。

每個(gè)專家使用了一個(gè)小的預(yù)定義專家容量，以減少硬件限制，因?yàn)閳D像通常有許多塊。然而，低容量往往會(huì)造成塊被丟棄（類似于令牌溢出）。

為了保持容量低，網(wǎng)絡(luò)將重要性分?jǐn)?shù)分配給補(bǔ)丁，并首先處理這些補(bǔ)丁，以便溢出的補(bǔ)丁通常不太重要。這被稱為批量?jī)?yōu)先路由。

因此，即使令牌百分比下降，我們?nèi)匀粦?yīng)該看到重要的補(bǔ)丁被路由。

優(yōu)先路由通過關(guān)注最重要的部分，允許處理更少的補(bǔ)丁。

在 V-MoE 中，優(yōu)先級(jí)評(píng)分器幫助區(qū)分重要程度不同的補(bǔ)丁。然而，補(bǔ)丁被分配給每個(gè)專家，未處理補(bǔ)丁中的信息丟失。

軟 MoE 旨在通過混合補(bǔ)丁從離散補(bǔ)丁（標(biāo)記）分配過渡到軟補(bǔ)丁（標(biāo)記）分配。8

在第一步中，我們將輸入 x（補(bǔ)丁嵌入）與可學(xué)習(xí)的矩陣Φ相乘。這為我們提供了路由信息，告訴我們某個(gè)標(biāo)記與特定專家的相關(guān)程度。

通過隨后對(duì)路由信息矩陣（按列）進(jìn)行 softmax 操作，我們更新每個(gè)補(bǔ)丁的嵌入。

更新的補(bǔ)丁嵌入實(shí)際上是所有補(bǔ)丁嵌入的加權(quán)平均值。

視覺上，就像所有的補(bǔ)丁都被混合在一起了。這些組合補(bǔ)丁隨后被發(fā)送給每位專家。在生成輸出后，它們?cè)俅闻c路由矩陣相乘。

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

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

活躍參數(shù)與稀疏參數(shù)在 Mixtral 8x7B 中

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

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

換句話說，我們?nèi)匀恍枰谀脑O(shè)備上加載整個(gè)模型（包括所有專家）（稀疏參數(shù)），但在進(jìn)行推理時(shí)，我們只需要使用子集（活動(dòng)參數(shù)）。MoE 模型需要更多的 VRAM 來加載所有專家，但在推理時(shí)運(yùn)行更快。

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

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

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

結(jié)論

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

本文轉(zhuǎn)載自??AI大模型世界??，作者： Maarten ????