Research map / Large language models / Reasoning and chain of thought

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

ICML 2026 · Large language models › Reasoning and chain of thought · approach: KV cache compression · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes RLKV, which uses GRPO-based RL to learn per-head gating adapters (with L1 sparsity) that identify reasoning-critical KV heads, keeping full KV cache for those heads and compressing the rest to sink+local windows.

Problem

KV cache compression for reasoning LLMs degrades sharply because long chain-of-thought generation is fragile to information loss; existing token-dropping and retrieval-head-based methods do not identify which heads sustain reasoning.

Approach

Use RL with verifiable rewards as a probe: learnable per-head gates mix full and local attention, GRPO optimizes gates against answer correctness with an L1 sparsity penalty, and high-gate heads get full KV cache while others are compressed to constant-size cache.

Key results

• RLKV achieves near-lossless performance at 20-60% KV cache reduction across models and tasks (e.g., Llama-3.1-8B-R1 Math500 84.6 at sparsity 0.4, +1.6 over full).
• On LongReason 64K subset, RLKV at sparsity 0.4 gets 52.5 on Llama-3.1-8B-R1 vs 49.25 full, while R-KV collapses to 0.0.
• End-to-end SGLang speedup of 2.06x at 60% sparsity on Math500 with Llama-3.1-8B-R1 (accuracy 73.8 vs 79.4 full).

Surprising findings

• Token-dropping R-KV collapses to repetitive loops under compression on reasoning models, while retrieval-head-based DuoAttention mainly fails by overlength generation.
• RLKV-selected reasoning-critical heads differ from retrieval heads, and compressing them causes repetitive and incorrect errors rather than overlength errors.

Limitations

• Near-lossless threshold varies by benchmark and model; high sparsity (0.8) causes severe degradation for most methods.
• Dual-pool SGLang integration lacks radix-cache reuse and paged attention for compressed heads; measured speedup falls short of theoretical limit.
• Evaluated only on small-scale (4B-8B) reasoning models and mostly math/code tasks.

Experimental setup

Datasets: DeepScaleR GSM8K MATH500 AIME24 MBPP MMLU-Pro LongReason

Benchmarks: GSM8K MATH500 AIME24 MBPP MMLU-Pro (Chemistry, Computer Science, Law, Physics) LongReason

Metrics: accuracy (Pass@1) end-to-end throughput speedup error mode rates response length

Baselines: H2O R-KV DuoAttention KVZip

Builds on: GRPO DuoAttention R-KV H2O AReaL

Closest papers by meaning

• R-KV: Redundancy-aware KV Cache Compression for Reasoning Models (NeurIPS 2025 · similarity 0.91)
• Training Large Reasoning Models Efficiently via Progressive Thought Encoding (ICLR 2026 · similarity 0.90)
• Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning (ICLR 2025 · similarity 0.90)
• BeaconKV: Key-Value Cache Compression Guided by Beacon Queries for Efficient Large Reasoning Model Inference (ICML 2026 · similarity 0.89)
• Reasoning Path Compression: Compressing Generation Trajectories for Efficient LLM Reasoning (NeurIPS 2025 · similarity 0.89)
• ThinKV: Thought-Adaptive KV Cache Compression for Efficient Reasoning Models (ICLR 2026 · similarity 0.89)