论文
如何在RL 后训练中压缩KV缓存？用于内存高效对齐的 Shadow Mask 蒸馏
How to Compress KV Cache in RL Post-Training? Shadow Mask Distillation for Memory-Efficient Alignment
摘要
强化学习 (RL) 已成为解锁 大语言模型 (LLM) 高级推理功能的关键范例，涵盖 RLHF 和 RLAIF 等框架。无论具体的优化算法（例如，PPO、GRPO 或 Online DPO）如何，在线 RL 本质上都需要探索性轨迹生成（采样轨迹）阶段。然而，对于长上下文推理任务，由于过高的键值 (KV) 缓存占用空间，此采样轨迹阶段会造成严重的“内存墙”。虽然在采样轨迹期间应用 KV 缓存压缩可以减轻这种内存开销，但它会导致严重的 离策略 偏差。尽管现代 KV 压缩在标准推理过程中通常几乎是无损的，但即使是极小的近似误差也会因 RL 优化的固有不稳定性而急剧放大。具体来说，采样器在稀疏上下文下生成响应，而学习器使用完整、密集的上下文更新参数。现有的统计解决方案（例如重要性重新加权）难以纠正这种放大的偏差，受到高梯度方差和严重样本效率低下的影响。