Diffusion Reinforcement Learning via Centered Reward Distillation

Yuanzhi Zhu; Xi Wang; Stéphane Lathuilière; Vicky Kalogeiton

cs.CVMar 14, 2026

Diffusion Reinforcement Learning via Centered Reward Distillation

Authors: Yuanzhi Zhu, Xi Wang, Stéphane Lathuilière, Vicky Kalogeiton

Organizations: LIX, École Polytechnique, CNRS, IPP · Inria at Univ. Grenoble Alpes, CNRS, LJK

Abstract

Diffusion and flow models achieve State-Of-The-Art (SOTA) generative performance, yet many practically important behaviors such as fine-grained prompt fidelity, compositional correctness, and text rendering are weakly specified by score or flow matching pretraining objectives. Reinforcement Learning (RL) fine-tuning with external, black-box rewards is a natural remedy, but diffusion RL is often brittle. Trajectory-based methods incur high memory cost and high-variance gradient estimates; forward-process approaches converge faster but can suffer from distribution drift, and hence reward hacking. In this work, we present Centered Reward Distillation (CRD), a diffusion RL framework derived from KL-regularized reward maximization built on forward-process-based fine-tuning. The key insight is that the intractable normalizing constant cancels under within-prompt centering, yielding a well-posed reward-matching objective. To enable reliable text-to-image fine-tuning, we introduce techniques that explicitly control distribution drift: (i) decoupling the sampler from the moving reference to prevent ratio-signal collapse, (ii) KL anchoring to a CFG-guided pretrained model to control long-run drift and align with the inference-time semantics of the pre-trained model, and (iii) reward-adaptive KL strength to accelerate early learning under large KL regularization while reducing late-stage exploitation of reward-model loopholes. Experiments on text-to-image post-training with GenEval and OCR rewards show that CRD achieves competitive SOTA reward optimization results with fast convergence and reduced reward hacking, as validated on unseen preference metrics.

Figures & tables

Figure 1: Qualitative results produced by our RL fine-tuned SD3.5M [ 24 ] model with GenEval [ 30 ] reward (top) and OCR [ 11 ] reward (bottom).

Figure 2: For each prompt, $K$ samples $\{x_{i}\}_{i=1}^{K}$ are generated from the sampling model $p_{\mathrm{samp}}$ ( $p_{s}$ ). A reward model produces external rewards $r(c,x_{i})$ , and implicit model rewards $\widehat{R}_{\theta,i}$ are estimated via diffusion ELBO differences between the current model $p_{\theta}$ and a moving reference $p_{\mathrm{old}}$ ( $p_{o}$ ). Within-prompt centering yields $\Delta_{r,w}^{i}$ and $\Delta_{\widehat{R},w}^{i}$ , cancelling the prompt-dependent normalizer and enabling a well-posed matching objective. A initial KL penalty with respect to the fixed CFG-guided pretrained model $p_{\phi}^{\mathrm{\scalebox{0.8}{CFG}}}$ is imposed to prevent reward hacking.

Figure 3

Compositional Image Generation
Model	Task Metric		Image Quality	Preference Score
Model	GenEval $\uparrow$	OCR $\uparrow$	Aesthetics $\uparrow$	ImgRwd $\uparrow$	PickScore $\uparrow$	CLIPScore $\uparrow$	HPSv2.1 $\uparrow$
SD3.5-M	0.63	0.59	5.39	0.87	22.34	27.99	0.279
Flow-GRPO [ 53 ] (w/o KL)	0.95	—	4.93	0.44	21.16	—	—
Flow-GRPO [ 53 ] (w/ KL)	0.95	—	5.25	1.03	22.37	29.25	0.274
AWM [ 100 ] (w/o CFG)	0.91	—	5.10	0.39	21.76	26.99	0.235

Table 1: Performance on Compositional Image Generation, Visual Text Rendering, and Human Preference benchmarks. Task metrics are evaluated on the corresponding test prompts; image quality and preference scores are evaluated on DrawBench [ 79 ] prompts. Baseline results are taken from Flow-GRPO [ 53 ] or reimplemented when unavailable in original papers. ImgRwd : ImageReward .

Model	Overall	Single Obj.	Two Obj.	Counting	Colors	Position	Attr. Binding
Autoregressive Models
Show-o [ 98 ]	0.53	0.95	0.52	0.49	0.82	0.11	0.28
Janus-Pro-7B [ 14 ]	0.80	0.99	0.89	0.59	0.90	0.79	0.66
GPT-4o [ 43 ]	0.84	0.99	0.92	0.85	0.92	0.75	0.61
Diffusion and Flow Matching Models
SD-XL [ 73 ]	0.55	0.98	0.74	0.39	0.85	0.15	0.23

Table 2: GenEval Result. Best scores are in blue , second-best in green . Results for baselines are taken from Flow-GRPO [ 53 ] or the respective original papers, and reimplemented where unavailable. Obj.: Object; Attr.: Attribution.

Panel of Figure 4 — Figure 4: Ablations on slow old model decay rate $\eta_{\mathrm{old}}$ , and initial KL strength $\beta_{\mathrm{init}}$ .

Figure 5: Visual comparison corresponding to the ablations in Fig. 4 .

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

	Reward Model	OCR	GenEval	Multi Rewards
LoRA	$\alpha$	64	64	64
LoRA	$r$	32	32	32
Sampling	Sampling CFG	1.0	1.0	1.0
	Training Rollout Steps	10	10	15
	Evaluation Sampling Steps	40	40	40
	ODE Solver	‘dpmv2’	‘dpmv2’	‘dpmv2’

Appendix

Table 3: Hyperparameters and configurations used for main experiment runs. Model architecture is the same as the pretrained model. Largest Experiments in this work are conducted with 4 NVIDIA H100 GPUs.

i

\eta_{\mathrm{old}}

and

\eta_{\mathrm{init}}

denotes for Training (Optimizer) Steps .

Panel of Figure 6 — Figure 6: BoN results. The orange curve is the BoN performance, and the blue curve is the average performance of the current $N$ samples.

Visual Text Rendering
Model	Task Metric			Preference Score
Model	PickScore $\uparrow$	CLIPScore $\uparrow$	HPSv2.1 $\uparrow$	Aesthetics $\uparrow$	ImgRwd $\uparrow$
SD3.5-M	22.34	27.99	0.279	5.39	0.87
Flow-GRPO [ 53 ] (w/o KL)	23.41	—	—	6.15	1.24
Flow-GRPO [ 53 ] (w/ KL)	23.31	27.81	0.315	5.92	1.28
DiffusionNFT [ 109 ]	23.61	28.80	0.344	6.04	1.46

Appendix

Table 4: Performance on Human Preference benchmarks with Multi-reward Training. Task metrics, image quality and preference scores are all evaluated on DrawBench [ 79 ] prompts. Baseline results are taken from Flow-GRPO [ 53 ] and [ 15 ] , or reimplemented when unavailable in the original papers. Flow-GRPo is trained on PickScore only. ImgRwd : ImageReward .

Panel of Figure 7 — Figure 7: Effectiveness of adaptive initial KL on DiffusionNFT.

Panel of Figure 10 — Figure 10: Ablations on adaptive weighting in ELBO estimator

Figure 11: We visualize samples generated with our method at successive checkpoints (from left to right) throughout training, starting from step 0 and sampled at every 200-step interval.

Figure 12: Visualization of generation with different inference CFG values. The left 5 columns are generated with prompts from GenEval and the right 5 columns are generated with prompts from OCR testsets, respectively.

Figure 13: Qualitative visual generations. Each row shows generations from a different random seed for the same prompt. Prompts follow the template “a photo of a {} zebra”, where the placeholder is replaced by the colors red , blue , green , yellow , and purple from top to bottom.

Figure 14: Qualitative visual generations from our model trained with multiple rewards. The text prompts are selected from the DrawBench .

Explore similar work

RL-Native Distillation: Exploiting Scored Trajectories for Few-Step Image Generation
Aug 10, 2026Yuhan Li, Fangao Zeng, Sicong Kang +7Diffusion Model Distillation Few-Step Diffusion Sampling
Shanghai Jiao Tong University · Taobao & Tmall group of Alibaba
Diff-Instruct with Diffused Reward: Towards Principled One-step Generator RL
May 18, 2026Junyi Wu, Weijian Luo, Haoyang Zheng +2Diffusion Model Alignment One-Step Generative Modeling
Purdue University · hi-lab, Xiaohongshu Inc.
Finite Difference Flow Optimization for RL Post-Training of Text-to-Image Models
Mar 13, 2026David McAllister, Miika Aittala, Tero Karras +4RL for Image Generation RL for Diffusion Models
UC Berkeley · NVIDIA

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

RL-Native Distillation: Exploiting Scored Trajectories for Few-Step Image Generation

Yuhan Li, Fangao Zeng, Sicong Kang, Mengfei Xu+6

Shanghai Jiao Tong University · Taobao & Tmall group of Alibaba

Diffusion Model Distillation Few-Step Diffusion Sampling RL for Diffusion Models

Open paper AskPDF arXiv

Efficient text-to-image generation requires both reinforcement-learning (RL)-based reward alignment and few-step distillation, yet these procedures are typically performed sequentially, increasing training cost and risking the loss of reward gains during compression. We take an RL-native perspective: diffusion RL already generates reward-scored finite-step trajectories, whose intermediate states provide distillation supervision. Based on this insight, we propose REST (Reward-Enhanced Scored-Trajectory Distillation), a single-stage co-training framework in which a decoupled student learns from the evolving RL teacher's trajectories without changing teacher optimization. Advantage-Modulated Distillation (AMD) transforms rollout advantages into signed weights, strengthening imitation of preferred trajectories and aligning distillation priorities with task value. The resulting framework is general and lightweight, requires no extra image rollouts, no separate distillation dataset, and no adversarial training. Experiments on compositional generation, visual text rendering, and human-preference alignment demonstrate competitive few-step, CFG-free generation with RAM or DiffusionNFT teachers. With only four sampling steps, REST-RAM achieves a DrawBench PickScore of 23.97, outperforming both the 40-step RAM teacher (23.95) and RTDMD (23.71).

Abstract

Explore similar work

RL-Native Distillation: Exploiting Scored Trajectories for Few-Step Image Generation

Diff-Instruct with Diffused Reward: Towards Principled One-step Generator RL

Finite Difference Flow Optimization for RL Post-Training of Text-to-Image Models