DeepSeek前實習生魔改MoE,用迭代機制把記憶體需求砍了42%,團隊:“免費午餐”最佳化方法

MLNLP社群是國內外知名的機器學習與自然語言處理社群,受眾覆蓋國內外NLP碩博生、高校老師以及企業研究人員。
社群的願景是促進國內外自然語言處理,機器學習學術界、產業界和廣大愛好者之間的交流和進步,特別是初學者同學們的進步。
來源 | 量子位
作者 | 西風
DeepSeek MoE“變體”來了,200美元以內,記憶體需求減少17.6-42%!
名叫CoE(Chain-of-Experts),被認為是一種“免費午餐”最佳化方法,突破了MoE並行獨立處理token、整體引數數量較大需要大量記憶體資源的侷限。
與並行處理不同,CoE使專家能在單層內序列通訊,形成一種迭代機制,即專家能“溝通”,在其它專家輸出之上處理token。
研究團隊在實驗中發現,經過2次迭代的CoE,在相同的計算預算下將數學任務的驗證損失從1.20降低至1.12,僅僅透過重構資訊流就獲得了效能提升。
透過擴充套件CoE的迭代次數,在效能相當的情況下,記憶體使用比透過增加模型層數或擴充套件專家選擇數量的方法降低了17.6-42%。
另外,在專家組合自由度、專家使用效率等其它方面,CoE也都具有顯著優勢,專家組合增加823倍。
目前,研究團隊曬出了CoE技術Blog(完整論文即將釋出),引起不少網友圍觀。
翻看作者主頁,還發現作者Zihan Wang真曾在DeepSeek實習過😯
有網友看過這項研究表示:
MoEs中的IsoFLOP層迭代設計,非常奈斯。
還有網友已經開始預測下一代架構了。

CoE究竟長啥樣?以下是團隊釋出的Notion Blog介紹。

CoE專門針對稀疏MoE打造

CoE關鍵創新在於建立溝通性處理機制,改變了稀疏神經網路的資訊處理方式。
具體來說,是透過在單個層的迭代中將MoE輸出反饋為多次迭代的輸入來實現的。
CoE迭代處理機制可以形式化表示為下面這個樣嬸兒:
參考DeepSeek-V2的實現,研究團隊定義門控機制為如下:
團隊介紹,這種設計的好處在於每次迭代的專家選擇由前一次迭代的輸出決定,形成專家間的依賴關係和更動態的路由機制
而且序列資訊可以在迭代過程中累積,實現專家間的直接通訊。
實驗採取DeepSeek V2架構,在引數規模為500M的MoE模型上使用32K Tok的batch size訓練1000步,以此來驗證CoE的有效性。
結果CoE在效能、擴充套件策略、資源效率最佳化、專家組合自由度、專家使用效率方面具有顯著優勢。
除了開頭所展示的在相似的算力和記憶體要求下,CoE將loss從1.20下降至1.12,且有更陡峭的下降趨勢。
團隊進一步在“dense”(專家8選8)模型上也進行了測試,結果證明了序列處理在Sparse MoE上相比Dense模型更有效,CoE是一種專為(細粒度)稀疏混合專家模型(Sparse MoE)設計的方法
採取2次序列化處理並不能顯著提升Dense模型效能。
另外,在計算量和效果相似的情況下,CoE可以減小對記憶體的要求。如下,CoE-2(4/48)的效果與MoE(8/64)相近,但使用更少的總專家數量。loss match的情況下減小了17.6%的記憶體需求。
團隊還對比了在預算相似的情況下,擴充套件CoE迭代次數和擴充套件模型層數、擴充套件專家選擇個數的效果,結果擴充套件CoE迭代次數更優。
CoE-2(8/64),4層 vs MoE(8/64),8層/12層,8層MoE和CoE效果幾乎相同,但是對記憶體要求高72%,即CoE相對節省了42%記憶體。
團隊強調,獨立門控機制內殘差連線是CoE的關鍵架構創新,消融研究表明,移除任何元件都會顯著降低效能。
更多細節,感興趣的童鞋可以檢視技術報告原文~

誰造的?

CoE由一個5人組成的團隊提出。
Zihan Wang是美國西北大學計算機科學專業博士生,本科畢業於中國人民大學高瓴人工智慧學院,研究聚焦於基礎模型的自主性、效率以及長上下文理解。
Zihan Wang曾在DeepSeek工作過,是ESFT(Expert-Specialized Fine-Tuning)的論文一作。
ESFT透過僅調整與任務相關的部分高效地定製採用MoE,從而在減少資源和儲存使用的同時提升效率和效能。
CoE不是Zihan Wang第一次針對DeepSeek搞的“變體”。
之前他還曾基於verl復現了DeepSeek-R1(-Zero)框架——RAGEN (Reinforcement learning AGENt),在GitHub攬星近1k:
Zihan Wang師從Manling Li。Manling Li是西北大學計算機科學系助理教授,此前曾在吳佳俊教授的指導下工作,並得到李飛飛教授的指導。
RAGEN的貢獻者名單中也有Manling Li、吳佳俊、李飛飛的身影。
CoE技術報告:http://sandy-server-87f.notion.site/Chain-of-Experts-Unlocking-the-Communication-Power-of-MoEs-1ab9bb750b7980048d43e6aab3537ceaCoE GitHub連結:https://github.com/ZihanWang314/coe參考連結:[1]https://x.com/wzihanw/status/1896601518612021709[2]https://github.com/ZihanWang314/RAGEN?tab=readme-ov-file[3]https://github.com/deepseek-ai/ESFT

技術交流群邀請函

△長按新增小助手
掃描二維碼新增小助手微信
請備註:姓名-學校/公司-研究方向
(如:小張-哈工大-對話系統)
即可申請加入自然語言處理/Pytorch等技術交流群

關於我們

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


相關文章