顏水成袁粒提出新一代MoE架構(gòu)：專家吞吐速度最高提升2.1倍！

十三 2024-10-21 10:07:06 來源：量子位

北京大學(xué) 投稿

量子位 | 公眾號 QbitAI

比傳統(tǒng)MoE推理速度更快、性能更高的新一代架構(gòu)，來了！

這個通用架構(gòu)叫做MoE++，由顏水成領(lǐng)銜的昆侖萬維2050研究院與北大袁粒團(tuán)隊聯(lián)合提出。

總體來看，MoE++的創(chuàng)新之處在于引入了“零計算量專家”，得益于這個設(shè)計，使得新架構(gòu)有了三大優(yōu)勢：

• 降低計算成本：MoE++允許每個Token使用可變數(shù)量的FFN專家，甚至可以完全跳過當(dāng)前的MoE層。
• 提升性能：通過減少簡單Token所需的FFN專家數(shù)量，MoE++使更多專家能夠?qū)Ｗ⒂趶?fù)雜的Token，釋放出比傳統(tǒng)MoE更大的性能潛力。
• 零計算量專家的參數(shù)極小：可以在每個GPU上同時部署所有的零計算量專家，避免了分布式FFN專家部署帶來的通信開銷和專家負(fù)載不均的問題。

除此之外，MoE++還讓每個Token在選擇專家時參考前一層的路由路徑。

實驗結(jié)果表明，在0.6B到7B參數(shù)規(guī)模的LLMs上，MoE++在相同模型大小的情況下，相比傳統(tǒng)MoE，性能更優(yōu)，同時實現(xiàn)了1.1到2.1倍的專家吞吐速度。

并且這個模型權(quán)重也已開源！

那么MoE++具體是如何做到的，我們繼續(xù)往下看。

MoE++是如何做到的？

現(xiàn)有的大多數(shù)混合專家（MoE）方法通常為所有Token激活固定數(shù)量的FFN專家。

在許多研究中，每個Token會選擇Top-2 FFN專家，并將它們的輸出加權(quán)合成為下一層的輸入。然而，并非所有Token的預(yù)測難度都是相同的。

例如，像逗號等簡單符號可能只需要一個FFN專家來處理。

甚至在某些特殊情況下，某些Token如果與當(dāng)前MoE層的專家不匹配，繞過該層而不選擇Top-2 FFN專家反而可能更為高效。

基于這一見解，研究團(tuán)隊認(rèn)為，現(xiàn)有MoE使用的固定混合機(jī)制可能導(dǎo)致訓(xùn)練和推理效率下降，同時限制模型性能。

為了在提升速度的同時增強(qiáng)性能，研究團(tuán)隊提出了一種通用的異構(gòu)MoE框架，稱為MoE++。

具體來說，團(tuán)隊引入了三種零計算量專家：

• Zero專家，輸出空向量
• Copy專家，將輸入直接作為輸出
• Constant專家，用可訓(xùn)練的向量替代輸入作為輸出。

如圖1所示，與傳統(tǒng)MoE方法不同，MoE++允許每個Token使用可變數(shù)量的FFN專家，接受恒定向量的替換，甚至完全跳過當(dāng)前的MoE++層。

△圖1：MoE++和普通MoE的對比

這種異構(gòu)結(jié)構(gòu)通過擴(kuò)大網(wǎng)絡(luò)的組合空間，提升了模型的擬合能力，并顯著降低了計算成本。

此外，研究團(tuán)隊還將前一層的路由分?jǐn)?shù)整合到當(dāng)前層的專家選擇中，使Token在選擇專家時能夠參考其先前的路由路徑，從而實現(xiàn)更穩(wěn)定的專家分配。

研究團(tuán)隊認(rèn)為，新設(shè)計的MoE架構(gòu)應(yīng)滿足以下標(biāo)準(zhǔn)：

1. 設(shè)計應(yīng)盡量簡化，以高效處理簡單的Token
2. 為了確保公平比較，新增參數(shù)應(yīng)保持在可忽略的范圍

在這些原則的指導(dǎo)下，研究團(tuán)隊引入了零計算量專家，每個專家僅執(zhí)行最基本的操作。

如圖2（a）所示，團(tuán)隊設(shè)計了三種類型的零計算量專家：Zero專家、Copy專家和Constant專家，分別對應(yīng)丟棄、跳過和替換操作。

△圖2：MoE++的核心組成部分

Zero專家

最簡單的零計算量專家是丟棄當(dāng)前輸入的Zero專家。

本質(zhì)上，Zero專家的存在可以將Top-2 MoE++層降級為Top-1 MoE++層。

具體來說，當(dāng)Zero專家被激活時，Top-2 MoE++層的輸出將等同于另一個專家的單獨輸出。

這樣，Zero專家的引入提升了模型在處理簡單Token和復(fù)雜Token時的靈活性。

Copy專家

Copy專家直接將輸入作為輸出，直觀上相當(dāng)于跳過當(dāng)前MoE++層。

具體而言，當(dāng)輸入Token與現(xiàn)有專家的匹配較差時，選擇繞過MoE++層可能更為有利。

Constant專家

Constant專家通過可訓(xùn)練向量替換輸入Token。

然而，完全替換會導(dǎo)致輸入Token信息的丟失。

為此，研究團(tuán)隊引入了可訓(xùn)練的權(quán)重矩陣，用于動態(tài)預(yù)測替換的比例。由于Constant專家的計算開銷極小，因此仍被歸類為零計算量專家。

路由分?jǐn)?shù)殘差

由于MoE++包含異構(gòu)專家，因此與普通MoE相比，路由器的設(shè)計變得更加關(guān)鍵。

為此，如圖2（b）所示，研究團(tuán)隊提出了一種路徑感知路由器，它在選擇合適的專家時考慮了前一層所采用的路徑。

具體來說，MoE++將前一層的路由分?jǐn)?shù)通過一個可學(xué)習(xí)的轉(zhuǎn)換矩陣合并到當(dāng)前層的專家選擇中。

這些路由分?jǐn)?shù)殘差使每個Token在選擇專家時可以考慮其先前的路由路徑。

為什么MoE++比MoE更好（Why）？

對于這個問題，主要可以總結(jié)三點原因。

首先就是靈活的計算量分配。

MoE++通過為簡單Token分配較少的FFN專家，優(yōu)化了計算資源的分配，從而使更多FFN專家能夠?qū)Ｗ⒂谔幚砀咛魬?zhàn)性的Token。

正如圖3所示，研究團(tuán)隊發(fā)現(xiàn)，在MoE++中，動詞激活的FFN專家數(shù)量最多，其次是名詞，而拆分后的詞片激活的FFN數(shù)量最少。

這表明，MoE++能夠讓語義較少的Token使用更少的FFN專家，從而釋放更多專家來處理語義更豐富的Token。因此，MoE++不僅減少了計算開銷，還提升了整體性能。

△圖3：在MoE++中不同Token所需要的平均FFN專家數(shù)量

其次是穩(wěn)定的路由。

MoE++將前一層的路由分?jǐn)?shù)合并到當(dāng)前層的專家選擇中。

這些路由分?jǐn)?shù)殘差使每個Token在選擇專家時考慮其先前的路由路徑。

如圖4所示，路由分?jǐn)?shù)殘差有效地建立了不同MoE++層之間的聯(lián)系，減小了路由分?jǐn)?shù)的方差。

同時，路由分?jǐn)?shù)殘差不改變路由分?jǐn)?shù)的均值和取值范圍。因此，路由分?jǐn)?shù)殘差有助于在MoE++中實現(xiàn)異構(gòu)專家架構(gòu)的穩(wěn)定路由。

△圖4：路由分?jǐn)?shù)殘差對路由分?jǐn)?shù)分布的影響

最后是更低的計算復(fù)雜度。

如下表所示，MoE++具有比普通MoE更低的理論計算復(fù)雜度。

實驗結(jié)果

從0.6B的參數(shù)量逐漸擴(kuò)展到7B參數(shù)量的大量實驗結(jié)果表明，MoE++方法明顯優(yōu)于普通MoE方法。

與相同大小的普通MoE模型相比，MoE++的專家吞吐量提高了約15% ~ 111%，同時具有更高的性能。

現(xiàn)有的LLMs模型通常需要大量的訓(xùn)練預(yù)算，比如OpenMoE-8B/32E使用1.1T Tokens，TinyLlama-1.1B使用3T Tokens。

研究人員也將MoE++模型的訓(xùn)練預(yù)算擴(kuò)展到1T Tokens。

研究人員發(fā)現(xiàn)MoE++模的性能與具有2到3倍激活參數(shù)的稠密模型相當(dāng)。

值得注意的是，MoE++優(yōu)于OpenMoE-8B/32E，這是一個從零開始訓(xùn)練的更大的MoE模型，使用更多的訓(xùn)練Tokens。

這些結(jié)果表明，MoE++框架是一種很有前途的LLMs框架方案。

任務(wù)級專家負(fù)載分布的可視化

研究人員還探索了MoE++模型中跨不同任務(wù)的專家負(fù)載分。這些可視化揭示了幾個有趣的發(fā)現(xiàn):

1. 專家負(fù)載在不同層之間存在相關(guān)性，特別是在相鄰層之間。例如，當(dāng)?shù)趈層激活很大比例的FFN專家時，第j + 1層很可能也會以同樣大的比例激活FFN專家。
2. 與中間層相比，淺層和最后一層的專家分配模式在不同任務(wù)之間的差異更大。這表明該模型主要通過其淺層和最終層而不是中間層來適應(yīng)不同的任務(wù)。未來的工作可以集中在這些層中設(shè)計更復(fù)雜的結(jié)構(gòu)，以增強(qiáng)模型對不同任務(wù)的適應(yīng)性。
3. 不同任務(wù)中每個Token激活的FFN專家數(shù)量存在顯著差異，但并不一定是更簡單的任務(wù)激活更少的FFN專家。例如，ARC Challenge任務(wù)通常比ARC Easy任務(wù)激活更多的FFN專家。這些結(jié)果表明，MoE++模型根據(jù)知識內(nèi)容和Token級別的復(fù)雜性來分配專家，而不是根據(jù)整體任務(wù)難度來分配專家。
4. 在所有專家類型中，Zero專家的平均激活次數(shù)最高，更簡單的任務(wù)顯示出更高的平均激活次數(shù)。例如，ARC Easy任務(wù)比ARC Challenge任務(wù)激活更多的零專家。這表明Zero專家的激活水平可能可以作為模型任務(wù)難度的一個指標(biāo)。
5. 在MoE++模型的所有層中，不同任務(wù)主題的專家分配差異顯著，這表明MoE++模型通過采用不同的專家分配模式來處理不同主題的任務(wù)。

論文地址：
https://arxiv.org/abs/2410.07348

GitHub地址：
https://github.com/SkyworkAI/MoE-plus-plus

Huggingface地址：
https://huggingface.co/Chat-UniVi/MoE-Plus-Plus-7B