用自适应采样修正低精度强化学习的训练偏差,提升效率与稳定性。
AIS: Adaptive Importance Sampling for Quantized RL
- 基于实时诊断动态调整梯度修正强度,缓解精度不匹配问题。
- 在数学推理和规划任务中保持BF16基线性能,同时实现1.5至2.76倍加速。
- 适合追求高效低精度强化学习训练的研究者或工程团队。
大语言模型的强化学习受生成回溯成本制约,促使采用低精度回溯(如FP8)配合BF16训练器以提升吞吐量并减轻内存压力。这引入了回溯-训练不匹配,导致策略梯度偏差,可能使推理基准训练崩溃。我们发现该不匹配具有非平稳性:训练初期提供随机探索优势,暴露训练器未充分采样的轨迹;但随着策略集中,同一扰动转为不稳定偏差源。为此,我们提出自适应重要性采样(AIS),一种基于每批次动态调整干预强度的校正框架。AIS结合权重可靠性、发散严重性和方差放大三个实时诊断,生成混合系数,在未校正与完全重要性加权梯度间插值,抑制不稳定的偏差成分,同时保留探索益处。我们将AIS集成到GRPO,并在基于扩散的LLaDA-8B-Instruct及自回归Qwen3-8B与Qwen3.5-9B上评估其在数学推理与规划基准上的表现。结果表明,AIS在多数任务上达到与BF16基线相当的性能,同时维持1.5至2.76倍的FP8回溯加速。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。