近年來(lái)隨著大語(yǔ)言模型的爆火，旨在取代 Transformer 的高效模型架構(gòu)及其預(yù)訓(xùn)練成為大模型領(lǐng)域的研究熱點(diǎn)，主要包括線性序列建模（如 Linear Attention、SSM、Linear RNN 等）和混合專(zhuān)家（Mixture-of-Experts, MoE）兩部分。這兩部分分別都有了長(zhǎng)足的進(jìn)步，但兩者的結(jié)合卻鮮少有人研究，兩者結(jié)合后的 Linear-MoE 架構(gòu)開(kāi)源實(shí)現(xiàn)更是完全缺失。

值得一提的是，近期廣受好評(píng)的 MiniMax-01 模型（使用 Lightning Attention-MoE）和騰訊混元 TurboS 模型（使用 Mamba2-MoE）均屬于 Linear-MoE 架構(gòu)。

來(lái)自上海人工智能實(shí)驗(yàn)室團(tuán)隊(duì)的最新成果 Linear-MoE，首次系統(tǒng)性地實(shí)現(xiàn)了線性序列建模與 MoE 的高效結(jié)合，并開(kāi)源了完整的技術(shù)框架，包括 Modeling 和 Training 兩大部分，并支持層間混合架構(gòu)。為下一代基礎(chǔ)模型架構(gòu)的研發(fā)提供了有價(jià)值的工具和經(jīng)驗(yàn)。

• 論文標(biāo)題：Linear-MoE: Linear Sequence Modeling Meets Mixture-of-Experts
• 論文地址：https://arxiv.org/abs/2503.05447
• 項(xiàng)目地址：https://github.com/OpenSparseLLMs/Linear-MoE

線性序列建模的崛起

過(guò)去兩年，線性序列建模技術(shù)取得了顯著進(jìn)展，其核心優(yōu)勢(shì)在于線性時(shí)間復(fù)雜度的訓(xùn)練和恒定內(nèi)存占用的推理。

這類(lèi)模型主要分為三大類(lèi)：線性注意力（Linear Attention）、狀態(tài)空間模型（SSM）和線性 RNN（Linear RNN），代表性工作包括 Lightning Attention、GLA、Mamba2、RWKV 等。

已有研究工作表明，這些模型實(shí)際上可以通過(guò)統(tǒng)一的遞歸形式進(jìn)行表達(dá)，如下表所示。這也反映出，盡管三類(lèi)方法分別出自不同的技術(shù)流派，但已逐漸收斂至統(tǒng)一的表達(dá)形式。

混合專(zhuān)家 MoE 成為事實(shí)標(biāo)準(zhǔn)

另外一方面，從國(guó)際上的 GPT-4 系列、Gemini 系列、Claude 系列到國(guó)內(nèi)的 DeepSeek 系列、Qwen 系列、騰訊混元 LLM、字節(jié)豆包、MiniMax-01、Moonshot-Kimi 等，都在步伐一致地 All in MoE。其重要性不言而喻，本文不做過(guò)多展開(kāi)。

Linear-MoE：模型架構(gòu)與高效訓(xùn)練

Linear-MoE 的核心貢獻(xiàn)在于構(gòu)建了一個(gè)從 Modeling 到 Training 的完整系統(tǒng)，支持線性序列建模層與 MoE 層的靈活組合，同時(shí)兼容傳統(tǒng)的 Softmax Attention Transformer 層，支持形成混合架構(gòu)。其設(shè)計(jì)亮點(diǎn)包括：

• 模塊化架構(gòu)：

• LSM 層（線性序列建模層）：支持各類(lèi)線性序列建模方法（如 Lightning Attention、Gated-DeltaNet、Mamba2 等）。
• MoE 層：集成多種 MoE 實(shí)現(xiàn)（如 Qwen-MoE、DeepSeek-MoE、Mixtral-MoE），以及一種 Dense 實(shí)現(xiàn)（Llama3）。
  

• 高效訓(xùn)練技術(shù)：

• 基于 Megatron-Core 框架開(kāi)發(fā)，確保系統(tǒng)穩(wěn)定性和可擴(kuò)展性。
• 支持張量并行、流水線并行、專(zhuān)家并行、LASP（線性注意力序列并行）和 MegaBlocks 等優(yōu)化技術(shù)，顯著提升訓(xùn)練效率。

實(shí)驗(yàn)驗(yàn)證

大規(guī)模實(shí)驗(yàn)驗(yàn)證了 Linear-MoE 的優(yōu)越性：

• 訓(xùn)練穩(wěn)定：混合模型（線性 + Softmax Attention）比純線性模型表現(xiàn)出更穩(wěn)定的訓(xùn)練曲線。
• 訓(xùn)練效率：借助專(zhuān)家并行和 MegaBlocks，系統(tǒng)在超大參數(shù)規(guī)模下仍保持高吞吐量。
• 推理優(yōu)勢(shì)：線性模型的推理速度比傳統(tǒng)架構(gòu)快 2-5 倍，內(nèi)存占用降低 50% 以上。
• 性能擴(kuò)展：在不同規(guī)模的基準(zhǔn)測(cè)試中，Linear-MoE 展現(xiàn)出良好的性能線性增長(zhǎng)。

開(kāi)源生態(tài)與未來(lái)展望

目前 Linear-MoE 已全面開(kāi)源，支持多種主流的線性序列建模方法和 MoE 實(shí)現(xiàn)。這一工作不僅填補(bǔ)了線性建模與 MoE 結(jié)合的技術(shù)空白，還為社區(qū)提供了可復(fù)現(xiàn)的高效訓(xùn)練方案。未來(lái)將進(jìn)一步探索 Linear-MoE 在長(zhǎng)上下文理解、Vision-Language 模型架構(gòu)中的應(yīng)用潛力。