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

Feb 4, 2026cs.LG

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

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.

Jaemoo Choi, Yuchen Zhu, Wei Guo +6

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

Diffusion-Based Reinforcement Learning Methods Reward Gradients Large Language Model Reinforcement Learning Diffusion Models Policy Gradient Flow-Grpo

Ask PDF

Sep 29, 2025cs.LG

Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models

Reinforcement Learning (RL) has emerged as a central paradigm for advancing Large Language Models (LLMs), where both pre-training and RL post-training stages are grounded in the same log-likelihood formulation. In contrast, recent RL approaches for diffusion models, most notably Denoising Diffusion Policy Optimization (DDPO), optimize an objective different from the pretraining objectives--score/flow matching loss. In this work, we establish a novel theoretical analysis: DDPO is an implicit form of score/flow matching with noisy targets, which increases variance and slows convergence. Building on this analysis, we introduce Advantage Weighted Matching (AWM), a policy-gradient method for diffusion. It uses the score/flow-matching loss and reweights each sample by its advantage. In effect, AWM raises the influence of high-reward samples and suppresses low-reward ones while keeping the modeling objective identical to pretraining. This simple yet effective design yields substantial benefits: on the GenEval, OCR, and PickScore benchmarks, AWM delivers up to a

34\times

speedup over Flow-GRPO (which builds on DDPO), when applied to Stable Diffusion 3.5 Medium and FLUX, without compromising generation quality. Code is available at https://github.com/scxue/advantage_weighted_matching

Shuchen Xue, Chongjian Ge, Shilong Zhang +2

UCAS · Adobe Research · HKU +1

Diffusion Models Diffusion Policies

Ask PDF

Jul 8, 2026cs.LG

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

Reinforcement learning from human feedback (RLHF) has emerged as a powerful paradigm for aligning generative models with human preferences. However, applying RLHF to diffusion models remains highly feedback inefficient, as existing approaches typically require large amounts of human or reward model evaluations. This limitation reduces the practicality of diffusion RLHF in realworld settings where feedback is the primary bottleneck. In this paper, we propose two complementary strategies that substantially improve the feedback efficiency of diffusion RLHF while preserving generalization to unseen prompts. Our key observation is that reward information in diffusion trajectories is unevenly distributed: not all denoising timesteps or trajectories contribute equally to learning from a reward signal. By emphasizing informative timesteps and trajectories during optimization, we obtain more effective gradient updates. First, we introduce a per-timestep weighting scheme that reweights denoising steps during policy optimization. We theoretically connect this weighting to the optimal convergence properties of proximal policy optimization (PPO) and approximate the resulting weighting trend empirically. Second, we introduce a replay mechanism that prioritizes informative trajectories, enabling the model to reuse past samples instead of repeatedly querying new rewards. Together, these strategies significantly improve the feedback efficiency of diffusion RLHF. Under identical hyperparameter settings, our approach achieves up to a 6

\times

improvement in sample efficiency compared to widely used diffusion RLHF baselines.

Eric Zhu, Abhinav Shrivastava, Soumik Mukhopadhyay

Carnegie Mellon University · University of Maryland, College Park

Diffusion-Based Reinforcement Learning Methods Reinforcement Learning From Human Feedback Reward Gradients Timestep Proximal Policy Optimization

Ask PDF