跳转到主要内容
本文发布于 2025-02-19
MoBA通过将专家混合系统(Mixture of Experts, MoE)的思想与稀疏注意力(sparse attention)相结合，为大语言模型中的长文本处理方式带来革命性变化。 高效处理和生成长序列的能力对大语言模型（LLMs）十分重要，然而传统注意力机制在扩展到长上下文时面临平方复杂度计算挑战。现有的解决方案通常会引入特定任务的稀疏先验（比如滑动窗口注意力，Sink注意力等），或需要对注意力机制进行彻底的改变（比如线性注意力）。在本文中，我们介绍了混合块注意力（Mixture of Block Attention, MoBA），它将专家混合（Mixture of Experts, MoE）的思想与稀疏注意力（Sparse Attention）结合，在保持原始Transformer架构的灵活性和性能的同时，尝试解决上述挑战。

​
从全注意力（Full Attention）到混合块注意力（MoBA）

MoBA是一种受专家混合（MoE）和块稀疏注意力（Block Sparse Attention）启发的注意力架构。在传统注意力中，每个 query 都会和完整的上下文（KV）进行计算，MoBA 使每个 query 能够有选择地专注于一部分（KV），保持性能的同时降低了计算成本。 总的来说，MoBA有下面关键创新，感兴趣的读者可以参考我们的论文了解详细信息https://arxiv.org/abs/2502.13189
  1. 块划分和动态路由：上下文被划分为块，MoE 式的门控网络动态地为每个 query 选择最相关的块。
  2. 保持因果性：MoBA 确保 query 不能关注未来的块，保持了语言模型的自回归特性。当前块应用了 causal 掩码以防止信息泄露。
  3. 细粒度块划分：类似于MoE中细粒度专家划分，上下文的细粒度划分增强了性能，允许更细致的注意力模式。
  4. MoBA与全注意力的混合：MoBA可以无缝地在全注意力和稀疏注意力之间切换。这种灵活性使模型能够利用现有的预训练全注意力模型。
  5. 高性能实现：我们通过整合FlashAttention和MoE的优化技术，提供了MoBA的高性能实现。

​
将 MoBA 扩展到1000万上下文

  1. Scaling Law 实验：MoBA在 LM loss 方面实现了与全注意力相当的性能，即使对于序列末尾的 LM loss 差异也很小。
  2. 混合策略：尝试了训练混合和分层混合两种策略。训练混合是在训练期间将MoBA与全注意力交替使用（例如，用MoBA训练大多数语料，仅用少量语料激活全注意力），这样可以实现了与全注意力几乎相同的性能，同时提高了训练效率。分层混合策略，是在层间交替使用 MoBA 和全注意力，这在 SFT 期间进一步保证了效果。
  3. 现实世界任务评估：MoBA在各种长上下文基准测试中表现出色，包括大海捞针和RULER，取得和全注意力相同效果。
  4. 效率和可扩展性：MoBA显著减少了注意力计算时间，与全注意力相比，100万上下文长度实现了6.5倍的速度提升，在扩展到 1000 万上下文时实现了16倍的速度提升。
MoBA 与全注意力（使用 Flash Attention 实现）的效率对比。(a) 1M 模型加速评估：在序列长度从 8K 到 1M 增加的情况下，MoBA 与 Flash Attention 在 1M 模型上的计算时间扩展。(b) 固定稀疏率扩展：在序列长度从 8K 到 10M 增加的情况下，MoBA 与 Flash Attention 的计算时间扩展对比，保持恒定的稀疏率 95.31%（固定 64 个 MoBA 块，具有方差块大小和固定的 top-k=3）。

​
结论和未来工作

通过动态选择相关的上下文块并保持与传统 Transformer 架构的兼容性，MoBA在性能和效率之间取得了平衡。它能够在全注意力和稀疏注意力模式之间无缝切换，使其成为长上下文任务的有效解决方案。未来的工作会进一步优化 MoBA 的块选择策略，并探索其在复杂推理任务中的潜力。

​
参考资料

有关MoBA的更多详细信息，请参考技术报告和MoBA的GitHub仓库 https://github.com/MoonshotAI/MoBA