让大模型推理更省内存，还能保持准确率
Thought-Aware KV Cache Compaction for Reasoning via Adaptive Attention Matching
- 按推理步骤重要性分段压缩，动态分配存储空间
- 相同内存下比传统方法准确率更高，峰值内存降65%
- 适合需要长时间推理的数学题或复杂任务场景
推理型语言模型生成的思维链（CoT）序列在解码过程中导致键值缓存线性增长，成为内存瓶颈。现有压缩方法将推理路径视为平铺的词元序列，采用统一压缩策略，忽略了CoT推理中不同步骤重要性差异显著的层次结构。本文提出思考感知注意力匹配（TAM），通过三个机制利用该结构：(i) 思考分段，将轨迹分解为推理块；(ii) 自适应预算分配，根据每个片段的重要性和大小分配压缩预算；(iii) 关键词元保护，保留高注意力推理锚点。我们证明了该分配规则在凸误差模型下是最优的，且连续压缩下的累积误差有界。在AIME 2024和MATH-500数据集上使用Qwen3-4B进行实验，结果表明，在相同内存开销下，TAM优于均匀压缩，周期性压缩可将峰值内存控制在3.1–3.2 GB（降低65%），同时保持竞争力的准确率。
原文摘要 · Abstract (English)
Reasoning language models generate lengthy chain-of-thought (CoT) sequences whose key-value (KV) cache grows linearly and becomes a memory bottleneck during decoding. Existing compaction methods treat reasoning trajectories as flat token sequences and apply uniform compression, ignoring the hierarchical structure of CoT reasoning where different steps vary drastically in importance. We propose \textbf{Thought-Aware Attention Matching (TAM)}, which exploits this structure through three mechanisms: (i)~thought segmentation that decomposes the trajectory into reasoning blocks, (ii)~adaptive budget allocation that assigns compression budget based on each segment's importance and size, and (iii)~pivotal token protection that preserves high-attention reasoning anchors. We prove that the allocation rule is optimal under a convex error model and that cumulative error under sequential compaction remains bounded. Experiments on AIME 2024 and MATH-500 with Qwen3-4B show that TAM improves accuracy over uniform compaction at the same memory footprint, with periodic compaction bounding peak memory to 3.1--3.2\,GB (a 65\% reduction) while maintaining competitive accuracy.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。