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

Yitong Li; Junsong Chen; Shuchen Xue; Pengcuo Zeren; Siyuan Fu; Dinghao Yang; Yangyang Tang; Junjie Bai; Ping Luo; Song Han; Enze Xie

Source-linked AI summary

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

Yitong Li, Junsong Chen, Shuchen Xue, Pengcuo Zeren, Siyuan Fu, Dinghao Yang, Yangyang Tang, Junjie Bai, Ping Luo, Song Han, Enze Xie

arXiv:2604.06916v1Apr 8, 2026cs.LGcs.AIcs.CV

TL;DR

Diffusion RL gains from larger rollout groups, but candidate generation is expensive and direct quantized rollouts can degrade optimization. Sol-RL uses NVFP4 to explore and rank a massive pool, then regenerates selected contrastive samples in BF16 for policy optimization, achieving broad alignment gains with up to 4.64× faster convergence.

Problem
Larger rollout groups improve diffusion-RL alignment, but scaling candidate generation is computationally expensive and direct quantized targets risk degradation and instability.
Method
Sol-RL decouples rollout exploration from optimization by using high-throughput NVFP4 to rank a large pool and BF16 to regenerate selected high-contrastive samples.
Results
Across diverse diffusion models and reward metrics, Sol-RL maintains high-precision alignment quality and accelerates convergence by 1.91× to 4.64×.
Takeaways & Limitations
The framework enables efficient rollout scaling while reserving high-fidelity computation for the samples used in policy updates.
Takeaways & Limitations
Direct quantized rollouts remain unsuitable as optimization targets because they introduce an off-policy gap and distorted low-fidelity semantics.

Abstract

from arXiv · show

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.64\times$, unlocking the power of massive rollout scaling at a fraction of the cost.

1. Introduction

Diffusion RL benefits from larger rollout pools and selective training, but candidate generation becomes the computational bottleneck. Sol-RL addresses this efficiency–integrity dilemma by using FP4 for exploration and BF16 for optimization.

Motivation: Scaling rollout candidates and selecting highly contrastive samples improves alignment but shifts the bottleneck from policy optimization to candidate generation.Only a small subset is ultimately used for optimization, making large-scale candidate generation costly.
Sol-RL: Sol-RL uses FP4 quantization to generate a massive candidate pool and extract a high-contrastive subset before optimization.The framework separates candidate exploration from policy optimization.
Sol-RL: The selected samples are regenerated in BF16 precision so policy optimization uses high-fidelity training targets.This design reserves high-precision computation for samples retained after exploration.
Results: 4.64× convergence speedup is achieved while maintaining the alignment quality of the high-precision pipeline across SD3.5, FLUX.1, and SANA.The evaluation spans diverse foundation models and reward metrics.

2. Preliminaries

GRPO reduces critic-related overhead by estimating advantages from relative rewards within rollout groups. Larger groups provide more informative comparisons and stable statistics, but increase rollout and reward-evaluation costs; FP4 offers a throughput-oriented acceleration route.

Group Relative Policy Optimization: GRPO evaluates groups of candidate responses and computes advantages from relative rewards within each group.This avoids relying on a learned value-network critic for the baseline.
Group Relative Policy Optimization: GRPO optimizes a PPO-style clipped surrogate objective while regularizing the policy toward a reference model through a direct KL term.The objective is computed over the sampled group.
Rollout Scaling: Increasing group size can improve within-group comparisons and stabilize group statistics, but substantially raises rollout and reward-evaluation costs.The quality of the policy update depends strongly on the informativeness of sampled candidates.
FP4 Quantization: FP4 uses 4-bit floating-point arithmetic with block-level shared scaling; NVFP4 groups 16 elements under an E4M3 scale.The passage describes FP4 as a throughput-oriented acceleration paradigm with minimal precision degradation.

3. Methodology

Sol-RL decouples fast FP4 exploration from high-fidelity BF16 optimization. NVFP4 rollouts efficiently rank large candidate pools, while selected contrastive seeds are regenerated in BF16 to avoid direct quantized-target degradation and instability.

Rollout Scaling: Larger rollout pools and selective training improve exploration but shift computational cost to candidate generation, while most generated samples are discarded.This redundancy motivates low-bit inference acceleration.
Sol-RL Framework: Sol-RL separates high-throughput FP4 exploration from selective BF16 rollout to accelerate scaling while avoiding quantization-induced corruption.The pipeline caption reports up to 2.4× acceleration versus naive scaling with 2% computational overhead.
Quantized Rollout Risks: Directly using quantized rollouts for optimization can cause severe alignment degradation and training instability because of an off-policy distribution shift.The quantized policy’s trajectories differ from those of the high-precision target policy.
Quantized Rollout Risks: Low-bit samples can distort semantics when used as direct regression targets, causing high-precision policies to mimic noisy, low-fidelity outputs.This is especially problematic for diffusion objectives based on denoising score matching.
Proxy Reward Ranking: NVFP4 exploration preserves enough semantic structure and intra-group reward ordering to serve as a proxy for selecting contrastive samples.The dense diagonal rank distribution supports reliable identification of Top-K and Bottom-K candidates.
Stage 1: FP4 Exploration: NVFP4 generates a large pool using reduced-step ODE sampling, then filters it to K high-contrastive seeds using proxy rewards.The exploration stage can use examples such as N=96, six steps, and K=24.
Stage 2: BF16 Optimization: Selected seeds are regenerated in BF16, shielding the diffusion vector field from low-precision quantization during target generation and policy updating.The two-stage pipeline confines FP4 to exploration and reserves BF16 for optimization.

4. Experiments

Experiments across SANA, FLUX.1, and SD3.5-L evaluate Sol-RL’s two-stage rollout design, alignment performance, scaling behavior, semantic fidelity, and efficiency. Sol-RL consistently improves alignment under matched budgets while preserving high-fidelity optimization.

Experimental Setup: Sol-RL is evaluated on SANA, FLUX.1, and SD3.5-L using ImageReward, CLIPScore, PickScore, and HPSv2 objectives.The experiments use NVFP4 through NVIDIA Transformer Engine on 8 NVIDIA B200 GPUs.
Experimental Setup: The rollout pipeline generates 96 NVFP4 candidates in 6 steps, selects the top-12 and bottom-12, then regenerates 24 samples in BF16.The policy is optimized on these 24 high-fidelity rollouts.
Main Results: Sol-RL consistently surpasses DiffusionNFT across diverse models and reward metrics under identical GPU-hour budgets.On FLUX.1, it also achieves superior alignment across all evaluated metrics against FlowGRPO, DanceGRPO, AWM, and DiffusionNFT.
Main Results: 1.91× to 4.64× faster convergence is achieved relative to DiffusionNFT’s equivalent performance, with a higher final alignment level.Figure 4 reports this pattern across foundation models and reward metrics under matched wall-clock budgets.
Main Results: Sol-RL improves SANA’s complex detail rendering and semantic alignment when optimizing across multiple rewards compared with the unfine-tuned base model.The visual comparison uses HPSv2, PickScore, CLIPScore, and OCR rewards.
Ablation Experiments: Four-step FP4 exploration produces suboptimal alignment, while extending exploration beyond T=6 yields no further final-reward improvement.The results attribute the low-step deficit to insufficient semantic layouts and inaccurate intra-group ranking.
Ablation Experiments: Increasing the exploration pool N from 24 to 96 while fixing K=24 continuously improves final scores.The broader pool provides more opportunities to discover high-contrastive samples.
Efficiency and Training Integrity: NVFP4 rollouts preserve semantic structure and achieve on-par Inception Score and CLIP scores with BF16 baselines.Sol-RL maintains the alignment fidelity of BF16 brute-force sampling by regenerating the selected subset in BF16 for policy updates.

5. Related Work

Related work spans diffusion reinforcement learning objectives, rollout-based GRPO methods, diffusion preference optimization, model quantization, and quantized reinforcement-learning rollouts. These lines establish both the scalability context and the off-policy risks motivating Sol-RL.

Diffusion Reinforcement Learning and Preference Optimization: ReFL, DRaFT, and Adjoint Matching optimize diffusion-model rewards through predicted or sampled images while addressing approximation or memory costs.These methods represent gradient-based diffusion reinforcement-learning approaches with different efficiency mechanisms.
Diffusion Reinforcement Learning and Preference Optimization: DDPO and DPOK formulate diffusion RL as multi-step decision-making, while Flow-GRPO and DanceGRPO combine this formulation with GRPO.GRPO replaces PPO’s value-model baseline with the group-wise mean reward.
Diffusion Reinforcement Learning and Preference Optimization: Forward-process methods include offline and online reward-weighted denoising objectives, Diffusion-DPO, FMPG, and AWM.These approaches optimize diffusion models without all relying on the same rollout formulation.
Model Quantization and Quantized RL: Quantization research spans INT8, 4-bit, and 2–3-bit methods designed to reduce inference cost while maintaining model fidelity.Representative techniques include GPTQ, AWQ, learnable equivalent transformations, and sparse-quantized representations.
Model Quantization and Quantized RL: FlashRL and QeRL demonstrate speedups from quantized rollouts, while prior work links quantized sampling to off-policy distribution shifts and numerical discrepancies.QuRL addresses these vulnerabilities with adaptive clipping.

6. Conclusion

The conclusion frames Sol-RL as a two-stage solution to diffusion RL’s efficiency–stability dilemma. NVFP4 handles large-scale exploration, while BF16 regeneration preserves optimization stability and alignment fidelity.

Conclusion: Sol-RL confines NVFP4 generation to large-scale exploration and reserves BF16 computation for regenerating selected high-contrastive samples.This decouples exploration efficiency from optimization stability.
Conclusion: The framework integrates rollout scaling and selective training with NVFP4 throughput gains while avoiding direct optimization on low-bit samples.The conclusion presents this as the paper’s central efficiency–stability resolution.

A. Theoretical Justification

Sol-RL’s theoretical analysis models FP4 rollout errors as bounded perturbations and shows that rollout scaling preserves contrastive reward signals despite quantization noise.

Bounded Perturbation: FP4 rounding and low-precision solver arithmetic are modeled as an effective perturbation of the high-precision diffusion trajectory.The analysis assumes Lipschitz continuity of the vector field and reward model to bound trajectory and reward deviations.
Bounded Perturbation: Δ provides a static upper bound on cross-precision reward discrepancy that is independent of candidate pool size N.This separates quantization error from the benefits of increasing rollout scale.
Extreme-Value Guarantee: The low-precision screening mechanism selects empirical best and worst candidates using proxy rewards whose errors are bounded by Δ.The selected candidates’ true reward range is then lower-bounded using these perturbation limits.
Extreme-Value Guarantee: Increasing N eventually lets the extreme contrastive bounds overpower constant quantization noise, preserving the gradient signals needed for alignment.The analysis connects the retained gradient signal directly to rollout scaling.
Extreme-Value Guarantee: The quantization penalty contracts the reward margin by at most 4Δ, while the extreme-value advantage grows with rollout scale N.The expected advantage increases according to the extreme-value behavior of the reward distribution.

B.1. Training Hyperparameters

The experiments use deterministic ODE sampling across SANA, FLUX.1, and SD3.5, with model-specific guidance settings and a two-stage NVFP4 rollout configuration.

Training Hyperparameters: All rollouts use deterministic ODE sampling across the three diffusion models.Classifier-free guidance is disabled for SANA and SD3.5, while FLUX.1 receives a guidance embedding of 1.0.
Training Hyperparameters: Stage 1 draws N = 96 noise vectors per prompt and generates candidates with NVFP4 using T = 6 denoising steps.Candidates are scored and filtered by selecting top-K/2 and bottom-K/2 samples.

B.3. Reward Models and Evaluation

The study trains against four independent reward objectives and evaluates on held-out prompts using all four metrics.

Reward Models: The evaluation uses ImageReward, CLIPScore, PickScore, and HPSv2 as alignment objectives and metrics.These cover overall visual quality, semantic alignment, pairwise preference, and human preference scoring.
Evaluation: Reward models are used independently during training, while evaluation uses a held-out prompt subset and all four metrics.Training prompts come from the PickScore training split, with a separate subset reserved for evaluation.

C. Additional Analysis of NVFP4 Exploration

Additional analysis supports NVFP4 exploration by assuming preserved intra-group rankings and reporting strong global and extreme-candidate agreement with BF16.

Ranking Consistency: The framework assumes NVFP4 may perturb absolute rewards slightly while preserving intra-group relative candidate rankings.This ranking property is the basis for using NVFP4 as an exploration proxy rather than an optimization target.
Ranking Consistency: Spearman’s ρ and Kendall’s τ assess global agreement between NVFP4 and BF16 reward rankings.The analysis interprets ρ above 0.80 as very strong positive correlation and τ above 0.70 as highly consistent ordinal association.
Extreme-Candidate Filtering: NVFP4 identifies contrastive extremes with over 96% Top-4 precision and less than 4% Bottom-4 false inclusion.Top/Bottom-k Match measures exact overlap of highest and lowest candidates selected under BF16 versus NVFP4.
Ranking Consistency: Table 8 compares global reward correlations and Top/Bottom-k exact match rates between FP4-accelerated samples and BF16 baselines.The table is used to assess whether FP4 exploration preserves intra-group relative rankings.