FP4 Explorar, BF16 Treinar: Aprendizagem por Reforço de Difusão via Escalonamento 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
Resumo
O pós-treinamento baseado em Aprendizagem por Reforço (RL) emergiu recentemente como um paradigma promissor para alinhar modelos de difusão de texto para imagem com as preferências humanas. Em estudos recentes, o aumento do tamanho do grupo de rollouts (simulações) produz melhorias de desempenho pronunciadas, indicando um espaço substancial para ganhos adicionais de alinhamento. No entanto, a escalagem de rollouts em modelos de difusão fundamentais de grande escala (por exemplo, FLUX.1-12B) impõe uma pesada carga computacional. Para aliviar este gargalo, exploramos a integração da quantização FP4 nos rollouts de RL de Difusão. Contudo, identificamos que os pipelines quantizados ingênuos introduzem inerentemente riscos de degradação de desempenho. Para superar este dilema entre eficiência e integridade do treinamento, propomos o Sol-RL (Speed-of-light RL), uma nova estrutura de Aprendizagem por Reforço em Duas Etapas capacitada por FP4. Primeiro, utilizamos rollouts de alta produtividade em NVFP4 para gerar um pool massivo de candidatos e extrair um subconjunto altamente contrastivo. Em segundo lugar, regeneramos essas amostras selecionadas em precisão BF16 e otimizamos a política exclusivamente nelas. Ao desacoplar a exploração de candidatos da otimização da política, o Sol-RL integra os mecanismos algorítmicos de escalonamento de rollouts com os ganhos de produtividade em nível de sistema do NVFP4. Este projeto algorítmico-hardware sinérgico acelera eficazmente a fase de rollout, reservando ao mesmo tempo amostras de alta fidelidade para otimização. Demonstramos empiricamente que a nossa estrutura mantém a integridade do treinamento do pipeline de precisão BF16, enquanto explora totalmente os ganhos de produtividade possibilitados pela aritmética FP4. Extensas experiências com SANA, FLUX.1 e SD3.5-L corroboram que a nossa abordagem oferece um desempenho de alinhamento superior em múltiplas métricas, enquanto acelera a convergência do treinamento em até 4,64 vezes, desbloqueando o poder da escalagem massiva de rollouts a uma fração do custo.
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.