测试压缩键值缓存对长推理任务的影响,发现追踪重要信息的策略更有效。
Hold Onto That Thought: Assessing KV Cache Compression On Reasoning
- 通过追踪高频关键信息,提升推理时缓存压缩效率。
- 在低缓存预算下可生成更长的推理序列,但增加推理开销。
- 适用于需要多步思考的数学推理任务,如GSM8K、MATH500。
大语言模型在长上下文任务中表现优异,但常受内存限制制约。其中,用于加速注意力计算的键值缓存(KV cache)随上下文长度线性增长。已有多种压缩算法通过剔除不重要标记来缓解缓存膨胀,但多数方法针对预填充阶段(如处理长提示),其在需长时间解码的推理任务中表现鲜有评估。尤其在如GSM8K和MATH500等复杂短提示任务中,模型依赖多步推理与自我反思,形成数千token的思维链。本文首次系统评测多种主流压缩策略在长推理任务上的表现。对于非推理模型Llama-3.1-8B-Instruct,发现无单一策略普适,性能受数据集类型显著影响。但在推理模型上,H2O及我们改进的SnapKV变体表现最优,表明对关键信息的追踪对推理轨迹至关重要。此外,在低缓存预算下,某些淘汰策略反而能生成更长推理链,揭示缓存大小与推理成本间的权衡关系。
原文摘要 · Abstract (English)
Large language models (LLMs) have demonstrated remarkable performance on long-context tasks, but are often bottlenecked by memory constraints. Namely, the KV cache, which is used to significantly speed up attention computations, grows linearly with context length. A suite of compression algorithms has been introduced to alleviate cache growth by evicting unimportant tokens. However, several popular strategies are targeted towards the prefill phase, i.e., processing long prompt context, and their performance is rarely assessed on reasoning tasks requiring long decoding. In particular, short but complex prompts, such as those in benchmarks like GSM8K and MATH500, often benefit from multi-step reasoning and self-reflection, resulting in thinking sequences thousands of tokens long. In this work, we benchmark the performance of several popular compression strategies on long-reasoning tasks. For the non-reasoning Llama-3.1-8B-Instruct, we determine that no singular strategy fits all, and that performance is heavily influenced by dataset type. However, we discover that H2O and our decoding-enabled variant of SnapKV are dominant strategies for reasoning models, indicating the utility of heavy-hitter tracking for reasoning traces. We also find that eviction strategies at low budgets can produce longer reasoning traces, revealing a tradeoff between cache size and inference costs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。