Efficient Diversity-Preserving Diffusion Alignment via Gradient-Informed GFlowNets
Authors: Zhen Liu
Tim Z. Xiao
Weiyang Liu
Yoshua Bengio
Dinghuai Zhang
Venue: ICLR 2025 Poster (2025)
Keywords: Reward Finetuning
Diffusion Models
Alignment
GFlowNet
Generative Models
Primary area: generative models
TL;DR: Efficient, diversity-preserving and prior-preserving reward finetuning and alignment of large-scale continuous diffusion models, with a novel gradient-informed RL objective derived from the GFlowNet framework

Abstract: While one commonly trains large diffusion models by collecting datasets on target downstream tasks, it is often desired to align and finetune pretrained diffusion models with some reward functions that are either designed by experts or learned from small-scale datasets. Existing post-training methods for reward finetuning of diffusion models typically suffer from lack of diversity in generated samples, lack of prior preservation, and/or slow convergence in finetuning. In response to this challenge, we take inspiration from recent successes in generative flow networks (GFlowNets) and propose a reinforcement learning method for diffusion model finetuning, dubbed Nabla-GFlowNet (abbreviated as $\nabla$-GFlowNet), that leverages the rich signal in reward gradients for probabilistic diffusion finetuning. We show that our proposed method achieves fast yet diversity- and prior-preserving finetuning of Stable Diffusion, a large-scale text-conditioned image diffusion model, on different realistic reward functions.

Published as a conference paper at ICLR 2025
EFFICIENT DiVERSITY-PRESERVING DIFFUSION ALIGNMENT VIA GRADIENT-INFORIE GFLOWNETS
Zhen Liu 1:33t Tim Z. Xiao Weiyang Liu Yoshua Bengio'
Dinghuai Zhang!
Mila, Université de Montréal Max Planck Institute for Intelligent Systems Tübingen
The Chinese University of Hong Kong (Shenzhen)
UUniversity of Tübingen
5 University of Cambridge Microsoft Research CCorresponding author
"Equal contribution
Project page: nabla-gfn. github.
ABSTRACT
While one commonly trains large diffusion models by collecting datasets on tar- get downstream tasks, it is often desired align and finetune pretrained diffusion models with some reward functions that are either designed by experts or learned from small-scale datasets. Existing post-training methods for reward finetuning of diffusion models typically suffer from lack of diversity in generated samples, lack of prior preservation, and/or slow convergence in finetuning. In response to this challenge, we take inspiration from recent successes in generative flow networks (GtlowNets) and propose a reinforcement learning method for diffusion model finetuning, dubbed Nabla-GFlowNee (abbreviated as V-GFlowNet),, that leverages the rich signal in reward gradients for probabilistic diffusion finetuning. We show that our proposed method achieves fast yet diversity- and ing Stable Diffusion, a large-scale text-conditioned image diffusion model, on different realistic reward functions.
finetun-
INTRODUCTION
Diffusion models [17, 61, 51] are a powerful class of generative models that model highly complex data distributions with sequential denoising steps. They prove capable of modeling distributions of images [51, 9, 49], videos [19, 18, 4], 3D objects [81, 47, 34, 35, 13], molecules [72, 20, 32], lan- guages [37, 53, 36] and many others. State-of-the-arts diffusion models for downstream applications are typically large in network size and demand a significant amount of data to train.
is often desirable that one finetunes pretrained diffusion models with some reward either learned from human preferences as in reinforcement learning from human feedback (RLHF) [5, 43] or de- signed by human experts [58, 41]. While existing methods achieve fast reward convergence [71, 6], many of them are diversity-lacking, prior-ignoring. and/or computationally expensive.
To address this issue, we take the best from both traditional reinforcement learning (RL) and direct reward maximization approaches and propose novel reinforcement learning objective, dubbed V- DB, that leverages the rich information reward gradients. Our method rooted both generative flow networks (GFlowNets) [2], a framework and language for generative modeling a direct acyclic transition graph, and soft RL [15, 16]. Compared to existing gradicnn-inffoman methods, our probabilistic finetuning method achieves better sample diversity. We further propose V-GFlowNet with the objective residual VDDB that, by leveraging the structure of diffusion models, allows us to perform fast, diversity preserving and prior-preserving amortized finetuning of diffusion models with rather long sampling sequences.
We summarize major contributions below:
We propose V-DB and the pretrained-monl model-aware variant residual VDDB, the objectives that efficiently leverage the rich information in reward gradients in principled and probabilistic way.
Our V-GFlowNet the first GFlowNet method that considers first-order information in reward signals, and is closely connected to soft reinforcement learning.
We empirically show that with the proposed residual V-DB objective, we may achieve diversity- and prior-preserving yet fast finetuning and alignment of diffusion models.

---

Published as a conference paper at ICLR 2025
Residual V-DB
Qualitative Results
Vay log Prinia 66
Veu. Ay Prowilad))
Ver. log Polrelen-a R
Fly Pypekea
Vap Ing M(ave Ry--1 - Ty
24 Reyl
Forward-looking Residual V-DB
Squiree
Verui Aug (1
-May log. Pag.Baley (t
SURG0) REEalad)
Vep log M(re)
Py--1
Number of finetuning steps
Figure I: Left: Illustration of the proposed residual V-DB objective, along with its forward-looking variant. The two "forces" on each image in each transition 2ı + I141 out trajectory т (20p, 21, E3,,....2 2T) are expected to sum to zero. Green and blue terms represent forward and reverse residual policy scores (respec- tively), orange terms represent signals from terminal rewards and pink terms represent flow scores or residual flow scores, each of which is defined in Section 3. Notice that the reward term the in the final transition different from the others. Right: Generated image from model finetuned with the proposed residual V-DB on the Aesthetic Score reward. The text prompt for each row shown on the left. The leftmost figure is the image generated by the pretrained model while the rightmost one is from the model finetuned for 200 iterations.
PRELIMINARIES
2.1 DIFFUSION MODELS AND RL-BASED FINETUNING
Diffusion models [17, 59, 61] are class of hierarchical latent models that model the generation process as a sequence of denoising steps. Different from the convention in diffusion model literature, for convenience adopt this paper the reverse time of arrow where Ty means samples from the data distribution and the sampling process starts from f. = 0. Under this convention, the probability of the generated samples is:
PF(OT) = J
Po(xo) II Pe ( (e
(1)
Here Po(ro) is fixed initial distribution, (F(TT) is the likelihood of the model generating data IT. and the noisy states 2, in the intermediate time steps are constructed by a pre-defined noising process, and the forward policy Pp-(malarei) is the denoising step of the diffusion
model". Take DDPM [17] as an example: the corresponding noising process ((==-1/2;) N ((Vaa-1/aaa /I - ac-1/a;l) and induces q(x((=T) = N ((VTT(WTT vI - -a--f), where ] is noise schedule set. With this noising process defined, the training loss is:
E tr-Unnffrmm (f..mm((((................ (...))))))....)................ w(t)
nol Vor- 1T + v1 -1-or- oT (T--
FT
(2)
where D is dataset, w(t) is certain schedule weighting function, and ( t) is data prediction model that predict the clean data Vy given a noisy data x, at time step t.
The sequential sampling process of diffusion models is Markovian and one could construct an Markov desicion process (MDP) to describe its denoising process. Similar to existing approaches [3, 80]. given a Markovian diffusion sampling algorithm, sample (xz, 1) at each diffusion inference step as a state and a denoising step from &t to s441 transition. Given such MDP defined, one may finetune diffusion models with techniques like DDPO [3] by collecting on-policy sample trajectories { (T)] and optimize the forward (denoising) policy with a given terminal reward H(xT). We give more detailed overview of MDP construction and DDPO in Appendix F.
construct MDP by treating the noisy
'To clarify, the "forward" and "backward" directions in the GFlowNet literature typically the opposite of those the diffusion model literature.
2

---

Published as a conference paper at ICLR 2025
2.2 RL FINETUNING IN THE GELOWNET FRAMEWORK
It well known [15, 50] in the soft RL literature that, by maximizing Ex - (s0 in an MDP (with state St, terminating state sy, action or, terminal reward R(-), discount rate 7 pol- icy Th and entropy function H.), we have P(æ) x exp(3R(xz) for any terminating state z. Finding optimal policies to sample from R(-) can therefore be achieved with appropriate designs of MDPs. While we may derive our objective with the soft RL language, here we follow recent work [80] and adopt the GFlowNet [2] framework for generative models that sample terminal states with respect to reward with MDP defined on directly acyclic graph (DAG). The GFlowNet framework uniquely motivates the introduction of a special intermediate reward associated to reference backward pol- icy, which is not obvious in the soft RL framework. Furthermore, in modeling distributions with a DAG-based MDP, the GFlowNet framework is more general and perhaps more natural than soft RL. We discuss the connections between two frameworks in Appendix D.
In GFlowNet, the forward sampling policy Pp on DAG from an initial state so to a terminal state 8y paired with a backward "reference" policy (ss(sss). One can view Pp as distributing probability mass from to its successors, and Pn as moving mass backward from s' to its ancestors. When PF and Pn are matched, they induce an unnormalized density F(s) at each state, called the flow function. One way to specify the matching between Pp and Pn is through the following:
Detailed Balance (DB). A valid GFlowNet with a forward policy P=(s'|s), a backward policy (ss(ss'), and flow function F(s) satisfies the following DB condition for all transition (s + s')
P=((1(s)
PB (s(sss))F(s))
(3)
Hence we have the following GelowNet DB loss in the logarithm probability space:
LoB (8, 8)) log Py (s'1s) + log F(s) - log Pb(s)s') - log F(s')) 2
(4)
with an extra terminal constraint F(sy) = R(s)) to incorporate the target reward information.
In the context of time-indexed sampling processes such as diffusion models, the transition graph of states - (z f) is naturally acyclic, it adheres to the arrow of time [75]. With slight abuse of terminology that we use z, to represents the tuple (Iz, t) when necessary, for inde-iddexed settings
the forward policy is (x((e(1(22) the backward policy is PB (2 (z)(x(41), and the flow function F(xx)?. The corresponding DB condition is therefore
(5) To finetune a diffusion model with DB losses [80]. one can simply set Pr(x++ (2xe) to be the sam- pling process and fix (ate (12241) ) to be the noising process used by the pretrained diffusion model. If Pg is set to be fixed, we can similarly derive this condition using soft Q-learning [15], with details provided in Appendix D.
( (2)411)z) (xx) = P2((T))244)
3 REWARD FINETUNING VIA GRADIENT-INFORMER GFLOWNETS
3.1 VDDB: THE GRADIENT-INFORMEE DETAILED BALANCE
In setting, we do not have access any dataset of images, but are given external positive- valued reward function R(-) to which a generative model is trained to adapt. While typical GFlowNettbased algorithm can effectively achieve this objective with diversity in generated sam- ples, it only leverages the zeroth-order reward information and does not leverage any differentiability of the reward function. Yet, whenever the reward gradients are available, it is often beneficial to corporate them into the finetuning objective since they help navigate the finetuned model on the optimization landscape and may significantly accelerate the optimization process. We are therefore motivated to develop V-GFlowNet. method that builds upon GFlsw algorithms to take full advantage of the reward gradient signal. To achieve this, we take derivatives on the logarithms of both sides of the DB condition (logarithm of Equation 5) respect to I2+1 and obtain necessary condition, which we call the forward VDBB condition:
V Veeei log Pp ( (x:) = V Verui log Pa (1)( + Varus log F(2:41),
(6)
We write (x)) as F(x,) for the sake of notation simplicity. This is because the derivative is taken with respect to T++1-
3

---

Published as a conference paper at ICLR 2025
and hence the corresponding forward V-DB objective LD (2t, T241) to be
Vzt+1 log PF Pa (2<+1 - #141 log Pa(x)latt) - Varti log (a(-1))
(7)
with the terminal flow loss on the logarithm scale
LevB--erminal (xT) 22 log F(IT) VV., log R(xr)
(8)
where B is a temperature coefficient and serve as a hyperparameter in the experiments. No- tice that, by taking derivatives on the logarithms, we obtain the (conditional) score function Ves+. log (1(+++1)2)) of the finetuned diffusion model. Indeed, the V-DB loss is closely related Fisher divergence, also known as Fisher information score [25] (see Appendix B.I).
Similarly, by taking the derivative of both sides in Equation 5 with respect to one obtains the reverse V-DB objective:
L DB (e-1) Vz, log Po Vz, log (B((s) + Va, log log F(a))
(9)
Such V-GFlowNet objectives constitute valid GelowNet algorithm (see the proof Section B.2): Proposition 1. If T (12, De (xt, (t+1) 0 for any denoising transition (#t, T(+1) over the state space and D (ET) for all terminal state IT, then the resulting forward policy generate samples IT with probability proportional to the reward function R(xT)².
The original detailed balance condition propagates information from the reward func- tion to each state flow function in the sense of F(I++1) + (F(x;), 7r(2+41 (rx)), assuming the backward (noising) policy is fixed (i.e., there is learning component in the diffusion noising pro- cess). In our case, if we take a close look at Equation 7, we can see that De (x+.)) could propagate the information from F(1++1) to the forward policy P (zre) but not F(xz).
Remark 2.
Remark 3. Compared to previous GelowNet works which use scalar-output network to parame- terize the (log-) flow function, in V-GFlowNet we can directly use a U-Net [52]-like architecture that (whose output and imput shares the same number of dimension) to parameterize V log F(-). which potentially provides more modeling flexibility. Furthermore, it is possible to initialize V log F(.) with layers from the pretrained model so that it can learn upon known semantic information.
3.2 RESIDUAL DD FOR REWARD FINETUNING OF PRETRAINED MODELS
With the V-DB losses, one can already finetune a diffusion model to sample from the reward dis- tribution R(x). However, the finetuned model may eventually over-optimize the reward and thus forget the pretrained prior (e.g., how natural images look like). Instead, similar to other amortized inference work [82, 68), we consider the following objective with an augmented reward:
Pp(TT) R(xr)² " pt (xr),
(10)
where R(x) the positive-valued reward function, 8 is the temperature coefficient, P (xT) is the marginal distribution of the pretrained model and PF(*T) is the marginal distribution of the finetuned model (as defined in Equation 1).
Because both the finetuned and pretrained model share the same backward policy PB (the noising process of the diffusion models), we can remove the P3 term and obtain the forward residual V-DB condition by subtracting the forward V-DB equation for the pretrained model from the that of the finetuned model:
V 2o44 log PF (Z (1x) V Vetta log P (x=-1/2))= - Vees log F(Z(41) V meys log F# (#++1).
Ven.. log Pr (r( la:): residual policy function
Vzaya log (ar+1): residual flow function
(1I)
With the two residual terms defined above, we obtain the forward residual V-DB objective:
DD-rss (we, x++1) Vz. log Pp ( (mxa) V log ((c41)))
(12)
We use the notation of indicate quantities of the pretrained model.
4

---

Published as a conference paper at ICLR 2025
Similarly, we have the reverse residual V-DB loss:
DDB-res (Ie, (e++) = Vx, log Pr( /2x) + Va, log ((xz)
(13)
The terminal loss of the residual V-DB method stays the same form as in Equation 8
DDB-remmnnal (IT)
V log F(IT) BV or log R(xT)
(14)
for any denoising transition (Ie, F441) over the state space and L terminal = 0 for all terminal state XT, then the re- s policy generate samples 2T with probability proportional to R(FT) (IT).
Proposition 4.
If Da 2++1) L DB8 F++1)
Remark 5. We point out that perform the same way of deriving Equation 11, iee., subtraction be- tween GFlowNet conditions from the finetuned and pretrained model, on the DB condition without gradient, we can obtain a residual DB condition PF (#c+1|xr)) Pf (œ1+1 (1441/2))FF F((et))))))))) Mul- tiplying this condition across time and eliminate the term of intermediate Fxxx) will lead to the objective derived in the relative GFlowNet work [67] as shown in Section B.4, which is a prior paper that proposes work reward finetuning GFlowNets with given pretrained model.
Remark 6. One may completely eliminate the need for any residual flow score func- tion with the residual V-DB conditions of both directions: Va log Pf(x1+1/zz) = -VR++i log Pr((s+2l The bidirectional residual V-DB condition can be analogously un- derstood as the balance condition of two forces from 2, and Z+42 acting on 2441: if not balanced, one can locally find some other Uz++11 that makes both transitions more probable.
Flow reparameterization through forward-looking (FL) trick. Though mathematically valid, the bidirectional pair of V-DB conditions suffers from inefficient credit assignment for long se- quences, a problem commonly observed in RL settings [63, 66]. Instead, we may leverage the priors we have from the pretrained diffusion model to speed up the finetuning process and consider the individual conditions for the forward and reverse directions. Specifically, we employ the forward- looking (FL) technique for GelowNets [44, 80] and parameterize the residual flow score function with "baseline" of the "one-step predicted reward gradient":
V log (xx)) A B V., log R(ag(a))) + go(
(15)
predicted reward
where 7/e is the scalar to control the strength of forward looking with the constraint ?T 1 and (xxx) is the actual neural network with parameters satisfying the terminal constraint go = 0. Here (.) is the one-step clean data prediction defined in Equation 2. With the FL technique, one achieves faster convergences since it sets better initialization for the residual flow score function than naïve zero random initialization [44].
We therefore obtain the forward-looking version of residual VDDB losses of both directions:
L Le (11)) V set1 log ( | V Ze log R(Co(E=+1)) + (0(#441)
(16)
LDDB-FL-SS (t+,) (2,,
a, log PF + BT V., m, log ((to(z))) + go (17)
Moreover, the corresponding terminal loss objective now becomes
V øy log F(xr) By ør log R(zT) || | (ZT) (
LVDB-FLLAemmna
(18)
which indicates that the actual parameterized flow network go should take near-zero value for terminal states Jr. The total loss on collected trajectory T (20, 21,0,111-0 T1...., IT) is therefore
[[wa(t)L+nann (ot, 7141) wB ( (11))1+1) + LVDB FL (xr) (19)
where (t( and (t)()) scalar weights control the relative importance of each term.
We summarize the resulting algorithm in Algorithm in Appendix
Choice of FL scale. Naïvely setting % = 1 can be aggressive especially when the reward scale and the learning rate are set to a relatively high value. Inspired by the fact that diffusion models (x(2)) can be seen as R(x) smoothed with Gaussian kernel, we propose to set 7 Or--t-
5

---

Published as a conference paper at ICLR 2025
4 EXPERIMENTS AND RESULTS
4.1 BASELINES
For gradient- free methods, we consider DAG-DB [80] (ie., GFlowNet finetuning with the DB ob- jective) and DDPO [3]. Since the original DB objective aims to finetune with R (x) instead of p# (IT) (IT), we also consider the forward-looking residual DB loss, defined
Fe (F,, (e+1) (log ( log (( go(
+ R(to( + go(
(20)
For other gradient-aware finetuning methods, consider ReFL [71] and DRaFT [6]. ReFL samples a trajectory and stops at some random time step i, with which maximizes R(āg(Tz) where (n(() is the one-step sample prediction function and xı is the computed via denoising V (Ia-1) with y being the stop-gradien operation. Differently. DRaFT samples time step T K (typically K = 1) and expand the computational graph of DDPM from ((T-KK) to 2T so that R(ET) can be backprop- agated to TT-K. with all x; in the previous time steps removed from this computational graph. A variant of DRaFT called DRaFT-LV performs a few more differentiable steps of "noising-denoisiin on the sampled TT before feeding it into the reward function R(-). We follow the paper of DRaFT and use only one "noising-denoisee
step: -'T-1 1 ~ (B and r7 (=
4.2 REWARD FUNCTIONS, PROMPT DATASETS AND METRICS
For the main experiments, we consider two reward functions: Aesthetic Score [28]. Human Prefer- ence Score (HPSv2) [69, 70] and ImageReward [71], all of which trained on large-scale human pref- erence datasets such as LaIoN-aesthetic [28] and predict the logarithm of reward values. For base experiments with Aesthetic Score, we use a set of 45 simple animal prompts as used in DDPO [3]; for those with HPSv2, we photo+painting prompts from the human preference dataset (HPDv2) [69]. To measure the diversity of generated images, we follow Domingo- Eaninng-Ennrchh et al. [10] and compute the variance of latent features extracted from a batch of generated images (we use batch of size 64). Using the same set of examples, we evaluate the capability of prior preservation, we compute the per-prompt FID score between images generated from the pretrained model and from the finetuned model and take the average FID score over all evaluation prompts.
4.3 EXPERIMENT SETTINGS
For methods, we use 50-step DDPM sampler [17] to construct the MDP. StableDiifusiin-vll [51] is used as the base model. For finetuning diffusion model policies, we use LoRA [21] with rank The residual flow score function in residual V-DB is set to be a scaled-down version of the StableDiffusion U-Net, whereas the flow function (in DAG-DB and residual DB) is set to be similar network but without the U-Net decoding structure (since the desired output is a scalar instead of image vector). Both networks are initialized with tiny weights in the final output layers.
As the landscape of R(x) can be highly non-smooth, we approximate Vz, log (x))(z)) with Ee~A((..)) V.., log (2.) e) where is tiny constant. For StableDiffusion [51], since the diffusion process runs in the latent space, the reward function is instead (eowV(0.e) V.,, log ((decode((o(zz) + e)) in which decode(-) is the pretrained (and frozen) VAE decoder and is set to 2 10-3, slightly smaller than one pixel (i.e., 1/255). We approximate this expectation with 3 independent samples for each transition in each trajectory. For all experiments, we try 3 random seeds. We set wF (t) = for all t's and unless otherwise specified wp(t) =
To stabilize the training process of our method (residual VDDB), we follow the official repo of DAG-DB [80] and uses the following output regularization: (z)co(zz)) - cot (Ir where Øt the diffusion model parameters in the previous update step". For residual V-DB, We set the output regularization strength À = 2000 Aesthetic Score experiments and = 5000 in HPSv2 and ImageReward experiments. For all experiments with residual VDDB, we set the learning rate to
FEsenniiall Fisher divergence between the pretrained and the finetuned distributions conditioned Similar regularizations with KI. divergence has been seen in RI. algorithms like TRPO [54] and PPO [55].

---

Published as a conference paper at ICLR 2025
Pretrained Reward - 5.62
Residual V-DB
Residual
Reward 7.86
DAG-DB
Reward = 8.03
ReFL Reward 6.67
DRaFT-LV Reward 6.69
DDPO
Figure 2: Comparison between images generated by models finetuned with different methods for maximum of 200 update steps. For each method, pick the model trained that produces images with the highest rewards without semantic collapse among all model cheekpoints, methods like RcFL and DRaFT-LV casily collapses (as illustrated in Fig. 3). For each method, we show the average reward the corresponding presented images.
x 10-3 3 and ablate over a set choices reward temperature B, in a range such that the reward gradients more significant than the residual policy score function V., log PF(x (ax,) of the pretrained model. For HPSv2 and ImageReward experiments, we set (9 to be 500000 and 10000, respectively. For each epoch, we collect 64 generation trajectories for each of which we randomly shuffle the orders of transitions. We use the number of gradient accumulation steps to 4 and for each 32 trajectories we update both the forward policy and the residual flow score function. For residual V-DB in most of the experiments, we for training sub-sample 10% of the transitions in each collected trajectory by uniformly sample one transition each of the uniformly split time-step intervals but ensuring that the final transition step always included.
For residual DB and DAG-DB, we set learning rate to x 10-4 with the output regularization strength X = The sampling and training procedures are similar to the residual DD experiments. For both ReFL and DRaFT, we use a learning rate 10-4, For ReFL,, we follow the official repo and similarly set the random stop time steps to between 35 and 49, For DRaFT, since the official code released, follow the settings AlignProp [48], similar concurrent paper. We set the loss for both ReFL and DRaFT E ReLU RR decode(2z)) where the ReLU function is introduced for training stability in the case of the ImageReward reward function.
Residwal V-DB
Residual DB
ReFL
DRaFT-LV
Epoch 10
Epoch 20
Ejoch 40
Fpoch 70
Figure 3: Our V-GFlowNet finetuning yields stable output compared to other baselines.
4.4 EXPERIMENTAL RESULTS
General experiments. In Figure and Table 1, we show the evolution reward, DreamSim diver- sity and FID scores of all methods with the mean eurves and the corresponding standard deviations

---

Published as a conference paper at ICLR 2025
Pretrained Residual V-DB
Residual
DDPO
ReFi.
DRaFT-LV
Prompt: The image features castle surrounded by dreamy garden with roses and a clondy sky in the background.
Prompt: A painting a Rladerunner interior room in Africa with detailed artwork.
Prompt: A broken videogame console with colorful and compelling painting.
Figure 4: Qualitative results on HPSv2.
Pretrained
Residual V-DB
Residual DB
DDPO
ReFL
DRaFT-LV
Prompt: IT-guy trying fix hardware PC tower being tangled by the PC cables like Laokoon. Marble, copy after Hellenistic original from ca. 200 BC. Found in the Baths of Trajan, 1506.
Prompt: zebra underneath broccoli.
Prompt: blue colored dog.
Figure 5: Qualitative results on ImageReward.
(on 3 random seeds). Our proposed residual V-DB is able to achieve comparable convergence speed, measured in update steps, that of the gradient-free baselines while those diversity-aware baselines fail to do so. In Figure 7, we plot the diversity-reward tuples and FID-reward tuples for models evaluated different checkpoints (every update steps) and show that our method achieves Pareto improvements on diversity preservation, prior preservation and reward. The gradient-informer base- lines, ReFL. and the DRaFT variants, generally behave worse than residual V-DB due their mode- seeking nature. Qualitatively, we show that the model finetuned with residual V-DB on Aesthetic Score generate more aesthetic and more diverse samples both style and subject identity (Fig. and Fig. 24), while the other baselines exhibit mode collapse or even catastrophic forgetting of pre- trained image prior. We also demonstrate some images generated by the diffusion model finetumed with residual DB HPSv2 in Figure 4 and on ImageReward in Figure 5. Furthermore. qual- itatively show that residual VDDB robust while with the gradient-informer baseline methods are prone training collapse (Fig. 3). Due to limited space, here only show the most important abla-
8

---

Published as a conference paper at ICLR 2025
Aesthetic Score
HPSy2
ImageReward Diversity
Method
Reward (t)
Diversity DreamSim
FID
Reward
DreamSim
Diversity
FID (4)
Reward (t. 10
DreamSim
FID
<t. 10 B)
(
<t. 10
10 B)
(T. 10 3)
(4)
Base Model
5.83 0:01
35,91 + 0.00
216
2.38 0.13
37.75 a 021
563 5
-0.38 # 0.12 41.09 0.03 468 +:
DDPO
6.68 0.14
32,96 + 1.01
312
2.52 0.04
3,49
0.03
681 +
0.27 0.38
38.51
1.49
714
ReFL
9.53 0.46
8.20 3.06
1765
3.67 0.06
19.84 1,50
1191 46
1.36 0.30
36.50 0.52
597 +
DRaFT-1
10.16 1: 0.13
4.24 1: 0.45
1665 à 182
3.70 0.06
18.96 A 1.35
1222 84
1.59 0.25
37.27
0.49 531 +:
DRaFT-LV
10.21 0.34
6.39 1.66
1854
3.75 0.08
21.13
1.19
1164 43
0.25
37.56
529
DAG-DD DB
7.73 D:07
15.88 1: 0.70
595 87
2.52 0.06
32.50 1: 0.59
866 1: 41
4.70 0.75
26.78
0.64 809 1: 34
residual DR
7.20
19.38 1: 4.07
1 587
2.55 0.06
32.49 1: 0.47
0.54
772 1:
V V-GFoowNet
(=g=0)
7.86 0.06
29.21
0.18
318
3.53 0.03
24.40
0.56
973
5.33 0.62
35.85 0.66 638
GFlowNet
(=g=1)
7,90 = 0:09 29.67 1 0.51
317 =
15 3.53 1: 0.04 24.39 1: 0.87 1000
39
2.23
0.34
39.85
0.67 501 1:
Table 1: Comparison between the models finetuned with our proposed method V V-GFlowNet (with the objective residual V-DB) and the baselines. All models are finetuned with 200 update steps. For each method, the model with the best mean reward is used for evaluation. Note that while baselines like DDPO can achieve better scores on some metrie, often comes with the price of much worse performance on some other.
tion studies and leave the rest to Appendix J. For the same reason, leave plots for experiments HPSv2 and ImageReward to Appendix and qualitative comparisons on diversity to Appendix M.
Effect of reward temperature. We perform ablation study on Aesthetic Score with 8 E {5000, 7000, 10000) in residual V-DB. Not surprisingly, a higher reward temperature leads to faster convergence at the cost of worse diversity- and worse prior-preservation, as observed in Figure 8.
Effect of sub-sampling. Typically, sub-sampling results in worse gradient estimates. We empir- ically study how sub-sampling may effect the performance and show the results in Figure 10 and 11 in the appendix. We empirically do not observe huge performance drop due to the subsampling strategy, potentially because the rich gradient signals in both the reward and the flow are sufficient.
Effect of attenuating scaling predicted reward. In Fig. 16 and 17, we show that the model trained without attenuation converges faster, but at the cost of worse diversity and prior-preservation.
Reward finetuning with other sampling algorithms. To show that our method generalizes to different sampling diffusion algorithms, we construct another MDP based on SDE-DPM-Solvver [38], with 20 inference steps. In Fig. 18 and 19, we obverse that our residual VDDB can still achieve a good balance between reward convergence speed, diversity preservation and prior preservation.
Residual V-DB
Residual DB
DDPO
ReFl.
DRaFT-LV
DreanSii Divenity
0.
2000
Rowad
0.5
1500
6.
500
0.0
01
50
1001
1550
1000 150
200
50
1000
150
Updlstc Step
Updase Stop
Updovss Stop
Figure 6: Convergence curves of different metrics for different methods throughout the finetuning process on Aesthetic Score. Finetuning with our proposed residual V-DB converges faster than the non -e--grrdent-iffommee methods and with better diversity- and prior-preserving capability.
5 RELATED WORK
Reward finetuning of diffusion models. The demand for reward finetuning is commonly seen in alignment, where one obtains utilize human reference reward function to align the behavior of generative models [22, 1, 69] for better instruction following capability and better AI safety. With a reward function, typically obtained by learning from human preference datasets [84, 62]. may use reinforcement learning (RL) algorithms, for instance PPO [55]. to adapt not only autoregressive language models [43] which naturally admits Markov decision process (MDP), but also diffusion

---

Published as a conference paper at ICLR 2025
Daaddty ve Revad, Aucheti Scace
FED et. Recocd Archii Sexa
Roidal 9-08
Retikal
DEP0
keFl.
DRAFT1
Drauddan Divests
Reikeal
Figure 7: Tradc-offs between reward, diversity preservation and prior preservation for different reward fine- tuning methods. Dots represent the evaluation results of models checkpoint saved after every iterations of finctuning, where ones with larger reward, larger diversity scores and smaller FID scores considered better.
B=10000
B-7000
B=5000
DesunStn Diessty
0.36
Rownnd
0.34
0 250
0.21
150
30
Uplate Step
Opbhac
Updatc Skep
Figure 8: Higher temperature B leads to faster convergence but with less diversity and less prior preservation.
models [3, 12]. Specifically, one can construct MDPs from some diffusion sampling algorithm [59, 38] by considering each noisy image at some inference step as a state and each denoising step as an action. Besides RL algorithms, there exist some other approaches, including stochastic optimal control [10, 65], GFlowNets [80] and some other ones akin to RI. methods [30, 11]. While most of the aforementioned approaches train with only black-box rewards, once have access differentiable reward function we may accelerate the finetuning process with reward gradient signals. For instance, methods exist construct computational graph from sampled generation trajectories to directly optimize for rewards [6, 48, 71]. yet with these methods models are not trained to correctly sample according the reward function. While one may also generate samples from the reward function without finetuning using plug-in guidance methods for diffusion models [9, 60, 26, 14] as an alternative, but the generated distributions are often very biased. Besides, reward finetuning for diffusion models is typically memory consuming as many methods require a large computational graph rolled out from long generation trajectories, for which it typical to employ parameter- efficient finetuning techniques [21, 49, 33).
NetswNets. Generative flow network (GFlowNet) [2] is a high-level algorithmic framework that introduces sequential decision-making into generative modeling [75], bridging methodology be- tween reinforcement learning [79, 45, 46, 44, 29] and energy-based modeling [76]. GFlowNets are versatile since many models, such as diffusion models, can be easily treated as specifications of GFlowNets [75, 27, 78, 56]. GFlowNets perform amortized variational inference [40] and gen- erate samples with probability proportional to given density or reward. GFlowNets are used in various applications including but not limited drug discovery [23, 24, 57], structure learning [8], phylogenetic inference [83], combinatorial optimization [74, 77], and prompt adaptation [73].
6 CONCLUDING REMARKS
We propose --FFloNNet, fast, diversity and prior prior-preserving diffusion alignment method, by leveraging gradient information in the probabilistic framework of GFlowNets that aims to sample according to a given unnormalized density function or reward. Specifically, we develop V-DB, the gradient-informed version of the Detailed Balance objective, and the variant of residual V-DB with which one may finetune a diffusion model with reward results show that V-GFlowNet achieves a better trade-off between convergence speed, diversity in generated samples and prior preservation. We hope that our method sheds lights on future studies on more efficient post-training alignment strategies of diffusion models well as related applications.
rrior-preserving way. Our empirical
10

---

Published as a conference paper at ICLR 2025
REFERENCES
[1] Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05886 2022.
[2] Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J Hu, Mo Tiwari, and Emmanuel Ben- gio. Gflownet foundations. Journal of Machine Learning Research, 2023. 1, 3, 10
[3] Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. In ICLR, 2024. 2, 6, 10
[4] Andreas Blattmann, Tim Dockhorn, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Do- minik Lorenz, Yam Levi, Zion English, Vikram Voleti, Adam Letts, et al. Stable video diffu- sion: Scaling latent video diffusion models to large datasets. arXiv preprint arXiv:3311.15127, 2023.
[5] Paul Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. In NeurlPS, 2017.
[6] Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffusion models on differentiable rewards. In ICLR, 2024. 1, 6, 10
[7] Tristan Deleu. Generative flow networks: Theory and applications to structure learning. arXiv preprint arXiv:2501.05499, 2025. 20
[8] Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. In UAI, 2022. 10
[9] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In NeurIPS, 2021. 1, 10
[10] Carles Dwwingo--miiiiinneeeo Enrich, Michal Drozdzal, Brian Karrer, and Ricky T. Chen. Adjoint match- ing: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. In ICLR, 2024. 6, 10,
[11] 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. Transactions on Machine Learning Research, 2023.
10
[12] Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, P. Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. In NeurlPS, 2023. 10
[13] Gege Gao, Weiyang Liu, Anpei Chen, Andreas Geiger, and Bernhard Schölkopf. Graph- dreamer: Compositional 3d scene synthesis from scene graphs. In CVPR, 2024. 1
[14] Yingqing Guo, Hui Yuan, Yukang Yang, Minshuo Chen, and Mengdi Wang. Gradient guidance for diffusion models: optimization perspective. In NeurIPS, 2024. 10
[15] Tuomas Haamoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In ICML, 2017. 1, 21
[16] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critie: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In JCML, 2018.
[17] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In NeurlPS, 2020. 1, 6, 22
11

---

Published as a conference paper at ICLR 2025
[18] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al. Imagen video: High definition video generation with diffusion models. arXiv preprint arXiv:2210.02303, 2022.
[19] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. In NeurlPS, 2022.
[20] Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling, Equivariant diffusion for molecule generation 3d. In ICML, 2022. 1
[21] Edward 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. In ICLR, 2022. 6, 10
[22] Borja Ibarz, Jan Leike, Tobias Pohlen, Geoffrey Irving, Shane Legg, and Dario Amodei. Re- ward learning from human preferences and demonstrations atari. In NeurIPS, 2018.
[23] Moksh Jain, Emmanuel Bengio, Alex Hernandez- Garcia, Jarrid Rector-Brooks Bonaven- ture FP Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang. Michael Kilgour, Dinghuai Zhang, et al. Biological sequence design with gflownets. In ICML, 2022. 10
[24] Moksh Jain, Tristan Deleu, Jason Hartford, Cheng-Hao Liu, Alex Hernández-García, and Yoshua Bengio. Gflownets for ai-driven scientific discovery. ArXiv, 2023, 10
[25] Oliver Johnson. Information theory and the central limit theorem. World Scientific, 2004. 4
[26] Lingkai Kong, Yuanqi Du, Wenhao Mu, Kirill Neklyudov, Valentin De Bortol, Haorui Wang, Dongxia Wu, Aaron Ferber, Yi-An Ma, Carla P Gomes, et al. Diffusion models as constrained samplers for optimization with unknown constraints. In AISTATS, 2024. 10
[27] Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-Garef.. Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of continuous generative flow networks. In ICML, 2023. 10
[28] LAION.
Laion aesthetic score predictor. la att.onnetteeticss/// 2024 Accessed: 2024-09-27. 6
https:///adon.nll/lloo.
[29] Flaine Lau, Stephen Zhewen Lu, Ling Pan, Doina Precup, and Emmanuel Bengio. Qgfn: Controllable greediness with action values. NeurIPS, 2024. 10
[30] Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuqing Du, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, and Shixiang Shane Gu. Aligning text-to-image models using human feedback. arXiv preprint M.Xi::2302.12990 2023. 10
[31] Yaron Lipman, Ricky TQ Chen, Heli Ben- Dam-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In ICLR, 2023. 22
[32] Shengchao Liu, Divin Yan, Weitao Du, Weiyang Liu, Zhuoxinran Li, Hongyu Guo, Christian Borgs, Jennifer Chayes, and Anima Anandkumar, Manifold-constrainee nucleus-level denois- ing diffusion model for structure-based drug design. arXiv preprint arXiv:2409. 105844 2024.
[33] Weiyang Liu, Zeju Qiu, Yao Feng, Yuliang Xiu, Yuxuan Xue, Longhui Yu, Haiwen Feng. Zhen Liu, Juyeon Heo, Songyou Peng, et al. Parameter-efficient orthogonal finetuning via butterfly factorization. In ICLR, 2024. 10
[34] Zhen Liu, Yao Feng, Michael J Black, Derek Nowrouzezahraii, Liam Paull, and Weiyang Liu. Meshdiffusion: Score-based generative 3d mesh modeling. In ICLR, 2023. 1
[35] Zhen Liu, Yao Feng, Yuliang Xiu, Weiyang Liu, Liam Paull, Michael J Black, and Bernhard Schölkopf. Ghost on the shell: An expressive representation of general 3d shapes. In ICLR, 2024.1
12

---

Published as a conference paper at ICLR 2025
[36] Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In ICML, 2024.
[37] Justin Lovelace, Varsha Kishore, Chao Wan, Eliot Shekhtman, and Kilian Q Weinberger. La- tent diffusion for language generation. In NeurlPS, 2024.
[38]
Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver++: Fast solver for guided sampling of diffusion probabilistic models. In ICLR, 2023. 10, 22
[39] Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in gflownets. In NeurIPS, 2022. 19
[40] Nikolay Malkin, Salem Lahlou, Tristan Delcu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference. In ICLR, 2023. 10
[41] Azalia Mirhoseini, Anna Goldie, Mustafa Yazgan, Joe Jiang, Ebrahim Songhori, Shen Wang, Young-Joon Lee, Eric Johnson, Omkar Pathak, Sungmin Bac, et al. Chip placement with deep reinforcement learning. arXiv preprint arXiv:2004. 10746, 2020.
[42] Ofir Nachum, Mohammad Norouzi, Kelvin Xu, and Dale Schuurmans. Bridging the gap be- tween value and policy based reinforcement learning. In NeurIPS, 2017. 21
[43] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, al. Training language models to follow instructions with human feedback. In NeurIPS, 2022. 1,9
[44] Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of gflownets with local credit and incomplete trajectories. In ICML, 2023. 5, 10
[45] Ling Pan, Dinghuai Zhang, Aaron Courville, Longbo Huang, and Yoshua Bengio. Generative augmented flow networks. In ICLR, 2023. 10
[46] Ling Pan, Dinghuai Zhang, Moksh Jain, Longbo Huang, and Yoshua Bengio. Stochastic gen- erative flow networks. UAI, 2023.
[47] Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall, Dreamfusion: Text-to-3d using 2d diffusion. In ICLR, 2023.
[48] Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text- to-image diffusion models with reward backpropagation. arXiv preprint 1rX1v:2310.03739,
2023. 7, 10
[49] Zeju Qiu, Weiyang Liu, Haiwen Feng, Yuxuan Xue, Yao Feng, Zhen Liu, Dan Zhang, Adrian Weller, and Bernhard Schölkopf. Controlling text-to-image diffusion by orthogonal finetuning. In NeurlPS, 2023. 1, 10
[50] Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. In NeurlPS, 2023. 3
[51] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High-resolution image synthesis with latent diffusion models. In CVPR, 2022. 1,6
[52] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. MICCAI, 2015.
[53] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kam- yar Ghasemipour, Raphael Gontijo Lopes, Bureu Karagol Ayan, Tim Salimans, et al. Photo- realistic text-to-image diffusion models with deep language understanding. In NeurIPS, 2022.
[54] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In ICML, 2015.
13

---

Published as a conference paper at ICLR 2025
[55] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1007.06347, 2017. 7. 6, 9,
[56] Marcin Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Rector- Brooks, Alexandre Adam, Yoshua Bengio, and Nikolay Malkin. Improved off- oof-co training of diffusion samplers. In NeurlPS, 2024. 10
[57] Max W Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, and Tommaso Biancalani. Towards understanding and improving gflownet training. In ICML, 2023. 10
[58] David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, L. Sifre, George den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Vedavyas Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy P. Lilli- crap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nature, 529:484 489, 2016. URL https: / api. semant icscholar. org/ Corpus 515925. 1
[59] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In ICLR, 2021. 2, 10, 22
[60] Jiaming Song, Qinsheng Zhang, Hongxu Yin, Morteza Mardani, Ming- Yu Liu, Jan Kautz, Yongxin Chen, and Arash Vahdat. Loss-guided diffusion models for plug-and-play controllable generation. In ICML, 2023. 10
[61] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In ICLR, 2021. 1,2
[62] Nisan Stiennon, Long Ouyang, Jeff Wu, Daniel M. Ziegler, Ryan Lowe, Chelsea Voss, Alec Radford, Dario Amodei, and Paul Christiano. Learning to summarize from human feedback. In NeurlPS, 2022. 9
[63] Richard S. Sutton. Learning to predict by the methods of temporal differences. Learning, 3:9 44, 1988.
Machine
[64] Daniil Tiapkin, Nikita Morozov, Alexey Naumov, and Dmitry Vetrov. Generative flow net- works regulie-re rl. In AISTATS, 2024. 20
[65] 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-time diffusion models as entropy-regularie control. arXiv preprint arXiv:2402.15/99 15194, 2024. 10
[66] Hado van Hasselt, Yotam Doron, Florian Strub, Matteo Hessel, Nicolas Sonnerat, and Joseph Modayil. Deep reinforcement learning and the deadly triad. arXiv preprint arXiv: 02644, 2018.
[67] Siddarth Venkatraman, Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, et al. Amortizing in- tractable inference in diffusion models for vision, language, and control. In NewrlPS, 2024. 5, 19
[68] Luhuan Wu, Brian Trippe, Christian Naesseth, David Blei, and John P Cunningham. Practical and asymptotically exact conditional sampling in diffusion models. In NewrIPS, 2024. 4
[69] Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: A solid benchmark for evaluating human preferences of text- to-image synthesis. arXiv preprint 23030223006003341 2023. 6,
[70] Xiaoshi Wu, Keqiang Sun, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score: Better aligning text-to-image models with human preference. In ICCV, 2023. 6
14

---

Published as a conference paper at ICLR 2025
[71] Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagereward: Learning and evaluating human preferences for text-to-image generation. In NeurlPS. 2024. 1, 6, 10
[72] Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In /CLR, 2022.
[73] Taeyoung Yun, Dinghuai Zhang, Jinkyoo Park, and Ling Pan. Learning to sample effective and diverse prompts for text-to-image generation. arXiv preprint arXiv:2502.11477 2025. 10
[74] David Zhang, Corrado Rainone, Markus Peschl, and Roberto Bondesan. Robust scheduling with gflownets. In ICLR, 2023. 10
[75] Dinghuai Zhang. Ricky T. Q. Chen, Nikolay Malkin, and Yoshua Bengio. Unifying generative models with GFlowNets and beyond. arXiv preprint arXiv:2209.0266612 2022. 3, 10
[76] Dinghuai Zhang. Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. In ICML, 2022.
10
[77] Dinghuai Zhang, Hanjun Dai, Nikolay Malkin, Aaron C Courville, Yoshua Bengio, and Ling Pan. Let the flows tell: Solving graph combinatorial problems with gflownets. In NeurIPS, 2023. 10
[78] Dinghuai Zhang, Ricky T. Q. Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffusion generative flow samplers: Improving learning signals through partial trajectory op- timization. In ICLR, 2024. 10
[79] Dinghuai Zhang, Ling Pan, Ricky T. Q. Chen, Aaron Courville, and Yoshua Bengio. Distri- butional GFlownets with quantile flows. Transactions on Machine Learning Research, 2024.
10
[80] Dinghuai Zhang, Yizhe Zhang, Jiatao Gu, Ruixiang ZHANG, Joshua M. Susskind, Navdeep Jaitly, and Shuangfei Zhai. Improving GFlownets for text-to-image diffusion alignment. Trans- actions on Machine Learning Research, 2025. 2, 5, 6, 10
[81] Longwen Zhang, Ziyu Wang, Qixuan Zhang, Qiwei Qiu, Anqi Pang, Haoran Jiang, Wei Yang, Lan Xu, and Jingyi Yu. Clay: A controllable large-scale generative model for creating high- quality 3d assets. ACM Transactions on Graphics (TOG), 2024.
[82] Stephen Zhao, Rob Brekelmans, Alireza Makhzani, and Roger Baker Grosse. Probabilistic inference language models via twisted sequential monte carlo. In ICML, 2024. 4
[83] Ming Yang Zhou, Zichao Yan, Elliot Layne, Nikolay Malkin, Dinghuai Zhang, Moksh Jain, Mathieu Blanchette, and Yoshua Bengio. PhyloGFN: Phylogenetie inference with generative flow networks. In ICLR, 2024. 10
[84] Daniel M Ziegler, Nisan Stiennon, Jeffrey Wu, Tom B Brown, Alec Radford, Dario Amodei, Paul Christiano, and Geoffrey Irving. Fine-tuning language models from human preferences. arXiv preprint arXiv: 190::099999999 2019. 9
15

---

Published as a conference paper at ICLR 2025
Appendix
Table of Contents
A Overall algorithm
17
B Algorithmic details
18
B.1
V-DB objective as statistical divergence
18
B.2 Proof of Proposition
18
B.3 Proof of Proposition
B.4 Relationship between residual DB and Trajectory Balance
Generalization the learnable Pa setting
Commons differences between GFlowNet soft RI
20
D.I GFlowNet the lens of soft RL
20
D.2 Derivation of V-DB with soft Q-learning
21
E Corrections for non-ideal pretrained models
21
F MDP construction for diffusion models
22
Application flow matching models
22
Finetuning convergence wall time
23
Results of ablation study
24
J Additional ablation experiments
27
K More samples (Aesthetic Score)
29
L More samples (HPSv2)
30
M Qualitative comparison on diversity of samples generated by different methods
32
16

---

Published as a conference paper at ICLR 2025
A OVERALL ALGORITHM
Algorithm V-GFlowNet Diffusion Finetuning with residual V-DB
1: Inputs:
Pretrained diffusion model fa* with the MDP constructed from some Markovian sam- pler, reward function R(-)
2: Initialization: Model to finetune fo with 8 residual flow score function g+(+).
3: Sample the initial batch of trajectories prev ( (T)i)im with the current finetuned diffusion model fo-
4: Set (1t 0.
5: while not converged do
6: Sample batch of trajectories Peurr {(21 IT)) with the finetuned
diffusion model.
7: Subsample the time steps to train with: the full set T; = (1...., T or the sampled T; Sample-N(11,--. ..,, T}) where Sample-N is unbiased sampling algorithm to randomly pick N samples.
8: Compute the loss
(e LV (Zt, 2411: o, R)
+ FI () + ( - Sor (2x)
9:
Set 81 +
10:
Update the diffusion model and the residual flow score function.
11: Set D Dpevv <- Dear-
end while
13:
return finetuned model fo.
17

---

Published as a conference paper at ICLR 2025
B ALGORITHMIC DETAILS
B.I V-DB OBJECTIVE AS A STATISTICAL DIVERGENCE
LyDB (Equation 7) is analogous Fisher divergence (up to constant scale) if we always on-policy samples to update the diffusion model and the flow function:
DFisher ( PB(z; (1++1)
F(x))
1
V Vonn log (3(2))2441) V Ves4i log
E 2 2 P
F(x))
E
LID
(21)
B.2 PROOF OF PROPOSITION
Proof. When the training objectives equal for all states, we would have
V se+1 log (p(x++1 (xz) V ..+1 log Pa (B( Futs log F4+1(1+41)
(22)
V me log (r( (x2)) = V... log (2( V., log (1
(23)
VaP log FxT)) vV log R(xT) for any trajectory (za, .......,rr)
(24)
Through indefinite integral, these indicate that there exist function Ca(as) satisfies
C,(xt) (#++1 (x)) F F((x.) Pr (#t41l2)) = (c-1((4)) (ETA((eea)) F(or) R(
(25)
(26)
(27)
Therefore, we have
C2(21) (141(1441) (x())) (41(()+1))
(28)
The right hand side does not depend He, therefore, the left hand side is constant. So have
((())) x F((x)), ll.
(29)
The probability of generating a data xy then equals
PF(xT) / Po(zol@) II P
(30)
/ Fo(zo)
II Fet1 (T44 dxoor-1
(31)
C+(1a)
/ Fa(wo) II
Fet1 ( dxorr-1
(32)
F((x))
x F(er) / II (1((
(33)
F F(TT) R(TT),
(34)
which proves the validity of the V-GFlowNet algorithm.
B.3 PROOF OF PROPOSITION
When the training objectives equal 0 for all states, we have
V Vz, log Pr( 1/xx) = Vayta log ((re+1)
(35)
V 2e log ( = - Vz, log F
(36)
V 2x log F(xr) ,V log R(xr)
(37)
18

---

Published as a conference paper at ICLR 2025
for any trajectory (fu....,.T)) Through indefinite integral, these indicate that there exist a function Ca(a)) satisfies
C((z)) (x)) F(I++1)
(38)
Pr (x)) C+1(2241)
(39)
F(xT) x R(ar).
(40)
Thus we have
C.(rt)) (141(1441) (x4))) C++1(2(+1))
(41)
The right hand side does not depend Tr, therefore, the left hand side is constant. So we have
C((xa) x Fz(z)) Vt.
(42)
The probability of generating data Fy then equals
/ Po 2 II Pr (x (x))dro:T-1
PF(ØT)
(43)
= f Pf (zol (2) II P# (x))) II Pr(xttila doyo0----1
(44)
f pøt (zu (2) II p (x4+)l=))
Fet1( droor-1
(45)
Cı(z;)
(T(TT) / pof (2o1e)
II p# (Ie (
(46)
x [(IT)* p#
p (KT).
(47)
Hence zero V-DB losses in both forward and reverse direction for all (Ire P141) pairs yield the ideal tilted distribution R(xT)* p# (ZT).
B.4 RELATIONSHIP BETWEEN RESIDUAL DB AND TRAJECTORY BALANCE
Here we illustrate the Remark 5. different but equivalent condition for GFlowNets states:
Trajectory Balance (TB) [39]. The following TB condition must hold for any transition sequence (so, .1.. SN) where Ze the unique starting state the MDP and s is a terminal state, given GFlowNet with the forward policy (r [s) and the backward policy PB (s]s'):
Z Pp('ss)
log R(xr) II (s((ss')
(48)
where Z = F(2o) is the total flow and R(ør) F(zr) the reward. The proof is immediate with telescoping product of the DB condition.
With an (ideal) pretrained model p# and the satisfication of the finetuning objective of Equation 10, can prove the conclusion in [67]:
2 Pr('ss)) Bog R(xr),
log 2 P
(49)
which also an immediate result telescoping products of the residual DB condition (which leads to Equation 20):
19

---

Published as a conference paper at ICLR 2025
log (F( log F(x++1) log F(xe).
(50)
While Equation 49 and residual DB are mathematically equivalent, implementation-wiss TB in Equation 49 demands the whole sampling sequence be stored in the memory for gradient com- putation, or one has resort to the time-costly technique of gradient checkpointing. In comparison, with DB-based methods one may amortize the computational cost into flows different time steps and therefore allow diffusion finetuning with flexible sampling sequence, of which the distribution approximation capacity and generation performance are generally greater.
C GENERALIZATION TO THE LEARNABLE PB SETTING
the backward process P3 learnable, the resulted Pa does not cancel the backward propcess P the pretrained model in the derivation of the residual V-DB condition. Let Let ( [1++1; 0B)
log Pn (2 8B) log P (z, (1.1x We instead have the following residual V-DB losses:
Vayji log Pr (1++1 De; 0) - Vxcn log (n(malaeet;;0n) = Vacia log ((I+1:0)
(51)
Va, log (1 (re 0) Vx, log (z 2. log ())
(52)
Therefore, we have the general residual V-DB losses:
Loorward general ( x1+1)
M log Pp (zz Vorni log (1(2)(2441000) V log f(g
(53)
and
Treverse, general (Ie, 1z+1)
V log (1 0) V., log ( V., log (xx
(54)
D COMMONS AND DIFFERENCES BETWEEN GEL AND SOFT RL
D.1 GFLOWNET IN THE LENS OF SOFT RL
We restate the finding [64, 7] that, by constructing a special MDP, we may prove that GFlowNet and soft RL are equivalent when PB is fixed.
Theorem 7. Let M (G.r) be an MDP define on DAG & and reward function r(s. (defined for any transition + s' G). Let PB be an arbitrary backward transition probability on G. Suppose for any complete trajectory T = (S $1,..., ST), we have the intermediate reward
R log (a(St St+1)
tt T -t=
r(st, 8t+
log R(sT) + log PI((s ((41)
Then the optimal policy with respect to the following objective (where H is the entropy function)
T
r(s St41) + H (-(( st)) Σ
π arg max ETAT
samples terminating states x with probability proportional to R(x)². Remark 8. The introduction of the intermediate reward terms log Pg (st(st+)) (with fixed PB) es- sentially anchors the optimal solution around the reference backward policy, which is rather natural from the generative model perspective and beneficial in creating better optimization landscapes. Remark 9, Soft RI. does not cover the case where Pn be jointly learned.
20

---

Published as a conference paper at ICLR 2025
D.2 DERIVATION OF V-DB WITH SOFT Q-LEARNING
With soft Q-learning [15] (a special case of path consistency learning [42]) in the MDP defined in Sec D.I with fixed Pn, aims
Asqu(s + + Q(s, (r(8,))))
= ++
(((())
(55)
where the soft value function is V(8)) log Eo' exp (3Q(s' ss)) and Qss, is the soft Q-value function. One show that
exp(3Q(s,5')
E. Pr(s'(s) (s) E. (xx((3((s,5))
(56)
Therefore, at optimality where Asqr(s + 8') = 0 for all s', can plug in the definitions above and obtain
log Ps(''(s) - log PB(s)s') = V(s') - V(s)
(57)
which is in the same form the DB condition. Taking gradients of the above condition with respect to both and s' leads to equivalent V-DB conditions.
E CORRECTIONS FOR NON-IDEAL PRETRAINED MODELS
For non-ideal pretrained models in which F does not match PB (i.e., the DB condition is violated), the original residual V-DB objective is apparently biased. We may introduce an additional Icarnable term h F1+1; 4/) compensate for this error, with which we have the corrected DB condition for the pretrained model:
log Pf (1++1l34)) - log (n( loog F# (ar++) - log ## (x:) h(x,, Te+1)
(58)
Hence, we have the following residual V-DB losses the case of the imperfect pretrained model:
L Da ((+1) Verus log (2 Verii log (I++1) + V V h(T,, T(415 1))
(59)
L LeD (It, 2x41) Vz, log Pr(a)+1(x)) Vay log (( + (a, (11
(60)
and the bidirectional V-DB loss for learning h in the pretrained model:
I pe pretrained 21t1, = V #ot1 log Pf (1e41 (2x) + V seti log Pf (2442)
Vzan log Pr (x(0:12te))) Vaan log Pn (n(
V bix,, x14154) + Vau (inhh(0444 (1331,
(61)
21

---

Published as a conference paper at ICLR 2025
F MDP CONSTRUCTION FOR DIFFUSION MODELS
Many typical inference algorithms (or samplers) of diffusion models, including but not limited to DDPM [17], DDIM [59] and SDE-DPM-solvert+ [38], can be abstracted into loop with the fol- lowing two steps (with the convention of GFlowNets on time indexing):
1. Computation of predicted clean data. -i = f(x,,)) t).
2. Back-projection of the predicted clean data. Fe41 ~ NoliT, L 0c+1l)
The MDP for diffusion models can therefore be simply constructed as:
State: (ær,t). t).
Transition: N
(0nn~~(((on1111 ((([(
1), -1),0),..)
Starting state: Zo P(xo) = NN(0, ool).
Terminal state: (.T.
Terminal reward: R(TT,T), T")
Intermediate reward: R(2;, t) 0.
Since the intermediate rewards are always zero, for terminal rewards we simply write R(xr) without the second argument of T.
With this MDP defined, is straightforward to apply policy gradient methods to optimize the return of the sampled on-policy trajectories. For instance, DDPO employs PPO [55] to perform updates.
G APPLICATION TO FLOW MATCHING MODELS
Flow matching models [31] popular and powerful class of generative models that sample points via simulation but of an ordinary differential equation (ODE) instead of an SDE. Specifically, flow matching model defines velocity field v(x, t) and, by starting from randomly initialized x(0) (O0 I). generate samples = (x1 with á = v(z, t). Since this is a deterministic process, our V-GFlowNet does not naïvely apply.
However, is shown that flow matching model above can be turned into an equivalent family of SDEs with the same probability flow, with an arbitrary diffusion term (t)) [10]:
dX,, [0(x:.))+
x(t)²
(o(X,.1)- àz (xxx) dt d(), Xo I),,
+
23
Br B.) a.
a4
(62)
This is essentially a diffusion model but with a non-linear noising process. Since we can always fix this noising process (i.e., Pg in our setting) during finetuning, we can use V-GFlowNet to obtain finetuned diffusion model from this base diffusion model.
22

---

Published as a conference paper at ICLR 2025
H FINETUNING CONVERGENCE IN WALL TIME
We further show the convergence speed measured relative wall time on single node with 80GB-mem A100 GPUs in Fig. 9.
Residual V-DB
Residual DB
DDPO
ReFL
DRaFT-LV
Drexabli Diveesiyy
0.4
2000
0.3
1500
1000
0.0
400
Relutive Wall Time (mia)
Relutive Wall Time (mis)
Relative Wall Time (miin)
Figure 9: Convergence curves of different metrics for different methods throughout the finetuning process on Aesthetic Score, with the r-axis being the relative wall time. All methods are benchmarked on single node with 8 80GB-men mem A100 GPUs.
23

---

Published as a conference paper at ICLR 2025
1
RESULTS OF ABLATION STUDY
10%
20%
30%
0,400
DeangNin Diresitty
62375
500
0.350
275
0.325
0.300
200
50
100
130
200
100
150
200
50
100
150
200
Updite Stop
Update Step
Up Step
Figure 10: Ablation study on effect of subsampling rate on the collected trajectories computing residual V-DB loss.
Ds Rewndt AwdwiiSSerr
FID Beenet Ardiei Stew
18%
28%
Figure 11: Pareto frontiers reward, diversity and pron-preservation (measured by FID) of models trained with different subsampling rate. In expectation, higher subsampling rates seem to slightly help in increasing diversity.
Residual V-DB
Residual DB
DDPO
ReFL
DRaFT-LV
Dreuadio Direstty
2000
0.3
E 1500
0.2
1000
50
200
150
200
500
Updáte Step
Uppdatc Step
Update Stop
Figure 12: Convergence curve metrics of different methods throughout the finetuning process on the HPSv2 reward model.
24

---

Published as a conference paper at ICLR 2025
Reiiaal 2.0
Rosiaal DEPO ReFt DRaF FLL
Doonddes Dhocut
DIPO
RaFL. DItAF-AN
Booan
Figure 13: Pareto frontiers for reward, diversity and (measured by FID) on HPSv2.
Residual V-DB
Residual DB
DDPO
ReFL
DRaFT-LV
Drein Divesity
8.35
700
0.30
Updote
Uptate
Updame
Figure 14: Convergence curve of metrics of different methods throughout the finetuning process on the Im- ageReward reward model.
Dosuddy Ravan, IragsBooust
FTD Ranand, Braan
Rasseal DG
Barisal
DOPO
RaFi
DRaFT1
Rawand
Rewond
Figure 15: Pareto frontiers for reward, diversity and prior-preservation (measured by FID) on IewgeReward.
With attenuated rewards
Without attermated rewards
0.40
Breanši Divesitty
2,30
8.28
0.20
Step
Figure 16: Convergence curve of metrics of different methods throughout the finetuning process on Aesthetic Score with time-depee dependeenddnt attenuation of predicted rewards. Both models trained with & 10000. With decayed predicted rewards, convergence speed is slower due less aggressive prediction on reward signal, model with reward attenuation achives hetter diversity and
results.
25

---

Published as a conference paper at ICLR 2025
Dsrrity en Kevel, Andaeicker
HD m. Rewell Ardade Seee
Rovan
Figure 17: Pareto frontiers for reward, diversity and prior-preservation (measured FID) Aesthetic Score with time-dependent scaling of predicted rewards.
Residual V-DB
DDPO
DRaFT-IV
13
0.375
2500
10
DregnSii Diveniee
0.390
2000
0.325
1000
0.300
500
0.275
100
200
100
150
200
100
150
200
Uplsse Step
Update Step
Update Step
Figure 18: Convergence of metries of different methods throughout the finetuning process Aesthetic Score with the MDP constructed by SDE DPM Solver++ (with 20 inference steps).
Dacrate en Anthdickeer
Reiidal nn nneo BRUTTL
Prode taoopaa Dieane
Roviaal Ploe DOPO
Rovari
Revax
Figure 19: Pareto frontiers for reward, diversity and prior-preservatie (measured by FID) on Aesthetic Score with the MDP constructed by Sol-DPM-Solver++ (with 20 inference steps).
26

---

Published as a conference paper at ICLR 2025
J ADDITIONAL ABLATION EXPERIMENTS
Effect of different prior strengths. It is often interest to have slightly weaker prior instead of simply increasing the reward strength, of which the sampling objective is R(Fr)* P (zr)". While it is hard to obtain exact estimate of P (IT)" where 7) (0, 1] is the prior strength, one can approximate it with the weighted score function V log P t). Denoting Fy, (2x) the corresponding flow, have the modified residual V-DB conditions.
Verti log (2( log Pf (1++1)z))= Vieei log F(x++1) Verei log F# (I++1).
V Ven log (at))))))) residual policy score function
Varas log F(2+41): residual flow score function
(63)
Va log PF (z)) -1 nx log P (1)+)
V xe log F(xx)- zt log Fof (1++1)
Vx, log Pr(zesila)): reverse residual policy score function
log F(mx): residual flow function
(64)
We experiment with choices of prior strengths and observed in Fig. 20 and 21 that lower v) lead to better diversity-rewand trade-off and faster reward convergence.
n=
n= 0.7
n= 0.3
0.3%
500
DreamSin PreamSinjPPrastt Diresity
0.390
200
Step
Uplate Step
Figure 20: Convergence curve of metrics of different methods throughout the finetuning process on Aesthetic Score with different prior strength 7j.
Daendly es Reved
n. Revend Anthels
n- 0.3
Dreasdin Dovaaa)
n=1
gm0.7
9-3.3
Rovare
Ranani
Figure 21: Pareto frontiers for reward, diversity and piior-preservation (measured FID) on Aesthetic Score with different prior strength 17.
Effeet of 2nd-order gradients in finetuning. In Fig. 22 and 23, show the comparison between models with and without 2nd-r gradients, where both models are trained with = 10000. Em- pirically, 2nd-order gradients hurts the trade-off between reward convergence, diversity preservation and prior preservation.
27

---

Published as a conference paper at ICLR 2025
2d
-rdder Grudient
0.375
450
DreunSinn Divesgy
0.390
100
200g
Figure 22: Convergence curve of metrics of different methods throughout the finetuning process on Aesthetic Score with and without 2nd-order gradients.
Daoc day Reancl Audieti
FID Reund Ardieey Seea
pnde Ordden.
Decartis Divanity
Wichedd -asle Gradex.
Wite 2dd under Gradezs
Roweri
Brain
Figure 23: Pareto frontiers for reward, diversity and crtorppreservation (measured FID) Aesthetic Score with and without 2nd-order gradients.
28

---

Published as a conference paper at ICLR 2025
K
MORE SAMPLES (AESTHETIC SCORE)
Figure 24: Additional uncurated samples from the model finetuned with residual VDDB the reward model of Acsthetic Score.
29

---

Published as a conference paper at ICLR 2025
L
MORE SAMPLES (HPSV2)
Prompi: An altm-res listic illustration of bird god swinging gold metal shick weapon, with blue man and yellow bird mouth, and intricate trddisional Chinese elements..
Residual V-DB
Residual DB
DDPO
DRaFT-LV
Prompt: female uncher elf leads group adventurers through forest crystal trees fintaay painting.
Residual V-DB
Residual DR
DDPO
DRaFT-IV
Figure 25: More comparison between samples generated by residual DD and the bascline methods. The model finetuned with residual VDDB is capable of following the instructions while generating diverse samples.
30

---

Published as a conference paper at ICLR 2025
Prompt: A surrcal cut with smill and intricate details.
Residual V-DB
Residual DB
DDPO
DRaFT-LV
Prompt: Redhend punk girl playing electric guitar no oil painting masterpiece
Residual V-DB
Residual
DRaFI-LL
Figure 26: HPSy2 samples, Continued.
31

---

Published as a conference paper at ICLR 2025
M QUALITATIVE COMPARISON ON DIVERSITY OF SAMPLES GENERATED BY DIFFERENT METHODS
Below show uncurated samples generated by models finetuned with different methods with propmts cat, bird, rabbit, bird, kangaroo.
Residual V-DB
Residual
DDPO
ReFL
DRaFT-LV
Figure 27: Uncurated samples generated by models finetuned with different methods with prompt cat.
32

---

Published as a conference paper at ICLR 2025
Residual V-DB
Residual DB
DDPO
ReFL
DRaFT-LV
Figure 28: Uncurated samples generated by models finetuned with different methods with prompt bird.
33

---

Published as a conference paper at ICLR 2025
Residual V-DB
Residual
DDPO
ReFL
DRaFT-LV
Figure 29: Uncurated samples generated by models finetuned with different methods with prompt rabbit.
34

---

Published as a conference paper at ICLR 2025
Residual V-DB
Residual DB
DDPO
ReFL
DRaFT-LV
Figure 30: Uncurated samples generated models finetuned with different methods with prompt deer.
35

---

Published as a conference paper at ICLR 2025
Residual V-DB
Residual DB
DDPO
ReFL
DRaFT-LV
Figure 31: Uncurated samples generated by models finetuned with different methods with prompt kangaroo.
36


Reviews:
[{"review_id": "puDI282thk", "reviewer": "Reviewer_ZKLE", "summary": "The paper addresses the problem of fine-tuning the pretrained diffusion models on a target reward while aiming for 1) preserving diversity of generated images and 2) fast convergence. It proposes Nabla-GFlowNet to do so, inspired by generative Flow Nets (GFlow-Nets) that sample with unnormalized density of the reward function. Experiments on different benchmarks show that the proposed method generally achieves the best diversity vs reward trade-off compared to baselines.", "questions": "- I suggest that the authors include the qualitative comparative results for the HPS-v2 reward in the paper.\n\nIf the authors address my concerns, I am willing to increase my score.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The proposed idea is based on the generative flow nets, which makes it intuitive and straightforward. \n- The Nabla-GFlowNet can leverage the first order information of the reward function (gradient) while the baselines only use the zero-order information. \n- The experimental results show that the proposed method can generally achieve the best diversity vs. reward trade-off frontiers.", "weaknesses": "- I think the \"predicted reward\" estimation in Eq. 15 can be severely unreliable, especially for the high-noise time-steps of the diffusion model. The predicted clean image will be noisy, and if the reward function is calculated by a model that has been trained on not noisy images, the predicted reward will be inaccurate.\n\n- The parameter \\lambda and the output regularization described in Page 7 seems to be crucial to the model's performance, but they are not the paper's contribution.\n\n- The qualitative samples that are compared with the baselines are only for the aesthetic score. I think a qualitative comparison on the HPSv2 can be more valuable and insightful about the model's performance.", "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": "2T7149GaTr", "reviewer": "Reviewer_BEKY", "summary": "This paper introduces an interesting method to address the challenges of fine-tuning multistep sampling in diffusion models. It employs GFlowNets to incorporate a middle term, $F(x_t)$ or $g_\\phi(x_t)$, which allows the reward score to effectively influence different timesteps.  This method successfully eliminates the need to train a reward model that accepts noisy input. This paper implements their idea in both theoretical and practical contexts. Section 3.1 covers the theoretical aspect, while sections 3.2 and 3.3 address the practical application. Experiments also show that this method can enhance reward tuning in diffusion models.", "questions": "See the weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. This paper presents a new method for addressing the challenges of fine-tuning multistep sampling in diffusion models using GFlowNets. This method effectively eliminates the need to train a reward model that processes noisy input.\n2. This paper implements their idea in both theoretical and practical contexts. Section 3.1 covers the theoretical aspect, while sections 3.2 and 3.3 address the practical application.", "weaknesses": "The main weakness is in the experiment part.\n1. The function $g_\\phi(x_t)$ is an interesting and reasonable choice for achieving the fitness task; however, it results in approximately zero vectors, with a terminal constraint of $g_\\phi(x_T) = 0$. It remains unclear whether Unet is a suitable option for this purpose.\n2. The regularization term appears significant, with $\\lambda=1000$ in the Aesthetic Score experiments and $\\lambda=100$ in the HPSv2 experiments. However, Section 3.2 states that it \"may eventually over-optimize the reward and thus neglect the pretrained prior.\" Is Section 3.2 more helpful for preventing over-optimization than the regularization term?\n3. The method is interesting and useful, but the paper's claim of \"Diversity-Preserving\" in the title creates a gap. Does this imply that the method theoretically ensures better diversity, or is it based solely on observational results? From my perspective, this method prioritizes improving information backpropagation for fine-tuning multistep sampling rather than enhancing diversity.\n\nIf the authors can resolve my issue, I will contemplate raising my score.", "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": "slTxSOAXfV", "reviewer": "Reviewer_y62f", "summary": "The authors propose Nabla-GFlowNet (∇-GFlowNet) to efficiently finetune pretrained diffusion models. This approach addresses issues of limited sample diversity and slow convergence in existing methods by leveraging reward gradients through ∇-DB and its variant, residual ∇-DB. Empirical results show that residual ∇-DB enables fast, diversity-preserving finetuning of models like StableDiffusion on various realistic reward functions.", "questions": "Please see weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The paper offers a comprehensive theoretical deduction of the proposed method, thoroughly explaining how the objectives nabla-DB and residual nabla-DB are derived. \n- By introducing residual ∇-DB, the authors extend the applicability of their work to pretrained large-scale models, which is crucial.\n- The paper enhances the quantitative evaluation of diversity in generated samples. By employing a broader range of metrics and more extensive comparisons.", "weaknesses": "- The current experimental setting appears somewhat outdated. To enhance the study's relevance, please consider using more recent schedulers and pre-trained models instead of DDPM or Stable Diffusion 1.5.\n- The qualitative results shown in Figure 2 are confusing. Additional explanation is needed to clearly demonstrate the superiority of ∇-DB, as DDPO and DAG-DB also exhibit strong performance.\n- A user study would be helpful for evaluating diversity.", "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": "xRoqJKjfG5", "reviewer": "Reviewer_5N8Z", "summary": "The authors propose a method, called $\\nabla$-GFlowNet which which is modification of GFLowNets. They define a new objective $\\nabla$-DB, which is a gradient informed version the Detailed Balance objective. The paper then goes on to propose a residual version of this loss which at optimality samples proportionally to the argumented distribution $r(x_T)p^\\sharp (x_T)$, which maintains diversity of generations. The paper then presents experiments which shows a good pareto frontier of reward vs diversity on two tasks.", "questions": "- It is known that optimizing the KL constrained optimization problem, with the closed form solution of the augmented distribution, can lead to a biased result [0, 1]. Does this method have this bias problem? If not, how do you get around it?\n\n[0] https://arxiv.org/abs/2409.08861\n[1] https://arxiv.org/abs/2402.15194", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- This paper offers a good way at maintaining diversity in generations while aligning generations to a reward model.\n- The paper is theoretically founded and shows that their new training objective maintains the validity of GFlowNets while taking into account gradient information.\n- The work has a number of good ablations and experiments to show diversity and reward tradeoff.", "weaknesses": "- The problem of bias is well known in optimizations for diffusion models (elaborated in question section). Some treatment of this problem problem would be desirable since it seems that this method may be optimizing for the biased distribution", "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": "fNr2sAY9Bo", "reviewer": "Reviewer_5qSt", "summary": "This paper introduces Nabla-GFlowNet with a new objective, $ \\nabla$-DB, and its variant, residual $ \\nabla$-DB, for finetuning pretrained diffusion models. The method aims to enhance sample diversity and finetuning efficiency by utilizing reward gradients. Experiments on two reward functions, Aesthetic Score and Human Preference Score (HPSv2), demonstrate improved performance.", "questions": "Please see weaknesses.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The paper has a clear motivation, addressing the challenge of preserving diversity and improving efficiency in finetuning diffusion models, which is crucial for real-world applications.\n- The paper introduces a unique application of GFlowNet principles to diffusion model finetuning, specifically focusing on the reward gradient to preserve sample diversity.", "weaknesses": "- In the quantitative evaluation, the proposed method performs well on the smaller dataset of Aesthetic Score, nearly doubling the baseline in the DreamSim metric. However, on the larger dataset of HPSv2, its performance is similar to baselines, showing no clear advantage. The effectiveness of the method remains inconclusive due to the differences in dataset size.\n- While the authors mention “fast and efficient finetuning” as a contribution, only Figure 4 shows comparable convergence speed to the baseline. It would be helpful to include details on training resource consumption, such as GPU usage and computational cost, to substantiate this claim.\n- Figure 3 only compares the results of the pretrained model and the proposed method, lacking visual comparisons with other baselines.\n- Figure 5 lacks sufficient information in the title and annotations to clarify what each point represents (e.g., training iteration). To improve clarity, consider adding an explanation in the figure caption or in Section 4.4 to specify the meaning of each point.\n- The conclusion section is missing, which may limit the clarity of the paper’s overall findings and contributions.", "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)