FP4 Verkennen, BF16 Trainen: Diffusie Versterkend Leren via Efficiënte Rollout-schaalvergroting
FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
April 8, 2026
Auteurs: Yitong Li, Junsong Chen, Shuchen Xue, Pengcuo Zeren, Siyuan Fu, Dinghao Yang, Yangyang Tang, Junjie Bai, Ping Luo, Song Han, Enze Xie
cs.AI
Samenvatting
Op reinforcement learning gebaseerde nafundering is recent naar voren gekomen als een veelbelovend paradigma voor het afstemmen van tekst-naar-beeld diffusiemodellen op menselijke voorkeuren. Uit recente studies blijkt dat het vergroten van de rollout-groepsgrootte aanzienlijke prestatieverbeteringen oplevert, wat duidt op een substantiële ruimte voor verdere afstemmingswinst. Het opschalen van rollouts op grootschalige fundamentele diffusiemodellen (zoals FLUX.1-12B) legt echter een zware computationele last op. Om dit knelpunt te verlichten, onderzoeken we de integratie van FP4-kwantisatie in Diffusion RL-rollouts. We constateren echter dat naïeve gekwantiseerde pijplijnen inherent risico's op prestatieverlies met zich meebrengen. Om dit dilemma tussen efficiëntie en trainingsintegriteit te overwinnen, stellen we Sol-RL (Speed-of-light RL) voor, een nieuw FP4-gestuurd Two-stage Reinforcement Learning-framework. Ten eerste gebruiken we high-throughput NVFP4-rollouts om een enorme kandidaatpool te genereren en hieruit een sterk contrastieve subset te extraheren. Ten tweede regenereren we deze geselecteerde samples in BF16-precisie en optimaliseren we het beleid uitsluitend hierop. Door kandidaatverkenning te ontkoppelen van beleidsoptimalisatie, integreert Sol-RL de algoritmische mechanismen van rollout-schaling met de systeemniveau doorvoerwinsten van NVFP4. Dit synergetische algoritme-hardware-ontwerp versnelt de rollout-fase effectief en reserveert tegelijkertijd hoogwaardige samples voor optimalisatie. We tonen empirisch aan dat ons framework de trainingsintegriteit van de BF16-precisiepijplijn behoudt, terwijl het de door FP4-rekenkunde mogelijk gemaakte doorvoerwinsten volledig benut. Uitgebreide experimenten met SANA, FLUX.1 en SD3.5-L bevestigen dat onze aanpak superieure afstemmingsprestaties levert op meerdere metrieken, terwijl de trainingsconvergentie tot 4,64 keer wordt versneld, waardoor de kracht van massale rollout-schaling tegen een fractie van de kosten wordt ontsloten.
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.