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

Abstract: Reinforcement learning (RL) post-training for 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.

• The paper introduces SIS that converts off-policy tokens to on-policy tokens using token-level rejection sampling, significantly reducing variance.
• It integrates with various RL post-training objectives (e.g., GRPO, DAPO, GSPO) while eliminating heuristic clipping, thereby stabilizing LLM training.
• Empirical results demonstrate consistent performance gains, robust training stability, and tighter theoretical error bounds across diverse LLM architectures.

Selective Importance Sampling for Mitigating Off-Policy Drift in LLM RL Post-Training

Introduction

Reinforcement learning (RL) has established itself as the canonical approach for post-training LLMs, enabling substantial improvements in complex reasoning and multi-modal agentic behaviors. However, large-scale RL setup for LLMs typically relies on reusing off-policy data for efficiency, which induces distributional mismatch between data-generating (behavior) policies and the current policy being updated. Standard solution frameworks for off-policy correction incorporate importance sampling (IS) at sequence or token level, but these approaches suffer severe variance and instability in long-horizon settings. Hard or soft clipping of correction ratios is ubiquitous but inevitably weakens the raw learning signal by discarding the contribution of highly off-policy data.

This work introduces Selective Importance Sampling (SIS), a plug-in method that fundamentally alters the off-policy correction regime: instead of suppressing the variance by clipping, SIS utilizes token-level rejection sampling to convert as many off-policy tokens as possible into provably on-policy tokens, assigning them unit importance weight. The SIS-corrected gradient estimator thereby inherits the best properties of IS while reducing both variance and bias. SIS is computationally lightweight, easily integrates with diverse RL post-training objectives (e.g., GRPO, DAPO, GSPO), requires negligible wall-clock overhead, and displays significant empirical advantages and theoretical guarantees across dense and MoE transformer backbones.

Methodology

Off-Policy Correction as Importance Sampling

LLM RL post-training typically maximizes an expected reward objective under the current policy, but the standard "rollout-then-update" pipeline only provides samples generated from a behavior (stale) policy, yielding off-policy data. To correct for this, IS rescales each sampled trajectory or token by the ratio of current to behavior policy probabilities. However, in the token-level product form, IS ratios in long sequences compound exponentially, leading to high variance, exploding gradients, and training instabilities.

SIS: Token-Level Rejection Sampling

SIS reframes off-policy tokens as proposals from the behavior distribution and applies a per-token rejection sampling test. For each token at position $t$, the algorithm computes:

• The IS ratio $$w_t = \frac{\pi_\theta(y_t | x, y_{<t})}{\pi_{\text{old}}(y_t | x, y_{<t})}$$.
• The maximal envelope constant $$M_t = \max_{v \in V} \frac{\pi_\theta(v | x, y_{<t})}{\pi_{\text{old}}(v | x, y_{<t})}$$, restricting maximization to top-K vocabulary for efficiency.
• A Bernoulli acceptance test with probability $w_t / M_t$.

Tokens that pass are treated as exact on-policy data ($w_t=1$); rejected tokens retain the IS ratio. The method thus obviates the need for a priori clipping and directly bridges the policy distribution mismatch.

Efficient Top-K Approximation

While rejection sampling theoretically requires a max over the full vocabulary, empirical LLMs concentrate probability mass in the top-K tokens. Thus SIS approximates the envelope using only tokens in the top-K set, maintaining both high faithfulness ($K$ as low as 5–50 suffices) and amortized computational cost $O(K)$. The residual error induced by this approximation is shown to be negligible in practice.

Policy Gradient Modifications & Algorithm-Agnostic Integration

The SIS mechanism replaces only the per-token IS correction factor in any policy-gradient-based RL objective. Integration is trivial: for any baseline (e.g., GRPO, DAPO, GSPO), SIS substitutes $w_{t}$ with the modified coefficient (1 for accepted tokens), retaining all other components such as clipping, KL regularization, or group normalization.

Theoretical Analysis

A crucial theoretical contribution of this work is the characterization of the error between sequence-level and token-level gradient estimators under off-policy drift. The authors derive tight bounds on the deviation, showing that the cumulative log-importance deviation governs the approximation error. Critically, SIS strictly contracts this deviation—since it converts a large fraction of off-policy tokens to on-policy, the effective variance and mismatch are provably reduced. The acceptance-rejection test introduces no additional bias for accepted tokens, leading to both tighter convergence bounds and empirically more stable updates.

Empirical Evaluation

Benchmarks and Models

Experiments encompass a wide spectrum of settings: dense and mixture-of-experts (MoE) LLMs (Qwen3-8B, Qwen3-14B, Qwen3-30B-A3B), strong math benchmarks (MATH500, AMC23, AIME24/25), standard QA and agentic search tasks (NQ, TriviaQA, HotpotQA, Musique). SIS is further validated on Llama3.2 backbones to demonstrate architecture-agnostic efficacy.

Main Results

SIS consistently increases average accuracy on all mathematical and agentic tasks when plugged into any of the underlying RL algorithms. The strongest performance gains ($\mathbf{+6.37}$ math, $\mathbf{+3.05}$ agent) are observed on MoE models where the off-policy drift is most pronounced. SIS maintains robustness under increasing reuse of stale rollouts, a regime where vanilla baselines degrade severely. Under clipping-free conditions, SIS alone provides stabilization that outperforms clip-based methods.

Training Stability and MoE Drift

Stress testing under increased rollout reuse ($$w_t = \frac{\pi_\theta(y_t | x, y_{<t})}{\pi_{\text{old}}(y_t | x, y_{<t})}$$0 up to 16 steps) and in the presence of MoE routing divergence confirms that SIS substantially improves training stability: accuracy consistently stays above baselines at every staleness level, entropy collapse is mitigated, and gradient spikes are dampened or eliminated. When combined with router-alignment approaches (e.g., Rollout Routing Replay), SIS yields further gains in both accuracy and stability.

Token-Level Selectivity and Hyperparameter Robustness

Ablation studies on the top-K envelope hyperparameter show that SIS is highly insensitive: sweeping $$w_t = \frac{\pi_\theta(y_t | x, y_{<t})}{\pi_{\text{old}}(y_t | x, y_{<t})}$$1 from 10 to 100 alters accuracy by less than 1 point. Token-level analyses reveal selective acceptance of reasoning-relevant tokens (e.g., mathematical expressions) and rejection of off-distribution formatting/web artifacts, indicative of SIS’s effective discrimination at the distributional level.

Practical and Theoretical Implications

SIS introduces a principled mechanism for mitigating off-policy error in RL post-training of LLMs—its plug-in nature means it can be deployed with virtually all existing policy gradient RL objectives. It achieves a unique balance: reducing variance without discarding gradient contributions or introducing additional sensitivity to hyperparameters. The approach is orthogonal to, and thus can be combined with, other system-oriented improvements like router alignment or sparse rollouts.

Theoretically, SIS establishes new approximation-error bounds and formalizes the utility of token-level acceptance-rejection sampling for policy optimization under distribution drift. The method suggests that fine-grained, locality-aware IS correction can yield superior alignment and stability over legacy sequence-level or heuristic clipping approaches.

Conclusion

Selective Importance Sampling reframes off-policy correction for LLM RL post-training, shifting the focus from variance-suppressing tricks to principled, token-level acceptance-rejection conversion. SIS delivers consistent empirical gains, stronger training stability, and is provably superior in terms of reducing sequence-to-token approximation error. Its negligible compute cost and agnostic interface facilitate immediate and broad adoption. Further, SIS opens up novel directions for the theoretical study of distributional control in RL optimization for LLMs and the development of alignment algorithms that respect both statistical rigor and practical constraints.