arXiv:2512.03324cs.LGcs.AI2025-12被引 15

通过学习令牌重要性,动态保留关键信息,提升长文本推理效率。

Cache What Lasts: Token Retention for Memory-Bounded KV Cache in LLMs

  • 用轻量保留门实时计算每个令牌的重要性得分
  • 在低内存下仍优于传统淘汰策略,部分场景超越全缓存模型
  • 可解释性强,自动发现人类直觉的高效记忆机制

由于自注意力的二次复杂度和不断增长的键值(KV)缓存,内存与计算仍是长序列大模型推理的核心瓶颈。现有内存受限推理策略如量化、卸载或启发式缓存淘汰,要么协调成本高,要么依赖不可靠的重要度代理。本文提出TRIM-KV,一种新方法:在生成时通过轻量级保留门学习每个令牌的内在重要性。每个门输出一个随时间衰减的标量保留分数,反映该令牌在特定层和头中的长期价值。当内存超限,低分令牌被剔除,确保缓存始终包含最关键内容。TRIM-KV通过从冻结大模型蒸馏并结合容量损失进行高效训练,仅需微调门控参数,推理开销几乎不变。在数学推理(GSM8K、MATH-500、AIME24)、程序生成(LongProc)、对话长记忆(LongMemEval)以及长上下文理解(LongBenchV2 和 SCBench)任务中,TRIM-KV持续优于强基线,尤其在低内存条件下表现突出。令人惊讶的是,它在某些场景甚至超越全缓存模型,表明选择性保留可作为正则化手段,抑制无意义令牌带来的噪声。定性分析显示,学习到的保留分数与人类直觉一致,无需显式设计即可自然恢复诸如锚点令牌、滑动窗口和摘要压缩等启发式策略。除了效率提升,保留分数还揭示了各层与头的具体功能,为大模型可解释性开辟新路径。

原文摘要 · Abstract (English)

Memory and computation remain core bottlenecks in long-horizon LLM inference due to the quadratic cost of self-attention and the ever-growing key-value (KV) cache. Existing strategies for memory-bounded inference, such as quantization, offloading, or heuristic KV eviction, either incur high orchestration costs or rely on unreliable attention-based proxies of importance. We propose TRIM-KV, a novel approach that learns each token's intrinsic importance at creation time via a lightweight retention gate. Each gate predicts a scalar retention score that decays over time, reflecting the long-term utility of the token for a specific layer and head. Tokens with low scores are evicted when the memory budget is exceeded, ensuring that the cache always contains the most critical tokens. TRIM-KV is trained efficiently through distillation from a frozen LLM combined with a capacity loss, requiring only gate fine-tuning and adding negligible inference overhead. Across mathematical reasoning (GSM8K, MATH-500, AIME24), procedural generation (LongProc), conversational long-memory benchmarks (LongMemEval), and long-context understanding (LongBenchV2 and SCBench), TRIM-KV consistently outperforms strong eviction and learnable retrieval baselines, especially in low-memory regimes. Remarkably, it even surpasses full-cache models in some settings, showing that selective retention can serve as a form of regularization, suppressing noise from uninformative tokens. Qualitative analyses further reveal that learned retention scores align with human intuition, naturally recovering heuristics such as sink tokens, sliding windows, and gist compression without explicit design. Beyond efficiency, retention scores provide insights into layer- and head-specific roles, suggesting a new path toward LLM interpretability.

大模型推理缓存优化可解释性注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。