Turning Off-Policy Tokens On-Policy: A Plug-in Approach for Improving LLM Alignment
2026-07-06 • Computation and Language
Computation and LanguageArtificial IntelligenceMachine Learning
AI summaryⓘ
The authors address a challenge in reinforcement learning for large language models, where training uses off-policy data causing high variance in importance sampling. They propose Selective Importance Sampling (SIS), inspired by rejection sampling, to treat some tokens as if they came from the current policy to reduce variance. SIS modifies how importance scores are calculated during policy updates, improving training stability with minimal extra compute. Experiments show SIS enhances performance and robustness across different models and tasks.
reinforcement learninglarge language modelsoff-policy dataimportance samplingrejection samplingpolicy gradienttoken-level correctionvariance reductionpost-training
Authors
Yu Li, Xiuyu Li, Mingyang Yi, Jiaxing Wang, zhangliangxu, Zhaolong Xing, Zhen Chen
Abstract
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.