论文地址：https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf

MoBA 通过将专家混合系统(Mixture of Experts, MoE)的思想与稀疏注意力(sparse attention)相结合，为大语言模型中的长文本处理方式带来革命性变化。

高效处理和生成长序列的能力对大语言模型（LLMs）十分重要，然而传统注意力机制在扩展到长上下文时面临平方复杂度计算挑战。现有的解决方案通常会引入特定任务的稀疏先验（比如滑动窗口注意力，Sink注意力等），或需要对注意力机制进行彻底的改变（比如线性注意力）。混合块注意力（Mixture of Block Attention, MoBA）将专家混合（Mixture of Experts, MoE）的思想与稀疏注意力（Sparse Attention）结合，在保持原始Transformer架构的灵活性和性能的同时，尝试解决上述挑战。

MoBA借鉴了Mixture of Experts (MoE)的思想，将其应用于注意力机制中。具体来说，MoBA将上下文序列划分为多个块（block），并通过门控机制动态选择每个查询（query）需要关注的最相关的块。这种方法允许模型在全注意力和稀疏注意力之间无缝切换，从而在保持性能的同时显著提高计算效率。

MoBA的核心创新点包括：

1. 块划分与选择策略：将上下文划分为多个块，每个查询只关注最相关的块。
2. 门控机制：通过计算查询与每个块的亲和度分数，动态选择最相关的块。
3. 因果性保持：确保模型在自回归生成过程中不会关注未来的块，保持因果性。

1、技术介绍

MoBA是一种受专家混合（MoE）和块稀疏注意力（Block Sparse Attention）启发的注意力架构。在传统注意力中，每个 query 都会和完整的上下文（KV）进行计算，MoBA 使每个 query 能够有选择地专注于一部分（KV），保持性能的同时降低了计算成本。

回顾全注意力： 

MoBA使每个 query 只注意到长度为 N 的完整 KV 的一个子集： 

MoBA的关键创新在于块划分和选择策略。我们将长度为 N 的完整上下文划分为 n 个块，并用 B 表示块大小，用  表示第 i 个块。通过应用 MoE 的 top-k 门控机制，我们使每个 query 能够有选择地专注于若干 KV 块，而不是整个上下文：

其中  是门控网络的输出,  是第  个块的重要性打分。

相关实现的伪代码：

二、效率对比实验

MoBA 与全注意力的效率对比。(a) 1M 模型加速评估：在序列长度从 8K 到 1M 增加的情况下，MoBA 与 Flash Attention 在 1M 模型上的计算时间扩展。(b) 固定稀疏率扩展：在序列长度从 8K 到 10M 增加的情况下，MoBA 与 Flash Attention 的计算时间扩展对比，保持恒定的稀疏率 95.31%

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

三、结论

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

YUTransformer是谁？

个人介绍：技术博客名为YUTransformer，8年AI老兵，从事NLP、大语言模型、多模态大模型等相关算法的研发和落地，拥有丰富的算法经验，先后在百度、平安、小鹏汽车从事算法落地的工作，借助平台将个人的一些算法研究和经验分享出来，一起推动技术的进步！加我vx（yx116169)入vlm/vla大模型群和llm群