用强化学习动态淘汰无关缓存,让大模型长文本推理更省内存。
DistillCache: KL-Guided Adaptive KV-Cache Eviction for Memory-Efficient LLM Inference
- 基于KL散度奖励的强化学习框架,自适应选择保留哪些缓存。
- 7B模型在25%缓存下保持94.2%准确率,比现有方法高2.7个百分点。
- 适合追求高效长文本推理的开发者,尤其在资源受限场景。
基于Transformer的大语言模型在多项任务中表现优异,但其键值(KV)缓存随序列长度线性增长,严重制约长文本推理的内存使用。现有启发式淘汰方法(如H$_2$O和SnapKV)依赖静态注意力或位置信号,难以捕捉词元的未来预测影响。本文提出DistillCache,一个将KV缓存淘汰建模为序列决策问题的强化学习框架。该方法利用模型内部丰富信号(注意力统计、值范数、熵与位置),通过REINFORCE算法训练轻量级策略网络,并以每步的KL散度作为奖励,以保留全缓存输出分布。在70亿参数的指令微调Transformer(Mistral-7B-Instruct-v0.3)上,DistillCache在LongBench上以25%缓存预算维持94.2%的全缓存准确率,优于强启发式基线(H$_2$O、SnapKV)最高2.7个百分点;在重实现条件下,也超越同期基于强化学习的方法(ForesightKV、RLKV)最多1.4个百分点。在推理基准测试中,其性能媲美最优并发方法,且在激进压缩下更优。同时,相比全缓存,吞吐量最高提升2.1倍,兼具实用效率。结果表明,学习得到的、分布感知的策略在内存高效的长上下文推理中极具潜力。
原文摘要 · Abstract (English)
Transformer-based large language models (LLMs) achieve strong performance across many tasks, but their Key-Value (KV) cache grows linearly with sequence length, creating a severe memory bottleneck for long-context inference. Existing heuristic eviction methods (e.g., H$_2$O and SnapKV) rely on static attention or positional signals that often fail to capture a token's future predictive influence. We propose DistillCache, a reinforcement learning framework that formulates KV-cache eviction as a sequential decision problem. DistillCache learns a lightweight policy network using rich internal model signals (attention statistics, value norms, entropy, and position) and trains it with REINFORCE via a per-step KL-divergence reward to preserve the full-cache output distribution. On a 7B-parameter instruction-tuned Transformer (Mistral-7B-Instruct-v0.3), DistillCache retains 94.2% of full-cache accuracy on LongBench at a 25% cache budget, outperforming both strong heuristic baselines (H$_2$O, SnapKV) by up to 2.7 absolute points and, under our re-implementations, concurrent RL-based methods (ForesightKV, RLKV) by up to 1.4 points on long-context tasks. On reasoning benchmarks, DistillCache is competitive with the best concurrent method and surpasses it under aggressive compression. It also delivers up to 2.1x full-cache throughput while maintaining competitive practical efficiency. These results highlight the effectiveness of learned, distribution-aware policies for memory-efficient long-context LLM inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。