How to Compress KV Cache in RL Post-Training? Shadow Mask Distillation for Memory-Efficient Alignment
cs.LG, cs.AI
Submitted: 2026-05-07
Updated: 2026-09-16
Comments: The authors are withdrawing this paper because errors were identified in the experimental results
License: http://creativecommons.org/licenses/by/4.0/
The gist: Reinforcement Learning (RL) has emerged as a crucial paradigm for unlocking the advanced reasoning capabilities of Large Language Models (LLMs), encompassing frameworks like RLHF and RLAIF.
Terminology
Abstract
Reinforcement Learning (RL) has emerged as a crucial paradigm for unlocking the advanced reasoning capabilities of Large Language Models (LLMs), encompassing frameworks like RLHF and RLAIF. Regardless of the specific optimization algorithm (e.g., PPO, GRPO, or Online DPO), online RL inherently requires an exploratory trajectory generation (rollout) phase. However, for long-context reasoning tasks, this rollout phase imposes a severe ``memory wall'' due to the exorbitant Key-Value (KV) cache footprint. While applying KV cache compression during rollouts mitigates this memory overhead, it induces a critical off-policy bias. Although modern KV compression is often nearly lossless during standard inference, even minuscule approximation errors are drastically amplified by the inherent instability of RL optimization. Specifically, the sampler generates responses under a sparse context, whereas the learner updates parameters using the full, dense context. Existing statistical solutions, such as importance reweighting, struggle to correct this magnified bias, suffering from high gradient variance and severe sample inefficiency.
Sources
- Llama 2: Open Foundation and Fine-Tuned Chat Models
- Qwen Technical Report
- Hey, That's My Data! Token-Only Dataset Inference in Large Language Models
- Training a Helpful and Harmless Assistant with Reinforcement Learning from Human Feedback
- Adversarial Attack-Defense Co-Evolution for LLM Safety Alignment via Tree-Group Dual-Aware Search and Optimization
- DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models
- LongRoPE: Extending LLM Context Window Beyond 2 Million Tokens
- Fast Transformer Decoding: One Write-Head is All You Need
- SnapKV: LLM Knows What You are Looking for Before Generation
- Parity Quantum Computing as YZ-Plane Measurement-Based Quantum Computing
- SGLang: Efficient Execution of Structured Language Model Programs
- Proximal Policy Optimization Algorithms
- KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache
- Near-Lossless Model Compression Enables Longer Context Inference in DNA Large Language Models
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
- ByzSFL: Achieving Byzantine-Robust Secure Federated Learning with Zero-Knowledge Proofs
- Efficient Attentions for Long Document Summarization
- Training Verifiers to Solve Math Word Problems
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks