Technical Summary: Adaptive Importance Sampling for Quantized RL (AIS)

1. Problem Statement

Reinforcement Learning (RL) for Large Language Models (LLMs) is currently bottlenecked by the computational cost of trajectory rollout generation. To mitigate this, recent industry practices (e.g., DeepSeek V3.2, veRL) employ low-precision quantization (e.g., FP8) for the rollout engine while maintaining a higher-precision (BF16) trainer. This creates a rollout–training mismatch: the quantized rollout policy ($\hat{\pi}_{\theta_{old}}$) generates trajectories that differ distributionally from the full-precision learner policy ($\pi_{\theta_{old}}$) used to compute gradients.

The paper identifies this mismatch as non-stationary and acting as a "double-edged sword":

• Early Training: The quantization noise acts as a stochastic exploration bonus, exposing the gradient to trajectories the trainer might otherwise under-sample.
• Late Training: As the policy sharpens, the same noise becomes a destabilizing source of bias, causing training collapse on reasoning benchmarks (e.g., a 6.63% drop in AIME25 accuracy for Qwen3-8B).

Standard Importance Sampling (IS) attempts to correct this bias but often fails due to prohibitively high variance in long sequences. Conversely, Truncated Importance Sampling (TIS) stabilizes training but uses a fixed correction strength that cannot adapt to the evolving nature of the mismatch, leading to either over-correction (suppressing useful exploration) or under-correction (allowing instability).

2. Methodology: Adaptive Importance Sampling (AIS)

The authors propose Adaptive Importance Sampling (AIS), a framework that dynamically adjusts the strength of the importance sampling correction on a per-batch basis. Instead of a fixed threshold, AIS computes a single mixing coefficient, $\alpha(x) \in [0, 1]$, which interpolates between the uncorrected gradient and the fully importance-weighted gradient.

2.1 The Mixing Coefficient

The core of AIS is a factorized gating mechanism that combines three real-time diagnostics derived from statistics already computed during the policy update:

1. Weight Reliability ($\alpha_{ess}$): Measures the Effective Sample Size (ESS) ratio of the truncated importance weights. If the weight distribution is concentrated (low ESS), the correction is deemed unreliable and downweighted.
   $$\alpha_{ess} = (1 + CV^2(\bar{w}))^{-1/2}$$
2. Divergence Severity ($\alpha_{mis}$): Measures the mean absolute log-probability discrepancy between the rollout and training policies. Correction is only engaged when this mismatch exceeds a saturation threshold $\delta$.
   $$\alpha_{mis} = \min\left(1, \frac{\bar{d}}{\delta}\right)$$
3. Variance Amplification ($\alpha_{var}$): Monitors whether large importance weights coincide with large-magnitude advantages, which would inflate gradient variance beyond acceptable limits. If the variance amplification ratio $\Delta\sigma$ exceeds a tolerance $\gamma$, a penalty is applied.
   $$\alpha_{var} = \max\left(0, \frac{\Delta\sigma - \gamma}{\gamma}\right)$$

The final mixing coefficient is computed as:
$$\alpha = \text{clip}(\alpha_{ess} - \beta \cdot \alpha_{var}, 0, 1) \cdot \alpha_{mis}$$

2.2 Integration with GRPO

AIS is integrated into Group Relative Policy Optimization (GRPO). The method introduces a rectification weight $\tilde{w}_i$ that adjusts the advantage term in the GRPO objective:
$$\tilde{w}_i = 1 + \alpha \cdot (\min(\rho_{rt}, C) - 1)$$
where $\rho_{rt}$ is the rollout-training ratio and $C$ is the truncation threshold. This formulation ensures that when the mismatch is negligible ($\alpha \to 0$), the objective reduces exactly to vanilla GRPO, maintaining backward compatibility.

3. Key Contributions

• Characterization of Non-Stationarity: The paper establishes that the rollout-training mismatch is not uniformly harmful; it provides an exploration bonus early in training but becomes a bias source later. This challenges the view of quantization mismatch solely as an adversary.
• Adaptive Correction Framework: AIS is the first method to dynamically modulate the strength of importance sampling correction based on real-time diagnostics (reliability, divergence, and variance), rather than using a static threshold.
• Diffusion Model Application: This work presents the first study of low-precision rollout for RL on diffusion-based LLMs (specifically LLaDA-8B-Instruct), demonstrating that the mismatch issue and the AIS solution generalize beyond autoregressive architectures.
• Efficiency-Accuracy Trade-off: The method achieves a balance where it recovers the accuracy of full-precision BF16 training while retaining the throughput and memory benefits of FP8 rollouts.

4. Experimental Results

The authors evaluated AIS on Qwen3-8B, Qwen3.5-9B (autoregressive), and LLaDA-8B-Instruct (diffusion) across mathematical reasoning (GSM8K, MATH500, AIME25, AMC, Olympiad) and planning (Sudoku, Countdown) benchmarks.

• Accuracy Recovery:
  • Autoregressive Models: On Qwen3-8B, naive FP8 rollout dropped AIME25 accuracy from 36.70% (BF16) to 26.81%. AIS recovered this to 43.33%, surpassing the BF16 baseline. Similar improvements were seen on Qwen3.5-9B, where AIS matched or exceeded BF16 on 11 of 12 benchmark pairs.
  • Diffusion Models: On LLaDA-8B, FP8 rollout caused significant degradation (e.g., -19.00% on Countdown at 512 tokens). AIS recovered performance to match or exceed BF16 levels across various sequence lengths and fine-tuning regimes (Full FT and LoRA).
• General Capabilities: AIS preserved performance on general benchmarks (MMLU, PIQA, ARC, etc.), indicating no overfitting to the RL reward signal.
• Efficiency Gains:
  • Speedup: FP8 rollout provided 1.5× to 2.76× speedup in end-to-end rollout time, depending on batch size and model architecture (higher gains observed in diffusion models due to tensor-core utilization).
  • Memory: FP8 rollout halved the rollout model footprint (e.g., ~7.95 GB vs. ~15.9 GB for LLaDA-8B), enabling larger batch sizes that were previously out-of-memory (OOM) under BF16.

5. Significance and Claims

The paper claims that AIS is the first study of low-precision rollout for RL on diffusion LLMs. Its primary significance lies in resolving the "rollout-training mismatch" without sacrificing the efficiency gains of quantization.

The authors assert that:

1. Management over Elimination: The mismatch should be managed, not eradicated. AIS successfully suppresses the destabilizing bias of late-stage training while preserving the exploratory benefits of quantization noise in early stages.
2. Scalability: The method is computationally efficient, adding negligible overhead because the diagnostics are scalar reductions over existing statistics.
3. Robustness: AIS generalizes across different model families (autoregressive vs. diffusion) and training paradigms (Full FT vs. LoRA).

The paper concludes that AIS enables high-throughput, low-precision RL training that matches or surpasses full-precision baselines, making it a viable solution for scaling RL in resource-constrained environments.