Welche Köpfe sind für das Schlussfolgern entscheidend? RL-gesteuerte KV-Cache-Komprimierung

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

Große Sprachmodelle, die für logisches Denken ausgelegt sind, zeigen komplexe Denkprozesse durch die erweiterte Generierung von Gedankenketten (Chain-of-Thought), was während der Dekodierungsphase einen beispiellosen Overhead im Key-Value (KV)-Cache verursacht. Bestehende Methoden zur KV-Cache-Kompression schneiden bei solchen Denkmodellen schlecht ab: Token-Verwerfungsmethoden beeinträchtigen die Integrität des Denkprozesses, indem sie kritische Informationen entfernen, während Methoden zur Neuverteilung von Attention-Heads versehentlich für das Denken entscheidende Heads komprimieren, da sie für Abfrageaufgaben konzipiert sind. Dies führt zu einer erheblichen Leistungsverschlechterung bei steigenden Kompressionsraten. Wir stellen die Hypothese auf, dass KV-Heads in Denkmodellen eine funktionale Heterogenität aufweisen – einige Heads sind für die Konsistenz der Gedankenkette entscheidend, während andere komprimierbar sind. Um diese Erkenntnis zu validieren und zu nutzen, schlagen wir RLKV vor, ein neuartiges Framework zur Identifizierung von für das Denken kritischen Heads, das Reinforcement Learning einsetzt, um den Zusammenhang zwischen der Cache-Nutzung jedes Heads und der Denkqualität direkt zu optimieren. Da RLKV während des Trainings Belohnungen aus tatsächlich generierten Beispielen ableitet, identifiziert es auf natürliche Weise Heads, die für Denkprozesse relevant sind. Wir weisen diesen Heads dann einen vollständigen KV-Cache zu, während wir für andere Heads einen komprimierten, konstanten KV-Cache verwenden, um eine effiziente Inferenz zu ermöglichen. Unsere Experimente zeigen, dass nur ein kleiner Bruchteil der Attention-Heads für das Denken essenziell ist, wodurch unser KV-Kompressionsansatz Baseline-Methoden übertrifft und eine Reduzierung des Caches um 20–50 % bei nahezu verlustfreier Leistung im Vergleich zu unkomprimierten Ergebnissen erreicht.

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.