原作：?塞繆爾·弗蘭德

引言：為最強(qiáng)大的語言模型鋪平道路的核心技術(shù)

?

使用 Dall-E 生成的圖像

稀疏混合專家模型 (MoE) 已成為最新一代 LLMs 的核心技術(shù)，例如 OpenAI 的 GPT-4、Mistral AI 的 Mixtral-8x7 等。簡而言之，稀疏MoE是一種非常強(qiáng)大的技術(shù)，因?yàn)槔碚撋希试S我們以 O(1) 的計(jì)算復(fù)雜度擴(kuò)展任何模型的容量！

然而，正如通常的情況一樣，問題在于細(xì)節(jié)，要讓稀疏的 MoE 正常工作就需要確保這些細(xì)節(jié)完全正確。

在本文中，我們將深入探討稀疏MoE領(lǐng)域的一個(gè)核心貢獻(xiàn)，即Switch Transformer（Fedus等人，2022年），它首次展示了利用這項(xiàng)技術(shù)實(shí)現(xiàn)了令人印象深刻的擴(kuò)展特性，在Transformer模型的訓(xùn)練中實(shí)現(xiàn)了7倍以上的加速。我們將討論以下內(nèi)容：

• 硬路由：僅針對每個(gè)token執(zhí)行單個(gè)專家模型所帶來的有利擴(kuò)展特性，
• Switch Transformer 架構(gòu)：MoE 如何適應(yīng)更廣泛的 Transformer 架構(gòu)背景，
• token動(dòng)態(tài)路由：如何使用容量因子來權(quán)衡計(jì)算效率和建模精度，以及
• 實(shí)驗(yàn)結(jié)果：Switch Transformer的令人印象深刻的擴(kuò)展特性。

?讓我們開始吧。

?硬路由

回顧一下，MoE 的關(guān)鍵思想是使用專家的線性組合 E(x) 對給定輸入 x 的輸出 y 進(jìn)行建模，每個(gè)專家的權(quán)重由門控 G(x) 控制，

其中門控是輸入x與可學(xué)習(xí)的權(quán)重矩陣W相乘后的softmax函數(shù)：

因此，在訓(xùn)練 MoE 模型時(shí)，學(xué)習(xí)目標(biāo)有兩個(gè)：

1. 專家將學(xué)習(xí)將給定的輸入處理成最佳的輸出（即預(yù)測）
2. 門控將通過學(xué)習(xí)矩陣W來將正確的訓(xùn)練樣本分配給正確的專家。

MoE 的這種原始表述可以追溯到30多年前 Geoffrey Hinton 在 90 年代的原創(chuàng)作品，如今已被稱為軟路由（soft routing）。之所以稱為“軟”，是因?yàn)楸M管不同的專家可能具有非常不同的權(quán)重（其中一些幾乎不明顯），我們?nèi)匀粚⑺袑＜业妮敵鼋M合在最終結(jié)果中，無論他們的貢獻(xiàn)多么小。

相比之下，在硬路由（hard routing）中，我們只對由門控決定的單個(gè)最合適的專家進(jìn)行向前傳遞，也就是說，我們近似地

其中I是使G取得最大值的索引。

這種方法的動(dòng)機(jī)是在計(jì)算成本上進(jìn)行建模精度的權(quán)衡：如果一個(gè)專家的權(quán)重只有0.01，那么是否真的值得對該專家進(jìn)行向前傳遞呢？

順便說一句，硬路由實(shí)際上是 top-k 路由的一種特殊情況，最初由 Shazeer 等人 (2017) 提出，其中 k=1。雖然Shazeer等人假設(shè)k=1在實(shí)踐中可能效果不佳，因?yàn)檫@會(huì)嚴(yán)重限制通過專家反向傳播的梯度，但Switch Transformer證明并非如此。

（注意，“gate”和“router”這些術(shù)語通常可以互換使用——在稀疏MoE的背景下，它們實(shí)際上意味著相同的東西。）

具有 5 個(gè)專家和 3 個(gè)token的硬路由插圖，每個(gè)token都有一個(gè) 3 維嵌入。此處，路由器將token 1 分配給專家 1，將token 2 分配給專家 4，將token3 分配給專家 2。圖片來自 Fedus et al 2022。

Switch Transformer架構(gòu)

在常規(guī)的（密集）Transformer塊中，我們將一個(gè)前饋神經(jīng)網(wǎng)絡(luò)（FFN）層堆疊在一個(gè)自注意力層之上，并在它們之間添加殘差連接。在高層次上，自注意力層選擇模型關(guān)注的輸入，而FFN層處理該輸入并將輸出傳遞到下一個(gè)塊，依此類推。例如，BERT-Large是一個(gè)將24個(gè)這些Transformer塊堆疊起來的Transformer模型。

這是這個(gè)密集的 Transformer 塊的可視化效果：

密集的變壓器塊。圖片來源：Fedus 等人 2022。

在 Switch Transformer 中，我們用多個(gè) FFN“專家”替換單個(gè) FFN 模塊，并讓硬路由器向?qū)＜曳峙鋞oken，如下所示：

Switch Transformer 中的稀疏Transformer塊。圖片來源：Fedus 等人 2022。

最重要的是，這兩種架構(gòu)具有相同的計(jì)算復(fù)雜度，但與前者相比，后者實(shí)際上具有4倍的建模容量（即神經(jīng)元數(shù)量）！這要?dú)w功于稀疏 MoE 的硬路由：雖然我們現(xiàn)在有4個(gè)FFN而不是一個(gè)，但對于任何給定的token，只有其中一個(gè)會(huì)處于活動(dòng)狀態(tài)。

token動(dòng)態(tài)路由和容量因子

到目前為止，我們只考慮了路由器如何在專家之間分配token，但沒有具體說明這些專家位于何處。在實(shí)踐中，我們通常將專家分布在多臺(tái)機(jī)器上，這是一種模型并行化形式，也稱為專家并行化。

專家并行化有一個(gè)重要的物理限制：我們不能將比專家所在機(jī)器的內(nèi)存允許的更多token路由到該專家！

token動(dòng)態(tài)路由作為容量因子的函數(shù)。如果因子為 1，我們就必須丟棄一個(gè)token（紅色箭頭）。在容量因子為1.5的情況下，我們需要添加過多的填充（白色矩形）。來自 Fedus 等人 2022

在這種背景下，我們將專家容量定義為

capacity = f x T/E

其中T是token數(shù)量，E是專家總數(shù)，f是一個(gè)被稱為容量因子的自由超參數(shù)。（如果我們使用k>1的top-k路由，則還會(huì)在方程的右側(cè)添加k的因子——然而，在這里我們考慮的是使用硬路由的Switch Transformer，所以我們將使用k=1。）

例如，對于T=6，E=3和f=1，我們允許路由器向每個(gè)專家發(fā)送最多 2 個(gè)token，如上圖左側(cè)面板所示。如果我們發(fā)送更多token，我們就需要丟棄額外的token，如果我們發(fā)送較少token，我們就需要對輸入進(jìn)行填充，以確保一致性——畢竟，在GPU上可以運(yùn)行的計(jì)算圖必須是靜態(tài)的，不能是動(dòng)態(tài)的。

因此，容量因子f引入了一個(gè)權(quán)衡：如果太大，我們會(huì)通過過多的填充浪費(fèi)計(jì)算資源（即上圖中的白色矩形）；如果太小，我們會(huì)因?yàn)閠oken丟失而犧牲模型性能（如上圖中的紅色箭頭所示）。

在Switch Transformer論文中，作者發(fā)現(xiàn)較低的容量因子可以獲得更好的性能，例如，將 f 從 2.0 降低到 1.0，在 100k 步后，對數(shù)困惑度從 -1.554 提高到 -1.561。這表明我們因丟棄token而遭受的懲罰并不像計(jì)算資源利用不足所帶來的懲罰那么嚴(yán)重。或者更簡單地說：最好優(yōu)化資源利用率，即使以丟棄token為代價(jià)。

Switch Transformer的可擴(kuò)展性特征

最重要的是，只需增加專家數(shù)量，Switch Transformer可以以近似恒定的計(jì)算復(fù)雜度進(jìn)行擴(kuò)展。這是因?yàn)楦嗟膶＜也粫?huì)導(dǎo)致更多的專家向前傳遞——得益于硬路由——而且與整個(gè)模型架構(gòu)相比，添加更多專家時(shí)需要由門控完成的額外計(jì)算可以忽略不計(jì)，特別是與計(jì)算量大的 Transformer 塊相比。

下圖顯示了通過用128個(gè)稀疏專家替換單個(gè)FFN模塊，作者能夠達(dá)到與T5語言模型相同的性能，但速度快了7倍！所有這些模型都在相同的機(jī)器上進(jìn)行訓(xùn)練，使用32個(gè)TPUv3核心，并且每個(gè)訓(xùn)練樣本具有相等的FLOPs。

Switch Transformer通過硬路由實(shí)現(xiàn)了7倍的加速，同時(shí)保持了恒定的計(jì)算復(fù)雜度。

乍一看，這似乎像是魔法。在相同的FLOPs下，怎么可能更快呢？這就好像一輛汽車以相同的速度行駛，卻比原來快7倍到達(dá)目的地！

答案就是，Switch Transformer 利用了稀疏性，特別是稀疏 MoE。雖然我們通過增加專家數(shù)量增加了建模容量，但由于硬路由，我們保持了恒定的FLOPs，也就是說，在每個(gè)訓(xùn)練迭代中我們并沒有執(zhí)行所有的專家，而是只選擇適合每個(gè)token的最佳專家進(jìn)行計(jì)算。

最后

?讓我們回顧一下：

• 稀疏MoE是一項(xiàng)突破性的技術(shù)，因?yàn)樗刮覀兡軌蛞越芆(1)的計(jì)算復(fù)雜度擴(kuò)展建模容量，并實(shí)現(xiàn)了Google的Switch Transformer、OpenAI的GPT-4、Mistral AI的Mixtral-8x7b等重大突破。
• 硬路由意味著我們只對最合適的專家運(yùn)行向前傳遞，而不是所有專家，從而節(jié)省了FLOPs。
• Switch Transformer 用使用硬路由的稀疏 MoE 層替換了 T5 Transformer 中的 FFN 層。
• 容量因子決定了每個(gè)專家允許的token數(shù)量，是在token丟棄和機(jī)器利用率之間權(quán)衡的一個(gè)手段。實(shí)驗(yàn)證明，即使以犧牲token為代價(jià)，最大化機(jī)器利用率也是正確的目標(biāo)。
• 由于其良好的可擴(kuò)展性特征，與T5（密集）Transformer模型相比，Switch Transformer在訓(xùn)練速度上實(shí)現(xiàn)了7倍的加速。

雖然Switch Transformer已經(jīng)在語言模型領(lǐng)域取得了突破，但我相信我們才剛剛開始看到它對行業(yè)的全面影響。稀疏 MoE 帶來的強(qiáng)大擴(kuò)展特性有可能為跨領(lǐng)域的應(yīng)用帶來巨大的建模改進(jìn)，而不僅僅局限于LLMs。這是機(jī)器學(xué)習(xí)領(lǐng)域激動(dòng)人心的時(shí)刻！文章來源地址http://www.zghlxwxcb.cn/news/detail-828169.html

到了這里，關(guān)于【譯】稀疏混合專家模型的崛起： Switch Transformers的文章就介紹完了。如果您還想了解更多內(nèi)容，請?jiān)谟疑辖撬阉鱐OY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！