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

作者：量子位 2024-10-21 12:30:00

人工智能新聞

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

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

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

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

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

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

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

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

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

MoE++是如何做到的？

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

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

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

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

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

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

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

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

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

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

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

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

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

設計應盡量簡化，以高效處理簡單的Token
為了確保公平比較，新增參數應保持在可忽略的范圍

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

如圖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++比MoE更好（Why）？

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

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

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++模型中跨不同任務的專家負載分。這些可視化揭示了幾個有趣的發現:

專家負載在不同層之間存在相關性，特別是在相鄰層之間。例如，當第j層激活很大比例的FFN專家時，第j + 1層很可能也會以同樣大的比例激活FFN專家。
與中間層相比，淺層和最后一層的專家分配模式在不同任務之間的差異更大。這表明該模型主要通過其淺層和最終層而不是中間層來適應不同的任務。未來的工作可以集中在這些層中設計更復雜的結構，以增強模型對不同任務的適應性。
不同任務中每個Token激活的FFN專家數量存在顯著差異，但并不一定是更簡單的任務激活更少的FFN專家。例如，ARC Challenge任務通常比ARC Easy任務激活更多的FFN專家。這些結果表明，MoE++模型根據知識內容和Token級別的復雜性來分配專家，而不是根據整體任務難度來分配專家。
在所有專家類型中，Zero專家的平均激活次數最高，更簡單的任務顯示出更高的平均激活次數。例如，ARC Easy任務比ARC Challenge任務激活更多的零專家。這表明Zero專家的激活水平可能可以作為模型任務難度的一個指標。
在MoE++模型的所有層中，不同任務主題的專家分配差異顯著，這表明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

責任編輯：張燕妮來源：量子位

模型框架