Гипермасштабирование во время вывода с компрессией кэша ключей и значений

Inference-Time Hyper-Scaling with KV Cache Compression

June 5, 2025

Авторы: Adrian Łańcucki, Konrad Staniszewski, Piotr Nawrot, Edoardo M. Ponti

cs.AI

Аннотация

Масштабирование на этапе вывода жертвует эффективностью ради повышения точности рассуждений за счет генерации более длинных или более параллельных последовательностей. Однако в трансформерных больших языковых моделях (LLM) стоимость генерации ограничивается размером кэша ключей-значений (KV), а не количеством сгенерированных токенов. Поэтому мы исследуем гипермасштабирование на этапе вывода: сжатие кэша KV позволяет генерировать больше токенов в рамках того же вычислительного бюджета и дополнительно повысить точность масштабированного вывода. Успех этого подхода, однако, зависит от способности методов сжатия сохранять точность даже при высоких коэффициентах сжатия. Чтобы сделать гипермасштабирование практичным, мы представляем Dynamic Memory Sparsification (DMS) — новый метод разрежения кэшей KV, который требует всего 1K шагов обучения для достижения 8-кратного сжатия, сохраняя при этом более высокую точность по сравнению с разреженным вниманием без обучения. Вместо преждевременного удаления кэшированных токенов DMS откладывает их вытеснение, неявно объединяя представления и сохраняя критически важную информацию. Мы демонстрируем эффективность гипермасштабирования на этапе вывода с использованием DMS на нескольких семействах LLM, показывая, что оно повышает точность при сопоставимом времени выполнения вывода и нагрузке на память. Например, мы улучшаем Qwen-R1 32B в среднем на 9.1 балла на AIME 24, на 7.6 на GPQA и на 9.6 на LiveCodeBench для различных вычислительных бюджетов.

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.