AIS: Adaptive Importance Sampling for Quantized RL

• URL: http://arxiv.org/abs/2605.13907v1
• Date: Wed, 13 May 2026 03:36:57 GMT
• Title: AIS: Adaptive Importance Sampling for Quantized RL
• Abstract summary: Reinforcement learning for large language models (LLMs) is dominated by the cost of rollout generation.<n>This introduces a rollout-training mismatch that biases the policy gradient and can cause training to collapse outright on reasoning benchmarks.<n>We show that the mismatch is nonstationary and acts as a double-edged sword: early in training it provides an exploration bonus, exposing the gradient to trajectories the trainer would otherwise under-sample.<n>We propose Adaptive Importance Sampling (AIS), a correction framework that adjusts the strength of its intervention on a per-batch basis.
• Score: 21.387834718338496
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Reinforcement learning (RL) for large language models (LLMs) is dominated by the cost of rollout generation, which has motivated the use of low-precision rollouts (e.g., FP8) paired with a BF16 trainer to improve throughput and reduce memory pressure. This introduces a rollout-training mismatch that biases the policy gradient and can cause training to collapse outright on reasoning benchmarks. We show that the mismatch is non-stationary and acts as a double-edged sword: early in training it provides a stochastic exploration bonus, exposing the gradient to trajectories the trainer would otherwise under-sample, but the same perturbation transitions into a destabilizing source of bias as the policy concentrates. To solve this, we propose Adaptive Importance Sampling (AIS), a correction framework that adjusts the strength of its intervention on a per-batch basis. AIS combines three real-time diagnostics, namely weight reliability, divergence severity, and variance amplification, into a single mixing coefficient that interpolates between the uncorrected and fully importance-weighted gradients, suppressing the destabilizing component of the mismatch while preserving its exploratory benefit. We integrate AIS into GRPO and evaluate it on the diffusion-based LLaDA-8B-Instruct and the autoregressive Qwen3-8B and Qwen3.5-9B across mathematical reasoning and planning benchmarks. AIS matches the BF16 baseline on most tasks while retaining the 1.5 to 2.76x rollout speedup of FP8.

Related papers

• DISA: Offline Importance Sampling for Distribution-Matching LLM-RL [56.9445657766829]
  We introduce DISA, which moves this calibration problem outside the RL loop.<n> DISA draws proposal trajectories offline, estimates the partition function via importance sampling, and freezes the resulting partition-function estimate.<n>On two open-weight backbones across six math and three code benchmarks, DISA matches or exceeds the online-coupled distribution-matching baseline Flow.
  arXiv  Detail & Related papers  (2026-05-17T07:14:44Z)
• Too Correct to Learn: Reinforcement Learning on Saturated Reasoning Data [55.84428098924793]
  We propose Constrained Uniform Top-K Sampling (CUTS), a parameter-free decoding strategy enforcing structure-preserving exploration.<n>We integrate this into Mixed-CUTS, a training framework synergizing exploitative and exploratory rollouts to amplify intra-group advantage variance.<n> Notably, Mixed-CUTS improves Pass@1 accuracy on the challenging AIME25 benchmark by up to 15.1% over standard GRPO.
  arXiv  Detail & Related papers  (2026-04-20T16:43:28Z)
• Bridging SFT and RL: Dynamic Policy Optimization for Robust Reasoning [10.4404175363648]
  Post-training paradigms for Large Language Models (LLMs) face a fundamental dilemma.<n>SFT provides stability (low variance) but suffers from high fitting bias, while RL enables exploration (low bias) but grapples with high gradient variance.<n>Existing unified optimization strategies often employ naive loss weighting, overlooking the statistical conflict between these distinct gradient signals.<n>We propose textbfDYPO (Dynamic Policy Optimization), a unified framework designed to structurally mitigate this conflict.
  arXiv  Detail & Related papers  (2026-04-10T03:42:16Z)
• Improving Search Agent with One Line of Code [68.58667107354253]
  Tool-based Agentic Reinforcement Learning (TARL) has emerged as a promising paradigm for training search agents to interact with external tools.<n>We propose textbfSearch textbfAgent textbfPolicy textbfOptimization (textbfSAPO), which stabilizes training via a conditional token-level KL constraint.
  arXiv  Detail & Related papers  (2026-03-10T04:07:39Z)
• Controllable Exploration in Hybrid-Policy RLVR for Multi-Modal Reasoning [88.42566960813438]
  CalibRL is a hybrid-policy RLVR framework that supports controllable exploration with expert guidance.<n>CalibRL increases policy entropy in a guided manner and clarifies the target distribution.<n>Experiments across eight benchmarks, including both in-domain and out-of-domain settings, demonstrate consistent improvements.
  arXiv  Detail & Related papers  (2026-02-22T07:23:36Z)
• VESPO: Variational Sequence-Level Soft Policy Optimization for Stable Off-Policy LLM Training [18.849117699859622]
  Training stability is a central challenge in reinforcement learning for large language models.<n>We propose Variational sEquence-level Soft Policy Optimization (VESPO)<n> Experiments on mathematical reasoning benchmarks show that VESPO maintains stable training under staleness ratios up to 64x and fully asynchronous execution.
  arXiv  Detail & Related papers  (2026-02-11T09:48:08Z)
• Information Fidelity in Tool-Using LLM Agents: A Martingale Analysis of the Model Context Protocol [69.11739400975445]
  We introduce the first theoretical framework for analyzing error accumulation in Model Context Protocol (MCP) agents.<n>We show that cumulative distortion exhibits linear growth and high-probability deviations bounded by $O(sqrtT)$.<n>Key findings include: semantic weighting reduces distortion by 80%, and periodic re-grounding approximately every 9 steps suffices for error control.
  arXiv  Detail & Related papers  (2026-02-10T21:08:53Z)
• Beyond Precision: Training-Inference Mismatch is an Optimization Problem and Simple LR Scheduling Fixes It [24.70923739848818]
  We show that gradient noise and training-inference mismatch escalate in tandem as training progresses.<n>We find that the mismatch can be effectively suppressed by shrinking the update size.<n>We propose a simple yet effective solution: a specialized Learning Rate scheduler.
  arXiv  Detail & Related papers  (2026-02-02T09:00:53Z)
• Enhancing the Outcome Reward-based RL Training of MLLMs with Self-Consistency Sampling [90.87033586963828]
  Outcome-reward reinforcement learning (RL) is a common and increasingly significant way to refine the step-by-step reasoning of multimodal large language models (MLLMs)<n>We propose Self-Consistency Sampling (SCS) to correct this issue.<n>Based on Qwen2.5-VL-7B-Instruct, SCS improves accuracy by up to 7.7 percentage points on six multimodal benchmarks with negligible extra computation.
  arXiv  Detail & Related papers  (2025-11-13T18:59:57Z)
• Think Right: Learning to Mitigate Under-Over Thinking via Adaptive, Attentive Compression [68.69801176669843]
  We propose an online post-training RL method that prunes redundant steps and estimates difficulty.<n> TRAAC (Think Right with Adaptive, Attentive Compression) achieves an average absolute accuracy gain of 8.4%.<n>Although our models are trained on math datasets, they show accuracy and efficiency gains on out-of-distribution non-math datasets.
  arXiv  Detail & Related papers  (2025-10-02T02:00:20Z)
• Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models [35.36024202299119]
  We introduce textbfAdvantage Weighted Matching (AWM), a policy-gradient method for diffusion.<n>AWM raises the influence of high-reward samples and suppresses low-reward ones while keeping the modeling objective identical to pretraining.<n>This unifies pretraining and RL conceptually and practically, is consistent with policy-gradient theory, reduces variance, and yields faster convergence.
  arXiv  Detail & Related papers  (2025-09-29T17:02:20Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.