📊 statistics

AIS: Adaptive Importance Sampling for Quantized RL

本文介绍了自适应重要性采样（AIS），这是一个框架，能够动态校正大语言模型强化学习中由低精度（FP8）采样轨迹引起的策略梯度偏差，从而在保持训练稳定性并实现与全精度基线相当的性能的同时，显著提升推理速度。

原作者： Jiajun Zhou, Wei Shao, Lingchao Zheng, Yuwei Fan, Ngai Wong

发布于 2026-05-15

📖 1 分钟阅读☕ 轻松阅读

原作者： Jiajun Zhou, Wei Shao, Lingchao Zheng, Yuwei Fan, Ngai Wong

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是论文"AIS：用于量化强化学习的自适应重要性采样”的解释，已用通俗易懂的语言和日常类比进行翻译。

宏观图景：“快但不可靠”的学生

想象你正在训练一位才华横溢的学生（大型语言模型）来解决复杂的数学问题。为了让训练更快、更便宜，你决定让学生使用低分辨率教科书（FP8 量化）进行练习，而老师则使用高分辨率教科书（BF16 精度）来批改作业。

问题所在：

速度提升： 低分辨率教科书更轻便、阅读速度更快。学生可以以 1.5 到 2.7 倍的速度进行练习（生成“轨迹”），并且内存占用减半。
故障： 因为教科书模糊不清，学生有时会犯小错，或者以老师未曾预料的方式看待问题。
- 早期阶段： 这些模糊的错误实际上很有帮助！它们就像“幸运的意外”，迫使学生尝试他们原本不会想到的奇怪解决方案。这就像学生胡乱猜测，却意外找到了正确答案。
- 后期阶段： 随着学生变得更好，这些模糊的错误变得危险起来。它们开始混淆老师，导致学生学到错误的教训，最终在困难测试中“崩溃”（完全失败）。

旧方案：
以前的方法试图通过以下方式解决这个问题：

对所有事情都使用沉重、缓慢的高分辨率教科书（太慢了）。
应用一种“一刀切”的过滤器来修正模糊的教科书。这就像给学生戴上眼罩：有时它有帮助，但通常要么太严格（扼杀了好运气），要么太宽松（让坏错误通过）。

新方案：AIS（智能教练）

作者提出了自适应重要性采样（AIS）。将 AIS 想象成一位智能教练，他实时观察学生的练习，并随时调整规则。

AIS 不使用静态过滤器，而是利用三个“传感器”来决定在多大程度上修正学生的作业：

可靠性传感器（学生的猜测值得信任吗？）：
- 类比： 如果学生正在胡乱猜测，教练就知道“修正”（告诉他们正确答案应该是什么）风险太大。教练会说：“暂时不要相信这个修正。”
发散度传感器（模糊的教科书有多不同？）：
- 类比： 如果模糊的教科书看起来和真实的一模一样，教练会说：“无需修正；你没问题。”但如果模糊的教科书完全错误，教练会说：“我们需要立即修正这个问题。”
方差传感器（修正是否让事情变得混乱？）：
- 类比： 如果修正过于极端，导致学生的大脑晕头转向（高方差），教练就会退一步以保持稳定。

工作原理：
教练混合了两种方法：

方法 A： 让学生拿着模糊的教科书自由发挥（有利于早期的探索）。
方法 B： 使用高分辨率数学严格修正学生（有利于后期的精确度）。

教练为每一批练习计算一个“混合分数”。

早期训练： 分数倾向于方法 A。教练让“模糊噪声”帮助学生探索新想法。
后期训练： 分数转向方法 B。教练收紧规则，防止学生犯危险错误。

结果：快速、廉价且准确

团队在两种类型的 AI 模型上测试了这位“智能教练”：

标准模型（Qwen）： 这种“自回归”模型一次生成一个词。
扩散模型（LLaDA）： 一种较新的类型，通过“去噪”（就像将静电转化为清晰图像）来生成文本。

发生了什么？

速度： 他们保留了使用快速、模糊教科书带来的1.5 倍到 2.7 倍的速度提升。
内存： 他们节省了约50% 的内存。
准确性： AI 的表现与全程使用缓慢、沉重的教科书一样好（有时甚至更好）。
- 为什么更好？ 论文指出，这种“模糊噪声”充当了有益的“探索奖励”，帮助 AI 找到了比完美精确的学生独自可能找到的更好的解决方案。

总结

这篇论文解决了 AI 训练中的一个棘手问题：我们如何在快速训练 AI 的同时不使其变笨？

他们发现，使用 AI 的“模糊”版本进行练习对速度很有利，但需要一个智能、自适应的教练来知道何时让 AI 胡乱猜测，何时介入并修正。他们的新方法AIS就充当了这位教练，使 AI 能够以两倍的速度训练，同时不损失其智能。

技术摘要：量化强化学习中的自适应重要性采样（AIS）

1. 问题陈述

大语言模型（LLM）的强化学习（RL）目前受限于轨迹 rollout 生成的计算成本。为缓解这一问题，近期的工业实践（如 DeepSeek V3.2、veRL）在 rollout 引擎中采用低精度量化（如 FP8），同时保持较高精度（BF16）的训练器。这导致了rollout 与训练的不匹配：量化 rollout 策略（ $\hat{\pi}_{\theta_{old}}$ ）生成的轨迹，在分布上不同于用于计算梯度的全精度学习策略（ $\pi_{\theta_{old}}$ ）。

该论文指出这种不匹配是非平稳的，并起到“双刃剑”的作用：

训练早期：量化噪声充当随机探索奖励，使梯度接触到训练器原本可能采样不足的轨迹。
训练后期：随着策略变得锐利，同样的噪声变成了导致偏差的不稳定源，导致在推理基准测试上训练崩溃（例如，Qwen3-8B 在 AIME25 上的准确率下降了 6.63%）。

标准重要性采样（IS）试图校正这种偏差，但由于长序列中极高的方差，往往难以奏效。相反，截断重要性采样（TIS）虽然稳定了训练，但使用了固定的校正强度，无法适应不匹配性质的演变，导致要么过度校正（抑制有用的探索），要么校正不足（允许不稳定性）。

2. 方法论：自适应重要性采样（AIS）

作者提出了自适应重要性采样（AIS），这是一个在每个批次基础上动态调整重要性采样校正强度的框架。AIS 不使用固定阈值，而是计算一个单一的混合系数 $\alpha(x) \in [0, 1]$ ，该系数在未校正梯度与完全重要性加权梯度之间进行插值。

2.1 混合系数

AIS 的核心是一个因子化门控机制，它结合了三个实时诊断指标，这些指标均源自策略更新过程中已计算的统计量：

权重可靠性（ $\alpha_{ess}$ ）：衡量截断重要性权重的有效样本量（ESS）比率。如果权重分布集中（低 ESS），则校正被视为不可靠并降低权重。
$\alpha_{ess} = (1 + CV^2(\bar{w}))^{-1/2}$
发散严重程度（ $\alpha_{mis}$ ）：衡量 rollout 策略与训练策略之间平均绝对对数概率的差异。仅当此不匹配超过饱和阈值 $\delta$ 时，才启用校正。
$\alpha_{mis} = \min\left(1, \frac{\bar{d}}{\delta}\right)$
方差放大（ $\alpha_{var}$ ）：监控大重要性权重是否与大幅值优势值重合，这会将梯度方差放大至超出可接受范围。如果方差放大比率 $\Delta\sigma$ 超过容忍度 $\gamma$ ，则施加惩罚。
$\alpha_{var} = \max\left(0, \frac{\Delta\sigma - \gamma}{\gamma}\right)$

最终混合系数计算如下：
$\alpha = \text{clip}(\alpha_{ess} - \beta \cdot \alpha_{var}, 0, 1) \cdot \alpha_{mis}$

2.2 与 GRPO 的集成

AIS 被集成到组相对策略优化（GRPO）中。该方法引入了一个校正权重 $\tilde{w}_i$ ，用于调整 GRPO 目标函数中的优势项：
$\tilde{w}_i = 1 + \alpha \cdot (\min(\rho_{rt}, C) - 1)$
其中 $\rho_{rt}$ 是 rollout-训练比率， $C$ 是截断阈值。该公式确保当不匹配可忽略不计（ $\alpha \to 0$ ）时，目标函数精确退化为原始 GRPO，从而保持向后兼容性。

3. 主要贡献

非平稳性表征：该论文确立了 rollout-训练不匹配并非均匀有害；它在训练早期提供探索奖励，但在后期成为偏差源。这挑战了将量化不匹配单纯视为对手的观点。
自适应校正框架：AIS 是首个基于实时诊断（可靠性、发散性和方差）动态调节重要性采样校正强度的方法，而非使用静态阈值。
扩散模型应用：这项工作首次研究了基于扩散的 LLM（具体为 LLaDA-8B-Instruct）的低精度 rollout 强化学习，证明了不匹配问题及 AIS 解决方案不仅适用于自回归架构。
效率与准确性权衡：该方法实现了平衡，既恢复了全精度 BF16 训练的准确性，又保留了 FP8 rollout 的吞吐量和内存优势。

4. 实验结果

作者在Qwen3-8B、Qwen3.5-9B（自回归）和LLaDA-8B-Instruct（扩散）模型上评估了 AIS，测试基准涵盖数学推理（GSM8K、MATH500、AIME25、AMC、Olympiad）和规划（Sudoku、Countdown）。

准确性恢复：
- 自回归模型：在 Qwen3-8B 上，朴素 FP8 rollout 将 AIME25 准确率从 BF16 的 36.70% 降至 26.81%。AIS 将其恢复至43.33%，超越了 BF16 基线。在 Qwen3.5-9B 上也观察到了类似的改进，AIS 在 12 个基准测试对中的 11 个上匹配或超过了 BF16。
- 扩散模型：在 LLaDA-8B 上，FP8 rollout 导致了显著的性能下降（例如，在 512 个 token 的 Countdown 任务上下降了 19.00%）。AIS 使性能恢复，在各种序列长度和微调方案（全量微调 Full FT 和 LoRA）下均匹配或超过了 BF16 水平。
通用能力：AIS 在通用基准测试（MMLU、PIQA、ARC 等）上保持了性能，表明没有对 RL 奖励信号过拟合。
效率提升：
- 加速：FP8 rollout 根据批次大小和模型架构不同，提供了1.5 倍至 2.76 倍的端到端 rollout 时间加速（在扩散模型中观察到更高的增益，这是由于张量核心的利用）。
- 内存：FP8 rollout 将 rollout 模型的内存占用减半（例如，LLaDA-8B 约为 7.95 GB 对比 BF16 的约 15.9 GB），使得在 BF16 下曾出现内存不足（OOM）的更大批次大小成为可能。

5. 意义与主张

该论文声称，AIS 是首个针对基于扩散的 LLM 进行低精度 rollout 强化学习的研究。其主要意义在于解决了"rollout-训练不匹配”问题，同时未牺牲量化的效率增益。

作者断言：

管理而非消除：不匹配应被管理，而非根除。AIS 成功抑制了训练后期的破坏性偏差，同时保留了量化噪声在早期的探索益处。
可扩展性：该方法计算高效，增加的开销可忽略不计，因为诊断指标是对现有统计量的标量归约。
鲁棒性：AIS 泛化于不同的模型家族（自回归与扩散）和训练范式（全量微调与 LoRA）。

该论文得出结论，AIS 实现了高吞吐量的低精度 RL 训练，其表现匹配或超越了全精度基线，使其成为在资源受限环境中扩展 RL 的可行解决方案。