What Matters in RL for Diffusion Models?

The dominant role of noise in RL post-training for diffusion models

What Matters in RL for Diffusion Models?

This page describes What Matters in RL for Diffusion Models? The Dominant Role of Noise, submitted to NeurIPS 2026.

Reinforcement learning has become a powerful paradigm for post-training generative models — improving reasoning in language models and improving alignment in diffusion models. But a basic question has gone unanswered: what actually drives learning in RL-based diffusion training?

For language models, the signal is clear — it comes from the response. But diffusion models introduce a second axis of stochasticity: noise initialization. Each generated image is shaped both by the text prompt and by the noise sample that seeds the denoising trajectory. These two axes both produce variation in reward, but their relative contributions have never been disentangled — until now.

PickScore training curve: 2x Noise Margin vs Flow-GRPO Advantage variance vs improvement scatter plot

We present the first systematic study separating prompt-level and noise-level sources of optimization signal in RL training for diffusion models. The key finding: noise dominates. Reward variance and policy-gradient informativeness are driven overwhelmingly by differences among trajectories generated from the same prompt — not by differences across prompts. Prompt-level variation contributes substantially less after group-wise normalization.

The diagram below shows where the two sources of stochasticity enter generation, and why group-wise normalization leaves the noise axis as the dominant source of signal.

Two axes of stochasticity — and which one carries the signal Prompt p one text prompt noise seeds z₁ z₂ z₃ denoising trajectory image x₁ image x₂ image x₃ r₁ = 0.82 r₂ = 0.31 r₃ = 0.55 Group-wise normalization Aᵢ = rᵢ − r̄ Noise axis (same prompt, different seed): the spread among r₁, r₂, r₃ survives normalization and supplies nearly all of the gradient signal. Prompt axis (different prompts): their differences enter only through the group mean r̄ — and are subtracted away.

Figure: Group-wise normalization cancels prompt-level reward differences, leaving noise-induced spread within each group as the training signal.

This has a direct practical implication: if noise is what matters, training should be structured to maximize informative noise-induced variation. We introduce two strategies:

  • Structured noise oversampling: generate more candidate trajectories per prompt, exploiting the noise axis more efficiently under the same compute budget.
  • Margin-based trajectory selection: select training pairs with the largest reward margin, prioritizing trajectories that carry the most learning signal.

The diagram below shows how the two strategies restructure each training batch.

1 · Structured noise oversampling 2 · Margin-based selection FlowGRPO baseline prompt 4 trajectories Noise oversampling (2×), same compute prompt 8 trajectories from the same prompt Fewer prompts, more seeds per prompt: each group probes the noise axis more densely and produces a wider within-group reward spread. reward largest margin Δr FlowGRPO update on the selected pair trajectories with small reward margins are discarded

Figure: Noise oversampling widens each prompt's candidate pool; margin-based selection keeps the pair with the largest reward gap for the FlowGRPO update.

Both strategies plug directly into FlowGRPO, the standard RL framework for flow-based diffusion models. Across three tasks — compositional image generation, OCR-based text rendering, and human preference alignment — noise margin selection consistently improves over the baseline. The 2x Noise Margin variant reaches FlowGRPO’s final PickScore in 2.2× fewer training steps, and ends with a higher final score (23.14 vs. 22.92). Prompt-level filtering, by contrast, yields substantially weaker gains — confirming the asymmetry.

Method diagram Results across tasks

The broader takeaway is that RL for diffusion models is governed by trajectory-level exploration and selection over the noise-induced space — a fundamentally different structure from RL alignment in autoregressive language models, where the prompt is the primary axis of variation. Understanding this distinction matters for how we design data collection, sampling strategies, and reward-weighting schemes for the next generation of RL-trained generative models.

Next Project

Improving Activation Steering

 Improving Activation Steering