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

2026-07-06

AI Summary

This work addresses the off-policy data challenge in reinforcement learning (RL) post-training of large language models, where the prevailing “generate-then-update” paradigm induces high variance in traditional importance sampling over long sequences, degrading alignment performance. The authors propose Selective Importance Sampling (SIS), the first method to enable token-level off-policy to on-policy conversion: leveraging a rejection-sampling–inspired mechanism, SIS applies importance weight correction only to rejected tokens while treating the rest as approximately on-policy samples. Theoretically, SIS substantially reduces the bias gap between token-level and sequence-level gradient estimates. Designed as a plug-and-play module, SIS seamlessly integrates with diverse RL post-training algorithms and consistently improves all optimization objectives across dense and Mixture-of-Experts (MoE) architectures on mathematical and agent-based benchmarks, while significantly enhancing robustness to off-policy data.

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.

Problem

• off-policy
• importance sampling
• variance explosion
• reinforcement learning
• large language models

Innovation

• Selective Importance Sampling
• off-policy correction
• rejection sampling
• LLM alignment
• reinforcement learning post-training