摘要
扩展有效上下文长度对于推动大语言模型 (LLM) 向通用人工智能 (AGI) 发展至关重要。然而，传统注意力机制中固有的计算复杂度呈二次方增长，带来了令人望而却步的开销。现有的方法要么强加了强偏置结构，例如特定于任务的 sink 或窗口注意力，这些方法通常是任务特定的，要么彻底地将注意力机制修改为线性近似，但这些方法在复杂推理任务中的性能仍未得到充分探索。
在这项工作中，我们提出了一种遵循“更少结构”原则的解决方案，该方案允许模型自主决定关注的位置，而无需引入预定义的偏见。我们引入了块注意力混合（MoBA），这是一种创新的方法，它将专家混合（MoE）的原则应用于注意力机制。 这种新颖的架构在长文本任务中表现出卓越的性能，同时提供了一个关键优势：能够在完整注意力和稀疏注意力之间无缝切换，从而提高效率，且不会有损性能。MoBA 已经部署用于支持 Kimi 的长文本请求，并展示了在大语言模型（LLM）高效注意力计算方面的显著进展。
1 引言
现有的方法通常利用预定义的结构约束来利用这种稀疏性，例如基于汇聚节点的方法 (G. Xiao et al. 2023) 或滑动窗口注意力机制 (Beltagy et al. 2020)。虽然这些方法可能有效，但它们往往是高度特定于任务的，可能会阻碍模型的整体泛化能力。 替代方法，例如稀疏注意力机制，以 Quest (Tang et al. 2024)、Minference (H. Jiang et al. 2024) 和 Re-trievalAttention (Di Liu et al. 2024) 为例，在推理时选择 Token 的子集。 尽管这些方法可以减少长序列的计算量，但它们并没有实质性地减轻长上下文模型的高昂训练成本，因此难以有效地将大语言模型扩展到数百万 Token 的上下文。 最近，另一种有希望的替代方案以线性注意力模型的形式出现，例如 Mamba (Dao and Gu 2024)、RKWV (Peng, Alcalá, et al. 2023)、Peng, Goldstein, et al. 2024 和 ReNet (Sun et al. 2023)。 这些方法用线性近似代替传统的基于 softmax 的注意力，从而减少了长序列处理的计算开销。 然而，由于线性和传统注意力之间存在显着差异，因此调整现有的 Transformer 模型通常会产生高昂的转换成本 (Mercat et al. 2024, J. Wang et al. 2024, Bick et al. 2025, M. Zhang et al. 2024) 或需要从头开始训练全新的模型 (A. Li et al. 2025)。 更重要的是，它们在复杂推理任务中的有效性证据仍然有限。
因此，出现了一个关键的研究问题：我们如何设计一种稳健且适应性强的注意力架构，该架构保留原始 Transformer 框架，同时遵守“弱结构化”的原则，允许模型确定在哪里关注，而不依赖于预定义的偏差？ 理想情况下，这种架构将在完整和稀疏注意力模式之间无缝过渡，从而最大限度地提高与现有预训练模型的兼容性，并实现高效的推理和加速训练，而不会影响性能。
因此，我们引入了块注意力混合模型 (Mixture of Block Attention, MoBA)，这是一种新颖的架构，它建立在专家混合模型 (Mixture of Experts, MoE) (Shafer et al. 2017) 的创新原则之上，并将其应用于 Transformer 模型的注意力机制。MoE 主要用于 Transformer 的前馈网络 (Feedforward Network, FFN) 层 (Lepikhin et al. 2020, Fedus et al. 2022, Pop et al. 2022)，但 MoBA 开创性地将其应用于长上下文注意力，允许为每个查询 Token 动态选择历史上相关的键和值数据块。这种方法不仅提高了 大语言模型 LLM 的效率，而且使它们能够处理更长、更复杂的 提示词 prompt，而不会显著增加资源消耗。MoBA 通过将上下文划分为块，并采用门控机制来选择性地将键/值对路由到相关上下文，从而解决了传统注意力机制的计算问题。这种块稀疏注意力机制有效地降低了计算成本，为高效处理长序列铺平了道路。该模型动态选择信息量最大的键数据块的能力提高了性能和效率，尤其有利于涉及大量序列长度的任务。在本文中，我们详细介绍了 MoBA 的架构，首先是其块划分和路由策略，其次是其与传统注意力机制相比的计算效率。我们进一步展示了实验结果，证明了 MoBA 在需要处理长序列的任务中表现出卓越的性能。我们的工作为高效的注意力计算贡献了一种新颖的方法，突破了 大语言模型 LLM 在处理复杂和冗长的输入方面所能达到的界限。
2 方法
在这项工作中，我们介绍了一种新颖的架构，称为块注意力混合（MoBA），它通过动态选择历史片段（块）来进行注意力计算，从而扩展了 Transformer 模型的能力。MoBA 的灵感来自专家混合（MoE）和稀疏注意力技术。前一种技术主要应用于 Transformer 架构中的前馈网络（FFN）层，而后一种技术已被广泛应用于扩展 Transformer 以处理长上下文。我们的方法创新性地将 MoE 原理应用于注意力机制本身，从而可以更高效、更有效地处理长序列。
2.1 预备知识：Transformer 中的标准注意力
我们首先回顾 Transformer 中的标准注意力。为简单起见，我们回顾单个查询 Token \( $q \in \mathbb{R}^{1 \times d} $\) 注意到 \( $N$ \) 个键和值 Token 的情况，分别表示为 \( $K, V \in \mathbb{R}^{N \times d}$ \)。标准注意力计算如下：
$$ \text{Attn}(q, K, V) = \text{Softmax}(q K^T)V, $$其中 \( $d$ \) 表示单个注意力头的维度。为清楚起见，我们专注于单头场景。扩展到多头注意力涉及连接来自多个此类单头注意力操作的输出。
2.2 MoBA 架构
与标准注意力机制中每个查询 Token 都关注整个上下文不同，MoBA 使得每个查询 Token 仅关注键和值的子集：
$$ \text{MoBA}(q, K, V) = \text{Softmax}(qK[I]^T)V[I], $$其中 \( $I \subseteq [N]$ \) 是所选键和值的集合，\($q$\) 代表查询，\($K$\) 代表键，\($V$\) 代表值，\($N$\) 代表上下文长度。
MoBA 的关键创新在于块划分和选择策略。我们将长度为 \( $N$ \) 的完整上下文划分为 \( $n$ \) 个块，其中每个块代表后续 Token 的子集。不失一般性，我们假设上下文长度 \( $N$ \) 可被块数 \( $n$ \) 整除。我们进一步表示 \( $B = \frac{N}{n}$ \) 为块大小，并且
$$ I_i = ((i-1) \times B + 1, i \times B) $$表示第 \( $i$ \) 个块的范围。通过应用混合专家模型 (MoE) 中的 top-\( k \) 门控机制，我们使每个查询能够有选择地关注来自不同块的 Token 子集，而不是整个上下文：
$$ I = \bigcup I_i. $$$$ g_i= \begin{cases}1 & s_i \in \operatorname{Topk}\left(\left\{s_j \mid j \in[n]\right\}, k\right) \\ 0 & \text { otherwise }\end{cases} $$
一个运行示例
我们在图 1(a) 中提供了一个 MoBA 的运行示例，其中我们有两个查询 Token 和四个 KV 块。路由器（门控网络）动态地选择每个查询要参与的最高的两个块。如图 1(a) 所示，第一个查询被分配给第一个和第二个块，而第二个查询被分配给第三个和第四个块。
在自回归语言模型中保持因果关系非常重要，因为它们通过基于先前 Token 的下一个 Token 预测来生成文本。这种顺序生成过程确保了一个 Token 不能影响它之前的 Token，从而保持了因果关系。MoBA 通过两个特定的设计来保持因果关系：
因果关系：不关注未来的块
MoBA 通过限制注意力范围到当前和过去的块，确保查询 Token 不会被路由到任何未来的块。MoBA 遵循语言建模的自回归特性。更正式地说，将 pos($q$) 表示为查询 $q$ 的位置索引，对于任何块 $i$，如果 $pos(q) < i < B$，则我们设置 $s_i = −∞$ 和 $g_i = 0$。
当前块注意力与因果掩码
我们将“当前块”定义为包含查询 Token 本身的块。路由到当前块也可能违反因果关系，因为跨整个块的均值池化可能会无意中包含来自未来 Token 的信息。为了解决这个问题，我们强制每个 Token 必须被路由到其各自的当前块，并在当前块注意力期间应用因果掩码。这种策略不仅避免了来自后续 Token 的任何信息泄露，而且鼓励关注局部上下文。更正式地说，对于查询 Token 的位置 pos($q$) 位于区间 $[j, j + 1)$ 内的块 $j$，我们设置 $g_j = 1$。从 Mixture-of-Experts (MoE) 的角度来看，MoBA 中的当前块注意力类似于现代 MoE 架构中共享专家的角色，其中在专家选择时添加静态路由规则。
接下来，我们将讨论 MoBA 的一些其他关键设计选择，例如其块分割策略以及 MoBA 和完全注意力的混合。
细粒度块分割
细粒度专家分割在提高模型性能方面的积极影响已在专家混合（MoE）文献中得到充分证明（Dai et al. 2024; A. Yang et al. 2024）。 在这项工作中，我们探索了将类似的细粒度分割技术应用于 MoBA 的潜在优势。MoBA 受 MoE 的启发，沿上下文长度维度而不是 FFN 中间隐藏维度进行分割。因此，我们的研究旨在确定当以更细的粒度划分上下文 Token 时，MoBA 是否也能受益。
MoBA 和完整注意力机制的混合
MoBA 的设计是为了替代完整注意力机制，在不增加任何额外指令的情况下保持相同的参数量。这使得模型能够在完整注意力机制和 MoBA 之间平滑过渡。具体来说，在初始化阶段，每个注意力层都可以选择完整注意力机制或 MoBA，并且如果需要，可以在训练期间动态更改此选择。先前的工作（X. Zhang et al. 2024）已经研究了将完整注意力机制过渡到滑动窗口注意力机制的类似想法。
与滑动窗口注意力机制和注意力Sink机制的比较
滑动窗口注意力机制 (SWA) 和注意力Sink机制是两种流行的稀疏注意力架构。我们证明了这两种机制都可以看作是MoBA的特殊情况。对于滑动窗口注意力机制 (Beltagy et al. 2020)，每个查询Token只关注其相邻的Token。这可以解释为MoBA的一种变体，其门控网络持续选择最近的块。类似地，注意力Sink机制 (G. Xiao et al. 2023) 中，每个查询Token倾向于关注初始Token和最近Token的组合，这可以看作是MoBA的一种变体，其门控网络始终选择初始块和最近的块。上述讨论表明，MoBA比滑动窗口注意力机制和注意力Sink机制具有更强的表达能力。此外，它表明MoBA可以通过结合特定的门控网络来灵活地逼近许多稀疏注意力架构。
总的来说，MoBA的注意力机制允许模型自适应地和动态地关注上下文中信息量最大的块。这对于涉及长文档或序列的任务尤其有利，因为关注整个上下文可能是不必要的且计算成本高昂。MoBA选择性地关注相关块的能力使得能够更细致和高效地处理信息。
2.3 实现
我们提供了一个 MoBA 的高性能实现，通过结合 FlashAttention (Dao, D. Fu et al. 2022) 和 MoE (Rajbhandari et al. 2022) 的优化技术。图 2 展示了 MoBA 的高效率，我们在第 3.4 节详细介绍了关于效率和可扩展性的实验。
我们的实现包括五个主要步骤：
- 根据门控网络和因果掩码，确定查询 Token 到 KV 块的分配。
- 根据分配的 KV 块，排列查询 Token 的顺序。
- 计算每个 KV 块及其分配的查询 Token 的注意力输出。这一步可以通过具有不同长度的 FlashAttention 进行优化。
- 将注意力输出重新排列回其原始顺序。
- 使用在线 Softmax（即平铺）组合相应的注意力输出，因为查询 Token 可能会利用其当前块和多个历史 KV 块。
算法工作流程在算法 1 中进行了形式化，并在图 1b 中进行了可视化，说明了如何基于 MoE 和 FlashAttention 实现 MoBA。首先，KV 矩阵被划分为块（第 1-2 行）。接下来，根据公式 6 计算门控分数，该分数衡量查询 Token 和 KV 块之间的相关性（第 3-7 行）。将 top-k 算子应用于门控分数（以及因果掩码），从而产生稀疏的查询到 KV 块的映射 $v1$，以表示查询到 KV 块的分配（第 8 行）。然后，根据查询到 KV 块的映射排列查询 Token，并计算分块的注意力输出（第 9-12 行）。值得注意的是，对历史块的注意力（第 11 行和第 14 行）和当前块的注意力（第 10 行和第 13 行）是分开计算的，因为需要在当前块的注意力中保持额外的因果关系。最后，注意力输出被重新排列回其原始顺序，并与在线 softmax 组合（第 16 行）(Milakov et al. 2018, H. Liu et al. 2023)。
3 实验
3.1 缩放定律实验和消融研究
在本节中，我们进行缩放定律实验和消融研究，以验证 MoBA 的一些关键设计选择。
关于 LM 损失的可扩展性。 为了评估 MoBA 的有效性，我们通过比较使用完整注意力或 MoBA 训练的语言模型的验证损失来执行缩放定律实验。 遵循 Chinchilla 缩放定律 (Hoffmann et al. 2022)，我们训练了五个不同大小的语言模型，并具有足够数量的训练 Token，以确保每个模型都达到其训练最佳状态。 缩放定律实验的详细配置见表 1。MoBA 和完整注意力模型均以 8K 的序列长度进行训练。 对于 MoBA 模型，我们将块大小设置为 512，并选择前 3 个块进行注意力，从而产生稀疏度高达 \( $\frac{3}{512^2} \approx 81.25\%$ \)的稀疏注意力模式。 特别是，MoBA 可以作为完整注意力的替代方案，这意味着它不会引入新参数或删除现有参数。 这种设计简化了我们的比较过程，因为所有实验中唯一的区别在于注意力模块，而所有其他超参数（包括学习率和批量大小）保持不变。 如图 3a 所示，MoBA 和完整注意力的验证损失曲线显示出非常相似的缩放趋势。 具体来说，这两种注意力机制之间的验证损失差异保持在 1-3 的范围内。 这表明，尽管 MoBA 具有稀疏度高达 75% 的稀疏注意力模式，但它实现了与完整注意力相当的缩放性能。
长文本可扩展性
然而，大语言模型（LLM）损失可能会受到数据长度分布的影响而产生偏差（An et al. 2024），这种分布通常由短序列主导。为了充分评估MoBA的长文本能力，我们评估了尾部Token的LLM损失（简称为尾部LLM损失），它计算序列中最后几个Token的LLM损失。我们只对达到最大序列长度的序列计算此损失，以避免可能由非常短的序列引起的偏差。关于尾部Token缩放的详细讨论可以在附录A.1中找到。这些指标提供了对模型生成序列最后部分的能力的深入了解，这对于涉及长文本理解的任务尤其有帮助。因此，我们采用了一种修改后的实验设置，将最大序列长度从8k增加到32k。这种调整导致MoBA的注意力模式更加稀疏，实现了高达\( $\frac{1 - 31212/3786}{1} = 95.31\%$ \)的稀疏度。如图3b所示，尽管在所有实验中，MoBA的最后一个块的LLM损失略高于完全注意力，但损失差距正在逐渐缩小。这个实验表明了MoBA的长文本可扩展性。
细粒度块分割的消融研究
我们进一步对 MoBA 的块粒度进行了消融研究。我们使用一个具有 1.5B 参数和 32K 上下文长度的模型进行了一系列实验。调整块大小和 top-k 的超参数，以保持一致的注意力稀疏性。具体来说，我们将 32K 上下文划分为 8、16、32、64 和 128 个块，并相应地选择 2、4、8、16 和 32 个块，确保在这些配置中注意力稀疏性为 75%。如图 4 所示，MoBA 的性能受到块粒度的显著影响。具体来说，最粗粒度的设置（从 8 个块中选择 2 个）与更细粒度的设置之间存在 >1.2 的性能差异。这些发现表明，细粒度分割似乎是一种通用技术，可以增强 MoE 系列模型（包括 MoBA）的性能。
3.2 MoBA 和完整注意力的混合
正如第 2 节所讨论的，我们将 MoBA 设计为完全注意力的灵活替代方案，使其能够以最小的开销轻松切换到完全注意力或从完全注意力切换出来，并实现相当的长上下文性能。在本节中，我们首先展示完全注意力和 MoBA 之间的无缝过渡可以成为高效长上下文预训练的解决方案。然后，我们讨论分层混合策略，主要针对监督微调 (SFT) 的性能。
MoBA/完全注意力混合训练
我们训练了三个模型，每个模型具有 15 亿个参数，在 300 亿个 Token 上进行训练，上下文长度为 32K 个 Token。 对于 MoBA 的超参数，块大小设置为 2048，top-k 参数设置为 3。详细的训练方案如下：
- MoBA/完全注意力混合：该模型使用两阶段方案进行训练。在第一阶段，MoBA 用于训练 90% 的 Token。在第二阶段，该模型切换到完全注意力，用于剩余 10% 的 Token。
- 完全注意力：该模型在整个训练过程中都使用完全注意力进行训练。
- MoBA：该模型仅使用 MoBA 进行训练。
我们通过基于位置的语言模型（LM）损失来评估它们的长期上下文性能，这是一种细粒度的指标，用于评估序列中每个位置的 LM 损失。与通过平均所有位置的 LM 损失来计算的原始 LM 损失不同，基于位置的 LM 损失将每个位置的损失单独分解。先前的研究（Xiong et al. 2023; Reid et al. 2024）也提出了类似的指标。
如图 5a 所示，仅使用 MoBA 配方会导致训练 Token 的基于位置的损失更高。重要的是，我们的 MoBA/完整混合配方达到的损失几乎与完整注意力相同。这一结果突出了 MoBA/完整训练配方在平衡训练效率和模型性能方面的有效性。更有趣的是，我们没有观察到在 MoBA 和完整注意力之间切换时出现明显的损失峰值，这再次证明了 MoBA 的灵活性和鲁棒性。
分层混合
MoBA 的灵活性鼓励我们深入研究一种更为精妙的策略——MoBA 与完整注意力机制的分层混合。我们重点研究这种策略在监督微调 (SFT) 期间的应用。之所以研究这种策略，是因为我们观察到 MoBA 在 SFT 期间有时会产生次优的性能，如图 5b 所示。我们推测这可能与 SFT 中使用的损失掩码有关——在 SFT 期间，提示词 Token 通常不包含在损失计算中，这对于像 MoBA 这样的稀疏注意力方法来说，可能会构成稀疏梯度的挑战。因为它可能会阻碍梯度从最初由未掩码的 Token 计算得出，并贯穿整个上下文的反向传播。为了解决这个问题，我们提出了一种混合方法——将最后几个 Transformer 层从 MoBA 切换到完整注意力机制，而其余层继续采用 MoBA。如图 5b 和图 5c 所示，这种策略可以显著降低 SFT 损失。
3.3 大语言模型评估
我们对 MoBA 在各种真实世界的下游任务中进行了彻底的评估，将其性能与全注意力模型进行了比较。为了便于验证，我们的实验从 Llama 3.1 8B 基础模型开始，该模型用作长上下文预训练的起点。该模型，称为 Llama-8B-1M-MoBA，最初以 128K Token 的上下文长度进行训练，并且在持续预训练期间，我们将上下文长度逐渐增加到 256K、512K 和 1M Token。为了简化这种过渡，我们在 256K 持续预训练阶段的开始使用位置插值方法 (S. Chen et al. 2023)。这项技术使我们能够扩展
图 7：Llama-8B-1M-MoBA 在“大海捞针”基准测试中的表现（上下文长度高达 100 万 Token）。有效上下文长度从 128K Token 扩展到 100 万 Token。完成 100 万 Token 的连续预训练后，MoBA 被激活并用于 100B Token 的训练。我们将块大小设置为 4096，并将 top-k 参数设置为 12，这导致注意力稀疏度高达 $1 - (4096^12 / 4096^12) = 95.31%$。为了保留部分完全注意力机制的能力，我们采用了逐层混合策略——最后三层保持完全注意力，而其余 29 层完全注意力层则切换到 MoBA。对于监督微调，我们采用了类似的策略，逐渐将上下文长度从 32K 增加到 1M。基线完全注意力模型（称为 Llama-8B-1M-Full）也遵循类似的训练策略，如图 6 所示，唯一的区别在于整个训练过程中都使用完全注意力机制。这种方法使我们能够在相同的训练条件下直接比较 MoBA 与完全注意力模型的性能。
评估是在几个广泛使用的长文本基准上进行的。特别是在所有评估任务中，MoBA 仅用于预填充；为了获得更好的性能，我们在生成过程中切换到完全注意力机制。如表 2 所示，Llama-8B-1M-MoBA 的性能与 Llama-8B-1M-Full 的性能高度相当。尤其值得注意的是，在最长的基准 RULER 中，MoBA 的稀疏度高达 $1 - (4096^12 / 4096^32) = 62.5%$，Llama-8B-1M-MoBA 几乎与 Llama-8B-1M-Full 的性能相匹配，得分分别为 0.7818 和 0.7849。对于高达 100 万 Token 的上下文长度，我们使用传统的“大海捞针”基准来评估模型。如图 7 所示，即使上下文长度扩展到 100 万 Token，Llama-8B-1M-MoBA 也表现出令人满意的性能。
3.4 效率和可扩展性
上述实验结果表明，MoBA 不仅在语言模型损失方面，而且在现实世界的任务中，都取得了可比的性能。 为了进一步研究其效率，我们比较了在第 3.3 节中训练的两个模型——Llama-8B-1M-MoBA 和 Llama-8B-1M-Full 中注意力层的前向传播时间。 我们只关注注意力层，因为所有其他层（例如，FFN）在两个模型中都具有相同的 FLOPs。 如图 2a 所示，MoBA 在所有上下文长度上都比完整注意力更有效，这表明它具有亚二次计算复杂度。 特别是，当预填充 100 万个 Token 时，它实现了高达 6.5 倍的加速比。 我们还通过逐渐将上下文长度增加到 1000 万个 Token 来探索 MoBA 的长度可扩展性。 为了保持恒定的注意力稀疏性，我们保持 top-k 值和 MoBA 块的数量不变，同时按比例增加块大小。 为了达到 10M 的上下文长度，我们将张量并行性（Shoeybi et al. 2019）扩展到查询头级别。 具体来说，我们在分布式查询头之间广播键和值张量，有效地解决了 GPU 内存限制，同时保持了计算效率。 如图 2b 所示，当扩展到更长的序列时，MoBA 表现出比标准 Flash Attention 更优越的效率。 具体来说，在 10M Token 时，MoBA 实现了 16 倍的注意力计算时间减少的加速比。 顶部图中的插图，侧重于较短的序列（32K 到 512K），表明即使两种方法在较小规模上表现相当，但随着序列越来越长，MoBA 的计算优势变得越来越明显，突显了其在处理极长序列方面的特殊优势。
总的来说，MoBA 的高效率可以归因于两个关键创新：(1) 块稀疏注意力机制，以及 (2) 结合了混合专家模型 (MoE) 和 FlashAttention 的优化实现，如第 2.3 节所述。这些技术有效地解决了完整注意力的二次复杂度限制，将计算复杂度降低到更经济的亚二次规模。
4 相关工作
高效注意力机制的发展 (Tay, Dehghani, et al. 2020) 一直是处理更长序列和降低计算成本的关键研究领域。随着上下文长度的增长，高效注意力技术已成为解决自注意力机制的二次复杂度，同时保持模型性能的有希望的解决方案。
静态稀疏模式
诸如 Sparse Transformer (Child et al. 2019), Star-Transformer (Gu et al. 2019), BlockBERT (Que et al. 2019), Longformer (Beltagy et al. 2020), GMAT (Gupta et al. 2020), ETC (Ainslie, Ontanon, et al. 2020), BigBird (Zaheer et al. 2020), LongTS (M. Guo et al. 2021) 和 D. Jing et al. 2023) 等诸多努力都致力于大语言模型 (LLM) 中静态注意力模式的设计。 他们对静态注意力模式的选择包括跨步式和固定式注意力、窗口注意力、全局 Token 注意力、随机注意力、Delta 注意力、块稀疏注意力或它们的任意组合。 在多模态模型的领域中，也已经开发了静态稀疏注意力机制，例如用于 2D 图像的轴向注意力 (Ho et al. 2019) 和用于 3D 视频的时空注意力 (Zheng et al. 2024)。
动态稀疏模式
与静态模式不同，动态稀疏注意力技术自适应地确定要关注哪些 Token（Token）。Reformer (Kitaev et al. 2020) 和 Routing Transformer (Roy et al. 2021) 分别采用局部敏感哈希 (LSH) 和 K-means 来聚类 Token（Token），并关注于聚类而不是完整的上下文。Memorizing Transformers (Yuhai Wu et al. 2022) 和 Unformer (Bertsch et al. 2024) 动态地关注由 k 最近邻 (kNN) 算法选择的 Token（Token）。CoLTs (Ainslie, Lei, et al. 2023) 设计了一个路由模块来选择最重要的查询，而 Sparse Sinkhorn Attention (Tay, Bahri, et al. 2020) 学习置换输入序列中的块，从而实现动态块稀疏注意力计算。
无需训练的稀疏注意力 (Training-free Sparse Attention)
除了之前讨论的研究训练稀疏注意力机制的方法之外，还有一些策略旨在结合稀疏注意力机制，以提高模型推理的两个主要阶段的效率——无论是预填充阶段还是解码阶段，亦或是两者同时进行优化。在预填充优化阶段，可以使用完整的提示词进行注意力分析，从而探索更为复杂的稀疏注意力模式。例如，MoA (T. Fu et al. 2024)、Minference (H. Jiang et al. 2024) 和 SeerAttention (Y. Gao et al. 2024) 已经研究了诸如 a 字形、垂直斜线以及动态块稀疏性等稀疏注意力配置。在解码器优化方面，大量工作致力于压缩和修剪 KV 缓存，以在文本生成的速度和质量之间取得平衡。该领域值得关注的研究包括 HZO (Z. Zhang et al. 2024)、StreamingLLM (G. Xiao et al. 2023)、TOVA (Oren et al. 2024)、FastGen (Ge et al. 2023) 和 Quest (Tang et al. 2024)。特别是，Quest 可以被看作是 MoBA 的一种变体，它具有更小的块大小以及一种特殊的块表示函数，该函数结合了最小池化和最大池化。另一项与 MoBA 密切相关的工作是 Longheads (Y. Lu et al. 2024)，它可以被视为具有 Top-1 门控网络的 MoBA，这意味着每个查询都会选择最相关的 KV 块来进行注意力计算。
超越传统注意力架构：
另一项研究着眼于探索与传统注意力机制不同的新型模型架构。由于架构的改变，这些方法需要从头开始训练模型，因此无法重用预训练的 Transformer 模型。该领域的研究探索了受卷积神经网络 (CNN)、循环神经网络 (RNN)、状态空间模型 (SSM) 或线性注意力 (Katharopoulos et al. 2020) 启发而设计的架构。这些模型的例子包括 Hyena (Poli et al. 2023)、Performer (Choromanski et al. 2020)、Linformer (S. Wang et al. 2020)、RKVM Peng, Alcaide, et al. 2023)、Mamba (Ou et al. 2023)、ReNet (Sun et al. 2023) 等。总而言之，高效注意力技术的应用场景非常广泛，涵盖了从静态到动态的稀疏模式、从训练到推理的优化目标，以及从传统注意力机制到创新替代方案的各种架构。每种方法都有其独特的优势和不足，具体技术的选择通常取决于应用场景的特定需求，例如最大序列长度、 计算资源以及效率与性能之间的理想平衡。随着该领域研究的不断深入，这些方法有望在帮助 大语言模型 处理更加复杂的任务，同时保持效率和可扩展性方面发挥关键作用。
5 结论
在本文中，我们介绍了一种名为块注意力混合（MoBA）的新型注意力架构。该架构的灵感来源于专家混合（MoE）的原则，旨在提升大型语言模型（LLM）在处理长文本任务时的效率和可扩展性。MoBA 通过将上下文分割成多个块，并采用动态门控机制来选择性地路由查询模型性能，从而解决了传统注意力机制所面临的计算挑战。此外，MoBA 还支持在完全注意力和稀疏注意力之间进行无缝切换。通过大量的实验，我们证明了 MoBA 在显著提升效率和可扩展性的同时，能够达到与完全注意力相媲美的性能。
总而言之，MoBA 在各种基准测试中均表现出色。更重要的是，MoBA 具有高度的灵活性，可以与现有模型轻松集成，无需付出大量的训练成本，因此它是一种实用的持续预训练解决方案，能够有效增强 LLM 的长文本处理能力。总结来说，MoBA 代表了高效注意力机制的一大进步，在性能和效率之间实现了良好的平衡。未来的研究方向包括进一步优化 MoBA 的块选择策略，探索其在其他模态中的应用，以及研究其在提升复杂推理任务中的泛化能力方面的潜力。