月之暗面-KIMI-发布最新架构MoBA

MoBA：用于长上下文大语言模型的混合块注意力机制 - 技术报告

摘要

扩展有效上下文长度对于推动大语言模型（LLM）迈向通用人工智能（AGI）至关重要。然而，传统注意力机制固有的二次计算复杂度增长带来了难以承受的开销。现有的方法要么强加高度偏向的结构（如sink注意力或窗口注意力，这些方法通常针对特定任务），要么彻底修改注意力机制，采用线性近似，但其复杂推理任务中的表现仍需进一步探索。

在本文中，我们提出了一种遵循"少结构"原则的解决方案，允许模型自主决定注意力分配，而不是引入预定义的偏向。我们引入了混合块注意力（MoBA） ，这是一种创新方法，将专家混合（MoE）的原理应用于注意力机制。这种新颖的架构在长上下文任务中表现出色，同时具有一个关键优势：能够在全注意力和稀疏注意力之间无缝过渡，提高效率的同时不会影响性能。MoBA 已部署用于支持 Kimi 的长上下文请求，并在 LLM 的高效注意力计算方面取得了重大进展。我们的代码可在 https://github.com/MoonshotAI/moba 获取。

1. 引言

追求通用人工智能（AGI）推动了大型语言模型（LLM）向前所未有的规模发展，其目标是处理模仿人类认知的复杂任务。实现 AGI 的一个关键能力是处理、理解和生成长序列，这对于从历史数据分析到复杂推理和决策过程的各种应用至关重要。对扩展上下文处理的需求日益增长，这不仅体现在像 Kimi (MoonshotAI 2023)、Claude (Anthropic 2023) 和 Gemini (Reid et al. 2024) 等模型对长输入提示理解能力的流行上，也体现在对长链式思维（CoT）输出能力的最新探索中，例如 Kimi k1.5 (Team et al. 2025)、DeepSeek-R1 (D. Guo et al. 2025) 和 OpenAI o1/o3 (Guan et al. 2024)。

然而，由于原始注意力机制（Vaswani et al. 2017）的计算复杂度呈二次增长，扩展 LLM 中的序列长度并非易事。这一挑战激发了大量旨在提高效率而不牺牲性能的研究工作。一个主要方向是利用注意力分数固有的稀疏性。这种稀疏性在数学上源于 softmax 操作，其中各种稀疏注意力模式已被研究（H. Jiang et al. 2024），在生物学上也观察到（Watson et al. 2025），与记忆存储相关的脑区中观察到稀疏连接。

现有方法

现有的方法通常利用预定义的结构化约束，例如：

Sink-based 注意力 (G. Xiao et al. 2023) 或 滑动窗口注意力 (Beltagy et al. 2020) 等方法利用注意力分数的稀疏性来提高效率。
- 优点：在特定任务上有效。
- 缺点：高度任务特定，可能阻碍模型的整体泛化能力。
动态稀疏注意力机制，例如：
- Quest (Tang et al. 2024)、Minference (H. Jiang et al. 2024) 和 RetrievalAttention (Di Liu et al. 2024) 等方法在推理时选择标记的子集。
  - 优点：可以减少长序列的计算量。
  - 缺点：并未显著降低长上下文模型的大量训练成本，难以高效地将 LLM 扩展到数百万个标记的上下文。
线性注意力模型 ，例如 Mamba (Dao and Gu 2024)、RWKV (Peng, Alcaide, et al. 2023; Peng, Goldstein, et al. 2024) 和 RetNet (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)。
  - 在复杂推理任务中的有效性证据仍然有限。

MoBA：混合块注意力机制

为了解决上述问题，我们提出了一种新的注意力架构：混合块注意力（MoBA），它：

遵循"少结构"原则：不依赖预定义的偏向，让模型自主决定注意力分配。
基于专家混合（MoE）原理：将 MoE 应用于注意力机制本身，而不是仅应用于前馈网络（FFN）层。
实现全注意力和稀疏注意力的无缝过渡：在保持性能的同时提高效率。

2.1 标准 Transformer 中的注意力机制回顾

在 Transformer 中，单个查询标记 q ∈ R 1 × d q \in \mathbb{R}^{1 \times d} q∈R1×d 关注 N N N 个键和值标记，分别表示为 K , V ∈ R N × d K, V \in \mathbb{R}^{N \times d} K,V∈RN×d。标准注意力计算如下：

Attn ( q , K , V ) = Softmax ( q K T ) V \text{Attn}(q, K, V) = \text{Softmax}(qK^T)V Attn(q,K,V)=Softmax(qKT)V

其中 d d d 表示单个注意力头的维度。我们为了清晰起见关注单头场景。多头注意力的扩展涉及将多个这样的单头注意力操作的输出连接起来。

2.2 MoBA 架构

与标准注意力机制中每个查询标记关注整个上下文不同，MoBA 允许每个查询标记仅关注键和值的一个子集：

MoBA ( q , K , V ) = Softmax ( q K [ I ] T ) V [ I ] \text{MoBA}(q, K, V) = \text{Softmax}(qK[I]^T)V[I] MoBA(q,K,V)=Softmax(qK[I]T)V[I]

其中 I ⊂ [ N ] I \subset [N] I⊂[N] 是选定的键和值集合。

MoBA 的关键创新在于块划分和选择策略。 我们将长度为 N N N 的完整上下文划分为 n n n 个块，每个块代表后续标记的一个子集。不失一般性，我们假设上下文长度 N N N 可被块数 n n n 整除。我们进一步用 B = N n B = \frac{N}{n} B=nN 表示块大小，并用

I i = [ ( i − 1 ) B + 1 , i B ] I_i = [(i-1)B+1, iB] Ii=[(i−1)B+1,iB]

表示第 i i i 个块的区间。通过应用 MoE 的 top-k 门控机制，我们使每个查询能够有选择地关注来自不同块的不同标记子集，而不是整个上下文：

I = ⋃ I i 其中 g i > 0 I = \bigcup I_i \quad \text{其中} \quad g_i > 0 I=⋃Ii其中gi>0

模型采用门控机制（公式 4 中的 g i g_i gi）为每个查询标记选择最相关的块。MoBA 门控首先计算亲和力分数 s i s_i si，衡量查询 q q q 与第 i i i 个块的关联度，并对所有块应用 top-k 门控。更正式地说，第 i i i 个块的门控值 g i g_i gi 计算如下：

g i = { 1 如果 s i ∈ Topk ( { s i ∣ i ∈ [ n ] } , k ) 0 其他情况 g_i = \begin{cases} 1 & \text{如果} \quad s_i \in \text{Topk}(\{s_i \mid i \in [n]\}, k) \\ 0 & \text{其他情况} \end{cases} gi={10如果si∈Topk({si∣i∈[n]},k)其他情况

其中，Topk(., k) 表示包含为每个块计算的亲和力分数中最高的 k k k 个分数的集合。在这项工作中，分数 s i s_i si 通过查询 q q q 与键 K [ I i ] K[I_i] K[Ii] 沿序列维度的均值池化的内积计算：

s i = ⟨ q , mean_pool ( K [ I i ] ) ⟩ s_i = \langle q, \text{mean\_pool}(K[I_i]) \rangle si=⟨q,mean_pool(K[Ii])⟩

保持因果关系： 自回归语言模型必须保持因果关系，因为它们通过基于先前标记的下一个标记预测来生成文本。这种顺序生成过程确保标记不会影响其之前的标记，从而保持因果关系。MoBA 通过以下两种特定设计来保持因果关系：

因果关系：不关注未来块。 MoBA 确保查询标记不会被路由到任何未来块。通过将注意力范围限制在当前和过去的块，MoBA 遵循了语言建模的自回归特性。更正式地说，用 pos(q) 表示查询 q 的位置索引，我们对任何块 i i i 设置 s i = − ∞ s_i = -\infty si=−∞ 和 g i = 0 g_i = 0 gi=0，只要 pos(q) < i x B。
当前块注意力和因果掩码。 我们将"当前块"定义为包含查询标记本身的块。对当前块的路由也可能违反因果关系，因为对整个块的均值池化可能会无意中包含来自未来标记的信息。为了解决这个问题，我们强制每个标记必须路由到其各自的当前块，并在当前块注意力期间应用因果掩码。这种策略不仅避免了后续标记的信息泄漏，还鼓励对局部上下文的关注。更正式地说，我们对块 i i i 设置 g i = 1 g_i = 1 gi=1，其中查询标记的位置 pos(q) 在区间 I j I_j Ij 内。从专家混合（MoE）的角度来看，MoBA 中的当前块注意力类似于现代 MoE 架构中共享专家的角色 (Dai et al. 2024; A. Yang et al. 2024)，其中在专家选择时添加了静态路由规则。

其他关键设计选择：

细粒度块分割： MoE 文献中已经充分记录了细粒度专家分割在提高模型性能方面的积极影响 (Dai et al. 2024; A. Yang et al. 2024)。在这项工作中，我们探索了将类似的细粒度分割技术应用于 MoBA 的潜在优势。MoBA 受 MoE 启发，沿着上下文长度维度而不是 FFN 中间隐藏维度进行分割。因此，我们的调查旨在确定 MoBA 是否也能在将上下文划分为更细粒度的块时受益。更多实验结果见第 3.1 节。
MoBA 和全注意力的混合： MoBA 旨在替代全注意力，保持相同的参数数量，不进行任何添加或删除。这一特性激励我们进行全注意力和 MoBA 之间的平滑过渡。具体来说，在初始化阶段，每个注意力层可以选择全注意力或 MoBA，如果需要，在训练过程中可以动态更改。类似地将全注意力过渡到滑动窗口注意力的想法已在之前的工作中进行了研究 (X. Zhang et al. 2024)。更多实验结果见第 3.2 节。
与滑动窗口注意力和注意力汇聚的比较： 滑动窗口注意力（SWA）和注意力汇聚是两种流行的稀疏注意力架构。我们证明，两者都可以看作是 MoBA 的特例。对于滑动窗口注意力 (Beltagy et al. 2020)，每个查询标记仅关注其相邻标记。这可以解释为 MoBA 的一种变体，其门控网络始终选择最近的块。同样，注意力汇聚 (G. Xiao et al. 2023)，其中每个查询标记关注初始标记和最近标记的组合，可以看作是 MoBA 的一种变体，其门控网络始终选择初始和最近块。上述讨论表明，MoBA 比滑动窗口注意力和注意力汇聚具有更强的表达能力。此外，它表明 MoBA 可以通过结合特定的门控网络，灵活地近似许多静态稀疏注意力架构。

总体而言，MoBA 的注意力机制使模型能够自适应、动态地关注上下文中信息量最大的块。这对于涉及长文档或序列的任务特别有利，因为关注整个上下文可能是不必要且计算成本高昂的。MoBA 选择性地关注相关块的能力使得信息处理更加细致和高效。

2.3 实施

我们提供了 MoBA 的高性能实现，通过结合 FlashAttention (Dao, D. Fu, et al. 2022) 和 MoE (Rajbhandari et al. 2022) 的优化技术。图 2 展示了 MoBA 的高效率，而我们将效率与可扩展性的详细实验推迟到第 3.4 节。我们的实施包括五个主要步骤：

根据门控网络和因果掩码确定查询标记对 KV 块的分配。
根据其分配的 KV 块对查询标记的顺序进行排列。
计算每个 KV 块及其分配的查询标记的注意力输出。此步骤可以通过具有不同长度的 FlashAttention 进行优化。
将注意力输出重新排列回其原始顺序。
使用在线 Softmax（即平铺）组合相应的注意力输出，因为查询标记可能会关注其当前块和多个历史 KV 块。

算法流程在算法 1 中进行了形式化，并在图 1b 中进行了可视化，说明了如何基于 MoE 和 FlashAttention 实现 MoBA。首先，KV 矩阵被划分为块（第 1-2 行）。接下来，根据公式 6 计算门控分数，该分数衡量查询标记与 KV 块的关联度（第 3-7 行）。对门控分数应用 top-k 运算符（连同因果掩码），从而得到一个稀疏的查询到 KV 块映射矩阵 G G G 来表示查询到 KV 块的分配（第 8 行）。然后，根据查询到 KV 块映射对查询标记进行排列，并计算块级注意力输出（第 9-12 行）。值得注意的是，对历史块的注意力（第 11 行和第 14 行）和当前块注意力（第 10 行和第 13 行）是分开计算的，因为在当前块注意力中需要保持额外的因果关系。最后，注意力输出被重新排列回其原始顺序，并使用在线 Softmax 组合（第 16 行）。

3. 实验

3.1 扩展性定律实验和消融研究

在本节中，我们进行扩展性定律实验和消融研究，以验证 MoBA 的一些关键设计选择。

关于 LM 损失的扩展性： 为了评估 MoBA 的有效性，我们进行扩展性定律实验，通过比较使用全注意力和 MoBA 训练的的语言模型的验证损失。根据 Chinchilla 扩展性定律 (Hoffmann et al. 2022)，我们训练了五个不同大小的语言模型，并使用足够的训练标记以确保每个模型达到其训练最优。扩展性定律实验的详细配置见表 1。MoBA 和全注意力模型都使用 8K 的序列长度进行训练。对于 MoBA 模型，我们将块大小设置为 512，并选择 top-3 块进行注意力，从而形成稀疏度高达 1 - 512x3 = 81.25% 的稀疏注意力模式。特别是，MoBA 作为全注意力的替代品，意味着它不会引入新参数或删除现有参数。这种设计简化了我们的比较过程，因为所有实验中唯一的区别在于注意力模块，而所有其他超参数，包括学习率和批量大小，都保持不变。如图 3a 所示，MoBA 和全注意力的验证损失曲线显示出非常相似的扩展趋势。具体来说，这两种注意力机制之间的验证损失差异保持在 1e-3 的范围内。这表明，尽管 MoBA 的稀疏注意力模式的稀疏度高达 75%，但它实现了与全注意力相当的扩展性能。

长上下文扩展性： 然而，LM 损失可能会因数据长度分布 (An et al. 2024) 而产生偏差，后者通常以短序列为主。为了全面评估 MoBA 的长上下文能力，我们评估了尾部标记的 LM 损失（简称尾部 LM 损失），它计算序列中最后几个标记的 LM 损失。我们仅对达到最大序列长度的序列计算此损失，以避免可能由非常短的序列引起的偏差。关于尾部标记扩展性的详细讨论见附录 A.1。

这些指标可以深入了解模型生成序列最后部分的能力，这对于涉及长上下文理解的任务特别有用。因此，我们采用修改后的实验设置，将最大序列长度从 8k 增加到 32k。这一调整导致 MoBA 的注意力模式更加稀疏，稀疏度水平达到 1 - 512x3 = 95.31%。如图 3b 所示，尽管 MoBA 在所有五次实验中与全注意力相比表现出略高的最后一个块 LM 损失，但损失差距正在逐渐缩小。这个实验意味着 MoBA 的长上下文扩展性。

细粒度块分割的消融研究： 我们进一步消融了 MoBA 的块粒度。我们使用具有 32K 上下文长度的 1.5B 参数模型进行了一系列实验。块大小和 top-k 的超参数经过调整，以保持一致的注意力稀疏度。具体来说，我们将 32K 上下文划分为 8、16、32、64 和 128 个块，并相应地选择 2、4、8、16 和 32 个块，确保这些配置之间的注意力稀疏度为 75%。如图 4 所示，MoBA 的性能受块粒度的影响很大。具体来说，最粗粒度设置（从 8 个块中选择 2 个块）与具有更细粒度的设置之间存在 1e-2 的性能差异。这些发现表明，细粒度分割似乎是提高 MoE 系列模型（包括 MoBA）性能的一种通用技术。

3.2 MoBA 和全注意力的混合

如第 2 节所讨论的，我们设计 MoBA 作为全注意力的灵活替代品，以便它可以轻松地与/从全注意力无缝过渡，并在保持长上下文性能的同时实现高效的预训练。在本节中，我们首先展示了全注意力和 MoBA 之间的无缝过渡可以作为高效长上下文预训练的解决方案。然后我们讨论了分层混合策略，主要针对监督微调（SFT）的性能。

MoBA/全混合训练： 我们训练了三个模型，每个模型都有 1.5B 参数，在 30B 个标记的 32K 标记上下文上训练。对于 MoBA 的超参数，块大小设置为 2048，top-k 参数设置为 3。详细的训练配方如下：

MoBA/全混合： 该模型采用两阶段配方进行训练。在第一阶段，MoBA 用于训练 90% 的标记。在第二阶段，模型切换到全注意力以完成剩余 10% 的标记。
全注意力： 该模型在整个训练过程中使用全注意力。
MoBA： 该模型专门使用 MoBA 进行训练。

我们通过位置级语言模型（LM）损失来评估它们的长上下文性能，这是一种细粒度指标，用于评估序列中每个位置处的 lm 损失。与计算所有位置的平均 LM 损失的普通 LM 损失不同，位置级 LM 损失将每个位置的损失分别分解。类似指标已由之前的研究 (Xiong et al. 2023; Reid et al. 2024) 建议，他们注意到位置级 LM 损失相对于上下文长度遵循幂律趋势。如图 5a 所示，MoBA 专用配方导致尾部标记的位置级损失更高。重要的是，我们的 MoBA/全混合配方达到了与全注意力几乎相同的损失。这个结果突出了 MoBA/全混合训练配方在平衡训练效率与模型性能方面的有效性。更令人感兴趣的是，在 MoBA 和全注意力之间的切换过程中，我们没有观察到明显的损失峰值，这再次证明了 MoBA 的灵活性和稳健性。

分层混合： MoBA 的这种灵活性促使我们探索一种更复杂的策略------MoBA 和全注意力的分层混合。我们调查这种策略时，特别关注其在监督微调（SFT）中的应用。这种策略的动机源于我们的观察，即 MoBA 在 SFT 期间有时会导致次优性能，如图 5b 所示。我们推测，这可能是由于在 SFT 中采用的损失掩码------提示标记通常被排除在损失计算之外，这可能会给 MoBA 等稀疏注意力方法带来稀疏梯度挑战。因为它可能会阻碍从未掩码标记初始计算的梯度的反向传播，从而影响整个上下文。为了解决这个问题，我们提出了一种混合方法------将最后几个 Transformer 层从 MoBA 切换到全注意力，而其余层继续使用 MoBA。如图 5b 和图 5c 所示，这种策略可以显著降低 SFT 损失。

3.3 大语言模型评估

我们对 MoBA 进行了全面的评估，评估其在各种现实世界下游任务中的性能，并将其与全注意力模型进行比较。为了便于验证，我们的实验从 Llama 3.1 8B Base Model 开始，该模型用作长上下文预训练的开始。该模型称为 Llama-8B-1M-MoBA，最初使用 128K 标记的上下文长度进行训练，我们逐渐将上下文长度增加到 256K、512K 和 1M 标记。在 256K 连续预训练阶段开始时，我们使用位置插值方法 (S. Chen et al. 2023) 来缓解这种过渡。这种技术使我们能够将有效上下文长度从 128K 标记扩展到 1M 标记。在完成 1M 连续预训练后，激活 MoBA 以完成 100B 标记。我们将块大小设置为 4096，top-K 参数设置为 12，从而使注意力稀疏度达到 1 - 4096x12 = 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 - 4096x12 = 62.5% 的情况下运行，Llama-8B-1M-MoBA 的得分几乎与 Llama-8B-1M-Full 接近，得分为 128K 0.7818 相比 0.7849。对于高达 1M 标记的上下文长度，我们使用传统的"草堆中的针"基准测试来评估模型。如图 7 所示，即使在扩展的 1M 标记上下文长度下，Llama-8B-1M-MoBA 也表现出令人满意的性能。

3.4 效率和可扩展性

上述实验结果表明，MoBA 不仅在语言模型损失方面，而且在现实世界任务中也取得了可比较的性能。为了进一步研究其效率，我们比较了第 3.3 节中训练的两个模型------Llama-8B-1M-MoBA 和 Llama-8B-1M-Full 的注意力层的正向传递时间。我们仅关注注意力层，因为所有其他层（例如 FFN）在两个模型中具有相同的 FLOPs。如图 2a 所示，MoBA 在所有上下文长度上都比全注意力更高效，展示了次二次计算复杂度。特别是，在预填充 1M 标记时，它实现了高达 6.5 倍的速度提升。

我们还通过将上下文长度逐渐增加到 1000 万个标记来探索 MoBA 的长度可扩展性。为了保持恒定的注意力稀疏度，我们在按比例增加块大小的同时保持 top-k 值和 MoBA 块的数量。为了达到 10M 上下文长度，我们将张量并行性 (Shoeybi et al. 2019) 扩展到查询头级别，具体来说，我们将键和值张量广播到分布式查询头，有效地解决了 GPU 内存限制，同时保持了计算效率。如图 2b 所示，MoBA 在扩展到更长序列时表现出比标准 FlashAttention 更高的效率。具体来说，在 10M 标记时，MoBA 实现了注意力计算时间减少 16 倍的速度提升。顶部图形中的插图图，关注较短的序列（32K 到 512K），显示尽管两种方法在小规模时表现相当，但随着序列变长，MoBA 的计算优势变得越来越明显，这突出了其在处理极长序列方面的特殊优势。

总体而言，MoBA 的高效率可归因于两个关键创新：(1) 块稀疏注意力机制；(2) 结合专家混合（MoE）和 FlashAttention 的优化实施，如第 2.3 节所述。这些技术有效地解决了全注意力的二次复杂度限制，将计算复杂度降低到更经济的次二次规模。

4. 相关工作

在自然语言处理领域，特别是在大型语言模型（LLM）的兴起过程中，高效注意力机制的发展一直是关键的研究领域。随着对处理更长序列和降低计算成本的需求不断增长，高效注意力技术已成为减少自我注意力机制的二次复杂度同时保持模型性能的解决方案。

静态稀疏模式： 人们投入了大量努力，例如 Sparse Transformer (Child et al. 2019)、Star-Transformer (Q. Guo et al. 2019)、BlockBERT (Qiu et al. 2019)、Longformer (Beltagy et al. 2020)、GMAT (Gupta et al. 2020)、ETC (Ainslie, Ontanon, et al. 2020)、BigBird (Zaheer et al. 2020)、LongT5 (M. Guo et al. 2021) 和 LongNet (J. Ding et al. 2023)，致力于 LLM 中静态注意力模式的设计。他们对静态注意力模式的选择可以包括跨步和固定注意力、窗口注意力、全局标记注意力、随机注意力、扩张注意力、块稀疏注意力或它们的任意组合。在多模态模型领域，静态稀疏注意力机制也已得到开发，例如用于 2D 图像的轴向注意力 (Ho et al. 2019) 和用于 3D 视频的空间-时间注意力 (Z. Zheng et al. 2024)。

动态稀疏模式： 与静态模式不同，动态稀疏注意力技术自适应地确定关注哪些标记。Reformer (Kitaev et al. 2020) 和 Routing Transformer (Roy et al. 2021) 分别采用局部敏感哈希（LSH）和 K-means 来聚类标记，并关注集群而不是整个上下文。记忆化 Transformer (Yuhuai Wu et al. 2022) 和 Unlimiformer (Bertsch et al. 2024) 动态关注由 k-最近邻（kNN）算法选择的标记。CoLT5 (Ainslie, Lei, et al. 2023) 设计了路由模块以选择最重要的查询和键。稀疏 Sinkhorn 注意力 (Tay, Bahri, et al. 2020) 学习对输入序列的块进行置换，允许动态块稀疏注意力计算。

免训练稀疏注意力： 除了前面讨论的方法之外，还有一些策略旨在将稀疏注意力机制纳入到提高模型推理的两个主要阶段（即预填充阶段或解码阶段，或两者）的效率。在预填充优化阶段，可以利用完整的提示符进行注意力分析，这允许探索更复杂的稀疏注意力模式。例如，MoA (T. Fu et al. 2024)、Minference (H. Jiang et al. 2024) 和 SeerAttention (Y. Gao et al. 2024) 研究了 A 形、垂直斜线形和动态块稀疏性等稀疏注意力配置。在解码优化方面，大量的工作致力于压缩和修剪 KV 缓存，以在文本生成的质量和速度之间取得平衡。在这方面值得注意的努力包括 H2O (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)，它可以看作是 MoBA，其 top-1 门控网络意味着每个查询选择最相关的 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)、RWKV Peng, Alcaide, et al. 2023, Mamba (Gu et al. 2023)、RetNet (Sun et al. 2023) 等。

总之，高效注意力技术的格局是多样的，涵盖了从静态到动态的稀疏模式，从训练到推理的优化目标，以及从传统注意力机制到创新替代品的架构。每种方法都有其独特的优点和权衡，选择哪种技术通常取决于应用程序的具体要求，例如最大序列长度、计算资源以及效率与性能之间的期望平衡。随着该领域研究的不断发展，这些方法预计将在使 LLM 能够处理日益复杂的任务的同时，保持效率与可扩展性方面发挥关键作用。

5. 结论

在本文中，我们介绍了混合块注意力（MoBA），这是一种受专家混合（MoE）原理启发的创新注意力架构，旨在提高大型语言模型（LLM）处理长上下文任务的能力。MoBA 通过将上下文划分为块并采用动态门控机制，将查询标记路由到最相关的 KV 块，从而解决了传统注意力机制带来的计算挑战。这种方法不仅降低了计算复杂度，而且保持了模型性能。此外，它允许在全注意力和稀疏注意力之间无缝过渡。通过广泛的实验，我们证明了 MoBA 在提高计算效率的同时，实现了与全注意力相当的性能。我们的结果还表明，MoBA 可以有效地扩展到长上下文，保持低 LM 损失和在各种基准测试中的高性能。此外，MoBA 的灵活性使其能够与现有模型集成，而不会产生大量的训练成本，这使其成为增强 LLM 长上下文能力的实用的持续预训练解决方案。总之，MoBA 代表了高效注意力机制的重大进步，在性能与效率之间提供了一种平衡的方法。未来的工作可能会探索 MoBA 的块选择策略的进一步优化，研究其应用于其他模式，并研究其提高复杂推理任务中的泛化能力的潜力。