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

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

cs.LGAug 14, 2026

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

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

Organizations: State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · ByteDance Seed

Abstract

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.

Figures & tables

Method	$\mathbf{v}_{\text{base}}$ ( $\eta_{t}$ )	$\widehat{\nabla V}_{t}$	$s(t)$	$\hat{w}_{1}(A,t)$	$\hat{w}_{2}(t)$	Sampler
Flow-GRPO [ 24 , 41 ]	$\mathbf{v}_{\text{old}}$ ( $\eta_{t}$ )	$s(t)\cdot A{\bm{\upepsilon}}$	$\sqrt{\!\tfrac{1-t}{2\eta_{t}t\Delta t}}$	$A\cdot\frac{(1+\eta_{t})^{2}}{4\eta_{t}}\frac{1-t}{t}$	$1+\eta_{t}$	Flow-SDE
GRPO-Guard [ 37 ]	$\mathbf{v}_{\text{old}}$ ( $\eta_{t}$ )	$s(t)\cdot A{\bm{\upepsilon}}$	$\sqrt{\!\tfrac{1-t}{2\eta_{t}t\Delta t}}$	$0$	$\frac{(1+\eta_{t})}{s(t)\Delta t}$	Flow-SDE
TempFlow-GRPO [ 14 ]	$\mathbf{v}_{\text{old}}$ ( $\eta_{t}$ )	$s(t)\cdot A{\bm{\upepsilon}}$	$\sqrt{\!\tfrac{1-t}{2\eta_{t}t\Delta t}}$	$A\cdot\frac{(1+\eta_{t})^{2}}{4}\sqrt{\frac{2(1-t)\Delta t}{\eta_{t}t}}$	$(1+\eta_{t})/s(t)$	Branched-SDE
AWM [ 40 ]	$\mathbb{E}[\mathbf{v}\mid\mathbf{x}_{t}]$ ( $1$ )	$-s(t)\cdot A(\mathbf{v}\!-\!\mathbf{v}_{\text{old}})$	$\tfrac{1-t}{t}$	$A/\Delta t$	$\tfrac{2}{s(t)\Delta t}$	Flow-SDE / ODE
DiffusionNFT [ 45 ]	$\mathbb{E}[\mathbf{v}\mid\mathbf{x}_{t}]$ ( $1$ )	$-s(t)\cdot A(\mathbf{v}\!-\!\mathbf{v}_{\text{old}})$	$\tfrac{1-t}{t}$	$1/\Delta t$	$\tfrac{2}{\beta s(t)\Delta t}$	DPM-ODE [ 26 ]
Ours (Sec. 3.3 )	$\mathbf{v}_{\text{old}}$ ( $\eta_{t}$ )	$\widehat{\nabla V}_{t}^{\,\mathrm{KDE}}$ (Eq. ( 19 ))	$\frac{1-t}{t}$	$(1-t)^{\alpha_{1}}/\Delta t$	$\frac{t^{\alpha_{2}}}{s(t)\Delta t}$	Flow-SDE

Table 1: Diffusion RL methods as rows in the five-knob design space of Eq. ( 18 ). We use the FM parameterisation throughout;

\mathbf{v}

and

{\bm{\upepsilon}}

denote the FM conditional velocity and the per-step Brownian increment. See Appx. D for detailed derivations of each row.

Panel of Figure 1 — Figure 1 : (a) Comparison between different objectives on Pickscore reward. Original objective Eq. ( 12 ) with $\eta_{t}=1$ (Orange line); one-sample estimator with weighting in Eq. ( 17 ) (Blue line); reweighted version of Eq. ( 17 ) (Green line). (b,c) Empirical variance and squared bias of the value-gradient estimators. We set $\eta_{t}=1$ and evaluate SD3.5-Medium with PickScore reward on four prompts, using 128 trajectories per prompt and 40 denoising steps. At every intermediate state, the reference value gradient is approximated from all 128 trajectories; each estimator is computed with $G=24$ then compared with this common reference. The KDE training configuration uses $h=32.768$ .

Panel of Figure 2 — Figure 2 : On-policy weight (row 1) and off-policy (row 2) weight ablations for the stochastic, deterministic, and our KDE estimator. The results for the same estimator are in the same column.

Panel of Figure 3 — Figure 3 : (a,b,c). Training time comparison between our method and several baselines. (d). Comparison between Flow-GRPO, DiffusionNFT, and our method on Qwen-Image, HPSv3 reward. * means using the official implementation with different training configs compared to other baselines.

Figure 4 : Ablation studies on the components of our final recipe.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5 : Sampling noise schedule ablations for the stochastic one-sample estimator, deterministic one-sample estimator, and our KDE estimator. For each estimator, we run experiments for sampler noise schedule $\eta_{t}\in\{0.005,0.225\}$ on the Pickscore reward.

Panel of Figure 6 — Figure 6 : Visualization of different weights

Hyperparameter	Value
Group size $G$	$24$ for SD 3.5, $16$ for Qwen-Image
Train denoising steps $T$	$10$
Inference denoising steps	$40$ for SD 3.5, $50$ for Qwen-Image
Rollout SDE noise schedule	0.225 for stochastic estimator, 0.005 for deterministic and KDE estimator
Optimizer	AdamW, $\mathrm{lr}=3\!\times\!10^{-4}$ , $\beta_{1}\!=\!0.9,\beta_{2}\!=\!0.999$
KL coefficient $\beta$	$10^{-4}$

Appendix

Table 2: Hyperparameters used throughout Sec. 4 .

Training time (GPU hours)	$18.58$	$36.44$	$54.28$	$72.12$
DiffusionNFT	$22.85\pm 0.02886$	$23.06\pm 0.01618$	$23.16\pm 0.02153$	$23.22\pm 0.01274$
Ours	$22.95\pm 0.05240$	$23.28\pm 0.02239$	$23.48\pm 0.005547$	$23.56\pm 0.01461$

Appendix

Table 3: PickScore across three random seeds (mean

\pm

standard deviation). Our method achieves a higher mean at every checkpoint with small run-to-run variation.

Method	Training time per epoch (seconds)
DiffusionNFT	$115$
KDE estimator (ours)	$97$

Appendix

Table 4 : Measured wall-clock training time per epoch under matched SD3.5-Medium settings.

Figure 7 : Qualitative comparison across GRPO-Guard, DiffusionNFT and our method.

Figure 8 : Qualitative comparison across GRPO-Guard, DiffusionNFT and our method.

Figure 9 : Qualitative comparison across GRPO-Guard, DiffusionNFT and our method.

Explore similar work

Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design
Feb 4, 2026Jaemoo Choi, Yuchen Zhu, Wei Guo +6T2I Generation Diffusion Model Fine-Tuning
Georgia Institute of Technology · National University of Singapore · Nanjing university.
Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models
Sep 29, 2025Shuchen Xue, Chongjian Ge, Shilong Zhang +2Score Matching Diffusion Model Alignment
UCAS · Adobe Research · HKU +1
Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF
Jul 8, 2026Eric Zhu, Abhinav Shrivastava, Soumik MukhopadhyayDiffusion Model Alignment RL for Diffusion Models
Carnegie Mellon University · University of Maryland, College Park

Feb 4, 2026, cs.LGJ/K move · Enter open · S save

Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design

Jaemoo Choi, Yuchen Zhu, Wei Guo, Petr Molodyk+5

Georgia Institute of Technology · National University of Singapore · Nanjing university.

T2I Generation Diffusion Model Fine-Tuning Policy Gradient Methods RL for Diffusion Models Diffusion Models

Open paper AskPDF arXiv

Reinforcement learning has been widely applied to diffusion and flow models for visual tasks such as text-to-image generation. However, these tasks remain challenging because diffusion models have intractable likelihoods, which creates a barrier for directly applying popular policy-gradient type methods. Existing approaches primarily focus on crafting new objectives built on already heavily engineered LLM objectives, using ad hoc estimators for likelihood, without a thorough investigation into how such estimation affects overall algorithmic performance. In this work, we provide a systematic analysis of the RL design space by disentangling three factors: i) policy-gradient objectives, ii) likelihood estimators, and iii) rollout sampling schemes. We show that adopting an evidence lower bound (ELBO) based model likelihood estimator, computed only from the final generated sample, is the dominant factor enabling effective, efficient, and stable RL optimization, outweighing the impact of the specific policy-gradient loss functional. We validate our findings across multiple reward benchmarks using SD 3.5 Medium, and observe consistent trends across all tasks. Our method improves the GenEval score from 0.24 to 0.95 in 90 GPU hours, which is 4.6 times more efficient than FlowGRPO and

2\times

more efficient than the SOTA method without reward hacking.

Abstract

Explore similar work

Rethinking the Design Space of Reinforcement Learning for Diffusion Models: On the Importance of Likelihood Estimation Beyond Loss Design

Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models

Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF