🔥 今日值得一读 长思维链内存砍到四分之一,准确率不降反平!
AbridgeKV: Compressing Long Reasoning into the KV Cache of Learned Pause Tokens
arXiv CL (cs.CL) 🔥 重点 #KV缓存#推理优化#长上下文#训练方法 🕐 今天 12:00
👨‍💼 主理人解读 · 为什么值得关注
长思维链推理显存吃紧时,可让模型把中间推理压进暂停token的KV,降低缓存增长。

📖 AI 总结

AbridgeKV 针对长思维链解码中的内存瓶颈提出了一种新方案。由于 KV 缓存随测试时计算量线性增长,而直接驱逐缓存条目会导致模型生成更长甚至达到长度上限的输出,作者转而训练模型压缩自身输出:将每个已完成的生成文本块压缩为若干学习到的“暂停令牌”的 KV 条目,这些令牌不产生输出,随后释放原文本块的缓存,后续令牌改为关注暂停令牌的 KV 条目。该方法的三个关键设计是:暂停令牌的 KV 条目位于其所替换文本块内部的位置、按各层条目实测价值分配缓存预算、以及在较大文本块上执行压缩。在 AIME、HMMT、AMC 等六个竞赛的 446 道题、32K 令牌生成任务上,仅保留四分之一生成缓存即可在准确率和输出长度上匹配同等训练的全 KV 模型;同等预算下最强已发表驱逐方法输出远长于它,最接近的训练式压缩器在同数据重训后损失 10 分。在 vLLM 中部署时,每序列缓存缩小 3.5 倍,饱和 H100 完成一道题的速度提升 3.2 倍。

🔑 关键词速览

KV Cache键值缓存,在自回归生成中存储已生成 token 的键和值向量,以避免重复计算,但会随序列长度线性增长。
Pause Tokens暂停 token,一种特殊的输入 token,具有可训练的嵌入但不产生输出,用于在模型中引入计算暂停或压缩信息。
Chain-of-Thought (CoT)思维链,一种提示技术,让模型生成中间推理步骤,以提高复杂推理任务的性能。
Eviction逐出,指从缓存中移除某些条目以释放内存的策略,但可能导致模型性能下降。
vLLM一个高效的大语言模型推理和服务引擎,支持分页注意力等优化技术。
infoAI 公众号二维码

📱 每天一份 AI 前沿日报

关注公众号,每天 09:00 推送 · 不错过任何重磅