Welke Heads zijn Belangrijk voor Redeneren? RL-Gestuurde KV Cache Compressie

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

October 9, 2025

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

cs.AI

Samenvatting

Grote taalmodellen voor redenering vertonen complexe redeneergedragingen door middel van uitgebreide keten-van-gedachten-generatie, wat een ongekende Key-Value (KV) cache-overhead veroorzaakt tijdens de decodeerfase. Bestaande methoden voor KV-cachecompressie presteren slecht op redeneermodellen: methoden die tokens weglaten, breken de redeneerintegriteit door kritieke informatie te verwijderen, terwijl methoden die heads herverdelen ten onrechte redeneerkritieke heads comprimeren omdat ze zijn ontworpen voor retrievaultaken, wat leidt tot aanzienlijke prestatievermindering naarmate de compressiepercentages toenemen. We veronderstellen dat KV-heads functionele heterogeniteit vertonen in redeneermodellen—sommige heads zijn kritiek voor de consistentie van de keten-van-gedachten, terwijl andere comprimeerbaar zijn. Om dit inzicht te valideren en te benutten, stellen we RLKV voor, een nieuw framework voor het identificeren van redeneerkritieke heads, dat reinforcement learning gebruikt om de relatie tussen het cachegebruik van elke head en de redeneerkwaliteit direct te optimaliseren. Omdat RLKV beloningen produceert op basis van daadwerkelijk gegenereerde voorbeelden tijdens de training, identificeert het van nature heads die relevant zijn voor redeneergedragingen. We wijzen vervolgens volledige KV-cache toe aan deze heads, terwijl we gecomprimeerde constante KV-cache toepassen op andere heads voor efficiënte inferentie. Onze experimenten tonen aan dat slechts een klein deel van de aandacht-heads essentieel is voor redenering, waardoor onze KV-compressieaanpak de basislijnmethoden overtreft en een cache-reductie van 20-50% bereikt met bijna verliesvrije prestaties in vergelijking met ongecomprimeerde resultaten.

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.