Improving GFlowNets for Text-to-Image Diffusion Alignment
Authors: Dinghuai Zhang
Yizhe Zhang
Jiatao Gu
Ruixiang ZHANG
Joshua M. Susskind
Navdeep Jaitly
Shuangfei Zhai
Venue: Accepted by TMLR ()

Abstract: Diffusion models have become the de-facto approach for generating visual data, which are trained to match the distribution of the training dataset. In addition, we also want to control generation to fulfill desired properties such as alignment to a text description, which can be specified with a black-box reward function.
Prior works fine-tune pretrained diffusion models to achieve this goal through reinforcement learning-based algorithms. Nonetheless, they suffer from issues including slow credit assignment as well as low quality in their generated samples.
In this work, we explore techniques that do not directly maximize the reward but rather generate high-reward images with relatively high probability --- a natural scenario for the framework of generative flow networks (GFlowNets). 
To this end, we propose the Diffusion Alignment with GFlowNet (DAG) algorithm to post-train diffusion models with black-box property functions. 
Extensive experiments on Stable Diffusion and various reward specifications corroborate that our method could effectively align large-scale text-to-image diffusion models with given reward information.

Published in Transactions on Machine Learning Research (12/2024)
Improving GFlowNets for Text-to-lmage Diffusion Alignment
Dinghuai Zhang; Yizhe Zhang, Jiatao Gu, Ruixiang Zhang, Josh Susskind, Navdeep Jaitly, Shuangfei Zhai Apple
Reviewed on OpenReview: https: // openreviee. net/foruu forum? id- XDxxDby330 XDbM3g444
Abstract
Diffusion models have become the de-facto approach for generating visual data, which are trained to match the distribution of the training dataset. In addition, also want to control generation to fulfill desired properties such as aligmment to text description, which can be specified with black-box reward function. Prior works fine-tune pretrained diffussio models to achieve this goal through reinforcement learning- based algorithms. Nonetheless, they suffer from issues including slow credit assignment as well as low quality in their generated samples. In this work, we explore techniques that do not directly maximize the reward but rather generate high-reward images with relatively high probability natural scenario for the framework of generative flow networks (GFlowNets). To this end, propose the Diffusion Aligmment with GFlowNet (DAG) algorithm to post-train diffusion models with black-box property functions. Extensive experiments on Stable Diffusion and various reward specifications corroborate that our method could effectively align large-scale text-to-image diffusion models with given reward information. Our code public available at https://githhb..co hww.tthhu...oomaaaa http:///ttiibbo/appppp/epep happla/ml-ddffuuuttttaee aalddiffisoon-alimmmntt
1 Introduction
Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020) have drawn significant attention in ma- chine lcarning due to their impressive capability to generate high-quality visual data and applicability across a diverse range of domains, including text-to- image synthesis (Rombach et al., 2021), 3D genera- tion (Poole et al., 2022), material design (Yang et al., 2023), protein conformation modeling (Abramson et al., 2024), and contimuus control (Janner et al., 2022). These models, through a process of gradually denoising a random distribution, learn to replicate complex data distributions, showcasing their robust- ness and flexibility. The traditional training of diffu-
Figure 1: Generated samples before (top) and after (botlom) the proposed training with Aesthetic reward. sion models typically relies large datasets, from which the models learn to generate new samples that mimic and interpolate the observed examples.
However, such a dataset-dependenn approach often overlooks the opportunity to control and direct the generation process towards outputs that not only resemble the training data but also possess specific, desirable properties (Lee et al., 2023). These properties are often defined through explicit reward functions that assess certain properties, such as the aesthetic quality of images. Such requirement is crucial in fields where adherence to particular characteristics is necessary, such as alignment or drug discovery. The need integrate explicit guidance without relying solely on datascts presents a unique challenge for training mothodologics.
Work done during internship Apple MLR.. Correspondence

---

Published in Transactions on Machine Learning Research (12/2024)
Previous works have utilized methods such as reinforcement learning (RL) (Black et al., 2023; Fan et al., 2023) to tackle this problem. Nonetheless, these methods still sulfer from issues like low sample efficiency.
In this work, we propose a novel approach, diffusion aligument with GFlowNets (DAG), that fine-tunes diffusion models to optimize black-box reward functions directly. Generative flow networks (Bengio et al., 2023, GFlowNets), initially introduced for efficient probabilistic inference with given densities in structured spaces, provide unique framework for this task. Though initially proposed for composite graph-like structures, prior works have extended the GFlowNet framework to diffusion modeling (Zhang et al., 2022a; Lahlou et al., 2023). This work further investigates GFlow:Net-inspired algorithms for the task of text-to-image diffusion alignment. By aligning the learning process to focus on generating samples with probability proportional to reward functions rather than maximizing them, our method allows the diffusion model to directly target and generate samples that are not only high in quality but also fulfill specific predefined criteria. Besides developing a denoising diffusion probabilistic model-specific GFlowNett algorithm, we also propose new KL-based way to optimize our models. In summary, our contributions are follows:
We propose Diffusion Aligmment with GFlowNet (DAG), a GFlowNet-based algorithm using the denoising structure of diffusion modols, to improve largo-scale text-to-image alignment with a black-box reward function.
We propose a KL-based objective for optimizing GFlowNets that achieves comparable or better sample efficiency. We further called the resulting algorithm for the aligmment problem DAG-KL.
Our methods achieve better sample efficiency than the reinforcement learning baseline within the same number of trajectory rollouts, as well as a better reward-diversity trade-off, across a number of different learning largets.
2 Preliminaries
2.1
Diffusion models
Denoising diffusion model (Vincent, 2011; Sohl- Dickstickttein et al,, 2015; Ho et al., 2020; Song et al., 2020) is a class of hierarchical latent variable models. The latent variables are initialized from a white noise xy ~ ~N001 I) and then go through a sequential denoising (reverse) process po (x(x_1(x)). Therefore, the resulting generated distribution takes the form of
= - po (x(xxa)) =
/ p((x)) II po (x-- ( dx
T
Pa (x()))
On the other hand, the variational posterior q(x1r(X)) also called a diffusion or forward process, can be factorized as a Markov chain 4(x,
composed by a series of conditional Gaussian distributions g(x/|x;-1] 1)= = (xxx;; 1e/a--1x-111 (1- a?/a? ///211)) I], where {ax, ot is a set of pre-defined signal-noise schedule. Specifically, in Ho et al. (2020) we have ot + t 1. The benefit of such a noising process is that its marginal has a simple close form: (xx/(xo) = qxxxx dx = (xxx;; XxXo, II).
Given a data distribution Pdaca (-), the variational lower bound of model log likelihood can be written in the following simple denoising objective:
Lecnoising (0) E. 1x [[|x] (o ote, f)][3]
(2)
where (x(x;,() is deep neural network to predict the original clean data Xo given the noisy in- put X; 09Xo + oie, which N(x;-1 (07_ ax + (aî_ xo t)) (ofor-1, of/a?_ I). In practice, the network can also be parameterized with noise prediction or v-prediction (Salimans & Ho, 2022). The network architecture usually has a U-Net (Ronneberger et al., 2015) structure.
be used to parameterize the denoising process po(x(-1(x;)
In multimodal applications such as text-to-image tasks, the denoising diffusion model would have conditioning e in the sense of po (xx: c) p(xx) IT= po p c) dx The data prediction network, (* t, c) in this case, will also take conditioning input. We ignore the notation of e. without loss of generality.
2

---

Published in Transactions on Machine Learning Research (12/2024)
2.2 GFlowNets
Generative flow network (Bengio et al., 2021, GFlowNet) is high-level algorithmic framework of amortized inference, also known as training generative models with given unnormalized target density function. Let g = (S) be a directed acyclic graph, where S is the set of states and A C S x S are the set of actions. We assume the environmental transition is deterministic, i.e., one action would only lead to one next state. There is a unique initial state So E S which has no incoming edges and a set of terminal states Sy without. outgoing edges. A GFlowNet has a stochastic forward policy P'p(s'|s ()) for transition (s+5) s') as a conditional distribution over the children of given state s, which can be used to induce distribution over trajectories via P(r)= = I Pp(8n+1|s,) where = (8o, $1,...,8x) On the other hand, the buckward policy (s(s|s') is a distribution over the parents of a given state s'. The terminating distribution defined by Pr(x) = Er-xx (r(T) is the ultimate terminal state distribution generated by the GFlowNet. The goal of training GF is to obtain a forward policy such that Pr(() x R(.), where R(-) is a black-box reward function or unnormalized density that takes only non-negative values. Notice that we do not know the normalizing factor Z = Exx R(x). We can use the trajectory flow function F(T) (P) to take in the effect of the normalizing factor, and the corresponding state flow function F(s) Ers F(T) to model the unnormalized probability flow of intermediate state
Detailed balance (DB) The GFlowNet detailed balance condition provides a way to learn the above mentioned GFlowNet modules. For any single transition (s + s'), the following DB criterion holds:
F(s) (s) = ((s)))())s' V(s + s's A.
(3)
Furthermore, for any terminating state x, we require F(x) = R(x). In practice these constraints can be transformed into tractable training objectives, will be shown in Section 3. Based GFlowNet theories in Bengio et al. (2023), if the DB criterion is satisfied for any transition, then the terminating distribution Pr(-) will be the same desired target distribution whose density is proportional to R(-).
3 Methodology
3.1 Denoising Markov decision process
The denoising process for text-to-image diffusion models can be easily reformulated as a multi-step Markov decision process (MDP) with finite horizon (Fan et al., 2023; Black et al., 2023) as follows:
s: (x) p(so) = (xx:: 0,1) @ p(c). To (s)) = po (x (r-cc), a. XT- r(sr, ax) = R($++1, only if i-== T- 1, p(s++1 as) = on, @ de.
(4)
(5)
Here St, a is the state and action at time step l under the context of MDP. The state space is defined to be the product space (denoted by @) of x in reverse time ordering and conditional prompt c. The RL policy x is just the denoising conditional distribution. In this MDP, when time t has not reached the terminal step, we define the reward r(s,, ax) to be O. 5 here denotes the Dirac distribution.
Remark 1 (diffusion model as GFlowNet).. This formulation has a direct connection to the GFlowNet MDP definition in Section 2.2, which has been pointed out by Zhang et al. (2022a) and developed Lahlou et al. (2023); Zhang et al. (2023b); Venkatraman* et al. (2024). To be specific, the action transition ar) + 1 is a Dirac distribution and can be directly linked with the (s, + Sc+1) edge transition in the GFlowNet language. More importantly, the conditional distribution of the denoising process pe -xxx----(x---) corresponds to the GFlowNet forward policy Pr (s:), while the conditional distribution of the diffusion process (x---)))) corresponds to the GelowNet backward policy Pn (8+4)). Besides, X₁ is a GFlowNet terminal state if and only if = O.
Denoising diffusson
GFlowNet
The above discussion could be summarized in the right table. In the fol- lowing text, we use the denoising diffusion notation instcad of GFlowNet notation as is familiar to more broad audience. What's more, we ignore conditioning for the sake of simplicity.
(xT- c) p(xr-1- 1- (x--1,) -(x--1|xr--1)
Pp(s141|5)) Po(si(s(11)
3

---

Published in Transactions on Machine Learning Research (12/2024)
3.2 Diffusion alignment with GFlowNets
In this section, we describe our proposed algorithm, diffusion aligument with GelowNets (DAG). Rather than directly optimizing the reward targets as in RL, we aim to train the generative models so that in the end they could generate objects with a probability proportional to the reward function: po x R(xo). To achieve this, we construct the following DB-based training objective based on Equation 3, by regressing its one side to another in the logarithm scale for any diffusion step transition (xx:4(-1)
(DD((,, xI-) (log (xx, 1) + log - log Fx 1) log ((x
(6)
We additionally force Fx 1 0) = R(xo) to introduce the reward signal. Here 0, are the parameters of the diffusion U-Net model and the GFlowNett state flow function (which is another neural network), respectively. One can prove that if the optimization is perfect, the resulting model will generate a distribution whose density value is proportional to the reward function R(-) (Bengio et al., 2023; Zhang et al,, 2023b).
One way to parameterize the state flow function F is through the so-called forward-looking (Pan al., 2023b, FL) tech- nique in the way of (x t) = (x,, R(x)), where Fg is the actual neural network to be Iearned. Intuitively, this is equivalent to initialixing the state flow function to be the reward function in a functional way; therefore, learning of the state flow would become an casicr task. Note that to ensure F (xo: 0) = R(xo), we need to force (x = for all Xo at the terminal step.
Algorithm 1 GFlowNets (DAG-DB & DAG-KL)
Diffusion alignment with
Require: Denoising policy po(xx-1(xx,tt) nois- ing policy q(x((x(11) flow function
(x,,,,), black-box reward function R(-)
1: repeat
2: Rollout -= {xi]= with pe(x
3: For each transition (x,,x- x,-) 1) E r:
4: if algorithm is DAG-DB then
Incorporating denoising diffusion-speciffe structure However, the intermediate state x, is noisy under our context, and thus not appropriate for being evaluated by the given reward function, which would give noisy result. What's
normal DB-based update 6: Update O and & with Equation &
5:
7:
clsc if algorithm is DAG-KL then
8:
# KL-based update
more, what we are interested here is to "foresee" the reward of the terminal state Xo taken from the (partial) trajectory Xe.0 starting from given X2. As a result, we can do the FL technique utilizing the particular structure of diffusion model as in F& (xx, t) (x (Sx (Xe, t)), where Xo is the data prediction network. We notice that a similar technique has been used to improve classifier guidance (Bansal et al., 2023). In short, our innovation in FL technique is
9:
Update & with Equation 8 Update with Equation 14 end if
10:
11:
12: until some convergence condition
Fx 1) ()xx)) = Fx 1) Fo (( I)).
Then the FL-DB training objective becomes
() (x) EFL (xL, Xx-1) = log Fo 1) (x,-1,t- 1))g(x((x111)
(8)
Since in this work the reward function is a black-box, the gradient flow would not go through (x f) when we take the gradient of f. We summarize the algorithm in Algorithm and refer to it as DAG-DB.
Remark 2 (GPU memory and the choice of GFlowNet objectives). Similar to the temporal difference-A RL (Sutton, 1988), it is possible to use multiple connected transition steps rather than a single transition step to construct the learning objective. Other GFlowNet objectives such as Malkin et al. (2022); Madan et al. (2022) use partial trajectories with series of transition steps to construct the training loss and provide a different tradc-off between variance and bias in credit assignment. However, for large-scale sctups, this is not easy to implement, as computing policy probabilities for multiple transitions would correspondingly increase the GPU memory and computation multiple times. For example, in the Stable Diffusion setting, we could only use a batch size of 8 on each GPU for single transition computation. If we want to use a two transition based training loss, we would need to docrease the batch size by half to 4. Similarly, we will have to shorten the trajectory length by large margin if we want to use trajectory balance. This may influence the image
4

---

Published in Transactions on Machine Learning Research (12/2024)
generation quality and also make it tricky to compare with the RL baseline, which can be implemented with single transitions and does not need to decrease batch size or increase gradient accumulation. In practice, we find that single transition algorithms (such as our RL baseline) perform reasonably well.
Remark 3. In practice it makes sense to pursue a "tilt" target (Xo) x (pretrained(X) (Xx))((X))) R(xx)) to ensure that generated images are reasonable. This can be casily achioved within our framework by simply replacing the ((x/(xx-1) term with Ppretrained (x,-1 xx) in Equation We refer to Liu et al. (2024) for further details.
3.3 A KL-based GFlowNet algorithm with REINFORCE gradient
GFlowNet detailed balance is an off- ofolcol algorithm that uses training data from arbitrary distributions. In this section, we derive a different KL.-based on-policy objective, which has been rarcly investigated GFlowNet literature. We can reformulate DB (Equation 6) from square loss form a KL divergence form
min DKL. (xx(xx-1xx)) F (x: 1, (x((x(/x--1))
(x(x),1)
In theory, when DB is perfectly satisfied, the right term (x - ()x((x (x)))//>///) t) is a normalized density; in practice, it could be an unnormalized one but does not affect the optimization. Next, define
(x1,xx111 = stop-gradient (10g (x (0)0(xxx11 xx)
(10)
Fx 1)9((1 (x(x)
then the KL value of Equation becomes Epo (x (x)) [b(x,, x1- We have the following result for deriving a practical REINFORCE-style objective.
Proposition 4. The KL term in Equation 9 has the same expected gradient with l(x,, x,- 1) log po (
VODIL (20(x)-1x00] 1x 1)9(x)1x(-1))
(11)
Fx 8)
We defer its proof Section A. and make the following remarks:
Remark (gradient equivalence to detailed balance). Recalling Equation 6, since we have VOLDB (x,, xt-1) b(x,, Ve log po (xxx1xxxx) [x.), it is clear that this KL-based objective would lead to the same expected gradient 9 with Equation 6, if x,-1 Po(-(x.) (i.e., samples being on-policy). Nonetheless, this on-policy property may not be true in practice since the current model is usually not the same as the model used for rollout trajectorics after a few optimization steps.
Remark 6 (analysis of b(x,, x()))). The term b(x,, x1- scems to serve as the traditional "reward" in the RL framework. Previous works (Tiapkin et al., 2024; Mohammadpour et al., 2024; Deleu et al., 2024) have shown that GFlowNet could be interpreted as solving a maximum entropy RL problem in a modified MDP, have an extra non-zero reward that equals the logarithm
where any intermodiate transition modified of GFlowNet backward policy: Tmod (X,, (x,,xt xt-1) = log qxx//xx1),+ > 0. What's more, the logarithm of the transition flow function (i..,, either side of the detailed balance constraint of Equation 3) could be interpreted
the Q-function of this maximun entropy RL problem the modified MDP. Therefore, we could take more in-depth analysis of -bxx,,x,-1)
log (x t- ()((x((x-11) + (- log Po (x/-1/x;)) log F (x),
(12)
mudified Q- Q-uncction
entropy
constant
The last term tx) is constant for the x,-1 ~ Pe pe(-Ixx (-1x,) process.
Consequently, we can see that min- imizing this KL is effectively equivalent to having a REINFORCE gradient with the modified Q-function plus an endropy regularization. (Haarnoja et al., 2018). What's more, the entropy of forward policy (x.-1-po((1x.) log pe [x=-||x)]] is actually also a constant (the diffusion model's noise schedule is fixed), which does not affect the learning.
Note that this REINFORCE style objective in Equation is on-policy; the data has come from the same distribution as the current model. In practice, the model would become not exactly on-policy after few
5

---

Published in Transactions on Machine Learning Research (12/2024)
Figure 2: Top: samples from the original Stable Diffusion model. Middle: the proposed method trained with compressibility reward; these images have very smooth texture. Down: the proposed method trained with incompressibility reward; the texture part of images contains high frequency noise.
optimization steps, under which scenario we need introduce the probability ratio po (xxx-11 P via importance sampling:
Po (x, (x,)
Vepo (xc-1 (x))
E Po (12x))
b(x,, x(-1)
(13)
Paol (x,- X2,
Therefore, we can define a new objective to be
po xx) Peai (X
(xL, X1-1 b(xx, x1-1) clip
e, (1+))
(14)
where Xt-1 a (1)))) Here we also introduce clip operation to remove too drastic update, following PPO (Schulman et al., 2017). In this way, the overall gradient along trajectory becomes
po(xx- (xt) Po.m (x- 6, (,11 11,,))))))__
-61+0) 1
T Eepaas (Xa:r))
Σ b(x X clip
(15)
We use this to update the policy parameter and use FL-DB to only update Φ. We call this "diffusion alignment with GFlowNet and REINFORCE gradient" mothod to be DAG-KL. Note that when calculating b(x,, xt-1), we also adopt the diffusion-speeifie FL technique developed in Section 3.2. We also put the algorithmic pipeline of DAG-KT, in Algorithm 1.
4 Related Works
Diffusion alignment People have been modeling human values a reward function areas such as game (Ibarz et al., 2018) and language modeling (Bai al., 2022) make the model more aligned. In diffusion models, early researchers used various kinds of guidance (Dhariwal & Nichol, 2021; Ho & Salimans, 2022; Kong et al., 2024) to achieve the goal of steerable generation under the reward. This approach is as simple plug-and-play but requires querying the reward function during inference time. Another way is to post-train the model to incorporate the information from the reward function, which has a different. setup from guidance methods; there is also work showing that this outperforms guidance methods (Uehara et al., 2024). Lee al. (2023); Dong al. (2023) achieve this through maximum likelihood estimation on model-gane godellgenerated samples, which are reweighted by the reward function. These works could be thought of doing RT in one-step MDPs. Black et al. (2023); Fan et al. (2023) design RI. algorithm by taking the

---

Published in Transactions on Machine Learning Research (12/2024)
Aesthetic
ImageReword
HP5
Nathod
Method
0.31
Nethed
7.5
DAG D8
DAG KL
0.6
DAG OE
DAG KL
0.35
D4G KL
D4G 08
DUPC
0.G
DOPO
0.34
DUPO
1.0
Reward
6.5
0.33
0.2
0.28
6.0
0.0
0.25
55
0.24
100
150
50
150
200
100
150
Training stop
Training stop
Training stap
Figure 3: Sample efficiency results of our proposed methods and our RL baseline (DDPO). The number of training steps is proportional to the number of sampled trajectories. The experiments are conducted on reward functions including aesthelic score, ImageReward, and HPSv2.
diffusion generation process as a MDP (Section 3.1). In this work, we focus on black-box rewards where it is appropriate to use RL or GelowNet methods. Furthermore, there are methods developed specifically for differentiable rewards setting (Clark et al., 2023; Wallace et al., 2023b; Prabhudesai el al., 2023; Wu et al., 2024; Xu et al., 2023; Uehara et al., 2024; Marion al., 2024). Besides, Chen al. (2023) study the effect of finetuning text encoder rather than diffusion U-Net. There is also work that relies on preference data rather than an explicit reward function ( (Wallace el al., 2023a; Yang et al., 2024). Kim et al. (2024a) investigate how to obtain a robust reward based on multiple different reward functions.
is a family of generalized variational inference algorithms that treats the data sampling process as a sequential decision- maaiinn one. It is useful for generating diverse and high-quality samples in structured scientific domains (Jain et al., 2022; 2023b; Liu et al., 2022; Jain et al., 2023a; Shen et al., 2023; Zhang al., 2023e; Pan et al., 2023a; Kim et al., 2023; 2024b). A series of works have studied the connection between GelowNets and probabilistic modeling methods (Zhang et al., 2022b; Zimmermann et al., 2022; Malkin et al., 2023; Zhang et al., 2022a; Ma et al.; Zhang et al., 2024a), and between Gelow and control methods (Pan et al., 2023c;d;b; Zhang et al., 2024b; Tiapkin et al., 2024). GFlowNets also have wide application in causal discovery Deleu et al. (2022), phylogenetic inference (Zhou et al., 2024), and combinatorial optimization (Zhang et al., 2023a;d). A concurrent work (Venkatraman* et al., 2024) also studies GFlowNet on diffusion aligument which is similar to this work but has different scope and different developed algorithm. Specifically, this work is aiming for posterior approximate inference that the reward function is treated as likelihood information, and develops trajectory balance (Malkin al., 2022) based algorithm on length modified trajectories.
GFlowNet GetowNets
5 Experiments
Compresabilty
Irooppesstiity
Experimental setups We choose Stable Dif- fusion v1.5 (Rombach et al., 2021) our base generative model. For training, we use low- rank adaptation (Hu et al., 2021, LoRA) for parameter efficient computation. As for the reward functions, we do experiments with the LAION Aesthetics predictor, a neural aesthetic score trained from human feedback to give input image aesthetic rating. For text- image alignment rewards, we choose ImageRe- ward (Xu et al., 2023) and human preference score (HPSv2) (Wu et al., 2023). They are both CLIP (Radford et al., 2021)-type models, taking a text- ixtt-image pair as input and output a scalar score about to what extent the image follows the text description. We also test with the (in)coprrssiiblityy reward, which computes the file size if the input image is stored in hardware storage. As for the prompt distribution, we use a set of 45 simple animal prompts from Black et al.
Method DAD
Mehael
DOPO
130
Triniing stop
Sample efficiency results of our proposed methods and our RL baseline (DDPO) on learning from compressibil- ity and incompresiibility rewards.
Figure 4:
7

---

Published in Transactions on Machine Learning Research (12/2024)
A counter top with food sitting on some towels
Personal computer desk room with large glass double doors
A bathroom has toilet and scale
Several cars drive down the road on cloudy day
Stable Diffusion
DDPO
DAG-DB
DAG-KL
Figure 5: Text-image alignment results. We display four prompts and the corresponding generation visualiza- tion from the original Stable Diffusion (1st row), DDPO (2nd row), DAG-DB (3rd row), and DAG-KL (4th row) models to compare their alignment abilities. See Figure for more results.
(2023) for the Aesthetics lask; we use the whole imagenet classes for the (in)compressibiliiy task; we use the DrawBench (Saharia et al., 2022) prompt set for the ImageReward task; we use the photo and painting prompts from the human preference dataset (HPDv2) (Wu et al., 2023) for the HPSv2 task. We notice that prompt set containing hundreds of prompts which is more than some previous
in our experiments, work such as Black et al. (2023).
Effectiveness of the proposed methods We first demonstrate that proposed methods could generated images that have meaningful improvements corresponding to the rewards being used. In Figure 1, compare the images from the original Stable Diffusion pretrained model and our proposed method. After our post- training, the gencrated images become more vibrant and vivid; we also notice that these images have slightly higher saturation, which we believe aligned with the human preference on good-looking pictures. We also visualize the experiment results on compressibility and incompresiibility tasks Figure The second row shows the generated images from the model trained with the compressibility reward, which have low details and smooth textures, and also have very limited colors. On the other hand, the model trained with incompressibility reward would generate images with high frequency texture, shown in the third row. These results indicate that our method could effectively incorporate the reward characteristies into the generative models. We defer more experimental details to Section B.2.
Algorithmic comparisons The main bascline we compare with is denoising diffusion policy optimiza- tion (Black et al., 2023, DDPO), an RL algorithu that is specifically designed for denoising diffusion alignment. and has been shown to outperform other align-from caan--foo--ll---e blace-box-reward methods including (Lee et al., 2023; Fan et al., 2023). We show the reward curves w.r.t. the training steps of the acsthetic, ImageReward, and HPSv2 rewards in Figure 3. Here, the number of training steps corresponds proportionally the number of trajectories collected (see appendix for more details). Both our proposed methods, DAG-DB and

---

Published in Transactions on Machine Learning Research (12/2024)
Task
Compressibility
Incompressibility
Aesthetic
ImageReward
HPSv2
Metric
Rewardf
Div.t
Rewardf
Div.t
Rewardf
Div.f
Rewardf
Div.f
Reward?
Div.f
DDPO
-44.74
27.32
197.83
25.72
6,33
13.59
0.29
26,04
0.29
19.49
DAG-DB
337.23
38.94
294.68
35.25
6.63
14.09
0.51
44.91
0.30
17.95
DAG-KL
-34.67
41.30
218.93
27.68
6.50
13.63
0.48
30.37
0.30
21.85
Table 1: Comparison on average reward and diversity metrics across a variety of tasks. Our methods consistently achieve better trade-offs between these two directions. See Section and Section B.2 for details.
"A helmet-wearing monkey skating"
DDPO samples
your Saaa
"Anthropomooppic Virginia opossum playing guilar"
DDPO samples
Figure Visualization of alignment with regard to training progress. Left: the generated images from the proposed method become more aligned to the text prompt over the course of training. Right: samples from the DDPO baseline.
DAG-KL, achieve faster credit assignment than the DDPO bascline by a large margin. We additionally put corresponding curve plots for compressibility and incompressibility rewards in Figure 4, which also demonstrate the advantages of our proposed methods.
What's 's more, provide a diversily comparison in Table 1. For the RL baseline, we lake the last checkpoint; for our proposed methods, we take the earliest checkpoint with a larger reward value than the chosen RL algorithm checkpoint. This is to show that our mothods could achieve results with both better reward and diversity performance. For diversity measurement, we calculate the FID score between two batches of independently generated samples from that model. We use this as a diversity metric, so the larger the better. The results indicate that Netsweee- methods achieve better reward-diversity trade-off due to their distribution matching (rather than reward maximizing) formulation. We defer related details to Scction B.2.
Apart from quantitative comparisons, we also visualize the alignment improvement for models trained in the HPSv2 task. In Figure 5 and Figure 10 in Appendix, we exhibit generation results for different prompts across the original Stable Diffusion, DDPO, DAG-DB, and DAG-KL models. For example, in the first "a counter top with food sitting on some towels" example, images from the original Stable Diffusion either do not have food the food not on towels, which also the case for DDPO generation. This is improved for both DAG-DB and DAG-KL generation in that they capture the location relationship correctly. In the "personal computer desk room with large glass double doors" example, both the original and DDPO models cannot generate any double doors in the image, and DAG-DB model sometimes also fails. In contrast, the DAG-KL model seems to understand the concept well. Generation with other prompts also has similar results.
In Figure 6, we visualize the gradual alignment improvement of our DAG-KL method with regard to the training progress for the HPSv2 task. We show the images of our methods at 0%, 25%. 50%, 75%, and 100% training progress. the example of "a helmet- wearing monkey skating", the DDPO bascline could generate a skating monkey but seems to fail to generate helmet. For the proposed med the model gradually learns
9

---

Published in Transactions on Machine Learning Research (12/2024)
to handle the concept of helmet over the course of training. In the "anthropomorphic Virginia opossum playing guitar" example, the baseline understands the concept of guitar well, but the generated images are not anthropomorphic, while our method manages to generate anthropomorphic opossumss decently.
6 Conclusion
We propose Diffusion Aligument GFlowNet (DAG), a family of algorithms designed to fine-tune pretrained diffusion models based on external reward functions. Our approach advances the GFlowNet framework accommodate the specific properties of text-to-image diffusion models. We introduce two variants, DAG-DB and DAG-KL, each tailored to optimize the model's performance with respect to different objectives. Through extensive experiments on Stable Diffusion models, we demonstrate that DAG achieves more effective balance between maximizing reward values and maintaining output diversity than previous methods.
References
Josh Abramson, Jonas Adler, Jack Dunger, Richard Evans, Tim Green, Alexander Pritzel, Olaf Ronneberger, Lindsay Willmore, Andrew Ballard, Joshua Bambrick, et al. Accurate structure prediction biomolecular interactions with alphafold 3. Nature, pp. 1-3, 2024.
Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, Nicholas Joseph, Saurav Kadavath, Jackson Kernion, Tom Conerly, Sheer El-Showk, Nclson Elhage, Zac Hatficld-Doddd Danny Hernandez, Tristan Hume, Scott Johnston, Shauna Kravec, Liane Lovitt, Neel Nanda, Catherine Olsson, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, Ben Mann, and Jared Kaplan. Training a helpful and harmless assistant with reinforcement learning from human feedback, 2022.
Arpit Bansal, Mlog-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Universal guidance for diffusion models. 2023 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops (CVPRW), pp. 843-852, 2023.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow network based generative models for non-iterative diverse candidate generation. Newral Information Processing Systems (NeurIPS), 2021.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward Hu, Mo Tiwari, and Emmanuel Bengio. GelowNet. foundations. Journal of Machine Learning Research, (24):1-76, 2023.
Kevin Black, Michael Janner, Yilun Du, 1lya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. Arxiv, abs/2335.13301, 2023.
Chaofeng Chen, Annan Wang, Haoning Wu, Liang Liao, Wenxiu Sun, Qiong Yan, and Weisi Lin. Enhancing diffusion models with text-encoder reinforcement learning. ArXiv, abs/2311.15657 2023. URL https:
org/Corpusii 265457291.
Kevin Clark, Paul Vicol, Kevin Swersky, and David Fleet. Directly fine-tuning diffusion models differentiable rewards, 2023.
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. Uncertainly in Artificial Intelligence (UAT), 2022.
Tristan Delen, Padideh Nouri, Nikolay Malkin, Doina Precup, and Yoshua Bengio. Discrcte probabilistic inference as control in multi-path environments. ar preprint arXiv:2402. 10309, 2024.
Prafulla Dhariwal and Alex Nichol. Diffusion models beat gans image synthesis, 2021.
10

---

Published in Transactions on Machine Learning Research (12/2024)
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, 2023.
Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, P. Abbeel, Mohammad Ghavanzadch, Kangwook Lee, and Kimin Lec. Dpok: Reinforcement learning for fine-tuning text-to-image diffusion models. ArXie, 0ts/2305.16386 2023.
Roy Fox, Ari Pakman, and Naftali Tishby. Taming the noise in reinforcement learning via soft updates, 2017.
Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. International Conference on Machine Learning (ICML), 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. International Conference on Machine Learning (ICML), 2018.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance, 2022.
Jonathan Ho, Ajay Jain, and P. Abbeel. Denoising diffusion probabilistic models. ArXiv, abs/ 2////000011222 2020. URL https://api... http::///p...emattiieeee semammiiccanttiihcaaatt..t.... ottprattttt....ooooorrraaaaaaaaa htg//oorpusID:219955566
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.
Borja Ibarz, Jan Leike, Tobias Pohlen, Geoffrey Irving, Shane Legg, and Dario Amodei. Reward learning from human preferences and demonstrations in atari, 2018.
Moksh Jain, Emmanuel Bengio, Alex Hernandoz-Garcia, Jarrid Rector-Brooks, Bonaventure F.P. Dossou, Chanakya Ekbote, Jie Fu, Tianyu Zhang, Micheal Kilgour, Dinghuai Zhang, Lena Simine, Payel Das, and Yoshua Bengio. Biological sequence design with GFlowNets. International Conference on Machine Learning (ICML), 2022.
Moksh Jain, Tristan Deleu, Jason S. Hartford, Cheng-Hao Liu, Alex Garenn, and Yoshua Bengio. Gflownets for ai-driven scientific discovery. ArXiv, abs/2302.00615, 2023a. URL https: seaticssmatticcchoolar www.hhoorrrrrroooooooooooo orgo CoopppsSSD:266458311
Moksh Jain, Sharath Chandra Raparthy, Alex Kernandez-Garcia, Jarrid Rector-Brooks, Yoshua Bengio, Santiago Miret, and Emmanuel Bengio. Multi-objective GFlowNets. International Conference on Machine Learning (ICML), 2023b.
Michael Janner, Yihun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flcxible behavior synthesis. arXiv preprint vrXic:2205.09900 2022.
Kyuyoung Kim, Jongheon Jeong, Minyong An, Mohammad Ghavamzadeh, Krishnamurthy Dvijotham, Jinwoo Shin, and Kimin Lee. Confidence-aware reward optimization for fine-tuning text-to-image models, 2024a.
Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, and Jinkyoo Park. Local search gflownets. ArXio, abs/2310.02710, 2023.
Minsu Kim, Joohwan Ko, Tacyoung Yun, Dinghuai Zhang, Ling Pan, Woochang Kim, Jinkyoo Park, Eammanuel Bengio, and Yoshua Bengio. Learning to scale logits for temperature-ccoddtoonal gflownets, 2024b.
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. arXio preprint arXiv:2402.18013 2024.
Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Heráándoz-Garcda Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of continuous generative flow networks. International Conference on Machine Learning (ICMT.,, 2023.
11

---

Published in Transactions on Machine Learning Research (12/2024)
Kimin Lee, Hao Liu, Moonkyung Ryu, Olivia Watkins, Yuging Du, Craig Boutilier, P. Abbeel, Mohammad Ghavanzadeh, and Shixiang Shane Gu. Aligning text-to-image models using human leedback. ArXin, abs/2302.12192, 2023.
Dianbo Liu, Moksh Jain, Bonaventure F. P. Dossou, Qianli Shen, Salem Lahlou, Anirudh Goyal, Nikolay Malkin, Chris C. Emezue, Dinghuai Zhang, Nadhir Hassen, Xu Ji, Kenji Kawaguchi, and Yoshua Bengio. Gflowout: Dropout with generative flow networks. In International Conference on Machine Learning, 2022.
Zhen Liu, Tim Z. Xiao, Weiyang Liu, Yoshua Bengio, and Dinghuai Zhang. Efficient diversity-preservingg diffusion alignment via pradient-informed gfownets, 2024.
Jiangyan Ma, Emmanuel Bengio, Yoshua Bengio, and Dinghuai Zhang. Baking symmetry into gflownets.
Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. International Conference on Machine Learning (ICML), 2022.
Nikolay Malkin, Moksh Jain, Emmanucl Bongio, Chen Sun, and Yoshua Bengio. Trajoctory balance: Improved credit assignment in GFlowNets. Neural Information Processing Systems (NeurIPS), 2022.
Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference. International Conference Learning Representations (ICLR), 2023.
Pierre Marion, Anna Korba, Peter Bartlett, Mathieu Blondel, Valentin De Bortoli, Arnaud Doucet, Felipe Llinares-López, Courtney Paquette, and Quentin Berthet. Implicit diffusion: Efficient optimization through stochastic sampling, 2024.
Sobhan Mohammadpour, Emmanuel Bengio, Emma Frejinger, and Pierre-Lue Bacon. Maximum entropy gflownets with soft q-learning, In International Conference on Artificial Intelligence and Statistics, pp- 2593-2601. PMLR, 2024.
Ling Pan, Moksh Jain, Kanika Madan, and Yoshua Bengio. Pro-training and finc-tuning generative flow networks, 2023a.
Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of GFlowNets with local credit and incomplete trajectories. International Conference on Machine Learning (ICML), 2023b.
Ling Pan, Dinghuai Zhang, Aaron Courville, Longbo Huang, and Yoshua Bengio. Generative augmented flow networks. International Conference on Learning Representations (ICLR), 2023c.
Ling Pan, Dinghuai Zhang, Moksh Jain, Longbo Huang, and Yoshua Bengio. Stochastic generative flow networks. Uncertainly in Artificial Intelligence (UAT), 2023d.
Ben Poole, Ajay Jain, Jonathan T Barron, and Ben Mildenhall. Dreamfusion: Text-to-3d using 2d diffusion. arXiv preprint 2Xiu:2209.1498. 2022.
Mihir Prabhudesai, Anirudh Goyal, Deepak Pathak, and Katerina Fragkiadaki. Aligning text-to-inage diffusion models with reward backpropagation, 2023.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Tlya Sutskever. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, 2021.
Robin Rombach, A. Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10674 10685, 2021.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. ArXiv, abs/1505.04555 2015. URL https://apl. semanticscholar. .orpooolrrrrrrrrooooeeeeeeooeeee 3719281.
12

---

Published in Transactions on Machine Learning Research (12/2024)
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding, Advances in neural information processing systems, 35:36479 36494, 2022.
Tim Salimans and Jonathan Ho. Progrossive distillation for fast sampling of diffusion models, 2022.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint 1ix::707.06344 2017.
Max W. Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, and Tommaso 2023. URL
Biancalani. Towards understanding and improving gllownet. training. ArXiv,
attps://api.smmattcccchoo
semanticscholat org /orp/orrpssDD 2586776887.
Jascha Narain Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli, Deep unsu- pervised learning using nonequilibrium thermodynamics. ArXív, abs/1503.03556 033585555,5 2015. URL https: sea.pissmmaniiccooll... org/ CorpusID: 14888175.
Yang Song, Jascha Narain Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ArXiv, abs/201 22211.111144 2020. URL https://api. semanticscholar DatprttttcootoooorrrrrrstaaDDDD 227209335.
Richard S Sutton. Learning to predict by the methods of temporal differences. Machine Learning, 3(1):9444, 1988.
Daniil Tiapkin, Nikita Morozov, Alexey Naumov, and Dmitry P Vetrov. Generative flow networks entropy-rese reguapy-regularized rl. In International Conference on Artificial Intelligence and Statistics, pp. 4213 111--22.. PMLR, 2024.
Masatoshi Uehara, 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 entropx-regulariieed control, 2024.
Siddarth Venkatraman Moksh Jain", Luca Scimeca*, Minsu Kim*, Marcin Sendera*, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, Alexandre Adam, Jarrid Rector -rooos, Yoshua Bengio, Glen Berseth, and Nikolay Malkin. Amortizing intractable inference in diffusion models for vision, language, and control. 2024.
Pascal Vincent. A connection between score matching and denoising autocncoders. Newral Computation, 23: 1661 1674, 2011. CRL https: htt///pi..smmanicccchhllllla wtw//orpps :tg/CoopusID:5566643.. 5560643.
Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkan, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization, 2023a.
Bram Wallace, Akash Gokul, Stefano Ermon, and Nikhil Vijay Naik. End-to-end diffusion latent optimization improves classifier guidance. 2023 IEEE/CVF International Conference on Computer Vision (ICCV), pp. 7246 7256, 2023b.
Xiaoshi Wu, Yiming Hao, Keqiang Sun, Yixiong Chen, Feng Zhu, Rui Zhao, and Hongsheng Li. Human preference score v2: solid benchmark for evaluating human preferences of text-to-image synthesis. ArXin, abs/2306.0920 2023.
Xiaoshi Wu, Yiming Hao, Manyuan Zhang, Keqiang Sun, Zhaoyang Huang, Guanglu Song, Yu Liu, and Hongsheng Li. Deep reward supervisions for tuning text-to-image diffusion models, 2024.
Jiazheng Xu, Xiao Liu, Yuchen Wu, Yuxuan Tong, Qinkai Li, Ming Ding, Jie Tang, and Yuxiao Dong. Imagere- ward: Learning and evaluating human preferences for text-to-image generation. ArXiv, abs/2304.05977, 2023.
13

---

Published in Transactions on Machine Learning Research (12/2024)
Kai Yang, Jian Tao, Jiafei Lyu, Chunjiang Ge, Jiaxin Chen, Qimai Li, Weihan Shen, Xiaolong Zhu, and Xiu Li. Using human leedback to fine-tune diffusion models without any reward model, 2024.
Mengjiao Yang, KwangHwan Cho, Amil Merchant, Pieter Abbeel, Dale Schuurmans, Igor Mordatch, and Ekin Dogus Cubuk. Scalable diffusion for materials generation. arXiv preprint arXiv:231 1.09235, 2023.
David W. Zhang, Corrado Rainone, Markus F. Peschl, and Roberto Bondesan. Robust scheduling with gflownets. ArXiv, abs/2302.05446, 2023a. URL https:///pp..emmaiiiccccoo....e ott.raatttccccoooaaoooo/ooooo//o 256827133.
Dinghuai Zhang, Ricky T. Q. Chen, Nikolay Malkin, and Yoshua Bengio. Unifying generative models with GFlowNets and beyond. orXiv preprinl Xrx::22009. 0260662, 2022a.
Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. International Conference on Machine Learning (ICML), 2022b.
Dinghuai Zhang, Ricky Tian Qi Chen, Cheng-Hao Liu, Aaron C. Courville, and Yoshua Bengio. Diffusion generative flow samplers: Improving lcarning signals through partial trajectory optimization. ArXiv, abs/2330.00269,, 2023b.
Dinghuai Zhang, Aaron Courville, Yoshua Bengio, Qinqing Zheng, Amy Zhang, and Ricky T. Q. Chen. Latent state marginalization as low-cost approach for improving exploration, 2023c.
Dinghuai Zhang, Hanjun Dai, Nikolay Malkin, Aaron C. Courville, Yoshua Bengio, and Ling Pan. Let the flows tell: Solving graph combinatorial optimization problems with gflownets. ArXiv, abs/ 215/2305.17010, 2023d.
Dinghuai Zhang, L. Pan, Ricky T. Q. Chen, Aaron C. Courville, and Yoshua Bengio. Distributional gflownets with quantile flows. arXiv preprint arXiv:2302. 05793, 2023e.
Dinghuai Zhang, Ricky T. Q. Chen, Cheng-Haa Hao Liu, Aaron Courville, and Yoshua Bengio. Diffusion generative flow samplers: Improving learning signals through partial trajectory optimization, 2024a.
Dinghuai Zhang, Ling Pan, Ricky T. Q. Chen, Aaron Courville, and Yoshua Bengio. Distributional gflownets with quantile flows, 2024b.
Mingyang Zhou, Zichao Yan, Elliot Layne, Nikolay Malkin, Dinghuai Zhang, Moksh Jain, Mathicu Blanchette, and Yoshua Bengio. Phylogfn: Phylogenetic inference with generative flow networks, 2024.
Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In Aaai, volume 8, pp. 1433-1438. Chicago, II, USA, 2008.
Heiko Zimmermann, Fredrik Lindsten, J.-W. van de Meent, and Christian Andersson Naesseth. A variational perspective on generative flow networks. ArXic, abs/2210.07992, 2022. URL https: //api sewantiscchhol org/ CorpusID:252907677.
14

---

Published in Transactions on Machine Learning Research (12/2024)
A Proof
A.1 Proof of Proposition 4
Proof. Recalling that b(x,, x(-1) = stop- gradient (log P 11x))
VoDKI. (po(x--1xx)) (x 1)g(Xx1 (xx()
FΦ(x,,t) t)
(x., (xx0x))))((11111111 (xx)
-o (x((x-11xx)) log
F Fx - ()( (x11) dx--1
F6 (xp, (x, (x+)
(xr. (x ( dx
Vopo x log
dx + po( Vo log
(x 1---11990 (xx (x
F (xI-1, t ()g((x1/2x00))))
= po(xx-1| Vo log xx)0(xxx x1-1 dx + po 1x xx) Vologgo x,- x2) dx
=Vo -
Fx (po(- [b(xe, Xt-) Vog logpo(x(-1
B More about DAG
B.1 RL optimal solutions of the denoising MDP
Training a standard RI. algorithm within this diffusion MDP in Section 3.1 to perfection means the model would only gemerate single trajectory with the largest reward value. This usually comes with the disadvantage of mode collapse in generated samples in practice. One direct solution is soft maximum entropy RL (Ziebart et al., 2008; Fox et al., 2017; Haarnoja et al., 2017; Zhang et al., 2023c), whose optimal solution is trajectory- level distribution and the probability of generating cach trajectory is proportional to its trajectory cumulative reward, i.e., (Xo:T) x E. R& (x,) = R(xx). However, in theory this means (x(x)) = (Xx,r) dx x f R(xo) dx1.7 R(xo)) dx1:T, which is not a well-defined finite term for unbounded continuous spaces. In contrast, the optimal solution of GFlowNet is (xo) x R(xo).
B.2 Experimental details
Regarding training hyperparametors, we follow the DDPO github repository implementation and describe them below for completeness. We use classifier-Iree guidance (Ho & Salimans, 2022, CFG) with guidance weight being
We use a 50-step DDIM schedule. We use NVIDIA Ax 80GB GPUs for each task, and use batch size of 8 per single GPU. We do 4 step gradient accumulation, which makes the essential batch size to be 256. For each "epoch", sample 512 trajectories during the rollout phase and perform 8 optimization steps during the training phase. We train for 100 epochs. We use a 3 10-4 learning rate for both the diffusion model and the flow function model without further tuning. We use the Adamm optimizer and gradient clip with the norm being 1. We set == =1x x 10-4 in Equation 14. We use bfloat16 precision.
The GFlowNet framework requires the reward function to be always non-negative, so we just take the exponential of the reward be used the GelowNet reward. We also set the reward exponential to 3 = 100 (L.., setting the distribution temperature to be 1/100). Therefore, log R(-) R.riiiina))).. Note that in GFlowNet training practice, only need to use the logarithm of the reward rather than the original reward value. We lincarly anncal 33 from 0 to its maximal value in the first half of the training. We found that this almost does not change the final result. but is helpful for training stability.For DAG-KL, we put the final B coefficient on the KL gradient term. We also find using a KL regularization DKL (xe(x;11 Ippes (x-1/xx)) to be helpful for stability (this also mentioned in Fan et al. (2023)). In practice, essentially adding a lg regularization term on the output of the U-Net after CFG between the current model and previous rollout model. We simply use a coefficient. on this term without further tuning.
15

---

Published in Transactions on Machine Learning Research (12/2024)
Aesthetic
Method
Ours Nave FL
7.0
Without
Revare
6.0
5.0
150
Training stop
DDPO
DAG
Figure 8: Samples on C1 diffusion alignment experiments. The reward function is the probability of the generated image falling into the categorics of car, truck, ship, and plane calculated by a pretrained classifier. The RL baseline shows mode collapse behaviors while the target distribution is actually multimodal.
Figure 9: Ablation study for the forward-looking (FL) usage in Sec- tion 3.2 on the Aesthetic reward task.
We use Stable Diffusion v1.5 base model and use LoRA for post-training, following Black et al. (2023). For the architecture of the state flow function, we take similar structure to the downsample part of the U-Net. The implementation is based on the hugging face diffuuers package. We use 3 "CrossAttnDowmnBookk blocks and "DownBlock2D" and do downsampling on all of them. We set the layers per block to be 1, and set their block out channels to be 64, 128, 256, 256. We use a final average pooling layer with kernel and stride size to output a scalar given inputs including latent image, time step, and prompt embodding. We do not report diversity metric in previous GFlowNer Net. literature, as the average pairwise Euclidean distance in high dimensional space (64 64 X 4 10, 000 dim.) is not meaningful metric.
For computing diversity in Table 1, we take trained model and independently generate two batches of images based on corresponding prompls, with the batch size being 2048. For our proposed methods, we choose the carliest checkpoint with a reward larger than the DDPO final rewards. We save our checkpoints for every 10 epochs. We use an Inception network to compute the mean and covariance features and calculate the Frechet distance; then we use the resulting FID as the diversity metric. Additionally, we show in Figure that DAG-KL achieves better Pareto frontier the reward-diversity trade-off the HPSv2 task.
B.3 CIFAR-10 toy example
35.0
J0
10
DAG-D6
DAG-KL
32.5
We also include toy experiment on CIFAR-10 pretrained DDPM We train a ResNet18 classifier and set the reward function to be the probability of the generated image falling into the categories of car, truck, ship, and plane. We use the same hyperparameters with the Stable Diffusion setting, except only 1 GPU with a 256 batch size for each run without gradient accunulation. We illustrate the generation results in Figure 8. We use DAG-DB here, and the DAG-KL generation is similar and non non-distinguishable with it. We can see that in this relative toy task, the RL bascline casily optimizes the problem extremes and behaves mode collapse some extent (only generating samples of particular plane). While for our methods, the generation results are diverse and cover different classes of vehicles. Both methods achieve average log probability larger than -0.01, which means the probability of falling into target categories are very close to 1.
DDPO
20
27.5
AP
25.0
2600
28p
30
A0
90
.50
20.0
.90
70
17.5
BIFD
600
7
90
0.28
0.28
0.30
0.32
0.34
0.38
Reward
Figure Scatter plot about model diversity and reward when they have been trained for different number of epochs (shown beside each of the points). It be that our method (DAG-KL) achieves a better reward-diversity trade-off.
httpess/wangiigaaec../////////g
16

---

Published in Transactions on Machine Learning Research (12/2024)
The clock the side of the metal building is gold and black
Kitchen with wooden kitchen island and checkered floor
pink bicycle leaning against. fence river
An empty kitchen with lots of
tile blue counter top space
Figure 10: More text- image alignment results. We display four different prompts and the corresponding generation visualization from the original Stable Diffusion (Ist row), DDPO (2nd row), DAG-DB (3rd row), and DAG-KL (4th row) models to compare their aligument ability.
B.4 More results
In Figure 9, we perform ablation study the proposed denoising diffusion-specific way of forward-looking tochnique for the Aesthetic score task. Specifically, compare the left and right hand sides of Equation 7, where we use "naive FL" to refer to the left hand side, and "ours" for the right hand side, as it is just the DAG-DB method in the main text. We also show the performance of not using the forward-looking technique and call it "without FL" in the figure. We can see that both the variants cannot achiove effective learning compared to our proposed method.
In Figure 10, we put more visualization comparisons about the text-image alignment performance from the models trained the HPSv2 reward, with similar form to Figure
In Section 5 of our paper, we have discussed the reason to use the DDPO baseline it outperforms other align-from-blacke black-box methods including DPOK (Fan et al., 2023). For completeness, we conduct an experiment of DPOK on the Aesthetic score experiment. We keep the same experimental setup as the ones in Section 5; while DAG-DB, DAG-KL, and DDPO achives a score of 7.6587, 7.6396, and 6.3178 respectively, DPOK obtains a result of 6.1876. This result shows that our method is superior to other basclines.
17


Reviews:
[{"review_id": "S2CCU3SAQw", "reviewer": "Reviewer_d4Ze", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n* The proposed method establishes a meaningful connection between Generative Flow Networks (GFlowNets) and Diffusion Models, * which could inspire further research integrating these two frameworks for generative tasks.\n* The method achieves a superior balance between reward maximization and sample diversity compared to existing state-of-the-art methods, demonstrating its effectiveness in aligning quality and diversity goals in generative modeling.\n\nWeaknesses:\n\nThe main contributions of this work are not entirely clear. While the introduction of the KL-based on-policy objective appears to be a key innovation, the authors should further emphasize other novel aspects of the method beyond this. Clarifying additional contributions, such as specific algorithmic improvements, novel evaluation metrics, or new insights into generative modeling with GFlowNets, would strengthen the clarity and impact of the work.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper, the authors propose a GFlowNet-based algorithm to post training diffusion models in a black box manner. A new objective has been proposed to improve the sample efficiency which outperformance the reinforcement learning baseline within the same number of trajectory.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "It is recommended that the authors report the quality of the generated images. Metrics like Frechet Inception Distance (FID) or other relevant image quality metrics should be included to provide a more comprehensive evaluation of the method's performance.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "KeZlmxapzr", "reviewer": "Reviewer_LRTV", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- The paper is well-organized and well-written. The motivation for using GFlowNets for diffusion alignments is well clarified, and the method is formulated clearly step-by-step.\n- DAG-DB and DAG-KL show better training efficiency with better and faster convergence.\n- DAG shows strong image-to-text results with improved alignment, generation diversity, and reward diversity trade-off. \n\n### Weaknesses\n- Lack of comparisons with more diffusion model alignment methods.\n- A paragraph for the paper summarization is missing.\n-  In Remark 2, experiments with the GFlowNets objective using longer partial trajectory lengths (e.g., SubTB with length two or more) are suggested. This would better illustrate the impact of training batch size and flow consistency over varying trajectory lengths.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes DAG, a text-to-image diffusion model alignment method with GFlowNets rather than typical RL. DAG treats the denoising process as a MDP and improves the detailed balance (DB) objective in GFlowNets for better diffusion alignment. The authors further find the connection between DAG-DB and REINFORCE and propose DAG-KL, which updates the diffusion model and flow function independently with better alignment and diversity trade-off. The text-image samples show that the proposed DAG-DB and DAG-CL have better text-image alignments even for text with multiple concepts than DDPO.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Check the listed weaknesses above. DPKO [1] (extra RL baseline) and Diffusion-DPO [2] (implicit reward baseline) could be potential choices for extra baselines. \n\n[1] Fan, Ying, et al. \"Reinforcement learning for fine-tuning text-to-image diffusion models.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[2] Wallace, Bram, et al. \"Diffusion model alignment using direct preference optimization.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2024.", "broader_impact_concerns": "No broader impact concerns needed.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "eKKTzM7fau", "reviewer": "Reviewer_6d2o", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Rigorous theoretical derivation.\n2. Clear and polished presentation.\n3. Strong qualitative results.\n4. Reasonable quantitative performance.\n\nWeakness:\nThe paper compares performance against only one baseline. Including additional comparisons with alternative alignment methods would provide a more comprehensive evaluation and further strengthen the paper’s claims.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper reformulates denoising diffusion models within the GFlowNets framework, training the model to generate samples with probabilities proportional to a reward function rather than directly maximizing it. It also reformulates the Detailed Balance (DB) equation into a KL divergence form, deriving a REINFORCE-style objective for efficient single-transition updates. This allows for larger batch sizes during training, potentially improving performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "How does the training efficiency and performance of the proposed method compare to recent method such as AlignProp?\n\nPrabhudesai, Mihir, et al. \"Aligning text-to-image diffusion models with reward backpropagation.\" arXiv preprint arXiv:2310.03739 (2023).", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper incorporate the diffusion model and GFlowNets frameworks, which makes post-training of diffusion models under black-box reward functions feasible. The reviewers are unanimous on the positive sides of the proposed method. Some reviewers raised some questions on the evaluation and comparing with more baselines. The authors address the issues with either more experimental results or clarifications on the concerns. The reviewers are satisfied with the rebuttal with no further concerns. I agree with the reviewers, but request the authors to incorporate additional results and further clarifications to be incorporated into the final version.
Decision: Accept (poster)