漫畫混合專家(MoE)
當(dāng)查看大型語言模型(LLMs)的最新發(fā)布時,你經(jīng)常會看到標(biāo)題中帶有“MoE”。這個“MoE”代表什么?為什么這么多LLMs都在使用它?
在這份視覺指南中,我們將花時間探索這個重要組件——專家混合(MoE),通過超過 50 個可視化進(jìn)行深入探討!
在此視覺指南中,我們將介紹 MoE 的兩個主要組件,即專家和路由器,以及它們在典型的基于LLM架構(gòu)中的應(yīng)用。
查看與LLMs相關(guān)的更多可視化內(nèi)容,并支持此通訊,請參閱我寫的關(guān)于大型語言模型的書籍!
什么是混合專家(MoE)?
混合專家(MoE)是一種使用許多不同的子模型(或“專家”)來提高LLMs質(zhì)量的技術(shù)。
兩個主要組件定義了 MoE:
- 專家 - 每個 FFNN 層現(xiàn)在都有一個“專家”集合,其中可以選擇子集。這些“專家”通常是 FFNN 本身。
- 路由器或網(wǎng)關(guān)網(wǎng)絡(luò) - 確定哪些令牌發(fā)送給哪些專家。
在具有 MoE 的LLM的每一層中,我們發(fā)現(xiàn)(某種程度上是專業(yè)的)專家:
“專家”并非在特定領(lǐng)域如“心理學(xué)”或“生物學(xué)”中專業(yè)化。最多,它只是詞法方面的專家:
更具體地說,他們的專長在于處理特定語境中的特定標(biāo)記。
路由器(網(wǎng)關(guān)網(wǎng)絡(luò))選擇最適合給定輸入的專家(們):
每位專家不是一個完整的LLM,而是LLM架構(gòu)的子模型部分。
專家Expert, MoE中的E
為了探索專家所代表的內(nèi)容以及他們的工作方式,讓我們首先考察 MoE 應(yīng)該取代的內(nèi)容;密集層。
密集層Dence layer
混合專家(MoE)都始于相對基本的功能LLMs,即前饋神經(jīng)網(wǎng)絡(luò)(FFNN)。
請記住,標(biāo)準(zhǔn)解碼器 Transformer 架構(gòu)在層歸一化之后應(yīng)用 FFNN:
一個 FFNN 允許模型使用由注意力機(jī)制創(chuàng)建的上下文信息,進(jìn)一步將其轉(zhuǎn)化為捕捉數(shù)據(jù)中更復(fù)雜關(guān)系的能力。
FFNN 然而,其規(guī)模卻迅速增長。為了學(xué)習(xí)這些復(fù)雜關(guān)系,它通常會在接收到的輸入基礎(chǔ)上進(jìn)行擴(kuò)展:
稀疏層
在傳統(tǒng)的 Transformer 中,F(xiàn)FNN 被稱為密集模型,因為所有參數(shù)(其權(quán)重和偏差)都被激活。沒有任何東西被留下,所有內(nèi)容都用于計算輸出。
如果我們仔細(xì)觀察密集模型,注意輸入如何在一定程度上激活所有參數(shù):
相比之下,稀疏模型僅激活其總參數(shù)的一部分,并且與專家混合模型密切相關(guān)。
為了說明,我們可以將我們的密集模型切割成多個部分(所謂專家),重新訓(xùn)練它,并在給定時間只激活一組專家:
底層理念是每個專家在訓(xùn)練期間學(xué)習(xí)不同的信息。然后,在運(yùn)行推理時,僅使用與特定任務(wù)最相關(guān)的特定專家。
當(dāng)被問及問題時,我們可以選擇最適合執(zhí)行特定任務(wù)的專家:
專家學(xué)到了什么?
如我們所見,專家比整個領(lǐng)域?qū)W習(xí)到更細(xì)粒度的信息。1 因此,有時將他們稱為“專家”被視為誤導(dǎo)。
編碼器模型在 ST-MoE 論文中的專家專業(yè)化。
解碼器模型專家似乎沒有相同的類型專業(yè)化。但這并不意味著所有專家都是平等的。
一個很好的例子可以在 Mixtral 8x7B 論文中找到,其中每個標(biāo)記都涂上了第一個專家的選擇。
此視覺圖還表明,專家往往關(guān)注語法而非特定領(lǐng)域。
因此,盡管解碼器專家似乎沒有特定的專長,但他們似乎一直被用于某些類型的標(biāo)記。
專家架構(gòu)
盡管將專家視為密集模型中切割成塊的隱藏層聽起來很酷,但他們本身往往是完整的全連接神經(jīng)網(wǎng)絡(luò)(FFNN):
由于大多數(shù)LLMs具有多個解碼器塊,給定文本在生成之前將通過多個專家:
所選專家可能在標(biāo)記之間有所不同,從而導(dǎo)致采取不同的“路徑”:
如果我們更新解碼器塊的可視化,它現(xiàn)在將包含更多的 FFNN(每個專家一個)而不是:
解碼器塊現(xiàn)在具有多個 FFNN(每個都是一個“專家”),它在推理期間可以使用。
路由機(jī)制
現(xiàn)在我們有一批專家,模型是如何知道使用哪些專家的?
在專家之前,添加了一個路由器(也稱為網(wǎng)關(guān)網(wǎng)絡(luò)),該路由器經(jīng)過訓(xùn)練,能夠選擇為給定令牌選擇哪個專家。
路由器
路由器(或網(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 會將標(biāo)記分配給所有專家,而稀疏 MoE 只會選擇少數(shù)專家。
在當(dāng)前狀態(tài)LLMs下,當(dāng)你看到“MoE”時,它通常是一個稀疏 MoE,因為它允許你使用專家子集。這在計算上更便宜,這是LLMs的一個重要特性。
專家選擇
門控網(wǎng)絡(luò)可以說是任何 MoE 最重要的組成部分,因為它不僅決定了推理期間選擇哪些專家,還包括訓(xùn)練階段。
在其最基本的形式中,我們將輸入(x)乘以路由權(quán)重矩陣(W):
然后,我們對輸出應(yīng)用 SoftMax,為每位專家創(chuàng)建概率分布 G(x):
路由器使用這種概率分布來選擇給定輸入的最佳匹配專家。
最后,我們將每個路由器的輸出與每個選定的專家相乘,并將結(jié)果相加。
讓我們把所有內(nèi)容整合起來,探索輸入如何通過路由器和專家流動:
路由的復(fù)雜性
然而,這個簡單的函數(shù)往往導(dǎo)致路由器選擇相同的專家,因為某些專家可能比其他專家學(xué)習(xí)得更快:
不僅選定的專家分布不均,而且一些專家?guī)缀醵紱]有經(jīng)過培訓(xùn)。這導(dǎo)致在訓(xùn)練和推理過程中都存在問題。
相反,我們希望在訓(xùn)練和推理過程中專家之間保持同等重要性,我們稱之為負(fù)載均衡。從某種意義上講,這是為了防止對同一專家的過度擬合。
負(fù)載均衡
為了平衡專家的重要性,我們需要將路由器視為主要組件,以決定在特定時間選擇哪些專家。
保持頂部 K
一種路由器負(fù)載均衡的方法是通過一個簡單的擴(kuò)展稱為 KeepTopK 2。通過引入可訓(xùn)練(高斯)噪聲,我們可以防止總是選擇相同的專家:
然后,除了您想要激活的前 k 個專家(例如 2 個)之外,其余專家的權(quán)重將被設(shè)置為-∞:
通過將這些權(quán)重設(shè)置為 -∞,這些權(quán)重上的 SoftMax 輸出將導(dǎo)致概率為 0:
The KeepTopK 策略是許多 LLMs 盡管有許多有希望的替代方案仍在使用的策略。請注意,KeepTopK 也可以在不添加額外噪聲的情況下使用。
Token 選擇
The KeepTopK 策略將每個令牌路由到幾位選定的專家。這種方法稱為 Token Choice 3,允許將特定令牌發(fā)送給一位專家(top-1 路由):
或超過一個專家(top-k 路由):
一個主要好處是它允許權(quán)衡和整合專家各自的貢獻(xiàn)。
輔助損失
為了在訓(xùn)練期間獲得更均勻的專家分布,網(wǎng)絡(luò)常規(guī)損失中添加了輔助損失(也稱為負(fù)載均衡損失)。
它增加了一個約束,迫使專家們具有同等重要性。
該輔助損失的第一個組成部分是對整個批次中每個專家的路由器值進(jìn)行求和:
這為我們提供了每位專家的重要性得分,它代表了無論輸入如何,某位專家被選中的可能性。
我們可以用這個來計算變異系數(shù)(CV),它告訴我們專家之間的重要性分?jǐn)?shù)差異有多大。
例如,如果重要性得分差異很大,CV 將會很高:
相比之下,如果所有專家的重要性評分都相似,則 CV 值將較低(這正是我們追求的目標(biāo)):
使用此 CV 得分,我們可以在訓(xùn)練過程中更新輔助損失,使其盡可能降低 CV 得分(從而給予每個專家同等的重要性):
最后,將輔助損失作為一個單獨的損失項用于訓(xùn)練優(yōu)化。
專家能力
不平衡不僅存在于被選中的專家中,也存在于發(fā)送給專家的代幣分布中。
例如,如果輸入標(biāo)記不成比例地發(fā)送給某個專家而不是另一個,這也可能導(dǎo)致欠訓(xùn)練:
這里,不僅僅是關(guān)于使用哪些專家,而是他們被使用了多少。
該問題的解決方案是限制特定專家可以處理的令牌數(shù)量,即專家容量 4。當(dāng)專家達(dá)到容量時,產(chǎn)生的令牌將被發(fā)送到下一個專家:
如果兩位專家都已達(dá)到其容量,則令牌不會被任何專家處理,而是發(fā)送到下一層。這被稱為令牌溢出。
簡化 MoE 與 Switch Transformer
之一是處理 MoE(如負(fù)載均衡)訓(xùn)練不穩(wěn)定問題的第一個基于 transformer 的 MoE 模型——Switch Transformer。5 它簡化了大部分架構(gòu)和訓(xùn)練過程,同時提高了訓(xùn)練穩(wěn)定性。
交換層
開關(guān)變換器是一種 T5 模型(編碼器-解碼器),它用切換層替換了傳統(tǒng)的 FFNN 層。切換層是一個稀疏 MoE 層,為每個標(biāo)記選擇一個專家(Top-1 路由)。
路由器在計算選擇哪個專家時沒有特殊技巧,只是將輸入乘以專家的權(quán)重后取 softmax(與我們之前做的一樣)。
這種架構(gòu)(top-1 路由)假設(shè)路由器只需要 1 位專家就能學(xué)習(xí)如何路由輸入。這與我們之前所見的假設(shè)不同,我們之前假設(shè)應(yīng)該將令牌路由到多個專家(top-k 路由)以學(xué)習(xí)路由行為。
容量系數(shù)
容量系數(shù)是一個重要值,因為它決定了專家可以處理多少個標(biāo)記。Switch Transformer 通過引入直接影響專家容量的容量系數(shù)來擴(kuò)展這一點。
專家能力組成部分簡單明了:
如果我們提高容量系數(shù),每位專家將能夠處理更多的標(biāo)記。
然而,如果容量系數(shù)過大,我們會浪費計算資源。相比之下,如果容量系數(shù)過小,由于標(biāo)記溢出,模型性能將下降。
輔助損失
為進(jìn)一步防止丟失標(biāo)記,引入了輔助損失的簡化版本。
而不是計算變異系數(shù),這種簡化的損失衡量了分配的標(biāo)記分?jǐn)?shù)與每個專家的路由概率分?jǐn)?shù)之比:
由于目標(biāo)是讓 N 個專家之間具有統(tǒng)一的令牌路由,我們希望向量 P 和 f 的值為 1/N。
α 是一個超參數(shù),我們可以用它來微調(diào)訓(xùn)練過程中這個損失函數(shù)的重要性。值過高會超過主損失函數(shù),值過低對負(fù)載均衡的幫助很小。
專家混合模型
MoE 并非僅限于語言模型的技巧。視覺模型(如 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 的第一個實現(xiàn)之一。7 它采用我們之前看到的 ViT,并將編碼器中的密集 FFNN 替換為稀疏 MoE。
這允許 ViT 模型,通常比語言模型尺寸更小,通過添加專家進(jìn)行大規(guī)模擴(kuò)展。
每個專家使用了一個小的預(yù)定義專家容量,以減少硬件限制,因為圖像通常有許多塊。然而,低容量往往會造成塊被丟棄(類似于令牌溢出)。
為了保持容量低,網(wǎng)絡(luò)將重要性分?jǐn)?shù)分配給補(bǔ)丁,并首先處理這些補(bǔ)丁,以便溢出的補(bǔ)丁通常不太重要。這被稱為批量優(yōu)先路由。
因此,即使令牌百分比下降,我們?nèi)匀粦?yīng)該看到重要的補(bǔ)丁被路由。
優(yōu)先路由通過關(guān)注最重要的部分,允許處理更少的補(bǔ)丁。
在 V-MoE 中,優(yōu)先級評分器幫助區(qū)分重要程度不同的補(bǔ)丁。然而,補(bǔ)丁被分配給每個專家,未處理補(bǔ)丁中的信息丟失。
軟 MoE 旨在通過混合補(bǔ)丁從離散補(bǔ)丁(標(biāo)記)分配過渡到軟補(bǔ)丁(標(biāo)記)分配。8
在第一步中,我們將輸入 x(補(bǔ)丁嵌入)與可學(xué)習(xí)的矩陣Φ相乘。這為我們提供了路由信息,告訴我們某個標(biāo)記與特定專家的相關(guān)程度。
通過隨后對路由信息矩陣(按列)進(jìn)行 softmax 操作,我們更新每個補(bǔ)丁的嵌入。
更新的補(bǔ)丁嵌入實際上是所有補(bǔ)丁嵌入的加權(quán)平均值。
視覺上,就像所有的補(bǔ)丁都被混合在一起了。這些組合補(bǔ)丁隨后被發(fā)送給每位專家。在生成輸出后,它們再次與路由矩陣相乘。
路由矩陣影響令牌級別的輸入和專家級別的輸出。
因此,我們得到“軟”補(bǔ)丁/令牌,而不是處理離散輸入。
活躍參數(shù)與稀疏參數(shù)在 Mixtral 8x7B 中
莫埃(MoE)引人入勝的一大原因是其計算需求。由于在特定時間只使用專家子集,我們能夠訪問比實際使用的更多參數(shù)。
盡管給定的 MoE 有更多參數(shù)需要加載(稀疏參數(shù)),但由于我們只在推理時使用一些專家(激活參數(shù)),因此激活的參數(shù)較少。
換句話說,我們?nèi)匀恍枰谀脑O(shè)備上加載整個模型(包括所有專家)(稀疏參數(shù)),但在進(jìn)行推理時,我們只需要使用子集(活動參數(shù))。MoE 模型需要更多的 VRAM 來加載所有專家,但在推理時運(yùn)行更快。
讓我們通過一個例子來探討稀疏參數(shù)與活躍參數(shù)的數(shù)量,Mixtral 8x7B。9
這里,我們可以看到每個專家的大小是 5.6B,而不是 7B(盡管有 8 個專家)。
我們不得不加載 8x5.6B(46.7B)個參數(shù)(包括所有共享參數(shù)),但推理時我們只需要使用 2x5.6B(12.8B)個參數(shù)。
結(jié)論
這標(biāo)志著我們與專家混合模型的旅程結(jié)束!希望這篇帖子能讓你更好地理解這種有趣技術(shù)的潛力。現(xiàn)在,幾乎所有模型集都至少包含一個 MoE 變體,感覺它將在這里停留。
本文轉(zhuǎn)載自??AI大模型世界??,作者: Maarten ????
