卡内基梅隆团队：稀疏注意力如何优化大模型？

卡内基梅隆大学的研究团队Ranajoy Sadhukhan、Zhuoming Chen（共同第一作者）、Haizhong Zheng、Yang Zhou、Emma Strubell和Beidi Chen在2025年6月5日发表了一项突破性研究，题为《Kinetics: Rethinking Test-Time Scaling Laws》。这项研究对大语言模型(LLM)在测试阶段的资源分配策略提出了全新见解，有兴趣深入了解的读者可访问他们的项目网站：https://infini-ai-lab.github.io/Kinetics。

想象一下，你有一辆赛车（大语言模型），希望它跑得又快又好。传统观点认为，小型赛车只要加足够的油（计算资源）也能跑得和大车一样好。但这项研究发现，事实并非如此简单。

研究团队发现，之前人们对小型语言模型的效率评估存在严重高估。以往的研究主要关注计算量（浮点运算，即FLOPs），忽略了内存访问这一关键瓶颈。就像一辆车不仅需要油（计算能力），还需要足够宽的道路（内存带宽）。在实际应用中，特别是当模型需要生成长篇推理或多次尝试解题时，内存访问常常成为真正的限制因素。

团队通过分析从0.6B到32B参数规模的多种模型，提出了一个名为"Kinetics"的全新缩放定律。这个定律不仅考虑计算量，还将内存访问成本纳入考量，更全面地指导资源分配。就像物理学中的动能公式（Ek = 1/2mv?）一样，Kinetics揭示了测试阶段资源分配的本质规律。

这一新定律得出了一个令人惊讶的结论：测试阶段的计算资源在超过某个参数阈值（实证研究中约为14B）的模型上使用比在小型模型上更有效。这是因为在测试阶段，真正的成本瓶颈不是模型参数数量，而是注意力机制（attention）的计算和内存访问。

受此启发，研究团队提出了一个基于稀疏注意力（sparse attention）的全新缩放范式。这种方法通过降低每个生成token的成本，使模型能在相同资源预算下生成更长的文本或并行处理更多样本。就像让赛车只关注真正重要的道路信息，而不是浪费精力观察所有细节。

实验结果令人震惊：在美国数学邀请赛（AIME）问题解决任务上，采用稀疏注意力的模型在低成本场景下准确率提高了60多个百分点，在高成本场景下也有超过5个百分点的提升。这一优势在所有测试的模型上都得到了验证，包括最先进的混合专家模型（MoEs）。

一、为什么我们需要重新思考测试阶段的资源分配？

近年来，大语言模型展现出惊人的推理能力，特别是在复杂任务解决上取得了巨大突破。这些突破大多依赖于测试阶段的特殊策略，比如长链思维（Long-CoT）和多次尝试选最佳（Best-of-N）。想象你在解一道难题，不仅要写下详细的思考过程（长链思维），还可能尝试多种不同方法（多次尝试），最后选择最好的那个答案。

然而，这些策略带来了巨大的推理成本。就像一辆赛车参加比赛，不仅要考虑油耗（计算能力），还要考虑轮胎磨损和空气阻力（内存访问）。之前的研究主要关注"油耗"（浮点运算），却忽略了"轮胎磨损"（内存访问成本），导致资源分配决策出现偏差。

研究团队通过详细分析发现，在长链思维和多次尝试的场景下，注意力机制的计算和内存访问成本远远超过模型参数的计算成本，比例高达10-1000倍！这就像赛车比赛中，轮胎磨损比油耗更容易成为限制因素。

以Qwen3系列模型在AIME24数学竞赛上的表现为例，按照传统的计算量（FLOPs）评估方法，小型模型通过生成更长的推理或更多次尝试，看起来能达到与大模型相似的性能。但当把内存访问成本也考虑进去后，实际情况截然不同——大模型的效率明显更高，资源节省高达3倍！

二、Kinetics缩放定律：全面考虑计算与内存访问成本

想象你正在烹饪一道复杂菜肴。传统观点认为，只要有足够的烹饪时间（计算资源），一个小厨房（小模型）也能做出与大厨房（大模型）一样好的菜。但实际上，厨房的操作空间（内存带宽）同样至关重要，尤其是制作复杂菜肴时。

研究团队提出的Kinetics缩放定律，就像是一本考虑了厨房空间限制的烹饪指南。它综合考虑了两方面成本：

首先是计算成本，包括： - 模型参数计算：就像是按照菜谱步骤操作的时间 - 注意力机制计算：就像是不断观察和调整菜肴的时间

其次是内存访问成本，包括： - 模型参数访问：就像是查阅菜谱的时间 - KV缓存访问：就像是在厨房中取放各种食材和工具的时间

研究发现，在生成长文本时，注意力机制相关的计算和内存访问成本成为主要瓶颈。具体来说，成本与生成长度的平方成正比（就像动能与速度的平方成正比）。这意味着生成两倍长度的文本，成本不是增加两倍，而是四倍！

这一发现彻底改变了我们对资源分配的理解。以AIME24数学竞赛为例，团队发现：

在长链思维场景下，只有14B和32B这样的大模型才真正受益于生成10K以上的长文本。对于1.7B和4B等小模型，当生成长度低于5K时，更明智的选择是直接换用更大的模型。

在多次尝试场景下，之前的理论认为小模型多次尝试是最经济的选择。但Kinetics表明，使用14B模型并减少尝试次数往往更有效率。研究还发现14B似乎是个临界规模——小于这个规模的模型，资源最好用于增大模型；而达到或超过这个规模后，增加尝试次数或生成更长文本更为有效。

三、为什么小模型效率被高估？解密内存与计算的不平衡

要理解为什么小模型效率被高估，我们需要深入两个关键因素：

首先是KV内存大小与模型参数的不成比例关系。想象你有不同大小的厨房（模型），按理说厨房越大，需要的储物空间（KV缓存）也应该越大。但实际上，小厨房往往需要相对更多的储物空间。

以Qwen3系列为例，0.6B模型处理32K长度的文本需要3.5GB的KV缓存，而模型本身只有1.2GB。相比之下，32B模型仅需要8GB的KV缓存。研究发现，模型参数每增加一倍，KV缓存大小仅增加1.18倍左右。这一现象在OPT、Qwen2.5和LLaMA3等多个模型系列中都得到了验证。

其次是从线性到二次方的成本模型转变。在传统FLOPs计算下，生成8K长度的文本用14B模型（通常不足以解决复杂任务）的成本等同于用4B模型生成24K长度（足以完成大多数任务）。但在考虑KV缓存后，同样成本下14B@8K只相当于4B@9K。这个更严格的限制使得小模型很难仅通过增加生成长度来弥补其能力不足。

通过详细的等成本分析（Iso-Cost分析），研究团队发现：在Kinetics模型下，成本随生成长度呈二次方增长，而KV缓存与模型参数次线性增长。因此，当总预算较低时，增加模型规模比增加生成长度更有效。相比之下，传统基于FLOPs的模型会导致截然不同的结论。

四、稀疏注意力：测试阶段缩放的关键突破

既然确定了注意力机制是测试阶段的主要瓶颈，研究团队提出了一个创新方案：稀疏注意力。这就像是让厨师只关注真正重要的食材和步骤，而不是平等对待所有细节。

在传统注意力机制中，模型会考虑所有历史token的信息，这会导致计算和内存访问成本随文本长度呈二次方增长。稀疏注意力则只关注最相关的token，大幅降低成本。

研究团队提出了"Sparse Kinetics"缩放定律，探索了在稀疏注意力条件下的最佳资源分配策略。结果表明，在稀疏注意力下，计算资源应该优先用于测试阶段策略（如增加生成长度或尝试次数），而不是减少稀疏度。随着测试阶段投入更多计算资源，更高的稀疏度变得越来越重要。

实验结果令人震惊：

在AIME24数学竞赛上，与传统密集注意力相比，稀疏注意力模型在低成本场景下提高了高达60个百分点的问题解决率，在高成本场景下也保持了5个以上百分点的优势。从效率角度看，密集注意力模型需要超过10倍的计算资源才能达到相同的解题率。

研究还发现，稀疏注意力使小模型重新变得有价值。相比传统密集注意力，稀疏注意力显著提高了0.6B、1.7B和4B等小型模型的效率，使它们能在更广泛的场景下应用。

更重要的是，研究团队发现KV预算与生成token之间存在明确的权衡关系。在Best-of-N场景下，每当计算成本增加一倍，最佳KV预算仅增加1.18倍，而生成的token总数则增加1.74倍。这表明，增加生成更多token通常比扩大KV缓存更有效。

五、从理论到实践：实现稀疏注意力的块稀疏方法

为了验证理论的实用性，研究团队实现了一个名为"块稀疏注意力"（block top-k attention）的简单方法。理想的稀疏注意力（逐token选择最相关的K个token）在理论上效果最佳，但在实际中计算成本过高。块稀疏注意力通过选择整块相关token来平衡效率和效果。

块稀疏注意力的工作原理有点像购物时先选择哪些商店值得逛，而不是逐个商品比较。它首先将历史token分成多个块，计算每个块的平均表示，然后选择最相关的K'个块进行注意力计算。这种方法利用了注意力模式的时间局部性，并且能与现有的分页注意力技术无缝集成。

实验表明，块稀疏注意力虽然不如理想的token级稀疏方法，但仍然表现出色：在低成本场景下提高了45个百分点的准确率，并能以仅用传统方法1/8.58的资源达到同等准确率。更重要的是，它在实际硬件上取得了显著加速——在H200 GPU上实现了2.2×到33.3×的吞吐量提升。

研究还比较了块稀疏注意力与本地注意力（只关注邻近token的固定稀疏模式）。虽然本地注意力在实现上更简单，但其性能明显较差，除了在极低准确率场景外，甚至无法超过密集注意力。

对于混合专家模型（如Qwen3-30B-A3B），稀疏注意力的优势更为明显。这类模型通过激活部分参数降低计算成本，但注意力机制仍是主要瓶颈。实验表明，块稀疏注意力能在低成本场景下提高42-53个百分点的解题率，在高成本场景下也保持约8个百分点的优势。

六、更广阔的影响与未来展望

这项研究的意义远超出学术范畴。随着大语言模型在各领域的应用日益广泛，如何高效分配计算资源变得至关重要。Kinetics缩放定律为整个行业提供了新的视角，可能重塑从模型架构设计到硬件系统开发的多个环节。

与预训练阶段的缩放已经出现收益递减不同，测试阶段的缩放仍有巨大潜力，特别是通过增加生成token和优化推理路径。稀疏注意力可能是释放这一潜力的关键。

研究团队希望这项研究能引导并鼓励未来模型架构、测试阶段策略和硬件系统的协同设计，充分释放下一波大语言模型缩放的潜力。虽然本研究主要针对NVIDIA GPU进行分析，但其核心原理——扩展内存带宽比扩展计算能力更具挑战性和成本——适用于各种硬件平台。

值得注意的是，研究也存在一些局限。实验主要集中在Qwen3和DeepSeek-R1-Distilled-Qwen两个预训练推理模型系列，未充分探索训练和后训练策略的影响。此外，研究假设云端服务环境，计算资源充足且批处理规模较大，而在本地部署场景（如使用ollama）中，由于VRAM有限，模型参数访问可能成为主要成本，此时小模型可能更适合。

未来研究方向包括探索更多高级稀疏注意力算法，以及设计能适应稀疏注意力的新型测试阶段策略。研究还发现生成长度与最佳尝试次数之间存在强相关性，这可能作为调整尝试次数和KV预算的动态信号。稀疏注意力的减少推理成本，将为更多推理尝试和更长生成创造可能，在固定资源预算内提供更大的策略灵活性。

总的来说，这项研究不仅揭示了测试阶段资源分配的新视角，还提出了一个实用的解决方案——稀疏注意力，让我们能以更低的成本获得更好的性能。正如研究团队所言，我们可能正站在大语言模型推理新范式的起点，即使在预训练已遇到瓶颈的情况下，测试阶段的优化仍可能带来显著性能提升。