為DeepSeekMoE模型帶來「免費午餐」加速,專家鏈可大幅提升LLM的資訊處理能力

MLNLP

社群是國內外知名的機器學習與自然語言處理社群,受眾覆蓋國內外NLP碩博生、高校老師以及企業研究人員。


社群的願景是促進國內外自然語言處理,機器學習學術界、產業界和廣大愛好者之間的交流和進步,特別是初學者同學們的進步。
來源 | 機器之心
作者 | 王子涵
我們都知道,DeepSeek-R1 的訓練過程使用了一種名為專家混合模型(Mixture-of-Experts, MoE)的技術,而當前的 MoE 技術依然還有顯著的最佳化空間。
近日,美國西北大學計算機科學博士生王子涵(Zihan Wang)等人在這個方向上取得了突破,提出了一種名為專家鏈(CoE)的技術。實驗表明,CoE 在效能、擴充套件策略、資源效率和專家使用效率等多個方面都顯著超越先前的 MoE 模型。
目前,核心研究者王子涵已經撰寫併發布 CoE 的中英文部落格並已經在 GitHub 開源相關程式碼,他也在部落格中表示研究論文將在近期釋出。以下為相關連結:
  • 程式碼:https://github.com/ZihanWang314/coe
  • 中文報告:https://sandy-server-87f.notion.site/1ab9bb750b79801bbfebf01ae9a77b3f
  • 英文報告:https://sandy-server-87f.notion.site/Chain-of-Experts-Unlocking-the-Communication-Power-of-MoEs-1ab9bb750b7980048d43e6aab3537cea
作者簡介:王子涵,美國西北大學計算機科學博士一年級學生,導師為 Manling Li。他本科畢業於中國人民大學高瓴人工智慧學院,並曾參與過 DeepSeek-V2 模型的開發工作。

引言

我們提出專家鏈 (Chain-of-Experts,CoE) 架構,一種透過在單層內實現專家間序列通訊的創新方法,從根本上改變稀疏神經網路的資訊處理方式。
MoE 設計中存在專家間獨立處理以及視訊記憶體需求高的問題。與先前 MoE 獨立處理每個 token 不同,CoE 引入迭代機制使專家能夠 "溝通",在其他專家的輸出之上處理 token。
實驗證明 CoE 在多個方面顯著超越先前 MoE 模型。效能顯著提升,CoE 兩次迭代在 Math 任務上將驗證 loss 從 1.20 降至 1.12;擴充套件策略優勢,CoE 中擴充套件兩次迭代效能相當於將專家選擇數量擴大 3 倍,並優於擴充套件模型層數;資源效率最佳化,相似表現下減少了 17.6-42% 的記憶體需求;專家組合自由度提升了 823 倍專家使用效率增強,促進了專家間的直接通訊和更好的專家分化。這些優勢構成了一種罕見的 "免費午餐" 效應,為大規模語言模型的高效擴充套件開闢了新途徑。

Chain-of-Experts:釋放 MoE 專家的溝通潛能

大型語言模型(LLMs)不斷推動人工智慧可能性的邊界,但高效擴充套件這些模型仍然是一個重大挑戰。專家混合(Mixture of Experts,MoE)模型作為一種有前景的方法出現,透過每個 token 僅啟用部分引數來解決這一挑戰,理論上實現更高效的擴充套件。然而,MoE 模型存在以下侷限性:
  1. 獨立標記處理:MoE 模型通常並行獨立處理 token,專家之間沒有溝通。
  2. 記憶體效率低下:由於具有稀疏啟用模式,MoE 整體引數數量較大,需要大量記憶體資源。
Chain-of-Experts (CoE) 介紹
我們的研究引入了 Chain-of-Experts (CoE),這是一種改變稀疏神經網路處理資訊方式的新方法。

Chain-of-Experts 的形式化表述

CoE 的關鍵創新在於建立溝通性處理機制,超越了先前 MoE 模型中的獨立標準處理。我們如何實現這一點?透過在單個層的迭代中將 MoE 輸出反饋為多次迭代的輸入。
1、先前 MoE 的輸出表示
在先前的 MoE 層中,不考慮共享專家,輸出可以表示為:
其中 x 是輸入,

 是門控值(路由函式分配給專家

 的權重),K 是每個 token 選擇的專家數量,N 是每層專家總數,

 是專家 i 的路由引數。

考慮共享專家時,路由函式計算不變,MoE 的輸出變為
其中 M 是共享專家的總數,

是共享專家。

2、CoE 的迭代處理機制
在 CoE 中,我們引入迭代處理機制,可以形式化表示為:
其中 

 是第 t 次迭代的中間表示,

是第 t 次迭代中專家 i 的門控值,C 是迭代的總次數,

是第 i 個專家,

 是殘差連線指示變數,當使用內殘差時

,否則為0。

考慮到共享專家,可以形式化表示為:
3、門控機制詳細說明
參考 DeepSeek-V2 的實現,我們定義門控機制為:
其中,我們為每次迭代使用獨立門控,採用不同的路由函式,K/C 是每次迭代選擇的專家數量(以保持總計算量不變),

 是第 t 次迭代中專家 i 的路由引數。

方法分析
1、CoE 的優勢
這種迭代式專家處理方法具有兩個關鍵優勢:
  1. 每次迭代的專家選擇由前一次迭代的輸出決定,形成專家間的依賴關係和更動態的路由機制
  2. 序列資訊可以在迭代過程中累積,實現專家間的直接通訊
透過這種方式,CoE 能夠在保持計算效率的同時,顯著提高模型的表現,特別是在複雜的場景中(如實驗採用的數學任務)。
2、路由函式如何影響決策過程
在 CoE 模型中,路由函式

對決策過程的影響是多層次的:

  1. 動態專家選擇在每次迭代 t 中, 

    基於前一次迭代的輸出

     來決定當前迭代中哪些專家應該被啟用以及它們的權重。

  2. 資訊流控制:路由函式控制了資訊如何在專家網路中流動,透過在每次迭代中重新評估路由決策,確保資訊能夠流向最相關的專家。
  3. 適應性處理:隨著迭代的進行,路由函式可以調整其路由策略,根據先前迭代的處理結果來最佳化後續專家的選擇。

實驗結果與發現

我們採取 DeepSeekV2 結構,在 500M 級別 MoE 下使用 32K Tok 的 batch size 訓練 1000 步,完成了一系列實驗來驗證 CoE 的有效性。單次訓練使用一張 H100 約為 30 分鐘,使用一張 4090 約為 2 小時。更多實驗細節見後。
我們的實驗方法命名為

,例如 CoE-2 (4/64) 代表使用 CoE,迴圈次數為 2,每次選取 4 個專家,每層一共有 64 個專家。我們在所有實驗中設定共享專家的數量為 1,在迭代次數為 2 的情況下,這會提高 5% 左右的計算需求。

效果優勢
1、 在算力預算與視訊記憶體預算相似時,CoE 效果明顯更優。
在 Math 任務上進行預訓練 (步數 1000),CoE-2 (4/64) 的效果明顯優於 MoE (8/64),在相似的算力和記憶體要求下,將 loss 從 1.20 下降至 1.12,且有更陡峭的下降趨勢。
我們進一步在 “dense”(專家 8 選 8)模型上也測試了 CoE,證明了序列處理在 Sparse MoE 上相比 Dense 模型更有效,CoE 是 (Fine-Grained) Sparse MoE 的專屬方法。如圖所示,採取 2 次序列化處理對 Dense 模型效能沒有顯著益處。
資源效率
1、在計算量和效果相似的情況下,CoE 可以減小對視訊記憶體的要求。
例如,CoE-2 (4/48) 的效果與 MoE (8/64) 相近,但使用更少的總專家數量。
如下圖所示,loss match 的情況下減小了 17.6% 的視訊記憶體需求。
(注,我們記錄了 GPU memory allocated,與 micro_bsz=1 的時候理論最小視訊記憶體相關,實際上因為實驗中使用了較大的 batch size,使用的視訊記憶體會更大)
2、在預算相似的情況下,CoE 是更好的計算擴充套件方式。
我們對比了 CoE 和其他計算擴充套件方法:拓展模型層數與拓展專家選擇個數。
a) 拓展迭代次數 (CoE) > 拓展模型層數
CoE-2 (8/64),4 層 vs MoE (8/64),8 層 / 12 層,8 層 MoE 和 CoE 效果幾乎相同,但是對 Memory 要求高 72%,即 CoE 相對節省了 42% memory。
b) 拓展迭代次數 (CoE)> 拓展專家選擇個數
CoE-2 (8/64),4 層 vs MoE (16/64) 或 MoE (24/64),4 層,Memory 和 Compute 要求一致,CoE 效果更好。
3、架構設計的關鍵發現
a. 獨立門控機制
獨立門控機制可以提升模型效能,能從一定程度上解釋專家的分化,表明同一專家在不同迭代階段處理不同型別的資訊。
我們實現了共享門控的 CoE 變種,使用共享的

 而非

,具體而言:

發現效果差於 CoE,甚至差於沒有拓展規模的 MoE,如下圖所示,體現獨立門控機制的重要性。
b. 殘差連線
內殘差連線 (inner residual) 比外殘差連線 (outer residual) 更有效,這表明每組專家對 token 的序列處理實際上是在提高模型的有效深度,每個專家的處理都在更好地學習殘差。
我們實現了外殘差連線的變種,將 

設定為 0,層內迭代處理結束後再連線殘差,具體而言:

發現效果差於 CoE 且與 MoE 相當,如下圖所示,體現內殘差連線的重要性。

理論觀察:免費午餐效應

最顯著的是,CoE 或許提供了我們稱之為 "免費午餐" 加速。透過重構資訊在模型中的流動方式,與以往 MoE 方法相比,我們以更少的計算開銷實現了更好的結果。我們認為這種效果可能來自於三個因素:
1. 專家選擇的自由度提高。

,選擇的可能性變多

2. CoE 統一了序列處理 (sequential processing) 和專家通訊 (expert communication) 兩個概念:
  1. 不同專家之間可以序列處理,提高了 Transformer 的有效深度。兩次前向傳播送到不同專家時,增強了專家的序列特徵 — 處理同一個 token 的路由專家數相同,但是是序列地處理。
  2. 一個專家在迭代過程中有機會處理一個 token 多次,可能有助於促進專家的分化:一個專家與其他專家合作處理後,可以在後續迭代對 token 進行補充處理。
實驗細節
1、資料集
實驗使用的資料集為 MetaMathQA (連結:https://huggingface.co/datasets/meta-math/MetaMathQA),該資料為從 GSM8K 和 MATH 資料集中增強而來,不含網路獲取的通用資料。
2、模型配置
我們基於 DeepSeek-V2-Lite 架構,設定總引數為 544MB(不含 embedding)的模型,實現了 CoE 方法,主要配置引數如下:
  • 基礎模型引數:隱藏層大小:1024,隱藏層數量:4,注意力頭數:8
  • MoE 相關引數:路由專家總數:63,共享專家數量:1,每個 token 選擇的專家數量:8 (routed) + 1 (Shared),MoE 中間層大小:704,MoE 層頻率:1 (每層都是 MoE 層)
  • CoE 特有引數:迭代次數:2,殘差連線方式:內部連線,獨立門控機制:開啟 (每次迭代使用不同的路由函式)
3、訓練設定
  • 批次大小 (batch size):64,序列長度:512,訓練步數:1000,最佳化器:AdamW,學習率:3e-4,學習率排程:10% Warmup,betas: [0.9, 0.95]   weight_decay: 0.01   warmup_steps_ratio: 0.1   clip_grad: 1.0。
4、系統設定
訓練採取改版後的 veRL 架構(微調了原版實現),程式碼參考 https://github.com/ZihanWang314/coe,實驗都在多個單張 H100 的伺服器上進行。每組實驗在 1 小時以內完成。

結論

Chain-of-Experts 代表了高效、高效能語言模型發展的重要一步。透過在稀疏神經網路中實現溝通處理,CoE 解決了當前 MoE 架構中的基本限制,同時以更少的計算需求提供優越的效能。
這項研究為高效擴充套件語言模型開闢了新途徑,可能使先進的人工智慧能力更加普惠和可持續。
未來工作
1. 擴充套件法則研究:提高模型尺寸、預訓練步數、批次大小,測試不同資料集等。目前使用數學資料集是因為其推理複雜性具有挑戰性,但計劃在通用資料集上進行測試。
2. 進一步評估:
  • 測試模型的實際任務表現而非僅限於評估損失
  • 擴充套件迴圈次數:目前只測試了迴圈次數 = 2,探索 3、4 次等是否有效
  • 進一步評估共享專家對 CoE 的作用
3. 架構創新:
  • 基於迴圈告訴我們一個專家可以使用多次的原理,探索是否可以實現所有層共享專家的 MoE Transformer,並分層選擇專家
  • 例如,DeepSeekV2Lite 有 27 層,每層 64 個專家,目前每層選 6 個路由專家有約 7×10⁷ 種組合,而如果每一層過 MoE 時,都從全部 27×64 個專家中選擇,一共將有 3×10¹⁶ 種可能的組合
侷限性
  1. 雖然理論上 TFLOPs 保持一致,但實際訓練時間會略微增加,因為單次迭代選擇更少專家減小了矩陣乘法並行度。未來工作可能需要進行底層適配。
  2. 該方法需要從頭預訓練而不能簡單地適配現有模型,限制了大眾研究者的應用。
  3. 模型更大時,多節點 EP 通訊開銷也需要納入考慮;當前架構基於 DeepSeekV2,與 V3 的模型架構相容,但 V3 中實現的 MTP、FP8 等訓練架構尚未測試,相容性未知。
這篇部落格文章基於我們關於 Chain-of-Experts (CoE) 的研究論文。有關更多技術細節和實驗結果,請關注我們即將釋出的完整論文。

貢獻者

Zihan Wang*
Rui Pan*
Lu Yin*
Manling Li*
Shiwei Liu*
* 共同作者。ZW 和 SL 發起了專案。LY 探索了 MoE 內層處理機制。ZW 提出了序列專家和 CoE 概念,實現了程式碼並進行了實驗。SL 提出了獨立門控機制。RP 探索了層間重要性並貢獻了資料集。ML 進行了實驗指導並提供了雲計算服務。ZW 和 ML 準備了報告和視覺化材料。
參考文獻
https://arxiv.org/pdf/2309.08520 MoE 縮放法則
https://arxiv.org/pdf/2101.03961 Switch Transformer
https://arxiv.org/pdf/2502.05171 迴圈深度
https://arxiv.org/pdf/1909.11942 ALBERT (密集模型上的迴圈)
https://arxiv.org/abs/2401.06066 DeepSeekMoE
https://arxiv.org/abs/2405.04434 DeepSeek V2
https://arxiv.org/abs/2412.19437 DeepSeek V3
https://arxiv.org/abs/1512.03385 ResNet
技術交流群邀請函
△長按新增小助手
掃描二維碼新增小助手微信
請備註:姓名-學校/公司-研究方向
(如:小張-哈工大-對話系統)
即可申請加入自然語言處理/Pytorch等技術交流群

關於我們

MLNLP 社群是由國內外機器學習與自然語言處理學者聯合構建的民間學術社群,目前已經發展為國內外知名的機器學習與自然語言處理社群,旨在促進機器學習,自然語言處理學術界、產業界和廣大愛好者之間的進步。
社群可以為相關從業者的深造、就業及研究等方面提供開放交流平臺。歡迎大家關注和加入我們。

相關文章