오프라인 토큰을 온라인 토큰으로 변환하는 방법: LLM 정렬 개선을 위한 플러그인 방식

Turning Off-Policy Tokens On-Policy: A Plug-in Approach for Improving LLM Alignment

대규모 언어 모델(LLM)의 강화 학습 기반 추가 훈련은 효율적인 '탐색 후 업데이트' 패러다임을 따르며, 이는 필연적으로 오프라인 훈련 데이터를 생성합니다. 이를 해결하기 위해 중요 샘플링(IS)이 제안되었지만, 토큰 수준의 비율이 긴 시퀀스에서 누적되면서 심각한 분산 문제를 야기합니다. 본 논문에서는 이러한 오프라인 토큰을 온라인 토큰으로 '전환'하여 보정용 중요도 점수가 필요 없도록 하는 방법을 제시합니다. 이 아이디어를 바탕으로, 거부 샘플링에서 영감을 받은 선택적 중요 샘플링(SIS)을 제안합니다. 구체적으로, SIS는 오프라인 모델을 제안 분포로 보고 토큰 수준의 거부 테스트를 수행합니다. 수락된 토큰은 온라인 데이터로 간주되어 단위 중요도 점수를 받고, 거부된 토큰은 표준 IS 보정을 적용합니다. 제안하는 SIS는 이론적으로 토큰 수준과 시퀀스 수준의 오프라인 기울기 추정값 간의 격차를 줄이는 것으로 증명되었습니다. SIS는 정책 손실에서 중요 비율만 수정하는 플러그인 형태로 구현되어 무시할 만한 계산 오버헤드를 가지며, 다양한 강화 학습 기반 추가 훈련 알고리즘과 함께 사용할 수 있습니다. 수학 및 에이전트 벤치마크를 사용하여 Dense 및 MoE LLM에 대한 실험 결과, SIS는 모든 목표에서 일관되게 성능을 향상시키며, 특히 오프라인 데이터 환경에서 더 강력한 안정성을 제공합니다.

Reinforcement learning (RL) post-training for large language models (LLMs) follows a efficient paradigm of "rollout then update", which inevitably results in off-policy training data. To resolve this, Importance sampling (IS) is proposed, while the token-level ratios compound over long sequences, causing severe variance exploded. A natural idea is "transferring" these off-policy token into on-policy token, so that the importance scores for correction are unnecessary. Following this idea, we propose Selective Importance Sampling (SIS), which is inspired by rejection sampling. Concretely, SIS implements by viewing off-policy model as proposal distribution, and implement a token-level rejection test: accepted tokens are viewed as on-policy, so that receive unit importance score, while rejected tokens retain the standard IS correction. Our proposed SIS is theoretically proved reducing the gap between token-level and sequence-level off-policy gradient estimators. The SIS acts as a plug-in that only modifies the importance ratio in the policy loss, adding negligible wall-clock overhead, and can be combine with a vast vary of RL post-training algorithms. Experiments on dense and MoE LLMs across math and agent benchmarks show that SIS consistently improves all objectives, while providing substantially stronger robustness under off-policy data.