Quais Cabeças Importam para o Raciocínio? Compressão do Cache KV Orientada por RL

Which Heads Matter for Reasoning? RL-Guided KV Cache Compression

October 9, 2025

Autores: Wenjie Du, Li Jiang, Keda Tao, Xue Liu, Huan Wang

cs.AI

Resumo

Modelos de linguagem de grande escala que realizam raciocínio exibem comportamentos complexos de raciocínio por meio da geração estendida de cadeias de pensamento (chain-of-thought), criando uma sobrecarga sem precedentes no cache de chave-valor (KV) durante a fase de decodificação. Os métodos existentes de compressão do cache KV têm desempenho inferior em modelos de raciocínio: métodos de descarte de tokens comprometem a integridade do raciocínio ao eliminar informações críticas, enquanto métodos de realocação de cabeças comprimem erroneamente as cabeças essenciais para o raciocínio, já que foram projetados para tarefas de recuperação, resultando em uma degradação significativa de desempenho à medida que as taxas de compressão aumentam. Nossa hipótese é que as cabeças KV exibem heterogeneidade funcional em modelos de raciocínio—algumas cabeças são críticas para a consistência da cadeia de pensamento, enquanto outras são compressíveis. Para validar e explorar essa percepção, propomos o RLKV, uma nova estrutura de identificação de cabeças críticas para o raciocínio, que utiliza aprendizado por reforço para otimizar diretamente a relação entre o uso do cache de cada cabeça e a qualidade do raciocínio. Como o RLKV gera recompensas a partir de amostras reais geradas durante o treinamento, ele identifica naturalmente as cabeças relevantes para os comportamentos de raciocínio. Em seguida, alocamos o cache KV completo para essas cabeças, enquanto aplicamos um cache KV comprimido e constante às demais, para uma inferência eficiente. Nossos experimentos revelam que apenas uma pequena fração das cabeças de atenção é essencial para o raciocínio, permitindo que nossa abordagem de compressão do cache supere os métodos de base, alcançando uma redução de 20-50% no cache com desempenho quase sem perdas em comparação com os resultados não comprimidos.

English

Reasoning large language models exhibit complex reasoning behaviors through the extended chain-of-thought generation, creating unprecedented Key-Value (KV) cache overhead during the decoding phase. Existing KV cache compression methods underperform on reasoning models: token-dropping methods break reasoning integrity by discarding critical information, while head-reallocating methods mistakenly compress reasoning-critical heads since they are designed for retrieval tasks, resulting in significant performance degradation as compression rates increase. We hypothesize that KV heads exhibit functional heterogeneity in reasoning models-some heads are critical for chain-of-thought consistency while others are compressible. To validate and exploit this insight, we propose RLKV, a novel reasoning-critical head identification framework, which uses reinforcement learning to directly optimize the relationship between each head's cache usage and reasoning quality. As RLKV produces rewards from actual generated samples during training, it naturally identifies heads relevant to reasoning behaviors. We then allocate full KV cache to these heads while applying compressed constant KV cache to others for efficient inference. Our experiments reveal that only a small fraction of attention heads is essential for reasoning, enabling our KV compression approach to outperform baseline methods while achieving 20-50% cache reduction with near lossless performance compared to uncompressed results.