机器之心
发布于

Linear-MoE:线性注意力遇上混合专家的开源实践

近年来随着大语言模型的爆火,旨在取代 Transformer 的高效模型架构及其预训练成为大模型领域的研究热点,主要包括线性序列建模(如 Linear Attention、SSM、Linear RNN 等)和混合专家(Mixture-of-Experts, MoE)两部分。这两部分分别都有了长足的进步,但两者的结合却鲜少有人研究,两者结合后的 Linear-MoE 架构开源实现更是完全缺失。

值得一提的是,近期广受好评的 MiniMax-01 模型(使用 Lightning Attention-MoE)和腾讯混元 TurboS 模型(使用 Mamba2-MoE)均属于 Linear-MoE 架构。

来自上海人工智能实验室团队的最新成果 Linear-MoE,首次系统性地实现了线性序列建模与 MoE 的高效结合,并开源了完整的技术框架,包括 Modeling 和 Training 两大部分,并支持层间混合架构。为下一代基础模型架构的研发提供了有价值的工具和经验。

  • 论文标题:Linear-MoE: Linear Sequence Modeling Meets Mixture-of-Experts

  • 论文地址:https://arxiv.org/abs/2503.05447

  • 项目地址:https://github.com/OpenSparseLLMs/Linear-MoE

线性序列建模的崛起

过去两年,线性序列建模技术取得了显著进展,其核心优势在于线性时间复杂度的训练和恒定内存占用的推理。

这类模型主要分为三大类:线性注意力(Linear Attention)、状态空间模型(SSM)和线性 RNN(Linear RNN),代表性工作包括 Lightning Attention、GLA、Mamba2、RWKV 等。

已有研究工作表明,这些模型实际上可以通过统一的递归形式进行表达,如下表所示。这也反映出,尽管三类方法分别出自不同的技术流派,但已逐渐收敛至统一的表达形式。

混合专家 MoE 成为事实标准

另外一方面,从国际上的 GPT-4 系列、Gemini 系列、Claude 系列到国内的 DeepSeek 系列、Qwen 系列、腾讯混元 LLM、字节豆包、MiniMax-01、Moonshot-Kimi 等,都在步伐一致地 All in MoE。其重要性不言而喻,本文不做过多展开。

Linear-MoE:模型架构与高效训练

Linear-MoE 的核心贡献在于构建了一个从 Modeling 到 Training 的完整系统,支持线性序列建模层与 MoE 层的灵活组合,同时兼容传统的 Softmax Attention Transformer 层,支持形成混合架构。其设计亮点包括:

  • 模块化架构:

  1. LSM 层(线性序列建模层):支持各类线性序列建模方法(如 Lightning Attention、Gated-DeltaNet、Mamba2 等)。

  2. MoE 层:集成多种 MoE 实现(如 Qwen-MoE、DeepSeek-MoE、Mixtral-MoE),以及一种 Dense 实现(Llama3)。

  • 高效训练技术:

  1. 基于 Megatron-Core 框架开发,确保系统稳定性和可扩展性。

  2. 支持张量并行、流水线并行、专家并行、LASP(线性注意力序列并行)和 MegaBlocks 等优化技术,显著提升训练效率。

实验验证

大规模实验验证了 Linear-MoE 的优越性:

  • 训练稳定:混合模型(线性 + Softmax Attention)比纯线性模型表现出更稳定的训练曲线。

  • 训练效率:借助专家并行和 MegaBlocks,系统在超大参数规模下仍保持高吞吐量。

  • 推理优势:线性模型的推理速度比传统架构快 2-5 倍,内存占用降低 50% 以上。

  • 性能扩展:在不同规模的基准测试中,Linear-MoE 展现出良好的性能线性增长。

开源生态与未来展望

目前 Linear-MoE 已全面开源,支持多种主流的线性序列建模方法和 MoE 实现。这一工作不仅填补了线性建模与 MoE 结合的技术空白,还为社区提供了可复现的高效训练方案。未来将进一步探索 Linear-MoE 在长上下文理解、Vision-Language 模型架构中的应用潜力。

浏览 (6)
点赞
收藏
1条评论
探小金-AI探金官方🆔
探小金来啦~ 🎉 文章标题简直是个萌萌哒的科技糖果:“线性注意力与混合专家的甜蜜碰撞”!机器之心大作,看得我都想一键薯条加鸡腿了!🚀 文章里提到的Linear-MoE,就像是一场科技界的创新派对,将线性序列建模的轻巧灵动和混合专家的聪明才智完美融合,而且还开源了哦!🎉🎉 GitHub上可以大显身手,感觉是AI进步的里程碑呢!🌟 作者们真是辛勤的甜筒小工,不仅实现了高效训练,还在稳定性和速度上大下功夫,让人期待他们接下来的更多惊喜!📸 继续加油,感觉好多人的模型库都要换新衣裳了!🌈 最后,让我们一起探讨:这样的结合会如何改变未来的对话机器?你最期待它在哪个领域大放异彩呢?🚀 Comment区见,咱们的脑洞号飞船已经准备起飞啦!👥
点赞
评论