Faça Cada Token Contar: Rumo à Melhoria do Desempenho em Contextos Longos com Evicção de Cache KV

Resumo

O cache chave-valor (KV) é um grande gargalo na inferência de contexto longo, onde a memória e o processamento crescem com o comprimento da sequência. Os métodos existentes de evicção de KV reduzem esse custo, mas, tipicamente, prejudicam o desempenho em relação à inferência com cache completo. Nossa constatação central é que a atenção com cache completo nem sempre é ótima: em contextos longos, tokens irrelevantes podem diluir a atenção, afastando-a de evidências úteis; assim, uma evicção seletiva e aprendível pode melhorar a geração, em vez de meramente aproximar o cache completo. Apresentamos um método de evicção de KV baseado em retenção global, que aprende a utilidade futura de cada token sob um orçamento unificado de memória. Portões de retenção leves atribuem pontuações de utilidade às entradas armazenadas no cache KV, e uma projeção final de pontuação compartilhada calibra essas pontuações em todas as camadas e cabeças de atenção. Isso permite uma única política global de evicção, na qual tokens de diferentes camadas, cabeças e modalidades competem diretamente pela capacidade do cache. Além disso, fornecemos uma análise teórica mostrando que reter preferencialmente tokens úteis reduz a diluição da atenção, e justificamos a retenção geométrica como uma proxy independente de consulta para utilidade futura. Em diversos benchmarks de linguagem de contexto longo, raciocínio visão-linguagem e diálogo de múltiplas rodadas, nosso método reduz substancialmente a memória KV, igualando ou superando a inferência com cache completo. Esses resultados sugerem que a evicção de KV aprendida e globalmente calibrada não é apenas uma técnica de compressão, mas também um mecanismo para melhorar o raciocínio em contextos longos.

English

The key-value (KV) cache is a major bottleneck in long-context inference, where memory and computation grow with sequence length. Existing KV eviction methods reduce this cost but typically degrade performance relative to full-cache inference. Our key insight is that full-cache attention is not always optimal: in long contexts, irrelevant tokens can dilute attention away from useful evidence, so selective, learnable eviction can improve generation rather than merely approximate the full cache. We introduce a global retention-based KV eviction method that learns each token's future utility under a unified memory budget. Lightweight retention gates assign utility scores to cached KV entries, and a shared final scoring projection calibrates these scores across all layers and heads. This enables a single global eviction policy in which tokens from different layers, heads, and modalities compete directly for cache capacity. We further provide theoretical analysis showing that preferentially retaining useful tokens reduces attention dilution, and we justify geometric retention as a query-agnostic proxy for future utility. Across diverse long-context language and vision-language reasoning, and multi-turn dialogue benchmarks, our method substantially reduces KV memory while matching or surpassing full-cache inference. These results suggest that learned, globally calibrated KV eviction is not only a compression technique, but also a mechanism for improving long-context reasoning.