撞车DeepSeek NSA，Kimi杨植麟署名的新注意力架构MoBA发布，代码也公开

knXLdFdf2

昨天下午，DeepSeek 发布了一篇新论文，提出了一种改进版的注意力机制 NSA；加上还有创始人兼 CEO 梁文锋亲自参与，一时之间吸引眼球无数。
但其实就在同一天，月之暗面也发布了一篇主题类似的论文，并且同样巧合的是，月之暗面创始人兼 CEO 杨植麟也是该论文的署名作者之一。并且不同于 DeepSeek 只发布了一篇论文，月之暗面还发布了相关代码。且这些代码已有一年的实际部署验证，有效性和稳健性都有保证。
这篇论文提出了一种名为 MoBA 的注意力机制，即 Mixture of Block Attention，可以直译为「块注意力混合」。据介绍，MoBA 是「一种将混合专家（MoE）原理应用于注意力机制的创新方法。」该方法遵循「更少结构」原则，并不会引入预定义的偏见，而是让模型自主决定关注哪些位置。

• 论文标题：MoBA: Mixture of Block Attention for Long-Context LLMs
  
• 论文地址：https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf
  
• 项目地址：https://github.com/MoonshotAI/MoBA
  

与 DeepSeek 的 NSA 注意力机制新论文一样，月之暗面这篇 MoBA 论文也收获了诸多好评。

下面我们就来解读一下这篇论文。
最近一段时间，大语言模型（LLM）发展到了前所未有的规模，我们希望它们能够处理模拟人类认知的复杂任务。具体来说，实现 AGI 的一个关键能力是处理、理解和生成长序列的能力，这对于从历史数据分析到复杂推理和决策过程的各种应用至关重要。这种对扩展上下文处理能力的需求不仅体现在长输入提示理解的普及上，还体现在最近 Kimi k1.5、DeepSeek-R1 和 OpenAI o1/o3 中展示的，对长链思维（CoT）输出能力的探索中。
然而，由于传统注意力机制带来的计算复杂度呈二次函数增长，扩展 LLM 的序列长度并非易事。这一挑战催生了一波研究浪潮，其中一个重要方向是利用注意力分数的固有稀疏性。这种稀疏性既源于数学上的 softmax 操作，也源于生物学上的观察，即人们已在与记忆存储相关的大脑区域中观察到了稀疏连接。
现有方法通常利用预定义的结构约束（例如基于 sink 的注意力机制或滑动窗口注意力机制）来利用这种稀疏性。虽然这些方法可能有效，但它们往往高度依赖于特定任务，可能会限制模型的整体泛化能力。另一种方法是采用动态稀疏注意力机制，例如 Quest、Minference 和 RetrievalAttention，这些方法在推理时选择部分 token 进行计算。尽管这些方法可以减少长序列的计算量，但它们并未显著缓解长上下文模型的高昂训练成本。
最近，线性注意力模型作为一种有前景的替代方案出现，例如 Mamba、RWKV 和 RetNet。这些方法用线性近似替代了传统的基于 softmax 的注意力机制，从而降低了长序列处理的计算开销。然而，由于线性注意力与传统注意力机制存在显著差异，适配现有的 Transformer 模型通常需要高昂的转换成本，或者需要从头开始训练全新的模型。更重要的是，它们在复杂推理任务中的有效性证据仍然有限。
因此，一个关键的研究问题随之而来：我们如何设计一种稳健且适应性强的注意力架构，既能保留原始 Transformer 框架，又能遵循「更少结构」的原则，使模型能够在不依赖预定义偏差的情况下自主决定关注哪些部分？理想情况下，这种架构能够在完全注意力和稀疏注意力模式之间无缝切换，从而最大限度地与现有预训练模型兼容，并在不牺牲性能的前提下实现高效推理和加速训练。
为此研究人员提出了块注意力混合机制（Mixture of Block Attention, MoBA），这是一种全新架构，它基于混合专家系统（MoE）的创新原理，并将其应用于 Transformer 模型的注意力机制中。MoE 此前主要应用于 Transformer 的前馈网络（FFN）层，而 MoBA 首次将其引入长上下文注意力机制中，允许动态选择与每个查询 token 相关的历史关键块和值块。这种方法不仅提高了 LLM 的效率，还使其能够处理更长、更复杂的提示，而无需按比例增加资源消耗。
相比之下，MoBA 通过将上下文划分为块并采用门控机制选择性地将查询 token 路由到最相关的块，解决了传统注意力机制的计算效率低下的问题。这种块稀疏注意力显著降低了计算成本，为更高效地处理长序列铺平了道路。模型能够动态选择最具信息量的关键块，从而提高了性能和效率，这对于涉及大量上下文信息的任务尤为有益。
在该工作中，月之暗面详细介绍了 MoBA 的架构，首先是其块划分和路由策略，其次是与传统注意力机制相比的计算效率。他们也进行了实验，证明了 MoBA 在处理长序列任务中的卓越性能。
MoBA：将 MoE 原理应用于注意力机制

MoBA 通过动态选择历史片段（块）进行注意力计算，扩展了 Transformer 模型的能力。
MoBA 的灵感来源于混合专家（Mixture of Experts，简称 MoE）技术和稀疏注意力技术。前者主要应用于 Transformer 架构中的前馈网络（Feedforward Network，简称 FFN）层，而后者已被广泛用于扩展 Transformer 以处理长文本上下文。
本文创新点在于月之暗面将 MoE 原理应用于注意力机制本身，从而实现对长序列更高效和更有效的处理。
MoBA 主要包括如下部分：

• 可训练的块稀疏注意力：将完整上下文分割成若干块，每个查询 token 都会学习关注最相关的 KV 块，从而实现长序列的高效处理。
  
• 无参数门控机制：引入了一种新颖的无参数 top-k 门控机制，为每个查询 token 选择最相关的块，确保模型只关注最具信息量的内容。
  
• 完全注意力与稀疏注意力的无缝切换：MoBA 被设计为完全注意力机制的灵活替代方案，能够在完全注意力和稀疏注意力模式之间实现平滑过渡。
  

总体而言，MoBA 注意力机制使模型能够自适应且动态地关注上下文中最有信息量的块。这在处理涉及长文档或序列的任务时尤其有益，因为在这些任务中，关注整个上下文可能是不必要的，并且计算代价高昂。MoBA 选择性地关注相关块的能力，使得信息处理更加细致且高效。
月之暗面通过整合来自 FlashAttention 和 MoE 的优化技术，实现了 MoBA 的高性能版本。图 2 展示了 MoBA 的高效率。

MoBA 实现过程包含五个步骤：

• 根据门控网络和因果掩码，确定查询 token 对 KV 块的分配。
  
• 根据分配的 KV 块，安排查询 token 的顺序。
  
• 计算每个 KV 块的注意力输出和分配给它的查询 token，这一步可以通过不同长度的 FlashAttention 来优化。
  
• 将注意力输出重新排列回其原始顺序。
  
• 使用在线 Softmax 合并相应的注意力输出，因为一个查询 token 可能会关注其当前块以及多个历史 KV 块。
  

下图为该算法的形式化描述，并在图 1b 中进行了可视化，展示了如何基于 MoE 和 FlashAttention 实现 MoBA。

MoBA 的实验表现

为了验证 MoBA 的实际效果，月之暗面进行了一系列 Scaling Law 和消融实验。
Full Attention 与 MoBA 对比实验
该团队训练了 5 个模型来进行比较，表 1 给出详细配置。

结果见图 3a，可以看到，Full Attention 与 MoBA 的 scaling 趋势非常相似。具体来说，这两种注意力机制之间的验证损失差异在 1e − 3 的范围内保持一致。这表明，尽管 MoBA 的稀疏注意力模式稀疏度高达 75%，但它实现了与完全注意力相当的扩展性能。

此外，该团队也验证了 MoBA 的上下文可扩展性。在实验中，该团队将最大序列长度从 8k 提升到了 32k。结果见图 3b，可见尽管在所有五项实验中，MoBA 最后的块 LM 损失与 Full Attention 相比略高，但损失差距正在逐步缩小。该实验表明 MoBA 具有长上下文可扩展性。
该团队也通过消融实验对 MoBA 的细粒度块分割的有效性进行了验证，结果见图 4。

可以看到，MoBA 的性能受块粒度的显著影响。具体来说，最粗粒度设置（从 8 个块中选择 2 个块）与更细粒度的设置之间的性能差异为 1e-2。这些发现表明，细粒度分割似乎是增强 MoE 系列（包括 MoBA）模型性能的普适技术。
将 MoBA 与 Full Attention 融合到一起会如何？
MoBA 本就是作为 Full Attention 的灵活替代而设计的，可以最小的开销轻松地互相切换，并实现相当的长上下文性能。月之暗面实验表明，Full Attention 和 MoBA 之间的无缝过渡可以成为高效长上下文预训练的解决方案。然后他们还讨论了分层混合策略，其作用主要是提升监督微调（SFT）的性能。
在实验中的混合训练时，模型使用的是两阶段方案：第一阶段使用 MoBA 完成 90% 的 token 训练，第二阶段模型切换到 Full Attention 使用剩余 10% 的 token 进行训练。另外，当然也使用了纯 MoBA 和纯 Full Attention 方法作为对比。
结果见图 5a，仅使用 MoBA 时会导致 trailing token 的位置损失更高。重要的是，混合方案的损失几乎与 Full Attention 一样。这一结果凸显了混合训练方案在平衡训练效率和模型性能方面的有效性。更有趣的是，在 MoBA 和 Full Attention 之间切换时没有观察到明显的损失峰值，这再次证明了 MoBA 的灵活性和稳健性。

月之暗面也研究了分层混合的方案。这种方案当然更加复杂，研究动机是他们观察到 MoBA 有时会导致 SFT 期间的性能不佳，如图 5b 所示。
他们推测这可能归因于 SFT 中使用的损失掩蔽（loss masking）—— 提示词 token 通常被排除在 SFT 期间的损失计算之外，这可能会对 MoBA 等稀疏注意力方法造成稀疏梯度难题。因为它可能会阻碍梯度的反向传播，而这些梯度最初是根据未掩蔽的 token 计算出来的，并贯穿整个上下文。
为了解决这个问题，他们提出了另一种混合方法 —— 将最后几层 Transformer 从 MoBA 换成 Full Attention，而其余层继续采用 MoBA。如图 5b 和图 5c 所示，此策略可以显著减少 SFT 损失。
大语言模型评估
研究团队基于 Llama 3.1 8B 模型开发了 Llama-8B1M-MoBA，通过持续预训练将上下文长度从 128K 逐步扩展到 1M token。模型采用块大小 4096 和 top-K 参数 12 的设置，实现了高达 95.31% 的注意力稀疏度。为平衡性能，他们保留最后三层为 Full Attention，其余 29 层使用 MoBA。

评估结果显示，MoBA 模型与 Full Attention 模型（Llama-8B-1M-Full）性能相当。特别是在最具挑战性的 RULER 基准测试中，即使在 62.5% 的稀疏度下，MoBA 模型仍达到了 0.7818 的得分，几乎与 Full Attention 模型的 0.7849 持平。对于需要处理 100 万 token 的「大海捞针」测试集中，模型同样展现出优秀的能力。

效率和可扩展性
实验证明，MoBA 在保持性能的同时显著提升了效率：在处理 1M token 时，比 Full Attention 模型快 6.5 倍。在扩展到 1000 万 token 时，较标准 Flash Attention 实现了 16 倍加速；通过块稀疏注意力机制和优化实现，MoBA 将计算复杂度从二次方降低到了亚二次方级别。
这使得 MoBA 的优势在序列长度增加时更为明显，特别适合处理超长文本场景。
更多详情，请参见原论文。