¿Qué cabezas importan para el razonamiento? Compresión de la caché KV guiada 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

Resumen

Los modelos de lenguaje de gran escala que realizan razonamiento exhiben comportamientos de razonamiento complejos a través de la generación extendida de cadenas de pensamiento, lo que crea una sobrecarga sin precedentes en la caché de clave-valor (KV) durante la fase de decodificación. Los métodos existentes de compresión de la caché KV tienen un rendimiento inferior en modelos de razonamiento: los métodos que descartan tokens rompen la integridad del razonamiento al eliminar información crítica, mientras que los métodos que reasignan cabezas comprimen erróneamente las cabezas críticas para el razonamiento, ya que están diseñados para tareas de recuperación, lo que resulta en una degradación significativa del rendimiento a medida que aumentan las tasas de compresión. Nuestra hipótesis es que las cabezas KV exhiben heterogeneidad funcional en los modelos de razonamiento: algunas cabezas son críticas para la consistencia de la cadena de pensamiento, mientras que otras son comprimibles. Para validar y explotar esta idea, proponemos RLKV, un marco novedoso de identificación de cabezas críticas para el razonamiento, que utiliza aprendizaje por refuerzo para optimizar directamente la relación entre el uso de la caché de cada cabeza y la calidad del razonamiento. Como RLKV genera recompensas a partir de muestras generadas durante el entrenamiento, identifica naturalmente las cabezas relevantes para los comportamientos de razonamiento. Luego, asignamos la caché KV completa a estas cabezas mientras aplicamos una caché KV comprimida constante a las demás para una inferencia eficiente. Nuestros experimentos revelan que solo una pequeña fracción de las cabezas de atención es esencial para el razonamiento, lo que permite que nuestro enfoque de compresión KV supere a los métodos base mientras logra una reducción del 20-50% en la caché con un rendimiento casi sin pérdidas en comparación con los resultados sin compresión.

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.