推論時のハイパースケーリングとKVキャッシュ圧縮
Inference-Time Hyper-Scaling with KV Cache Compression
June 5, 2025
著者: Adrian Łańcucki, Konrad Staniszewski, Piotr Nawrot, Edoardo M. Ponti
cs.AI
要旨
推論時のスケーリングは、より長いまたはより並列なシーケンスを生成することで効率を犠牲にして推論精度を向上させます。しかし、Transformer LLMにおいて、生成コストは生成されるトークン数ではなく、キー・バリュー(KV)キャッシュのサイズによってボトルネックとなります。そこで、我々は推論時のハイパースケーリングを探求します。KVキャッシュを圧縮することで、同じ計算予算内でより多くのトークンを生成し、スケーリングされた推論の精度をさらに向上させることができます。ただし、このアプローチの成功は、高い圧縮率でも精度を維持できる圧縮手法の能力にかかっています。ハイパースケーリングを実用的にするために、我々はDynamic Memory Sparsification(DMS)を導入します。これは、わずか1Kのトレーニングステップで8倍の圧縮を達成し、トレーニング不要のスパースアテンションよりも優れた精度を維持する、KVキャッシュをスパース化する新しい手法です。DMSは、キャッシュされたトークンを早期に破棄するのではなく、トークンの削除を遅らせ、暗黙的に表現を統合し、重要な情報を保持します。我々は、DMSを用いた推論時のハイパースケーリングの有効性を複数のLLMファミリーで実証し、同等の推論実行時間とメモリ負荷で精度を向上させることを示します。例えば、Qwen-R1 32Bにおいて、AIME 24で平均9.1ポイント、GPQAで7.6ポイント、LiveCodeBenchで9.6ポイントの精度向上を達成しました。
English
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 8times 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 runtime and memory load. For
instance, we enhance Qwen-R1 32B by an average of 9.1 points on AIME 24, 7.6 on
GPQA, and 9.6 on LiveCodeBench across compute budgets.