技术Blog-4 | 新一代InfLLM：可训练的稀疏注意力机制

本文深入剖析 MiniCPM4 采用的稀疏注意力结构 InfLLM v2。作为新一代基于 Transformer 架构的语言模型，MiniCPM4 在处理长序列时展现出令人瞩目的效率提升。传统Transformer的稠密注意力机制在面对长上下文时面临着计算开销迅速上升的趋势，这在实际应用中造成了难以逾越的性能瓶颈。为此，MiniCPM4 创新性地引入了 InfLLM v2，一种 可训练的稀疏注意力机制，极大程度降低了计算开销，并结合定制化推理算子 CPM.cu，MiniCPM4 在 预填充 和 解码 阶段都实现了切实的加速效果。

➤ GitHub 链接：

🔗 https://github.com/openbmb/minicpm

➤ 技术报告链接：

🔗https://github.com/OpenBMB/MiniCPM/blob/main/report/MiniCPM_4_Technical_Report.pdf

🔗 https://arxiv.org/pdf/2506.07900

➤ HuggingFace 链接：

🔗https://huggingface.co/collections/openbmb/minicpm-4-6841ab29d180257e940baa9b

➤ ModelScope 链接：

🔗https://www.modelscope.cn/collections/MiniCPM-4-ec015560e8c84d

InfLLM v2：原生稀疏注意力模型

面对当前不断增长的处理长序列处理需求，当前许多研究致力于设计无需训练的稀疏注意力机制，以动态选择相关上下文词元进行长上下文理解与生成。无需训练的方法往往只能利用大模型自涌现的稀疏特性，导致稀疏性不够理想，只能应用于预填充加速阶段。为解决该问题，MoBA 和 NSA 在预训练阶段就引入稀疏注意力进行训练，能够进一步提高注意力层的稀疏性。然而，MoBA 使用查询块的设计，使其无法在解码阶段实现加速。NSA 引入了三种不同的注意力组件来捕获长距离信息，这对短序列带来了大量计算开销。

为解决以上问题，基于稀疏注意力模型 InfLLM，MiniCPM4 设计了一种可训练的稀疏注意力 InfLLM v2，以降低预填充和解码阶段的计算和内存访问成本。与现有方案相比，InfLLM v2 展现出三大核心优势：

• 真正的全流程加速：不同于 MoBA 仅能加速预填充阶段，InfLLM v2 通过词元级查询分组的设计，在解码阶段也能实现单 token 级别的加速，避免了训练推理不一致的问题。
• 零额外参数负担：相比 NSA 需要三倍键值存储成本和额外注意力组件，InfLLM v2 不引入任何额外参数，短序列处理时自动退化为标准注意力，确保了短文本场景的原生效率。
• 60% 计算成本削减：通过创新的高效 Top-K 实现和 LSE 近似方法，InfLLM v2 在上下文选择过程中相比 NSA 减少了 60% 的计算开销。

InfLLM v2的整体框架

传统稠密注意力计算过程，每个词元都需要与过往所有词元进行逐一的相关性计算，带来了巨大的计算与访存开销。因此，为避免逐词元的计算与访问，在 2024 年 2 月，清华大学就提出了 InfLLM 算法。

InfLLM 提出将上下文进行分块分区域处理。注意力层进行计算之前，对每个查询词元逐块地进行上下文相关性计算，并选取少量最相关的上下文块参与最终的注意力计算。即实现了注意力层的智能化选择机制，只对最有相关性的重点区域进行注意力计算“抽查”。

➤ InfLLM 算法：

🔗https://openreview.net/forum?id=bTHFrqhASY

InfLLM 是一种无需训练的动态稀疏注意力机制，能够被应用于所有的 Transformer 架构模型中实现长文本推理加速。MiniCPM4 中对 InfLLM 算法进行进一步的改进，实现了稀疏度更高、能够同时加速预填充与解码阶段的原生稀疏注意力机制 InfLLM v2：形式上，InfLLM v2 将键值划分为等大小的块，每个块包含 m 个词元。例如，如果总共有 l 个词元，块大小为 m，那么将产生大约 l/m 个块。每个块包含了对应位置的 m 个连续的键向量和值向量。InfLLM v2 的稀疏注意力计算包含两个阶段：

• 动态上下文块选择：基于当前的查询词元，我们从所有块中动态选择最相关的块。为此，我们需要计算查询词元与每个块之间的相关性得分，然后选择得分最高的若干个块。
• 稀疏注意力计算：基于第一阶段选择的块，我们计算查询词元与这些选定块内所有词元之间的注意力。这样既保证了计算效率，又能捕捉到最重要的上下文信息。

动态上下文块选择

InfLLM v2 最关键的组件是查询词元与键值块之间的相关性得分计算。为了避免逐个词元的相关性计算，InfLLM 从每个块中选择代表性词元作为块表示，将相关性得分定义为查询词元与这些代表性词元的点积。虽然这种方法能够捕捉块内的重要语义信息，但代表性词元的选择涉及词元级别的计算和内存访问，成为 InfLLM 的效率瓶颈之一。因此，在 InfLLM v2 中，我们引入细粒度语义核来捕捉块语义，避免词元级别的内存访问。此外，我们要求同组内的查询头共享相同的键值块，以降低内存访问成本。

语义核：InfLLM v2 改进了相关性得分计算方法。由于稀疏注意力机制需要最小化内存访问的不连续性，InfLLM v2 需要对键值序列进行粗粒度的块划分，块大小 m 通常是一个相对较大的值。如果我们只用一个向量来表示每个块的语义，将不可避免地遇到信息丢失的问题。为了实现更准确的相关性计算，InfLLM v2 引入细粒度的"语义核"来构建每个键值块的表示。具体而言，InfLLM v2 以更细的粒度对键值序列进行划分，产生若干语义核。为了确保输入序列中的每个语义片段都包含在完整的语义核内，这些核之间需要相互重叠。

形式上，InfLLM v2 将输入序列划分为大小为 p 的语义核，相邻核之间的步长为 s。因此，键向量序列被划分为多个语义核，总数约为序列长度 l / 步长s。由于语义核只参与相关性计算，因此只需要保留键向量， 节约了标准注意力计算当中对值向量的存储和计算开销。InfLLM v2 使用平均池化操作来计算每个语义核的表示，即对语义核内的所有键向量求平均。查询词元与语义核之间的相关性得分定义为：查询向量与语义核平均表示的点积，再经过 softmax 归一化。查询词元与块之间的相关性得分则表示为：与该块有交集的所有语义核的相关性得分的最大值。基于块相关性得分，InfLLM v2 选择得分最高的 k 个块。然后，InfLLM v2 计算查询词元与这些选定块内所有词元之间的注意力，生成最终输出。

值得注意的是，考虑到初始词元以及局部窗口内的词元通常对最终输出贡献很大，InfLLM v2 将每个查询词元与初始键值块以及局部窗口块之间的相关性得分设置为无穷大。这一机制确保每个查询词元都能关注初始块和局部窗口内的块。当文本长度较短，不超过k个块的总长度时，InfLLM v2 会退化为标准的稠密注意力机制。

Top-K块共享：当前大语言模型架构通常采用分组查询注意力层，即多个查询共享一个键值头。在 InfLLM v2 中，我们要求同一组内的查询头共享相同的 top-k 相关块。这样可以尽可能地减少内存访问。具体而言，在每个查询头与语义核计算完相关性之后，我们在组内对相关性分数进行平均，并将这个平均值作为整个查询组的相关性分数。

高效Top-K实现：Top-K 选择包含三个连续步骤。第一步：计算查询词元与每个语义核之间的相关性分数，然后进行 softmax 归一化。第二步：在查询组维度上对每个语义核的相关性分数进行聚合。第三步：基于聚合后的分数，为每个查询词元选择得分最高的 K 个上下文块。这一操作构成了稀疏注意力机制的计算瓶颈。传统的稠密注意力机制通常使用 FlashAttention 算法来减少显存使用并加速注意力计算。FlashAttention 利用 online softmax 技术来最小化注意力计算过程中对高带宽内存（HBM）的访问操作。

然而，Top-K 选择与注意力计算有本质区别，因为它需要获得每个查询组与每个语义核之间相关性分数的精确数值。由于相关性计算、softmax 归一化和查询组维度的分数聚合等操作不满足交换律，因此无法使用 online softmax。这导致 Top-K 选择需要对语义核进行两遍内存访问和计算：第一遍用于计算对数求和指数（LogSumExp，简称LSE），第二遍用于计算最终的相关性分数。Top-K 选择是 InfLLM v2 在长上下文处理中的瓶颈。为了降低计算成本，我们提出了一种高效的 LSE 近似方法。与计算查询词元和所有语义核之间的点积结果不同，我们尝试通过引入粗粒度语义核来近似 LSE 值。粗粒度语义核的核大小远大于普通语义核的大小。然后我们计算查询词元与粗粒度语义核之间相关性分数的 LSE。这种方法的计算和内存访问成本仅为原始方法的一小部分，具体比例为普通语义核步长与粗粒度语义核步长的比值。

可训练稀疏注意力的设计原则

随着需要长上下文处理和深度推理能力的应用不断发展，可训练稀疏注意力机制在提高预训练和推理效率方面展现出巨大潜力。在本节中，我们讨论 InfLLM v2 的几个关键特性和设计原则。我们希望这些讨论能够推动可训练稀疏注意力机制的未来发展。

复杂度分析

InfLLM v2 使每个词元只需要与得分最高的 k 个键值块进行注意力计算，可以大幅降低注意力的计算和访存开销。在此我们主要分析 InfLLM v2 的计算复杂度和访存复杂度。在第一阶段，我们需要计算查询词元与每个语义核之间的相关性分数。对于上下文长度为l的查询词元，共有约 l/s 个语义核（其中 s 是步长）。因此，该查询词元需要进行相应次数的向量乘法和内存访问。在第二阶段，我们需要计算查询词元与 k 个键值块之间的相关性。在此过程中，该查询词元需要进行 2km 次向量乘法和访存（其中 m 是块大小）。相比于需要 2l 次向量操作和访存的稠密注意力机制，当序列很长（l>>m）时，InfLLM v2 可以将计算开销和访存开销降低至原来的 1/s。我们可以看到，第一阶段的计算复杂度仍然是序列长度的平方关系，而第二阶段的计算复杂度是线性的。因此，如果想要进一步提高稀疏注意力的效率，我们应该在相关块检索上投入更多努力。

查询和键值词元的不同粒度

在 InfLLM v2 中，我们让每个查询词元与不同的键值块进行注意力计算。这里，查询词元的计算单元是词元级别的，而键值的计算单元是块级别的。许多以往的稀疏注意力工作也对查询序列进行分块，使一个块内的所有查询词元共享相同的 top-k 键值块。然而，查询分块操作只能加速长序列预填充，而无法加速解码过程，因为解码过程需要逐个词元生成，大多数情况下查询词元无法组成一个完整的块。因此，训练过程中查询采用块级别单元会在解码过程导致训练推理不一致的问题，进而降低模型效果。可训练的上下文选择在稀疏注意力机制中，top-k 选择这一操作是不可导的。这意味着语义核的表示无法通过第二阶段的稀疏注意力计算得到优化。因此，在 InfLLM v2 中，我们选择使用平均池化这样一种无需参数的操作来构建语义核的表示。同时，平均池化还有一个优点就是能够保证语义核的表示仍然与词元级别的键向量在相同的语义空间。这样的话，模型训练就能够通过优化词元级别的键向量来间接优化语义核的表示。NSA 使用了压缩注意力机制，让第一阶段上下文选择的输出也加到注意力输出上，从而使得块表示能够得到优化。但这一操作会对短文本增加大量额外开销。而 InfLLM v2 采用的平均池化操作对短文本效率没有影响，具有更好的实用价值。

超参数推荐

稀疏注意力机制的高效训练与推理需要算法设计与算子设计能够高度协同。因此，从算法效果和硬件约束来看， InfLLM v2 中的超参数设计需要满足一定的条件。在语义核大小的选择上，越小的核大小越能够获得更精准的相关性分数计算，然而越小的核大小会带来更大的计算开销。因此，为了在效果与效率之间取得较好的平衡，我们选择将核大小设置为 32，步长设置为 16。受限于 GPU 张量核心的矩阵乘法累加指令，一个查询组中必须包含 16 个以上的头，从而保证硬件能够被充分使用，确保模型能够高效训练与推理。

总结

随着 DeepResearch、Code Agent 等应用对大模型长文本处理能力的需求愈发旺盛、深思考能力受到越来越多的关注，让大模型能够高效地理解与生成长文本成为了当前大模型研究的重点。稀疏注意力能够有效地降低大模型长文本处理的开销，对未来模型上下文窗口大小的提升具有重要意义。MiniCPM4 是首个开源的原生稀疏注意力模型，将模型参数、训练算子、推理算子全面开源，希望能够促进稀疏注意力机制研究的发展。

文章来自于“OpenBMB开源社区”，作者“OpenBMB”。