DiffusionRVGE: Reinforcement Learning via Regression-Based Value Gradient Estimation

Jiayuan Wang; Tongyao Pang

cs.AIOct 16, 2026

DiffusionRVGE: Reinforcement Learning via Regression-Based Value Gradient Estimation

Authors: Jiayuan Wang, Tongyao Pang

Organizations: Qiuzhen College, Tsinghua University, Beijing 100084, China · Yau Mathematical Sciences Center, Shuangqing Complex, Tsinghua University, Beijing 100084, China

Abstract

Reinforcement learning improves diffusion models using rewards on generated images, but repeated denoising and gradient estimation make training expensive. We introduce \method, an efficient method that constructs useful update directions by regressing centered rewards on centered final denoised images. It requires neither reward derivatives nor adjoint backpropagation through the sampling trajectory. A unified first-variation formulation identifies the value-gradient target with or without a trajectory-KL penalty. A population analysis under Gaussian noising explains the regression geometry, while a subspace analysis motivates clean-image coordinates. The implementation adapts rollout bandwidth and query weights to reward continuity: local rollouts suit smooth rewards, whereas broader groups provide informative contrasts for discrete rewards. Trained jointly on five rewards, one SD3.5-Medium model matches or exceeds DiffusionNFT's joint model on seven of eight metrics, reaching GenEval $0.94$ , OCR $0.92$ , and PickScore $23.82$ at about one fifth of its estimated training compute. In single-reward training, \method reaches matched reward thresholds $2.1$ -- $6.7\times$ faster than reproduced DiffusionNFT on identical hardware. Ablations support the clean-image estimator, direct direction transfer, and reward-dependent rollout choices. Code is available at https://github.com/jy-wang-tsinghua/DIFFUSIONRVGE.

Figures & tables

Model	Iter.	GenEval	OCR	PickScore	ClipScore	HPSv2.1	Aesthetic	ImgRwd	UniRwd
SD3.5-M	–	0.63	0.59	22.34	0.285	0.279	5.36	0.85	3.03
+ FlowGRPO † (GenEval)	$>$ 5k	0.95	0.66	22.51	0.293	0.274	5.32	1.06	3.18
+ FlowGRPO † (OCR)	2k	0.66	0.92	22.41	0.290	0.280	5.32	0.95	3.15
+ FlowGRPO † (PickScore)	4k	0.54	0.68	23.50	0.280	0.316	5.90	1.29	3.37
+ DiffusionNFT (joint)	1.7k	0.94	0.91	23.80	0.293	0.331	6.01	1.49	3.49
+ Ours (joint)	1.0k	0.94	0.92	23.82	0.303	0.339	5.95	1.55	3.50

Table 1: Five-reward comparison on the metric panel of DiffusionNFT Table 1. Gray marks the rewards used to train each model. Baseline rows are published results and retain their published inference protocols; † denotes official checkpoints. DiffusionNFT lists

1.7

k iterations for its five-stage joint model, whose described stages total

1{,}600

(Appendix F.1 ); ours uses

1{,}000

updates (404 GPU-hours). Bold marks the best value in each column.

Model	Reward	Iter. (GPU-h)	GenEval	OCR	PickScore	ClipScore	HPSv2.1	Aesthetic	ImgRwd	UniRwd
SD3.5-M	–	–	0.63	0.59	22.34	0.285	0.279	5.36	0.85	3.03
FlowGRPO	GenEval	4k	0.97	0.30	21.78	0.277	0.248	5.15	0.74	2.87
FlowGRPO	OCR	1k	0.66	0.96	21.94	0.280	0.257	5.18	0.31	2.86
FlowGRPO	PickScore	4k	0.54	0.60	23.62	0.257	0.295	6.42	1.17	3.17
DiffusionNFT	GenEval	1k	0.98	0.36	21.92	0.271	0.251	5.33	0.68	2.91
DiffusionNFT	OCR	150	0.54	0.97	21.63	0.281	0.246	5.19	0.37	2.81

Table 2: Single-reward results on the panel of DiffusionNFT Table 2, with a separate OCR quality-floor run. Gray marks the primary training reward; blue marks the three floor components (PickScore, ClipScore, HPSv2.1). Published rows retain their original protocols. Metric precision matches Table 1 : two decimals for GenEval, OCR, PickScore, Aesthetic, ImageReward, and UnifiedReward; three for ClipScore and HPSv2.1.

Figure 1: Single-reward learning on

8\times

A800 hardware: each curve optimizes the named reward. All curves use the aligned DrawBench

200\times 1

, OCR

1{,}016

, and official GenEval

553\times 4

datasets. FDFO appears only in the PickScore panel because FDFO does not evaluate the non-continuous OCR and GenEval rewards.

Reward	Target	Ours	DiffusionNFT	Speed-up
PickScore	23.20	14.3	29.7	$2.1\times$
	23.50	38.2	98.5	$2.6\times$
OCR	0.80	2.2	14.9	$6.7\times$
	0.95	4.1	25.6	$6.2\times$
GenEval	0.80	11.5	30.3	$2.6\times$
	0.90	22.8	102.0	$4.5\times$

Table 3: Measured GPU-hours for single-reward training to shared reward thresholds in Figure 1 .

Figure 2: PickScore cosine with the real gradient through the numerical ODE. Left and middle: noisy-state and direct clean-image regression. Right: paired clean-image-minus-noisy-state cosine. Shading denotes 95% confidence intervals from resampling prompts.

Figure 3: PickScore training with direct clean-image directions or their norm-matched Jacobian pullback to

\mathbf{x}_{t}

, shown against compute (left) and updates (right). Both arms use the same trajectory bank and actor settings.

Figure 4: Rollout comparison across three rewards. Training curves for trajectory-bank and independent-seed group regression. Batch sizes differ in GenEval and OCR, so these panels compare practical recipes at equal compute.

Stopped weight / selection	10 GPU-h	20 GPU-h	28 GPU-h
$a_{i}=1$ (all 24)	0.72	0.84	0.86
$a_{i}=1-r_{i}$	0.71	0.88	0.89
$a_{i}=\mathbf{1}\{r_{i}<1\}$	0.74	0.88	0.88
Random 6	0.81	0.86	0.87
Highest 6	0.62	0.67	0.75
Lowest 6 (ours)	0.81	0.87	0.92

Table 4: GenEval actor-weight ablation. Scores are interpolated within the recorded validation curves at equal GPU-hours and rounded to two decimals; bold marks the best unrounded value in each column. Only the selected query trajectories receive actor gradients; all sibling clean images remain regression observations.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: GenEval query weights at a

28

-GPU-hour budget, shown against compute and updates.

Figure 6: The all-24 continuation. The lowest-six advantage persists beyond the initial budget but narrows near

72

GPU-hours.

GPU-hours	28	35	44	52	61
Lowest 6	0.92	0.92	0.93	0.94	0.95
All 24	0.86	0.88	0.85	0.88	0.88

Appendix

Table 5: Longer GenEval comparison at fixed compute, using interpolation within the saved curves.

Figure 7: PickScore training across seeds, shown against GPU-hours (left) and updates (right).

Setting	PickScore	OCR	GenEval	OCR floor
Estimator	bank	bank	group	group
Prompts / update	24	8	48	48
Trajectories / prompt	24	12	24	24
Trained trajectories / prompt	24	12	6 lowest	6 lowest
Rollout	SDE, $\eta=0.03$	SDE, $\eta=0.5$	ODE	ODE
Reward evaluations / trajectory	11	11	1	1

Appendix

Table 6: Configurations at the reported checkpoints. OCR floor uses Eq. ( 43 ) with

\lambda_{\mathrm{floor}}=10

and the margin schedule

m=0\to 0.03\to 0.06

at updates

40

and

100

Updates	Update type	Prompt mix	Learning rate
0–200	bank, PickScore	1 / 0 / 0	$3\times 10^{-4}$
200–500	group	0.20 / 0.45 / 0.35	$3\times 10^{-4}$
500–600	group	0.20 / 0.60 / 0.20	$1.35\times 10^{-4}$
600–700	bank	1 / 0 / 0	$7.25\times 10^{-5}$
700–850	group	0 / 0 / 1	$7.25\times 10^{-5}$
850–1000	group	0.20 / 0.60 / 0.20	$8.75\times 10^{-5}$

Appendix

Table 7: Stages of the five-reward joint model. The prompt mix gives the fractions of Pick-a-Pic, GenEval, and OCR prompts per update; the first stage is the first training seed of the single-reward PickScore run in Table 2 .

Figure 8: PickScore visual results comparing the DiffusionRVGE PickScore checkpoint with the public Flow-GRPO PickScore checkpoint on DrawBench prompts. Within each pair, models use the same prompt, seed, and 40-step ODE sampler.

Figure 9: OCR visual results comparing the DiffusionRVGE quality-floor checkpoint with the public Flow-GRPO Text checkpoint on OCR test prompts. Within each pair, models use the same prompt, seed, and 40-step ODE sampler; panel labels abbreviate the prompts.

Figure 10: GenEval visual results comparing the DiffusionRVGE single-reward GenEval checkpoint with the public Flow-GRPO GenEval checkpoint on GenEval prompts. Within each pair, models use the same prompt, seed, and 40-step ODE sampler.

Figure 11: DrawBench visual results comparing the five-reward joint DiffusionRVGE checkpoint with the public DiffusionNFT multi-reward checkpoint. Each pair uses the same full prompt and seed with 40-step ODE sampling; panel labels abbreviate the prompts.

Figure 12: OCR test visual results comparing the five-reward joint DiffusionRVGE checkpoint with the public DiffusionNFT multi-reward checkpoint. Each pair uses the same full prompt and seed with 40-step ODE sampling; panel labels abbreviate the prompts.

RL for Image Generation

Reinforcement learning for image generation trains visual models with reward feedback to improve prompt alignment, quality, and control.

See all 93 papers

Designing Reinforcement Learning for Diffusion Models: A Unified Path-Space View
Aug 14, 2026Yixian Xu, Yuanrui Zhang, Shengjie Luo +2Diffusion Model Alignment Policy Gradient Methods
State Key Laboratory of General Artificial Intelligence, School of Intelligence Science and Technology, Peking University · ByteDance Seed
Scaling Reinforcement Learning for Diffusion Models via Velocity Matching
Aug 24, 2026Jaemoo Choi, Wei Guo, Yuchen Zhu +4Diffusion Model Fine-Tuning RL for Diffusion Models
Georgia Institute of Technology · NVIDIA
Diffusion Reinforcement Learning via Centered Reward Distillation
Mar 14, 2026Yuanzhi Zhu, Xi Wang, Stéphane Lathuilière +1RL for Diffusion Models Diffusion Model Alignment
LIX, École Polytechnique, CNRS, IPP · Inria at Univ. Grenoble Alpes, CNRS, LJK

Aug 14, 2026J/K move · Enter open · S save

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

Yixian Xu, Yuanrui Zhang, Shengjie Luo, Liwei Wang+1

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

Diffusion Model Alignment Policy Gradient Methods RL for Diffusion Models Variance Reduction

Open paper AskPDF arXiv

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.

Abstract

RL for Image Generation

Explore similar work

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

Scaling Reinforcement Learning for Diffusion Models via Velocity Matching

Diffusion Reinforcement Learning via Centered Reward Distillation