提速4.6倍！NVIDIA×港大×MIT联手重新定义扩散模型训练速度上限

当强化学习后训练的大规模 rollout 已经被证明能够提升图像生成模型的偏好对齐能力，推理负担就成了制约训练速度的核心瓶颈。来自 NVIDIA、港大和 MIT 的团队提出的 Sol-RL，通过「FP4 先探索、BF16 再训练」的后训练框架，将达到等效 reward 水平的收敛速度最高提升到 4.64x，在训练速度与对齐效果之间给出了一条更具工程可行性的解法。

Sol-RL（Speed-of-light RL）是一种将 NVFP4 推理融合进 Diffusion 强化学习微调的高效训练框架。该方法并不是用量化推理结果直接训练，而是让 NVFP4 rollout 承担高吞吐的大规模探索任务，先在海量候选中筛出最有对比性的样本的初始噪声集合，再让 BF16 对这些关键初始噪声进行高保真再生成并完成策略优化。实验结果表明，该设计在 SANA、FLUX.1 和 SD3.5-L 模型上都带来了明显收益：在相同 GPU-hour 预算下，达到等效 reward 水平的收敛速度最高提升 4.64x，同时基本保持了 BF16 高精度 pipeline 的训练保真度。本文将深入探讨 Sol-RL 的核心思路、方法设计、实验结果与实际意义。

• 论文名称：FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
• 论文地址：https://arxiv.org/abs/2604.06916
• 项目主页：https://nvlabs.github.io/Sana/Sol-RL/
• 代码链接：https://github.com/NVlabs/Sana/

效果展示

Sol-RL 在速度和性能上表现出色，下图展示了经过 Sol-RL 优化后 SANA 模型的整体效果，以及与基础模型相比在复杂细节、语义一致性和整体审美上的改善。

研究背景

在文生图模型的后训练阶段，强化学习正在成为提升人类偏好对齐能力的重要路径。现有研究不断表明，扩大 rollout，也就是为同一个 prompt 生成更多候选图像，再从中筛选高对比样本（例如其中表现最好的 n 个和最差的 n 个构成的 2n 个图像）进行优化，能够显著提升模型的对齐效果。对于基于 GRPO 的 Diffusion RL 来说，更多候选意味着更强的对比信号，也意味着更可靠的梯度更新。

但问题在于：rollout 一旦增大，训练瓶颈很快就会从参数更新转移到海量候选样本生成。尤其在 FLUX.1、SD3.5-L 这类参数量较大、需要多次迭代推理的文生图 Diffusion 模型上，想真正把 rollout scaling 做起来，rollout 阶段生成候选样本的计算成本会迅速上升。为了提高海量 rollout 效率，团队引入 nvfp4 量化推理，但论文进一步指出，如果直接把低比特 rollout 样本当作训练目标，训练稳定性和最终效果都会受到明显影响。因此，关键并不只是「能不能用 FP4」，而是「FP4 应该在训练流程里扮演什么角色」。

核心创新

Sol-RL 的核心创新在于其双阶段解耦框架。与直接把低精度计算贯穿整个训练流程不同，Sol-RL 把 rollout 样本的「探索」和「生成」拆开处理，让不同精度承担不同职责。

量化 rollout 不适合直接训练：论文首先指出，直接将 FP4 量化推理样本用作优化目标，会带来训练退化和不稳定性的隐患，如下图实验中红色曲线所示。

FP4 适合作为探索代理：尽管 FP4 样本在像素层面会引入偏差，但研究发现，给定相同的初始噪声，在同组候选样本的相对 reward 排序上，NVFP4 推理样本依然保持了和 BF16 推理样本足够高的一致性。因此，它非常适合用来承担「大规模探索」和「候选筛选」的任务，也就是在大量初始噪声中快速识别出哪些会产生「最好」或者「最坏」的最终图像。

方法概述

Sol-RL 的整体流程可以概括为两个阶段。第一阶段，框架使用 NVFP4 rollout 和更少的采样步数快速生成一个大规模候选池，并根据 reward 对候选样本做排序，筛选出对应图像得分最高 / 最低的初始噪声种子集合。第二阶段，框架并不会直接拿这些低精度样本做训练，而是只保留第一阶段筛出的关键种子，再用 BF16 精度重新生成高保真样本，并仅基于这些高保真样本完成策略优化。这样一来，FP4 负责快速找方向，BF16 负责对其中一小部分高对比度样本做高质量生成用于训练，效率和稳定性被重新组织到同一个框架中。这样的流程让高成本 BF16 计算只集中在真正会影响梯度更新的部分，而不是平均浪费在大量最终不会参与训练的候选图像上。

实验结果

实验结果表明，Sol-RL 在多个基础模型和多个 reward 指标上都展现出明显优势。在相同 GPU-hour 预算下，Sol-RL 在 SANA、FLUX.1 和 SD3.5-L 上持续优于基线方法，并将达到等效 reward 水平的收敛速度最高提升至 4.64x。如下图所示，Sol-RL 在相同 wall-clock 预算下能够更快达到基线性能，在有限时间内达到更高对齐质量。

进一步做时间拆解，相对于直接使用高精度进行 rollout scaling，Sol-RL 在 rollout 阶段的加速最高达到 2.41x，训练迭代时间最高提升 1.62x。Sol-RL 通过探索 - 重生成的两阶段设计，显著缓解了完全使用 BF16 rollout scaling 带来的计算瓶颈，并且避免了直接使用低精度样本进行训练带来的不稳定性，相对于全程 NVFP4 低精度推理版本仅带来约为 2% 的额外开销。

结论与展望

当大规模 rollout 已经被证明能够持续提升生成模型偏好对齐能力，接下来的关键问题就是如何以更低成本释放这种扩展带来的收益。Sol-RL 给出的答案是：让低精度负责探索，让高精度负责优化。这也意味着，FP4 在生成式模型后训练中的角色被重新定义了。它不再只是一个推理加速工具，而是开始成为强化学习探索阶段的有效代理。对于文生图后训练、偏好对齐、低比特量化以及系统级优化方向的研究者和工程团队来说，这个方案为生成模型后训练提供了一条更现实的落地路径，具备持续关注的价值。