用FP4加速扩散强化学习,实现高效高质对齐。
FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
- 分两阶段:先用FP4快速生成大量候选样本,再用BF16精修优化。
- 在FLUX.1等模型上提升训练速度4.64倍,保持高精度效果。
- 适合需要大规模采样但算力受限的图像生成对齐任务。
基于强化学习的后训练近期成为对齐文本到图像扩散模型与人类偏好的一种有前景的方法。研究表明,增大采样组规模能显著提升性能,表明仍有巨大改进空间。然而,在大规模基础扩散模型(如FLUX.1-12B)上扩展采样会带来沉重计算负担。为缓解这一瓶颈,我们探索将FP4量化引入扩散强化学习采样流程。但发现直接量化会引发性能下降风险。为此,我们提出Sol-RL(光速强化学习)框架,一种基于FP4的两阶段强化学习方法。首先利用高吞吐的NVFP4采样生成海量候选样本,并提取高度对比性的子集;其次在BF16精度下重生成这些样本并仅对其优化策略。通过解耦候选探索与策略优化,该框架融合了采样扩展的算法优势与NVFP4的系统吞吐提升。实验表明,该方法在保持BF16精度训练完整性的同时,充分释放了FP4算术带来的吞吐增益。在SANA、FLUX.1和SD3.5-L上的广泛实验验证,本方法在多个指标上实现更优对齐效果,训练收敛速度最高提升4.64倍,以极低成本解锁大规模采样潜力。
原文摘要 · Abstract (English)
Reinforcement-Learning-based post-training has recently emerged as a promising paradigm for aligning text-to-image diffusion models with human preferences. In recent studies, increasing the rollout group size yields pronounced performance improvements, indicating substantial room for further alignment gains. However, scaling rollouts on large-scale foundational diffusion models (e.g., FLUX.1-12B) imposes a heavy computational burden. To alleviate this bottleneck, we explore the integration of FP4 quantization into Diffusion RL rollouts. Yet, we identify that naive quantized pipelines inherently introduce risks of performance degradation. To overcome this dilemma between efficiency and training integrity, we propose Sol-RL (Speed-of-light RL), a novel FP4-empowered Two-stage Reinforcement Learning framework. First, we utilize high-throughput NVFP4 rollouts to generate a massive candidate pool and extract a highly contrastive subset. Second, we regenerate these selected samples in BF16 precision and optimize the policy exclusively on them. By decoupling candidate exploration from policy optimization, Sol-RL integrates the algorithmic mechanisms of rollout scaling with the system-level throughput gains of NVFP4. This synergistic algorithm-hardware design effectively accelerates the rollout phase while reserving high-fidelity samples for optimization. We empirically demonstrate that our framework maintains the training integrity of BF16 precision pipeline while fully exploiting the throughput gains enabled by FP4 arithmetic. Extensive experiments across SANA, FLUX.1, and SD3.5-L substantiate that our approach delivers superior alignment performance across multiple metrics while accelerating training convergence by up to $4.64\times$, unlocking the power of massive rollout scaling at a fraction of the cost.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。