Fine-Tuning Diffusion Models via Intermediate Distribution Shaping
Authors: Gautham Govind Anil
Shaan Ul Haque
Nithish Kannen
Dheeraj Mysore Nagaraj
Sanjay Shakkottai
Karthikeyan Shanmugam
Venue: ICLR 2026 Poster (2026)
Keywords: diffusion
fine-tuning
reinforcement learning
Primary area: generative models
TL;DR: We propose intermediate distribution shaping for KL regularized reward maximization of diffusion models via P-GRAFT and refinement for pre-trained flow models via Inverse Noise Correction.

Abstract: Diffusion models are widely used for generative tasks across domains. Given a pre-trained diffusion model, it is often desirable to fine-tune it further either to correct for errors in learning or to align with downstream applications. Towards this, we examine the effect of shaping the distribution at intermediate noise levels induced by diffusion models. First, we show that existing variants of Rejection sAmpling based Fine-Tuning (RAFT), which we unify as GRAFT, can implicitly perform KL regularized reward maximization with reshaped rewards. Motivated by this observation, we introduce P-GRAFT to shape distributions at intermediate noise levels and demonstrate empirically that this can lead to more effective fine-tuning. We mathematically explain this via a bias-variance tradeoff. Next, we look at correcting learning errors in pre-trained flow models based on the developed mathematical framework. In particular, we propose inverse noise correction, a novel algorithm to improve the quality of pre-trained flow models without explicit rewards. We empirically evaluate our methods on text-to-image(T2I) generation, layout generation, molecule generation and unconditional image generation. Notably, our framework, applied to Stable Diffusion v2, improves over policy gradient methods on popular T2I benchmarks in terms of VQAScore and shows an 8.81% relative improvement over the base model. For unconditional image generation, inverse noise correction improves FID of generated images at lower FLOPs/image.

Published as a conference paper at ICLR 2026
FINE-TUNING DIFFUSION MODELS VIA INTERMEDIATE DISTRIBUTION SHAPING
Gautham Govind Anil'
Shaan UI Haque2+1 Nithish Kannen' Dheeraj Nagaraj
Sanjay Shakkottai" Karthikeyan Shanmugam'
Google DeepMind Georgia Institute of Technology, Atlanta 3 University of Texas at Austin
ABSTRACT
Diffusion models are widely used for generative tasks across domains. Given a pre-trained diffusion model, it is often desirable to fine-tune it further either correct for errors in learning or to align with downstream applications, Towards this, we examine the effect of shaping the distribution at intermediate noise levels induced by diffusion models. First, we show that existing variants of Rejection sAmpling based Fine-Tuning (RAFT), which we unify as GRAFT, can implicitly perform KL regularized reward maximization with reshaped rewards. Motivated by this observation, we introduce P-GRAFT to shape distributions at intermediate noise levels and demonstrate empirically that this can lead to more effective fine- tuning. We mathematically explain this via a variina tradeoff. Next, we look correcting learning errors in pre-trained flow models based on the developed mathematical framework. In particular, we propose inverse noise correction, a novel algorithm to improve the quality of pre-trained flow models without explicit rewards. We empirically evaluate our methods on text-to-image(221 generation, layout generation, molecule generation and unconditional image generation. Notably, our framework, applied to Stable Diffusion v2, improves over policy gradient methods on popular T2I benchmarks in terms of VQAScore and shows an 8.81% relative improvement over the base model. For unconditional image generation, inverse noise correction improves FII) of generated images at lower FI.OPs/image.
INTRODUCTION
Pre-trained generative models often require task-specific adaptations based on reward feedback a standard strategy to leverage RL algorithms such as Proximal Policy Optimization (PPO) (Schulman et al.) 2017) with KL regularized rewards. While such methods have found great success in the context of language modeling (Bai et al. 2022 Ouyang et al. 2022), their adoption to diffusion models is not straightforward. In particular, unlike autoregressive models, marginal likelihoods required for the implementation of KI. regularization in PPO are intractable for diffusion models.
Hence, in practice, KL regularization is ignored (Black et al. 2023) or relaxations such as trajectory KL regularization (Fan et al., 2023) are considered. However, ignoring the KL term results in unstable training in large-scale settings (Deng et al. 2024), whereas using the trajectory KL constraint gives subpar results (Black et al., 2023). Further, fine-tuning with trajectory KL also results in the initial value function bias problem (Domingo-EEnrich et al. 2024 Uehara et al. 2024).
Apart from policy gradient methods, recent research has also focused on fine-tuning methods based on rejection sampling such as RSO (Liu et al. (003a)), RAFT (Dong et al.) (023) and Reinforce-Rei (Xiong et al. 2025).. Further, fine-tuning based on Best-of-N (BoN) sampling and its relation to
policy gradient methods have also been explored, but in the context of autoregressive models (Amini
et al. 2024 Gui et al. 2024).
Given the intractability of marginal KL for diffusion models, we explore conceptual connections between rejection sampling based fine-tuning methods and KL regularized reward maximization.
'Equal contribution
I Work done while at Google DeepMind
1 "ooresponneecee cheera janearaanaaaaaaaaa ja jggoog ccm

---

Published as a conference paper at ICLR 2026
We start by re-stating known results about equivalence between rejection sampling strategies (such as classical rejection sampling from MCMC literature and Best-of-N) and KL regularized reward maximization. We do by unifying these strategies under common Generalized Rejection Sampling (GRS) framework, which also subsumes additional rejection sampling strategies. More importantly, note that fine-tuning using GRS, which we term Generalized Rejection sAmpling Fine-Tuning (GRAFT) enables marginal KL constraint for diffusion models, despite the marginal likelihoods being intractable. Building on this observation, we leverage the additional structure of diffusion (and flow) models to examine the effects of shaping the intermediate distributions induced by these models as opposed to the final data distribution. In particular, we make the following contributions:
(a) Leveraging properties of diffusion models, we propose Partial-GRAs (P-GRAFT) a framework which fine-tunes only till an intermediate denoising step by assigning the rewards of final generations to partial, noisy generations. We show that this leads to better fine-tuning empirically and provide a mathematical justification via bias-variance tradeoff. Empirically, we demonstrate significant quality gains across the tasks of text-to-image generation, layout generation and molecule generation.
(b) Motivated by P-GRAFT, we analyze the effect of learning the initial noise distribution of flow models. In particular, we go beyond reward-basee fine-tuning and introduce Inverse Noise Correction- an adapter-based parameter-efficiin explicit rewards. Intuitively, leverage the reversibility of flow models to learn appropriate initial noise distribution. We empirically demonstrate improved quality as well as FLOOPs for unconditional image generation.
method to correct errors in pre-trained flow models even without
(c) In particular, SDv2 fine-tuned using P-GRAFT demonstrates significant improvements in VQAS- core over policy-gradient methods as well SDXL-Base across datasets. The proposed Inverse Noise Correction strategy provides significant FID improvement at reduced FLOOP//maae.
more comprehensive list of related work be found in Appendix A
PRELIMINARIES
2.1 KL REGULARIZED RL FOR GENERATIVE MODELING
Following (Stiennon et al. 2020), we introduce the following reward maximization problem in our setting: Consider state space X, reward function r X + R and a reference probability measure p over X. Let P(X) ) be the set of probability measures over X' and a E (0, x)). Define
the KL regularized reward maximization objective ! P(X)) + R by (p)) = Ex Ex-p(r(X)) (Kl(p)))) where KL(-II-) is the KL divergence. RL algorithms such as PPO are then used to solve for
RL arg pepp) (p).
(I)
Using the method of Lagrangian Multipliers, we can show that (x) (x)() In generative modeling literature, p is often the law of generated samples from a pre-trained model- fine-tuning is done on the model as to sample from the tilted distribution RL.
2.2 OPTIMAL DISTRIBUTION VIA REJECTION SAMPLING:
Classical rejection sampling from the Monte Carlo literature (Thomopollss 2012 can be used to sample from RL. We note this folklore result in our setting:
Lemma 2.1. Let r(x)= S Finax for some Tmax- Given a sample Y ~ P. we accept it with probability
P(Accept||) Y- = exp r(V) Tmax) Then, conditioned on Accept, Y a sample from RL
provides way to obtain exact samples from P A well known challenge with this method is sample inefficiency- as often in practice, a is small leading to small acceptance probability. Thus, in practice, methods such as Best-of-N (BoN) which always accept fixed fraction of samples are used. We now introduce Generalized Rejection sAmpling Fine Tuning (GRAFT), a framework to unify such rejection sampling approaches. More specifically, Lemma 2.3 shows that this still leads to sampling from the solution of the KL regularized reward maximization objective, but with reshaped rewards. We then discuss its utility in the context of diffusion models.
Lemma 2.1
2

---

Published as a conference paper at ICLR 2026
2.3 GRAFT: GENERALIZED REJECTION SAMPLING FINE TUNING
Assume (xE()) are M i..d. samples with law over space A. Given reward function r X - I. let the reward corresponding to ( be R2 == r(x(1)), the empirical distribution of (X(()))E(A) be (x(-) and the empirical CDF of (Ra)E(MM be FR(-) We introduce Generalized Rejection Sampling (GRS) to accept subset of high reward samples, A (Y()) S (EM))) where yū) denotes the 1th accepted sample.
Definition 2.2. Generalized Rejection Sampling (GRS); Let the acceptance function R x [0, 1] x X' [0, 1] + [0, 1] be such that A is co-ordinate wise increasing in the first two co-ordinates. The acceptance probability of sample is p; A A(R,, FR((R), X Px). Draw C. Ber(p,) A {,...., M), not necessarily independent of each other. Then, X A iff C, =1
Definition 2.2 subsumes popular rejection sampling approaches such as RAFT and BoN. We now show that GRS implicitly performs KL regularized reward maximization with the reshaped reward f()):
Lemma 2.3.
The law of accepted samples under GRS 2.2) given by p (L) x+((1) E is the solution to the following K1. regularized reward maximization problem:
arg max (Bx))(()) aKL (fll)));
((=) == log (B Fa(r(z)), PxIXX() =
Here, the expectation is with respect to the randomness in the empirical distributions Fn and Px-
f()) is monotonically increasing with respect to the actual reward since A is an increasing function of the reward and its empirical CDF. We now instantiate GRS with commonly used variants of A:
Top K Sampling: Let the reward distribution be continuous with CDF F(-). We accept the top K samples out of the M samples based on their reward values.
=(1
So if (R(r) S K if Fn(r) >
Corresponding Acceptance Function: A(r, FR, Px) Px)=
M
Lemma 2.3 shows that this acceptance function results the reshaped reward satisfying: A(m
log| [> k== F((=t)))
F(r())))*
Preference Rewards: Setting M 2 and 1 in the above formulation gives preference rewards,
i... (1) is accepted and (2) rejected if r(x(1)
((x(2) (and vice versa). Since F' is an increasing function, the reward is monotonically reshaped from r(z) to log F(r(z)).
Varying K from to M, varies the strength of the tilt in Top K sampling. In particular, K M corresponds to f(z) = (no tilt) and K corresponds to f(z) M log F(r(z)).
Binary Rewards with De-Duplication: Suppose r(X) E fo, (for eg.. corresponds unstable/ stable molecules in molecule generation). De-duplication of the generated samples might be necessary to maintain diversity. Given any structure function ff (for eg., extracts the molecule structure from a configuration), let (X, Px) = Ki f(x(0)= = i.c, the number of copies of X in the data.
ir= 0
Proposed Acceptance Function: A(r, A I, Px)
if rf==1
Nytx,Px)
Draw C. Ber(p:) without-replacement among the duplicate/similar samples (i.e, they are marginally Bernoulli but are not independent). Thus, exactly one out of the duplicate molecules are selected almost surely. Applying Lemma 2.3 we conclude that:
f(x)
ifr( =0
log
(y Px) [x0= =
ifr(x)= =1
We see that the shaped reward increases with diversity and with the value of the original reward. We this in the molecule generation experiments avoid mode collapse (Section
3

---

Published as a conference paper at ICLR 2026
2.4 IMPLICATIONS FOR DIFFUSION MODELS:
While specialized versions of Lemma 2.3 are known in the context of AR models (Amini et al. 2024). the result is particularly useful in the context of diffusion models. Note that given sample I along with a prompt the marginal likelihood p(æly) can be easily computed for AR models. For diffusion models, we only have access to conditional likelihoods along the denoising trajectory of the diffusion process whereas (L(plpp) is intractable. That is, if the denoising process is run from iN to to. we have access to p(Fr, (Ttz+1) ). A commonly used relaxation is the trajectory KL., (X.rr) which can be shown as an upper bound on the marginal KL. As discussed in DDomingo-nnrie et al., 2024), this constraint can lead to the initial value function bias problem since the KL regularization is with respect to the learned reverse process (as shown in (Domingo-Enrich
(L(p(Xo:r)) (
ct al. 2024)). It becomes necessary to learn an appropriate tilt even at time T. In this context, Lemma 2.3 offers simple yet effective alternative to implicitly achieve marginal KL regularization.
Based on GRS, we propose GRAFT: Generalized Rejection sAmpling Fine Tuning (Algorithm 7 given reference model p, we generate samples and perform the GRS strategy proposed in 2.2 A
dataset is generated from the accepted samples and standard training is done on the generated dataset.
3 PARTIAL-GRAFT FOR DIFFUSION MODELS
Having established that GRAFT implicitly samples from the KL regularized reward maximization objective, we now examine methods to further improve the framework. Continuous diffusion models typically start with Gaussian noise X2 at time T and denoise it to the output X, via a discretized continuous time SDE. With N denoising steps, the model constructs a denoising trajectory (Tx=7 and lo 0), denoted by XT:0. We now consider the effect of shaping the distribution of an intermediate state X,. For the rest of the discussion, we reserve and N to refer to discrete timesteps, and I. and T for continuous time. For any E [0, T denote the
Xtx - +....,, ..... Xto
marginal density of X, by (z))
We first extend GRS to Partial Generalized Rejection Sampling (P-GRS). Let x(",.., x(aa be partially denoised (denoised till time t) samples. Let their corresponding completely denoised samples be x Xx..,.0 (M) Rewards are computed using the completely denoised samples (i.c. R2 = r(X)) for the 8th sample). We denote the empirical distribution of (xx") (!... by Px, (+) and the empirical CDF of (R,,..., Ras} by fn()).
,....
Definition 3.1. Partial Generalized Rejection Sampling (P-GRS): Consider acceptance function A : R x [0, 1] x X x [0, - [0, 10 such that is co-ordinate wise increasing in the first two co-ordinates. The acceptance probability of sample i is p; == A(R,, Fn((R), x;(") Px,,). Draw
C2 ~ (pp( V i E [M], not necessarily independent of each other. Then, x." E Aff = 1.
Lemma 3.2. The law of the accepted samples under P-GRS ( Def. 3.1 given by (x( x|xx'" A) is the solution to the following Pruximal Policy Optimizution problem:
arg max [Ex - aKL (fllp))];
plx) == log (B[A(+(x)"), (. (D) Px) x))
The key difference is that the reshaped reward now depends on the expected value of the acceptance function given a partially denoised state X,. This tilts pt instead of po. It is straightforward to modify the reshaped rewards corresponding to GRS to that of P-GRS. We illustrate this by instantiating Lemma 3.2 for preference rewards, as done with GRS (Lemma 2.3 above.
Preference rewards: With P-GRS, pc( x: x(xx E A) Pc(x) exp f(x)) with f(2)
a
log ((F(( z].-
Based on Lemma 3.2 we introduce P-GRAFT: Partial GRAFT (Algorithms [ and 2). Here, fine-tuning is done on a (sampled) dataset of partially denoised vectors instead of fully denoised vectors. The fine-tuned model is only trained from times T to t, and used for denoising from noise only till time t. We switch to the reference model for further denoising. The resulting final distribution is given in Appendix [.3.I] We will now discuss the mathematical aspects of GRGAATT and provide a justification for its improved performance.
4

---

Published as a conference paper at ICLR 2026
3.1 A VAS-VARIANCE TRADEOFF JUSTIFICATION FOR P-GRAFT
We analyze P-GRAFT from a bias-variance tradeoff viewpoint. Let us associate reward r r(Xo) with X₁. As argued in Lemma 3.3 variance of r(Xo) conditioned on X, increases with Consequently, P-GRAFT obtains noisy rewards, seemingly making it less effective than GRAFT. However, we subsequently show that the learning problem itself becomes easier when is large since the score function becomes simpler (i.e, the bias reduces). Therefore, we can balance the trade-off between the two by choosing an "appropriate" intermediate time for the distributional tilt.
Lemma 3.3. The expected conditional variance E Var)((Xo)) is an increasing function of
Example: Consider molecule generation, where molecules generated by pre-trained diffusion model. The generated molecule can be stable (r(Xo) = 1) or unstable (r(Xo) = 0). Intuitively, Xę. for T, carries more information about r (XXo) than Xr. We reinforee this claim empirically by giving the following illustrative statistical test. Consider the two hypotheses:
Ho : r(Xo) is independent of X2;
H1 : r(xo) and X, are dependent.
Algorithm P-GRAFT:
Training
Input:
Trainable model po. Reference model p. Reward function r, Acceptance function A, Number of rounds
Algorithm 2 P-GRAFT: Inference
Input: Fine tuned model p. Reference model p. Intermediate timestep N,,
Ns, Intermediate timestep N₂
Per-step denoiser DEN
1: Initialize empty set D
2: for j to Nş do
1: Sample X ~ I),
3: Generate M trajectories: X ~ pr:o: iE [M]
2: for n = N to N, do
4: Obtain rewards: r(x;,)); i [M]
3: Xt.. <- DEN(P, X 11, bn+1
5: Perform P-GRS using acceptance function A on
4: end for
Xix; iE [M] to get accepted samples A
5: for .2 to 0 do
6: X,,, + DEN(P, X Xt (n11)
6: Perform D - D-DL A
7: end for
7: end for
8: reeturn X
8: Train po on D for L (IND..... |N}
9:
return p6
Given X,, we obtain 100 roll outs
xf",X, [X, for S S 100 and its empirical average f(X,)
Binomial
0.15
3N/4
Table I: Conditional variance.
En r( r(x))) 1/100. If r(x)) is independent of X, (under Ho). the law of f(X,) is the binomial distribution Bin(100, 100, 8) with 0 P(r(X,)) = 1) being the marginal probability of observing a stable molecule. We perform 1000 rep- etitions for the experiment above for various values of and plot the empirical distributions in Figure close to the Binomial distribution and for = (x/4 (when X, is close to the target) it is far. That is, Xewya already carries a lot of information about r(Xo). This is further supported by the expected conditional variances reported in Table
N/2
N/4
R [Var(r(Xo)) )]
0.05
0.1341
3N/4
0.1327
1.0
N/2
0.1312
(X.,)
N/4
0.0848
Figure 1: Law of (()))
! For f3N/4 (when X, close to (O, the distribution is
Bias reduces with increasing t: We follow the Stochastic Differential Equation (SDE) framework from Song et al. (2020a) for our analysis. Let the target distribution go be the law of accepted samples under P-GRS. Diffusion models consider the forward process to be the Ornstein-Uhlenbed Uhlenbeck Process given by d = -x dt + v2dB, where XJ go is drawn from the target distribution over Rd and B, is the standard Brownian motion in Idd. It is well known that x/ d ''6 + v1-ca 21-2-222
where 7 (O, I) independent of X/_
5

---

Published as a conference paper at ICLR 2026
Let di be the density of the law of X Diffusion models learn the score function [0, T] x R
(t, X + log (X) via score matching (see Appendix A for literature review on score matching). P-GRAFT, in contrast, attempts to learn V log qs between for & E [1, At time T, log QT (T()) R -X, , the score of the standard Gaussian distribution, which is easy to learn. When t = 0, the score log (x)) corresponds to the data distribution which can be very complicated. Diffusion models via Denoising Score Matching, based Tweedie's s formula introduced by (Vincent 2011). We show theory (Bakry al. 2013) that the score function V log qe (x( converges to 9 (X) exponentially in 4, potentially making the learning easier. Consider s (X, t) RR x IR R to be neural network with parameters then score matching objective is given by:
( 1 [ dell xx
so(X!
I-
In practice,
[()) is approximated with samples. By Tweedie's formula, we have: B X2 X - loog (X:)) Thus, for some constant C, independent of
(( + -= l V V log (x() so(X! - L-" di ffe dx qx(X) log q1 (X) so ()]??
As shown by [Benton et al. 2023), [(0) directly controls the quality of generations. Note that d:o is
the density of (N, I) and V log fxc (X) X. The theorem below is proved in Appendix D.5
Theorem 3.4. Define Hद for S 1: Hi = - Ses dX (x) lo (X) V log tox (x)]12 Then,
Hf S e 2t
1- H
Therefore, the score functions between time (t, T) are exponentially closer to the simple Gaussian score function compared the score functions between times (0, in the precise sense given in Theorem 3.4 This means that the score functions at later times should be casier to learn.
4 INVERSE NOISE CORRECTION FOR FLOW MODELS
In the analysis so far, we have looked at inter- mediate distribution shap- ing in the context of KL regularized reward max- imization. particular, have we have estab- lished bias-variance trade- offs for diffusion models models which use SDEs to sample from target distribution. We now consider rectified flow models, which follow a deterministic ODE start- ing from an initial (random) noise. Hence, the final generated sample completely determined by the initial noise. Further, the bias-variance results from the previous section indicate that the learning process is potentially easier at the initial noise level.
Naise Distribution (po)
Image Distribution (p""))
Leamedd Distribution (p:)
Inverse Noise Distribution (p!")
ve (z,
(2, 8)
Par (z,
Figure 2: Inverse Noise Correction Setup
Since the final sample distribution is completely determined by the initial noise distribution, ask the following question: can we correct for errors in the (learned) final distribution the pre-trained model by correcting the initial noise distribution? Along with results from the fars-nariance tradcoff discussion, we utilize another property of flow models: they admit an exact reversal. This property has been utilized extensively in the literature map images 'noise" for image editing |Ruut et al. (2024); Garibi et al. (2024) and as part of the 2-rectification Liu et al. (2022) to achieve straighter flows. We will combine these two ideas to answer the above question and develop framework for improving flow models even without explicil rewards, by correcting for errors in the pre-trained distribution. We now develop this idea from first principles.
We restrict our attention to flow models with optimal transport based interpolation (Lipman et al. 2022 Liu et al. 2022), which leam velocity field v(x, () Idd x [0, 1] IRd such that the following

---

Published as a conference paper at ICLR 2026
ODE's solution at time 1 has the target distribution pluta:
dX. = =(eett), Xo N(O,T.
(2)
Note that in the literature for flow models (unlike diffusion models), t == 0 corresponds to noise and I corresponds to the target, a convention we follow in this section.
The errors in learned model: Suppose we have pre-trained vector-field, corresponding to parameter
0 and solve the ODE equation with v(r, t) 00(x, t). Then, Law(X1) # p ydaia due to:
a) Discretization error of the ODE and b) Statistical error due to imperfect learning.
Despite these two errors, the trained ODE is still invertible. We will leverage reversibility to arrive at our algorithm. To this end, consider the time reversal of equation 2
dx dt
(are
1 t),
(0) d
(3)
Algorithm 3 Inverse Noise Correction: Training
Input: Dataset D == {x((), (2)) ; x(O) duta step-size 77, backward Euler steps No
Algorithm 4 Inference
Input:
Flow models uo, thw,, step-size Initial point Xo (D,,))
1: mg TRAIN FLOW(N(()) D).
2: for i = 1 to M do
3: (eex,()) + BWD Euler Us Ty x(0,,N) BWD__ulerr
I: X1 FWD Euler(ugr, Xo)
4: end for
2: X1 + FWD FWD__uler((p, X
5: Dataset Dr + (x(ev.(1)
X2oo ~ pl
3: return X1
6: Ugre- =TRAIN_ FRRIINFLOWWW((..)) Dray)
7: return 0g, Up
The Inverse Noise: Consider the forward Euler discretization of equation 2 with step-size ry:
f(++1)n <- Rin, + vyul (# in))
(4)
Let To., be the function which maps to to 21 i.e. 3t To, (oo). The foward Euler approximation To.y (21) s yi where 9., < (6-1)) T) () 1 (i 1)m) with do â. is not good enough as noted in the image inversion/ editing literature Rout et al. (2024); Wang et al. (2024); Garibi et al. (2024). This mitigated via numerical and control theoretic techniques. We utilize the "backward Fuler discretization' (equation 3 as used in Garibi et al. (2024)) to exactly invert equation
geev + Rntit 1y0 (
7 1))
(5)
This is an implicit equation since Jeey being calculated in the LHS also appears in the RHS. It is not apriori clear that this can be solved. Lemma 4.1 addresses this issue:
Lemma 4.1. Suppose vo Lipschitz in x under fn-norm and nL Then,
(I) any in equation 5 has unique solution which can be obtained by fixed point method.
(2) To is invertible and To., (x:(e) 2
That is, the mapping from noise to data given by the learned, discretized model invertible. We show some important consequences of this in Lemma 4.2) Define the following probability distributions.
Let datn LawaData) (i.e, target data distribution).
Pg Law( I0. P1 Law(
p6 Law
(eevv)
data
Law
(areev
We call pfo the inverse noise distribution. With perfect training and 0 discretization error, p!ay N I). However, due to these errors pf + (O.
Lemma 4.2. Under the assumption of Lemma 4.1 Dit, Pv, pdota and po = =N((O, satisfy:
(To.) -dota
2. TV po) (Vl,,, 3. (L(pollx KL KL dada).
7

---

Published as a conference paper at ICLR 2026
That is, the distance between the inverse noise and the true noise is the same as the distance between the generated distribution and the true target distribution. Item shows that if we can sample from the inverse noise distribution pf then
can the pre-trained model vo (:)::)) with discretization and obtain samples from the true target ydaaa In Kim et al. (2024), the authors note that even 2-rectification suffers when the inverse noise o/ is far from [Q, Ij. While 2-rectification aims to improve improve the computational complexity while maintaining quality by aiming to obtain straight flows, we introduce inverse noise correction to improve quality of generations a sample efficient way.
Inverse Noise Correction: Inverse Noise Correction is given in Algorithms 3 and 4 and illustrated
in Figure 2 Given samples from the target distribution, D, TRAIN FL FLOW(N (0, I), D) trains a rectified flow model between (N0, I) to the target distribution Liu et al. (2022). Now, suppose we are given dataset (X(1), (X.,),.., X(M) -yaat and a trained flow model vo which generates i1 pi
using equation starting with dn ~ Po- We obtain samples x ~ p/"" by backward Euler iteration in equation 5 Thereafter, we train another flow model Ugrw which learns to sample from p/""
starting from (V,
During inference, we sample point from Xo (O, and obtain sample X p/ using Ugm. Once we have the corrected noise sample, we generate images using the original flow model 0g which now starts from X instead of Xg. FWD "0__Eulee(e,,. 7, zu) obtains i1 via Euler iteration
(equation 4). Similarly, BWD Euler (vo 77. 40 .N)) obtains 21" by approximately solving backward Euler iteration (equation 55. They are formally described as Algorithms 5 and 6 in Appendix B
Theoretical Justification along the lines of Section 3.1 is given in Appendix D.8
5 EXPERIMENTS
We the notation P- to denote P-GRAFT with intermediate timestep N, as described in Algorithms [ and 2 For instance, -GRRAT(0.755.) would denote instantiating P-GRAFT with Nj 0.75.V, where N is the total number of denoising steps. Recall that ty corresponds pure noise and to corresponds to completely denoised sample.
5.1 TEXT-TO-IMAN IMAGE GENERATION
The objective is to fine-tune pre-trained model so that generated images better align with prompts. We consider Stable Diffusion v2 (Rombach et al., 2022) as the pre-trained model. The reward model used is VQAScore (Lin et al. (024) a prompt-image alignment score between 0 to 1, with higher scores denoting better prompt-alignment. We fine-tune (separately) on GenAl-Bench (Li et al. 2024a) as well as the train split of T2ICompBench++ (Huang et al. 2025). Evaluations are done on GenAl-Bencn
Setup:
validation split of T2ICompBench++ and GenEval (Ghosh et al. 2023). We use LoRA (Hu et al. 202I) for compute-efficient fine-tuning. Top K sampling (Section 2.3 is used for both GRAFT and P-GRAFT. Since LoRA fine-tuning is used, the model switching in can be done by simply turning off the LoRA adapter. More implementation details are given in Appendix E
are reported in Table - for fine-tuning on GenAl-Bench, we use Top - 10 of 100 samples and on T2ICompBench++. we use Top - of 4 samples. First, note that both GRAFT and P- GRAFT outperform base SDv2, SDXL-Base and DDPO. The best performance is obtained for P-GRAFT with N, = 0.25.N across all evaluations this clearly shows the bias-variance tradeoff in action. Further, both GRAFT and P-GRAFT also generalize to unseen prompts.
Results:
In particular, DDPO did not improve over the baseline even when trained with more samples and FLOPs as compared to GRAFT/P- PRAFT/P-GRAFT. Experiments with different sets of hyperparameters as well as adding other features such as KL regularization and per-prompe advantage estimator on top of DDPO also did not show any significant improvements over SDv2 (see Appendix E.3). We also conduct ablations to further verify the effectiveness of the proposed methods these include experiments on different values of (M, K) in Top K of M sampling, different LoRA ranks for fine-tuning as well as a reverse P-GRAFT strategy (where the fine-tuned model is used in the later denoising steps instead of initial steps). We find that P-GRAFT remains effective across different (M, K) and that performance is insensitive to the LoRA rank. Further, P-GRAFT significantly outperforms reverse P-GRAFT. More details ablations be found in Appendix E.
8

---

Published as a conference paper at ICLR 2026
Table 2: Text-to-Image Generation fine-tuning on SDv2: VQAScore (normalized to 100) reported on GenAl-Bench, 12ICompBench++ Val (denoted as T2I Val) and GenEval.
Fine-Tuned on GenAl Beach
Fine-Tuned on T2ICompBench++ - Train
Model
GenAl
T2I Val
GenEval
GenAl
T2I- Val
GenEval
SD v2
66.87=0.14
69.20
49.0..1
06.87
09.20=0.17
73.49 73.49
SDXL-Buse
69. 09.69-0.17
72. 22.98
73.90=0.4
19.17
72.98+0.16
73.90+0.40
DDPO
65.70
68.0320.16 16
72.13
64.65119.17
69.0620.15
69.60100.37
GRAFT
70.51=0.15
75.69
29.85
70.97
75.88
79.57
)
69.46-0-15
10.11
79.44.0033
69.51
74.30a0:15
78.5010.33
P-GRAFT P-GRAFT(0.25N)
71.00=0.14
75.45=0..14
80.60=0.51
70.73+0.14
75.37=0.12
79.25++0.30
71. 71.9440.14
76.12+0.13
80. 50.96+0.29
71.42=00:4
16.15
80.29
Table 3: Layout Generation: Fine-tuning results for unconditional and category- conditional generation on PubLayNet.
Table 4: Molecule Generation: Fine-tuning re- sults QM9, (Relative) number of sampling rounds required are also reported.
Unconditional
Class-conditional
Model
Mol: Stability
Sampling Rounds
Model
Baseline
90.50
Aligament
FID
Aligument
FID
GRAFT
90. 76:...20
9x
Baseline
0.094
8.32
0.088
4,08
P-GRAFT(O.5N)
90.4610.27
1x
GRAFT
0.064
10.08
0.068
5.04
P-GRAFTIO.22 25N)
92.61:0.13
1x
P.GRAFTO.5N)
0.071
0.072
4.55
P-GRAFTO0255)
0.053
9.91
0.064
4.67
5.2 LAYOUT AND MOLECULE GENERATION
Setup: All experiments are done on pre-trained models trained using IGD (Anil et al. 2025 a diffusion framework capable of handling both layout generation and molecule generation. For layouts, we experiment with improving the alignment of elements in the generated layout as measured by the alignment metric note that the reward is taken as 1- alignment since lower values for the metric indicate better alignment. For molecules, the objective is to generate a larger fraction of stable molecules molecules which are deemed stable are assigned a reward of whereas unstable molecules are assigned reward of 0. For molecule generation, we usee the de- de-duplication instantiation of GRAFT/P-GR GRATT (Section 2.3p to ensure diversity of generated molecules we use RDKit to determine whether two molecules are identical or not. We use PubLayNet (Zhong) (ct al., 2019) for layout generation, and QM9 (Ramakrishnan et al. 2014) for molecule generation. To the best of our knowledge, this is the first work which addresses fine-tuning in the context of discrete-continuous diffusion models. Ablations and experimental details are given in Appendix
Results: for layout generation are given in Table 3 Both G-GAFT and GRAFT uniformly improve performance across both unconditional and clas-coonitional generation, with P.GRAFT:0.25. giving the best performance. We also report FID scores computed between the generated samples and the test set of PubLayNet- this is a measure of how close the generated samples are to the pre-training distribution. As expected, the baseline has the lowest FID. Note that the FID score for P-GRAFT is smaller than GRAFT, indicating that P-GRAFT aligns more closely to the pre-training distribution. For molecule generation, results are given in Table 4 Again, the best performance is with P-GRAFT at 0. 0.25N. Note that improvement with GRAFT is marginal, despite being trained on 9x the number of samples used for P-GRAFT this points to the learning difficulty in later denoising steps.
5.3 IMAGE GENERATION WITH INVERSE NOISE CORRECTION
Setup: We consider unconditional image generation on CelebA-HQ (Karras et al. 2017) and LSUN- Church (Yu et al., 2015) at 256 x 256 resolution. We first train pixel-space flow models from scratch. A training corpus of inverse noise is then generated by running the trained flow models in reverse, employing the backward Euler method, on all samples in the dataset. A second flow model, which we refer to as the Noise Corrector model, is then trained to generate this inverse noise. Once the Noise Corrector is trained, this model is first used to transform standard Gaussian noise to the inverse noise. The pre- pre-trained model then generates samples starting from the inverse noise. FID with 50000 generated samples with respect to the dataset is used to measure the performance. We emphasize that the our goal is not to compete with of-te-of-the-an (SOTA) models rather demonstrate that

---

Published as a conference paper at ICLR 2026
Table Image Generation: Results for inverse noise correction on CelebA-HQ and LSUN-Church. The noise corrector samples the inverse noise starting from (0, I) for 'Sampling Steps', and the pre-trained model samples the image starting from the inverse noise.
Sampling Steps
FID
FLOPs/inage (x10*9)
Noise Corrector
Pre-Traiind Model (65M parameters)
CelebA-HQ (256 256)
LSUN-Church (256 x 256)
(16M purameters)
1000
11.93
8.40 8.63
0,869
200
13.39
1.374
100
100
8.94
7.90
0,903
200
200
8.02
7.26
1.806
our procedure can be used to improve the performance given flow model by simply learning the distributional shift of noise at t = 0. SOTA models are larger ( [Rombach et al. (2022) has == 300M parameters) and are more sophisticated
do not seek match their performance.
Table 5 shows that the Noise Corrector signiticantly improves FID scores across both datasets, Apart from quality gains, Noise Corrector also allows for faster generation running the Noise Corrector for 100 steps and then running the pre-trained model for 100 steps can outperforms the pre-trained model with 1000 steps. The Noise Corrector only has 0.25x the number of parameters, leading to further latency gains evidenced by FLOPs counts.
Results:
CONCLUSION
We establish GRAFT, framework for provably performing marginal KL regularized reward maxi- mization for diffusion models through rejection sampling. We then introduce P-GRAFT, a principled framework for intermediate distribution shaping of diffusion models and provide mathematical jus- tification for this framework. Both GRAFT and P-GRAFT perform well empirically, outperforming policy gradient methods the text-to-image generation task. Further, both frameworks also extend seamlessly to discrete-continuout diffusion models. Finally, we introduce Inverse Noise Correction, strategy improve flow models even without explicit rewards and demonstrate significant quality gains even with lower FLOPs/image.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1tXiv:1707.06377 2017.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training helpful and harmless assistant with reinforcement learning from human feedback. ar preprint arXiv:2204.05862. 2022.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang. Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730- 27744, 2022.
Kevin Black, Michael Janner, Yilun Du, llya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. arXiv preprint arXiv:2305. 13301, 2023.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning text-to-imaee diffusion models. Advances in Neural Information Processing Systems, 36:79858-79885, 2023.
Fei Deng, Qifei Wang, Wei Wei, Tingbo Hou, and Matthias Grundmann. Prdp: Proximal reward difference prediction for large-scale reward finetuning of diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7423-7433, 2024.
Carles Domingo-Enrich, Michal Drozdzal, Brian Karrer, and Ricky TQ Chen. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. ICLR 2025, arXiv:2409.0888 2024.
Masatoshi Uchara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali. Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Fine-tuning of continuous-ting diffusion models as entropy-regularizee control. arXiv preprint arXiv:2402.15194, 2024.
Tianqi Liu, Yao Zhao, Rishabh Joshi, Misha Khalman, Mohammad Saleh, Peter J Liu, and Jialu Liu. Statistical rejection sampling improves preference optimization. arXiv preprint N1Xiv:23090066575 2023a.
Hanze Dong, Wei Xiong, Deepanshu Goyal, Yihan Zhang, Winnie Chow. Rui Pan, Shizhe Diao, Jipeng Zhang, Kashun Shum, and Tong Zhang. Raft: Reward ranked finetuning for generative foundation model alignment. arXiv preprint 234402304 2023.
Wei Xiong, Jiarui Yao, Yuhui Xu, Bo Pang, Lei Wang, Doyen Sahoo, Junnan Li, Nan Jiang, Tong Zhang, Caiming Xiong, et al. A minimalist approach to IIm reasoning: from rejection sampling to reinforce. arXiv preprint arXiv:2504.1 11343, 2025.
Afra Amini, Tim Vieira, Elliont Ash, and Ryan Cotterell. Variational best-of-n alignment. arXiv preprint arXiv:2407.06057. 2024.
Lin Gui, Cristina Gårbacea, and Victor Veitch. Bonbon alignment for large language models and the sweetness of best-of-n sampling. Advances in Neural Information Processing Systems, 37:
2851-2885, 2024.
Nisan Stiennon, Long Ouyang, Jeffrey Wu, Daniel Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul F Christiano. Learning to summarize with human feedback. Advances in neural information processing systems, 33:3008-3021. 2020.
Nick T Thomopoulos. Essentials of Monte Carlo simulation: Statistical methods for building simulation models. Springer Science & Business Media, 2012.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011 I, 131466, 2020a.
11

---

Published as a conference paper at ICLR 2026
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):161-1167, 2011.
Dominique Bakry, Ivan Gentil, and Michel Ledoux. Analysis and geometry of Markov diffusion operators, volume 348. Springer Science Business Media, 2013.
Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Nearly d-linear conver- gence bounds for diffusion models via stochastic localization. arXiv preprint arXiv:2308.0366 2023.
Litu Rout, Yujia Chen, Nataniel Ruiz, Constantine Caramanis, Sanjay Shakkottai, and Wen-Sheng Chu. Semantic image inversion and editing using rectified stochastic differential equations. arXiv preprint arXiv:2410.10792, 2024.
Daniel Garibi, Or Patashnik, Andrey Voynov, Hadar Averbuch-Elor, and Daniel Cohen-Or. Renoise: Real image inversion through iterative noising. In European Conference on Computer Vision, pages 395 413. Springer, 2024.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. arXiv preprint MiXiv:220.003003 2022.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamm Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint #xXiv:2210.02747, 2022.
Jiangshan Wang, Junfu Pu, Zhongang Qi, Jiayi Guo, Yue Ma, Nisha Huang, Yuxin Chen, Xiu Li, and Ying Shan. Taming rectified flow for inversion and editing. arXiv preprint arXiv:2411 00e:2411.04746, 1.04746, 2024.
Beomsu Kim, Yu-Guan Hsieh, Michal Klein, Marco Cuturi, Jong Chul Ye, Bahjat Kawar, and James Thortton. Simple reflow: Improved techniques for fast flow models. arXiv preprint arXiv:2410.08815, 2024.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models, 2022. URL
abs/2112. 10752
Zhiqiu Lin, Deepak Pathak, Baiqi Li, Jiayao Li, Xide Xia, Graham Neubig, Pengchuan Zhang, and Deva Ramanan. Evaluating text-to-visual generation with image-to-text generation, 2024. URL
https::/arxiis hwwp://wwwwwcrc//////////// 01291
Baiqi Li, Zhiqiu Lin, Deepak Pathak, Jiayao Li, Yixin Fei, Kewen Wu, Tiffany Ling, Xide Xia, Pengchuan Zhang, Graham Neubig, and Deva Ramanan. Genai-bench: Evaluating and improving compositional text-to-visual generation, 2024a. URL petps::/arrivye httpaaawwwv.ggaaaaaaeeeea org abss////////// 24006
13743
Kaiyi Huang, Chengqi Duan, Kaiyue Sun, Enze Xie, Zhenguo Li, and Xihui Liu. T2i-compbench++: An enhanced and comprehensive benchmark for compositional text-to-image generation, 2025.
06350
Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt. Geneval: An object-focused framework for evaluating text-to-image alignment. Advances in Neural Information Processing Systems, 36: 52132-52152, 2023.
Edward J. Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang. and Weizhu Chen. Lora: Low-rank adaptation of large language models, 2021. URL https
//axxiv.org/aaaa 2w06 09685
Gautham Govind Anil, Sachin Yadav, Dheeraj Nagaraj, Karthikeyan Shammugam, and Prateek Jain. Intericaved gibbs diffusion for constrained generation. arXiv preprint arXiv:2502.134500 2025,
Xu Zhong, Jianbin Tang, and Antonio Jimeno Yepes. Publaynet: largest dataset ever for document lay- out analysis. In 2019 International Conference on Document Analysis and Recognition (ICDAR), 2019. doi: 10.1 109/ICDAR 2019.00016..
12

---

Published as a conference paper at ICLR 2026
Raghunathan Ramakrishnan, Pavlo O Dral, Matthias Rupp, and O Anatole Von Lilienfeld. Quantum chemistry structures and properties of 134 kilo molecules. Scientific data, 1(1):177, 2014.
Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of gans for improved quality, stability, and variation. arXiv preprint arXiv:1710.101%6, 2017.
Fisher Yu, Ari Seff, Yinda Zhang, Shuran Song, Thomas Funkhouser, and Jianxiong Xiao. Lsun: Construction of a large-scale image dataset using deep learning with humans in the loop. arXiv preprint arXiv: 1..Kvv1506.03.5 506.03365, 2015.
Ronald J Williams and Jing Peng. Function optimization using connectionist reinforcement learning algorithms. Connection Science, 3(3):241-268, 1991.
Hao Liu, Carmelo Sferrazza, and Pieter Abbeel. Chain hindsight aligns language models with feedback. arXiv preprint arXiv:2302.022 02666, 2023b.
Ziniu Li, Tian Xu, Yushun Zhang, Zhihang Lin, Yang Yu, Ruoyu Sun, and Zhi-Quar Luo. Remax: A simple, effective, and efficient reinforcement learning method for aligning large language models. arXiv preprint XrXiv:2310 10.01.11....... 2023.
Arash Ahmadian, Chris Cremer, Matthias Gallé, Marzich Fadace. Julia Kreutzer, Olivier Pietquin, Ahmet Üstün, and Sara Hooker. Back to basics: Revisiting reinforce style optimization for learning from human feedback in lims. arXiv preprint wrXiv:2402.14740, 2024.
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi. Haowei Zhang, Mingchuan Zhang, YK Li, Yang Wu, et al. Deepseekmath: Pushing the limits mathemat- ical reasoning in open language models. arXiv preprint arXiv:2402.03300, 2024.
Jian Hu, Jason Klein Liu, Haotian Xu, and Wei Shen. Reinforce++: An efficient rlhf algorithm with robustness to both prompt and reward models. arXiv preprint arXiv:2501 .03662, 2025.
Allen Z Ren, Justin Lidard, Lars L Ankile, Anthony Simeonov, Pulkit Agrawal, Anirudha Majumdar, Benjamin Burchfiel, Hongkai Dai, and Max Simchowitz. Diffusion policy policy optimization. Xrviv preprint XiX::22409000888 2024.
Yao Zhao, Rishabh Joshi, Tianqi Liu, Misha Khalman, Mohammad Salch, and Peter J Liu. Slic-hf: Sequence likelihood calibration with human feedback. arXiv preprint arXiv:2305. 1ttt2235550442. 10425, 2023.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in neural information processing systems, 36:53728-53741 2023.
Yu Meng, Mengzhou Xia, and Danqi Chen. Simpo: Simple preference optimization with reference- free reward. Advances in Neural Information Processing Systems, 37:
2024.
Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8228-8238, 2024.
Shufan Li, Konstantinos Kallidromitis, Akash Gokul, Yusuke Kato, and Kazuki Kozuka. Aligning dif- fusion models by optimizing human utility. Advances in Neural Information Processing Systems, 37:24897-24925. 2024b.
Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffusion models on differentiable rewards. arXiv preprint arXiv:2309.1744: 2023.
Mihir Prabhudesai, Russell Mendonca, Zheyang Qin, Katerina Fragkiadaki, and Deepak Pathak. Video diffusion alignment via reward gradients. arXiv preprint wxXiv:2407.08731 2024.
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score matching.
Journal of Machine Learning Research, 6(4), 2005.
Durk P Kingma and Yann Cun. Regularized estimation of image statistics by score matching.
Advances in neural information processing systems, 23, 2010.
13

---

Published as a conference paper at ICLR 2026
Syamantak Kumar, Dheeraj Nagaraj, and Purnamrita Sarkar. Dimension-free score matching and time bootstrapping for diffusion models. arXiv preprint arXiv:2502.10334, 2025.
Adam Block, Youssef Mrouch, and Alexander Rakhlin. Generative modeling with denoising auto- encoders and langevin sampling. arXiv preprint arXiv:2002.00107, 2020.
Shivam Gupta, Aditya Parulekar, Eric Price, and Zhiyang Xun. Improved sample complexity bounds for diffusion model training. Advances in Neural Information Processing Systems, 37: 40976-41012, 2024.
Minshuo Chen, Kaixuan Huang, Tuo Zhao, and Mengdi Wang. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data. In International Conference on Machine Learning, pages 4672-4712 PMLR, 2023.
Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. In Uncertainty in artificial intelligence, pages 574-584. PMLR, 2020b.
Valentin De Bortoli, Michael Hutchinson, Peter Wirnsberger, and Arnaud Doucet. matching. arXiv preprint arXiv:2402.08667. 2024.
Target score
Amir Hertz, Ron Mokady, Jay Tenenbaum, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. Prompt- to-prompt image editing with cross attention control. arXiv preprint (eXXiv:2208.06665 2022.
Gwanghyun Kim, Taesung Kwon, and Jong Chul Ye. Diffusionclip: Text-guided diffusion models for robust image manipulation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 2433,, 2022.
Seongmin Hong, Kyeonghyun Lee, Suh Yoon Jeon, Hyewon Bae, and Se Young Chun. On exact inversion of dpm-solvers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 7069-7078, 2024.
Ron Mokady, Amir Hertz, Kfir Aberman, Yael Pritch, and Daniel Cohen-Or. Null-text inversion for editing real images using guided diffusion models. Proceedings the IEEE/CVF conference computer vision and pattern recognition, pages 6038-6047 2023.
Sangyun Lee, Zinan Lin, and Giulia Fanti. Improving the training of rectified flows. Advances in neural information processing systems, 17:63082-61109 2024.
Yuanzhi Zhu, Xingchao Liu, and Qiang Liu. Slimflow: Training smaller one-step diffusion models with rectified flow. In European Conference on Computer Vision, pages 342-359, Springer, 2024.
Xingchao Liu, Xiwen Zhang, Jianzhu Ma. Jian Peng, et al. Instaflow: One step enough for high-quality diffusion based lext-to-image generation. In The Twelfth International Conference on Learning Representations, 2023c.
Luca Eyring. Shyamgopal Karthik, Alexey Dosovitskiy, Nataniel Ruiz, and Zeynep Akata. Noise hy- pernetworks: Amortizing test-time compute in diffusion models. arXiv preprint arXiv:2508.099 009668, 2025.
Santosh Vempala and Andre Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. Advances in neural information processing systems. 32, 2019.
Leonard Gross. Logarithmic sobolev inequalities. American Journal of Mathematies, 97(4):1061- 1083, 1975.
John Charles Butcher. Numerical methods for ordinary differential equations. John Wiley Sons, 2016.
Bradley Efron. Tweedie's 's formula and selection bias. Journal the American Statistical Association, 1064496);1602-1614 201 1.
14

---

Published as a conference paper at ICLR 2026
Patrick Platen, Suraj Patil, Anton Lozhkov, Pedro Cuenca, Nathan Lambert, Kashif Rasul, Mishig Davaadorj. Dhruv Nair, Sayak Paul, William Berman, Yiyi Xu, Steven Liu, and Thomas Wolf. Diffusers: State-of-the of-the-art diffusion models. httpp::/ggithb...o attpss//ggthhbb..cmmuuugggggoge
di diffusers 2022.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B. Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models, 2020. URL https: http:://arxiv..ooo org.oog///aae abs 220020000111111
Jose Lezama, Tim Salimans, Lu Jiang, Huiwen Chang, Jonathan Ho, and Irfan Essa Discrete predictor-corrector diffusion models for image synthesis. In The Eleventh International Conference on Learning Representations, 2022.
Yixiu Zhao, Jiaxin Shi, Feng Chen, Shaul Druckmann, Lester Mackey, and Scott Linderman. Informed correctors for discrete diffusion models. arXiv preprint arXiv:240.221243, 2024.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances Neural Information Processing Systems, 01122266-282799 2022.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. Advances in Neural Information Processing Systems, 37: 133345- 113445-13335,, 2024.
15

---

Published as a conference paper at ICLR 2026
APPENDIX
A RELATED WORK
Policy Gradient Methods: Majority of the existing literature on policy gradient methods in the context of generative modeling draw inspiration from Proximal Policy Optimization(PPO) (Schulman al. 2017) and REINFORCE (Williams and Peng. 1991). PPO based methods in the context of language modeling include Bai et al. ): Ouyang al. (2022); Liu et al. (202bb; [Stiennon et al.
(2020), whereas frameworks based on REINFORCE include (Li al. 2023 Ahmadian et al. 2024]
Shao et al. 2024| Hu al. 2025). Policy gradient methods have also been studied in the context of fine-tuning diffusion models (Black et al. 2023 Fan et al. 2023 Ren et al. 2024).
Offline MunenT Methods: Algorithms which utilize offline preference datasets for fine-tuning generative models have also been widely studied. the context of language modeling, these include methods like SLiC (Zhao et al., 2023), DPO (Rafailov al. 2023) and SimPO (Meng et al. (0224).
Such methods have also been explore in the context of diffusion models as well these include methods like Diffusion-DPO (Wallace al. 2024) and Diffusion-KTO (Li al. 2024b).
Recently, many works have explored rejection sampling methods in the context of autoregressive models these include RSO (Liu et al.) (202a)), RAFT (Dong et al., 2023) and Reinforce-Rei (Xiong et al. (202).. In particular, Reinforce-Pe Rej demonstrated that rejection sampling methods can match or even outperform policy gradient methods.
Rejection Sampling Methods:
Fine-Tuming Diffusion Models: Apart from the policy gradient methods discussed already, a host of other methods have also been proposed for fine-tuning diffusion models. Direct reward backpropagation methods include DRaFT (Clark et al. 2023) and VADER (Prabhudesai et al. (024). Note that these methods assume access differentiable reward. Uehara et al. (2024) approaches the problem from the lens of retropy-r control however, the method is computationally heavy and requires gradient checkpointing as well as optimizing an additional neural SDE. Domingo-Enrich et al. (2024) proposes memoryless forward process to overcome the initial value function bias problem for the case of ODEs. PRDPDDng et al. (2024) formulates a supervised learning objective whose optimum matches with the solution to PPO, but with trajectory KL constraint the supervised objective, with clipping, was found to make the training stable as compared to DDPO.
first introduced in (Hyvärinen and Dayan 2005). The algorithm used in this case is called Implicit Score Matching. Diffusion models primarily use Denoising Score Matching (DSM), which is based Tweedie's 's formula (Vincent)
Score Matching: Score matching for distribution estimation
2011; Kingma and Cun, 2010). The sample complexity of DSM has been extensively studied in
the literature (Kumar al.) 2025 Block al. 2020 Gupta et al. 2024 Chen al. M.ny
alternative procedures such as Sliced Score Matching (Song et al. 2020b) and Target Score Matching
(De Bortoli et al. 2024 have been proposed.
ODE Reversal in Flow Models: A prominent use case of ODE reversal in flow models is that of
image editing (Hertz al. 2022 Kim al. 2022 Hong et al. 2024 Mokady al. 2023 Rout et al.)
2024 Garibi et al. 2024). The reverse ODE has also been used to achieve straighter flows, allowing for faster generation, through 2-rectification/reellww algorithm (Liu et al., [0022] Lee al.) 2024] Zhu al. 2024 Liu al.) (203c).. Notably, concurrent work Eyring et al. (2025) also proposes a strategy
for aligning distilled models by fine-tuning at the noise level.
16

---

Published as a conference paper at ICLR 2026
B ODE SOLVER ALGORITHMS
In the backward Euler Algorithm 6 at each time instant j in the reverse procedure, we solve fixed point equation obtain high precision solution of Eq. equation The step-size 7) tuned empirically so that the recursion does not blow up. Once the step-size is carefully tuned, the iteration converges the solution an exponential rate. In practice, we observed that No is sufficient to obtain satisfactory results.
Algorithm 6 Backward Euler (BWD_Euler)
Input: Flow model ug. step size τ), sample ( from the dataset, Number of fixed point iterations No
Algorithm 5 Forward Euler (FWD_Euler)
Input: Flow model Ug, step-size 17, Initial point Xo
1: X (()
2: for j = 0 to [1/n]
3: X = Xiev
1: for to |1/1| do
2: X,+1 t X, + (X,,
4:
for t N₂ do
5: X + XSSv mos
m76 + 1))
3: end for
4: return X(1/n]
end for
7: X + X
end for
9: return X [1/ml
C GRAFT: ALGORITHM
While instantiations of GRAFT are well-known in the literature and straightforward implement, we provide the exact algorithm here for the sake of completeness.
Algorithm 7 GRAFT: Training
Input: Trainable po, Reference p, Reward function Acceptance function A, Number of sampling rounds Ns
Initialize empty set
2: for i = 0 to Ng do
3: Get M samples: (X)....... X(M)
(MM). ~
4: Obtain rewards: r(x()); E [M
5: Perform GRS using acceptance function A to get accepted samples A
6: Perform D+ D
7: end for
8: Train Po on D
9: return P8
17

---

Published as a conference paper at ICLR 2026
D PROOFS
D.1 LEMMA 2
Proof. Let B be any measurable Consider the following probability measure:
P( E B\X() AA)
Using Bayes' rule, this measure be rewritten
P(X B|X E P(X B, x() EA). P(X()) A)
Recall that ( is drawn from the distribution p. Then, from the definition of P(X E B, B, A), we have:
P(X E B, B, E PP((()) e = =x)d((x).
=
From Definition 2.2 we know that ( E Aiff C1 = 1. Therefore:
P(X E B, A)
PCC1 1 x))p(x))
E 1 = 1) 1) = =x] d7(x))
where I(-) denotes the indicator function. rewritten
Using the tower property expectations, this can be
/ E E [L(C] 1)(x()= = x, x(1 x(mO)) (100))1x0)) = =x] dp(x)
P(X E B, E A) =
{ R ( 1|10= =,, X x((n) |x
) = =s] dp(x)..
Note that in the conditional expectation here, X(),... X(M), are distributed according to po since (X Jj=1 are i.dd samples. Again, from Definition 2.2 we know that
P(C1 1|x =x, =x, ( A(r(z), P,(r())),, Px)
where FR and Px are computed using the samples プo From definition of Radon-Nikodym derivative, the distribution of the accepted samples can therefore be written as:
('(x) EE
2 p(z)
(6)
where 7 = 1/P(X(1) E A) is normalizing constant independent of æ. Now, from the method of Lagrangian Multipliers, mentioned in Section [ the solution to the RI. regularized reward maximization objective with reward function f(-) is given by:
(x) = Zz exp (f(x)) p(x)
where Zz is the normalization constant. Comparing equation and equation 7 n = RLL whenever:
f(x) log Fn((()),, Px|(() =x] =x)
a
18

---

Published as a conference paper at ICLR 2026
D.2 INSTANTIATIONS OF GRAFT
D.2.1 TOP-K OUT OF M SAMPLING
Substituting A(.) in:
log (B(xonyy
x1 (
we get:
log ( (Axcon 1(((x) Top r((())))((x())) (xx( dd)p((())) dp(a(ari)
where Top K(r(x), M(.
denotes the top-K fr(x), (xx(()))..... r()()))) Let UK denote the event where X among the M samples, where the other M samples i.id from p. This event be decomposed as:
(()(M))
samples in ranks in top-K
UK U
where Fg denotes the event where r(x) the t in the ranked (descending) ordering of rewards. Further, note that (Ex) are mutually exclusive events. Therefore:
p(Uk)= = E
Computing P(Ex): If x ranks 4tth when ranked in terms of rewards, there are A - 1 samples which have higher rewards than x and M k samples which have lower rewards than x. Thus, the required probability
be computed by finding the probability of having K -1 samples having higher rewards and the rest having lower rewards. Note that the ordering within the 1 group or M K group doesn't i matter. The probability of any sample having higher reward than r(x) - F(r(z)) and having a lower reward is F(r(z)). Therefore, the required probability can be computed as:
(M -1 P(Ex)
(NL (M-1) (1 (1-F))))) 1
And hence:
M - 1)0 -1- ( (((--((()))((F((()))**)))
K-1
P(Ux) =
Σ
k=0
Therefore:
f(x)
K-1
( (
M 111
log
k=
is straightforward to check that this an increasing function in r.
D.2.2 PREFERENCE REWARDS
Substituting A(.) in:
log (B =x,x(x
we get:
log ( 1(r(x(2) & r(E)))p(a(?1) log F(r(x))
19

---

Published as a conference paper at ICLR 2026
D.3 LEMMA 3.2
Proof. Let B be any measurable set. Consider the following probability measure:
P(X E BIX AA).
Using Bayes' rule, this measure can be rewritten as:
P(X!") B, B,X," EA). F(X!" E))
P(X!) BIX E E
Recall that x is drawn from the distribution p.. Then, from the definition of P(X!)) E B,x!" A), we have:
P(X B. E EA) = p(X,") EAX! = x))p((x).
From Definition 3.1
know that x E A iff C₁ 1. Therefore:
P(X()) B, x,) E
P(Ci =
= =1[x("==>)d)p (x)
/ E [ 1)(x =
(x)
where 1()) denotes the indicator function. Using the tower property expectations, this can be rewritten P(X E B. E
E (1(C) = 1)(x(' = x!" xx(? x( (xx" =x] dp)(x)
A)
E
B
E = = x, x!"), x(x) x...,,x x(ar)) x = x] dp)(
Note that in the conditional expectation here, x(2) ((M are distributed according to po
(i) L J==1 are i.idd samples. However, X (() is distributed according to Puje because of the since (x)
conditioning on x("). Again, from Definition 3.1 we know that
P(C = 1|xx') = (x
(0 ) = A(r(x))), (n(r(x( x.
Px) ,
where FR and Px are computed using the samples MM From the definition of Radon- Nikodym derivative, the density of the accepted samples can therefore be written as:
()(x) = Z,E
FR((X
X = 2)) (x(x))
(8)
where 21 = -/P(X A)) a normalizing constant independent of x. Now, from the method of Lagrangian Multipliers, as mentioned in Section 2 the solution to the RL regularized reward maximization objective (with reward function f(-)) is (where Zg is the normalization constant):
f(x:
RI. (x) Zz exp
(z)
(9)
Comparing equation 8 and equation g p
RL whenever:
f(x) log [ A(r(x) Fn(r(x()))) x,,""")) Px xx = x))
a
20

---

Published as a conference paper at ICLR 2026
D.3.1 DISTRIBUTION INDUCED BY P-GRAFT
P-GRAFT uses the fine-tuned model for early denoising steps and the reference model for later denoising steps. Let us call this effective model the "stitched" model. Let us denote the distribution the stitched model samples from as p². From the discussion above, the distribution the P-GRAFT fine-tuned model samples from at time t is pt Further, let Poj= denote the distribution of samples at time O, given a sample at time under the reference model. Then clearly:
p" (zo) (o) (x,)
= = / (z) exp
f(xx))
dr:
where Z is normalization constant independent of zo. general, an explicit solution to this integral not available.
To get more intuition regarding this distribution, we analyze two special cases. Assume that the acceptance probability depends only on r(Xo) and Pn(r(Xo)))
Case 1: The reward of sample time is independent of the latent time nr. The acceptance probability also independent of the latent T2-
From equation 10
F(x) = log E , = a))
log (E
since acceptance probability is independent of rt. Note that the expectation is over (2)
X. (M)
as discussed above. Hence, because of the independence assumption, log (E [A(-)] is independent of x(() and hence f(z) is independent of æ. Let us denote this quantity Z1. Then:
/
p²
Z po1((2ol2)) exp( dz,
/
= Z exp(21)
Z (exp(Zi))t
= p((x))
where we have used the fact that p(rr)) the normalized reference distribution. Hence, if acceptance probability of sample at time 0 is independent of the latent z,, the stitched model results in no rilt whatsoever.
Case 2: The reward of sample at time 0 is completely determined by the latent at time t 2z--
deterministic mapping ra exists such that r(Xo) rx(X)) A Xs, where Xo (-
Therefore, from equation 10
F(x) log (E (1) F,((xx" = =x)
= log (E ((((((())
Note that the expectation here is only over (2) ..... (M) because of the conditioning, And hence:
(xo)) = % f po exp f(x)) dra
= /
(ol( A(re(1x), Fn(re((x))) dr
21

---

Published as a conference paper at ICLR 2026
Using the fact that (x2)) r(zu), we have:
p" (2o) = Z Pol F FR(((2))))) dr
ZE A((((o)) FR(((Z0))))
/ (o Me (x.))da
Zp(ro)E A(r(zz)) (((((z))
Comparing p²(x0) with Lemma 2.3 we see that in this case, P-GRAFT results in sampling from the exact same distribution GRAFT.
In other cases, P-GRAFT interpolates between the reference distribution and the distribution induced by GRAFT.
D.4 PROOF OF LEMMA 33
Proof. Let s > t. Note that X, - Xx + Xo forms Markov chain. By the law of total variance, we have for any random variables Y, Z::
Var IEVar(Z)Y) + Var(B[z/\]))
2 EVar(Z|Y)
(10)
Given Xx, Suppose Z, Y be jointly distributed as the law of Xx). Then, have X, almost surely:
Va)()(X))|x..) 2 x -XVa(((
(1I)
In the last line, we have used the Markov property show that the law of rXx))|X,, X, is the same as the law of r (Xo)]Xe almost surely. We conclude the result by taking expectation over both the sides.
D.5 PROOF OF THEOREM 3.4
Proof. We will follow the exposition in Vempala and Wibisono (2019) for proofs. qe converges to Joo as t + co. By (Vempala and Wibisono 2019 Lemma 2) applied to the forward process, we conclude that:
kL((all00) -- d.x. (Xxx(X)[[Voogg(( V log goo (X)II
di
J." dt f xXq. log (X) (x)] KL(gellaw) KL(grl1 (12)
For brevity, we call the LHS to be H.. Clearly,
- H KL(gell KL((gollg>) KL(gr (l
Notice that 2= is the density of the standard Gaussian random variable. Therefore, it satisfies the Gaussian Logarithmic Soboley inequality Gross 1975). Thus, we can apply (Vempala and Wibisono
2019 Theorem 4) to conclude that for every 0, (Lealla -2KL(gol (gx))))).)).... (.o). Thus.
H < c 21
e Hb.
22

---

Published as a conference paper at ICLR 2026
D.6 PROOF OF LEMMA 4.1
The uniqueness and the convergence of fixed point iteration for implicit Euler methods have been established under great generality in Butcher (2016). However, give simpler proof for specialized setting here.
1. Consider the update for the backward Euler iteration each time step ryi
reey
( n7{ 1))
up
((-1)
Let us define an operator
T
R + Rd such that
8.0
T at
(x)
T 1(- mes 1 - n7(i - 1))
0,0
as defined above is contractive operator under the condition n < 1. Then, one can use Banach fixed point theorem to establish uniqueness of the solution and obtain the solution through fixed point iteration. To this end, consider two
First, we will show that T -
0.n
point 21 and 22 in Rd and apply To.
them
- (1==ii-1( (x1)
T 111--111
(122 =1]lv0(I1, 71i - 1)) - uo(22, t - 1))ll2
S EnL |21 22112.
20t is contractive operator. Thus, by Banach fixed Since 7L conclude that To.
point theorem, the fixed point equation T Ta.
T= (x) has unique solution for each step = mp. To obtain the solution to the backward Euler update, we use the Banach fixed point method, i.e., start with Z(()) = reyy 1)
(or any arbitrary point in IRd)) and run the iteration
T(k++)) = -00 (z(1)). Then, lima- I(a)) rere
2. The invertibility of the operator Ta,, follows directly from the previous part. Since the solution for the backward Euler method is unique at each time step t = ryi, it implies that there exists a one-to-one mapping between sample points x/) and x7
D.7 PROOF OF LEMMA
Before starting the proof of this lemma, we will state the following well-known theorem from information theory.
Theorem D.I. [Date Processing Inequality] Let X and y be two sample spaces. Denote P(x) and
P(Y) as the set of all possible probability distributions on X and y, respectively. Let Px, Qx E P()) and Py x be transition kernel. Denote Py and Qy be the push through, i.e., Py (B) Sx PyJx (B|X X =x)dPx(1)) Then, for any S-divergence have
(x 2 Dj (P
(13)
1. By part 3 of Lemma 4.1
have that rev = To (a
To.,(= 2/ rev Suppose
=0 p", then by demmition, 2 p. This concludes the result.
2. Recall that TV-norm is an f-divergence. Furthermore, To,n is the push forward function from po and pfth to p1 and p", respectively. Thus, using D.I have
TV(n1" .Po) 2 TV(p", ,P1).
Additionally, Te an invertible mapping. Hence, T can also be viewed as the push forward function from pr and p" po and p!er respectively. Thus, again using DPI D. get
TV(p", P1) 2 TV(p
Combining both the bounds, we get the desired claim.
3. KL divergence is also a valid f f-divergence. Thus, repeating the arguments from the previous part, one gets the desired equality.
23

---

Published as a conference paper at ICLR 2026
D.8 THEORETICAL JUSTIFICATION FOR INVERSE NOISE CORRECTION
In this section, our goal is to provide a theoretical justification for inverse noise correction in the context of flow models. Specifically, less challenging to learn the score function corresponding to and thereby the velocity field vt governing the rectified flow. To this end, let X be sample from distribution px and Z, Y be standard normal random variables all independent of cach other. Consider the following two linear interpolations:
will argue that if (L [ox*IIo [O, I)) is small, then it is
X = = -Z
(14a)
+
(14b)
Denote ps and as the distribution of X, and Y, respectively. Then, it is easy to verify that they satisfy the following continuity equations:
Pt (v pt)
(15a)
④ V (u 9:) =0
(15b)
where uy (x) = E[X Z|X; xi and By (z) E Z|Y, = z²). Then, we have the following theorem which establishes the relation between KL-divergence of and qı in terms of the velocities "i" and vy The proof for the theorem is provided in Section D.9
Theorem D.2. Let Pt and be the distribution of Xt and Y, defined in equation 14 KL-divergence between pı and qı satisfy the following relation
Then, the
t pr(x) (1( -. (x)]] dadt.
(L(pllla)) KL (L
(16)
Now, consider the distribution of the inverse noise pl obtained by iterating equation 5 and substitute it with px in the theorem above. Suppose that the flow model is trained such that (L (P1) S
Then, by Lemma 4.2 it follows that ( (pp) I e. Combining this observation with equation 16 it is easy simply corresponds to learning a flow model from standard Gaussian to itself, we can explicitly compute oy as follows:
that the velocities vzx (xx) and vi (x) should be close to each other. Additionally, since
x
x
20 (x)
t
t (1-t)a t
(2t
+
Thus, of (x)) is linear funetion of and rational function of Because (L (po) S c, Theorem
D.2 suggests that learning ux" from data should be relatively easier as it is close to vi
D.9 PROOF OF THEOREM D.2
Then, the time derivative of the divdeieeee between Pt and gi given by dKL(p:llg.)
(x(xr)
J&c dt
pc(z)) log
dx
dt
(x)
| pe(x)) log
(t(2z)
+ (x( d dt
log(pe(x)) px(x) (x)x)) d dt log(u((1)) dr
9((x))
Pr(zz)
+ pr px(x) (x(x))
p((x)) log
log
dr
(x(x)
(x()))
We will consider each term separately as Tı, T2 and T3. For Ti using the continuity equation, we have
pt(z)) dx (((x))
T
JRd (e(x:) log
(r(z) dx
V. (vz (x)) log
Ja
(x()))
pe(o)) ( (1)( Vlog
(x)
dr
(Integration by parts)
(xr)
24

---

Published as a conference paper at ICLR 2026
Note that Jaa (x()d2 1 for all t [0, 1]. Thus for T2, we obtain
T2 | P.(r)dr==
dt [ p,())dr - 0.
d
For the final term Ts. we again use the continuity equation to get
(x(x)
T3 =
P((
de
(x)))
Rd
(v)
9t)
(x())
dr
Rd
(x))
| 91(x)
pr(x))
V
,2/ (x) dx
(Integration by parts)
JRa
(x(x))
(e(
(1()
V Vogg
.1z (xr) dr.
(x(x))
Ry
Combining all the terms above, we get
dKL(Pellge) dt
Pe(x)) (x)2)
pr(x)
log
15 (x) - 2i (x)
dr.
(17)
JR"
To obtain an expression for score function in terms of the velocity vector, we use Tweedic's formula Eiron 12011) which leads to
E z|X, 2]
E[X X, = x]
E[X Xx = x) - T
(x (1 1)2 (ollgg(())) -
(Tweedie's Formula)
t(1 - t)
I +
V (x).
(18)
[
Similarly, we obtain
1 log (z).
x = E|Y - 2| = x] 7
(19)
Plugging in the expressions for the score functions into equation dKL(p:llg:) dt
we obtain
[ pt(x))
L 1- (1z - of (x)) ve (z) - ux(()) dx
(x()))
1 / Pc(
11
(x) (x))(())))))){)))
oy (z) dx
Rdd
(L(pll)9)) KL(Pollgo) f"'
Jed p((x) 10 (z) of (z) drdt.
Recall that Po = qu 91 = I) and p1 = xxx Thus, we get the desired claim
KL (LLx,, ()) = //''
1 { (x( 1|0 (x) 2. (x)) drdi.
t
25

---

Published as a conference paper at ICLR 2026
E TEXT-TO-IMAGE GENERATION
E. ABLATIONS
E. DIFFERENT CHOICES OF AND M
We report results for various choices of K and M for Top - of M sampling for GenAl-Bench in Tables 6 and Note that these models are also trained GenAl-Bench. We also report the (mean)scor seeaeessooeeeeee separately for the "Basic" and "Advanced" split in the prompt set. Results for Top 10 of 100 sampling for T21-CompBenchh++ is given in Table 9 Models for Table | were trained the train split of T2J-CompBench++. All results are consistent with the developed theory: both GRAFT and P-GRAFT outperform base SDv2 and P-GRAFT, for an appropriate choice of N always outperform GRAFT.
Table 6: VQAScore on GenAl-Bench for K = 1 and M = 4
Model
Basic
Advanced
Mean
SD v2
74.83
59.19
66.32
GRAFT
77.33
62.76
69.
P-GRAFT (0.8N)
76.30
62.18
68.62
P-GRAFT (0.5N)
78.57
63.38
70.32
Table 7: VQAScore GenAl-Bench for K
and M 100
Model
Basic
Advanced
Mean
SD
74.83
59.19
66.32
GRAFT
79.61
64.26
71.2
P-GRAFT (0.75.V)
76.02
62.91
68.89
P-GRAFT (0.5N)
78.68
64.5
70.97
P-GRAFT (0.25.V)
80.05
64.85
71.79
Table 8: VQAScore on GenAl-Bench for K
and M 100
Model
Basic
Advanced
Mean
SD v2
74.83
59.19
66.32
GRAFT
78.01
63.31
70.02
P-GRAFT (0.75.N)
77.36
63.33
69.73
P-GRAFT (0.5N)
78.18
64.28
70.62
GRGAFT (0.25.V)
78.77
65.29
71.44
Table 9: VQAScore on T2I- D1--CompBencc++++ (Val) for K = 10 and M 100
Model
Mean
SD
69.76
GRAFT
74.66
P-GRAFT (0.25N)
75.16
E.1.2 CONDITIONAL VARIANCE OF REWARD FOR TEXT-TO-IMAGE GENERATION
While experimental results in Table 2 already demonstrate the bias-variance tradeoff, provide further evidence of Lemma 3.3 in the context of text-to-image generation. We evaluate conditional variance of VQAReward scores of the base SDv2 model in GenAl-Bende We follow the methodology as described in Section 3.1 except that we generate 4 images per prompt for a total of 1600 prompts. The results given in Table 10 It can be seen that even N, 0.75N, the expected conditional
26

---

Published as a conference paper at ICLR 2026
variance of the reward is significantly smaller than at ły. This explains why even Ni 0.75N gives a significant gain over the base model as seen in T
Table 10: Expected conditional variance for T21 generation
N1
E [Var(r(Xo)) [xx))
N
0.0193
3N/4 N/2
0.0080
0.0039
N/4
0.0019
E.1.3 EFFECT LORA RANK
We increase the LoRA rank used for fine-tuning and check the impact on the performance. Table shows that increasing LoRA rank does not seem to affect performance, indicating that the default LoRA rank is sufficient. Ablations are done on GenAl-Benci with M = 100,
Table 11: Effect of LoRa Rank
Model
Rank
Mean Reward
70.97
70.87
P-GRAFT (0.5N)
8
70.57
10
70.84
4
71.79
71.84
P-GRAFT (0.25N)
71.49
10
71.63
E.1.4 REVERSE STITCHING
In P-GRAFT, we always use the fine-tuned model for the first (N N7) steps and then switch to the reference model. We experiment with reverse stitching strategy, where model for the earlier denoising steps and fine-tuned model for the later denoising steps. For switching timestep N,, denote this strategy as RP-GRAFT ()+ -ie. RP-GRAFT (0.75.V) indicates that the base model will be used from ty to f5.75N, after which the fine-tuned model will be used. From
the reference
Table 12, we observe that this strategy significantly worse when compared P-GRAFT this provides further evidence of the bias-variance tradeoff. Ablations are done with M 100, K = 1.
Table 12: Ablations reverse stitching
Model
Basic
Advanced
Mean
SDv2
74.83
59.19
66.32
GRAFT
79.61
64.26
71.20
RP-GRAFT (0. 75.)
79.23
62.63
70.20
RP-GRAFT (0.5N)
76.60
60.87
68.05
RP-GRAFT (0.25N)
75.74
59.76
67.05
E.2 IMPLEMENTATION DETAILS
Since require samples only from the pre-trained model, sampling and training can be done separately. Therefore, we first perform rejection sampling according to Top - of M for the chosen values of K and M. The selected samples are then used as the dataset for training. If not mentioned explicitly, hyperparameters can be assumed to be the default values for SD 2.0 in the Diffusers library (von Platen et al. 2022).
27

---

Published as a conference paper at ICLR 2026
Training GenAI-Bench:
The hyperparameters for sampling and training are given Table 13 and Table 14 respectively. Note that one training epoch is defined as complete pass the training dataset. The size of the training dataset depends on the chosen K and M. For instance, K 10 and M = 100 results in 10 images per-prone prompt, for total of 16000 images. One training epoch corresponds to single pass these 16000 images, which with a batch size of 8 corresponds to 2000 iterations per epoch.
Table 13: Sampling hyperparameters for GenAl-Bench
Sampling Steps
50 EuterrissreteSchcdule
Scheduler Guidance Scale
7.5
Table 14: Training hyperparameters for GenAl-Bent -ennl-Bench
Training Epochs
Image Resolution
768 x 768
Batch Size
Learning Rate
10
LR Schedule
Constant
LoRA Fine-Tuning
True
Training on 121-CompBench++:
The hyperparameters for sampling and training are given in Table and Table [6 respectively. We use different sampling schedulers for the two datasets to ensure that our results hold irrespective of the choice of the scheduler.
Table 15: Sampling hyperparameters for 121-chomBeecch++
Sampling Steps
50
Scheduler
DDIMScheduler
7 (DDIMScheduler specific hyperparameter) ) Guidance Scale
1.0
7.5
P-GRAFT Training:
Training and sampling using GRAFT is straightforward since standard training and inference scripts be used out-of-the box: the only additional step need is rejection sampling on the generated samples before training. For P-GRAFT, the following changes are to be made:
While sampling the training data, the intermediate latents should also be saved along with the final denoised image/latent. Rejection sampling is to be done on these intermediate latents, but using the rewards corresponding the final denoised images.
While training, note that training has to be done by noising the saved intermediate latents. This needs re-calibration of the noise schedule, since by default, training assumes that we start from completely denoised samples. The easiest way to re-calibrate the noise schedule is by getting set of values for the betas parameter, new betas as follows (where N, denotes the intermediate step P-GRAFT):
new_betas(0 (0, N7] +- 0 new_betas N,, N] <- betas[N,,
After re-calibrating the noise, we use new bew_aetas to get the corresponding new aa__aphas and new_a a ppass_cmpprod It is also necessary to note that while training, the denoiser has been trained to predict Xo given any noised state X, and not the saved intermediate
28

---

Published as a conference paper at ICLR 2026
Table 16: Training hyperparameters for 121-CompBench++
Training Epochs
10
Image Resolution
768 x 768
Batch Size
Learning Rate
10
LR Schedule
Constant
LoRA Fine-Tuning
True
latent Xtw, Let the corresponding saved completely denoised latent be Xo To ensure that the training is consistent, we train using the following strategy:
Sample 1(0,1)
Get X₁ + (V /enwww (Vnew_alphar wwwp_alphas_cmmg hew_aalphasccuppodl]] Xev, + (VVI new wwww_alphas_cmmm Supphhas_cmmrooff//
Get X1 - caappnss_cmmrod(((
VI alphas_ -lphasscummrrod comprod
Compute Loss using X, and e'
E.3 POLICY GRADIENT ALGORITIMS
DDPO(Black al. 2023) an on-policy policy gradient method for diffusion models that optimizes clipped importance-weighted objective over the denoising trajectory. The original paper reports results on experiments using at most 400 prompts. Both prompt sets we consider are significantly larger (1600 prompts for GenAl-Bench and 5600 (train) prompts for T21-CompBench++) This difference is crucial, since it has been shown in Deng et al.] (2024) that scaling DDPO to large prompt sets result in unstable training and subpar performance. We also observe this phenomenon, as evidenced by the results in T As menioned in the main text, we also augment DDPO with additional elements in attempt to improve performance. In particular, we study the following variants:
1. DDPO: Clipped ceperrannc-weighteee policy gradient.
2. DDPO+KL DDPO augmented with stepwise KL regularizer to the (frozen) reference model.
3. DDPO+KL+EMA DDPO with KL regularization well prompt-wise exponential- moving-average baseline for advantage estimation.
Baseline Implementation: We use the official PyTorrh implementation of DDPQ we further adapt the codebase to implement other variants. Pine-tuning is always done on SDv2 using LoRA on the UNet only with a LoRA rank of 16. For the results reported in Table 2 we retain the hyperparameters used in Black et al. (2023). In particular,
use PPO clip range of 10-4, gradient clipping norm of 1.0, Adam optimizer with 81 = 0.9, 32 0.999 and weight decay of 10-4, Following the original paper, we train with relatively high learning rate of 3 10 since LoRA fine-tuning is used. We sample 32 prompts per epoch and train with a batch size of 8, leading to 4 training iterations per epoch. However, note that each training iteration requires gradients across the whole denoising trajectory- - this means that within each training iteration, 50 gradient calls are needed, corresponding to 50 sampling steps. For GenAl-Bench, training is done for 500 such epochs, whereas for T21-Compbencht++ training is done for 800 epochs. With this setup, in Tables 17 and [8] we compare the sampling/compute requirements for DDPO and GRAFT. In particular, note that GRAFT/P GRAPT already outperforms DDPO with K = 1, = 4 despite DDPO being trained on 10x more samples and 50x more gradient calls.
Additional configurations with base hyperparameters: With the base hyperparameters described above, we also try augmenting DDPO with KI. and EMA as described above. The training curves are given in Figure
29

---

Published as a conference paper at ICLR 2026
eni
m
(a) DDPO+KL
(b) DDPO+KL+EMA
Figure 3: Training curves for the three baselines on GenAl Bench (1,600 prompts) with low value of clipping.
menert Linee
(a) DDPO
(b) DDPO+KL
(c) DDPO+KL=EMA
Figure 4: Training curves for the three policy-gradien baselines on GenAl Bench (1,600 prompts)
with high value of clipping.
We also try additional settings for hyperparameters apart from the so far. Sampling uses DDIM with T E [40, 50] steps and classifier-free guidance g = 5. Optimiza- tion uses Adam W with learning rates {2x10-5, 10-5j, batch sizes 8/8 (samplin//rainiing.. style clipping E (0.1,0.2) Following DDPO, we replay the scheduler to compute per-step log bog-prbbabilities on the same trajectories: & log ( 2e, c) and gold log peo (It-1 2t, c).
DDPO:
have reported
We use the clipped objective:
LDDPO = Emin(rr clip(re, 1-e, L+e) A)]
r2 = exp 2old),
(20)
with a centered batchwise advantage A. Specifically, we experiment with higher clipping range, E (0.1, 0.2) and use whitened batchwise advantage. l,, gold are obtained by replaying the DDIM scheduler on the same trajectory.
Result, On 1600 prompts, the learning curve exhibits a short initial rise followed by a sharp collapse after 150 steps (Fig. 4. The setting of 1600 heterogeneous prompts induces high variance and many ratios T saturate at the clipping boundary, producing low-magnitude effective gradients and the observed drop in reward.
DDPO+KL:
We augment equation 20 with per-step quadratic penalty to the frozen reference:
LDDPO = LDPPO + 8 T
B (0.02, 0.005).
(21)
1=1
Result. The KL term prevents divergence of the policy and eliminates the reward collapse after the first few steps. Even with this, average reward improvements remain limited. Larger contracts the policy towards the reference, whereas smaller 3 provides insufficient variance control, yielding small net gains.
DDPO+KI+EME (prompt-wise baseline):
To mitigate cross-prompt bias, maintain for each
prompt z, an EMA of reward and variance,
b(z) + (1=a)b(=) + ar,
(()) + (1-a)u(z) + a(r - b(z))²
and employ a whitened advantage inside equation 20 Ā b=z) + n, 17 ~ ~(0,00)
Vo(z)+E
Result. Training is the most stable among the three variants and exhibits smooth reward trajectories without collapse, yet the absolute improvement in mean reward is modest relative to the base policy.
30

---

Published as a conference paper at ICLR 2026
PRDP: We also tried implementing PRDP (Deng et al. 2024 using the PRDP loss function provided in the appendix of the paper since no official code was provided. However, we did not see any signifi- cant improvement compared to the baseline despite following the algorithm and hyperparameters closely. One potential reason for this could be that we use LoRA fine-tuning whereas the original paper uses full fine-tuning. Further, we rely gradient checkpointing for the implementation well since the backpropagation is through the entire sampling trajectory.
Table 17: Comparison of Sampling Cost and Training Cost for GenAl-Bench
Algorithm
Samples generated
Samples Trained on
Gradient Calls
GRAFT(K = 10, M 100) GRAFT (K 1, M.M- 4)
160k
16k
20k
6.4k
1.6k
2k
DDPO
16k
16k
100k
Table 18: Comparison of Sampling Cost and Training Cost for 121-CompBench++
Algorithm
Samples generated
Samples Trained on Gradient Calls
GRAFT (K = M, = 4)
22. 4k
5.6k 25.6k
7k 160k
DDPO
25.6k
E.4 COMPUTE FLOPS ANALYSIS OF P-GRAFT
We compare the compute cost of P-GRAFT and DDPO in terms of total UNet FLOPs. Let F,, denote the cost of one UNet forward pass at 64x64 latent resolution. Following Kaaolan et al. (222),, we approximate backward training step times of forward step. So if F, is the forward step compute, a forward + backward step will incur 3F,. We assume batch size of 1 for both algorithms for standardization.
For P prompts, M samples per prompt, top-K K retained, diffusion steps, Esit epochs. For the implementation
the standard stable diffusion training script that only samples a single timestep I E [o, T] during training:
FP-GRAFT = PMT Fo
+ Estt PK3F.
sampling
training
FDDPO = Eddpo Ngen T. 3)
trajectories
GenAI-Bench configuration. We use P=1600. M=100, K=10, T-40, Esn for P-GRAFT; Trajectories generated per epoch Ngen=128 and Number of Epochs Eddpo=50 for DDPO
Table 19: FLOPS in units of forward pass F, for GenAI-Bench.
Algorithm
Sampling
Training
Total
P-GRAFT (K=10, ==100)
6.40M
0.48M
6.88M
P-GRAFT (K=1, M=4)
0.256M
0.048M
0.304M
DDPO (E=50, N 128
0.256M
0.768M
1.024M
Discussion, P-GRAFT's total compute dominated by sample generation, while backpropagation is confined to fine-tuning the selected top- K samples. In contrast, DDPO backpropagates through all T' denoising steps online for every sample, creating a sequential bottleneck. Consequently, despite DDPO's nominal FLOPs appearing comparable or lower in our regime, its wall-clock time substantially longer due to stepwise backward passes that are less parallelizable. Moreover, as shown in Table 2. P-GRAFT achieves higher rewards under the reported budgets; and in the compute- matched case (K=1; Table (). P-GRAFT still outperforms DDPO, indicating that gains come from improved optimization and not just additional training compute.
31

---

Published as a conference paper at ICLR 2026
E.5
QUALITATIVE EXAMPLES
E.5.1
GENAI-BENCH
Prompt: Three flowers in the meadow, with only the red rose blooming; the others are not open.
SDv2
DDPO
GRAFT
P-GRAFT
Prompt: In the yoga room, all the mats are red.
SDv2
DDPO
GRAFT
P-GRAFT
Prompt: Three policemen working together to direct traffic at a busy intersection.
SDv2
DDPO
GRAFT
P-GRAFT
Prompt: There is apple and two bananas the table, neither of which is bigger than the apple.
SDv2
DDPO
GRAFT
P-GRAFT
Figure 5: Qualitative examples on GENAI-BENCH. All results are reported for the same seed across different algorithms.
32

---

Published as a conference paper at ICLR 2026
E.5.2
T21-COMPBENCH++
Prompt: a cubic dice and cylindrical pencil holder
SDv2
DDPO
GRAFT
P-GRAFT
Prompt: a bee the bottom of a airplane
SDv2
DDPO
GRAFT
P-GRAFT
Prompt: a green bench and a blue bowl
SDv2
DDPO
GRAFT
P-GRAFT
Prompt: a green acorn and a brown leaf
SDv2
DDPO
GRAFT
P-GRAFT
Figure 6: Qualitative examples on T2I-COMPBENCH, All results reported for the same seed across different algorithms.
33

---

Published as a conference paper at ICLR 2026
F LAYOUT AND MOLECULE GENERATION
F.1 INTERLEAVED GIBBS DIFFUSION (IGD)
Fine-tuning for both layout generation and molecule generation are done using models pre-trained using the Interleaved Gibbs Diffusion (IGD) (Anil et al.) 2023) framework. IGD performs well for discrete-continuout generation tasks with strong constraints between variables and hence is particularly useful for tasks like layout generation and molecule generation. Further, IGD offers generalizable framework which can be used for both tasks while other discrete-continuvwe diffusion frameworks exist, they are specialized particular task, often using domain specific adaptations.
On a high-level, IGD interleaves diffusion for discrete and continuous elements using Gibbs-style noising and denoising. Essentially, discrete elements are noised using flipping and trained using a binary classification loss. Continuous elements use typical DDPM-style noising and training. While the exact forward and reverse processes are different from DDPM-style processes which we have considered in the main text, the key results follow empirically and theoretically.
F.2 LAYOUT GENERATION
Problem Formulation:
A layout is defined set of N elements {er) Each element e, is represented by a discrete category 4 E N and a continuous bounding box vector pa E R.. Following (Anil et al. 2025),
the parameterization P. [ii Vis 4, wa] where (3a ys) represents the upper-left corner of the bounding box, and (4, w;) its length and width, respectively. Unconditional generation represents generation with no explicit conditioning for the elements, whereas Class- Conditional generation indicates generations conditioned on element categories.
Implementation Details: For pre-training, we follow the exact strategy used in (Anil et al. (025). Fine-tuning is also done with the same hyperparameters used for pre-training. Since the data and model sizes are significantly smaller compared to images, each round of rejection sampling is done on 32768 samples, of which the top 50% samples are selected. For each sampling round, 10000 training iterations are performed with a training batch size of 4096. The results reported in Table such sampling rounds. FID computation done by comparing against the test split PubLayNer.
3 are for 20
F.3 MOLECULE GENERATION
The task of molecule generation involves synthesizing molecules given a dataset of molecules. A molecule consists of n atoms denoted by {=i, P. 3 where 2; N is the atom's atomic number and p, E R the position. A diffusion model is trained to generate such molecules. In this work, we take such a pre-trained model, and try to increase the fraction of stable molecules, as deemed by RDKit.
Problem Formulation:
For pre-training, we follow the exact strategy used in (Anil et al. (025). Fine-tuning is also done with the same hyperparameters used for pre-training. Since the data and model sizes are significantly smaller compared to images, each round of rejection sampling is done on 32768 samples. We select all stable molecules, but with the de-duplication strategy described in Section 2.3 we find that this is crucial to maintain diversity of generated molecules. For each sampling round, 10000 training iterations are performed with a training batch size of 4096. The 1x in Table 4 corresponds to 10 such sampling rounds 9x therefore corresponds to 90 sampling rounds.
Implementation Details:
Uniqueness of Generated Molecules: To demonstrate that the fine-tuned models still generate diverse molecules, and do not collapse to generating a few stable molecules, we report the uniqueness metric computed across the generated molecules below. From Table 20 it is clear that the fine-tuned models still generate diverse samples since the uniqueness of the generated molecules remain close to the pre-trained model. Uniqueness is determined by RDKit.
Effect of de-duplication We also try out ablation where duplication - i.c., we train on all stable molecules irrespective of whether they are unique or not. The results are shown in Figure 7 without de-duplication, it can be seen that though stability is recovered,
GRAFT, but without the de-
34

---

Published as a conference paper at ICLR 2026
Table 20: Uniqueness of generated molecules
Model
Mol: Stability
Uniqueness
Baseline
90.5010.15
95.6010.11
GR
90.7620.2 20
96.0420.46
P-GRAFTO..5)
90.46+0.27
96.70+0.28
P-GRAFT(0.258) 25.N)
92 12.61.4 10.13
95.3210.07
380
2.
0.75
07
0.0
(a) Molecule Stability
(b) Molecule Uniqueness
Figure 7: Molecule Stability and Uniqueness without De-duplication
uniqueness is lost, indicating that the model produces only a small subset of molecules was initially able to produce.
Tone-Tuning without Predictor-Coorrector:: IGD makes use of a version predictor-corrector
method (Lezama et al. 2022 Zhao et al. 2024 Campbell al., 2022 Gat al. 2024) termed ReDeNoise at inference-time to further improve generations. The results reported so far make use of this predictor-corrector. While ReDeNoise improves performance significantly, comes the cost of higher inference-time compute. We report results of the baseline and fine-tuned version without ReDeNoise in Table 21 Both GRAFT and P-GRAFT still show improvement over the baseline, even without ReDeNoise.
Table 21: Results for Molecule Generation without ReDeNoise
Model
Mol: Stability
Sampling Steps
Baseline
84,00
GRAFT
87.13
9x
P-GRAFT (0.5.))
84,57
1x
P-GRAFT (0.25N)
88.36
X
35

---

Published as a conference paper at ICLR 2026
G INVERSE NOISE CORRECTION
G.1 IMPLEMENTATION DETAILS
The pre-trained models, corresponding to TRAIN FLOW function in Algorithm are trained using the NSCNpp architecture and hyperparameters from the official codebase of [fong et al. (2020a) with minor changes which we describe below. The noise corrector model is also trained with the same architecture except that the number of channels are reduced from the original 128 to 64 channels this leads to reduction in parameter count by R 4x. For pre-trained model, we train with num scales 2000, positional embeddings and batch size of 128. For the noise corrector model, we use the same hyperparameters except for num scales = 1000. FID with 50000 samples is used to measure the performance, is standard in the literature. Note that separate noise corrector model is trained for each choice of v) in Algorithm 3 i.e., for the results reported in Table separate noise corrector models trained for pre-trained steps of 100 and 200.
CelebA-HO For the baseline pre- trained flow model, we use the checkpoint after 330k iterations, since this gave the lowest FID. For noise corrector model training. we this checkpoint generate the inverse noise dataset and train on for 150k iterations.
the checkpoint after 350k iterations, since this gave the lowest FID. For noise corrector model training, we use this checkpoint to generate the inverse noise dataset and train it for 55k iterations. Note that Backward Euler (Algorithm ( suffered from numerical instability, which we hypothesize is due to plain backgrounds, when done on LSUN-Church. To alleviate this issue, we perturb the images with small Gaussian noise N(00, II).
LSUN-Church: For the baseline pre-trained flow model,
with d = 10
G.2 FLOPS COMPARISON
We present a comparison of the exact FLOPs used for inference:
14
Inverse Corrected
Pre-trained LDM-4
12
8
Inverse Corrected
Pre-trained
LDM-8
6
4
6.5
41
2
9
FLOPs (x (x
FLOPs (x (x10²2)
(a) Celeb HQ
(b) LSUN Church
Figure 8: FLOPs FID: The inverse corrected model achieves better FID despite incurring lower FLOPs. Corresponding LDM models have been added for both datasets for reference.
36


Reviews:
[{"review_id": "cpI6HePqZm", "reviewer": "Reviewer_iasd", "summary": "The paper recognizes the importance of reward-guided fine-tuning of diffusion models and the complexity of leveraging proximal PG methods due to intractable likelihoods. The authors introduce a unified rejection sampling framework, show that fine-tuning according to it leads to the same solution of proximal PG schemes, and introduce a method for controlling also the tilt of intermediate distributions of the diffusion process. They present a mathematical analysis for bias-variance trade-off, and propose a noise-correction scheme for reward-independent improvement of flow models. Ultimately, the present an experimental evaluation of the proposed contributions.", "questions": "- Did I misinterpret or misunderstand any of my points above within (a) or (b)?  \n- What is the core algorithmic intuitive idea for the method introduced in Sec. 5?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "- The paper correctly identifies a limitation of RL-based schemes for fine-tuning of diffusion models (i.e., intractable log likelihoods), a timely task with high-relevance.\n\n- The idea proposed within Sec. 5 seems interesting, but I could not fully understand its logic on an intuitive level.\n\n- The experimental sections includes diverse dataset spanning images, layout gen., and molecular design.", "weaknesses": "The paper presents an excessive amount of ideas, each not sufficiently justified, motivated, or explained. As the paper structure is very complex due to the multitude contributions, I will discuss weaknesses of each contribution a-d as listed in Sec. 1.\n\n**a)** \n1. The paper often seems to confuse a *problem*, in particular solving an entropy-regularized MDP, with an algorithm, i.e. PPO. I understand that the presented algorithm has solution corresponding to the optimal solution of entropy-regularized MDPs, which can be tackled by PPO, but what does 'GRAFT enables PPO' or 'implicitly perform PPO' even mean? This confusion seems repeated several times within the work. Concretely, I believe these statements are effectively wrong, as the propose method, while inducing the same solution, does *not* enable/perform PPO.\n2. It is well-known that a general class of inference-time schemes for diffusion induces this solution class (see [1], Eq. 1). One such case is [2] (see Sec. 3.2). So it seems there is already a variety of inference-time schemes solving this problem via diverse techniques (c.f., [1]). Moreover, there exist fine-tuning control-theoretic schemes that can solve this entropy-reg. problem as well (e.g., [3]) without value bias problem, as also mentioned by the authors. Ultimately, leveraging inference-time schemes solving this problem to then fine-tune a model, which seems to me the core algorithmic idea here, is already presented in [1, Sec. 9] arguably in a more performative fashion than the one presented within this paper (i.e. , via policy distillation rather than plain training). As a consequence, it seems to me that there isn't significant novelty within the presented contribution.\n\n**b)**\n1. The KL is typically enforced on the data level distribution as a way to preserve the high-probability set learned by the pre-trained model. The justifications presented within the paper (both in theory 4.1 and experiments) do not seem convincing to me regarding why we should instead enforce the KL-reg. at another time-step. In particular, it seems to me that the theoretical investigation in Sec. 4.1 does not really provide a concrete answer to this question. In a sense, it shows that this problem might be easier, but this does not imply (practical or theoretical) relevance. Since the authors here are presenting a novel problem setting, it would be essential to motivate it clearly.\n\n**c)**\n- (writing) this section (i.e. Sec. 5) is not clearly written to the point that I could not fully grasp the presented idea. The problem tackled seems not particularly related with the problems treated in the rest of the paper, and is not sufficiently formalized to properly understand the gains of the proposed methodology. I would strongly suggest to also introduce algorithmic aspects with an intuitive presentation of their workings before/after presenting their implementation.\n\n\n**Overall (writing/structure)**\nThe paper is poorly structured and lacks a solid narrative. Concretely, it presents multiple ideas without sufficient clarification of their motivation, and/or workings. The text often lacks conceptual explanations of new concepts/mechanisms and their implications. \n\n\n**References**:\n\n[1] Inference-Time Alignment in Diffusion Models with Reward-Guided Generation: Tutorial and Review, 2025\n\n[2] Derivative-Free Guidance in Continuous and Discrete Diffusion Models with Soft Value-Based Decoding, 2024\n\n[3] Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control, 2024", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QfBmSvP9E0", "reviewer": "Reviewer_sMVE", "summary": "The authors investigate:\n(1) Rejection sampling and it’s uses in sampling and fine-tuning diffusion models. The authors derive the sampling distribution for top-k sampling, with arbitrary reward functions and de-duplication criterion, for diffusion models. The authors then propose fine-tuning on the trajectories yielded by rejection sampling.\n\n(2) An inverse noise correction algorithm that trains a flow model from N(0, 1) to the t=1 distribution of the flow model, and then sample from the flow model.", "questions": "Will the authors clarify what their contributions are exactly, as\n\n1. The sampling distribution of top-k sampling is known:\n    1. As the authors acknowledge, the derivations are known from Amini et al 2024. The derivations there are not specific to any particular kind of method of generation, it applies to any generative model, including both AR, diffusion models. See theorem 2 in Amini et al 2024. \n2. Top-k sampling for fine-tuning has been proposed and shown in RAFT, see Dong et al\n\nQuestions regarding fine-tuning with Top-k sampling:\n\n1. What distribution does alg 2 (P-GRAFT inference) sample?  Alg 2 stiches together two different diffusion models, with the first one being fine-tuned and sampling from the reward tilted intermediate distribution, followed by the base model. \n\nQuestions regarding the motivation for inverse-noise correction:\n\n1. The motivation for the method is not clear from the text, either in the main paper or appendix. \n\nExperimental Clarifications: \n\n1. Can the authors clarify what objective was used to fine-tune the diffusion model in alg 1 (P-GRAFT train)\n2. The GenEval scores reported in the paper for SDv2 and SDXL-base are higher than those reported in the GenEval paper. Can the authors clarify what prompts were used for producing the numbers reported in table 2. \n\nMinor clarification regarding using the term PPO distribution: PPO is a method for learning the reward tilted distribution, p(x) exp(r(x)). This distribution can be learned using methods other than PPO.", "limitations": "", "rating": 2, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "See questions", "weaknesses": "See questions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Xp0ujdSjNo", "reviewer": "Reviewer_xSnM", "summary": "This paper introduces fine-tuning methods for diffusion models through intermediate distribution shaping. The authors first unify rejection sampling-based fine-tuning methods as GRAFT (Generalized Rejection sAmpling Fine-Tuning) and prove it implicitly performs PPO with reshaped rewards, enabling marginal KL regularization for diffusion models despite intractable likelihoods. They then propose P-GRAFT (Partial-GRAFT), which fine-tunes only until intermediate denoising timesteps by assigning final generation rewards to partial noisy states. This is justified through a bias-variance tradeoff: while variance of rewards conditioned on intermediate states increases with noise level (Lemma 4.3), the score function becomes easier to learn at higher noise levels (Theorem 4.4). Additionally, they introduce Inverse Noise Correction for flow models, which trains a model to correct the distribution shift in initial noise without explicit rewards. Empirically, P-GRAFT applied to Stable Diffusion v2 achieves 8.81% relative improvement in VQAScore over the base model on text-to-image benchmarks and outperforms policy gradient methods. For unconditional generation, Inverse Noise Correction improves FID at lower FLOPs/image.", "questions": "1. How does the optimal $N_I$ depend on task characteristics (reward smoothness, data modality, model capacity)? Can you provide theoretical or empirical guidance for selecting $N_I$?\n\n2. Theorem 4.4 shows score functions at later times are exponentially closer to Gaussian. Have you verified this empirically by measuring $H_t^T$ at different timesteps?\n\n3. The backward Euler method (Algorithm 6) requires solving a fixed-point equation. How many iterations $N_b$ are typically needed in practice, and what is the computational cost relative to forward sampling?\n\n4. Can P-GRAFT be combined with classifier-free guidance training? How does the guidance scale interact with the intermediate timestep choice?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The unification of rejection sampling methods (RAFT, RSO, BoN) under the GRAFT framework with explicit connection to PPO is novel and theoretically grounded (Lemma 3.2)\n* P-GRAFT represents a genuinely new approach to diffusion model fine-tuning by operating at intermediate noise levels rather than final outputs\n* The bias-variance analysis providing theoretical justification for intermediate distribution shaping is insightful and well-executed\n* Mathematical formulations are generally clear with good notation", "weaknesses": "1. **Limited theoretical analysis for P-GRAFT optimality**: There is no formal analysis of the optimal choice of $N_I$ or convergence guarantees. The paper relies heavily on empirical validation across different $N_I$ values, but theoretical guidance for selecting $N_I$ would strengthen the contribution.\n\n2. **Assumption dependencies**: Lemma 5.1 requires $\\eta L < 1$ for backward Euler, but the paper doesn't discuss how this constraint affects practical step size choices or what happens when the Lipschitz constant $L$ is large. The impact on computational cost is not analyzed.\n\n3. **Gap between theory and practice**: The theoretical analysis assumes continuous-time SDEs, but experiments use discrete-time DDPM/DDIM schedulers. The discretization error and its interaction with P-GRAFT is not discussed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)