Quali Teste Sono Importanti per il Ragionamento? Compressione della Cache KV Guidata da RL

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

October 9, 2025

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

cs.AI

Abstract

I modelli linguistici di grandi dimensioni dimostrano comportamenti di ragionamento complessi attraverso la generazione estesa di catene di pensiero, creando un sovraccarico senza precedenti della cache Key-Value (KV) durante la fase di decodifica. I metodi esistenti di compressione della cache KV si rivelano inefficaci sui modelli di ragionamento: i metodi di eliminazione dei token compromettono l'integrità del ragionamento scartando informazioni critiche, mentre i metodi di riallocazione delle head comprimono erroneamente le head cruciali per il ragionamento poiché sono progettati per task di recupero, portando a un significativo degrado delle prestazioni all'aumentare dei tassi di compressione. Ipotesizziamo che le head KV mostrino un'eterogeneità funzionale nei modelli di ragionamento: alcune head sono critiche per la coerenza della catena di pensiero, mentre altre sono comprimibili. Per validare e sfruttare questa intuizione, proponiamo RLKV, un nuovo framework di identificazione delle head critiche per il ragionamento, che utilizza l'apprendimento per rinforzo per ottimizzare direttamente la relazione tra l'uso della cache di ciascuna head e la qualità del ragionamento. Poiché RLKV genera ricompense da campioni effettivamente generati durante l'addestramento, identifica naturalmente le head rilevanti per i comportamenti di ragionamento. Assegniamo quindi la cache KV completa a queste head, applicando una cache KV costante e compressa alle altre per un'inferenza efficiente. I nostri esperimenti rivelano che solo una piccola frazione delle head di attenzione è essenziale per il ragionamento, consentendo al nostro approccio di compressione KV di superare i metodi di base e di ottenere una riduzione della cache del 20-50% con prestazioni quasi senza perdite rispetto ai risultati non compressi.

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.