摘要：这项由NVIDIA公司的Adrian Lańcucki、Konrad Staniszewski，以及爱丁堡大学的Piotr Nawrot和Edoardo M. Ponti共同完成的突破性研究，于2025年6月发表在计算机学习领域的顶级期刊上。有兴趣深入了解技术

想象一下，你的手机在运行复杂的AI应用时，不仅能保持原有的准确性，还能比以前快好几倍，这就是这项研究想要实现的效果。研究团队发现了一个关键问题：当前最先进的AI推理系统就像一个记性特别好但动作缓慢的学者，它需要记住之前处理过的每一个细节，这导致它的"大脑内存"很快就被填满，处理速度也越来越慢。

这个问题在AI领域被称为"KV缓存瓶颈"。简单来说，就像你在阅读一本厚书时，需要在脑海中记住前面所有的情节细节，随着故事越来越长，你需要记住的信息越来越多，最终大脑就会感到负担沉重，反应变慢。研究团队提出的解决方案就像是训练大脑学会筛选记忆，只保留最重要的信息，而巧妙地丢弃那些不太重要的细节，从而让思考过程变得更加高效。

研究团队首次提出了"推理时间超级扩展"的概念，这听起来很复杂，但实际上就像是给AI装了一个智能的记忆管理器。这个管理器不是简单粗暴地删除信息，而是会巧妙地延迟删除决定，确保重要信息不会过早丢失。通过这种方法，AI可以在相同的计算资源下处理更长的序列或同时处理更多的任务，就像一个效率极高的多任务处理专家。

更令人惊喜的是，这项技术不需要昂贵的重新训练过程。传统的AI优化方法就像要重新教会一个人说话，需要大量的时间和资源。而这项新技术只需要很少的调整步骤，就像给汽车换个更高效的燃油系统，就能显著提升性能。具体来说，研究团队开发的动态内存稀疏化技术只需要1000个训练步骤就能实现8倍的压缩率，而传统方法可能需要几万个步骤。

在实际测试中，这项技术在多个复杂的推理任务上都表现出色。比如在数学竞赛题目上，使用了这项技术的32B参数模型平均提升了9.1分，在科学问答任务中提升了7.6分，在编程挑战中提升了9.6分。这些提升不是在增加计算资源的前提下实现的，而是在保持相同计算预算的情况下获得的性能突破。

这项研究的意义远不止于技术突破本身。它为AI系统的普及应用开辟了新的可能性，让更多人能够在有限的硬件条件下享受到高质量的AI服务。无论是在智能手机上运行复杂的AI助手，还是在边缘设备上部署先进的推理系统，这项技术都为实现更高效、更实用的AI应用铺平了道路。

一、推理时间扩展的奥秘：让AI思考更深更广

要理解这项研究的核心价值，我们需要先了解什么是"推理时间扩展"。想象你在解决一个复杂的数学题，你有两种策略：一种是快速给出答案，另一种是花更多时间仔细思考，探索不同的解题路径，最终给出更准确的答案。推理时间扩展就是让AI选择后一种策略。

在AI领域，这种策略已经被证明非常有效。就像OpenAI的o1模型和DeepSeek的R1模型一样，它们之所以能在复杂推理任务上表现出色，正是因为它们会在给出最终答案之前进行更深入的"思考"。这种思考过程包括两个维度：一是顺序扩展，就像深入挖掘一个问题的各个层面；二是并行扩展，就像同时探索多种可能的解决方案。

然而，这种策略面临一个根本性的挑战。AI系统在进行推理时，需要维护一个叫做"键值缓存"的记忆系统。这个系统就像一个巨大的图书馆，存储着AI在推理过程中需要参考的所有信息。随着推理链条越来越长，或者同时进行的推理线程越来越多，这个图书馆就会变得越来越庞大，最终超出硬件设备的承载能力。

更严重的是，AI从这个庞大的记忆库中检索信息需要时间，就像在一个杂乱无章的图书馆中寻找特定书籍一样。当记忆库变得过于庞大时，检索时间会大大增加，导致整个推理过程变得缓慢。这就是为什么传统的推理时间扩展方法虽然能提高准确性，但代价是显著增加的计算时间和内存需求。

研究团队敏锐地意识到，问题的关键不在于生成的标记数量，而在于键值缓存的大小。这就像是发现了问题的根源：不是书的内容太多，而是图书馆的存储系统效率太低。基于这个洞察，他们提出了一个革命性的想法：如果能够巧妙地压缩这个记忆系统，让它在保持重要信息的同时大幅减少存储空间，那么就能在相同的硬件条件下实现更大规模的推理扩展。

这种方法的美妙之处在于，它不是简单地减少推理的深度或广度，而是通过提高记忆系统的效率来实现更好的性能。就像是发明了一种新的图书分类和存储方法，让同样大小的图书馆能够存储更多的书籍，同时让查找过程变得更加快速。

研究团队将这种方法称为"推理时间超级扩展"，因为它能够突破传统扩展方法的限制，在相同的计算预算下实现更大规模的推理。这不仅仅是一个技术改进，更是一个思维方式的转变：从关注如何增加计算资源转向如何更高效地利用现有资源。

二、动态内存稀疏化：AI记忆管理的艺术

研究团队开发的核心技术被称为"动态内存稀疏化"，简称DMS。这个技术就像是为AI的大脑设计了一套智能的记忆管理系统，能够自动决定哪些信息值得保留，哪些可以安全地丢弃，而且这个决定过程是动态和自适应的。

传统的记忆管理方法就像一个机械的清洁工，会按照预设的规则定期清理记忆，比如总是删除最旧的信息或者注意力权重最低的信息。虽然这种方法简单直接，但它忽略了一个重要事实：不同的信息在不同的时刻具有不同的重要性。有些看似不重要的信息可能在后续的推理过程中变得至关重要。

动态内存稀疏化采用了一种更加智慧的策略。它不会立即删除被标记为不重要的信息，而是会给这些信息一个"缓刑期"。在这个缓刑期内，AI系统可以继续使用这些信息进行推理，只有当信息真正不再需要时才会被删除。这就像一个有经验的图书管理员，不会急于把书送到仓库，而是会观察一段时间，确保没有读者需要这些书籍后才会移除。

这种延迟删除的策略带来了显著的好处。研究团队发现，立即删除信息会导致性能急剧下降，特别是对于需要长期依赖关系的复杂推理任务。而延迟删除不仅能够保持稳定的性能，还能大幅减少训练所需的数据量。在实验中，采用延迟删除策略的模型比传统方法快8倍达到相同的压缩比，而性能几乎没有损失。

更令人印象深刻的是DMS的训练效率。传统的内存压缩方法，如动态内存压缩技术，需要大量的训练数据和时间来学习如何有效地合并和压缩信息。而DMS采用了一种更简单但更有效的策略：学习何时删除信息，而不是学习如何合并信息。这种策略的简单性使得训练过程变得极其高效。

具体来说，DMS使用了一种叫做"Gumbel-sigmoid分布"的数学技巧来学习删除决策。这听起来很复杂，但实际上就像是训练AI学会掷骰子，但这个骰子的概率是可以调整的。在训练开始时，AI几乎不会删除任何信息，就像一个舍不得扔东西的人。随着训练的进行，AI逐渐学会识别哪些信息可以安全删除，就像逐渐学会整理房间的技巧。

训练过程使用了一种巧妙的目标函数，结合了知识蒸馏和压缩约束。知识蒸馏确保压缩后的模型能够模仿原始模型的行为，而压缩约束确保模型达到预设的压缩目标。这种组合就像是在学习模仿一位大师的技艺的同时，还要确保使用的材料不超过预算。

DMS的另一个创新之处在于它不需要添加额外的参数。研究团队巧妙地重新利用了现有模型中的一个神经元来做删除决策，就像是把汽车的一个装饰性按钮改造成了一个功能按钮。这种设计不仅节省了内存空间，还确保了技术的通用性，可以轻松应用到不同的模型架构上。

三、实验验证：理论照进现实的精彩表现

为了验证动态内存稀疏化技术的实际效果，研究团队进行了一系列全面而严格的实验。这些实验就像是对一个新发明的汽车进行各种路况测试，要确保它不仅在理想条件下表现良好，在各种复杂情况下也能稳定运行。

实验选择了多个具有代表性的AI模型家族，包括Qwen 2.5系列的1.5B、7B和32B参数模型，以及Llama 3.2 1B Instruct模型。这些模型就像不同排量的汽车，能够测试技术在不同规模下的适用性。特别值得注意的是，所有这些模型都使用了分组查询注意力机制，这使得键值缓存的压缩变得更加困难，因为多个查询头需要共享相同的键值信息。

测试任务的选择也极其用心，涵盖了AI推理能力的各个方面。在数学推理方面，研究团队选择了MATH-500和AIME 2024这两个极具挑战性的数据集。MATH-500包含了从高中到大学水平的数学题目，而AIME 2024则是美国数学邀请赛的题目，被认为是测试数学推理能力的金标准。在科学推理方面，GPQA Diamond数据集包含了物理、化学和生物领域的高难度问题，需要深厚的学科知识和严密的逻辑推理。编程任务则选择了LiveCodeBench，这是一个实时更新的编程挑战平台，能够测试AI的代码生成和问题解决能力。

实验设计的巧妙之处在于，研究团队没有简单地比较不同方法的绝对性能，而是构建了性能-效率的帕累托前沿。这就像是比较不同汽车的油耗和动力表现，找出在每个油耗水平下动力最强的车型。通过这种方式，可以客观地评估每种方法在不同效率要求下的最佳性能。

实验结果令人印象深刻。在所有测试的模型规模和任务上，使用DMS技术的模型都显著优于传统方法。具体来说，在AIME 24任务上，DMS平均提升了11.5分；在GPQA Diamond上提升了5.5分；在LiveCodeBench上提升了8.3分。这些提升不是在增加计算资源的前提下实现的，而是在保持相同的内存使用和运行时间预算下获得的。

更有趣的是，研究团队发现性能提升的幅度与模型规模和任务类型有关。在MATH 500这样相对简单的任务上，性能提升随着模型规模的增加而减少，这是因为大模型本身就能很好地处理这些任务。而在GPQA Diamond和LiveCodeBench这样更具挑战性的任务上，性能提升随着模型规模的增加而增加，说明DMS技术在处理复杂任务时的优势更加明显。

研究团队还进行了详细的消融实验，就像拆解一台机器来了解每个部件的作用。他们发现延迟删除策略是技术成功的关键因素。当使用立即删除策略时，无论滑动窗口的大小如何，性能都会快速下降。而采用延迟删除策略时，即使是很小的滑动窗口（16个标记）也能在4倍压缩比下保持良好的性能。

数据效率的比较更是令人震惊。DMS技术只需要传统动态内存压缩方法约八分之一的训练数据就能达到相同的压缩效果。在实际应用中，研究团队使用了比原始论文少60倍的训练数据，就成功训练出了高性能的压缩模型。这种极高的数据效率使得技术的实际部署变得更加可行和经济。

四、技术深度解析：从理论到实现的完整图景

动态内存稀疏化技术的核心在于重新定义了AI模型的注意力机制。在传统的Transformer架构中，每个新的输入标记都会生成对应的键值对，这些键值对会被永久保存在缓存中，供后续的注意力计算使用。这就像是在一个不断增长的图书馆中添加新书，但从不移除旧书，最终导致图书馆变得臃肿不堪。

DMS技术引入了一个智能的门控机制，为每个标记的键值对配备了一个"生存概率"。这个概率不是固定的，而是根据当前的隐藏状态动态计算得出的。具体来说，系统会分析当前输入的特征，然后预测这个标记在未来的推理过程中被需要的可能性。这就像是一个经验丰富的图书管理员，能够根据书籍的内容和当前的借阅趋势预测这本书在未来被借阅的可能性。

在训练阶段，这个概率是连续的，使用Gumbel-sigmoid分布来建模。这种分布的美妙之处在于它能够在连续空间中近似离散的决策过程，从而使得整个系统可以通过标准的梯度下降方法进行优化。训练过程使用了温度参数来控制决策的尖锐程度，低温度会产生更接近离散决策的结果。

滑动窗口机制是DMS技术的另一个关键创新。与传统的立即删除策略不同，DMS为每个删除决策提供了一个延迟期。当系统决定删除某个标记的键值对时，这个决策不会立即执行，而是会等待一定数量的时间步。这个等待期就是滑动窗口，其大小是一个可调节的超参数。

这种延迟删除策略的理论基础来自于对Transformer注意力模式的深入分析。研究表明，解码器模型倾向于更多地关注最近的标记，这种局部性偏好为延迟删除提供了理论支撑。通过给即将被删除的标记一个"缓刑期"，系统能够充分利用这些信息，避免过早删除可能仍然有用的内容。

在推理阶段，连续的删除概率会被四舍五入为离散的0或1，对应保留或删除的决策。被标记为删除的键值对会在滑动窗口期满后从缓存中移除，从而实现内存的动态管理。这种设计既保证了推理时的确定性，又维持了内存使用的可预测性。

训练目标函数巧妙地结合了知识蒸馏损失和压缩约束。知识蒸馏确保压缩后的模型能够模仿原始模型的输出分布，而压缩约束则通过单侧L1损失来控制实际压缩比。这种设计允许模型在训练过程中逐渐学习合适的压缩策略，避免了过度压缩导致的性能急剧下降。

实现层面的一个巧妙设计是避免增加额外的参数。研究团队重新利用了现有查询头中的一个神经元来计算删除概率，并在训练过程中逐渐将这个神经元的影响减至零，确保它不会干扰正常的注意力计算。这种设计不仅节省了内存，还保证了技术的向后兼容性。

性能优化方面，DMS技术充分考虑了现代GPU的内存层次结构。通过与PagedAttention等内存管理技术的结合，DMS能够有效地利用GPU的内存带宽，减少内存访问的延迟。删除的键值对可以直接被新的内容覆盖，避免了频繁的内存分配和回收操作。

五、广泛适用性验证：超越特定任务的通用价值

为了验证DMS技术不仅仅在推理时间扩展场景下有效，研究团队还进行了一系列广泛的适用性测试。这些测试就像是检验一种新药物不仅能治疗特定疾病，还能在各种健康状况下都表现良好。

首先，团队测试了DMS在等长度生成任务中的表现。这种测试模拟了更常见的AI应用场景，即在固定的序列长度下进行推理，而不是追求更长的推理链。结果显示，DMS在4倍压缩比下几乎完全保持了原始性能，而在8倍压缩比下平均只有约3.5分的性能下降。这种轻微的性能损失相对于显著的内存和计算节省来说是完全可以接受的。

更令人惊讶的是DMS在长文本处理任务上的表现。在针刺干草堆任务中，模型需要在长文档中找到特定的信息，这考验的是模型的长期记忆能力。在变量追踪任务中，模型需要跟踪多个变量在长序列中的状态变化。在这些任务中，DMS不仅没有因为压缩而失去性能，反而在某些情况下超越了原始模型。

这种反直觉的结果揭示了一个重要现象：适度的信息压缩可能实际上有助于模型专注于真正重要的信息。就像一个整理得当的工作空间比杂乱无章的环境更有利于提高工作效率，压缩后的记忆系统可能帮助模型避免了信息过载的问题。

研究团队还测试了DMS在不同模型架构和规模下的表现。从1B参数的小模型到32B参数的大模型，DMS都展现出了稳定的性能。特别是在使用分组查询注意力的模型上，DMS的表现尤其出色，这证明了技术对现代高效架构的良好适应性。

在多任务学习场景下，DMS也表现出了良好的泛化能力。模型在数学推理任务上训练的压缩策略，能够很好地迁移到科学问答和代码生成任务上。这种跨任务的泛化能力说明DMS学到的不是特定于任务的启发式规则，而是更通用的信息重要性判断原则。

对比其他高效注意力方法的实验结果更加突出了DMS的优势。传统的稀疏注意力方法如H2O和TOVA在高压缩比下表现出明显的性能退化，而Quest方法虽然保持了较好的准确性，但没有实现真正的内存节省。只有DMS能够在大幅减少内存使用的同时保持甚至提升性能。

研究团队还特别关注了DMS在边缘设备和资源受限环境下的适用性。通过模拟不同的硬件约束，他们发现DMS能够让原本无法运行大型模型的设备成功部署这些模型。这种能力对于AI技术的普及具有重要意义，特别是在移动设备和嵌入式系统的应用场景中。

长期稳定性测试显示，DMS在长时间运行过程中能够保持稳定的压缩性能，不会出现压缩比逐渐下降或性能突然退化的问题。这种稳定性对于实际部署至关重要，确保了系统能够在生产环境中可靠运行。

六、技术影响与未来展望：开启AI效率新纪元

动态内存稀疏化技术的成功不仅仅是一个技术突破，更代表了AI系统设计理念的重要转变。传统的AI优化方法主要关注如何增加计算资源或改进算法结构，而DMS技术开创了一个新的方向：通过智能的资源管理来提升系统效率。

这种理念转变的意义深远。在AI技术快速发展的今天，模型规模越来越大，计算需求越来越高，能耗问题也日益严重。DMS技术提供了一种可持续的发展路径，让我们能够在不大幅增加硬件投入的前提下，获得更好的AI性能。这就像是发明了一种新的燃油技术，让汽车能够用更少的燃料跑更远的距离。

从实际应用的角度来看，DMS技术的影响将是多方面的。在移动设备上，它可以让智能手机运行更复杂的AI应用，而不会快速耗尽电池或占用过多存储空间。在云计算环境中，它可以大幅降低服务器的内存需求，从而降低运营成本并提高服务的可扩展性。在边缘计算场景中，它使得在资源受限的设备上部署先进的AI模型成为可能。

教育和研究领域也将从这项技术中受益。许多学术机构和小型研究团队因为计算资源的限制无法进行大规模的AI实验，DMS技术可以显著降低这些实验的硬件门槛，让更多的研究者能够参与到AI技术的发展中来。这种技术民主化的效应可能会加速整个领域的创新步伐。

从技术发展趋势来看，DMS代表了一类新兴的"智能压缩"技术。与传统的静态压缩方法不同，这类技术能够根据内容和上下文动态调整压缩策略。我们可以预期，未来会有更多类似的技术出现，不仅仅针对内存使用，还可能涉及计算优化、网络传输等多个方面。

研究团队也指出了技术的当前限制和未来改进方向。目前的DMS技术主要针对标准的多头注意力机制，未来需要扩展到其他类型的注意力架构，如多头潜在注意力等。此外，DMS与其他效率优化技术的结合也是一个有前景的研究方向，比如与量化技术或知识蒸馏方法的结合。

更大规模模型和更长上下文的支持也是未来的重要发展方向。随着AI模型不断增长，对效率优化技术的需求也会更加迫切。研究团队计划在更大的模型规模和更长的序列长度上验证DMS技术的效果，并根据需要进行相应的技术改进。

安全性和可靠性方面，虽然DMS技术在实验中表现稳定，但在大规模生产环境中的长期表现还需要更多验证。特别是在关键应用场景中，需要确保压缩过程不会意外删除重要信息，影响系统的可靠性。

从更广泛的视角来看，DMS技术体现了AI发展的一个重要趋势：从单纯追求性能向追求效率和可持续性转变。随着AI技术越来越多地渗透到日常生活中，能效问题和资源利用效率将变得越来越重要。DMS这样的技术为构建更加绿色、更加普惠的AI生态系统提供了重要的技术基础。

来源：至顶网一点号