Стохастическое вытеснение KV-кэша с учетом ценности для моделей рассуждения

Аннотация

Модели рассуждений повышают точность за счет удлиненных цепочек размышлений, однако их длинные выходные данные создают узкое место по памяти и вычислениям. Методы вытеснения KV-кэша снижают эти затраты путем удаления неважных пар ключ-значение из кэша, но часто дают меньшую точность по сравнению с альтернативами на основе разреженного внимания с выборкой, которые сохраняют полный KV-кэш. Мы выявили ключевые факторы, определяющие точность вытеснения KV-кэша. Во-первых, небольшая доля состояний значений имеет аномально большие величины, и их вытеснение приводит к катастрофическим сбоям, когда модели входят в повторяющиеся циклы рассуждений. Во-вторых, введение стохастичности в процессе вытеснения улучшает точность за счет увеличения разнообразия кэша. Основываясь на этих выводах, мы предлагаем Value-aware Stochastic KV Cache Eviction (VaSE) — метод без обучения, который защищает состояния значений с большой величиной и способствует разнообразию решений о вытеснении. На шести задачах рассуждений модели Qwen3, использующие VaSE с 4-кратным сжатием KV-кэша, достигают более высокой средней точности, чем современные методы выборки при той же разреженности, превосходя самый сильный метод вытеснения более чем на 4%. В целом, VaSE устраняет разрыв между эффективностью и точностью, поддерживая FlashAttention2 и обеспечивая статический след памяти для моделей рассуждений.

English

Reasoning models improve accuracy through extended chains of thought, but their long outputs create a memory and compute bottleneck. KV cache eviction methods reduce this cost by evicting unimportant key-value pairs from the cache, yet they often yield worse accuracy than selection-based sparse attention alternatives, which keep the full KV cache. We identify key factors crucial to KV cache eviction accuracy. First, a small fraction of value states have abnormally large magnitudes, and evicting them causes catastrophic failure where models enter repetitive reasoning loops. Second, introducing stochasticity during eviction improves accuracy by increasing cache diversity. Based on these findings, we propose Value-aware Stochastic KV Cache Eviction (VaSE), a training-free recipe that protects large-magnitude value states and promotes diverse eviction decisions. Across six reasoning tasks, Qwen3 models using VaSE with 4x KV cache compression yield higher average accuracies than SOTA selection method at the same sparsity, while outperforming the strongest eviction method by more than 4%. Overall, VaSE bridges the gap between efficiency and accuracy, supporting FlashAttention2 and enabling a static memory footprint for reasoning models.