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

比傳統MoE推理速度更快、性能更高的新一代架構，來了！

這個通用架構叫做MoE++，由顏水成領銜的崑崙萬維2050研究院與北大袁粒團隊聯合提出。

總體來看，MoE++的創新之處在於引入了“零計算量專家”，得益於這個設計，使得新架構有了三大優勢：

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

實驗結果表明，在0.6B到7B參數規模的LLMs上，MoE++在相同模型大小的情況下，相比傳統MoE，性能更優，同時實現了1.1到2.1倍的專家吞吐速度。

並且這個模型權重也已開源！

那麼MoE++具體是如何做到的，我們繼續往下看。

現有的大多數混合專家（MoE）方法通常爲所有Token激活固定數量的FFN專家。

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

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

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

基於這一見解，研究團隊認爲，現有MoE使用的固定混合機制可能導致訓練和推理效率下降，同時限制模型性能。

爲了在提升速度的同時增強性能，研究團隊提出了一種通用的異構MoE框架，稱爲MoE++。

具體來說，團隊引入了三種零計算量專家：

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

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

這種異構結構通過擴大網絡的組合空間，提升了模型的擬合能力，並顯著降低了計算成本。

此外，研究團隊還將前一層的路由分數整合到當前層的專家選擇中，使Token在選擇專家時能夠參考其先前的路由路徑，從而實現更穩定的專家分配。

研究團隊認爲，新設計的MoE架構應滿足以下標準：

在這些原則的指導下，研究團隊引入了零計算量專家，每個專家僅執行最基本的操作。

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

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

Zero專家

最簡單的零計算量專家是丟棄當前輸入的Zero專家。

本質上，Zero專家的存在可以將Top-2 MoE++層降級爲Top-1 MoE++層。

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

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

Copy專家

Copy專家直接將輸入作爲輸出，直觀上相當於跳過當前MoE++層。

具體而言，當輸入Token與現有專家的匹配較差時，選擇繞過MoE++層可能更爲有利。

Constant專家

Constant專家通過可訓練向量替換輸入Token。

然而，完全替換會導致輸入Token信息的丟失。

爲此，研究團隊引入了可訓練的權重矩陣，用於動態預測替換的比例。由於Constant專家的計算開銷極小，因此仍被歸類爲零計算量專家。

路由分數殘差

由於MoE++包含異構專家，因此與普通MoE相比，路由器的設計變得更加關鍵。

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

具體來說，MoE++將前一層的路由分數通過一個可學習的轉換矩陣合併到當前層的專家選擇中。

這些路由分數殘差使每個Token在選擇專家時可以考慮其先前的路由路徑。

對於這個問題，主要可以總結三點原因。

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

MoE++通過爲簡單Token分配較少的FFN專家，優化了計算資源的分配，從而使更多FFN專家能夠專注於處理更具挑戰性的Token。

正如圖3所示，研究團隊發現，在MoE++中，動詞激活的FFN專家數量最多，其次是名詞，而拆分後的詞片激活的FFN數量最少。

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

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

其次是穩定的路由。

MoE++將前一層的路由分數合併到當前層的專家選擇中。

這些路由分數殘差使每個Token在選擇專家時考慮其先前的路由路徑。

如圖4所示，路由分數殘差有效地建立了不同MoE++層之間的聯繫，減小了路由分數的方差。

同時，路由分數殘差不改變路由分數的均值和取值範圍。因此，路由分數殘差有助於在MoE++中實現異構專家架構的穩定路由。

△圖4：路由分數殘差對路由分數分佈的影響

最後是更低的計算複雜度。

如下表所示，MoE++具有比普通MoE更低的理論計算複雜度。

從0.6B的參數量逐漸擴展到7B參數量的大量實驗結果表明，MoE++方法明顯優於普通MoE方法。

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

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

研究人員也將MoE++模型的訓練預算擴展到1T Tokens。

研究人員發現MoE++模的性能與具有2到3倍激活參數的稠密模型相當。

值得注意的是，MoE++優於OpenMoE-8B/32E，這是一個從零開始訓練的更大的MoE模型，使用更多的訓練Tokens。

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

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

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

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

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