来源:市场资讯

(来源:机器之心Pro)

打开网易新闻 查看精彩图片

当强化学习后训练的大规模 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 的核心思路、方法设计、实验结果与实际意义。

打开网易新闻 查看精彩图片

效果展示

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 在生成式模型后训练中的角色被重新定义了。它不再只是一个推理加速工具,而是开始成为强化学习探索阶段的有效代理。对于文生图后训练、偏好对齐、低比特量化以及系统级优化方向的研究者和工程团队来说,这个方案为生成模型后训练提供了一条更现实的落地路径,具备持续关注的价值。