PaperlayerView exact revision on arXiv

Source-linked AI summary

Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View

Yixian Xu, Yuanrui Zhang, Shengjie Luo, Liwei Wang, Di He

arXiv:2608.14430v1cs.LGcs.CVstat.ML

TL;DR

Diffusion-RL methods use structurally different reverse-trajectory and forward-matching losses without a unified account. This paper derives a path-space framework linking them through variance reduction and reports improved performance over prior baselines on SD3.5-M and Qwen-Image.

  • Problem

    Diffusion-RL methods use structurally different reverse-trajectory and forward-matching losses, while a unified framework for understanding and designing them remains lacking.

  • Method

    The paper applies path-space importance sampling to derive a variance-reduced value-gradient objective unifying existing methods, then proposes multi-sample KDE estimation and scale-bounded weights.

  • Results

    Experiments on SD3.5-M and Qwen-Image across PickScore, OCR, and GenEval verify the variance-reduction explanation and exceed prior diffusion-RL baselines.

  • Takeaways & Limitations

    Reverse-trajectory and forward-matching diffusion-RL algorithms can be treated as instances of one path-space objective and organized by shared design choices.

  • Takeaways & Limitations

    The framework does not determine which of four weight-shape families is optimal, so that choice remains empirical.

Abstract

from arXiv · show

Reinforcement learning (RL) post-training provides a direct way to align diffusion models with human preferences and task-specific rewards. However, current RL algorithms for diffusion models remain fragmented: reverse-trajectory methods rely on discretized likelihood ratios, whereas forward-matching methods train on reward-labeled noising versions of the rollout samples. This paper shows that these seemingly different losses arise from a single path-space principle. Starting from the regularized diffusion-RL objective, we use importance sampling between sampling SDEs to obtain an explicit policy-gradient estimator on trajectory space. The estimator contains the stochastic Itô integral underlying Flow-GRPO-type updates; we derive an equivalent variance-reduced value-gradient form that recovers the forward-matching structure of AWM and DiffusionNFT. This identifies the empirical gap between these method families as a variance-reduction effect rather than a difference in RL principle. The derivation yields a unified design space organized by value-gradient estimation, weight functions, and sampling choices. Within this space, we propose a multi-sample KDE value-gradient estimator that reuses rollout groups, together with scale-bounded weight families that retain stable existing recipes while excluding singular ones. Experiments on SD3.5-M and Qwen-Image models validate the variance-reduction explanation and show that the resulting recipe improves over prior diffusion-RL baselines.

1 Introduction

The introduction frames diffusion-RL methods as fragmented despite sharing a common path-space principle, then presents a unified framework, variance-reduced estimators, design guidance, and empirical validation. It proposes multi-sample value-gradient estimation and scale-bounded weights to improve efficiency and performance over prior baselines.

  • Motivation: Diffusion and flow models achieve high-quality visual generation but require RL post-training to align outputs with human preferences and task-specific objectives.Target objectives include aesthetic preference, compositional correctness, text rendering, and task-specific verifiability.
  • Problem: Existing diffusion-RL methods split between reverse-time likelihood-ratio approaches and forward-matching losses, leaving their relationship and principled design unclear.Reverse methods discretize denoising transitions, while forward methods use reward-labeled clean samples and forward-noised versions.
  • Unified framework: Path-space importance sampling between SDE-induced measures yields a tractable continuous-time policy-gradient estimator whose Itô-integral term underlies reverse-trajectory updates.The construction applies when the SDEs share the same diffusion coefficient.
  • Unified framework: Variance reduction converts the estimator into a value-gradient form that recovers AWM and DiffusionNFT under forward noising and Flow-GRPO-type methods under reverse sampling.The design space is organized by value-gradient estimators, time-dependent weights, and sampling choices.
  • Design choices: A multi-sample KDE estimator reuses GRPO rollout groups to reduce conditional variance without additional sampling, while scale-bounded weights recover stable existing recipes.These choices expose remaining tunable degrees of freedom for update and regularization weighting.
  • Empirical validation: Experiments on SD3.5-M and Qwen-Image across PickScore, OCR, and GenEval verify the variance-reduction explanation and show that the final recipe exceeds baselines.The evaluation also finds that proposed weight families track empirical optima and close the Flow-GRPO–DiffusionNFT efficiency gap.

2 Backgrounds

Diffusion and flow-matching models transform Gaussian noise into data through reverse denoising, while diffusion-RL post-training maximizes reward relative to a frozen reference policy. Existing methods either use reverse-trajectory likelihood ratios or forward reward-labeled noising paths, motivating a unified account.

  • Diffusion and Flow Matching: Flow-matching trains a vector field on the interpolation x_t = (1−t)x_0 + tϵ, then samples by simulating a reverse-time SDE that becomes a deterministic ODE when η_t = 0.The noise satisfies ϵ ∼ N(0, I), and η_t = 1 gives the time-reversal of the forward process.
  • Diffusion-RL Objective: The canonical diffusion-RL objective maximizes clean-sample reward while regularizing the trajectory policy against a frozen reference policy.Rewards are defined on clean samples x_0 for prompts c, and trajectories are represented as continuous stochastic paths.
  • Existing Diffusion-RL Methods: Flow-GRPO discretizes the reverse-time sampling SDE into Gaussian denoising transitions and applies clipped PPO using per-step likelihood ratios and group-normalized rewards.The current policy p_θ is compared with the old policy p_old through Euler–Maruyama per-step ratios, with clipping applied to each step.
  • Existing Diffusion-RL Methods: AWM and DiffusionNFT instead construct training signals from reward-labeled clean samples passed through the forward noising path, using advantage weighting or positive and negative velocity mixtures.AWM reuses the flow-matching target with advantage weighting, while DiffusionNFT defines v+ and v− around the old model using β.
  • Open Problem: These reverse-trajectory and forward-matching methods appear structurally unrelated, leaving the field without a unified principle for interpreting existing recipes or designing new diffusion-RL algorithms.The passage identifies the missing unified view as a central problem motivating the paper.

3 A Unified Path-Space View of Diffusion RL

The section derives a unified path-space template for diffusion RL by replacing intractable trajectory likelihoods and high-variance stochastic integrals with tractable, variance-reduced value-gradient forms. It organizes existing methods and new designs around value-gradient estimators, weight functions, and sampling choices, including KDE estimators and scale-bounded weights.

  • Path-Space Importance Sampling: Path-space importance sampling yields an explicit trajectory-space policy-gradient estimator despite the absence of closed-form diffusion path densities.The method compares SDE path measures sharing the same diffusion coefficient and samples full trajectories from a tractable proposal.
  • Variance Reduction: The stochastic-integral estimator contains noise exploration and an advantage-weighted quadratic, but its Itô-integral term dominates variance.The variance-reduced template replaces this stochastic integral with a deterministic integral depending on a value-gradient estimator.
  • A Unified Framework: Diffusion RL methods share a two-term form combining an off-policy quadratic in ∆vθ with an on-policy inner product using a chosen value-gradient estimator.The estimator, off-policy weight, and on-policy weight act as free parameters defining a unified design space.
  • Value-Gradient Estimation: A multi-sample KDE estimator is unbiased at h = 1 and reduces variance relative to the deterministic one-sample estimator, while h ≠ 1 generally introduces bias.The bandwidth h trades exact posterior matching for smoother finite-sample estimates.
  • Weight Functions: The scale-bounded principle requires both weighted value-gradient and weighted quadratic terms to remain uniformly bounded over t, while the optimal trust-region shape remains empirical.Methods outperforming Flow-GRPO satisfy the first boundedness condition, whereas Flow-GRPO violates it.

4 Experiments

Experiments on SD3.5-Medium and Qwen-Image validate the framework’s variance-reduction explanation and show that estimator-specific reweighting is essential for stable, fast training. The final KDE-based recipe improves convergence over prior diffusion-RL baselines across evaluated rewards.

  • Setup: Experiments use SD3.5-Medium and Qwen-Image at 512 × 512, with group size G = 24, 10 training denoising steps, and 40 evaluation steps.Models are fine-tuned with LoRA and without classifier-free guidance for fair comparison.
  • Variance reduction: The one-sample estimator with inherited weighting is unstable and fails, whereas its reweighted form converges remarkably faster, showing variance reduction alone is insufficient.The experiment motivates estimator-specific weight design.
  • Variance reduction: 13.8×: the stochastic estimator has 13.8× the variance of the deterministic one-sample estimator under matched SD3.5-Medium PickScore conditions.The measurement uses four prompts, 128 trajectories per prompt, and 40 denoising steps.
  • Weight ablations: The preferred weights depend on the estimator: stochastic estimators favor w2 = t, while larger α values stabilize deterministic estimators.Different estimators exhibit different weight preferences in the on-policy ablation.
  • Weight ablations: The off-policy weight w1(t) = (1 −t)^α balances regularization and policy-gradient terms, with low-noise-region regularization improving training stability.Alternative weight families either overemphasize regularization or cause instability.
  • Baseline comparison: On OCR reward, the KDE recipe converges 2× faster than AWM and 3× faster than DiffusionNFT.The comparison uses the KDE estimator with its selected optimal weighting across several rewards.
  • Ablation study: Replacing DiffusionNFT’s deterministic estimator with KDE improves stability but hurts performance, while adding the proposed on-policy weight stabilizes and improves training.The ablation shows that variance reduction alone does not explain the full gain.

5 Conclusion · A Organization of Appendix

The conclusion presents a unified path-space framework for diffusion-RL algorithms, deriving a variance-reduced general objective and design space. The appendix provides theoretical background, proofs, derivations, related work, experimental details, additional results, and qualitative examples.

  • 5 Conclusion: The work develops a unified path-space framework placing previous diffusion-RL algorithms on a common theoretical foundation.
  • 5 Conclusion: Starting from the original RL objective, the paper derives a general training objective based on its variance-reduced form.
  • 5 Conclusion: Previous reverse-trajectory and forward-matching methods are presented as instances of the general objective.
  • 5 Conclusion: The derivation yields a unified design space that includes value-gradient estimation and weight choices.
  • A Organization of Appendix: Appendix B provides background in stochastic calculus, while Appendix C contains proofs of the main-text theorems.
  • A Organization of Appendix: Appendix D gives the derivations in Table 1, and Appendix E discusses related works.
  • A Organization of Appendix: Appendix F provides experimental details for Section 4 and shows additional results, while Appendix G presents qualitative examples.

B Backgrounds for Stochastic Analysis … C.3 Proof of Thm. 2

The paper establishes the stochastic-calculus and path-space foundations for diffusion-RL analysis, then proves both the importance-sampling policy-gradient estimator and its variance-reduced value-gradient representation. The proofs use Girsanov’s theorem, Itô calculus, martingale identities, and the Kolmogorov backward equation under explicit regularity and trust-region conditions.

  • B Backgrounds for Stochastic Analysis: Stochastic processes are treated as measurable path-valued random objects, with Itô processes characterized by SDEs and their generators.The background introduces stochastic processes on a probability space and defines the generator for an Itô process.
  • B Backgrounds for Stochastic Analysis: Itô’s lemma supplies the stochastic chain rule used to differentiate smooth functions along SDE solutions.The lemma applies to C2 functions when x_t solves an SDE with drift and diffusion terms.
  • B Backgrounds for Stochastic Analysis: The value function V_t(x) = E[R(x_T, c)|x_t = x] satisfies a Kolmogorov backward equation governed by the SDE generator.This PDE is the basis for the later value-gradient derivation.
  • B Backgrounds for Stochastic Analysis: Girsanov’s theorem provides the Radon–Nikodym derivative between diffusion path measures with shared initial conditions and diffusion coefficients.The theorem requires sufficiently smooth drifts satisfying Novikov’s condition.
  • C.1 Setup: path-space importance sampling: The path-space setup represents the base, current, and reference velocity fields by corresponding path measures, transforming the original objective for importance sampling.The Radon–Nikodym derivative of Pθ with respect to Q is evaluated using Girsanov’s theorem and the time change r = 1 − t.
  • C.2 Proof of Prop. 1: Under a trust-region condition controlling the exponential likelihood increment and gradient moment, the policy gradient admits a stochastic-integral path-space estimator.The proposition assumes clipping or a small step size ensures | exp(M1(θ−)) − 1| < ϵ and a bounded expected gradient norm.
  • C.3 Proof of Thm. 2: The variance-reduced estimator conditions the terminal advantage on intermediate states through a value function and is proved using forward sampling time, martingales, and Itô’s product rule.The proof uses y_r = x_{1−r}, identifies the cross-variation of martingales, and carries the quadratic term unchanged.
  • C.3 Proof of Thm. 2: The proof concludes by taking expectations, changing variables back to FM-time, restoring the quadratic term, and establishing the stated identity.The resulting expression is exactly Eq. (12), while the subsequent argument uses the tower property, Jensen’s inequality, and Fubini’s theorem for Eq. (13).

C.4 Proof of Prop. 3

The section proves that the value-function gradient has a closed-form conditional-expectation expression involving the diffusion score and advantage. Under the reverse-forward-noising setting, this expression specializes to a flow-matching conditional-velocity form using group-relative advantages.

  • Closed-form gradient: The value-function gradient equals an advantage-weighted conditional score expectation minus a value-weighted unconditional score expectation.This is the closed-form expression established in Proposition 11.
  • Flow-matching specialization: With η_t ≡ 1 and v_base = E[v | x_t] ≈ v_old, the result specializes to the flow-matching conditional velocity v(x_t, x_0) := (x_t − x_0)/t.Under these conditions, q is the reverse of the forward noising process.
  • Proof setup: The proof uses q(x_t | x_0) = N(x_t; (1 − t)x_0, t^2I) and differentiates the value function under the reverse-noising sampling SDE.The assumption η_t ≡ 1 reduces the sampling SDE to the reverse of the forward noising process.
  • Advantage-weighted velocity form: Substituting the Gaussian score and replacing R with the group-relative advantage yields an expectation of A(x_0, c) times the centered conditional velocity.The advantage can be viewed as a normalized reward function, and the centered term is v(x_t, x_0) − E[v(x_t, x_0) | x_t].

C.5 Proof of Prop. 4

The proposition establishes unbiasedness for the stochastic and deterministic one-sample value-gradient estimators under their respective assumptions, and shows that the KDE estimator is unbiased with lower variance under stronger conditions.

  • Proof of (i): The discretized stochastic one-sample estimator is asymptotically unbiased as the discretization step tends to zero.The proof uses the value martingale, conditional expectations, the Itô isometry, and the continuous-time limit.
  • Proof of (ii): The deterministic one-sample estimator is unbiased when η_t ≡ 1 and v_old(x_t, t) equals E[v(x_t, x_0) | x_t].Under these assumptions, the conditional expectation reduces to the relevant gradient expression.
  • Proof of (iii): When η_t ≡ 1 and h = 1, the KDE estimator is unbiased and has smaller variance than the deterministic one-sample estimator.The proof identifies the KDE estimator as the conditional expectation of the deterministic anchor estimator and applies conditional variance.
  • Proof of (iii): At h = 1, Bayes’ rule yields a Rao–Blackwell structure in which the KDE estimator conditions on an endpoint while the remaining group endpoints remain independent proposal samples.Symmetry of the KDE estimator permits an auxiliary uniformly sampled index without changing its distribution.

D Results in Table 1 … F Experimental Details

The paper derives existing diffusion-RL methods within a unified path-space design space, showing how reverse-sampling and forward-matching objectives correspond to specific value-gradient estimators and weights. The appendices connect Flow-GRPO variants, AWM, and DiffusionNFT to this framework and situate them among related approaches.

  • D Results in Table 1: The section provides detailed derivations of the results reported in Table 1.
  • D.1 Flow-GRPO, GRPO-Guard and TempFlow-GRPO: Flow-GRPO is obtained by discretizing the stochastic-integral path-space estimator with Euler–Maruyama.The derivation produces the Flow-GRPO expressions in Table 1.
  • D.1 Flow-GRPO, GRPO-Guard and TempFlow-GRPO: GRPO-Guard differs from Flow-GRPO by dropping the regularization term and normalizing Term (A), which changes its weight functions.
  • D.1 Flow-GRPO, GRPO-Guard and TempFlow-GRPO: TempFlow-GRPO uses an ODE rollout with single-SDE-step branching, endpoint ODE rollouts, and reweighting by gt.The resulting trajectories form a branching structure called a branched-SDE sampler.
  • D.2 AWM and DiffusionNFT: AWM is equivalent to choosing vbase(xt, t) = E[v(xt, x0) | xt], ηt ≡ 1, and the deterministic one-sample estimator in the unified design space.Its reverse sampling process is exactly the time reversal of the corresponding forward noising path, so both constructions share the same joint distributions.
  • D.2 AWM and DiffusionNFT: DiffusionNFT is recovered by choosing vbase = E[v | xt], ηt = 1, the deterministic one-sample estimator, and weights ˆw1(t) = 1/∆t and ˆw2(t) = 2/(βs(t)∆t).
  • E Related Work: Related work spans reverse-sampling RL, forward-matching RL, stochastic-control and adjoint methods, reward-gradient optimization, and offline preference methods.The paper further decomposes the objective axis into value-gradient estimation and associated on- and off-policy weights.
  • F Experimental Details: Sampling-noise ablations compare stochastic one-sample, deterministic one-sample, and KDE estimators across ηt ∈ {0.005, 0.225} on the Pickscore reward, while another figure visualizes different weights.

F.1 Additional Results · F.2 Hyperparameters

Additional results show estimator-specific sensitivity to sampling-noise settings: stochastic GRPO estimators benefit from larger η_t, deterministic estimators can become unstable, and the proposed estimator remains stable across scales. The appendix also visualizes weight functions and records experimental hyperparameters, including disabled CFG and batch-level advantage normalization.

  • F.1 Additional Results: Stochastic estimators in GRPO-based methods perform better with the larger η_t sampling-noise setting.This result is shown in Fig. 5(a).
  • F.1 Additional Results: Deterministic estimators suffer training instability in both sampling-noise settings.The instability is reported across the two settings evaluated in Fig. 5.
  • F.1 Additional Results: The small noise scheduler performs slightly better for the deterministic estimator.This comparison is reported in Fig. 5(b).
  • F.1 Additional Results: The proposed estimator is insensitive to different sampling-noise scales and remains stable across all tested settings.The stability claim accompanies the sampling-noise schedule ablations in Fig. 5.
  • F.1 Additional Results: Figure 6 visualizes the different weight functions used in the experiments.The passage identifies the figure as a visualization of weight functions without reporting additional numerical results.
  • F.2 Hyperparameters: All experiments disable CFG and use batch-level advantage normalization.These implementation settings apply across the reported experiments.
  • F.2 Hyperparameters: Table 2 summarizes the hyperparameters used throughout Section 4.The supplied passage does not enumerate the table entries.

F.3 Robustness across random seeds · F.4 Runtime comparison · G Qualitative Examples

Across random seeds, the method maintains a higher PickScore with low variation, while its KDE-based update reduces runtime by avoiding an extra denoiser evaluation. Qualitative evaluation also favors the method over GRPO-Guard and DiffusionNFT.

  • F.3 Robustness across random seeds: Our method achieves a higher mean PickScore at every checkpoint across three random seeds, with small run-to-run variation.All other experimental configurations remain fixed across the repeated runs.
  • F.3 Robustness across random seeds: At the final checkpoint, our method reaches 23.56 ± 0.01461 PickScore versus 23.22 ± 0.01274 for DiffusionNFT.The relative standard deviations are 0.062% and 0.055%, respectively.
  • F.4 Runtime comparison: The KDE estimator computes pairwise distances in O(G2d) time and reuses cached rollout latents, velocities, and rewards without additional denoiser evaluation.With G = 24 and latent dimension d = 65,536, the computation is on the order of 107 arithmetic operations per timestep and is efficiently parallelized on GPUs.
  • F.4 Runtime comparison: Under matched SD3.5-Medium settings, our method saves 18 seconds per epoch, a 15.7% wall-clock reduction.The measured comparison is reported in Table 4.
  • G Qualitative Examples: Visualization grids compare GRPO-Guard, DiffusionNFT, and our method on identical test prompts and inference settings across three prompt splits.The grids support direct inspection of prompt alignment, perceptual quality, excessive saturation, and repetitive reward-hacking patterns.
  • G Qualitative Examples: In a blinded evaluation by eight evaluators over 30 prompts, our method receives a 51.4% preference share versus 29.4% for GRPO-Guard and 19.2% for DiffusionNFT.Evaluators judged overall visual quality, prompt alignment, and aesthetic preference.

H Limitations · I Social Impact

The study’s findings are limited by evaluation on two image models, one fixed group size, and three reward functions, leaving broader generalizability untested. Its social-impact discussion is brief, noting potential alignment with human preferences and safety constraints while emphasizing the work’s primarily theoretical and algorithmic focus.

  • H Limitations: Experiments use SD3.5-M and Qwen-Image, limiting direct evidence to these two diffusion-model architectures.Generalizability to other architectures, including video diffusion models, was not systematically evaluated.
  • H Limitations: All experiments use a fixed group size of G = 24, so the effects of alternative group sizes remain untested.The stated experimental configuration fixes G at 24 throughout the reported studies.
  • H Limitations: The reward evaluation covers PickScore, OCR, and GenEval but does not systematically test human-preference or verifiable-task rewards.These three rewards represent a range of reward types, but broader reward-function generalizability remains unresolved.
  • I Social Impact: The framework is intended to support alignment of generative models with human preferences and safety constraints.The paper presents this contribution as supporting the responsible development of generative AI systems.
  • I Social Impact: The paper does not claim a broader social impact because it mainly focuses on theory and algorithms.The social-impact discussion explicitly limits its scope to the work’s theoretical and algorithmic emphasis.
  • I Social Impact: Qualitative comparisons across GRPO-Guard, DiffusionNFT, and the proposed method are presented in Figure 7.The figure is identified as a qualitative comparison across the three methods.
  • I Social Impact: Qualitative comparisons across GRPO-Guard, DiffusionNFT, and the proposed method are also presented in Figures 8 and 9.Both figures are identified as qualitative comparisons across the same three methods.
Loading 2608.14430v1…