提出阴影掩码蒸馏，高效压缩推理记忆，解决强化学习长文本生成的内存瓶颈。
How to Compress KV Cache in RL Post-Training? Shadow Mask Distillation for Memory-Efficient Alignment
- 用阴影掩码蒸馏技术，实现低误差的KV缓存压缩
- 显著降低长上下文推理时的内存占用，提升训练效率
- 适合需要高效长文本对齐的RL微调场景
强化学习（RL）已成为释放大语言模型高级推理能力的关键范式，涵盖如RLHF和RLAIF等框架。无论采用何种优化算法（如PPO、GRPO或在线DPO），在线强化学习必然包含探索性轨迹生成（滚动）阶段。然而，在长上下文推理任务中，这一阶段因巨大的键值（KV）缓存占用而面临严峻的“内存墙”问题。尽管在滚动过程中应用KV缓存压缩可缓解内存开销，但会引入关键的偏离策略偏差。尽管现代KV压缩在标准推理中近乎无损，但即使微小的近似误差也会被强化学习优化的内在不稳定性大幅放大。具体而言，采样器在稀疏上下文中生成响应，而学习器则使用完整密集上下文进行参数更新。现有统计解决方案（如重要性重加权）难以校正这种放大的偏差，存在梯度方差高和样本效率差的问题。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。