基于KV缓存压缩的推理时超扩展
Inference-Time Hyper-Scaling with KV Cache Compression
AI总结：
本文提出动态记忆稀疏化DMS方法，通过高压缩比KV缓存压缩实现推理时超扩展，在相同计算预算下提升LLM推理准确率。
摘要：
推理时扩展通过生成更长或更多并行序列，以效率换取更高的推理准确率。然而，在Transformer大语言模型中，生成成本主要受键值（KV）缓存大小的制约，而非生成token的数量。因此，我们探索推理时超扩展：通过压缩KV缓存，可以在相同计算预算内生成更多token，从而进一步提高扩展推理的准确率。不过，该方法的成功关键在于压缩方法能否在高压缩比下仍保持准确率。为使超扩展切实可行，我们提出动态记忆稀疏化（Dynamic Memory Sparsification，DMS），这是一种新颖的KV缓存稀疏化方法，仅需1K训练步即可实现8倍压缩，同时保持优于免训练稀疏注意力的准确率。DMS不提前丢弃已缓存token，而是延迟token淘汰，隐式合并表示并保留关键信息。我们在多个LLM系列上验证了基于DMS的推理时超扩展的有效性，结果表明在相近推理延迟和内存负载下，该方法能提升准确率。例如，在相同内存读取次数下，我们将Qwen-R1 32B在AIME 24上平均提升12.0分，在GPQA上提升8.6分，在LiveCodeBench上提升9.7分。
英文摘要：
Inference-time scaling trades efficiency for increased reasoning accuracy by generating longer or more parallel sequences. However, in Transformer LLMs, generation cost is bottlenecked by the size of the key-value (KV) cache, rather than the number of generated tokens. Hence, we explore inference-time hyper-scaling: by compressing the KV cache, we can generate more tokens within the same compute budget and further improve the accuracy of scaled inference. The success of this approach, however, hinges on the ability of compression methods to preserve accuracy even at high compression ratios. To make hyper-scaling practical, we introduce Dynamic Memory Sparsification (DMS), a novel method for sparsifying KV caches that only requires 1K training steps to achieve 8$\times$ compression, while maintaining better accuracy than training-free sparse attention. Instead of prematurely discarding cached tokens, DMS delays token eviction, implicitly merging representations and preserving critical information. We demonstrate the effectiveness of inference-time hyper-scaling with DMS on multiple families of LLMs, showing that it boosts accuracy for comparable inference latency and memory load. For instance, we enhance Qwen-R1 32B by 12.0 points on AIME 24, 8.6 on GPQA, and 9.7 on LiveCodeBench on average for an equivalent number of memory reads.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
