推理时超比例扩缩:基于 KV 缓存压缩的推理时间放大
机器学习
2025-11-10 v2 计算与语言
摘要
推理时间放大通过生成更长或更并行的序列序列,以牺牲效率为代价来提高推理准确性。然而,在 Transformer LLM 中,生成成本受限于键值(KV)缓存的大小,而非生成标记的数量。因此,我们探索推理时间的超比例放大:通过压缩 KV 缓存,可在相同的计算预算内生成更多标记,从而进一步提高放大推理的准确性。然而,这一方法的成功取决于压缩方法在高压缩比率下仍能保持准确性的能力。为使超比例放大实用,我们引入动态内存稀疏化(DMS)——一种一种稀疏化 KV 缓存的方法,只需 1K 个训练步骤即可实现 8 倍压缩,同时比训练免费稀疏注意力在保持准确性方面表现更好。DMS 不会过早丢弃已缓存的标记,而是延迟标记驱逐,隐式地合并表示并保留关键信息。我们在多个 LLM 系列上展示了基于 DMS 的推理时间超比例放大有效性,表明它可在相同的推理延迟和内存负载下提升准确性。例如,我们在 AIME 24 上提升 12.0 分,在 GPQA 上提升 8.6 分,在 LiveCodeBench 上提升 9.7 分(平均值),且等效的内存读取次数下即可实现。
引用
@article{arxiv.2506.05345,
title = {Inference-Time Hyper-Scaling with KV Cache Compression},
author = {Adrian Łańcucki and Konrad Staniszewski and Piotr Nawrot and Edoardo M. Ponti},
journal= {arXiv preprint arXiv:2506.05345},
year = {2025}
}
备注
Accepted to NeurIPS 2025