Distribution-Based Regularization for Reward Fine-Tuning Diffusion and Flow Models

Romeo Passaro; Jeongjae Lee; Jong Chul Ye; Michael M. Bronstein; Yoshua Bengio; Alexander Tong; Danyal Rehman

cs.LGOct 17, 2026

Distribution-Based Regularization for Reward Fine-Tuning Diffusion and Flow Models

Authors: Romeo Passaro, Jeongjae Lee, Jong Chul Ye, Michael M. Bronstein, Yoshua Bengio, Alexander Tong, Danyal Rehman

Organizations: University of Oxford · AITHYRA · KAIST · Mila – Qu´ebec AI Institute ·

Abstract

Reward fine-tuning is an effective way to adapt pretrained diffusion and flow-matching models to downstream objectives, but aggressive optimization often leads to mode collapse and distributional drift. Existing methods attempt to control this drift with reverse-KL regularization, implemented as an $\ell_2$ penalty on the velocity field; however, this functions by indirectly constraining sampling dynamics rather than directly acting on the generated distribution. To address this, we propose DFR (Distributional Functional Rewards), which augments the task reward with a distribution-dependent, per-sample reward that acts explicitly on the generated distribution. Formally, we generalize the KL-regularized objective from expected rewards, which are linear in the model's distribution, to convex functionals of it. The first variation of this objective defines an effective, distribution-dependent reward that directly penalizes deviation from the pretrained distribution. This formulation yields additive per-sample rewards for the Fréchet distance (FD), maximum mean discrepancy (MMD), and DreamSim diversity, recovering several recently proposed rewards as special cases. Across various text-to-image fine-tuning tasks, DFR matches the peak reward of KL-tuned baselines while increasing diversity by up to 29% at equal or lower drift.

Figures & tables

Figure 1: Top : RAM achieves high rewards at the cost of sample diversity and distributional fidelity, whereas DFR matches—and often surpasses—its rewards while preserving diversity and FID relative to the pretrained distribution. Bottom : DFR more closely mirrors the pretrained SD3.5 distribution, generating a wider diversity of images relative to RAM across multiple seeds for the same prompt: “A zebra to the right of a fire hydrant.”

Trajectory-based
			GenEval $0.95$			GenEval $0.96$
Method	Peak	Speedup $(\uparrow)$	FID $(\downarrow)$	FD-DINOv2 $(\downarrow)$	DreamSim $(\uparrow)$	FID $(\downarrow)$	FD-DINOv2 $(\downarrow)$	DreamSim $(\uparrow)$
SD3.5-M	0.63	–	–	–	0.302/0.274
Flow-GRPO		1.0	114.32/125.61	1019/997	0.155/0.189
RSM §		6.25	132.82/154.29	1271/1310	0.142/0.183
ELBO-based

Table 1: Fine-tuning on GenEval, compared at matched GenEval (Section 4.1 ). Peak is the highest threshold it reaches or, if none, its mean GenEval at the evaluated checkpoints; “–” marks a threshold not reached on every seed. Speedup is the time to reach

0.95

relative to Flow-GRPO. Entries of the form

a/b

report results on the GenEval test prompts (

a

) and on DrawBench (

b

); RAM and DFR entries are means over three seeds with SD beneath, and existing methods are single checkpoints evaluated at

0.95

. Blue marks a DFR entry that wins against every RAM configuration reaching the same threshold; struck-through entries are computed with an encoder related to that instance’s reward and are not compared.

			OCR $0.96$
Method	Peak	Time $(\downarrow)$	FID $(\downarrow)$	FD-DINOv2 $(\downarrow)$	DreamSim $(\uparrow)$
SD3.5-M	0.59	–	–	–	0.258/0.274
RAM ( $\beta=2$ )	0.96	1.00	219.07/131.38 $\pm$ 24.04/7.71	1270/1040 $\pm$ 232/84	0.128/0.245 $\pm$ 0.010/0.009
RAM ( $\beta=3$ )	0.95	–	–	–	–
RAM ( $\beta=4$ )	0.96	1.85 $\pm$ 0.23	138.04/100.94 $\pm$ 6.73/4.36	667/751 $\pm$ 46/56	0.185/0.249 $\pm$ 0.004/0.005
RAM ( $\beta=5$ )	0.96	2.05 $\pm$ 0.09	136.78/97.16 $\pm$ 5.35/2.14	662/721 $\pm$ 39/21	0.192/0.252 $\pm$ 0.005/0.003

Table 2: Fine-tuning on OCR, compared at OCR reward

0.96

. Columns and markings are as in Table 1 , except that entries report results on the OCR test prompts (

a

) and on DrawBench (

b

), and Time is the time to reach

0.96

relative to RAM (

\beta=2

Figure 2: Drift versus diversity at matched reward. Tuning

\beta

moves RAM along a path (grey) that ends where a larger

\beta

no longer reaches the threshold. Every DFR instance lies above or to the left of the end of the path.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Method	$\beta$	$\gamma$	$f$	${{v}}_{\bullet}$	Time weight	Details
ORW-CFM-W2	their $\alpha$	$0$	$\exp(r/\alpha)$	${{v}}_{\theta}$	$1$	App. B.1
FPO ∗	$0$	$0$	$A$	${{v}}_{\theta}$	$\tilde{w}_{\text{FPO}}(t)$	App. B.2
AWM ∗,†	their $\beta$	$0$	$A$	${{v}}_{\theta}$	$w_{\text{AWM}}(t)$	App. B.3
RAM §	$1$	$0$	$A/\alpha$	${{v}}_{\text{old}}$	$2t$	App. B.4
DiffusionNFT	$0^{\ddagger}$	$1$	$\tilde{A}/\zeta$	${{v}}_{\text{old}}$	$1$	App. B.5
AdvantageFlow	their $\lambda$	$1.1$ or $1-\tilde{A}$	$\tilde{A}$	${{v}}_{\theta}$	$t^{2}$	App. B.5

Appendix

Table 3: Online RWR methods as instances of the unified gradient, Eq. ( 24 ). Methods marked with ∗ reduce exactly only at

\theta=\theta_{\text{old}}

(Appendix B.6 ); all other reductions are exact identities. Time weights are relative to

t\sim\mathcal{U}(0,1)

: sampling

t\sim p(t)

is equivalent to uniform sampling with weight

p(t)

Hyperparameter	Value
Base model	SD3.5-M
Fine-tuning	LoRA, rank 32, scaling 64
Optimizer	AdamW
Learning rate	$3\times 10^{-4}$
Adam $(\beta_{1},\beta_{2})$	$(0.9,\,0.99)$
Weight decay	$10^{-4}$

Appendix

Table 4: Training hyperparameters, shared by RAM and every DFR instance on both tasks. Only

\beta

and the DFR terms differ between configurations.

GenEval, threshold $0.95$
Method	Reward	Steps $(\downarrow)$	FID $(\downarrow)$	FD-DINOv2 $(\downarrow)$	DreamSim-DINO $(\uparrow)$	DreamSim-OpenCLIP $(\uparrow)$	ImgRwd $(\uparrow)$	PickScore $(\uparrow)$
SD3.5-M	0.63	–	–	–	0.405/0.372	0.302/0.274	0.840/0.859	0.872/0.863
Flow-GRPO	0.95	–	114.32/125.61	1019/997	0.224/0.274	0.155/0.189	1.222/1.013	0.883/0.862
RSM §	0.94	–	132.82/154.29	1271/1310	0.184/0.249	0.142/0.183	1.082/0.838	0.869/ 0.851
AWM ∗	0.95	–	158.88/183.71	1412/1578	0.124/0.173	0.095/0.132	1.028/ 0.573	0.869/ 0.840
DiffusionNFT	0.96	–	144.36/159.22	1352/1437	0.151/0.204	0.109/0.143	1.232/1.091	0.876/ 0.851

Appendix

Table 5: All metrics for models fine-tuned on GenEval, at thresholds

0.95

and

0.96

. Reward is the threshold for configurations that reach it and, otherwise, the mean GenEval at the evaluated checkpoints; Steps is the mean number of training steps to the threshold. Entries of the form

a/b

report results on the GenEval test prompts (

a

) and on DrawBench (

b

), as means over three seeds with SD beneath; existing methods are single checkpoints. Greyed-out rows ( † ) do not reach the threshold on every seed. Blue and struck-through entries are as in Table 1 ; in the ImgRwd (ImageReward) and PickScore columns, red marks a value significantly below that of SD3.5-M (Appendix D.2 ).

Method	Reward	Steps $(\downarrow)$	FID $(\downarrow)$	FD-DINOv2 $(\downarrow)$	DreamSim-DINO $(\uparrow)$	DreamSim-OpenCLIP $(\uparrow)$	ImgRwd $(\uparrow)$	PickScore $(\uparrow)$
OCR, threshold $0.96$
SD3.5-M	0.59	–	–	–	0.307/0.372	0.258/0.274	1.004/0.859	0.863/0.863
RAM ( $\beta=2$ )	0.96	67 $\pm$ 12	219.07/131.38 $\pm$ 24.04/7.71	1270/1040 $\pm$ 232/84	0.248/0.347 $\pm$ 0.016/0.010	0.128/0.245 $\pm$ 0.010/0.009	0.545 / 0.712 $\pm$ 0.251/0.133	0.849 / 0.855 $\pm$ 0.009/0.004
RAM ( $\beta=3$ ) †	0.95	100 $\pm$ 0	149.72/104.84 $\pm$ 5.66/3.13	728/798 $\pm$ 40/58	0.272/0.349 $\pm$ 0.019/0.011	0.167/0.246 $\pm$ 0.007/0.004	1.039/0.894 $\pm$ 0.054/0.048	0.869/0.862 $\pm$ 0.003/0.002
RAM ( $\beta=4$ )	0.96	123 $\pm$ 15	138.04/100.94 $\pm$ 6.73/4.36	667/751 $\pm$ 46/56	0.291/0.355 $\pm$ 0.010/0.008	0.185/0.249 $\pm$ 0.004/0.005	1.062/0.922 $\pm$ 0.064/0.045	0.867/0.862 $\pm$ 0.003/0.002
RAM ( $\beta=5$ )	0.96	137 $\pm$ 6	136.78/97.16 $\pm$ 5.35/2.14	662/721 $\pm$ 39/21	0.292/0.358 $\pm$ 0.003/0.004	0.192/0.252 $\pm$ 0.005/0.003	1.040/0.886 $\pm$ 0.043/0.030	0.867/ 0.861 $\pm$ 0.002/0.001

Appendix

Table 6: All metrics for models fine-tuned on OCR, at threshold

0.96

. Columns and markings are as in Table 5 , except that entries report results on the OCR test prompts (

a

) and on DrawBench (

b

Method	Speedup	Source
Flow-GRPO	1	reference
RSM	6.25	$5\times$ over TempFlow-GRPO ( Lee et al., 2026 ) , which reports $1.25\times$ over Flow-GRPO ( He et al., 2026 )
AWM	8	reported by Xue et al. (2026)
DiffusionNFT	24	reported by Zheng et al. (2026)
DiffusionNFT (multi)	2.28	reported by Choi et al. (2026)
RAM ( $\beta=1$ )	52.3	$50\times$ fewer steps, times the ratio of time per step ( Bergmeister et al., 2026 )

Appendix

Table 7: Speedups. Top: speedups over Flow-GRPO reported by the respective papers. Bottom: measured from our runs. Time is the number of steps to the threshold multiplied by the time per step, both relative to the RAM configuration stated in each block. GenEval speedups divide the reported speedup of RAM (

\beta=1

) by Time; OCR has no Flow-GRPO reference and reports Time only. Means with SD beneath are over seeds; “–” marks a threshold not reached on every seed.

Figure 3: ImageReward (left axis) and DreamSim diversity (right axis) of RAM (

\beta=1

) during training on GenEval, on the GenEval (solid) and DrawBench (dashed) prompts, evaluated every 30 steps; step 0 is SD3.5-M. Mean over three seeds with

\pm 1

SD bands.

Panel of Figure 4 — Figure 4: Target reward on held-out prompts during RAM training. Left: GenEval at $\beta=1$ and $\beta=4$ , with one line per seed and the target indicated at $0.95$ . At $\beta=1$ , no seed collapses within the training budget; at $\beta=4$ , every seed declines after peaking, and one collapses. Right: OCR for $\beta\in\{1,2,4,8,12\}$ , using 1 seed for each coefficient and a target of $0.96$ . The weakly regularized $\beta=1$ run collapses after reaching its peak, whereas the more strongly regularized runs remain stable over their observed training ranges, at the cost of slower reward improvement for $\beta=8$ and $\beta=12$ .

Figure 5: DFR qualitatively appears closer to the pretrained distribution; the colour of the glazing on the donut is preserved in DFR , while this does not hold for the RAM baseline.

Figure 6: DFR captures both the clock colors and the mix of Roman and Arabic numerals, preserving variations from the pretrained SD3.5 distribution, whereas the RAM baseline collapses these variations.

Figure 7: Sample images generated for diverse prompts across various seeds using the base SD3.5 model.

Figure 8: Sample images generated for diverse prompts across various seeds using the RAM baseline.

Figure 9: Sample images generated for diverse prompts across various seeds using DFR .

Generative Modeling

Generative modeling learns data distributions to create new samples, from images and text to molecules, simulations, and structured records.

See all 487 papers

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
Reward Score Matching: Unifying Reward-based Fine-tuning for Flow and Diffusion Models
Apr 19, 2026Jeongjae Lee, Jinho Chang, Jeongsol Kim +1Diffusion Model Alignment Diffusion Model Fine-Tuning
Graduate School of AI KAIST, Korea
$R_\text{dm}$ : Re-conceptualizing Distribution Matching as a Reward for Diffusion Distillation
Mar 30, 2026Linqian Fan, Peiqin Sun, Tiancheng Wen +2Diffusion Model Distillation Distribution Matching Distillation
KlingAI Research · Tsinghua University

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

Diffusion Reinforcement Learning via Centered Reward Distillation

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

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

RL for Diffusion Models Diffusion Model Alignment Diffusion Model Fine-Tuning T2I Diffusion Models Reward Hacking

Open paper AskPDF arXiv

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.

Abstract

Generative Modeling

Explore similar work

Diffusion Reinforcement Learning via Centered Reward Distillation

Reward Score Matching: Unifying Reward-based Fine-tuning for Flow and Diffusion Models

RdmR_\text{dm}Rdm​: Re-conceptualizing Distribution Matching as a Reward for Diffusion Distillation

$R_\text{dm}$ : Re-conceptualizing Distribution Matching as a Reward for Diffusion Distillation