Sparse-RL: 안정적인 희소 Rollout을 통한 LLM 강화 학습의 메모리 병목 현상 해결

Sparse-RL: Breaking the Memory Wall in LLM Reinforcement Learning via Stable Sparse Rollouts

강화 학습(RL)은 대규모 언어 모델(LLM)에서 복잡한 추론 능력을 이끌어내는 데 필수적인 기술이 되었습니다. 그러나 긴 시간 동안의 Rollout 과정에서 Key-Value (KV) 캐시를 저장하는 데 필요한 상당한 메모리 오버헤드는 중요한 병목 현상으로 작용하며, 특히 제한된 하드웨어 환경에서 효율적인 학습을 어렵게 만듭니다. 기존의 KV 압축 기술은 추론 과정에서 이러한 문제를 해결하는 데 도움이 될 수 있지만, 이를 강화 학습 훈련에 직접 적용하면 심각한 정책 불일치를 초래하여 성능이 급격히 저하될 수 있습니다. 이러한 문제를 해결하기 위해, 우리는 안정적인 희소 Rollout 환경에서 강화 학습 훈련을 가능하게 하는 Sparse-RL을 제안합니다. 우리는 불안정성이 밀집된 기존 정책, 희소 샘플링 정책, 학습 정책 간의 근본적인 정책 불일치에서 비롯된다는 것을 확인했습니다. 이 문제를 완화하기 위해, Sparse-RL은 압축으로 인한 정보 손실로 인해 발생하는 Off-policy 편향을 수정하기 위해 Sparsity-Aware Rejection Sampling과 Importance-based Reweighting을 통합합니다. 실험 결과는 Sparse-RL이 밀집된 기준 모델과 비교하여 Rollout 오버헤드를 줄이면서도 성능을 유지한다는 것을 보여줍니다. 또한, Sparse-RL은 본질적으로 희소성을 고려한 훈련을 구현하여 희소한 추론 배포 환경에서 모델의 견고성을 크게 향상시킵니다.

Reinforcement Learning (RL) has become essential for eliciting complex reasoning capabilities in Large Language Models (LLMs). However, the substantial memory overhead of storing Key-Value (KV) caches during long-horizon rollouts acts as a critical bottleneck, often prohibiting efficient training on limited hardware. While existing KV compression techniques offer a remedy for inference, directly applying them to RL training induces a severe policy mismatch, leading to catastrophic performance collapse. To address this, we introduce Sparse-RL empowers stable RL training under sparse rollouts. We show that instability arises from a fundamental policy mismatch among the dense old policy, the sparse sampler policy, and the learner policy. To mitigate this issue, Sparse-RL incorporates Sparsity-Aware Rejection Sampling and Importance-based Reweighting to correct the off-policy bias introduced by compression-induced information loss. Experimental results show that Sparse-RL reduces rollout overhead compared to dense baselines while preserving the performance. Furthermore, Sparse-RL inherently implements sparsity-aware training, significantly enhancing model robustness during sparse inference deployment.