PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2510.08525v3cs.CL

TL;DR

Reasoning models’ extended chain-of-thought is fragile to KV-cache information loss, and existing compression methods do not identify heads that preserve reasoning consistency. RLKV uses reinforcement learning to find those heads and guide cache allocation. It achieves 20–60% KV-cache reduction with near-lossless performance across diverse tasks and models.

  • Problem

    Extended chain-of-thought makes reasoning models’ KV caches vulnerable to information loss, while existing methods do not identify heads that maintain reasoning consistency and generation flow.

    Source: §1 · pp. 1–2

  • Method

    RLKV uses reinforcement learning on sampled reasoning generations to identify reasoning-critical heads and allocate full cache to them while compressing others.

    Source: §1 · p. 2

  • Results

    20–60% KV-cache reduction achieves near-lossless performance across diverse tasks and models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Analyses characterize reasoning-critical heads as primary carriers of chain-of-thought consistency and empirically distinct from retrieval heads.

    Source: §6 · p. 8

    Performance degradation may vary by downstream task and deployment setting, so real-world use should be carefully evaluated.

    Source: Impact Statement · p. 9

Abstract

from arXiv · show

Reasoning large language models exhibit complex reasoning behaviors via extended chain-of-thought generation that are highly fragile to information loss during decoding, creating critical challenges for KV cache compression. Existing token-dropping methods directly disrupt reasoning chains by removing intermediate steps, while head-reallocation methods, designed for retrieval tasks, fail to preserve the heads essential for generative reasoning. However, no existing method can identify which attention heads genuinely maintain reasoning consistency and control generation termination. To address this, we propose RLKV, which uses reinforcement learning as a probe to discover which heads contribute to reasoning quality by directly optimizing their cache usage against actual generation outcomes. This discovery naturally leads to an efficient compression strategy: we allocate full KV cache to reasoning-critical heads while aggressively compressing others with constant-size KV cache. Experiments reveal that a fraction of heads proves essential for reasoning, enabling 20--60% cache reduction with near-lossless performance across diverse tasks and models, and up to 2.06x end-to-end speedup at 60% reduction.

1. Introduction

Extended chain-of-thought makes reasoning models’ KV caches both larger and more fragile to information loss. RLKV uses reinforcement learning to identify reasoning-critical heads and allocate full cache to them while compressing others.

1. Introduction

  • Extended chain-of-thought generation makes KV cache compression fragile, and methods that work well for instruct models can severely degrade reasoning scenarios.Source: §1 · p. 1
  • R-KV and DuoAttention degrade on reasoning models despite stable instruct-model performance; their errors include repetitive loops and over-extended reasoning that misses solutions.The comparison attributes the degradation largely to reasoning models’ longer generations; retrieval-oriented head identification does not capture how compression affects actual reasoning outcomes.Source: §1 · pp. 1–2, Figure 1 · p. 2
  • RLKV samples reasoning generations, rewards their quality, and uses reinforcement learning with sparse per-head gating adapters to identify heads needing full cache.It then allocates full KV cache to those reasoning-critical heads and compressed constant KV cache to others.Source: §1 · p. 2
  • 20–60% KV cache reduction delivers near-lossless reasoning across diverse tasks and models, with up to 2.06× end-to-end speedup under SGLang continuous batching.Source: §1 · p. 2

2. Related Work

Related work reduces KV-cache overhead through architectural, system-level, and compression approaches, but existing compression methods do not capture reasoning behaviors that emerge during dynamically extending chain-of-thought generation.

Efficient LLM Inference

  • Architecture and system optimizations reduce KV-cache overhead through sharing, quantization, reuse, or storage changes, while sparse attention often still requires full KV-cache storage.These approaches are described as not exploiting inherent head-level sparsity patterns.Source: §2 · pp. 2–3

KV Cache Compression

  • Existing token-dropping and head-reallocation methods target long-context compression, but their eviction strategies or retrieval-head selection cannot capture reasoning behaviors during dynamically extending CoT.The cited learning objectives are synthetic recall-task output deviation for DuoAttention and next-token loss on long-context pretraining corpora for PruLong.Source: §2 · p. 3

Reinforcement Learning for Efficiency

  • Prior RL work includes architecture search and pruning, while reasoning-model RL has also targeted shorter CoT; RLKV instead uses gating values to identify reasoning-critical heads for cache compression.The paper describes this head-gating approach as orthogonal to methods that reduce CoT length.Source: §2 · p. 3

3. Methodology

RLKV identifies reasoning-critical KV heads by optimizing per-head gating against reasoning rewards, then assigns full or compressed KV access based on the learned scores. It uses data selection and adaptive regularization to stabilize this process as sparsity increases.

3. Methodology

  • RLKV defines reasoning-critical heads as those whose reasoning degrades under local KV access, and uses them to guide cache compression.The method applies RL with sparsity pressure and adds stabilization techniques to address the conflict between dense regularization and sparse rewards.Source: §3 · p. 3

3.1. Mixed Attention with Gating Adapters

  • Per-head gating adapters mix full and local attention, reducing the trainable search space to L × H parameters while the LLM weights remain frozen.The local attention mask retains constant initial sink tokens and recent tokens for numeric stability.Source: §3.1 · p. 4

3.2. RL for Reasoning Head Identification

  • Under RLVR, reward signals preserve high α values for reasoning-critical heads requiring full KV-cache access, while an L1 penalty drives compressible heads toward zero; RLKV removes the KL penalty.Qwen-2.5-7B-R1 has inherent limitations in achievable sparsification without compromising reasoning behavior because its KV group size is 7, versus 4 in the other models.Source: §3.2 · p. 4, Figure 3 · p. 4

3.3. Stabilization for RL Training

  • When adapters become sparse, reasoning performance can degrade and rewards become unstable, while the dense L1 penalty can drive adapters further toward zero without recovery.RLKV filters to problems the model initially solves correctly and curates 3k examples using curriculum sampling based on output-token length; adaptive penalty weighting decays regularization near mean reward 0.7 and removes it below a cutoff.Source: Figure 4 · p. 4, §3.3 · p. 5
  • At inference, RLKV discretizes gating scores into a top-k selection: selected heads retain full KV access, while remaining heads use compressed access retaining sink and recent tokens.The selection follows the target compression ratio.Source: §3.3 · p. 5

4. Experiments

RLKV outperforms compression baselines across reasoning and knowledge benchmarks, with near-lossless performance at 20–60% KV cache reduction. The experiments also examine long-context transfer, serving efficiency, training components, and compressed-window choices.

4.1. Experimental Settings

  • The evaluation covers three reasoning models, four reasoning benchmarks, four MMLU-Pro subsets, and LongReason, with dynamic budgets for H2O and R-KV.Training uses 3,000 mathematical problems; evaluation uses 16 sink and 64 local tokens for compressed KV caches.Source: §4.1 · p. 5

4.2. Main Results

  • Up to 20%, RLKV outperforms baselines across four reasoning benchmarks and four MMLU-Pro subsets, sometimes surpassing the full KV cache baseline.Figure 5 compares results at sparsity levels from 0.2 to 0.8.Source: §4.2 · p. 6, Figure 5 · p. 6
  • 20–60% KV cache reduction yields near-lossless performance across tasks and models.Table 1 reports accuracy, its delta, and sparsity at each task- and model-dependent threshold.Source: §4.2 · p. 6, Table 1 · p. 6
  • On LongReason, RLKV outperforms all baselines on both evaluated models, with a widening gap at higher sparsity and transfer from 8K training rollouts to 70K context.H2O is omitted for out-of-memory, and Qwen-2.5-7B-R1 is excluded because its full-attention accuracy is 0.5%.Source: §4.2 · p. 6, Table 2 · p. 6

4.3. Inference Efficiency

  • SGLang uses a paged cache for full-attention heads and a fixed-size circular sink-plus-local buffer for compressed heads, requiring O(max running requests × W) compressed-pool tokens.The reported speedup evaluation uses Math500 with continuous batching.Source: §4.3 · p. 7, Table 3 · p. 7

4.4. Ablation Studies

  • Adaptive penalty weighting prevents sparse rewards and dense L1 penalties from driving training toward collapse; self-distillation sampling stabilizes rewards, and β = 0.001 balances sparsity with reward strength.The ablations use Qwen-2.5-7B-R1 on Math500.Source: §4.4 · p. 7, Figure 6 · p. 7
  • At high sparsity, larger sink-plus-local windows extend the lossless range and soften degradation; the authors recommend modestly enlarging the window for high-sparsity targets.The window ablation uses Llama-3.1-8B-R1 on Math500.Source: §4.4 · pp. 7–8, Table 4 · p. 7

5. Qualitative Analyses

The analyses distinguish reasoning-critical heads from retrieval and random heads, and examine the different errors and response-length costs associated with compressing them.

Head sensitivity

  • Compressing reasoning-critical heads causes a sharper accuracy drop than compressing retrieval or random heads in two models; in Qwen-3-4B-Thinking, preserving them supports better high-sparsity accuracy.The analysis compares head sensitivity across Llama-3.1-8B-R1, Qwen-2.5-7B-R1, and Qwen-3-4B-Thinking.Source: §5 · p. 8

Error modes

  • Compressing retrieval heads primarily causes overlength errors, while compressing reasoning-critical heads triggers repetitive and incorrect errors, with overlength errors rare in the latter case.The error analysis uses instances solved correctly with full KV cache but failed under compression.Source: §5 · p. 8, Figure 8 · p. 8

Average response length

  • RLKV maintains strong accuracy with competitive response lengths, while DuoAttention often needs longer reasoning steps and token-dropping methods can appear shorter by solving only easier instances.Average length is measured on samples answered correctly under both full KV cache and each compressed setting.Source: §5 · p. 8, Figure 9 · p. 8

6. Conclusion

RLKV uses reinforcement learning to identify reasoning-critical heads and guide KV cache compression by directly optimizing each head’s cache usage against reasoning quality.

6. Conclusion

  • RLKV directly optimizes each head’s cache usage against reasoning quality through reinforcement learning to identify reasoning-critical heads for compression.The conclusion reports 20–60% KV cache reduction with near-lossless accuracy and up to 2.06× end-to-end speedup under SGLang continuous batching.Source: §6 · p. 8

Impact Statement

RLKV’s impact statement notes that its efficiency gains may come with performance degradation that varies across tasks and deployment settings. In the motivation study, compression methods are relatively stable on instruct models but degrade on reasoning models, with distinct error patterns across methods.

Impact Statement

  • RLKV’s efficiency approach may still degrade performance, with effects varying by task and deployment setting.The impact statement says these effects should be evaluated carefully in real-world systems.Source: Impact Statement · p. 9

A. Motivation Study

  • Compression methods remain relatively stable on instruct models but drop substantially on reasoning models as compression increases.The motivation study evaluates H2O, R-KV, and DuoAttention across three reasoning models, their instruct variants, and four benchmarks.Source: A. Motivation Study · p. 12
  • The study categorizes failures as repetitive, incorrect, or overlength errors.These mean repeated token sequences, wrong final answers, or generations exceeding normal length baselines, respectively.Source: A. Motivation Study · p. 12
  • Token-dropping methods show repetitive errors, while DuoAttention tends toward over-length errors on reasoning models.The reported patterns suggest token-dropping disrupts reasoning-critical information, whereas head reallocation better preserves sequence integrity but may not preserve reasoning capability.Source: Figure 11 · p. 13

B. Experiment Details

Experiments use filtered, length-stratified reasoning solutions for training and greedy decoding for evaluation. The appendix also details implementation choices, baseline adjustments, training costs, and constraints of the SGLang integration.

B. Experiment Details

  • Training data comprises 3,000 sampled problems from DeepScaleR, selected from correct greedy-decoded solutions and stratified by output length.The source dataset contains about 40,000 problems; sampled examples are divided across four token-length bins from 0–2k through 6k–8k.Source: B. Experiment Details · p. 13
  • Training runs for two epochs and 185 steps on two A100 GPUs, using AdamW for gating adapters and four GRPO samples per query.The reported configuration uses batch size 32, learning rate 0.01, and sampling temperature 1.0.Source: B. Experiment Details · p. 13
  • At inference, reasoning-critical heads retain full KV cache, while other heads keep the first 16 sink tokens and 64 recent local tokens.Training uses block-sparse attention approximation; sample generation uses masks for prefilling and custom Triton kernels for decoding in SGLang.Source: B. Experiment Details · pp. 13–14
  • SGLang v0.5.2 lacks radix-cache reuse and paged attention for the compressed-head buffer, and generation is not strictly deterministic.The authors state that both RLKV and the full-attention baseline share the configuration, so these constraints do not affect the reported relative speedup; absolute accuracy has minor engine variance.Source: B. Experiment Details · p. 14
  • H2O and R-KV receive the same sink and local-token overhead as RLKV, and their fixed budgets are converted to dynamic allocations.DuoAttention is replicated with its default settings on the evaluated models and the same inference settings as RLKV.Source: B. Experiment Details · p. 14
  • Adapter training costs 40, 22, and 36 GPU-hours for Llama-3.1-8B-R1, Qwen-2.5-7B-R1, and Qwen-3-4B-Thinking, respectively.Source: B. Experiment Details · p. 14
  • Evaluations use greedy decoding and Pass@1, with 8,192-token limits for GSM8K, Math500, MBPP, and MMLU-Pro, and 16,384 for AIME24.Math and code benchmarks use HuggingFace Transformers, while MMLU-Pro subsets use the SGLang integration.Source: B. Experiment Details · p. 14

C. Complete numerical results

Tables 6–8 provide complete numerical results for RLKV and baselines across three reasoning models, benchmarks, and KV cache budgets.

C. Complete numerical results

  • Tables 6–8 report results across benchmarks and KV cache budgets for three reasoning models, with parenthetical values showing differences from full KV cache.Positive differences indicate improvement and negative differences indicate degradation.Source: C. Complete numerical results · p. 15

D. Details of Error Modes Analyses

The appendix reports that its error-mode findings generally match the main experiments, with a stated exception for one model and setting.

D. Details of Error Modes Analyses

E. An Implicitly Unfair Comparison in Fixed-Budget Evaluation

Fixed budgets can impose very different compression ratios across reasoning samples, making reported compression optimistic. The authors evaluate with per-sample dynamic budgets and find that this changes R-KV’s comparison with H2O at most tested sparsity levels.

E. An Implicitly Unfair Comparison in Fixed-Budget Evaluation

  • Reasoning responses can greatly exceed prompt lengths, so a global fixed budget leaves shorter outputs uncompressed while compressing longer ones.This produces different per-sample compression ratios.Source: E. An Implicitly Unfair Comparison in… · p. 16
  • R-KV reports its compression rate as budget divided by average full length; on Math500 with Llama-3.1-8B-R1, that ratio is 66.2% at a 200-token budget and 3019 average full length.Many samples remain uncompressed, making the reported ratio optimistic.Source: E. An Implicitly Unfair Comparison in… · p. 16

F. Comparison of Fixed Budget and Dynamic Budget for R-KV and H2O

G. Limitations and Future Work

RLKV’s reasoning-critical head distributions vary across models and tasks, and its static gating leaves finer functional categories unresolved. The authors identify high-compression degradation and dynamic gating overhead as open issues, while describing extensions beyond KV cache compression and to multimodal reasoning as future directions.

G. Limitations and Future Work

  • RLKV suffers substantial performance degradation at sparsity ≥0.8; further compression without quality loss may require combining head reallocation with methods such as KV-cache quantization.The authors present this as a limitation at extremely high compression ratios.Source: G. Limitations and Future Work · p. 17
  • Reasoning-critical head distributions vary across models and tasks, and RLKV does not decompose this static pattern into finer functional categories.The distribution reflects heterogeneous reasoning mechanisms.Source: G. Limitations and Future Work · p. 17
  • RLKV fixes gating scores at training time; query-adaptive gating may be more flexible, but its decision overhead versus compression gains remains unresolved.Source: G. Limitations and Future Work · p. 17
  • The authors propose extending RLKV’s on-policy formulation to other efficiency methods and tasks, including multimodal and agentic reasoning.They describe these applications as future directions.Source: G. Limitations and Future Work · pp. 17–18

Related papers

Loading 2510.08525v3…