FP4 Erkunden, BF16 Trainieren: Diffusionsverstärkungslernen durch effiziente Rollout-Skalierung
FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
April 8, 2026
Autoren: Yitong Li, Junsong Chen, Shuchen Xue, Pengcuo Zeren, Siyuan Fu, Dinghao Yang, Yangyang Tang, Junjie Bai, Ping Luo, Song Han, Enze Xie
cs.AI
Zusammenfassung
Reinforcement-Learning-basierte Nachbearbeitung (Post-Training) hat sich kürzlich als vielversprechendes Paradigma erwiesen, um Text-zu-Bild-Diffusionsmodelle mit menschlichen Präferenzen in Einklang zu bringen. In aktuellen Studien führt eine Vergrößerung der Rollout-Gruppengröße zu deutlichen Leistungsverbesserungen, was auf ein erhebliches Potenzial für weitere Anpassungsgewinne hindeutet. Die Skalierung von Rollouts bei großen foundation Diffusionsmodellen (z.B. FLUX.1-12B) stellt jedoch eine erhebliche rechnerische Belastung dar. Um diesen Engpass zu entschärfen, untersuchen wir die Integration von FP4-Quantisierung in Diffusion-RL-Rollouts. Allerdings stellen wir fest, dass naive quantisierte Pipelines inhärent das Risiko von Leistungseinbußen bergen. Um dieses Dilemma zwischen Effizienz und Trainingsintegrität zu überwinden, schlagen wir Sol-RL (Speed-of-light RL) vor, ein neuartiges FP4-gestütztes Zwei-Stufen-Reinforcement-Learning-Framework. Zunächst nutzen wir hochdurchsatzfähige NVFP4-Rollouts, um einen umfangreichen Kandidatenpool zu generieren und eine stark kontrastive Teilmenge zu extrahieren. Anschließend regenerieren wir diese ausgewählten Samples in BF16-Präzision und optimieren die Policy ausschließlich auf Basis dieser. Durch die Entkopplung der Kandidatensuche von der Policy-Optimierung vereint Sol-RL die algorithmischen Mechanismen der Rollout-Skalierung mit den systemseitigen Durchsatzvorteilen von NVFP4. Dieser synergetische Algorithmus-Hardware-Entwurf beschleunigt die Rollout-Phase effektiv, während hochwertige Samples für die Optimierung reserviert bleiben. Wir zeigen empirisch, dass unser Framework die Trainingsintegrität einer BF16-Präzisionspipeline bewahrt und gleichzeitig die durch FP4-Arithmetik ermöglichten Durchsatzgewinne voll ausschöpft. Umfangreiche Experimente mit SANA, FLUX.1 und SD3.5-L belegen, dass unser Ansatz eine überlegene Anpassungsleistung über mehrere Metriken hinweg liefert und dabei die Trainingskonvergenz um bis zum 4,64-fachen beschleunigt, wodurch die Skalierung mit massiven Rollouts zu einem Bruchteil der Kosten ermöglicht wird.
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.64times, unlocking the power of massive rollout scaling at a fraction of the cost.