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

• URL: http://arxiv.org/abs/2604.06916v1
• Date: Wed, 08 Apr 2026 10:14:47 GMT
• Title: FP4 Explore, BF16 Train: Diffusion Reinforcement Learning via Efficient Rollout Scaling
• Abstract summary: Reinforcement-Learning-based post-training has emerged as a promising paradigm for aligning text-to-image diffusion models with human preferences.<n> scaling rollouts on large-scale foundational diffusion models (e.g., FLUX.1-12B) imposes a heavy computational burden.<n>We propose Sol-RL, a novel FP4-empowered Two-stage Reinforcement Learning framework.
• Score: 38.64059734487925
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.64\times$, unlocking the power of massive rollout scaling at a fraction of the cost.

Related papers

• QUADS: Stabilizing NVFP4 Reinforcement Learning for MoE via QUantization-error Alignment across Dual Sides [26.939321070753607]
  Rollout generation is a major bottleneck in Reinforcement Learning for Mixture-of-Experts (MoE) Large Language Models.<n>We propose QUantization-error Alignment across Dual Sides (QUADS) to stabilize NVFP4 RL for MoE.<n>QUADS achieves BF16-level accuracy, improves average pass@1 by 21.49 points over naive NVFP4 RL, and delivers 16% higher rollout throughput than FP8.
  arXiv  Detail & Related papers  (2026-07-17T10:21:27Z)
• Full-Stack FP4: Stable LLM Pretraining with Quantized Projections, Optimizers, and Attention [34.07319488267708]
  Recent NVFP4 pretraining methods mainly target transformer linear layers, leaving computation states, arithmetic and attention underexplored in 4-bit pipelines.<n>This critical gap blocks stable full-stack 4-bit pretraining, as the three core modules exhibit unique numerical failure patterns.<n>We propose Full-Stack FP4, the first complete NVFP4 pretraining framework resolving all three stability bottlenecks.
  arXiv  Detail & Related papers  (2026-07-05T17:31:36Z)
• Convex Optimization for Alignment and Preference Learning on a Single GPU [52.997197698288936]
  Fine-tuning large language models to align with human preferences has driven the success of systems such as Gemini and ChatGPT.<n>Direct Preference Optimization (DPO) offers a simpler alternative but has limitations such as inconsistent ranking accuracy and high dependence on GPU resources.<n>We propose the Convex Optimization for Alignment and Preference Learning Algorithm (COALA), a novel lightweight strategy with strong theoretical guarantees.
  arXiv  Detail & Related papers  (2026-05-22T05:25:00Z)
• DiFFPO: Training Diffusion LLMs to Reason Fast and Furious via Reinforcement Learning [37.20873499361773]
  We propose a unified framework for training masked diffusion large language models (dLLMs) to reason better (furious)<n>We first unify the existing baseline approach by proposing to train surrogate policies via off-policy RL, whose likelihood is much more tractable as an approximation to the true dLLM policy.<n>We also propose a new direction of joint training efficient samplers/controllers of dLLMs policy. Via RL, we incentivize dLLMs' natural multi-token prediction capabilities by letting the model learn to adaptively allocate an inference threshold for each prompt.
  arXiv  Detail & Related papers  (2025-10-02T16:57:24Z)
• Pretraining Large Language Models with NVFP4 [53.235038214986865]
  We introduce a novel approach for stable and accurate training of large language models (LLMs) using the NVFP4 format.<n>Our method integrates two-dimensional quantization scheme for consistent representations across both the forward and backward passes.<n>Our results show that the model trained with our NVFP4-based pretraining technique achieves training loss and downstream task accuracies comparable to an FP8 baseline.
  arXiv  Detail & Related papers  (2025-09-29T17:53:17Z)
• Towards High Data Efficiency in Reinforcement Learning with Verifiable Reward [54.708851958671794]
  We propose a Data-Efficient Policy Optimization pipeline that combines optimized strategies for both offline and online data selection.<n>In offline phase, we curate a high-quality subset of training samples based on diversity, influence, and appropriate difficulty.<n>During online RLVR training, we introduce a sample-level explorability metric to dynamically filter samples with low exploration potential.
  arXiv  Detail & Related papers  (2025-09-01T10:04:20Z)
• Shuffle-R1: Efficient RL framework for Multimodal Large Language Models via Data-centric Dynamic Shuffle [65.14124923451077]
  Reinforcement learning (RL) has emerged as an effective post-training paradigm for enhancing the reasoning capabilities of multimodal large language model (MLLM)<n>However, current RL pipelines often suffer from training inefficiencies caused by two underexplored issues: Advantage Collapsing and Rollout Silencing.<n>We propose Shuffle-R1, a simple yet principled framework that improves RL fine-tuning efficiency by dynamically restructuring trajectory sampling and batch composition.
  arXiv  Detail & Related papers  (2025-08-07T17:53:47Z)
• FP4 All the Way: Fully Quantized Training of LLMs [26.195547788434908]
  We demonstrate fully quantized training (FQT) of large language models (LLMs) using predominantly 4-bit floating-point (FP4) precision.<n>We investigate key design choices for FP4, including block sizes, scaling formats, and rounding methods.
  arXiv  Detail & Related papers  (2025-05-25T12:14:25Z)
• Optimizing Large Language Model Training Using FP4 Quantization [73.55459961002371]
  Quantized training presents a promising solution by enabling low-bit arithmetic operations to reduce costs.<n>This work introduces the first FP4 training framework for large language models (LLMs)
  arXiv  Detail & Related papers  (2025-01-28T18:04:50Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.