MOBA：长上下文 LLMs 的混合块注意机制
25年2月来自月之暗面、清华和浙大的论文“MOBA: Mixture Of Block Attention For Long-context LLMs”。
扩展有效上下文长度对于将大语言模型 (LLM) 推进到通用人工智能 (AGI) 至关重要。然而，传统注意机制固有的计算复杂度的二次增加，带来高昂的开销。现有的方法要么施加强偏结构，例如特定任务的汇聚（sink）注意或窗注意，要么将注意机制彻底修改为线性近似，其在复杂推理任务中的表现仍未得到充分探索。
这项工作提出一种遵循“较少结构”原则的解决方案，允许模型自主确定关注的位置，而不是引入预定义的偏差。其中引入了混合块注意 (MoBA)，这是一种将混合专家 (MoE) 原则应用于注意机制的方法。这种架构在长上下文任务上表现出色，同时提供一个关键优势：能够在完整和稀疏注意之间无缝转换，从而提高效率而不会影响性能。 MoBA 已经部署用于支持 Kimi 的长上下文请求，并展示 LLM 在高效注意计算方面的重大进步。
对通用人工智能 (AGI) 的追求推动大语言模型 (LLM) 的发展，达到了前所未有的规模，有望处理模仿人类认知的复杂任务。实现 AGI 的关键能力是处理、理解和生成长序列的能力，这对于从历史数据分析到复杂推理和决策过程的广泛应用至关重要。对扩展上下文处理日益增长的需求不仅体现在长输入提示理解的流行中，如 Kimi (MoonshotAI 2023)、Claude (Anthropic 2023) 和 Gemini (Reid 2024) 等模型所展示的那样，还体现在 Kimi k1.5 (Team 2025)、DeepSeek-R1 (Guo 2025) 和 OpenAI o1/o3 (Guan 2024) 中对长链思维 (CoT) 输出能力的最新探索中。
然而，由于与原始注意机制相关的计算复杂度呈二次增长，延长 LLM 中的序列长度并非易事 (Waswani，2017)。这一挑战引发一波旨在提高效率而不牺牲性能的研究浪潮。一个突出的方向是利用注意分数固有的稀疏性。这种稀疏性既来自数学——来自 softmax 运算，其中研究各种稀疏注意模式 (Jiang，2024)——也来自生物学 (Watson，2025)，其中在与记忆存储相关的大脑区域中观察到稀疏连接。
现有方法通常利用预定义的结构约束，例如基于汇聚 (G. Xiao，2023) 或滑动窗注意 (Beltagy，2020)，来利用这种稀疏性。虽然这些方法可能有效，但它们往往高度针对特定任务，可能会阻碍模型的整体通用性。或者，一系列动态稀疏注意机制，例如 Quest（Tang，2024）、Minference（Jiang，2024）和 RetrievalAttention（Liu，2024），在推理时选择 tokens 子集。虽然这些方法可以减少长序列的计算量，但它们并不能显着减轻长上下文模型的密集训练成本，因此很难将 LLM 有效地扩展到数百万个 tokens 的上下文。另一种有前途的替代方法最近以线性注意模型的形式出现，例如 Mamba（Dao & Gu，2024）、RWKV（Peng，2023；Peng，2024）和 RetNet（Sun，2023）。这些方法用线性近似取代基于典型 softmax 的注意机制，从而减少长序列处理的计算开销。然而，由于线性注意机制和传统注意机制之间存在很大差异，调整现有的 Transformer 模型通常会产生高昂的转换成本（Mercat，2024；Wang，2024；Bick，2025；Zhang，2024）或需要从头开始训练全新的模型（Li，2025）。更重要的是，它们在复杂推理任务中有效性的证据仍然有限。
高效的注意机制（Tay，2020）的开发一直是自然语言处理领域的一个重要研究领域，尤其是随着大语言模型 (LLM) 的兴起。随着处理更长序列和降低计算成本的需求不断增长，高效注意技术已成为一种有前途的解决方案，可在保持模型性能的同时降低自注意机制的二次复杂度。
静态稀疏模式：许多研究者致力于设计 LLM 中的静态注意模式，例如 Sparse Transformer（Child，2019）、Star-Transformer（Guo，2019）、BlockBERT（Qiu，2019）、Longformer（Beltagy，2020）、GMAT（Gupta，2020）、ETC（Ainslie，2020 年）、BigBird（Zaheer，2020）、LongT5（Guo，2021）和 LongNet（Ding，2023）。他们对静态注意模式的选择，可以包括步幅和固定注意、窗注意、全局 token 注意、随机注意、扩张注意、块稀疏注意或它们的任意组合。在多模态模型领域，静态稀疏注意机制也得到发展，例如针对二维图像的轴向注意 (Ho et al. 2019) 和针对三维视频的时空注意 (Zheng et al. 2024)。
动态稀疏模式：与静态模式不同，动态稀疏注意技术可以自适应地确定要关注哪些 tokens。Reformer (Kitaev 2020) 和 Routing Transformer (Roy 2021) 分别采用 LSH 和 K-均值对 tokens 进行聚类，并关注聚类而不是整个上下文。Memorizing Transformers (Wu 2022) 和 Unlimiformer (Bertsch 2024) 动态关注由 kNN 算法选择的 tokens。 CoLT5（Ainslie，2023）设计一个路由模块来选择最重要的 Q 和 K。稀疏 Sinkhorn 注意（Tay，2020）学习从输入序列中排列块，从而实现动态块稀疏注意计算。
无需训练的稀疏注意：除了前面讨论的训练稀疏注意模型的方法外，还有一些策略旨在结合稀疏注意机制来提高模型推理的两个主要阶段（预填充阶段或解码阶段，或两者）的效率。在预填充优化阶段，完整的提示可用于注意分析，从而可以探索更复杂的稀疏注意模式。例如，MoA（Fu，2024）、Minference（Jiang，2024）和 SeerAttention（Gao，2024）研究稀疏注意配置，例如 A-形、垂直-斜线（slash）和动态块稀疏性。在解码优化方面，人们投入大量工作来压缩和修剪 KV-缓存，在文本生成的质量和速度之间取得平衡。该领域的著名努力包括 H2O（Zhang，2024）、StreamingLLM（Xiao，2023）、TOVA（Oren，2024）、FastGen（Ge，2023）和 Quest（Tang，2024）。
超越传统注意架构：另一条研究路线研究偏离传统注意机制的新型模型架构。随着架构的变化，这些方法需要从头开始训练模型，并且无法重用预训练的基于 Transformer 模型。该领域的研究探索，受卷积神经网络 (CNN)、循环神经网络 (RNN)、状态空间模型 (SSM) 或线性注意 (Katharopoulos et al. 2020) 启发的架构，此类模型的示例包括 Hyena (Poli 2023)、Performer (Choromanski 2020)、Linformer (Wang 2020)、RWKV （Peng 2023）、Mamba（Gu 2023）、RetNet（Sun 2023）等。
一个关键的研究问题是：如何设计一个鲁棒且适应性强的注意架构，既保留原始 Transformer 框架，又遵循“较少结构”原则，让模型能够在不依赖预定义偏差的情况下确定哪儿关注？理想情况下，这种架构可以在全注意模式和稀疏注意模式之间无缝转换，从而最大限度地提高与现有预训练模型的兼容性，并在不影响性能的情况下实现高效推理和加速训练。
混合块注意 (MoBA) 是一种新架构，它建立在混合专家 (MoE) (Shazeer，2017) 之上，并将其应用于 Transformer 模型的注意机制。
Quest（Tang，2024）。特别是 Quest，它可以被视为MoBA，只是具有较小块大小和结合最小和最大池化的专门块表示函数。另一项与 MoBA 密切相关的工作是 Longheads (Y. Lu et al. 2024)，它可以被视为具有 top-1 门控网络的 MoBA，这意味着每个查询都会选择最相关的 KV 块进行注意。
与标准注意机制中每个 Q token 关注整个上下文不同，MoBA 允许每个 Q token 仅关注 K 和 V 的子集 I：
Attn(q, K, V) = Softmax(qK^T)V，
MoBA(q, K, V) = Softmax(qK[I]^T)V[I],
MoBA 的关键创新在于块划分和选择策略。将长度为 N 的完整上下文划分为 n 个块，每个块代表后续 tokens 的子集。不失一般性，假设上下文长度 N 可以被块数 n 整除。其中 B = N/n，第 i 个的块范围如下：
I_i = [(i−1) × B + 1, i × B]，
通过应用 MoE 的 top-k 门控机制，使每个 Q 能够有选择地关注来自不同块的 token 子集，而不是整个上下文：
该模型采用门控机制 g_i 来为每个 Q token 选择最相关的块。MoBA 门首先计算亲和力分数 s_i，以测量查询 q 与第 i 个块之间的相关性，然后在所有块中应用前 k 个门控。更正式地说，第 i 个块 g_i 的门值由以下公式计算
本研究中，分数 s_i 由 q 与 K[I_i] 沿序列维度均值池化之间的内积计算得出：
如图所示混合块注意 (MoBA) ： (a) MoBA 运行示例；(b) MoBA 集成到 Flash Attention 中。
在自回归语言模型中，保持因果关系非常重要，因为它们基于以前 tokens 做下一个 token 预测来生成文本。这种顺序生成过程确保 token 不会影响其之前的 tokens，从而保留因果关系。MoBA 通过两种特定设计来保留因果关系：
因果关系：注意不考虑未来块。MoBA 确保 Q tokens 不能被路由到任何未来块。通过将注意范围限制在当前和过去的块，MoBA 遵循语言建模的自回归性质。更正式地，将 pos(q) 表示为查询 q 的位置索引，为任何块 i 设置 s_i = −∞ 和 g_i = 0，使得 pos(q) < i × B。
当前块注意和因果掩码。将“当前块”定义为包含 Q token 本身的块。路由到当前块也可能违反因果关系，因为整个块的均值池化可能会无意中包含来自未来 tokens 的信息。为了解决这个问题，强制每个 token 必须被路由到其各自的当前块，并在当前块注意期间应用因果掩码。这种策略不仅可以避免后续 tokens 的任何信息泄漏，还可以鼓励对局部上下文的注意。更正式地，为块 i 设置 g_i = 1，其中 Q token pos(q) 的位置在区间 I_i 内。从混合专家 (MoE) 的角度来看，MoBA 中的当前块注意类似于现代 MoE 架构中共享专家的角色 (Dai 2024; Yang 2024)，其中在专家选择时添加静态路由规则。
MoBA 的关键设计选择如下。
细粒度块分割。细粒度专家分割，对提高模式性能的积极影响，已在混合专家 (MoE) 文献中得到充分证明 (Dai 2024; Yang 2024)。这项工作探索将类似的细粒度分割技术应用于 MoBA 的潜在优势。 MoBA 受到 MoE 的启发，沿着上下文长度维度而不是 FFN 中间隐藏维度进行分割。因此，研究旨在确定将上下文划分为具有更细粒度的块，MoBA 是否也能受益。
MoBA 和全注意的混合。MoBA 旨在替代全注意，保持相同数量的参数而不进行任何增加或减少。此功能启发在全注意和 MoBA 之间进行平滑过渡。具体而言，在初始化阶段，每个注意层都可以选择全注意或 MoBA，并且可以在训练期间动态更改此选择（如果需要）。在先前的工作中已经研究将全注意转换为滑动窗注意的类似想法（Zhang，2024）。
与滑动窗注意和注意汇聚（sink）的比较。滑动窗注意 (SWA) 和注意汇聚，是两种流行的稀疏注意架构。两者都可以看作是 MoBA 的特殊情况。对于滑动窗注意（Beltagy，2020），每个 Q token 仅关注其相邻 tokens。这可以解释为 MoBA 的一种变型，其门控网络不断选择最近的块。同样，注意汇聚（Xiao，2023），其中每个 Q token 关注初始 tokens 和最新 tokens 的组合，可以看作是 MoBA 的一种变型，其门控网络始终选择初始块和最新块。以上讨论表明，MoBA 比滑动窗注意和注意汇聚具有更强的表达能力。此外，它表明 MoBA 可以通过合并特定的门控网络灵活地近似许多静态稀疏注意架构。
总体而言，MoBA 的注意机制，允许模型自适应地动态地关注上下文中信息量最大的块。这对于涉及长文档或序列的任务特别有益，因为在这些任务中，关注整个上下文可能是不必要的，而且计算成本很高。 MoBA 能够选择性地关注相关块，从而实现更细致、更高效的信息处理。
MoBA 的实现如下。
通过结合 FlashAttention（Dao，2022）和 MoE（Rajbhandari，2022）的优化技术，提供 MoBA 的高性能实现。其包括五个主要步骤：
• 根据门控网络和因果掩码确定 Q tokens 到 KV 块的分配。
• 根据分配的 KV 块排列 Q tokens 的顺序。
• 计算每个 KV 块和分配给它的 Q token 的注意输出。此步骤可以通过长度不同的 FlashAttention 进行优化。
• 将注意输出重排列回其原始顺序。
• 使用在线 Softmax（即平铺 tiling）组合相应的注意输出，因为 Q tokens 可能会关注其当前块和多个历史 KV 块。
算法工作流程在算法 1 中形式化。首先，将 KV 矩阵划分为块（第 1-2 行）。接下来，计算门控分数，该分数测量Q tokens 和 KV 块之间的相关性（第 3-7 行）。在门控分数上应用 top-k 算子（与因果掩码一起），得到一个稀疏Q-到-KV-块的映射矩阵 G，以表示 Q 到 KV 块的分配（第 8 行）。然后，根据 Q-到-KV-块映射排列 Q tokens，并计算分块注意输出（第 9-12 行）。值得注意的是，对历史块的注意（第 11 和 14 行）和当前块注意（第 10 和 13 行）是分开计算的，因为当前块注意中需要保持额外的因果关系。最后，将注意输出重排列回其原始顺序，并与在线 softmax 相结合（第 16 行）（Milakov 2018；Liu 2023）。
将 MoBA 设计为全注意的灵活替代品，以便它能够以最小的开销轻松地从/到全注意切换，并实现相当的长上下文性能。
MoBA/全混合训练。在 30B 个 token 上训练三个模型，每个模型有 1.5B 参数，上下文长度为 32K 个 token。对于 MoBA 的超参，块大小设置为 2048，top-k 参数设置为 3。详细的训练配方如下：
• MoBA/全混合：此模型使用两阶段配方进行训练。在第一阶段，使用 MoBA 对 90% 的 tokens 进行训练。在第二阶段，模型切换到全注意，计算剩余 10% 的 tokens。
• 全注意：该模型在整个训练过程中都使用全注意进行训练。
• MoBA：该模型仅使用 MoBA 进行训练。
通过语言模型 (LM) 逐个位置损失来评估它们的长上下文性能，这是一种细粒度指标，用于评估序列中每个位置的 LM 损失。与对所有位置 LM 损失取平均值来计算的原始 LM 损失不同，逐个位置 LM 损失分别分解每个位置的损失。先前的研究（Xiong，2023；Reid，2024）提出类似的指标，他们注意到逐个位置 LM 损失相对于上下文长度呈幂律趋势。如图 a 所示，仅使用 MoBA 的配方会导致尾随 token 的位置损失更高。重要的是，MoBA/全混合配方达到的损失几乎与全注意相同。这一结果凸显 MoBA/全混合训练方案在平衡训练效率和模型性能方面的有效性。更有趣的是，在 MoBA 和全注意之间切换时，没有观察到明显的损失峰值，这再次证明 MoBA 的灵活性和稳健性。
分层混合。MoBA 的这种灵活性鼓励深入研究一种更复杂的策略——MoBA 和全注意的分层混合。特别是这一策略在监督微调 (SFT) 期间的应用。研究这一策略的动机源于MoBA 有时会导致 SFT 期间的性能不佳，如图 b 所示。这可能归因于 SFT 中使用的损失掩码——提示 tokens 通常被排除在 SFT 期间的损失计算之外，这可能会对 MoBA 等稀疏注意方法造成稀疏梯度挑战。因为它可能会阻碍梯度的反向传播，这些梯度最初是从未掩码的 tokens 中计算出来的，并贯穿整个上下文。为了解决这个问题，提出一种混合方法——将最后几层 Transformer 从 MoBA 切换到全注意机制，而其余层继续采用 MoBA。如图 b 和图 c 所示，此策略可以显著减少 SFT 损失。
对于监督微调，遵循类似的策略，逐渐将上下文长度从 32K 增加到 1M。基线全注意模型（称为 Llama-8B-1M-Full）也遵循类似的训练策略，如图所示，唯一的区别是整个过程中使用全注意。这种方法能够在等效训练条件下直接比较 MoBA 与全注意模型的性能。
