FP4 Exploration, BF16 Entraînement : Apprentissage par Renforcement par Diffusion via une Mise à l’Échelle Efficace des Rollouts
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
Résumé
La post-formation par apprentissage par renforcement est récemment apparue comme un paradigme prometteur pour aligner les modèles de diffusion texte-image avec les préférences humaines. Dans les études récentes, l'augmentation de la taille des groupes d'exploration entraîne des améliorations significatives des performances, indiquant un potentiel substantiel de gains d'alignement supplémentaires. Cependant, la mise à l'échelle des explorations sur les modèles de diffusion fondateurs à grande échelle (par exemple, FLUX.1-12B) impose une lourde charge computationnelle. Pour atténuer ce goulot d'étranglement, nous explorons l'intégration de la quantification FP4 dans les explorations de Diffusion RL. Pourtant, nous identifions que les pipelines quantifiés naïfs introduisent intrinsèquement des risques de dégradation des performances. Pour surmonter ce dilemme entre efficacité et intégrité de l'entraînement, nous proposons Sol-RL (Speed-of-light RL), un nouveau cadre d'apprentissage par renforcement à deux étages exploitant le FP4. Premièrement, nous utilisons des explorations NVFP4 à haut débit pour générer un vaste pool de candidats et en extraire un sous-ensemble hautement contrastif. Deuxièmement, nous régénérons ces échantillons sélectionnés en précision BF16 et optimisons la politique exclusivement sur ceux-ci. En découplant l'exploration des candidats de l'optimisation de la politique, Sol-RL intègre les mécanismes algorithmiques de mise à l'échelle des explorations avec les gains de débit au niveau système du NVFP4. Cette conception algorithmique-matérielle synergétique accélère efficacement la phase d'exploration tout en réservant des échantillons haute fidélité pour l'optimisation. Nous démontrons empiriquement que notre cadre maintient l'intégrité de l'entraînement du pipeline en précision BF16 tout en exploitant pleinement les gains de débit permis par l'arithmétique FP4. Des expériences approfondies sur SANA, FLUX.1 et SD3.5-L confirment que notre approche offre des performances d'alignement supérieures sur plusieurs métriques tout en accélérant la convergence de l'entraînement jusqu'à 4,64 fois, libérant ainsi la puissance de la mise à l'échelle massive des explorations à une fraction du coût.
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.