FP4 для исследования, BF16 для обучения: обучение с подкреплением на основе диффузионных моделей через эффективное масштабирование rollout

FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling

April 8, 2026

Авторы: Yitong Li, Junsong Chen, Shuchen Xue, Pengcuo Zeren, Siyuan Fu, Dinghao Yang, Yangyang Tang, Junjie Bai, Ping Luo, Song Han, Enze Xie

cs.AI

Аннотация

Посттренировка на основе обучения с подкреплением недавно появилась как перспективная парадигма для согласования диффузионных моделей текст-изображение с человеческими предпочтениями. В последних исследованиях увеличение размера группы rollout приводит к значительному улучшению производительности, что указывает на существенный потенциал для дальнейшего выравнивания. Однако масштабирование rollout на крупных фундаментальных диффузионных моделях (например, FLUX.1-12B) создает значительную вычислительную нагрузку. Для смягчения этого узкого места мы исследуем интеграцию FP4-квантования в процесс rollout обучения с подкреплением. Тем не менее, мы выявили, что наивные квантованные конвейеры неизбежно несут риски деградации производительности. Чтобы преодолеть эту дилемму между эффективностью и целостностью обучения, мы предлагаем Sol-RL (обучение с подкреплением на скорости света) — новую FP4-ускоренную двухэтапную структуру обучения с подкреплением. Сначала мы используем высокопроизводительные NVFP4 rollout для генерации обширного пула кандидатов и извлечения высококонтрастного подмножества. Затем мы повторно генерируем выбранные образцы с точностью BF16 и оптимизируем политику исключительно на них. Разделяя исследование кандидатов и оптимизацию политики, Sol-RL объединяет алгоритмические механизмы масштабирования rollout с преимуществами пропускной способности NVFP4 на системном уровне. Этот синергетический алгоритмически-аппаратный дизайн эффективно ускоряет этап rollout, сохраняя при этом образцы высокой точности для оптимизации. Мы эмпирически демонстрируем, что наша структура сохраняет целостность обучения конвейера с точностью BF16, полностью используя прирост пропускной способности, обеспечиваемый FP4-арифметикой. Многочисленные эксперименты на моделях SANA, FLUX.1 и SD3.5-L подтверждают, что наш подход обеспечивает превосходную производительность согласования по множеству метрик, ускоряя сходимость обучения до 4.64 раз и раскрывая потенциал массового масштабирования rollout при существенно меньших затратах.

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.