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

Jul 16, 2026

Upload by Aria

Reinforcement learning post-training for large language models follows a rollout-then-update workflow, producing unavoidable off-policy training data. Importance sampling corrects this mismatch but becomes unstable as token-level correction ratios compound, leading to exploding variance. Selective Importance Sampling (SIS) transfers off-policy tokens into on-policy tokens using a rejection-sampling inspired token acceptance test, assigning unit weight to accepted tokens while retaining standard IS for rejected ones. SIS is plug-in only modifying importance ratios, with negligible overhead, and is theoretically shown to reduce token- and sequence-level off-policy gradient mismatch. Experiments on dense and MoE LLMs across math and agent benchmarks show consistent objective gains and improved robustness under off-policy regimes.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages23

Reading Time58 min

Abstract

Introduction

Off-policy training challenge in RL post-training

Importance sampling and variance/clipping limitations

SIS method and token-level transfer idea

Contributions

Background

Off-Policy Policy Gradient

arXiv :2607 .04728v 1 [ cs .CL] 6 Jul 2026 TURNING OFF-POLICY TOKENS ON-POLICY: A PLUGIN APPROACH FOR IMPROVING LLM ALIGNMENT Yu Li 1 ,∗ , Xiuyu Li 1 ,∗ ,‡, Mingyang Yi 1 ,†, Jiaxing Wang2 , zhangliangxu2 , Zhaolong Xing2 , Zhen Chen2 [1](1 Renmin University of China 2 JD.com)[ Renmin University of China](1 Renmin University of China 2 JD.com)[ 2](1 Renmin University of China 2 JD.com)[ JD.com](1 Renmin University of China 2 JD.com)[ ](1 Renmin University of China 2 JD.com){liyu0929,[yimingyang](yimingyang}@ruc.edu.cn)[}](yimingyang}@ruc.edu.cn)[@ruc.edu.cn](yimingyang}@ruc.edu.cn) ABSTRACT Reinforcement learning (RL) post-training for large language models (LLMs) follows a efficient paradigm of “rollout then update”, which inevitably results in offpolicy 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 avast 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. 1 INTRODUCTION Reinforcement learning (RL) (Sutton et al., 1998) has become the dominant paradigm for posttraining large language models, driving substantial progress in reasoning (Guo et al., 2025; Chenet al., 2025), code generation (Jiang et al., 2025), tool use (Li et al., 2026b), and broader capabilities (Li et al., 2026a;c; Zhang et al., 2026a;b; Tu et al., 2026) . The RL process follows an order of “rollout then update”. However, owing to the considerations of efficiencies in sample reuse (Noukhovitchet al., 2025), asynchronous auto-regressive inference (Fu et al., 2026), researchers inevitably use off-policy data from stale policies to update model. The resulting distribution mismatch biases gradient estimates and degrades training stability (Ma et al., 2025), making effective off-policy gradient correction a fundamental challenge for RL-based post-training methods. Importance sampling (IS) (Tokdar & Kass, 2010) offers a principled correction to off-policy gradient by multiplying a correction ratio, but its variance can grow rapidly in long-horizon reasoning (Metelli et al., 2020), and is sensitive to some outliers (Schulman et al., 2017) . To mitigate this, mainstream methods (Schulman et al., 2017; Shao et al., 2024; Yu et al., 2025) rely on hard clipping. Although effective for stabilization, hard clipping suppresses gradients from highly off-policy samples and loses useful learning signal. Some recent methods (Zheng et al., 2025b; Chen et al., 2025; Gao et al., 2025a) mitigate this issue by soft clipping, but remain heuristic. Instead of working on clipping method, we address this problem by “transferring” the off-policy tokens into on-policy tokens, which fundamentally resolve the off-policy problem. Our method, Selec- ∗Equal contribution. †Corresponding author. ‡[Work done during an internship at JD.com](Work done during an internship at JD.com). tive Importance Sampling (SIS), a simple and effective method conduct this transfer via acceptancerejection sampling (Robert &

FAQs

What evidence supports SIS effectiveness and practicality?

SIS provides consistent gains across dense and MoE LLMs, improves robustness under challenging off-policy regimes, and acts as a plug-in that only modifies the importance ratio with negligible wall-clock overhead, with ablations showing low sensitivity to hyperparameters.

Jul 16, 2026

Upload by Aria