Quelles têtes comptent pour le raisonnement ? Compression du cache KV guidée par RL

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

October 9, 2025

papers.authors: Wenjie Du, Li Jiang, Keda Tao, Xue Liu, Huan Wang

cs.AI

papers.abstract

Les grands modèles de langage dotés de capacités de raisonnement manifestent des comportements de raisonnement complexes grâce à la génération étendue de chaînes de pensée, créant ainsi une surcharge sans précédent du cache clé-valeur (KV) lors de la phase de décodage. Les méthodes existantes de compression du cache KV sous-performent sur les modèles de raisonnement : les méthodes d'élimination de tokens compromettent l'intégrité du raisonnement en supprimant des informations critiques, tandis que les méthodes de réallocation de têtes compriment par erreur les têtes essentielles au raisonnement, car elles sont conçues pour des tâches de recherche, entraînant une dégradation significative des performances à mesure que les taux de compression augmentent. Nous émettons l'hypothèse que les têtes KV présentent une hétérogénéité fonctionnelle dans les modèles de raisonnement—certaines têtes sont cruciales pour la cohérence de la chaîne de pensée, tandis que d'autres sont compressibles. Pour valider et exploiter cette intuition, nous proposons RLKV, un nouveau cadre d'identification des têtes critiques pour le raisonnement, qui utilise l'apprentissage par renforcement pour optimiser directement la relation entre l'utilisation du cache de chaque tête et la qualité du raisonnement. Comme RLKV génère des récompenses à partir d'échantillons réels produits pendant l'entraînement, il identifie naturellement les têtes pertinentes pour les comportements de raisonnement. Nous allouons ensuite un cache KV complet à ces têtes tout en appliquant un cache KV compressé constant aux autres pour une inférence efficace. Nos expériences révèlent que seule une petite fraction des têtes d'attention est essentielle pour le raisonnement, permettant à notre approche de compression KV de surpasser les méthodes de référence tout en réalisant une réduction de 20 à 50 % du cache avec une performance quasi sans perte par rapport aux résultats non compressés.

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.