FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
| Authors | Yitong Li et al. |
| Year | 2026 |
| HF Upvotes | 24 |
| arXiv | 2604.06916 |
| Download | |
| HF Page | View on Hugging Face |
Abstract
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.
Engineering Breakdown
Plain English
This paper addresses a critical bottleneck in scaling reinforcement learning-based post-training for large diffusion models like FLUX.1-12B. The authors found that increasing rollout group size significantly improves alignment with human preferences, but the computational cost becomes prohibitive at scale. They propose Sol-RL, which integrates FP4 quantization into the RL training pipeline to reduce memory and compute requirements while maintaining performance integrity through a two-stage approach that mitigates the performance degradation risks of naive quantization.
Core Technical Contribution
Sol-RL introduces a two-stage FP4-quantized reinforcement learning framework specifically designed for diffusion model alignment that solves the efficiency-integrity trade-off. The core novelty is identifying and addressing how naive quantization degrades RL training quality, then systematically architecting around this failure mode through staged computation. This represents a significant methodological advance over simply applying existing quantization techniques to RL workflows, as it recognizes that reward signals and policy gradients behave differently under extreme quantization than standard forward-pass inference does.
How It Works
Sol-RL operates in two stages to balance computational efficiency with training stability. In stage one, the model performs quantized forward passes during rollout generation to sample candidate images efficiently, using FP4 weights to reduce memory footprint and accelerate inference. Stage two processes these rollouts through selective full-precision or strategic mixed-precision computations for gradient computation and reward signal propagation, preventing quantization noise from corrupting the learning signal. The key insight is that you can afford to quantize the sampling mechanism aggressively since it's primarily generating diverse candidates, but you must preserve precision in the backward pass where gradient information flows and reward gradients determine policy updates. This staged approach lets the system scale to larger batch sizes and more rollouts per training iteration without the computational constraints of fully precision-based RL.
Production Impact
For production systems training alignment models on large diffusion bases, Sol-RL directly enables scaling from small rollout groups (e.g., 4-8 samples) to much larger batches (32-128+) within the same GPU memory budget, potentially cutting training time by 50-70% while improving final model quality. This matters concretely: teams can either train the same model faster, or allocate compute to explore more diverse reward functions, better preference data, or longer training horizons. The trade-off is moderate—you gain compute efficiency at the cost of implementing a two-stage pipeline and tuning the precision thresholds for your specific model and reward function, plus validation overhead to ensure your quantization strategy doesn't silently degrade reward signal quality. Adoption requires changes to training harnesses: separate quantization configs for rollout vs. gradient phases, careful reward computation placement, and extensive evaluation metrics to catch performance regressions early.
Limitations and When Not to Use This
The paper does not fully characterize how performance degradation risks scale across different model sizes, reward function types, or preference dataset distributions—the approach may work well for one reward signal but fail silently on another. It also assumes access to a well-trained base diffusion model; the method doesn't address alignment from scratch or cold-start regimes where you lack good rollout diversity. The two-stage design introduces additional hyperparameter tuning surface (when to switch stages, which operations stay full precision), and there's likely a sweet spot that varies per use case—one set of thresholds may not generalize. Finally, the paper doesn't discuss how Sol-RL interacts with other efficiency techniques like LoRA or pruning, so composing multiple optimization strategies remains an open question.
Research Context
This work builds directly on recent advances in RL-based diffusion alignment (RLHF for text-to-image models) which showed that larger rollout groups correlate with better human preference alignment, but hit a hardware wall at scale. It sits at the intersection of model efficiency research (quantization, mixed precision) and RL training stability, drawing from extensive literature on post-training alignment and quantization-aware training. Sol-RL opens a new direction: systematic two-stage quantization strategies tailored to RL rather than generic inference, potentially applicable beyond diffusion to other large-scale RL post-training scenarios (language models, reward modeling). The work implies future research should explore automated precision scheduling, theoretical bounds on quantization noise propagation through reward signals, and methods to adapt quantization dynamically based on gradient statistics.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::