赖勋浩

MiniMax

徐伟琦

MiniMax

杨宇峰

MiniMax

陈乔睿

徐阳

MiniMax

曾伦彬

MiniMax

华中科技大学

李晓龙

MiniMax

孙浩海

MiniMax

朱海超

MiniMax

Vito Zhang

MiniMax

胡金凯

MiniMax

李佳瑶

MiniMax

高睿

MiniMax

南京大学

李泽坤

MiniMax

朱松泉

MiniMax

周靖凯

MiniMax

杭州电子科技大学

赵鹏宇

MiniMax

摘要

超长上下文能力正成为前沿大语言模型不可或缺的能力：智能体工作流、仓库级代码推理和持久化记忆都要求模型能够联合关注数十万到数百万个模型 token——然而 softmax 注意力的二次方成本使得这一需求在部署规模下难以实现。我们提出了 MiniMax 稀疏注意力（MSA），这是一种基于分组查询注意力（GQA）构建的分块稀疏注意力机制。一个轻量级的索引分支对键值块进行评分，并为每个 GQA 组独立选择 Top- 子集，从而在保持高效块级执行的同时实现组特定的稀疏检索；主分支则仅对所选块执行精确的块稀疏注意力。MSA 的设计遵循简洁与可扩展的原则，刻意精简，使其能够在广泛的 GPU 上高效部署。为了将稀疏性转化为实际的速度提升，我们将 MSA 与一条 GPU 执行路径协同设计，该路径使用无指数 Top- 选择和 KV 外稀疏注意力，以在块粒度访问下提高张量核心利用率。在一个经过原生多模态训练的 109B 参数模型上，MSA 的性能与 GQA 相当，同时在 1M 上下文长度下将每个模型 token 的注意力计算量减少了。结合我们协同设计的内核，MSA 在 H800 上实现了预填充和解码的端到端加速。我们的推理内核可在以下地址获取：https://github.com/MiniMax-AI/MSA。一个由 MSA 驱动的、达到生产级别的原生多模态模型已在以下地址公开发布：https://huggingface.co/MiniMaxAI/MiniMax-M3。

1 引言

大语言模型正迅速从短篇幅、单轮交互转向长周期、智能体驱动的工作流程，这些流程包含数百个交错的推理与行动步骤——编写和部署生产代码、浏览开放网络、编排多种工具以及生成结构化文档 [OpenAI, 2025, Anthropic, 2025, Google DeepMind, 2025, DeepSeek-AI, 2026, Moonshot AI, 2026, Zhipu AI, 2026]。然而，这些任务所需的超长上下文给训练和推理带来了严重的计算和内存瓶颈，其中二次复杂度的 softmax 注意力机制是主要症结所在，而生产级部署的延迟和吞吐量限制则进一步加剧了这一问题。

上下文长度是大语言模型的一个关键扩展维度，在模型质量与效率之间进行权衡仍然是一项艰巨的挑战。业界正在积极推动这一领域的前沿发展。混合架构 [MiniMax, 2025b, Qwen, 2026] 将部分 softmax 注意力层替换为更高效的替代方案，例如线性注意力 [Team et al., 2025a, Yang et al., 2025, Gu and Dao, 2023] 或滑动窗口注意力 [OpenAI et al., 2025, MiMo et al., 2026]。此外，另一条研究路线试图对 softmax 注意力本身进行稀疏化处理 [DeepSeek-AI et al., 2025, DeepSeek-AI, 2026, Team et al., 2025b, Lu et al., 2025]，以打破计算瓶颈。

我们提出 MiniMax 稀疏注意力（MSA），其设计遵循奥卡姆剃刀原则：经过大量消融实验后，我们仅保留最核心的组件。MSA 采用稀疏 softmax 注意力范式，以最大化复用现有软硬件基础设施。我们采用分块 token 选择机制，配合更小的 top- 参数，在放宽先前设计所施加的头维度约束的同时，能够在更广泛的 GPU 架构上实现高效执行。具体而言，一个超轻量级的索引分支通过最大池化评分，为每个注意力组选择 top- 块，同时始终保留最近的一个块以确保训练稳定性。

要将 MSA 的理论稀疏性转化为实际的端到端加速，需要将算法与其 GPU 执行路径进行协同设计。为此，我们设计了一个专为小 k 场景优化的免指数运算 TopK 内核，利用分块索引器在选择前绕过不必要的 softmax 计算。对于主注意力分支，我们以 KV 外序组织稀疏注意力：选中的 KV 块收集其关联的查询并拼接起来，以填充张量核心矩阵乘法单元，并采用预调度分块与两阶段合并策略来处理高度偏斜的块流行度，无需原子更新。在训练方面，我们进一步将稀疏 KL 损失所需的辅助 LSE 计算融合到前向传播中，并在反向传播中采用持久负载均衡。

为验证 MSA 是否同时保留了文本和多模态能力，我们将其与分组查询注意力进行对比，在一个 109B 参数的混合专家（MoE）模型上从头训练，预算为 3T token。MSA 在下游基准测试中与 GQA 表现相当，同时在 1M 上下文长度下实现了预填充和解码加速。

主要贡献。

• 我们提出 MSA，一种最小化、可扩展且加速的分块稀疏注意力机制，支持从头训练以及从预训练 GQA 检查点进行近乎无损的转换。

• 我们协同设计了高效的训练和推理内核，将 MSA 的理论计算节省转化为大规模实际运行时的加速。

• 我们进行了大规模的消融实验，将模型扩展至具有原生多模态训练的 109B 参数 MoE 模型，深入剖析了 MSA 在不同规模和模态下的行为表现。

2 预备知识

2.1 因果注意力与 GQA

我们用 表示序列长度， 表示隐藏维度， 表示注意力头维度。对于每个查询位置 和注意力头 ，因果 Softmax 注意力计算如下：

(1)

公式 (1) 的计算成本为 次 FLOPs，随序列长度 呈二次方增长。分组查询注意力 [Ainslie et al., 2023] 使用 个查询头，并将键值头数量减少至 ，将相邻的查询头绑定到单个共享的键值头上。因此，每个键值头定义一个 GQA 分组。

2.2 稀疏注意力作为两阶段过程

稀疏注意力层将因果注意力分解为两个阶段：一个索引器，用于选择要关注的键；以及一个在所选键上进行的稀疏注意力计算。对于每个查询位置 ，

(2)

其中 由 参数化（对于固定规则索引器为空；对于可训练索引器则为学习得到）， 表示选定的索引集合， 表示限制在此索引集合上的标准缩放点积 softmax 注意力。我们将第一阶段称为索引分支，第二阶段称为主分支。在多注意力头机制中，每个由位置 和注意力头 指定的查询，可以选择不同的键值索引集合，记作 ；公式 (2) 仅为符号简洁而省略了注意力头索引。

2.3 基于 GQA 的块稀疏注意力

逐注意力头的 token 级选择提供了最细的粒度，但这种细粒度的计算难以高效映射到 GPU 矩阵运算上。为了提高效率，基于 GQA 构建的稀疏注意力可以在每个 GQA 分组内共享索引结果。设 表示由第 个键值头服务的查询头集合。分组共享的索引集合可写为：

(3)

选择键值块而非单个 token 可以减少路由开销，并使稀疏注意力更加规整。对于块大小 ，定义：

(4)

对于查询位置和 GQA 组，该集合表示所选块的索引集。然后，对于该组中的任意查询头，稀疏注意力输出是在所选块中因果可见的 token 上，使用同一组的键值头计算得出的。MSA 遵循这种基于 GQA 的块稀疏公式，具体的索引器架构和训练目标将在下一节中描述。

3 MSA

我们引入了 MiniMax 稀疏注意力（MSA），这是一种基于 GQA 的稀疏注意力机制，包含两个分支，如图 1 所示。对于每个查询 token，一个轻量级的索引分支从因果上下文中选择一小部分键块，主分支则在这些块中的 token 上计算 softmax 注意力。索引分支仅向标准 GQA 添加了两个投影矩阵，以块粒度运行，并为每个 GQA 组独立进行选择。我们在第 3.1 节中描述其架构，在第 3.2 节中描述训练过程。

3.1 架构

MSA 在第 2.2 节中描述的两阶段稀疏注意力公式基础上，以 GQA 组和块粒度进行实例化（图 1）。对于每个查询 token，索引分支为每个 GQA 组选择大小为 的键块，主分支仅关注所选块中的 token，其预算最多为 。设 为输入隐藏状态。根据第 2.1 节，我们分别用 和 表示查询头和键值头的数量，因此每个键值头服务于 个查询头。

索引分支。

索引分支为每个 GQA 组引入一个索引查询头，以及一个跨组共享的单一索引键头：

(5)

对于查询 token 和组 ，索引分支首先对可见的键 token 进行评分，然后将这些分数聚合到块级别。使用第 2.3 节中定义的块划分，

(6)

这里 索引 GQA 组， 强制因果性，并且对于没有可见 token 的块，其分数被赋值为 。然后，索引分支选择前 个块索引：

(7)

这里 返回在 条件下最大的 个块的索引。我们始终包含包含位置 的本地块，并且 由组 中的所有查询头共享。

主分支。

给定索引分支选择的块索引集合，主分支仅关注所选块中因果可见的 token。对于任意查询头，它使用与 GQA 组关联的键值头，对这部分 token 应用标准缩放点积注意力：

(8)

其中 表示位置 和查询头 处的查询向量，而 和 分别表示第 个 GQA 组的键矩阵和值矩阵。符号 和 表示从所选块中收集因果可见的 token。块索引集合由 中的所有查询头共享，而每个头保留自己的查询投影。由于所选块最多包含 个因果可见的 token，每个查询的注意力成本从 降低到 ，该成本随序列长度增加而保持固定。

3.2 训练

公式 7 中的 top- 选择是不可微的，因此语言建模损失无法直接训练索引投影。我们通过 KL 对齐损失来训练索引分支，并使用三种机制来稳定稀疏训练：梯度分离、索引器预热和强制局部块。下面我们描述每个组件。

KL 损失。

KL 损失通过将索引分支的分数与主分支在所选 token 上的分数进行匹配，为索引分支提供直接的学习信号。将 记为所选块索引所诱导的因果可见 token，对于每个查询位置 和 GQA 组 ，我们在此 token 索引集合上定义索引分支分布 和主分支教师分布：

(9)

其中 是 token 级别的索引分数， 是查询头 的主分支分数。教师分布对每个查询头的主分支分布在概率层面进行平均。然后训练索引器以匹配 ，对所有查询位置和 GQA 组取平均：

(10)

其中 是序列长度，教师分布与梯度计算分离。该辅助损失将索引分布与主分支注意力模式对齐，使后续的块选择具有语义意义。

梯度分离。

为了将辅助目标与主干网络隔离，我们对索引分支的输入应用停止梯度：

(11)

公式 9 中的教师模型是分离的，因此不会影响主分支的投影；公式 11 进一步阻止它通过反向传播影响主干网络。在此规则下，只有索引分支的投影被更新，使得 KL 散度成为索引器的一个干净的对齐信号。

索引器预热。

我们采用两阶段训练计划来初始化索引分支，并避免早期的随机选择。在最初的几次迭代中，模型在两个分支中都运行全注意力，并使用索引损失训练新添加的索引投影。预热后，模型切换到稀疏注意力，并且注意力计算仅针对前 K 个选中的位置进行。在将预训练的全注意力检查点稀疏化时，也使用相同的计划，这有助于在索引投影控制主分支路由之前，先对齐新添加的索引投影。

局部块。

对于每个查询位置和 GQA 组，包含该查询位置的局部块在训练和推理期间始终被选为选中块的一部分。这种固定分配保留了一个块槽位，并将剩余的槽位留给索引分支选择，从而防止出现省略查询位置直接邻域的退化选择。

完整的层级训练过程总结在算法 1 中。

媒体内容 · 前往原文查看

0: 隐藏状态；块大小，选中块的数量。

1: // 索引投影

2: // 主分支投影；分离

3: // 注意力分数；按组，因果

4: // 选中的块索引；包含局部块

5: // 稀疏注意力；关注选中的块

6: // 输出

7: // 由索引投影引起的 token 上的 KL 散度

8: 返回

3.3 计算复杂度

在相同的隐藏维度、注意力头数、键值头数和序列长度下，GQA 和 MSA 的因果注意力 FLOPs 分别为：

(12)

GQA 的主注意力路径随完整上下文长度扩展，而 MSA 使用固定的选择预算加上轻量级的索引计算；因此，当块大小和选中块数量固定时，FLOPs 差距随序列长度增长而扩大。

4 内核设计

本节描述了我们稀疏预填充实现中使用的 GPU 内核，包括索引 TopK 内核、KV-外部稀疏注意力前向传播和稀疏 KL 损失反向传播。

4.1 索引与 TopK

免指数选择。

为了高效地选取 top-KV 块，索引模块直接对索引分数进行排序。由于 softmax 保持顺序不变，分数的相对排序得以保留（），因此 top- 索引保持不变。因此，前向传播跳过了 softmax 的 max/exp/sum 步骤，直接将原始分数传递给选择模块。

每线程寄存器 top-。

块大小和选择大小是与 top- 内核协同设计的：较大的 会增加注意力算术强度（第 4.2 节），而在此处保持较小的 可使每行候选块数量和 都低于通用 top- 内核的最佳点，后者通过多路分桶（基数选择）或按 规模（双调排序）来分摊开销。我们采用 ，。每个 warp 的 32 条通道以 1/32 的步幅流式处理输入行，并在共享内存中维护一个包含 个元素的最小堆。堆根缓存在寄存器中，插入操作采用延迟写入。最后，通过 轮 shuffle 合并将 32 个局部 TopK 结果合并。共享内存布局将每条通道映射到固定的存储体，从而避免冲突。

基准测试。

我们在 H800 GPU 上，使用 fp32 输入和未排序输出，将我们的方案与 torch.topk 以及 TileLang [Wang et al., 2025] 的基数选择 top- 进行了比较；延迟取预热后迭代的中位数。表 1 显示，我们的专用内核在所有测试设置中速度最快，在部署设置 下提升最为显著。

媒体内容 · 前往原文查看

序列长度	块数		torch	TileLang	我们的方案	对比 torch	对比 TileLang
K
K
K
K

4.2 稀疏注意力

我们重新审视了在查询和键/值长度相等的稀疏预填充场景下迭代顺序的选择。设 、、、、、、 分别表示查询头数、键值头数、GQA 比率、头维度、序列长度、KV 块大小以及每查询选取的块数。为简化起见，以下 IO 估算假设为 2 字节元素（bfloat16 大小的流量）。我们的内核也支持 fp8；使用 fp8 会按比例缩放绝对 IO 量，但不会改变 Q-outer 与 KV-outer 迭代之间的比较结果。

在外层循环中迭代查询，得到

				(13)
				(14)

因此 。

在外层循环中迭代 KV 块，并收集选择了每个块的查询，需要一个中间输出缓冲区：

				(15)
				(16)

因此 。

由于在实践中，我们选择 KV 外层迭代配合 Q 收集，以最大化算术强度。该内核以持久化网格方式在瓦片（tile）上执行。对于每个瓦片，从 TopK 选择中得到的反向稀疏索引会标识出相关的查询位置。这些查询通过 TMA 拷贝加载到共享内存中，每个查询 token 对应一次拷贝，由 warp 的 32 条通道并行分发。

预调度瓦片分块。

直接的每 CTA 对应一个瓦片的映射会被汇聚行（sink rows）主导——即几乎每个查询都选择的单个早期 KV 块——而同样的热点模式也可能出现在任何热门 KV 块上。因此，一个 GPU 调度器内核会沿着查询维度将每个 KV 瓦片分割成最多包含 个查询的块，将热门瓦片分散到多个共享相同负载的 CTA 上。由于每个查询的部分结果现在由多个 CTA 产生，调度器还会为每个（查询，块）对预先分配一个 中的槽位——与查询索引一起打包成一个 位句柄——这样注意力内核就可以将部分结果写入预先分配的偏移量，无需使用原子操作。合并内核会读取每个查询的槽位计数，以了解需要合并多少个部分结果。

两阶段前向传播。

KV 外层分割不允许内联的 softmax 归一化，因为每个查询的部分结果由不同的 CTA 产生。因此，前向传播被拆分为两个内核，中间通过 HBM 缓冲区（局部归一化的部分输出）和（每个部分结果的 logsumexp）连接。注意力内核执行上述工作列表，并将每个部分结果写入其预先分配的槽位。合并内核读取每个查询的有效槽位，计算 和 ，然后形成归一化的分割 K 权重 。它输出 以及最终的 。这两个内核使用程序化依赖启动（Programmatic Dependent Launch）来隐藏内核之间的启动延迟。

查询拼接。

KV 外层迭代通常将每个 KV 块仅与几个到几十个查询位置相关联。逐个处理这些位置会导致评分 MMA（矩阵乘法累加）填充不足：单个查询位置仅贡献查询头，使得 MMA 维度仅为 16。在 Q 外层迭代下，查询位置无法沿序列维度拼接，因为它们通常选择不同的 KV 子集。然而，在 KV 外层迭代下，为给定块收集的所有位置共享相同的 KV 操作数。因此，内核将查询位置及其关联的查询头（均位于同一 KV 头下）打包到评分 MMA 中。

4.3 稀疏 KL 散度损失

LSE 融合。

在初始实现中，我们使用专用内核来计算 KL 散度的前向传播，存储相关值以支持反向传播。然而，由于 KL 损失仅影响反向梯度，我们对此进行了优化：在主计算过程中将这些 LSE 值直接发射到全局内存中，从而完全跳过 KL 损失的前向传播。此外，在索引分支计算期间，我们保存每个块的 LSE，并对 top- 个块进行归约以获得最终结果。反向内核随后将这些标量直接加载到 softmax 中，消除了冗余的前向计算。

动态负载均衡。

在变长序列和数据依赖的稀疏性下，每个块的工作量可能相差数个数量级。该内核以持久化网格形式运行，其中 CTA（协作线程阵列）通过全局原子计数器来认领工作；每个块沿其收集的查询维度被划分为子块，子块的数量随每个块的查询数量变化，同时受最小子块粒度的约束，以分摊每个子块的开销。

5 实验

本节报告了两次 109B 规模的实验，用于验证最终 MSA 设计在基于文本与图像/视频数据混合训练的原生多模态模型上的效果。第一次实验从头训练了一个原生 MSA 模型，我们将其记为 MSA-PT。第二次实验从一个全注意力检查点出发，在将密集注意力替换为 MSA 后继续进行预训练，我们将其记为 MSA-CPT。两个模型采用与全注意力基线相同的架构族，但将密集注意力层替换为 MSA 层。

5.1 实验设置

模型结构

所有模型均采用相同的 41 层 MoE 骨干网络，总参数量约 109B，每个 token 激活 6B 参数。前 3 层为密集层，其余 38 层为 MoE 层。模型使用 200K token 词表，隐藏层维度为 。每个注意力模块采用 MSA，包含 64 个查询头、4 个 KV 头，头维度为 128，RoPE 维度为 64。每个 MoE 层使用 128 个路由专家、1 个共享专家，并采用 top-4 路由专家选择机制。在稀疏训练与评估阶段，两个 MSA 模型均使用块大小 ，每个查询和 GQA 组保留的键值块数为 。

训练预算

所有模型均在总计 3T token 的预算下进行训练。MSA-PT 从头开始训练：经过 40B token 的索引器预热后，在预训练剩余阶段保持稀疏训练。MSA-CPT 从基于 2.6T token 训练得到的 GQA 全注意力检查点出发，随后将密集注意力替换为 MSA 并继续训练 400B token：其中前 40B token 用于索引器预热，之后进行稀疏持续预训练。

评估

我们在相同的预训练评估套件上，使用相同训练预算下的匹配检查点，对 Full、MSA-PT 和 MSA-CPT 进行了评估。对于通用推理和问答任务，我们使用了 MMLU [Hendrycks et al., 2021]、MMLU-Pro [Wang et al., 2024a]、BBH [Suzgun et al., 2022]、GPQA Hard [Rein et al., 2023]、ARC Challenge [Clark et al., 2018]、TriviaQA [Joshi et al., 2017] 和 WinoGrande [Sakaguchi et al., 2020]。对于数学和代码任务，我们使用了 GSM8K [Cobbe et al., 2021]、MGSM [Shi et al., 2022]、MathVista [Lu et al., 2024]、OlymMATH [Sun et al., 2025]、HumanEval [Chen et al., 2021]、EvalPlus [Liu et al., 2023]、BigCodeBench [Zhuo et al., 2025] 和 MultiPL-E MBPP [Cassano et al., 2023]。我们还评估了多模态能力：图像基准包括 AI2D [Kembhavi et al., 2016]、ChartQA [Masry et al., 2022]、MMMU [Yue et al., 2024]、OCRBench v2 [Fu et al., 2025]、CharXiv [Wang et al., 2024b]、VisualWebBench [Liu et al., 2024] 和 CVBench [Tong et al., 2024]，而视频基准包括 EgoSchema [Mangalam et al., 2023]、LongVideoBench [Wu et al., 2024]、MLVU [Zhou et al., 2025]、MMVU [Zhao et al., 2025b]、VideoMME [Fu et al., 2024] 和 TemporalBench [Cai et al., 2024]。对于长上下文评估，我们使用了 RULER [Hsieh et al., 2024] 和 HELMET [Yen et al., 2025]。此外，我们还报告了在下游智能体任务上的困惑度，这些任务包括 -bench [Barres et al., 2025]、TheAgentCompany [Xu et al., 2024]、Humanity’s Last Exam [Phan et al., 2025] 和 SWE-bench [Jimenez et al., 2024]。

5.2 训练动态

图 2 比较了原生稀疏预训练与匹配的全注意力运行。在 3T token 的训练过程中，两条语言模型损失曲线几乎无法区分，这表明 MSA 相对于全注意力并未引入明显的优化退化。梯度范数曲线在整个训练过程中也保持在相同范围内，表明 MSA 不会导致异常的梯度波动或训练不稳定。这些结果表明，在大规模训练中，训练稀疏注意力模型与训练全注意力基线模型一样稳定。

图 3 展示了从训练好的全注意力检查点向稀疏持续预训练的过渡过程。索引器预热阶段在启用稀疏注意力之前迅速降低了 KL 散度损失。切换到稀疏持续预训练后，KL 散度损失保持较低水平。对于每个查询和分组查询注意力头，设 为由主分支分数诱导出的对应 Top- 块集合， 为索引分支的选择结果。块召回率为 ，而分数召回率为 ，其中 是主分支注意力概率在块 内各 token 上的总和。块召回率保持良好，表明重要块能被可靠地恢复。更高的分数召回率进一步表明，检索到的块占据了主分支注意力权重的绝大部分。这些动态过程共同表明，预热提供了一个干净的转换阶段，并且在稀疏持续预训练期间，持续预训练索引器保持了良好的对齐。

5.3 主要结果

表 2 在代表性预训练评估集上比较了 Full、MSA-PT 和 MSA-CPT。两种稀疏模型在整体上仍与全注意力基线保持竞争力，表明用 MSA 替换密集注意力并不会显著降低模型在通用语言、推理、多模态或面向智能体的困惑度表现。两种训练路线展现出不同的优势。MSA-PT 在整个预训练过程中学习稀疏模式，在多项数学、图像、视频和长上下文检索基准上取得了最强结果，这表明原生稀疏预训练能够使模型表征适应稀疏注意力模式。MSA-CPT 则更为保守：它保留了全注意力检查点的大部分行为，在大多数文本、代码和困惑度评估上保持接近，因此当已有一个训练好的密集检查点时，它是一种实用的转换路线。其余差距与具体基准相关，而非集中在单一能力领域。

媒体内容 · 前往原文查看

分组基准 完整 MSA-PT MSA-CPT 通用 MMLU 67.0 67.2 66.8 MMLU-Pro 38.5 38.8 39.1 BBH 67.7 66.6 66.1 GPQA Hard 25.9 26.3 26.3 ARC Challenge 82.7 82.5 82.9 TriviaQA 66.0 65.5 67.7 WinoGrande 58.3 60.9 62.0 数学 GSM8K 76.2 77.7 73.7 MGSM 44.1 46.0 44.2 MathVista 43.8 46.8 44.5 OlymMATH Easy P@100 23.0 26.0 22.0 代码 HumanEval 61.0 64.0 57.9 EvalPlus 59.4 61.8 60.0 BigCodeBench 44.8 44.0 45.7 MultiPL-E MBPP P@10 82.1 81.6 81.1 检索 RULER-8K 79.8 84.2 77.2 RULER-32K 75.0 77.5 75.7 图像 AI2D 68.3 70.6 67.3 ChartQA 75.0 75.4 71.4 MMMU 46.8 45.9 44.5 OCRBench v2 55.0 55.7 54.3 CharXiv 37.55 41.55 37.15 VisualWebBench 55.6 68.4 59.4 CVBench 57.0 59.7 58.8 EgoSchema 29.6 37.6 25.8 LongVideoBench 38.5 41.8 38.9 MLVU 44.14 46.94 43.68 MMVU 45.8 47.5 45.8 VideoMME 41.11 45.48 39.65 TemporalBench 49.4 53.4 50.6 PPL TAU2 1.155 1.148 1.150 AgentCompany 1.248 1.249 1.247 HLE 1.275 1.278 1.275 SWE 1.216 1.218 1.216

为了评估 MSA 在长上下文扩展后是否仍然有效，我们在 MSA-CPT 模型上进行了额外的扩展实验。从稀疏持续预训练检查点开始，我们运行了大约 140B 个模型 token 的长上下文训练，然后在 HELMET 和 RULER 上进行评估。结果报告在表 3 中。在扩展阶段之后，MSA-CPT 仍然接近全注意力基线。由于每个查询和 GQA 组仍然只关注键值模型 token，这些结果表明 MSA 可以在高度紧张的注意力预算下保持长上下文能力。

附录中提供了支持这些设计选择的额外消融实验。具体来说，B 节研究了索引分支的训练方案，包括梯度来源、KL 梯度分离、预热以及与滑动窗口稀疏基线的比较。C 节进一步考察了架构选择，例如块大小、强制汇聚、局部选择以及索引分支值头。这些消融实验为主要实验中使用的最终 MSA 设计提供了经验基础。

媒体内容 · 前往原文查看

基准测试子集 全注意力 MSA-CPT HELMET-128K 总体 46.53 45.93 -0.60 ICL 70.40 72.80 +2.40 Rerank/RAG 34.60 32.50 -2.10 RULER-128K 总体 72.00 72.12 +0.12 CWE/FWE 46.35 45.00 -1.35 MK/MQ/MV 96.63 98.87 +2.24 QA1/QA2 47.80 46.80 -1.00 VT 97.80 96.80 -1.00

5.4 效率

我们在实验模型配置上实例化了第 3.3 节的复杂度分析，并报告了理论上的注意力 FLOPs 减少量和实测的运行时间加速比。密集 GQA 和 MSA 使用相同的查询头数量、键值头数量、头维度和上下文长度；唯一的区别在于，密集 GQA 关注整个上下文，而 MSA 则先进行索引选择，然后在固定的 KV 预算上执行稀疏注意力。在我们的设置中，MSA 使用 和 ，对应于每个查询的选定 token 预算为 。

如图 4 所示，在我们的设置中，MSA 相对于 GQA 显著减少了每个 token 的注意力 FLOPs，且减少幅度在更长的上下文中随之增加。在 个 token 时，相同头配置下的 FLOPs 减少量达到 。实测的运行时间加速比遵循相同的缩放趋势，但预计不会与 FLOPs 减少量完全匹配。稀疏注意力引入了索引构建、top- 选择、反向索引物化、查询收集和负载均衡开销，其内存访问模式也不如密集注意力规则。因此，运行时间加速比小于理论 FLOPs 减少量，但它会随着上下文长度的增加而增大，因为密集基线会随完整序列长度扩展，而 MSA 则保持主要的注意力预算固定。

6 相关工作

长上下文效率催生了大量关于高效注意力的研究工作，这些工作大致可分为两个方向：用更廉价的线性或循环替代方案替换密集的softmax注意力，以及保留softmax注意力但限制其感受野。线性注意力[Katharopoulos等人，2020，Choromanski等人，2021]用线性复杂度的替代函数替换softmax核，而Mamba[Gu和Dao，2023]等状态空间模型则用基于隐藏状态的选择性循环替换注意力。混合堆叠架构[MiniMax，2025a，b]将线性模块与全注意力模块交错排列，在减少二次复杂度层数的同时保留部分精确softmax能力。固定模式注意力保留softmax注意力，但施加预定义的支持范围，包括局部窗口、全局token[Beltagy等人，2020，Zaheer等人，2020]，以及带滑动窗口的注意力汇聚点[Xiao等人，2024b]。这些方法通过部分或全部替换密集注意力，或使用与内容无关的注意力模式，来降低长上下文成本。

超越固定的稀疏模式，自适应稀疏注意力使得被关注的支撑集依赖于输入。现有方法的主要区别在于该支撑集何时构建，以及选择器是否作为模型的一部分进行训练。推理时稀疏化作用于预训练的全注意力骨干网络，仅在服务期间构建稀疏支撑集。H2O [Zhang et al., 2023] 和 SnapKV [Li et al., 2024] 在解码过程中利用累积的注意力统计信息对 KV 缓存进行剪枝，Quest [Tang et al., 2024] 对每个查询执行页面级重要性估计，MInference [Jiang et al., 2024] 和 FlexPrefill [Lai et al., 2025] 在预填充阶段调度每个注意力头的稀疏内核，而 InfLLM [Xiao et al., 2024a] 则维护注意力汇聚点、一个局部上下文窗口以及可检索的块。这些方法继承了全注意力的训练成本，并且至少有一个推理阶段的速度接近全注意力。原生训练的稀疏注意力设计在预训练期间训练索引器，构成了与 MSA 最接近的先前工作。NSA [Yuan et al., 2025] 针对 MQA/MHA 骨干网络，采用三个并行分支：用于粗略全局上下文的压缩注意力、对细粒度块进行的选择注意力，以及用于局部上下文的滑动窗口。InfLLM-V2 [Zhao et al., 2025a] 通过将无参数块选择与局部滑动窗口统一起来，实现了零样本的密集到稀疏切换。MoBA [Lu et al., 2025] 同样作用于 GQA，但使用非常大的 KV 块，这些块通过块平均键进行评分，并且仅通过语言建模梯度来训练其索引器。DSA [DeepSeek-AI et al., 2025] 在其 MQA 模式下基于 MLA 构建：一个基于多头 ReLU 的闪电索引器对每个 token 单独评分，所有查询头共享一个单一的 Top- 索引，并且选择是在 token 级别进行的。MSA 与这些相关工作在两个维度上有所不同，并且这两个维度是同时采用的：每个 GQA 组共享 Top- 与块级选择相结合，这实现了多组块粒度检索，同时保持了 KV 读取的连续性。

高效核对于稀疏注意力而言至关重要，它能将理论上的 FLOP 减少转化为实际运行时间的加速。FlashAttention [Dao 等人，2022] 和 FlashAttention-2 [Dao，2024] 引入了 IO 感知的分块 softmax 注意力机制，而 FlashDecoding [Dao 等人，2023] 将其扩展到了内存受限的解码场景。诸如 Flash-Sparse-Attention [Yan 等人，2025] 和 FlashMoBA [Xiao 等人，2025] 等开源块稀疏核，已使这种递归的块稀疏变体得以实现。MSA 的核（在第 4 节中描述）复用了 FlashAttention 的算法框架，并采用了一种循环排序，该排序针对 MSA 产生的、原生支持 GQA 且以块为粒度的访问模式进行了调优。

7 结论

我们提出了 MSA，这是一种与分组查询注意力协同设计的稀疏注意力机制。该架构在标准 GQA 层上附加了一个轻量级的索引分支：每个 GQA 组通过一个块级点积索引器独立选择一小部分键值块，而主分支则对所选块执行受限的 softmax 注意力。索引分支是一个纯粹的选择器，通过一个两阶段预热计划下的 KL 对齐损失与主分支进行训练，并对索引输入应用停止梯度，从而将辅助损失限制在索引投影上。在 109B-MoE 规模下，MSA 在大多数预训练和智能体基准测试中保持了 GQA 全注意力基线的能力，同时将每个 token 的注意力计算量减少了 8 倍，这是在长上下文推理成为关键部署约束的上下文窗口长度下的表现。

展望。MSA 的核心决策——按 GQA 分组独立选择、块级粒度以及使用 KL 对齐目标训练的索引器——与当前大多数开源前沿模型共享的 GQA 主干架构相兼容，因此该方案应能通过少量修改即可迁移。两个方向是自然的下一步：一是弥合剩余的长上下文检索差距，可通过更长的稀疏训练、推理时更大的选择预算或更丰富的索引器评分函数来实现；二是将相同的仅选择器设计扩展到预训练之外的场景，包括强化学习后训练和智能体部署，在这些场景中，长上下文成本是主要的运营约束。

参考文献

• J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, 和 S. Sanghai (2023) GQA：从多头检查点训练广义多查询 Transformer 模型。载于《2023 年自然语言处理经验方法会议论文集》(EMNLP)，引用自：§2.1。
• Anthropic (2025) Claude Opus 4.6 和 Sonnet 4.6 模型卡。注：https://www.anthropic.com/news/claude-4-6 引用自：§1。
• V. Barres, H. Dong, S. Ray, X. Si, 和 K. Narasimhan (2025) -bench：在双控制环境中评估对话智能体。arXiv 预印本 arXiv:2506.07982。引用自：§5.1。
• I. Beltagy, M. E. Peters, 和 A. Cohan (2020) Longformer：长文档 Transformer。arXiv 预印本 arXiv:2004.05150。引用自：§6。
• M. Cai, R. Tan, J. Zhang, B. Zou, K. Zhang, F. Yao, F. Zhu, J. Gu, Y. Zhong, Y. Shang, Y. Dou, J. Park, J. Gao, Y. J. Lee, 和 J. Yang (2024) TemporalBench：用于多模态视频模型的细粒度时间理解基准测试。外部链接：2410.10818，链接 引用自：§5.1。
• F. Cassano, J. Gouwar, D. Nguyen, S. Nguyen, L. Phipps-Costin, D. Pinckney, M. Yee, Y. Zi, C. J. Anderson, M. Q. Feldman, A. Guha, M. Greenberg, 和 A. Jangda (2023) MultiPL-E：一种可扩展的多语言神经代码生成基准测试方法。《IEEE 软件工程汇刊》第 49 卷第 7 期，第 3675–3691 页。引用自：§5.1。
• M. Chen、J. Tworek、H. Jun、Q. Yuan、H. P. de Oliveira Pinto、J. Kaplan、H. Edwards、Y. Burda、N. Joseph、G. Brockman、A. Ray、R. Puri、G. Krueger、M. Petrov、H. Khlaaf、G. Sastry、P. Mishkin、B. Chan、S. Gray、N. Ryder、M. Pavlov、A. Power、L. Kaiser、M. Bavarian、C. Winter、P. Tillet、F. P. Such、D. Cummings、M. Plappert、F. Chantzis、E. Barnes、A. Herbert-Voss、W. H. Guss、A. Nichol、A. Paino、N. Tezak、J. Tang、I. Babuschkin、S. Balaji、S. Jain、W. Saunders、C. Hesse、A. N. Carr、J. Leike、J. Achiam、V. Misra、E. Morikawa、A. Radford、M. Knight、M. Brundage、M. Murati、K. Mayer、P. Welinder、B. McGrew、D. Amodei、S. McCandlish、I. Sutskever 和 W. Zaremba（2021）《评估基于代码训练的大语言模型》。外部链接：2107.03374，链接，被 §5.1 引用。
• K. Choromanski、V. Likhosherstov、D. Dohan、X. Song、A. Gane、T. Sarlós、P. Hawkins、J. Davis、A. Mohiuddin、L. Kaiser、D. Belanger、L. Colwell 和 A. Weller（2021）《用 Performers 重新思考注意力机制》。发表于国际学习表征会议（ICLR），被 §6 引用。
• P. Clark、I. Cowhey、O. Etzioni、T. Khot、A. Sabharwal、C. Schoenick 和 O. Tafjord（2018）《你以为已经解决了问答问题？试试 ARC，AI2 推理挑战》。外部链接：1803.05457，链接，被 §5.1 引用。
• K. Cobbe、V. Kosaraju、M. Bavarian、M. Chen、H. Jun、L. Kaiser、M. Plappert、J. Tworek、J. Hilton、R. Nakano、C. Hesse 和 J. Schulman（2021）《训练验证器解决数学应用题》。外部链接：2110.14168，链接，被 §5.1 引用。
• T. Dao、D. Y. Fu、S. Ermon、A. Rudra 和 C. Ré（2022）《FlashAttention：具有 IO 感知能力的快速且内存高效精确注意力机制》。发表于神经信息处理系统进展，被 §6 引用。
• T. Dao、D. Haziza、F. Massa 和 G. Sizov（2023）《用于长上下文推理的 Flash-Decoding》。注释：https://crfm.stanford.edu/2023/10/12/flashdecoding.html，被 §6 引用。
• T. Dao（2024）《FlashAttention-2：通过更好的并行性和工作分区实现更快的注意力机制》。发表于国际学习表征会议（ICLR），被 §6 引用。
• DeepSeek-AI, A. Liu, A. Mei, B. Lin, B. Xue, B. Wang, B. Xu, B. Wu, B. Zhang, C. Lin, C. Dong, C. Lu, C. Zhao, C. Deng, C. Xu, C. Ruan, D. Dai, D. Guo, D. Yang, D. Chen, E. Li, F. Zhou, F. Lin, F. Dai, G. Hao, G. Chen, G. Li, H. Zhang, H. Xu, H. Li, H. Liang, H. Wei, H. Zhang, H. Luo, H. Ji, H. Ding, H. Tang, H. Cao, H. Gao, H. Qu, H. Zeng, J. Huang, J. Li, J. Xu, J. Hu, J. Chen, J. Xiang, J. Yuan, J. Cheng, J. Zhu, J. Ran, J. Jiang, J. Qiu, J. Li, J. Song, K. Dong, K. Gao, K. Guan, K. Huang, K. Zhou, K. Huang, K. Yu, L. Wang, L. Zhang, L. Wang, L. Zhao, L. Yin, L. Guo, L. Luo, L. Ma, L. Wang, L. Zhang, M. S. Di, M. Y. Xu, M. Zhang, M. Zhang, M. Tang, M. Zhou, P. Huang, P. Cong, P. Wang, Q. Wang, Q. Zhu, Q. Li, Q. Chen, Q. Du, R. Xu, R. Ge, R. Zhang, R. Pan, R. Wang, R. Yin, R. Xu, R. Shen, R. Zhang, S. H. Liu, S. Lu, S. Zhou, S. Chen, S. Cai, S. Chen, S. Hu, S. Liu, S. Hu, S. Ma, S. Wang, S. Yu, S. Zhou, S. Pan, S. Zhou, T. Ni, T. Yun, T. Pei, T. Ye, T. Yue, W. Zeng, W. Liu, W. Liang, W. Pang, W. Luo, W. Gao, W. Zhang, X. Gao, X. Wang, X. Bi, X. Liu, X. Wang, X. Chen, X. Zhang, X. Nie, X. Cheng, X. Liu, X. Xie, X. Liu, X. Yu, X. Li, X. Yang, X. Li, X. Chen, X. Su, X. Pan, X. Lin, X. Fu, Y. Q. Wang, Y. Zhang, Y. Xu, Y. Ma, Y. Li, Y. Li, Y. Zhao, Y. Sun, Y. Wang, Y. Qian, Y. Yu, Y. Zhang, Y. Ding, Y. Shi, Y. Xiong, Y. He, Y. Zhou, Y. Zhong, Y. Piao, Y. Wang, Y. Chen, Y. Tan, Y. Wei, Y. Ma, Y. Liu, Y. Yang, Y. Guo, Y. Wu, Y. Wu, Y. Cheng, Y. Ou, Y. Xu, Y. Wang, Y. Gong, Y. Wu, Y. Zou, Y. Li, Y. Xiong, Y. Luo, Y. You, Y. Liu, Y. Zhou, Z. F. Wu, Z. Z. Ren, Z. Zhao, Z. Ren, Z. Sha, Z. Fu, Z. Xu, Z. Xie, Z. Zhang, Z. Hao, Z. Gou, Z. Ma, Z. Yan, Z. Shao, Z. Huang, Z. Wu, Z. Li, Z. Zhang, Z. Xu, Z. Wang, Z. Gu, Z. Zhu, Z. Li, Z. Zhang, Z. Xie, Z. Gao, Z. Pan, Z. Yao, B. Feng, H. Li, J. L. Cai, J. Ni, L. Xu, M. Li, N. Tian, R. J. Chen, R. L. Jin, S. S. Li, S. Zhou, T. Sun, X. Q. Li, X. Jin, X. Shen, X. Chen, X. Song, X. Zhou, Y. X. Zhu, Y. Huang, Y. Li, Y. Zheng, Y. Zhu, Y. Ma, Z. Huang, Z. Xu, Z. Zhang, D. Ji, J. Liang, J. Guo, J. Chen, L. Xia, M. Wang, M. Li, P. Zhang, R. Chen, S. Sun, S. Wu, S. Ye, T. Wang, W. L. Xiao, W. An, X. Wang, X. Sun, X. Wang, Y. Tang, Y. Zha, Z. Zhang, Z. Ju, Z. Zhang, and Z. Qu (2025) DeepSeek-v3.2: 推动开源大语言模型前沿发展. 外部链接: 2512.02556, 链接 引用自: §1, §6.
• DeepSeek-AI (2026) DeepSeek-V4：迈向高效百万 token 上下文智能。注：技术报告（预览版）。引用自：§1, §1。
• C. Fu, Y. Dai, Y. Luo, L. Li, S. Ren, R. Zhang, Z. Wang, C. Zhou, Y. Shen, M. Zhang, 等 (2024) Video-MME：首个面向多模态大语言模型视频分析的综合评估基准。arXiv 预印本 arXiv:2405.21075。引用自：§5.1。
• L. Fu, Z. Kuang, J. Song, M. Huang, B. Yang, Y. Li, L. Zhu, Q. Luo, X. Wang, H. Lu, Z. Li, G. Tang, B. Shan, C. Lin, Q. Liu, B. Wu, H. Feng, H. Liu, C. Huang, J. Tang, W. Chen, L. Jin, Y. Liu, 和 X. Bai (2025) OCRBench v2：用于评估大型多模态模型在视觉文本定位与推理能力的改进基准。外部链接：2501.00321, 链接 引用自：§5.1。
• Google DeepMind (2025) Gemini 3.1 pro。注：https://deepmind.google/technologies/gemini/ 引用自：§1。
• A. Gu 和 T. Dao (2023) Mamba：基于选择性状态空间的线性时间序列建模。arXiv 预印本 arXiv:2312.00752。引用自：§1, §6。
• D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, 和 J. Steinhardt (2021) 测量大规模多任务语言理解能力。收录于国际学习表征会议，外部链接：链接 引用自：§5.1。
• C. Hsieh, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, 和 B. Ginsburg (2024) RULER：你的长上下文语言模型真正的上下文大小是多少？收录于第一届语言建模会议，外部链接：链接 引用自：§5.1。
• H. Jiang, Y. Li, C. Zhang, Q. Wu, X. Luo, S. Ahn, Z. Han, A. H. Abdi, D. Li, C. Lin, Y. Yang, 和 L. Qiu (2024) MInference 1.0：通过动态稀疏注意力加速长上下文大语言模型的预填充阶段。收录于神经信息处理系统进展，引用自：§6。
• C. E. Jimenez, J. Yang, A. Wettig, S. Yao, K. Pei, O. Press, 和 K. R. Narasimhan (2024) SWE-bench：语言模型能否解决真实的 GitHub 问题？收录于国际学习表征会议 (ICLR)，引用自：§5.1。
• M. Joshi、E. Choi、D. Weld 和 L. Zettlemoyer（2017）《TriviaQA：一个用于阅读理解的大规模远程监督挑战数据集》。载于《第55届计算语言学协会年会论文集（第1卷：长论文）》，R. Barzilay 与 M. Kan 编，加拿大温哥华，第1601–1611页。外部链接：链接，文献编号。引用自：§5.1。
• A. Katharopoulos、A. Vyas、N. Pappas 和 F. Fleuret（2020）《Transformer 即 RNN：具有线性注意力的快速自回归 Transformer》。载于《国际机器学习大会（ICML）》。引用自：§6。
• A. Kembhavi、M. Salvato、E. Kolve、M. Seo、H. Hajishirzi 和 A. Farhadi（2016）《一图胜千言》。载于《欧洲计算机视觉大会（ECCV）》。引用自：§5.1。
• X. Lai、J. Lu、Y. Luo、Y. Ma 和 X. Zhou（2025）《FlexPrefill：一种用于高效长序列推理的上下文感知稀疏注意力机制》。载于《第十三届国际学习表征大会，ICLR 2025，新加坡，2025年4月24-28日》。外部链接：链接。引用自：§6。
• Y. Li、Y. Huang、B. Yang、B. Venkitesh、A. Locatelli、H. Ye、T. Cai、P. Lewis 和 D. Chen（2024）《SnapKV：大语言模型在生成前已知你所寻》。载于《神经信息处理系统进展》。引用自：§6。
• J. Liu、C. S. Xia、Y. Wang 和 L. ZHANG（2023）《你的 ChatGPT 生成的代码真的正确吗？对大型语言模型代码生成能力的严格评估》。载于《神经信息处理系统进展》，A. Oh、T. Naumann、A. Globerson、K. Saenko、M. Hardt 与 S. Levine 编，第36卷，第21558–21572页。外部链接：链接。引用自：§5.1。
• J. Liu、Y. Song、B. Y. Lin、W. Lam、G. Neubig、Y. Li 和 X. Yue（2024）《VisualWebBench：多模态大语言模型在网页理解与定位方面已进化到何种程度？》。载于《第一届语言建模大会》。外部链接：链接。引用自：§5.1。
• E. Lu, Z. Jiang, J. Liu, Y. Du, T. Jiang, C. Hong, S. Liu, W. He, E. Yuan, Y. Wang, Z. Huang, H. Yuan, S. Xu, X. Xu, G. Lai, Y. Chen, H. Zheng, J. Yan, J. Su, Y. Wu, N. Y. Zhang, Z. Yang, X. Zhou, M. Zhang, and J. Qiu (2025) MoBA：面向长上下文大语言模型的混合块注意力机制。外部链接：2502.13189，链接 引用自：§1, §6。
• P. Lu, H. Bansal, T. Xia, J. Liu, C. Li, H. Hajishirzi, H. Cheng, K. Chang, M. Galley, and J. Gao (2024) MathVista：在视觉上下文中评估基础模型的数学推理能力。收录于第十二届国际学习表征会议，外部链接：链接 引用自：§5.1。
• K. Mangalam, R. Akshulakov, and J. Malik (2023) EgoSchema：一个用于超长视频语言理解的诊断性基准。收录于神经信息处理系统进展（NeurIPS）数据集与基准轨道，引用自：§5.1。
• A. Masry, D. X. Long, J. Q. Tan, S. Joty, and E. Hoque (2022) ChartQA：一个结合视觉与逻辑推理的图表问答基准。收录于计算语言学协会会议发现：ACL 2022，S. Muresan, P. Nakov, and A. Villavicencio 编，爱尔兰都柏林，第2263–2279页。外部链接：链接，文献标识码：Document 引用自：§5.1。
• MiMo, B. Xiao, B. Xia, B. Yang, B. Gao, B. Shen, C. Zhang, C. He, C. Lou, F. Luo, G. Wang, G. Xie, H. Zhang, H. Lv, H. Li, H. Chen, H. Xu, H. Zhang, H. Liu, J. Duo, J. Wei, J. Xiao, J. Dong, J. Shi, J. Hu, K. Bao, K. Zhou, L. Li, L. Zhao, L. Zhang, P. Li, Q. Chen, S. Liu, S. Yu, S. Cao, S. Chen, S. Yu, S. Liu, T. Zhou, W. Su, W. Wang, W. Ma, X. Deng, B. Mao, B. Ye, C. Cai, C. Wang, C. Zhu, C. Ma, C. Chen, C. Li, D. Zhu, D. Xiao, D. Zhang, D. Zhang, F. Liu, F. Yang, F. Shi, G. Wang, H. Tian, H. Wu, H. Qu, H. Yi, H. An, H. Guan, X. Zhang, Y. Song, Y. Yan, Y. Zhao, Y. Lai, Y. Gao, Y. Cheng, Y. Tian, Y. Wang, Z. Tang, Z. Tang, Z. Wen, Z. Song, Z. Zheng, Z. Jiang, J. Wen, J. Sun, J. Li, J. Xue, J. Xia, K. Fang, M. Zhu, N. Chen, Q. Tu, Q. Zhang, Q. Wang, R. Li, R. Ma, S. Zhang, S. Wang, S. Li, S. Gu, S. Ren, S. Deng, T. Guo, T. Lu, W. Zhuang, W. Zhang, W. Xiong, W. Huang, W. Yang, X. Zhang, X. Yong, X. Wang, X. Xie, Y. Jiang, Y. Yang, Y. He, Y. Tu, Y. Dong, Y. Liu, Y. Ma, Y. Yu, Y. Xiang, Z. Huang, Z. Lin, Z. Xu, Z. Chen, Z. Deng, Z. Zhang, and Z. Yue (2026) MiMo-v2-flash 技术报告. 外部链接: 2601.02780, 链接 引用自 §1.
• MiniMax (2025a) MiniMax-01: 利用闪电注意力机制扩展基础模型. arXiv 预印本 arXiv:2501.08313. 引用自 §6.
• MiniMax (2025b) MiniMax-M1: 利用闪电注意力机制高效扩展测试时计算. arXiv 预印本 arXiv:2506.13585. 引用自 §1, §6.
• 月之暗面 (2026) Kimi K2.6: 开放智能体基础模型. 备注: https://moonshotai.github.io/Kimi-K2/ 引用自 §1.
• OpenAI, :, S. Agarwal, L. Ahmad, J. Ai, S. Altman, A. Applebaum, E. Arbus, R. K. Arora, Y. Bai, B. Baker, H. Bao, B. Barak, A. Bennett, T. Bertao, N. Brett, E. Brevdo, G. Brockman, S. Bubeck, C. Chang, K. Chen, M. Chen, E. Cheung, A. Clark, D. Cook, M. Dukhan, C. Dvorak, K. Fives, V. Fomenko, T. Garipov, K. Georgiev, M. Glaese, T. Gogineni, A. Goucher, L. Gross, K. G. Guzman, J. Hallman, J. Hehir, J. Heidecke, A. Helyar, H. Hu, R. Huet, J. Huh, S. Jain, Z. Johnson, C. Koch, I. Kofman, D. Kundel, J. Kwon, V. Kyrylov, E. Y. Le, G. Leclerc, J. P. Lennon, S. Lessans, M. Lezcano-Casado, Y. Li, Z. Li, J. Lin, J. Liss, Lily, Liu, J. Liu, K. Lu, C. Lu, Z. Martinovic, L. McCallum, J. McGrath, S. McKinney, A. McLaughlin, S. Mei, S. Mostovoy, T. Mu, G. Myles, A. Neitz, A. Nichol, J. Pachocki, A. Paino, D. Palmie, A. Pantuliano, G. Parascandolo, J. Park, L. Pathak, C. Paz, L. Peran, D. Pimenov, M. Pokrass, E. Proehl, H. Qiu, G. Raila, F. Raso, H. Ren, K. Richardson, D. Robinson, B. Rotsted, H. Salman, S. Sanjeev, M. Schwarzer, D. Sculley, H. Sikchi, K. Simon, K. Singhal, Y. Song, D. Stuckey, Z. Sun, P. Tillet, S. Toizer, F. Tsimpourlas, N. Vyas, E. Wallace, X. Wang, M. Wang, O. Watkins, K. Weil, A. Wendling, K. Whinnery, C. Whitney, H. Wong, L. Yang, Y. Yang, M. Yasunaga, K. Ying, W. Zaremba, W. Zhan, C. Zhang, B. Zhang, E. Zhang, and S. Zhao (2025) GPT-OSS-120B & GPT-OSS-20B 模型卡。外部链接：2508.10925，链接 被 §1 引用。
• OpenAI (2025) 推出 GPT-5。注：https://openai.com/gpt-5/ 被 §1 引用。
• L. Phan, A. Gatti, Z. Han, N. Li, J. Hu, H. Zhang, 等 (2025) 人类的最后考试。arXiv 预印本 arXiv:2501.14249。被 §5.1 引用。
• Qwen (2026) Qwen3.5：通过原生多模态智能体加速生产力。外部链接：链接 被 §1 引用。
• D. Rein, B. L. Hou, A. C. Stickland, J. Petty, R. Y. Pang, J. Dirani, J. Michael, 和 S. R. Bowman (2023) GPQA：一个研究生级别的谷歌无法解答的问答基准。arXiv 预印本 arXiv:2311.12022。被 §5.1 引用。
• K. Sakaguchi、R. Le Bras、C. Bhagavatula 和 Y. Choi（2020）《WinoGrande：大规模对抗性维诺格拉德模式挑战》。载于《AAAI 人工智能会议论文集》第34卷第05期，第8732–8740页。外部链接：ISSN 2159-5399，链接，文献。被 §5.1 引用。
• F. Shi、M. Suzgun、M. Freitag、X. Wang、S. Srivats、S. Vosoughi、H. W. Chung、Y. Tay、S. Ruder、D. Zhou、D. Das 和 J. Wei（2022）《语言模型是多语言链式推理者》。外部链接：2210.03057，链接。被 §5.1 引用。
• H. Sun、Y. Chen、X. Wen、B. Hu、T. Shi、T. Wang、J. Wu、W. X. Zhou 和 J. Wen（2025）《挑战推理边界：面向大语言模型的奥赛级数学基准》。arXiv 预印本 arXiv:2503.21380。被 §5.1 引用。
• M. Suzgun、N. Scales、N. Schärli、S. Gehrmann、Y. Tay、H. W. Chung、A. Chowdhery、Q. V. Le、E. H. Chi、D. Zhou 和 J. Wei（2022）《挑战 Big-Bench 任务及链式推理能否解决它们》。外部链接：2210.09261，链接。被 §5.1 引用。
• J. Tang、Y. Zhao、K. Zhu、G. Xiao、B. Kasikci 和 S. Han（2024）《Quest：面向高效长上下文 LLM 推理的查询感知稀疏性》。载于《国际机器学习大会（ICML）》。被 §6 引用。
• K. Team、Y. Zhang、Z. Lin、X. Yao、J. Hu、F. Meng、C. Liu、X. Men、S. Yang、Z. Li、W. Li、E. Lu、W. Liu、Y. Chen、W. Xu、L. Yu、Y. Wang、Y. Fan、L. Zhong、E. Yuan、D. Zhang、Y. Zhang、T. Y. Liu、H. Wang、S. Fang、W. He、S. Liu、Y. Li、J. Su、J. Qiu、B. Pang、J. Yan、Z. Jiang、W. Huang、B. Yin、J. You、C. Wei、Z. Wang、C. Hong、Y. Chen、G. Chen、Y. Wang、H. Zheng、F. Wang、Y. Liu、M. Dong、Z. Zhang、S. Pan、W. Wu、Y. Wu、L. Guan、J. Tao、G. Fu、X. Xu、Y. Wang、G. Lai、Y. Wu、X. Zhou、Z. Yang 和 Y. Du（2025a）《Kimi Linear：一种富有表现力且高效的注意力架构》。外部链接：2510.26692，链接。被 §1 引用。
• M. Team, C. Xiao, Y. Li, X. Han, Y. Bai, J. Cai, H. Chen, W. Chen, X. Cong, G. Cui, N. Ding, S. Fan, Y. Fang, Z. Fu, W. Guan, Y. Guan, J. Guo, Y. Han, B. He, Y. Huang, B. Ji, C. Kong, Q. Li, S. Li, W. Li, X. Li, Y. Li, Y. Li, Z. Li, D. Liu, B. Lin, Y. Lin, X. Long, Q. Lu, Y. Lu, P. Luo, H. Lyu, L. Ou, Y. Pan, L. Pu, Z. Qu, Q. Shi, Z. Song, J. Su, Z. Su, A. Sun, X. Sun, P. Tang, F. Wang, F. Wang, S. Wang, Y. Wang, Z. Wang, Y. Wu, Z. Xiao, J. Xie, Z. Xie, X. Xu, Y. Yan, J. Yuan, J. Zhang, K. Zhang, L. Zhang, L. Zhang, X. Zhang, Y. Zhang, H. Zhao, W. Zhao, W. Zhao, Y. Zhao, Z. Zheng, C. Zhou, G. Zhou, J. Zhou, W. Zhou, Y. Zhou, Z. Zhou, Z. Zhou, Z. Liu, G. Zeng, C. Jia, D. Li, 和 M. Sun (2025b) MiniCPM4：终端设备上的超高效大语言模型。外部链接：2506.07900，链接 被 §1 引用。
• S. Tong, E. Brown, P. Wu, S. Woo, M. Middepogu, S. C. Akula, J. Yang, S. Yang, A. Iyer, X. Pan, A. Wang, R. Fergus, Y. LeCun, 和 S. Xie (2024) Cambrian-1：一次完全开放、以视觉为中心的多模态大语言模型探索。收录于《神经信息处理系统进展》，A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, 和 C. Zhang 编，第 37 卷，第 87310–87356 页。外部链接：文档，链接 被 §5.1 引用。
• L. Wang, Y. Cheng, Y. Shi, Z. Tang, Z. Mo, W. Xie, L. Ma, Y. Xia, J. Xue, F. Yang, 和 Z. Yang (2025) TileLang：一种面向 AI 系统的可组合分块编程模型。外部链接：2504.17577，链接 被 §4.1 引用。
• Y. Wang, X. Ma, G. Zhang, Y. Ni, A. Chandra, S. Guo, W. Ren, A. Arulraj, X. He, Z. Jiang, T. Li, M. Ku, K. Wang, A. Zhuang, R. Fan, X. Yue, 和 W. Chen (2024a) MMLU-Pro：一个更稳健且更具挑战性的多任务语言理解基准。收录于《神经信息处理系统进展》，A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, 和 C. Zhang 编，第 37 卷，第 95266–95290 页。外部链接：文档，链接 被 §5.1 引用。
• Z. Wang, M. Xia, L. He, H. Chen, Y. Liu, R. Zhu, K. Liang, X. Wu, H. Liu, S. Malladi, A. Chevalier, S. Arora, 和 D. Chen (2024b) 《CharXiv：揭示多模态大语言模型在真实图表理解中的差距》。收录于《神经信息处理系统进展》（NeurIPS）数据集与基准测试轨道，引用于 §5.1。
• H. Wu, D. Li, B. Chen, 和 J. Li (2024) 《LongVideoBench：一个用于长上下文交错视频语言理解的基准》。收录于《神经信息处理系统进展》，A. Globerson, L. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomczak, 和 C. Zhang 编，第 37 卷，第 28828–28857 页。外部链接：文档，链接。引用于 §5.1。
• C. Xiao, P. Zhang, X. Han, G. Xiao, Y. Lin, Z. Zhang, Z. Liu, S. Han, 和 M. Sun (2024a) 《InfLLM：利用高效上下文记忆实现大语言模型免训练长上下文外推》。arXiv 预印本 arXiv:2402.04617。引用于 §6。
• G. Xiao, J. Guo, K. Mazaheri, 和 S. Han (2025) 《优化混合块注意力》。arXiv 预印本 arXiv:2511.11571。引用于 §6。
• G. Xiao, Y. Tian, B. Chen, S. Han, 和 M. Lewis (2024b) 《利用注意力汇聚点实现高效流式语言模型》。收录于《国际学习表征会议》（ICLR）。引用于 §6。
• F. F. Xu, Y. Song, B. Li, Y. Tang, K. Jain, M. Bao, Z. Z. Wang, X. Zhou, Z. Guo, M. Cao, 等 (2024) 《TheAgentCompany：在具有实际后果的真实世界任务上基准测试大语言模型智能体》。arXiv 预印本 arXiv:2412.14161。引用于 §5.1。
• R. Yan, Y. Jiang, 和 B. Yuan (2025) 《Flash 稀疏注意力：更高效的原生可训练稀疏注意力》。arXiv 预印本 arXiv:2508.18224。引用于 §6。
• S. Yang, J. Kautz, 和 A. Hatamizadeh (2025) 《门控 Delta 网络：利用 Delta 规则改进 Mamba2》。外部链接：2412.06464，链接。引用于 §1。
• H. Yen, T. Gao, M. Hou, K. Ding, D. Fleischer, P. Izsak, M. Wasserblat, 和 D. Chen (2025) 《HELMET：如何有效且全面地评估长上下文模型》。收录于第十三届国际学习表征会议，外部链接：链接。引用于 §5.1。
• J. Yuan, H. Gao, D. Dai, J. Luo, L. Zhao, Z. Zhang, Z. Xie, Y. X. Wei, L. Wang, Z. Xiao, Y. Wang, C. Ruan, M. Zhang, W. Liang, 和 W. Zeng (2025) 原生稀疏注意力：硬件对齐且原生可训练的稀疏注意力机制。收录于第63届计算语言学协会年会（ACL）论文集，注：arXiv:2502.11089，被§6引用。
• X. Yue, Y. Ni, K. Zhang, T. Zheng, R. Liu, G. Zhang, S. Stevens, D. Jiang, W. Ren, Y. Sun, C. Wei, B. Yu, R. Yuan, R. Sun, M. Yin, B. Zheng, Z. Yang, Y. Liu, W. Huang, H. Sun, Y. Su, 和 W. Chen (2024) MMMU：面向专家级AGI的大规模多学科多模态理解与推理基准。收录于IEEE/CVF计算机视觉与模式识别会议（CVPR）论文集，被§5.1引用。
• M. Zaheer, G. Guruganesh, A. Dubey, J. Ainslie, C. Alberti, S. Ontañón, P. Pham, A. Ravula, Q. Wang, L. Yang, 和 A. Ahmed (2020) Big Bird：面向更长序列的Transformer架构。收录于神经信息处理系统进展，被§6引用。
• Z. Zhang, Y. Sheng, T. Zhou, T. Chen, L. Zheng, R. Cai, Z. Song, Y. Tian, C. Ré, C. Barrett, Z. Wang, 和 B. Chen (2023) H2O：面向大语言模型高效生成式推理的重击者预言机。收录于神经信息处理系统进展，被§6引用。
• W. Zhao, Z. Zhou, Z. Su, C. Xiao, Y. Li, Y. Li, Y. Zhang, W. Zhao, Z. Li, Y. Huang, A. Sun, X. Han, 和 Z. Liu (2025a) InfLLM-v2：面向无缝短长适配的密集-稀疏可切换注意力机制。CoRR abs/2509.24663。外部链接：Link, Document, 2509.24663，被§6引用。
• Y. Zhao, L. Xie, H. Zhang, G. Gan, Y. Long, Z. Hu, T. Hu, W. Chen, C. Li, J. Song, 等 (2025b) MMVU：衡量专家级多学科视频理解能力。arXiv预印本 arXiv:2501.12380，被§5.1引用。
• 智谱AI (2026) GLM-5.1：来自智谱AI的开放基础模型。注：https://github.com/THUDM/GLM-5，被§1引用。
• J. Zhou, Y. Shu, B. Zhao, B. Wu, Z. Liang, S. Xiao, M. Qin, X. Yang, Y. Xiong, B. Zhang, T. Huang, 和 Z. Liu (2025) MLVU：多任务长视频理解基准测试。收录于IEEE/CVF计算机视觉与模式识别会议（CVPR）论文集，第13691–13701页，被§5.1引用。
• T. Y. Zhuo、V. M. Chien、J. Chim、H. Hu、W. Yu、R. Widyasari、I. N. B. Yusuf、H. Zhan、J. He、I. Paul、S. Brunner、C. GONG、J. Hoang、A. R. Zebaze、X. Hong、W. Li、J. Kaddour、M. Xu、Z. Zhang、P. Yadav、N. Jain、A. Gu、Z. Cheng、J. Liu、Q. Liu、Z. Wang、D. Lo、B. Hui、N. Muennighoff、D. Fried、X. Du、H. de Vries 和 L. V. Werra（2025）《BigCodeBench：通过多样化函数调用与复杂指令对代码生成进行基准测试》。收录于第十三届国际学习表征大会，外部链接：Link，被 §5.1 引用。

附录 A 可视化

为了更好地理解学习型索引器选择了什么，我们在图 5 中可视化了每个注意力头的索引分支在所有查询块与键块对上的选择概率。我们展示了来自早期层（第 1 层）和后期层（第 18 层）的四个注意力头，分别对应四个不同的 GQA 组。跨层来看，学习到的稀疏模式恢复了密集注意力中预期的主要结构：所有注意力头都对局部对角线赋予高概率，持续选择汇聚列，并将剩余预算分配给少量长程相对位置。同时，非局部选择在不同 GQA 组之间并不相同。不同组在共享共同的局部和汇聚模式的同时，关注不同的长程条带，这表明学习到的索引器捕捉到了组特有的稀疏注意力模式，而非坍缩为单一的全局选择模式。

我们进一步研究了 MSA 模型中的注意力汇聚现象。即使没有明确强制索引器选择第一个键值块，我们观察到学习到的索引分支在所有层和所有注意力头中，都会自然地将高选择概率分配给初始块。图 6 展示了两个代表性层（第 4 层和第 24 层）的结果，每层各采样了八个注意力头。在这两层中，每个注意力头都将其注意力权重的很大一部分分配给了第一个 token。这证实了，即使在我们的稀疏注意力机制中，注意力焦点也会自然出现，并且普遍存在于不同的注意力头和层中。

附录 B 初步实验

本节展示了对一个试点模型进行的小规模消融研究。我们的目标是确定对稳定优化和强大下游性能至关重要的训练设计选择。这些结果作为第 3 节所述最终方案的经验基础。

B.1 实验设置

本节中的所有消融实验均使用一个 100 亿参数的试点 Transformer 模型，该模型与主论文中的 MSA 模型采用相同的架构族，但只有 16 层。该模型使用 20 万 token 的词汇表和隐藏层大小。每个注意力模块使用 GQA，包含 32 个查询头、4 个 KV 头、头维度为 128，RoPE 维度为 64。MoE 包含 64 个专家，采用 top-4 专家路由，专家内部维度为 1536。该模型总参数量为 105.3 亿，每个 token 的活跃参数量为 14.7 亿。优化器、学习率调度和分词器均与全尺寸配置一致。每次运行都在全尺寸训练所使用的同一预训练语料库的子集上进行训练。

B.2 索引分支的梯度来源

训练索引分支的一个核心挑战在于，公式 7 中的 top- 选择操作是不可微的。在普通的稀疏注意力前向传播中，被选中的块索引仅被用作离散的路由决策。因此，索引投影 和 无法从语言建模目标中获得有用的梯度，索引器也无法学习应该选择哪些块。有几种可能的方法可以为索引器引入训练信号。我们研究了两种机制，它们在保持稀疏注意力结构的同时，为索引分支提供梯度。

索引分支输出。第一种机制让索引分支贡献一个额外的注意力输出。具体来说，我们为索引分支附加一个值投影，并计算 。该输出通过一个单独的输出投影 添加到层输出中。这种设计通过索引分支对下一个 token 预测的贡献来训练它。

KL 损失。第二种机制通过将索引分支的选择分布与主分支在所选支持集上的分布进行匹配，来直接监督索引分支。我们使用公式 10 中定义的辅助损失。该损失作用于 和 ，并为索引选择提供了明确的训练信号。

为了分离这两种梯度源的影响，我们从零开始以三种配置训练模型，从第一步开始就使用稀疏注意力：

• 仅 LM 损失：索引分支输出被添加到层输出中，模型仅使用语言建模损失进行训练，

  (17)

• 仅 KL 损失：丢弃索引分支输出，索引器仅通过辅助 KL 损失进行训练，

  (18)

• LM 损失 + KL 损失：两种机制均启用，

  (19)

图 7 报告了每种配置相对于在同一数据上训练的全注意力 GQA 基线的逐基准差值。两种单信号配置表现出互补的弱点。仅使用 LM Loss 能保留短上下文能力，但在长上下文检索上表现不佳：由于没有针对 top- 选择本身的优化目标，索引器几乎未受到直接压力去选择相关块。仅使用 KL Loss 能改善检索，但会降低短上下文能力：从层输出中移除该损失会减少语言模型可用的注意力容量。LM Loss + KL Loss 在两个维度上取得了最佳平衡，这也是我们在本节剩余消融实验中所采用的配置。

基于这些结果，我们在本节剩余消融实验中使用 LM Loss + KL Loss 配置。我们稍后在 C.3 节中会展示，一旦在全规模设置中使用了 B.4 节引入的索引器预热方法，Index Branch 输出便不再必要。因此，最终方案保留了 KL 监督，但移除了 Index Branch 的值头及其加性输出路径。

B.3 将 KL 梯度限制在 Index Branch 内

辅助 KL 损失旨在训练 Index Branch 使其匹配 Main Branch 的选择分布。在默认的自动求导图下，KL 梯度会通过 Index Branch 的查询和键投影反向传播至隐藏状态，再通过残差流进入主干网络。在这种情况下，KL 损失成为主干网络的附加目标，而非索引器的局部监督信号。

我们观察到这种梯度路由存在两种失效模式。当 KL 系数较大时，偶发的 KL 梯度尖峰会传播到主干网络中，导致梯度范数尖峰，并在几百步内造成语言模型损失发散（图 8）。即使在稳定的系数下，标准短上下文基准在训练过程中也会逐渐退化（图 9）。我们将这种退化归因于一种自蒸馏效应：主干网络可以通过简化主分支注意力分布来降低 KL 损失，而不是通过改进索引分支来实现。

我们通过在索引分支输入处停止 KL 梯度来解决这两种失效模式（第 3.2 节）。这样，每一层的 KL 损失就成为其自身索引器的局部监督信号。采用这种分离操作后，在原本会导致发散的相同 KL 系数下，语言模型损失和梯度范数保持稳定（图 8），并且短上下文退化也被消除（图 9）。我们在后续所有运行中都使用了这种分离操作。

B.4 索引器预热

我们观察到，在训练的最初阶段，主分支注意力分布变化迅速。如图 10 所示，注意力熵从初始的平滑分布迅速下降到更尖锐的分布，然后进入较慢的表征学习阶段。这使得稀疏选择在初始化时变得脆弱。如果从第零步就启用 top- 选择，索引器必须追踪一个快速变化的目标，而此时它自身的选择几乎还是随机的。糟糕的早期选择会将主分支引导到无信息量的 token 上，从而削弱主干网络的学习以及索引器接收到的 KL 监督信号。

我们通过一个简短的索引器预热阶段来解决这个问题。在预热期间，主分支使用全注意力，而索引分支则通过 KL 散度损失，以全序列主分支分布为目标进行训练。这使得主干网络能够在不出现稀疏路由错误的情况下度过早期的锐化阶段，并为索引器在控制 token 选择之前提供一个有意义的初始化。经过若干步骤后，我们启用 top- 稀疏选择，并继续训练，此时 KL 散度损失仅作用于被选中的支撑集。

图 11 比较了有无此预热阶段的预训练运行结果。经过预热后的运行在短上下文任务上表现更优，长上下文检索能力也更强。这些结果表明，一个简短的全注意力预热阶段能为稀疏训练提供更好的初始化。因此，在通过持续预训练将全注意力检查点转换为稀疏注意力时，我们也采用了这种预热方法。

B.5 可学习的注意力汇聚点

图 6 的可视化结果显示，第一个 token 通常充当注意力汇聚点：即使稀疏选择器没有明确强制包含它，许多注意力头也会将相当一部分注意力权重分配给序列前缀。这引发了一个问题：这种汇聚点行为是否应该由一个显式的可学习机制来表示，而不是被序列中第一个真实的 token 所吸收。因此，我们测试了一种类似 GPT-OSS 风格的可学习注意力汇聚点。具体来说，每个注意力头都获得一个额外的可学习汇聚点 logit，该 logit 在注意力 softmax 中与正常的键位置进行竞争。

图 12 可视化了由此产生的注意力模式。可学习的汇聚节点在某些注意力头中吸收了大量的注意力权重，但它并未完全消除原始的首 token 汇聚现象。在几个注意力头中，尤其是那些可学习汇聚节点接收到的权重较小的头，首 token 仍然获得了大量注意力，并继续充当隐式汇聚节点。

我们还在图 13 中比较了有无可学习汇聚节点时的下游困惑度。可学习汇聚节点变体相比默认设计并未带来明确或一致的改进。考虑到其额外的参数、实现复杂性，以及它并未完全抑制首 token 汇聚行为，我们未将可学习注意力汇聚节点纳入最终方案。

B.6 动态稀疏选择 vs. 滑动窗口

为了评估动态选择的价值，我们将 MSA 与一个 FLOP 匹配的滑动窗口基线进行了比较。该基线移除了索引分支，并使用固定的稀疏模式：每个查询关注第一个键块，以及一个以该查询结尾、具有相同 token 预算的局部窗口。因此，这两种方法具有相同的选择预算，其区别仅在于所选 token 是由位置固定决定还是动态选择。

图14报告了下游智能体任务上的困惑度。在相同的稀疏选择预算下，滑动窗口模型在整个训练过程中的困惑度均高于MSA。尽管两种模型都受益于更多的训练token，但固定的局部窗口模式无法达到动态稀疏选择的困惑度水平。这表明，对于这些智能体任务而言，位置固定的稀疏模式不如基于内容的token选择合适。

附录C 额外消融研究

C.1 块大小

MSA主分支中的稀疏注意力计算以连续token块为单位处理键值对，这会影响模型性能和效率。较大的块可以提高内核效率，但由于选择粒度更粗，可能会降低检索质量。通过在保持选定token总数不变的情况下调整块大小，我们研究了这一权衡。与主实验相比，这些运行使用了更少的训练迭代次数和评估套件的一个子集。

如表4所示，在此设置下，改变块大小对模型质量的影响有限。不同取值下的PPL结果几乎没有变化，并且当块大小从32增加到64或128时，RULER分数没有出现明显的下降。这表明，在这些消融实验中，MSA可以使用更大的键值块来提高内核效率，同时质量损失有限。

媒体内容 · 前往原文查看

基准测试	块大小32	块大小64	块大小128
PPL
TAU2	1.176	1.176	1.176
AgentCompany	1.266	1.276	1.266
HLE	1.299	1.299	1.300
SWE	1.233	1.233	1.233
长上下文检索
RULER-8K	72.5	72.8	73.8
RULER-32K	66.1	65.3	64.6

C.2 强制汇聚与局部选择

在早期的稀疏训练实验中，我们明确强制选择器包含两类块：序列中的第一个块，以及查询位置附近的一个固定局部窗口。第一个块对应常见的注意力汇聚（attention-sink）模式，而局部窗口则保留了对于短程建模至关重要的邻近上下文，并为索引器提供密集监督。这一设计主要是作为一种稳定机制引入的：在索引器变得可靠之前，强制选择这些块可以降低稀疏分支在早期训练过程中遗漏基本上下文的可能性。

我们后来发现，这些先验知识并不需要硬编码。当移除对第一个块和固定局部窗口的强制选择后，训练好的模型仍然展现出这两种结构：当有用时，注意力会集中在序列前缀上，而邻近的 token 也仍然被频繁选中。如表 6 所示，移除强制汇聚（forced sink）和固定局部选择对标准模型质量影响很小：推理、代码和 PPL 指标几乎保持不变。长上下文检索能力也相当。这些结果表明，稀疏模型可以在没有硬编码选择规则的情况下，学习汇聚（sink）和局部选择模式。因此，最终方案不强制选择第一个块或大的局部窗口，仅强制选择特殊的不完整自块（incomplete self block）。

媒体内容 · 前往原文查看

基准测试	无强制	强制
通用知识与推理
MMLU	60.5	60.5
MMLU-Pro	32.5	33.4
BBH	58.2	58.2
ARC Challenge	78.1	77.9
数学
GSM8K	66.0	66.9
MGSM	35.8	36.3
代码
EvalPlus	54.0	53.6
BigCodeBench	35.6	35.7
MultiPL-E MBPP P@10	80.1	79.5
图像
ChartQA	73.5	73.7
MMMU	43.6	42.9
VideoMMMU	32.1	32.0
PPL
TAU2	1.175	1.175
AgentCompany	1.268	1.266
HLE	1.301	1.300
SWE	1.235	1.233
长上下文检索
RULER-8K	71.6	71.7
RULER-32K	61.5	65.8

媒体内容 · 前往原文查看

基准测试	带值头	无值头
通用知识与推理
MMLU	66.4	67.3
MMLU-Pro	39.0	39.1
BBH	65.3	65.9
ARC Challenge	82.2	82.4
数学
GSM8K	77.6	76.4
MathVista	45.2	43.6
MGSM	48.4	47.6
代码
HumanEval	60.4	59.1
EvalPlus	57.7	58.7
BigCodeBench	46.0	44.0
图像
AI2D	69.3	70.4
ChartQA	75.3	74.9
MMMU	44.9	43.4
OCRBench v2	53.2	53.9
MLVU	42.4	43.9
MMVU	44.9	43.7
PerceptionTest	45.0	47.3
长上下文检索
RULER-8K	84.1	83.0
RULER-32K	79.7	80.4

C.3 索引分支值头

我们的初步实验（B.2 节）表明，通过索引分支提供额外的注意力输出，有助于模型从第零步开始进行稀疏训练。然而，这个索引值头引入了额外的计算和复杂性。由于 B.4 节中的索引器预热已经改善了稀疏训练的初始化，我们进一步消融了值头是否仍然必要。

我们将原始的带值头设计与一个仅使用 KL 对齐信号训练索引器的无值头变体进行了比较。如表 6 所示，移除索引值头并未导致评估套件出现系统性性能下降。无值头变体在一些通用推理基准上略优，而带值头变体在部分数学和代码任务上仍保持小幅优势。在多模态基准和长上下文检索方面，差异也是混合的。

总体而言，结果表明，一旦使用了索引分支预热，索引值头就不再关键。它对下游质量的影响很小且因基准而异，两种变体均未持续占据主导地位。这表明，在早期方案中，其主要作用是提供额外的早期训练信号，而非在收敛时提供必要的能力。因此，最终设计基于效率考量去掉了索引值头。在推理时，top-索引器仅需块级最大值，完全避免了值聚合路径和指数计算。