FP4 Esplora, BF16 Allena: Apprendimento per Rinforzo Diffusivo tramite Scalabilità Efficiente dei Rollout
FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
April 8, 2026
Autori: Yitong Li, Junsong Chen, Shuchen Xue, Pengcuo Zeren, Siyuan Fu, Dinghao Yang, Yangyang Tang, Junjie Bai, Ping Luo, Song Han, Enze Xie
cs.AI
Abstract
L'addestramento post-allenamento basato sul reinforcement learning è recentemente emerso come un paradigma promettente per allineare i modelli di diffusione testo-immagine con le preferenze umane. Negli studi recenti, l'aumento della dimensione del gruppo di rollout produce miglioramenti prestazionali marcati, indicando un ampio margine per ulteriori guadagni di allineamento. Tuttavia, il ridimensionamento dei rollout su modelli di diffusione fondazionali su larga scala (ad es. FLUX.1-12B) impone un pesante onere computazionale. Per alleviare questo collo di bottiglia, esploriamo l'integrazione della quantizzazione FP4 nei rollout del Diffusion RL. Tuttavia, identifichiamo che le pipeline quantizzate naive introducono intrinsecamente rischi di degradazione delle prestazioni. Per superare questo dilemma tra efficienza e integrità dell'addestramento, proponiamo Sol-RL (Speed-of-light RL), un nuovo framework di Reinforcement Learning a due stadi potenziato da FP4. In primo luogo, utilizziamo rollout ad alto throughput NVFP4 per generare un pool massiccio di candidati ed estrarne un sottoinsieme altamente contrastivo. In secondo luogo, rigeneriamo questi campioni selezionati in precisione BF16 e ottimizziamo la policy esclusivamente su di essi. Disaccoppiando l'esplorazione dei candidati dall'ottimizzazione della policy, Sol-RL integra i meccanismi algoritmici del ridimensionamento dei rollout con i guadagni di throughput a livello di sistema di NVFP4. Questo design sinergico algoritmo-hardware accelera efficacemente la fase di rollout riservando al contempo campioni ad alta fedeltà per l'ottimizzazione. Dimostriamo empiricamente che il nostro framework mantiene l'integrità addestrativa della pipeline in precisione BF16 sfruttando appieno i guadagni di throughput abilitati dall'aritmetica FP4. Esperimenti estesi su SANA, FLUX.1 e SD3.5-L confermano che il nostro approccio fornisce prestazioni di allineamento superiori su più metriche, accelerando la convergenza dell'addestramento fino a 4.64 volte e sbloccando la potenza del ridimensionamento massiccio dei rollout a una frazione del costo.
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.