MOBA: MIXTURE OF BLOCK ATTENTION FOR LONG-CONTEXT LLMS

文章标题：MOBA：用于长上下文大语言模型的混合块注意力机制
作者/机构：Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, Zhiqi Huang, Huan Yuan, Suting Xu, Xinran Xu, Guokun Lai, Yanru Chen, Huabin Zheng, Junjie Yan, Jianlin Su, Yuxin Wu, Neo Y. Zhang, Zhilin Yang, Xinyu Zhou (Moonshot AI); Mingxing Zhang (Tsinghua University); Jiezhong Qiu (Zhejiang Lab/Zhejiang University)

A1 主要贡献

本文旨在解决大语言模型（LLM）在处理长序列时面临的核心挑战，即传统自注意力机制带来的二次方计算复杂度增长问题。现有方法通常引入强结构偏见（如窗口注意力）或从根本上改变注意力机制（如线性近似），前者可能损害模型的泛化能力，后者在复杂推理任务上的有效性尚未得到充分验证。

核心问题与研究目标：
如何设计一个鲁棒且适应性强的注意力架构，该架构能保留 Transformer 的原始框架，同时遵循“更少结构”原则，让模型自主决定注意力分配，而不是依赖预定义的偏见？理想情况下，该架构应能在全注意力（full attention）和稀疏注意力（sparse attention）模式之间无缝切换，以兼容现有预训练模型，并在不牺牲性能的前提下实现高效推理和加速训练。

主要创新点：
本文提出了混合块注意力（Mixture of Block Attention, MoBA），这是一种新颖的架构，它将混合专家（Mixture of Experts, MoE）的原理创新性地应用于 Transformer 模型的注意力机制。
1. MoE 原理应用于注意力：与主要将 MoE 用于前馈网络（FFN）层的传统做法不同，MoBA 首次将其应用于长上下文注意力，允许每个查询（query）动态地从历史键（key）和值（value）块中选择最相关的部分进行计算。
2. 动态块稀疏注意力：MoBA 将上下文划分为多个块，并使用一个门控网络（gating mechanism）来选择性地将查询路由到最相关的块。这种块稀疏注意力显著降低了计算成本。
3. 无缝模式切换与兼容性：该架构被设计为全注意力的直接替代品，参数量相同，能够在全注意力和稀疏注意力之间平滑过渡。这使其能够高效地进行长上下文模型的持续预训练，并与现有模型框架高度兼容。

MoBA 架构已成功部署于 Kimi 的长上下文服务中，展现了其在高效注意力计算方面的显著优势。

A2 方法细节

本文提出了一种名为混合块注意力（Mixture of Block Attention, MoBA）的新颖架构，通过动态选择历史片段（块）进行注意力计算，扩展了 Transformer 模型的能力。MoBA 的灵感来源于混合专家（MoE）和稀疏注意力技术。前者主要应用于 Transformer 架构中的前馈网络（FFN）层，而后者已被广泛用于扩展 Transformer 以处理长上下文。本文的方法创新性地将 MoE 原理应用于注意力机制本身，从而实现对长序列更高效、更有效的处理。

2.1 预备知识：Transformer 中的标准注意力

标准注意力机制回顾：首先回顾 Transformer 中的标准注意力。为简化起见，我们考虑单个查询词元 $q \in R^{1 \times d}$ 与 $N$ 个键和值词元（分别表示为 $K, V \in R^{N \times d}$）进行交互的情况。标准注意力的计算公式如下：


其中 $d$ 表示单个注意力头的维度。为清晰起见，我们专注于单头场景。多头注意力的扩展涉及将多个此类单头注意力操作的输出进行拼接。

图 1: 混合块注意力（MoBA）图示。(a) MoBA 的一个运行示例；(b) MoBA 与 Flash Attention 的集成。

2.2 MoBA 架构

基于块选择的稀疏注意力：与标准注意力中每个查询词元关注整个上下文不同，MoBA 使每个查询词元只关注键和值的一个子集：


其中 $I \subseteq [N]$ 是所选键和值的集合。

块划分与 Top-k 门控机制：MoBA 的关键创新在于其块划分和选择策略。我们将长度为 $N$ 的完整上下文划分为 $n$ 个块，每个块代表一部分连续的词元。不失一般性，我们假设上下文长度 $N$ 能被块数 $n$ 整除。我们进一步定义块大小为 $B = N/n$，并用以下公式表示第 $i$ 个块的范围：


通过应用来自 MoE 的 top-k 门控机制，我们使每个查询能够选择性地关注来自不同块的词元子集，而不是整个上下文：

门控网络与亲和度分数计算：模型采用一个门控机制（如公式4中的 $g_i$）为每个查询词元选择最相关的块。MoBA 门控首先计算亲和度分数 $s_i$，用于衡量查询 $q$ 与第 $i$ 个块之间的相关性，然后在所有块之间应用 top-k 门控。形式上，第 $i$ 个块的门控值 $g_i$ 计算如下：


其中 $Topk(\cdot, k)$ 表示在为每个块计算的亲和度分数中，值最高的 $k$ 个分数的集合。在本文中，分数 $s_i$ 是通过查询 $q$ 与 $K[I_i]$ 沿序列维度进行平均池化（mean pooling）后的向量之间的内积来计算的：

运行示例：我们在图1a中提供了一个 MoBA 的运行示例，其中有两个查询词元和四个 KV 块。路由器（门控网络）为每个查询动态选择前两个块进行关注。如图1a所示，第一个查询被分配到第一个和第二个块，而第二个查询被分配到第三个和第四个块。

因果性保持机制：在自回归语言模型中保持因果性至关重要，因为它们通过基于先前词元进行下一词元预测来生成文本。这个顺序生成过程确保一个词元不能影响它之前的词元，从而保持因果关系。MoBA 通过两种具体设计来保持因果性：
1. 对未来块不进行注意力计算：MoBA 确保查询词元不能被路由到任何未来的块。通过将注意力范围限制在当前和过去的块，MoBA 遵循了语言建模的自回归特性。形式上，记 $pos(q)$ 为查询 $q$ 的位置索引，对于任何满足 $pos(q) < i \times B$ 的块 $i$，我们设置 $s_i = -\infty$ 和 $g_i = 0$。
2. 当前块注意力与因果掩码：我们将“当前块”定义为包含查询词元本身的块。对当前块的路由也可能违反因果性，因为对整个块进行平均池化可能会无意中包含来自未来词元的信息。为了解决这个问题，我们强制每个词元必须被路由到其各自的当前块，并在当前块的注意力计算期间应用因果掩码。这个策略不仅避免了来自后续词元的任何信息泄露，还鼓励了对局部上下文的关注。形式上，对于查询词元位置 $pos(q)$ 位于区间 $I_i$ 内的块 $i$，我们设置 $g_i = 1$。从混合专家（MoE）的角度看，MoBA 中的当前块注意力类似于现代 MoE 架构（如【Dai et al. 2024, Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models, 2024, arXiv】; 【A. Yang et al. 2024, Qwen2. 5 Technical Report, 2024, arXiv】）中共享专家的角色，在专家选择时添加了静态路由规则。

其他关键设计：接下来，我们讨论 MoBA 的一些额外关键设计选择，例如其块分割策略以及 MoBA 与全注意力的混合使用。

细粒度块分割：细粒度专家分割在提升模型性能方面的积极影响已在混合专家（MoE）文献（【Dai et al. 2024, Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models, 2024, arXiv】; 【A. Yang et al. 2024, Qwen2. 5 Technical Report, 2024, arXiv】）中得到充分证明。在本文中，我们探索了将类似的细粒度分割技术应用于 MoBA 的潜在优势。受 MoE 启发，MoBA 沿上下文长度维度而非 FFN 中间隐藏维度进行分割。因此，我们的研究旨在确定当我们将上下文划分为更细粒度的块时，MoBA 是否也能受益。更多实验结果见3.1节。

MoBA 与全注意力的混合：MoBA 被设计为全注意力的替代品，保持相同数量的参数，不增不减。这一特性启发我们在全注意力和 MoBA 之间进行平滑过渡。具体来说，在初始化阶段，每个注意力层都可以选择全注意力或 MoBA，并且这个选择可以在训练期间根据需要动态改变。之前的工作【X. Zhang et al. 2024, Simlayerkv: A simple framework for layer-level KV cache reduction, 2024, arXiv】也研究过将全注意力过渡到滑动窗口注意力的类似想法。更多实验结果见3.2节。

与滑动窗口注意力和注意力池的比较：滑动窗口注意力（SWA）和注意力池（attention sink）是两种流行的稀疏注意力架构。我们证明两者都可以被视为 MoBA 的特例。对于滑动窗口注意力【Beltagy et al. 2020, Longformer: The long-document transformer, 2020, arXiv】，每个查询词元只关注其相邻的词元。这可以解释为 MoBA 的一个变体，其门控网络持续选择最近的块。类似地，对于注意力池【G. Xiao et al. 2023, Efficient streaming language models with attention sinks, 2023, arXiv】，每个查询词元关注初始词元和最近词元的组合，这可以看作是 MoBA 的一个变体，其门控网络总是选择初始和最近的块。以上讨论表明 MoBA 比滑动窗口注意力和注意力池具有更强的表达能力。此外，它还表明 MoBA 可以通过整合特定的门控网络灵活地逼近许多静态稀疏注意力架构。

自适应与动态聚焦：总的来说，MoBA 的注意力机制允许模型自适应地、动态地聚焦于上下文中信息最丰富的块。这对于涉及长文档或长序列的任务尤其有益，因为在这些任务中，关注整个上下文可能是不必要的，且计算成本高昂。MoBA 选择性地关注相关块的能力，使得信息的处理更加精细和高效。

2.3 实现

高性能实现方案：我们通过整合来自 FlashAttention 【Dao, D. Fu, et al. 2022, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, 2022, NeurIPS】 和 MoE 【Rajbhandari et al. 2022, Deepspeed-moe: Advancing mixture-of-experts inference and training to power nextgeneration ai scale, 2022, ICML】的优化技术，提供了 MoBA 的高性能实现。图2展示了 MoBA 的高效率，而关于效率和可扩展性的详细实验将在3.4节中讨论。我们的实现包括五个主要步骤：
- 根据门控网络和因果掩码确定查询词元到 KV 块的分配。
- 根据查询词元分配的 KV 块重新排列其顺序。
- 为每个 KV 块及其分配的查询词元计算注意力输出。这一步可以通过使用可变长度的 FlashAttention 进行优化。


图 2: MoBA 与全注意力（使用 Flash Attention 实现）的效率对比。(a) 1M 模型加速评估：在序列长度增加（8K-1M）时，MoBA 与 Flash Attention 在 1M 模型上的计算时间扩展对比。(b) 固定稀疏度比例扩展：在序列长度增加（8K-10M）时，MoBA 与 Flash Attention 的计算时间扩展对比，保持 95.31% 的恒定稀疏度（固定 64 个 MoBA 块，块大小可变，固定的 top-k=3）。
- 将注意力输出重新排列回其原始顺序。
- 使用在线 Softmax（即分块计算）合并相应的注意力输出，因为一个查询词元可能会关注其当前块和多个历史 KV 块。

算法流程：算法工作流在算法1中被形式化，并在图1b中进行了可视化，展示了如何基于 MoE 和 FlashAttention 实现 MoBA。
1. 首先，将 KV 矩阵划分为块（第1-2行）。
2. 接着，根据公式6计算门控分数，该分数衡量查询词元和 KV 块之间的相关性（第3-7行）。
3. 在门控分数（连同因果掩码）上应用 top-k 操作，产生一个稀疏的查询到 KV 块的映射矩阵 G，表示查询到 KV 块的分配（第8行）。
4. 然后，根据查询到 KV 块的映射重新排列查询词元，并计算分块的注意力输出（第9-12行）。值得注意的是，对历史块的注意力（第11和14行）和当前块的注意力（第10和13行）是分开计算的，因为在当前块的注意力中需要保持额外的因果性。
5. 最后，将注意力输出重新排列回原始顺序，并使用在线 softmax（第16行）进行合并【Milakov et al. 2018, Online normalizer calculation for softmax, 2018, arXiv】; 【H. Liu et al. 2023, Blockwise Parallel Transformer for Large Context Models, 2023, arXiv】。

算法 1 MoBA（混合块注意力）实现
需要：查询、键和值矩阵 $Q, K, V \in R^{N \times h \times d}$；MoBA 超参数（块大小 $B$ 和 top-k）；$h$ 和 $d$ 分别表示注意力头的数量和头的维度。另记 $n = N/B$ 为块的数量。
1: // 将 KV 分割成块
2: $\{ \tilde{K}_i, \tilde{V}_i \} = \text{split_blocks}(K, V, B)$，其中 $\tilde{K}_i, \tilde{V}_i \in R^{B \times h \times d}, i \in [n]$
3: // 计算门控分数以进行动态块选择
4: $\bar{K} = \text{mean_pool}(K, B) \in R^{n \times h \times d}$
5: $S = Q\bar{K}^\top \in R^{N \times h \times n}$
6: // 使用因果约束选择块（不对未来块进行注意力计算）
7: $M = \text{create_causal_mask}(N, n)$
8: $G = \text{topk}(S + M, k)$
9: // 组织注意力模式以提高计算效率
10: $Q_s, \tilde{K}_s, \tilde{V}_s = \text{get_self_attn_block}(Q, \tilde{K}, \tilde{V})$
11: $Q_m, \tilde{K}_m, \tilde{V}_m = \text{index_select_moba_attn_block}(Q, \tilde{K}, \tilde{V}, G)$
12: // 分别计算注意力
13: $O_s = \text{flash_attention_varlen}(Q_s, \tilde{K}_s, \tilde{V}_s, \text{causal=True})$
14: $O_m = \text{flash_attention_varlen}(Q_m, \tilde{K}_m, \tilde{V}_m, \text{causal=False})$
15: // 使用在线 softmax 合并结果
16: $O = \text{combine_with_online_softmax}(O_s, O_m)$
17: return $O$

A4 实验环境


表 1: 扩展定律实验的配置

A4 实验结果

3.1 扩展定律实验与消融研究

3.2 MoBA 与全注意力的混合使用

3.3 大型语言模型评估


表 2: MoBA 与全注意力在不同评估基准上的性能对比。

图 7: Llama-8B-1M-MoBA 在大海捞针基准测试上的性能（上下文长度高达 1M）。

3.4 效率与可扩展性

A5 结论

本文介绍了混合块注意力（MoBA），一种受混合专家（MoE）启发的创新注意力架构，旨在提升大语言模型（LLM）在长上下文任务中的效率和可扩展性。MoBA 通过将上下文分块并采用动态门控机制，将查询路由至最相关的 KV 块，从而在降低计算复杂度的同时保持了模型性能，并实现了全注意力与稀疏注意力之间的无缝切换。

大量实验证明，MoBA 在性能上与全注意力相当，同时显著提升了计算效率。结果表明，MoBA 能够有效扩展至长上下文，在各类基准测试中保持低 LM 损失和高性能。此外，MoBA 的灵活性使其能够以较低的训练成本集成到现有模型中，为提升 LLM 的长上下文能力提供了一个实用的持续预训练解决方案。

总而言之，MoBA 代表了高效注意力领域的一项重要进展，在性能和效率之间取得了平衡。未来的工作可能包括进一步优化 MoBA 的块选择策略，探索其在其他模态上的应用，以及研究其在提升复杂推理任务泛化能力方面的潜力。

A6 附录

A.1 长上下文可扩展性

关注尾部词元损失的动机：为了解决自然数据分布中偏向短上下文的偏差，我们策略性地根据实际位置将整个序列分割成离散的片段。例如，跨越位置 30K-32K 的片段仅反映与上下文长度超过 30K 的文档相关的损失，并且也对从 30K 到 32K 的位置进行了掩码。这种方法确保了在不同上下文长度上进行更均衡和有代表性的评估。

关键发现与简化评估：在探索长上下文可扩展性的过程中，我们做出了一个关键发现：尾部词元（trailing tokens）是导致全上下文基线与新提出的稀疏注意力架构之间性能差异的主要原因。因此，我们通过专注于尾部词元扩展来简化长上下文的扩展过程。这不仅简化了计算需求，还显著提高了研究长上下文场景的效率和有效性。这一发现对未来开发更高效、更可扩展的注意力机制具有重要意义。


图 8: 位置 0-16k 的扩展定律


图 8: 位置 16-32k 的扩展定律

表 3: 不同位置的损失扩展