论文
FP4探索、BF16训练:通过扩大轨迹采样提高扩散强化学习效率
FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
摘要
基于强化学习的 后训练 最近已成为将文本到图像扩散模型与人类偏好保持一致的有前途的范例。在最近的研究中,增加轨迹采样组规模可以显着提高性能,这表明进一步调整增益的空间很大。然而,大规模基础扩散模型(例如 FLUX.1-12B)的扩展会带来沉重的计算负担。为了缓解这一瓶颈,我们探索将 FP4 量化集成到 Diffusion RL 中。然而,我们发现幼稚的量化管道本质上会带来性能下降的风险。为了克服效率和训练完整性之间的困境,我们提出了 Sol-RL(光速 RL),这是一种新型的 FP4 授权的两阶段强化学习框架。首先,我们利用高吞吐量 NVFP4 部署来生成大量候选池并提取高度对比的子集。其次,我们以 BF16 精度重新生成这些选定的样本,并专门针对它们优化策略。通过将候选探索与策略优化解耦,Sol-RL 将轨迹采样扩展的算法机制与 NVFP4 的系统级吞吐量增益相集成。这种协同算法-硬件设计有效地加速了轨迹采样阶段,同时保留了高保真样本进行优化。我们凭经验证明,我们的框架保持了 BF16 精度管道的训练完整性,同时充分利用 FP4 算法实现的吞吐量增益。 SANA、FLUX.1 和 SD3.5-L 的广泛实验证实,我们的方法在多个指标上提供了卓越的对齐性能,同时将训练收敛速度加快了高达 4.64 美元\次$,以极低的成本释放了大规模部署扩展的力量。