Evicção Estocástica de Cache KV Ciente do Valor para Modelos de Raciocínio

Resumo

Modelos de raciocínio melhoram a precisão por meio de cadeias de pensamento estendidas, mas suas saídas longas criam um gargalo de memória e computação. Métodos de expulsão de cache KV reduzem esse custo ao remover pares chave-valor irrelevantes do cache, porém frequentemente produzem precisão inferior em comparação com alternativas de atenção esparsa baseadas em seleção, que mantêm o cache KV completo. Identificamos fatores-chave cruciais para a precisão da expulsão de cache KV. Primeiro, uma pequena fração dos estados de valor possui magnitudes anormalmente grandes, e sua expulsão causa falhas catastróficas, nas quais os modelos entram em loops repetitivos de raciocínio. Segundo, a introdução de estocasticidade durante a expulsão melhora a precisão ao aumentar a diversidade do cache. Com base nessas descobertas, propomos a Expulsão Estocástica de Cache KV Consciente de Valores (VaSE), uma receita livre de treinamento que protege estados de valor de grande magnitude e promove decisões de expulsão diversas. Em seis tarefas de raciocínio, modelos Qwen3 que utilizam VaSE com compressão de cache KV de 4x apresentam precisões médias mais altas do que o método de seleção de última geração na mesma esparsidade, superando o método de expulsão mais forte em mais de 4%. No geral, VaSE preenche a lacuna entre eficiência e precisão, suportando FlashAttention2 e permitindo uma pegada de memória estática para modelos de raciocínio.

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.