Compressão de Cache KV Consciente da Informação para Raciocínio Longo

Resumo

A capacidade de raciocínio avançou rapidamente em modelos de linguagem de grande escala (LLMs), levando a um aumento no tamanho do cache de chave-valor (KV) tanto nos estágios de pré-preenchimento quanto de decodificação. Os métodos existentes de compressão do cache KV baseiam-se principalmente em pesos de atenção para estimar a importância dos tokens. Embora a atenção capture efetivamente a relevância contextual, ela ignora sinais complementares da teoria da informação relacionados à incerteza preditiva e à informatividade dos tokens. Neste artigo, revisitamos a importância dos tokens a partir de uma perspectiva prospectiva e introduzimos a Influência Prospectiva, uma métrica que mede como tokens comprimidos afetam contextos futuros. Nossa análise revela que tokens selecionados por pontuações de atenção influenciam principalmente contextos próximos, enquanto tokens associados a alta incerteza preditiva exibem uma influência substancialmente mais forte sobre contextos futuros distantes. Com base nessa observação, propomos o InfoKV, uma estrutura de compressão de cache KV sensível à entropia que incorpora sinais da teoria da informação. Ela combina a incerteza preditiva em nível de token com a evolução da representação em camadas e integra as pontuações de entropia resultantes com as pontuações de atenção durante o raciocínio. Experimentos em benchmarks de raciocínio de contexto longo com Llama-3.1, Llama-3.2 e DeepSeek-R1 demonstram que o InfoKV supera consistentemente os métodos existentes de compressão KV baseados em atenção tanto em cenários de pré-preenchimento longo quanto de decodificação.

English

Reasoning capability has advanced rapidly in large language models (LLMs), leading to an increasing size of key-value (KV) cache in both prefilling and decoding stages. Existing KV cache compression methods mainly rely on attention weights to estimate token importance. While attention effectively captures contextual relevance, it overlooks complementary information-theoretic signals related to predictive uncertainty and token informativeness. In this paper, we revisit token importance from a forward-looking perspective and introduce Forward Influence, a metric that measures how compressed tokens affect future contexts. Our analysis reveals that tokens selected by attention scores mainly influence nearby contexts, whereas tokens associated with high predictive uncertainty exhibit substantially stronger influence on distant future contexts. Based on the observation, we propose InfoKV, an entropy-aware KV cache compression framework that incorporates information-theoretic signals. It combines token-level predictive uncertainty with layer-wise representation evolution and integrates the resulting entropy scores with attention scores during reasoning. Experiments on long-context reasoning benchmarks with Llama-3.1, Llama-3.2, and DeepSeek-R1 demonstrate that InfoKV consistently outperforms existing attention-based KV compression methods in both long prefilling and decoding scenarios.