清华大学团队让AI看懂长文如飞：突破传统注意力机制的革命性框架

在人工智能的世界里，有一个困扰研究者多年的难题，就像是让一个人同时阅读几十本书并记住所有细节一样困难。当大型语言模型需要处理很长的文本时，它们的"注意力"就会变得力不从心，就好比一个学生试图同时关注教室里每一个同学的发言，结果什么都听不清楚。

这项由清华大学赵韦林、周子涵等研究者与OpenBMB团队、哈尔滨工业大学合作开展的研究，于2025年9月发表在arXiv预印本平台（编号：arXiv:2509.24663v1），为这个问题带来了突破性的解决方案。研究团队开发出一种名为InfLLM-V2的全新注意力框架，能够让AI模型在处理长文本时既保持高效率又不丢失准确性，就像给学生配备了一副智能眼镜，能够自动聚焦到最重要的信息上。

传统的AI模型在处理长文本时面临着计算资源和内存的双重瓶颈。当文本长度增加时，模型需要的计算量会呈平方倍数增长，这就像是一个图书管理员需要同时管理的书籍数量翻倍时，他的工作量却要增加四倍。为了解决这个问题，研究人员此前提出了各种稀疏注意力机制，但这些方法往往存在一个关键缺陷：要么在训练时引入额外的参数负担，要么在短文本处理时效率低下，就像是为了节省时间而设计的快速通道，却只在特定情况下才有用。

这项研究的创新之处在于设计了一个"密集-稀疏可切换"的注意力框架。简单来说，就是让AI模型能够根据文本长度自动调整自己的注意力模式：处理短文本时使用全面的密集注意力，就像仔细阅读一篇短文章时关注每一个细节；处理长文本时则切换到稀疏注意力模式，重点关注最重要的信息片段，就像快速浏览长篇小说时抓住关键情节。

一、传统方法的困境与新思路的诞生

要理解这项研究的重要性，我们需要先了解AI模型处理文本的基本原理。现代大型语言模型都基于一种叫做Transformer的架构，其核心是自注意力机制。可以把这个机制想象成一个会议室里的讨论：每个与会者（文本中的每个词）都需要听取和回应其他所有人的发言。在小型会议中这样做很有效，但当参会人数增加到几百甚至几千人时，这种全员参与的讨论模式就变得不现实了。

研究团队发现，现有的稀疏注意力方法主要分为两类。第一类是训练无关的方法，它们利用注意力权重的自然稀疏性来加速推理，就像在嘈杂的会议中只听声音最大的几个人发言。然而，这种方法为了避免严重的性能下降，只能应用有限的稀疏程度，最终的效率提升也很有限。

第二类是可训练的稀疏注意力方法，其中最具代表性的是NSA（原生稀疏注意力）方法。NSA的设计思路是创建三个不同的注意力模块：压缩注意力、选择注意力和滑动注意力，然后通过一个门控机制将它们的输出进行组合。这就像是在会议室里设置三个不同的讨论小组，每个小组负责不同类型的信息处理，最后再将结果汇总。

然而，NSA方法存在一个根本性问题：它与标准的"短序列预训练、长序列微调"工作流程不匹配。这种工作流程是目前构建长文本处理模型的主流方式，就像是先教学生处理简单问题，再逐步提高难度。NSA需要引入三套独立的键值参数和三个注意力模块，这使得模型从处理短文本的单输出注意力架构突然切换到多输出注意力架构，造成了训练的不稳定性。

二、密集-稀疏可切换注意力框架的设计原理

InfLLM-V2的核心创新在于设计了一个真正无缝的转换机制。研究团队的方法建立在InfLLM这一无训练块稀疏注意力机制的基础上，但加入了三个关键创新。

首先是无缝的短到长适应能力。与NSA需要额外参数和多个注意力模块不同，InfLLM-V2通过直接重用现有的密集注意力参数来实现从密集到稀疏注意力的转换。这种设计就像是一个变形金刚，能够在不更换零件的情况下改变自己的形态和功能。当处理短文本时，模型使用标准的密集注意力；当文本长度超过某个阈值时，模型自动切换到稀疏模式，但所有的核心参数保持不变。

第二个创新是对短序列和长序列的统一效率优化。由于InfLLM-V2的转换过程不需要额外参数且引入的分布偏移很小，模型在短文本上保持了强大的性能，同时可以轻松切换回密集注意力模式以获得短序列的高效率。这就像是一辆能够在城市道路和高速公路之间自由切换驾驶模式的智能汽车，在每种环境下都能发挥最佳性能。

第三个创新是加速的块选择机制。研究团队发现，稀疏注意力之前的块选择步骤本身就会损害稀疏注意力的效率增益。他们提出了一种硬件感知的高效实现方案，有效地消除了这个瓶颈，释放了稀疏注意力的全部潜力。

具体来说，InfLLM-V2采用了统一的键值投影参数。研究团队发现NSA中使用三套独立的键值投影参数是不必要的，这不仅复杂化了从短到长序列的适应过程，还显著降低了短序列的计算速度。因此，他们提出使用单一共享的投影参数集，这些参数用预训练的密集注意力参数初始化，然后用于长序列的微调。

在计算对齐方面，除了确保稀疏和密集注意力共享相同的参数外，它们的计算过程也必须紧密对齐。在NSA中，三个注意力模块都会生成输出，然后通过额外的门控模块进行聚合。这迫使即使在短序列上也要计算所有三个模块，导致了大量的开销。为了缓解这个问题，InfLLM-V2将选择注意力和滑动注意力这两种稀疏模式合并，并消除了压缩注意力的输出，形成了统一的稀疏注意力模块。

三、块表示的三阶段压缩策略

在稀疏注意力的实现中，一个关键问题是如何有效地表示和选择重要的信息块。简单地用大块大小在一个阶段内压缩长序列会导致粒度信息的显著丢失。为了解决这个问题，研究团队实现了一个三阶段的、从粗粒度到细粒度的压缩过程。

第一阶段处理输入键序列，产生中间的粗粒度表示。通过将初始压缩块大小表示为lC1，步长表示为sC1，他们通过对连续块应用均值池化操作来实现这一点。这就像是先用广角镜头拍摄整个场景，获得一个总体的印象。

第二阶段采用块级稀疏注意力而不是词级方法，以提高稀疏注意力的效率。在使用分组查询注意力的模型中，他们通过强制组内所有头的块选择模式相同来实现这一点。他们在头组内进行求和以获得共享的重要性分数。这个阶段就像是用中焦镜头重新审视重要区域，获得更详细的信息。

第三阶段应用最大池化操作，它可以保留最显著的特征。聚合后的分数被定义并用于稀疏注意力。这就像是用微距镜头仔细观察最关键的细节，确保不遗漏任何重要信息。

通过这种三阶段方法，研究团队设置了特定的参数值：lC1等于块大小B的一半，sC1等于B的四分之一，第三阶段的参数l为5，s为4，这样可以达到与一阶段块大小B压缩相同的压缩比率。直观地说，他们基于块内的几个滑动子块计算整个块的稀疏分数。

四、高效实现的技术突破

为了实现真正的加速效果，研究团队面临的一个主要挑战是压缩分数计算引入的显著性能瓶颈。性能分析显示，计算压缩分数Scmp引入了明显的性能瓶颈。这种缓慢的一个主要来源是将第一阶段注意力分数SC1存储到慢速GPU高带宽内存中所需的大量输入输出操作。

需要写入的数据量是hqn?/sC1，其中n是完整序列长度。由于sC1远小于n，将完整注意力分数矩阵具体化到GPU高带宽内存会产生巨大的成本。受到FlashAttention的启发，研究团队的目标是通过确保注意力分数尽可能多地保留在快速GPU静态随机存取内存中来最小化这种输入输出。

他们的方法是融合头组求和，将第二阶段压缩所需的头组求和直接融合到基于静态随机存取内存的FlashAttention计算循环中。这样，他们只需将减少后的注意力分数Sshared存储到GPU高带宽内存中，其大小为hqn?/(sC1G)。

然而，另一个挑战来自于头组维度上的求和和序列维度上的在线softmax不是可交换操作这一事实。这种冲突阻止了直接融合。为了克服这个问题，研究团队实现了一个两遍方法。在第一遍中，他们在静态随机存取内存中计算softmax归一化所需的对数求和指数项。在第二遍中，他们利用对数求和指数来计算最终的注意力分数，在静态随机存取内存中执行头组求和，并将减少后的分数写入高带宽内存。

这种两遍方法的代价是它使计算工作量翻倍。因此，研究团队提出了对数求和指数近似，通过使用更粗粒度的注意力分数SC2来近似对数求和指数计算。通过设置sC2为4sC1，lC2为4lC1，计算开销从2倍减少到1.25倍。

五、实验验证与性能表现

研究团队在多个任务上评估了InfLLM-V2，从短文本到长文本，并展示了其效率优势。他们使用标准的分组查询注意力模型骨干，具有8B参数，隐藏大小为4096，查询头数为32，键值头数为2，头维度为128。

在预训练阶段，他们首先使用全注意力在短序列数据上预训练模型。预训练数据集包含8T个4k长度序列的标记，主要由FineWeb-Edu和Stack-v2组成。当转换到长文本微调时，他们切换到InfLLM-V2的稀疏模式。

在长文本理解任务上，研究团队在RULER、LongBench和LongPPL基准上比较了InfLLM-V2和不同基线方法。RULER是一个可配置平均长度的合成数据集，LongBench是长文本理解的双语基准，LongPPL是长序列的困惑度评估基准。实验结果显示，InfLLM-V2在所有稀疏方法中实现了最佳性能，结果与强大的全注意力基线高度竞争且接近匹配。

在长推理任务上，研究团队比较了几个主要的长推理任务，包括MATH-500、AIME和LiveCodeBench。他们在OpenMathReasoning和OpenCodeReasoning上微调了InfLLM-V2和基线方法。实验结果表明，InfLLM-V2达到了与全注意力相当的性能，证实了其在长输出场景中的有效性。

在效率方面，研究团队首先评估了在NVIDIA A100和NVIDIA 4090上的内核实现效率。当选择块数为16时，InfLLM-V2在A100上实现了比FlashAttention高达7.4倍的加速，在4090上实现了9.3倍的加速。相比之下，NSA在相同设置下的加速仅限于3.5倍。端到端推理速度实验显示，InfLLM-V2可以实现2.13倍的预填充加速和2.32倍的解码加速。

六、与现有方法的深度对比

通过与现有方法的详细比较，可以更清楚地看出InfLLM-V2的优势所在。训练无关的稀疏注意力方法，如InfLLM和MInference，虽然能够利用注意力层的内在稀疏性，但受限于稀疏度水平，为了避免严重的性能下降，最终的加速效果有限。

可训练的稀疏注意力方法中，SeerAttention采用自蒸馏后训练算法来训练选择查询块相关上下文的路由器。MoBA在短到长适应阶段采用块稀疏注意力结构，训练查询块和键值块之间的路由器进行上下文选择。这些方法将查询标记分割成块，只能加速预填充阶段。

NSA设计了三个注意力组件用于标记级稀疏性，有效地加速了预填充和解码过程。然而，NSA引入了大量额外参数，使其不适合高效的短到长适应，并在短序列处理上造成了显著的计算开销。

相比之下，InfLLM-V2专注于提出一种稀疏注意力机制，能够有效且高效地处理短序列和长序列，支持预填充和解码。通过共享键值投影和统一的计算流程，InfLLM-V2实现了真正的无缝转换，既保持了短序列的高效率，又在长序列上达到了优秀的性能。

特别值得注意的是，InfLLM-V2的一个独特优势是能够在密集模式和稀疏模式之间灵活切换。这种灵活性不仅提供了密集计算的选项，还可以进一步提升性能，甚至超越全注意力基线。实验中，使用LSE近似的InfLLM-V2稀疏变体没有损失任何性能，证实了加速技术的有效性。

七、技术实现的创新细节

InfLLM-V2的技术实现包含多个创新细节，这些细节共同确保了方法的实用性和高效性。在块选择机制方面，原始的选择注意力模块根据压缩注意力模块的注意力分数识别重要的标记块。对于位于块bi中索引为i的查询标记，注意力总是授予固定的初始块集合和局部块集合。然后在剩余块集合上应用top-k选择，完整的被关注块索引集合是这三个集合的并集。

通过将选择注意力和滑动注意力中的局部块和滑动注意力中的窗口合并，研究团队通过扩展统一稀疏注意力中的局部块数量来严格覆盖滑动注意力的区域。这种单输出设计更好地反映了密集注意力，有助于稀疏注意力模型的训练。

在硬件实现方面，研究团队遵循NSA的技术，将分组查询注意力的组大小设置为16，这是一个非常适合块稀疏注意力的配置。他们的分析显示，计算压缩分数的一个主要瓶颈来源是将第一阶段注意力分数存储到慢速GPU高带宽内存所需的大量输入输出操作。

为了解决这个问题，研究团队实现了融合头组求和技术，将第二阶段压缩所需的头组维度求和直接融合到基于静态随机存取内存的FlashAttention计算循环中。通过这种方式，他们只需要将减少后的注意力分数存储到GPU高带宽内存中，显著减少了内存访问开销。

八、实际应用前景与意义

InfLLM-V2的成功不仅仅是一个技术突破，更代表了AI模型处理长文本能力的重要进步。在实际应用中，这项技术可以显著改善多个领域的用户体验。

在文档处理领域，律师事务所和研究机构经常需要处理长达数百页的法律文件或研究报告。传统的AI模型在处理这些长文档时往往会丢失重要信息或处理速度极慢。InfLLM-V2使得AI助手能够快速阅读和分析这些长文档，同时保持对关键信息的准确理解。

在代码开发领域，软件工程师经常需要理解和维护包含数千行代码的大型项目。InfLLM-V2使得AI编程助手能够更好地理解整个代码库的上下文，提供更准确的代码建议和错误修复方案。这就像给程序员配备了一个能够同时记住整个项目所有细节的超级助手。

在教育领域，学生和研究者经常需要阅读和理解长篇学术论文或教材。基于InfLLM-V2的AI系统可以帮助他们快速提取关键信息，生成摘要，或回答特定问题，显著提高学习效率。

在客服和咨询领域，AI客服系统需要记住整个对话历史来提供连贯的服务。InfLLM-V2使得这些系统能够处理更长的对话历史，提供更个性化和上下文相关的回应。

研究团队已经基于InfLLM-V2框架训练并开源了MiniCPM4.1混合推理模型，为研究社区提供了可重现的实现。这个模型在多个基准测试中展现了优异的性能，证明了框架的实用性和可扩展性。

InfLLM-V2的另一个重要意义在于它为未来的研究指明了方向。这项工作表明，通过巧妙的设计，可以在不增加模型复杂性的情况下显著提升性能。这种"做减法"的设计哲学对于构建更高效、更实用的AI系统具有重要的指导意义。

说到底，InfLLM-V2代表了AI领域中一个重要的进步方向：不是简单地增加模型规模或计算资源，而是通过更智能的设计来解决实际问题。这项研究让我们看到，有时候最好的解决方案不是最复杂的，而是最优雅的。就像一个经验丰富的图书管理员知道如何快速找到读者需要的书籍一样，InfLLM-V2教会了AI模型如何更聪明地管理和利用信息。

这项技术的成功也提醒我们，在追求AI能力提升的道路上，效率和实用性同样重要。随着AI模型越来越多地应用到实际场景中，像InfLLM-V2这样既能保持高性能又能提高效率的技术将变得越来越重要。对于想要深入了解这项技术细节的读者，可以通过论文编号arXiv:2509.24663v1查询完整的研究报告。

Q&A

Q1：InfLLM-V2与传统稀疏注意力方法有什么不同？

A：InfLLM-V2的最大不同在于它能够在密集和稀疏注意力模式之间无缝切换，不需要额外参数。传统方法如NSA需要三套独立的键值参数和多个注意力模块，导致从短文本到长文本适应时出现训练不稳定。InfLLM-V2直接重用现有的密集注意力参数，确保了短到长序列处理的一致性。

Q2：InfLLM-V2在实际应用中能带来多大的速度提升？

A：实验结果显示，InfLLM-V2比传统的FlashAttention在A100上快7.4倍，在4090上快9.3倍。在端到端推理中，它实现了2.13倍的预填充加速和2.32倍的解码加速。同时，它在长文本理解任务上保持了98.1%的性能，在推理任务上保持了99.7%的性能。

Q3：普通开发者如何使用InfLLM-V2技术？

A：研究团队已经基于InfLLM-V2框架训练并开源了MiniCPM4.1混合推理模型，开发者可以通过huggingface.co/openbmb/MiniCPM4.1-8B获取。这个8B参数的模型提供了完整的实现代码，让开发者能够直接应用这项技术到自己的长文本处理任务中。