FP4 Exploración, BF16 Entrenamiento: Aprendizaje por Refuerzo de Difusión mediante Escalado Eficiente de Rollouts
FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
April 8, 2026
Autores: Yitong Li, Junsong Chen, Shuchen Xue, Pengcuo Zeren, Siyuan Fu, Dinghao Yang, Yangyang Tang, Junjie Bai, Ping Luo, Song Han, Enze Xie
cs.AI
Resumen
El post-entrenamiento basado en aprendizaje por refuerzo ha surgido recientemente como un paradigma prometedor para alinear los modelos de difusión de texto a imagen con las preferencias humanas. En estudios recientes, aumentar el tamaño del grupo de muestreo (rollout) produce mejoras de rendimiento notables, lo que indica un margen sustancial para obtener mayores ganancias de alineación. Sin embargo, escalar los rollouts en modelos de difusión fundamentales a gran escala (por ejemplo, FLUX.1-12B) impone una carga computacional significativa. Para aliviar este cuello de botella, exploramos la integración de la cuantización FP4 en los rollouts de Difusión por RL. No obstante, identificamos que las canalizaciones de cuantización simples introducen inherentemente riesgos de degradación del rendimiento. Para superar este dilema entre eficiencia e integridad del entrenamiento, proponemos Sol-RL (Speed-of-light RL), un novedoso marco de Aprendizaje por Refuerzo en Dos Etapas potenciado por FP4. Primero, utilizamos rollouts de alto rendimiento NVFP4 para generar un grupo masivo de candidatos y extraer un subconjunto altamente contrastivo. En segundo lugar, regeneramos estas muestras seleccionadas con precisión BF16 y optimizamos la política exclusivamente sobre ellas. Al desacoplar la exploración de candidatos de la optimización de políticas, Sol-RL integra los mecanismos algorítmicos del escalado de rollouts con las ganancias de rendimiento a nivel de sistema de NVFP4. Este diseño sinérgico algoritmo-hardware acelera eficazmente la fase de rollout mientras reserva muestras de alta fidelidad para la optimización. Demostramos empíricamente que nuestro marco mantiene la integridad del entrenamiento de la canalización de precisión BF16 mientras explota plenamente las ganancias de rendimiento habilitadas por la aritmética FP4. Experimentos exhaustivos en SANA, FLUX.1 y SD3.5-L corroboran que nuestro enfoque ofrece un rendimiento de alineación superior en múltiples métricas, al tiempo que acelera la convergencia del entrenamiento hasta 4.64 veces, desbloqueando el poder del escalado masivo de rollouts a una fracción del coste.
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.