Amortizing intractable inference in diffusion models for vision, language, and control
Authors: Siddarth Venkatraman
Moksh Jain
Luca Scimeca
Minsu Kim
Marcin Sendera
Mohsin Hasan
Luke Rowe
Sarthak Mittal
Pablo Lemos
Emmanuel Bengio
Alexandre Adam
Jarrid Rector-Brooks
Yoshua Bengio
Glen Berseth
Nikolay Malkin
Venue: NeurIPS 2024 poster (2024)
Keywords: diffusion
inverse problems
conditional generation
language models
infilling
discrete diffusion
offline RL
planning
GFlowNet
Primary area: probabilistic_methods
TL;DR: An asymptotically unbiased objective for sampling from the product of a diffusion prior with a constraint, applied to vision, language, and RL.

Abstract: Diffusion models have emerged as effective distribution estimators in vision, language, and reinforcement learning, but their use as priors in downstream tasks poses an intractable posterior inference problem. This paper studies *amortized* sampling of the posterior over data, $\mathbf{x}\sim p^{\rm post}(\mathbf{x})\propto p(\mathbf{x})r(\mathbf{x})$, in a model that consists of a diffusion generative model prior $p(\mathbf{x})$ and a black-box constraint or likelihood function $r(\mathbf{x})$. We state and prove the asymptotic correctness of a data-free learning objective, *relative trajectory balance*, for training a diffusion model that samples from this posterior, a problem that existing methods solve only approximately or in restricted cases. Relative trajectory balance arises from the generative flow network perspective on diffusion models, which allows the use of deep reinforcement learning techniques to improve mode coverage. Experiments illustrate the broad potential of unbiased inference of arbitrary posteriors under diffusion priors: in vision (classifier guidance), language (infilling under a discrete diffusion LLM), and multimodal data (text-to-image generation). Beyond generative modeling, we apply relative trajectory balance to the problem of continuous control with a score-based behavior prior, achieving state-of-the-art results on benchmarks in offline reinforcement learning. Code is available at [this link](https://github.com/GFNOrg/diffusion-finetuning).

Amortizing intractable inference in diffusion models for vision, language, and control
Siddarth Venkatraman" Mila, Université de Montréal
Moksh Jain" Mila, Université de Montréal
Luca Scimeca' Mila, Université de Montréal
Minsu Kim Mila, Université de Montréal KAIST
Marcin Sendera"
Mohsin Hasan
Mila, Université de Montréal Jagiellonian University
Mila, Université de Montréal
Luke Rowe Mila, Université de Montréal
Sarthak Mittal
Pablo Lemos
Mila, Université de Montréal
Mila, Université de Montréal Ciela Institute
Dreamfold
Emmanuel Bengio
Alexandre Adam
Jarrid Rector-Brooks
Recursion
Mila, Université de Montréal
Mila, Université de Montréal Dreamfold
Ciela Institute
Yoshua Bengio
Glen Berseth
Nikolay Malkin
Mila, Université de Montréal CIFAR
Mila, Université de Montréal CIFAR
Mila, Université de Montréal University of Edinburgh
siddarth.
venkatraman, moksh. jaishhjain,, scimeca minsu. kim, marcin. sendera, marcin.sendera,
@mill.qqebeec
,nikolay malkin
Abstract
Diffusion models have emerged as effective distribution estimators vision, lan- guage, and reinforcement learning, but their use priors in downstream tasks poses an intractable posterior inference problem. This paper studies amorlized sampling of the posterior a diffusion generative model prior pxx) and black-box constraint or likelihood function r(x). We state and prove the asymptotic correctness of data-free learning objective, relative trajectory balance, for training a diffusion model that samples from this posterior, problem that existing methods solve only approximately or in restricted cases. Relative trajectory balance arises from the generative flow net- work perspective on diffusion models, which allows the use of deep reinforcement learning techniques to improve mode coverage. We illustrate the broad potential of unbiased inference of arbitrary posteriors under diffusion priors across a collection of experiments: in vision (classifier guidance), language (infilling under discrete diffusion LLM), multimodal data (text-to-image generation). Beyond gener- ative modeling, apply relative trajectory balance to the problem of continuous control with a score- based behavior prior, achieving state-of-the-an results on benchmarks in offline reinforcement learning, Code is available at this link.
data, pest (x) pxx) in model that consists of
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Table I: Sources of diffusion priors and constraints.
Domain
Prior p(s)
Constraint e(x)
Posterice
Conditional image generation ($3.1)
Inuge diffusion model pix)
Classifice likellhood ptr [x
Class- Cladi-ounii distribution p(x [
Test-in-image genenation ($3.2)
Text-eo-mmige foundation model
HF reward model
Aligned lest-do-image model
Language intilling (83.3)
Discrete diffusion model
Autoregressive completion likellhoood
Intiling distribution
Offine RL policy extraction (53.4)
Diffusion model as behavioe policy
Boltzmann dist. of function
Optiall Kl-coosstaineed policy
Introduction
Diffusion models [68, 27, 72] are powerful class of hierarchical generative models, used to model complex distributions over images [51, 12, 63], text [4, 13, 40, 24, 23, 43], and actions in reinforcement learning [30, 83, 32] to name a few. In each of these domains, downstream problems require sampling product distributions, where a pretrained diffusion model serves as a prior p(x) that is multiplied by an auxiliary constraint r(x). For example, if p(x) is a prior over images defined by diffusion model, and r(x)= = p(c x) is the likelihood that an image belongs to class c, then class-conditional image generation requires sampling from the Bayesian posterior p(x | c) p(x)p(c x). In offline reinforcement learning, μ((a s) is a conditional diffusion model over actions serving as a behavior policy, KL-constrained policy improvement [55, 44] requires sampling from the normalized product of μ(a s) with a Boltzmann distribution defined by Q-function, (a 8) or p(a s) exp(BQ(s, (s, a)). In language modeling, various conditional generation problems [43, 23, 29] amount to posterior sampling under discrete diffusion model prior. Table I summarizes four such problems that the proposed method improves upon prior work.
The hicrarchical nature the generative process in diffusion models, which generate samples from p(x) by a deep chain of stochastic transformations, makes exact sampling from posteriors p (x) under black-box function r(x) intractable. Common solutions to this problem involve inference techniques based on linear approximations [73, 33, 31, 11] or stochastic optimization [22, 48]. Others estimate the 'guidance' term the difference drift functions between the diffusion models sampling the prior and posterior - by training a classifier on noised data [12]. but when such data is not available, one must resort to approximations or Monte Carlo estimates [70, 14, 10], which are challenging to scale to high-dimensional problems. Reinforcement learning methods that have recently been proposed for this problem [8, 16] are biased and prone to mode collapse (Fig. 1).
Contributions. Inspired by recent techniques in training diffusion models to sample distributions defined by unnormalized densities [89, 62, 78, 65], propose an asymptotically unbiased training ob- jective, called relative trajectory balance (RTB), for training diffusion models that sample from poste- rior distributions under diffusion model prior ($2.2). RTB is derived from the perspective of diffusion models as continuous generative flow networks [38]. This perspective also allows us to freely leverage off-policy training, when data with high density under the posterior is available ($2.3). RTB can be applied to iterative generative processes beyond standard diffusion models: our methods generalize discrete diffusion models and extend existing methods for autoregressive language models ($2.4).
Our experiments demonstrate the versatility of our approach in a variety of domains:
In vision, we show that RTB achieves competitive classifier- dlasiifierguideed image generation for uncondi- tional diffusion vision priors (53.1) and can be used improve caption-conditioned generation under text-to-iman -maae foundation model priors ($3.2).
language modeling, we report strong results for infilling tasks with discrete diffusion language models ($3.3).
Finally, we show that RTB achieves state-of-the-art results on continuous control benchmarks that leverage score-based behavior priors ($3.4).
2 Learning posterior samplers with diffusion priors
We consider the problem of posterior inference under a prior given by hierarchical generative model. In this section, we present the mathematical setting (62.1), our proposed RTB objective ($2.2), and training methods for RTB ($2.3). We will first discuss the case of a diffusion prior over Rd, and later discuss how the methods generalize arbitrary hierarchical priors ($2.4).

---

(a) Prior
(b) Posterior
(c) RTB (ours) (d) RL (KL reg.) (c) RL (no reg.)
(f) CG
Figure I: Sampling densities learned by various posterior inference methods. The prior is diffusion model sampling mixture of 25 Gaussians (a) and the posterior is the product of the prior with constraint that masks all but 9 of the modes (b). Our method (RTB) samples close to the true posterior (c). RL methods with tuned KL regularization yield inaccurate inference (d), while without KL regularization, they mode-collapse (c). A classifier guidance (CG) approximation (f) results in biased outcomes. For details, see SC.
2.1 Background and setting: Diffusion models as hierarchical generative models denoising diffusion model generates data xj by Markovian generative process:
(noise)
Xg
X
X
- XI = X (data),
where Ar 1 and 7 the number of discretization steps. The initial distribution (Xx) is fixed 7 (typically (0, I) and the transition from x,-1 X, is modeled as a Gaussian perturbation with time-dependent variance:
P x,) = (x)+ar X, + uy (X, of
The scaling of the mean and variance by Ar insubstantial for fixed T, but ensures that the diffusion process is well-defined in the limit T - assuming regularity conditions u, [53, 64]. The process given by (1, 2) is then identical to Euler-Maruyama integration the stochastic differential equation (SDE) dx, My (x,) ) dr o, dw,-
The likelihood of denoising trajectory Xo
XAr
X factors as
[
((xo, XAI.
x1) (xx))
p(xisi x(--1ar)
(3)
and defines marginal density over the data space:
f p(xn, KAr...., X1) dxg dx dxi-Ar-
p(x)) =
A reverse-time process, X1 XI- similarly defines a conditional density over trajectories:
xo, with densities q, can be defined analogously, and
q (xo NAI...... XI-Ai x1)
[ i=1
9(x(-1)5r Xixr).
In the training of diffusion models, as discussed below, the process & typically fixed to a simple distribution (usually discretized Ornstein- Uhlenbeck process), and the result of training is that and q are close as distributions over trajectories.
Diffusion model training as divergence minimization. Diffusion models parametrize the drift n, (x,) in (Equation 2) as a neural network 1; with parameters 0 and taking x, and as input. We denote the distributions over trajectories induced by (Equation 3, Equation 4) by Po to show their dependence on the parameter.
In the most common setting, diffusion models are trained maximize the likelihood of dataset. In the notation above, this corresponds to assuming (xx)) ) is fixed to an empirical measure (with the
'The time indexing suggestive of an SDE discretization is used for consistency with the diffusion samplers literature [89, 65]. The indexing X7 XT-1 Xg is often used for diffusion models trained from data.

---

points of a training dataset D assumed to be i..d. samples from (xx)). Training minimizes with respect the divergence between the processes q and Po:
(g ..r....,,. x1) (x.. Ksr....... ....)..,))) xI))
(6)
(xx((((x)) II (x()))) + -xxx D ( (x(x N.r..... xi-ar x1) ll (x Xar. .......,---. xx))
2 (xi(((())) Il (=()))) - |- log pa (x1)] const.
where the inequality- an instance of the data processing inequality for the KL divergence - shows that minimizing the divergence between distributions over trajectories is equivalent to maximizing lower bound on the data log-likelihood under the model Po-
As shown in [71], minimization of the KL. in (Equation 6) is essentially equivalent to the traditional
approach to training diffusion models via denoising score matching [80, 68, 27J. Such training exploits that for typical choices of the noising process q, the optimal 4, (X,) can be expressed in terms of the Stein score of y ((x)) convolved with a Gaussian, allowing an efficient stochastic regression objective for u,. For full generality of our exposition for arbitrary iterative generative processes, prefer to think of (Equation 6) as the primal objective and denoising score matching as an efficient means of minimizing it.
Trajectory balance and distribution-matchinn training. From (Equation 6) we also see that the bound is tight if the conditionals of Po and q on x; coincide, i.e., q is equal the posterior distribution of conditioned on x1. Indeed, the model pe minimizes (Equation 6) for a distribution with continuous density (x1) if and only if, for all denoising trajectories,
(x NAr..... x1) (x)) ((x)) Xar.... X1- x)).
(7)
This was named the trajectory balance (TB) constraint by [38] - by analogy with a constraint for discrete-space iterative sampling [46] and is a time-discretized version of a constraint used for enforcing equality of continuous-time path space measures in [52] (see [7] for asymptotic analysis).
In [61, 38], the constraint (7) was used for the training of diffusion models in a data-free setting, where instead of i..dd samples from ((x)) one has access a (possibly unnormalized) density (xx)) c(xi)/2 /Z. from which one wishes sample. These objectives minimize the squared log-ratio between the two sides of (7), which allows the trajectories Xo - XAr - ... - xi used for training to be sampled from any training distribution, such as 'exploratory' modifications of Pe or trajectories found by local search (MCMC) in the target space. The flexibility of off-policy exploration that this allows was studied by [65]. Such objectives contrast with on-policy, simulation- based approaches that require differentiatigg through the sampling process [e.g., 89, 78, 6, 79].
2.2 Intractable inference under diffusion priors
Consider a diffusion model Po. defining a marginal density Po (x1), and a positive constraint function r Rd - Rs0- We are interested in training a diffusion model post with drift function post that would sample the product distribution (x1) x (x(x))x(x)) If r(x)) p(y x:) is conditional distribution over another variable y, then ppost is the Bayesian posterior p ( y)-
Because samples from (x))(())) are not assumed to be available, one cannot directly train using the objective (6). Nor can one directly apply objectives for patributtionmmatchi training, such as those that enforce (7), since the marginal (xi) is not available. However, we make the following
observation (proof in SA).
Proposition 1 (Relative TB constraint). If Pa. p and the scalar Zo jointly satisfy the relative trajectory balance (RTB) constraint
Zo
P (Xo, XAr..... x1) = rxxxx (No. NAr.....X.
(8)
for every denoising trajectory XA
X1, then post (x1) or px(x))r(x)), iee, the diffusion model & samples the posterior distribution. Furthermore, if po also satisfies the TB constraint (7) with respect to the noising process 9 and some target density (xx), then p satisfies the TB constraint with respect to the target density (x,) x ((x))r(x)) and 7. ((x))((x)) dxj-
post
Note that the two joints appearing in (8) are defined as products over transitions, via (3).

---

Relative trajectory balance as loss. Analogously to the conversion of the TB constraint (7) into a trajectory-dependent training objective in [46, 38], we define the relative trajectory balance loss as the discrepancy between the two sides of (8), seen function of the vector & that parametrizes the posterior diffusion model and the scalar Z (parametrized via log Z for numerical stability):
P post (XO. XA
X1)
(Xo XAI 7 x Φ) = log r(xi)po No. XA
Optimizing this objective to 0 for all trajectories ensures that (8) is satisfied. While the RTB constraint (8) has similar form TB (7), RTB involves the ratio of two denoising processes. while TB involves the ratio of forward and a backward process. However, the name 'relative TB' is justified by interpreting the densities in a TB constraint relative to a measure defined by the prior model; see $2.4.
If assume (Xn) p
paal (xn) are fixed (e.g., to standard normal), then (9) reduces
Z log r(xi)
P (Xiar X(i-1)ar)
log ( X((-1)5e)
(10)
Notably, the gradient of this objective with respect to & does require differentiation (backpropaga- tion) into the sampling process that produced a trajectory Xg -.... X1. This offers two advantages over on- policy simulation-based methods: (1) the ability to optimize LRTB off-policy objective, ie., sampling trajectories for training from distribution different from p itself, discussed further in $2.3; (2) backpropagating only to a subset of the summands in (10), when computing and storing gradients for all steps in the trajectory is prohibitive for large diffusion models (see $H.1).
Comparison with classifier guidance. It is interesting to contrast the RTB training objective with the technique of classifier guidance [12] used for some problems of the same form. If r(x,) p(y x1) is conditional likelihood, classifier guidance relies upon writing u, (X,) - post (x,) explicitly terms of Vx, log p(y xt), by combining the expression the optimal drift u, in terms of the score of the target distribution convolved with Gaussian (cf. 12,1), with the 'Bayes' rule' for the Stein score: Vx, log (xx y) vs, log p (xx, + Vx, logp(y xx).
Classifier guidance gives the exact solution for the posterior drift when a differentiable classifier on noisy data, p(y x,) p(y x1) xx) x₂) dx), is available. Unfortunately. such classifier is not, in general, tractable to derive from the classifier on noiseless data, p(y x1), and cannot be learned without access to unbiased data samples. RTB is an asymptotically unbiased objective that recovers the difference in drifts (and thus the gradient of the log-convolved likelihood) in a data-free manner.
2.3 Training, parametrization, and conditioning
Training and exploration. The choice of which trajectories we use to take gradient steps with the RTB loss can have large impact sample efficiency. In on-policy training, we use the current policy p to generate trajectories likelihood of under pa. and a gradient updates on & minimize LRTB(T: Φ).
(Xo
x1), evaluate the reward log r(xi) and the
However, on-policy training may be insufficient to discover the modes of the posterior distribution. In this case, we can perform off-policy exploration to ensure mode coverage. For instance, given samples xj that have high density under the target distribution, we can sample noising trajectories X1- ... Xo starting from these samples and use such trajectories for training. Another effective off-policy training technique uses replay buffers. We expect the flexibility of mixing on- policy training with off-policy exploration to be strength of RTB over on-policy RL methods, as was shown for distribution searibution-matchinn training of diffusion models in [65].
Conditional constraints and amortization. Above we derived and proved the correctness of the RTB objective for an arbitrary positive constraint r(x)). If the constraints depend on other variables y for example, r r(x1:y) p(y x1)= then the posterior drift u can be conditioned on and the learned scalar log 7 replaced by model taking y as input. Such conditioning achieves amortized inference and allows generalization to new y not seen in training. Similarly, all of the preceding discussion casily generalizes priors that are conditioned on some context variable.
5

---

Efficient parametrization and Langevin inductive bias. Because the deep features learned by the prior model uo are expected to be useful in expressing the posterior drift _ we can choose to initialize u as a copy of ue and to fine-tune it, possibly in a parameter-efficien way (as described in each section of $3). This choice is inspired by the method of amortizing inference in large language models by inn-ttniig a prior model to sample an intractable posterior [29].
Furthermore, if the constraint r(xi) is differentiable, we can impose an inductive bias on the posterior drift similar the introduced for diffusion samplers of unnormalized target densities in [89] and shown to be useful for off-policy methods in [65]. namely, we write
post
/) = NN
(N, log rx,),
(11)
where NN1 and 2 are neural networks outputting a vector and a scalar, respectively. This parametrization allows the constraint to provide signal to guide the sampler at intermediate steps.
Stabilizing the loss. We propose two simple design choices for stabilizing RTB training. First, the loss in (9) can be replaced by the empirical variance over minibatch of the quantity inside the square, which removes dependence on log Zo and is especially useful in conditional settings, consistent with the findings of [65]. This amounts to relative variant of the VarGrad objective [61] (see (23) in SG). Second, we employ loss clipping: to reduce sensitivity to an imperfectly fit prior model, we do not perform updates on trajectories where the loss is close 0 (see $E,SF).
2.4
Generative flow networks and extension other hierarchical processes
RTB as TB under the prior measure. The theoretical foundations for continuous generative flow networks [38] establish the correctness of enforcing constraints such as trajectory balance (7) for training sequential samplers, such as diffusion models, to match unnormalized target densities. While we have considered Gaussian transitions and identified transition kernels with their densities with respect to the Lebesgue measure over Rd, these foundations generalize to more general reference measures. In $B, we show how the RTB constraint can be recovered as a special case of the TB constraint for certain choice of reference measure derived from the prior.
Extension to arbitrary sequential generation. While our discussion was focused on diffusion models for continuous spaces, the RTB objective can be applied to any Markovian sequential generative process, in particular, one that can be formulated as a generative flow network in the sense of [5, 38]. This includes, in particular, generative models that generate objects by a sequence of discrete steps, including autoregressive models and discrete diffusion models. In the case of discrete diffusion, where the intermediate latent variables x, lie not Rd but in the space of sequences, one simply replaces the Gaussian transition densities by transition probability masses in the RTB constraint (8) and objective (9). In the case of autoregressive models, where only one sequence of steps can generate any given object, the backward process q becomes trivial, and the RTB constraint for a model p to sample sequence from distribution with density r(x)pe()) is simply post (x) (x) for all sequences x. We note that sub-trajectory generalization of this Zop objective was used in [29] to amortize intractable inference in autoregressive language models.
3 Experiments
In this section, we present empirical results validate the efficacy of relative trajectory balance. Our experiments are designed to demonstrate the wide applicability of RTB to sample from posteriors for diffusion priors with arbitrary rewards on vision, language, and continuous control tasks.
3.1 Class-conditional posterior sampling from unconditional diffusion priors
We evaluate RTB in classifier-guided visual task where we wish learn diffusion posterior post (x c) (c given pretrained diffusion prior P (x) and a classifier r(x) p(c x).
Setup. We consider two 10-class image datasets, MNIST and CIFAR-10, using off-the-shelf unconditional diffusion priors from [27] and standard classifiers p(c x) for both datasets. We perform parameter-efficien fine-tuning of prsst initialized as copy of the prior Pa, using the RTB objective (see $E.1 for details). The RTB objective is optimized on trajectories sampled

---

Table 2: Classifier-guidee posterior sampling with pretrained unconditional diffusion priors. We report the mean-sid of each metric computed across all relevant classes for each experiment set, and highlight +5% from highest/lower experimental value. The FIID is computed between learned posterior samples and the true samples from the class in question. DP and LGD-MC fail to appropriately model the posterior distribution (high average logr(x)) while DDPO mode-collapses. RTB achieves comparable or superior performance to all other baselines, optimally balancing high
reward and diversity as measured by FID. See Table E. for conditional variants.
Dottact -
MNIST
MNIST eventdi
CIFAR-10
Algonihmm_ Metric
P[[kg((]] (D>
())
Divenity (f)
F[log/88] (T)
Diverity (T)
R[logr(8)] <T)
())
Divasity (
DPS
-2. 1997easss
0 1600mme
-1 -1e
1. 1408
3.370
-3,8005ee wn
m
0 0.2788-
LGD--MC
-2. 18990ic0
1.2873-4
1600are
-1.1720 19
1. 1445
0. 0.1600=2m
-3.088- ve
0 74022mm
0 0.2745-2
DDPD
-1,5-19910"
1.5224450
0 1350arrs
0.1314-2m0
-2 77700000
1 1.76%
0. 0.157542
DPOK
-0. 0.13399a22
1. 3.2
1462004
-0.07832200
2336
0.1631e200
2,441--22
0.5316wus5
1. 0.2415
RTR (ours)
1734
1474mues
-0. 151
1784
-2. 162
0:4777mxx
0. 0.2440-291
Prior
Posterior (seven)
Posterior (even)
Prior
Pusterior (dog)
1
3
0
5
7 7 7
7
6 0 4
8
8
9
8
7
7
117777
7
7
6
8
0
4
3
8
7
o
7
7
7
0
8
2
4
8
5
4
0
7
7
7 7
6
2
0
6
MNIST Figure 2: Samples from RTB fine-tuned diffusion posteriors.
CIFAR-10
on-policy from the current posterior model. We compare RTB with two RL-based fine-tuning techniques derived from DPOK [16] and DDPO [8] and with two classifier guidance baselines, namely DPS [11], and LGD-MC [70]. We consider three experimental settings: MNIST single- digit posterior (learning to sample images of each digit class c), CIFAR-10 single-class posterior (analogous to the previous), and MNIST multi-digit posterior. The latter is multimodal posterior, for which we set rxx)
(0.2,4668) p(c =i x) to generate even digits, and similarly for odd digits.
Samples from the RTB-fine-tuned posterior models are shown in Fig. 2. In Table 2 we report meanasdd of various metrics across all trained posteriors. We observe that models fine-tuned with RTB generate class samples with both the highest diversity (highest mean pairwise cosine distance in Inceptionv3 feature space) and closeness to true samples of the target classes (FID), while achieving high expected log r(x). Pure RL
Results.
(no KL regularization) displays mode collapse characteristics, achieving high rewards in exchange for significantly poorer diversity and FID scores (see also Fig. E. Classifier-guidas guidance-basee methods, like DP and LGD-MC, exhibit high diversity, but fail to appropriately model the posterior distribution (lowest log r(x)). Additional results can be found in $F.2.
3.2 Fine-tuning text-to-image diffusion model
Diffusion models for text-conditional image generation [e.g. 63] can struggle to consistently generate images x that adhere to complex prompts z, for example, those that involve composing multiple objects (e.g.. "A cat and a dog") or specify "unnatural" appearances (e.g., "A green-colored rabbit"). Fine- tuning pretrained text-to-image diffusion models (x z) as RI. policies to maximize some reward r(x1,z) z) based human preferences has become the standard approach to tackle this issue [8, 16, 77]. Simply maximizing the reward function can result in mode collapse as well as over-optimization of the reward. This is typically handled by constraining the fine-tuned model 5 to be close the prior p:
argmax B Eptx (2) [r(xx,))) DKL[ 1) || p(x) 2) s E. p
(12)
The optimal ē for (12) is (xx 2) x p(x] z) exp(Br(x) z)) for some inverse temperature B. The marginal KI. is intractable for diffusion models, so methods like DPOK [16] optimize an upper bound on the marginal KL in the form of a per-step KL penalty -y ΣT DKL | X(i-))A,, pxls X(i-liAr ] added to the reward. By contrast, RTB can avoid the bias in such an approximation and directly learn to generate unbiased samples from the posterior p(x) z).
Setup. We demonstrate how RTB can be used to fine-tune pretrained text-to-image diffusion models. We use the latent diffusion model Stable Diffusion v1-5 [63] as a prior over 512 x 512
7

---

(a) Four roses.
(b) green rabbit.
(c) A cat and dog.
(d) Cybernetic man.
Figure 4: Images generated from prior (top row), DPOK (middle row) and RTB (bottom row) for 4 different prompts. Images in the same column share the random DDIM seed. More images in HH.2.
ImageReward [87], which has been trained to match human preferences as well as prompt accuracy to attributes such as the number of objects, color, and compositionality, the reward log (xgr(x122 z). As reference, we present comparisons against DPOK with the default KL regularization = 0.01 and DPOK with y = 0.0, which is equivalent to DDPO [8]. We measure the final average reward and the diversity of the generated image, measured by the average pairwise cosine distance between CLIP embeddings [58] of a batch of generated images. Further details about the experimental setup and ablations discussed in $H.
images. Following DPOK [16],
Results. Fig. plots the diversity versus log re-
greer coloned rabo
A catiand a deg.
Four rase.
Oycernatic man
ward on a set of prompts from [16, 87). In terms of average log r(x,, z), RTB either matches or
outperforms DPOK, while generally achieving lower reward than DDPO. The CLIP diversity score for RTB and DPOK are on average higher than DDPO, which expected since it does not use KL regularization. For qualitative image assessments, refer to Fig. 4 and $H.2. Through this experiment, we show that RTB scales well to high dimensional, multimodal data, matching State-of-the-eee methods for fine-tuning text-to- image diffusion models.
tono kibiear
Divrsty
Dieesil
Dvass
Figure 3: Fine-tuning Stable Diffusion with Im- ageReward. We report mean Jogr(x1,35 z) and diver- sity. measured as the mean cosine distance between CLIP embeddings for batch of 100 generated im- ages.
3.3
Text infilling with discrete diffusion language models
To evaluate our approach on discrete diffusion models, we consider the problem of text infilling [91], which involves filling in missing tokens given some context tokens. While discrete diffusion models unlike their continuous counterparts can be challenging to train [4, 9, 49, 74], score entropy discrete diffusion [SEDD; 43] matches the language modeling performance autoregressive language models of similar scale. Non- rar-autorrgresivee generation in diffusion language models can provide useful inductive biases for infilling, such as the ability to attend to context on both sides of target token.
Setup. We use the ROCStories corpus [50], a dataset of short stories containing sentences each. We adopt the task setup from [29]. where the first 3 sentences of a story x and the last sentence are given, and the goal to generate the fourth sentence z such that the overall story is coherent and consistent. The fourth sentence can involve a turning point in the story and is thus challenging to fill in. We aim to model the posterior ppos (z x, y) p(z x) Prewaary y x, z) where p is SEDD language model prior (a conditional model over given x) and Preward is autoregressive language model tine-tuned with a maximum likelihood objective on a held-out subset of the dataset. As baselines, we consider simply prompting the diffusion language model with (Prompt (x)) and y (Prompt (x, y)). Additionally, to contextualize the performance, we also consider autoregressive language model baselines from [29], which studied this problem under autoregressive prior p(z x). SFT is trained on 50, 000 examples compared to 1000 for RTB, and serves as an upper bound on the performance in this task. See $F for further details about the experimental setup.
Full prompt for "Cybernetic man": "A half masked rugged laboratory engineer man with cybernetic enhancements seen from distance, scifi character portrait by greg rutkowski, esuthio, craig mullins."
8

---

Table 3: Results on the story infilling task with autoregressive and discrete diffusion language models. Metrics are computed with respect to reference infills from the dataset. All metrics are mean-std over 5 samples for each of the 100 test ex- amples. RTB with discrete diffusion prior performs better than best baseline with autoregressive prior.
Results. Following [29], we use three standard metrics to measure the similarity of the generated infills with the reference infills
from the dataset: BERTScore [90] (with De- BERTa (26]), BLEU-4 [54], and GLEU-4 [86]. Table 3 summarizes the results. We observe that the diffusion language model performs signif- icantly better than the autoregressive language model without any fine-tuning. RTB further improves the performance over prompting, and even outperforms the strongest autoregressive baseline of GelowNet fine-tuning. We provide some examples of generated text in $F.
Model
Algorithm: Metric -
BLEU-
GLEU-
BERTSsome
Prompting
0.010x0:02
0.022e0.00
0.005
Antorgg,
Supervised fine-tuning
0.012+0:01
0.023+3.01
0.013-30000
GPN fine-tuning (29)
0.01948000
0,031
0.102
Prompt (X)
0.01 1.0000
0.023+3.0
0.014+300
Discrete Prompt (x.3)
0.01440:003
0.027
0.092
diftusion
RTB
0.025-0009
0:045-0:00
0.156-3300
SFT Cupper bound)
0.031+0000
0.087
0.18222/005
3.4 KL-constrained policy search in offline reinforcement learning
The goal of RL algorithms is to learn policy x(a s). i.e., a mapping from states to actions a in an environment, that maximizes the expected cumulative discounted reward [75]. In the offline RL setting [39]. the agent has access to dataset D {(. a², s/.1 r))) /=1 of transitions (where each sample (S,, 42, St+1. r2) indicates that an agent taking action a, at state St transitioned to the next state Si+1 and received reward r,). This dataset is assumed to be generated by a behavior policy μ(a I s), which may be a diffusion model trained on D. Offline RL algorithms must learn a new policy π which achieves high return using only this dataset without interacting with the environment.
An important problem offline RL is policy extraction from trained Q-functions [55, 25, 44). For reliable extrapolation, one wants the policy to predict actions that have high Q-values, but also have high density under the behavior policy μr, as naive maximization can result in choosing actions with low probability under μ and thus unreliable predictions from the Q-function. This is formulated as KL-constrained policy search problem:
argmax Ex-d (Q(s.
(DKL s) Il p(a s))] <.
(13)
where dr is the distribution over states induced by following the policy μ.. The optimal policy in (13) is the product distribution x' (a s) or p(a s) exp( a)) for some inverse temperature B. If p[a | s) is a conditional diffusion model over continuous actions conditioned on state s, we use RTB to fine-tune diffusion behavior policy to sample from x² using as the prior and exp(BQ(s, a)) as the target constraint. We use @-function trained using IQL [36].
Setup. We test on continuous control tasks in the D4RL suite [18], which consists of offline datasets collected using a mixture of SAC policies of varying performance. We evaluate on the halfcheetah. hopper and walker2d MuJoCo [76] locomotion tasks, each of which contains three datasets of transitions: "medium" (collected from an early-stopped policy), "mmdiium both an expert and an early-stopped policy) and "medium-replay" (transitions stored in the replay buffer prior to early stopping). We compare against standard offline RL. baselines (Behavior Cloning (BC), CQL [37], and IQL [36]) and diffusion-based offline RL methods which are currently state-of- the-art: Diffuser [D; 30J, Decision Diffuser (DD; 2), D-QL [83], IDQL [25], and QGPO [44]. For algorithm implementation details, hyperparameters, and a report of baselines, see $G.
(collected from
Results. Table 4 shows that RTB matches statc-of-the-ar results across the D4RL tasks. In particular, RTB performs strongly in the medium-replay tasks, which contain the most suboptimal data and consequently the poorest behavior prior. We highlight that our performance is similar to QGPO [44], which learns intermediate energy densities for diffusion posterior sampling.
Other related work
Composing iterative generative processes. Beyond the approximate posterior sampling algorithms and application-specife techniques discussed in $1 and $3, several recent works have explored the use of hierarchical models, such as diffusion models, as modular components in generative processes. Diffusion models can be used to sample product distributions to induce compositional structure in images [41, 15). Amortized Bayesian inference [35, 60, 59, 20] another domain of sampling from product distributions where diffusion models are now being used [21]. Beyond product models, [19]

---

Table 4: Average rewards of trained policies on D4RL locomotion tasks (meanssid over 5 random seeds). Following past work, numbers within 5% of maximum in every row are highlighted.
Task J Algorithm
BC
CQL.
IQL.
D
DD
D-QI.
IDQL
QGPO
RTB (ours)
55.2
91.6
86.7
79.8
90.6-1.
96.1-03
95.9
93,5+0.3
74.93=1.7
hopper meppprr-ediimu-pperr
52.5
105.4
91.5
107.2
111.8:18
110.7=13
108.6
108.0 002.5
96.71
walker2d Birkkerdd-mdiiu---e medium-exxeeeeeee expert
107.5
108.8
109.6
108.4
108.8
109.720.
112.7
110.720.6
109.52e0.11
halfcheetah-mmedium
42.6
44.0
47.4
44.2
19.1=1.0
50.6=0.
51.0
14.4
53. 70.0.00
hopper-mediune
52.9
58.5
66.3
58.5
79.3=3.6
82.424.6
65.4
98.0+2.6
82.76=7.01
walker2d meddiiu
75.3
72.5
78.3
79.7
82.5a1.4
85.1=09
82.5
86.0
87.29a3.15
36.6
45.5
44.2
42.2
39.3:4.1
4..5
45.8
47.6:1.4
48. s
repper-mecfuum-rppllyy
18.1
95.0
94.7
96.8
100.0-0.7
100.7-0.6
92.1
96.942.6
100.40-a.21
walker2d maaker2d-mmddiimm-eppaaa
26.0
77.2
73.9
61.2
75.0=4.3
94,3=1.5
85.1
84,4
93.57=2.0)
studies ways to amortize other kinds of compositions of hierarchical processes, including diffusion models, while [67] proposes methods to sample the product of many iterative processes in application to federated learning. Finally, models without hierarchical structure, such as normalizing flows, have been used to amortize intractable inference in pretrained diffusion models [e.g.. 17]. In contrast, our method performs posterior inference by fine-tuning a prior model, developing a direction on flexible extraction of information from large pretrained models [29].
Diffusion samplers. Several prior works seek to amortize MCMC sampling from unnormalized densities by training diffusion models for efficient mode-mixing [6, 89, 78, 62, 79, 3]. Our work is most closely related to continuous GFlowNets [38], which offer an alternative perspective on training diffusion samplers using off off-policy flow consistency objectives [38, 88, 65]. Recently, Berner et al.
[7] have shown connections among existing families of diffusion sampling algorithms and analyzed their continuous-time limits.
Conclusions and future work
Relative trajectory balance provides a new approach to training diffusion models to generate unbiased posterior samples given diffusion prior and an arbitrary reward function. Through experiments on variety of domains vision, language, continuous control we demonstrated the flexibility and general applicability of RTB. RTB can be optimized with off-policy trajectories, and future work can explore ways to leverage off-policy training, using techniques such as local search [34, 65) to improve sample efficiency and mode coverage. Simulation- based objectives in the style of [89] are also applicable to the amortized sampling problems we consider and should be explored, simulation- free extensions, e.g., through objectives that are local in time [45]. The ability to handle arbitrary black-box likelihoods also makes RTB useful candidate for inverse problems in domains such as 3D object synthesis with likelihood computed via a renderer [e.g., 56, 82), imaging problems in astronomy [e.g., 11, medical imaging [e.g.. 731, and molecular structure prediction [e.g., 84).
should
Moreover, RTB could facilitate a breakthrough in modeling molecular dynamics-a notoriously challenging task due to the need to sample rare-event trajectories in chemical simulations by converting these problems into posterior inference over amplified distributions of rare-event samples. Notably, Scong et al. [66] have already explored preliminary version of this concept by employing TB with a reward multiplied by the prior likelihood, which is effectively equivalent to RTB.
Limitations. RTB learns the posterior through simulation- based training, which can be slow and eeemory-intensive Additionally, the RTB objective is computed on complete trajectories without any local rredit-assigmment signal, which can result in high variance in the gradients. Guarantees on the error incurred by imperfect fit of the prior model, amortization, and time discretization (analogous to [7]'s analysis for diffusion samplers) have not been obtained and should be considered in future work.
Broader impact. While our contributions focus on an algorithmic approach for learning posterior samplers with diffusion priors, we acknowledge that like other advances in generative modelling, our approach can potentially be used by nefarious actors to train generative models to produce harmful content and misinformation. At the same time, our approach can be also be used to mitigate biases captured in pretrained models and applied to various scientific problems.
10

---

Acknowledgments and Disclosure of Funding
The authors thank Adam Coogan, Yashar Hezaveh, Guillaume Lajoie, and Laurence Perreault Levasseur for helpful suggestions in the course of this project and Mandana Samiei for comments on a draft of the paper.
The authors acknowledge funding from CIFAR, NSERC, IVADO, UNIQUE, FACS Acuité, NRC Al4Discovery, Samsung, and Recursion.
The research was enabled in part by computational resources provided by the Digital Research ca), Mila (https://mill quebec), and NVIDIA.
Alliance of Canada
References
[1] Alexandre Adam, Adam Coogan, Nikolay Malkin, Ronan Legin, Laurence Perreault-Levasseurr Yashar Hezaveh, and Yoshua Bengio. Posterior samples of source galaxies in strong gravitational lenses with score-based priors. arXiv preprint arXiv:221 1.08812, 2022.
[2] Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? International Conference on Learning Representations (ICLR), 2023.
[3] Tara Akhound-Sadegh. Jarrid Rector- Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Niko- lay Malkin, and Alexander Tong. Iterated denoising energy matching for sampling from Boltzmann densities. International Conference on Machine Learning (ICML), 2024.
[4] Jacob Austin, Daniel D Johnson, Jonathan Ho, Daniel Tarlow, and Rianne Van Den Berg. Structured denoising diffusion models in discrete state-spaces. Neural Information Processing Systems (NeurIPS), 2021.
[5] Yoshua Bengio. Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet foundations. Journal of Machine Learning Research, 24(210):1-555 2023.
[6] Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion- based generative modeling. Transactions on Machine Learning Research (TMLR), 2024.
[7] Julius Berner, Lorenz Richter, Marcin Sendera, Jarrid Rector-Brooks, and Nikolay Malkin. From discrete-time policies to continuous- diffusion samplers: Asymptotic equivalences and faster training. arXiv preprint arXiv:2501.06/44 2025.
[8] Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine Training diffusion models with reinforcement learning. International Conference on Learning Representations (ICLR), 2024.
[9] Andrew Camphell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. continuous time framework for discrete denoising models. Neural Information Processing Systems (NeurIPS). 2022.
[10] Gabriel Cardoso, Yazid Janati el idrissi, Sylvain Le Corff, and Eric Moulines. Monte carlo guided denoising diffusion models for bayesian linear inverse problems. International Confer- ence on Learning Representations (ICLR), 2024.
[11] Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Mare Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. International Conference on Learning Representations (ICLR), 2023.
[12] Prafulla Dhariwal and Alexander Nichol. Diffusson models beat GANs image synthesis. Neural Information Processing Systems (NeurIPS), 2021.
[13] Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, al. Continu- ous diffusion for categorical data. arXiv preprint arXiv:2211.15000 2022.
11

---

[14] Zehao Dou and Yang Song. Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. International Conference on Learning Representations (ICLR), 2024.
[15] Yilun Du, Conor Durkan, Robin Strudel, Joshua B. Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC. International Conference Machine Learning (ICML), 2023.
[16] Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Reinforcement learning for fine- tuning text-to-image diffusion models. Neural Information Processing Systems (NeurIPS), 2023.
[17] Berthy T. Feng, Jamie Smith, Michael Rubinstein, Huiwen Chang, Katherine L. Bouman, and William T. Freeman. Score-based diffusion models as principled priors for inverse imaging. International Conference on Computer Vision (ICCV), 2023.
[18] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
[19) Timur Garipov, Sebastiaan De Peuter, Ge Yang, Vikas Garg, Samuel Kaski, and Tommi Jaakkola. Compositional sculpting of iterative generative processes. Neural Information Processing Systems (NeurlPS), 2023.
[20] Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv:1807.01622 2018
[21] Tomas Geffner, George Papamakarios, and Andriy Mnih. Compositional score modeling for simulation-based inference. International Conference on Machine Learning (ICML), 2023.
[22] Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-plas priors. Neural Information Processing Systems (NeurIPS), 2022.
[23] Ishaan Gulrajani and Tatsunori B Hashimoto. Likelihood-based diffusion language models.
Neural Information Processing Systems (NeurIPS), 2023.
[24] Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov. SSD-LM: Semi-autoregressive simplex- based diffusion language model for text generation and modular control. Association for Computational Linguistics (ACL), 2023.
[25] Philippe Hansen-Estne Eastr,, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. IDQL: Implicit Q-learning preprint arXiv:2304. 1057,, 2023.
actor-critic method with diffusion policies. arXiv
[26] Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding-enhanced bert with disentangled attention. International Conference on Learning Representations (ICLR), 2021.
[27] Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Neural Information Processing Systems (NeurIPS), 2020.
[28] 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. Interna- tional Conference on Learning Representations (ICLR), 2022.
[29] Edward J. Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, and Nikolay Malkin. Amortizing intractable inference in large language models. International Conference on Learning Representations (ICLR), 2024.
[30] Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffuusion for flexible behavior synthesis. International Conference on Machine Learning (ICML), 2022.
[31] Zahra Kadkhodaie and Eero P. Simoncelli. Solving linear inverse problems using the prior implicit in denoiser. Neural Information Processing Systems (NeurIPS), 2021.
12

---

[32] Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. Neural Information Processing Systems (NeurIPS), 2024.
[33] Bahjat Kawar, Gregory Vaksman, and Michael Elad. SNIPS: Solving noisy inverse problems stochastically. Neurul Information Processing Systems (NeurIPS), 2021.
[34] Minsu Kim, Taeyoung Yun. Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, and Jinkyoo Park. Local search GFlowNets. International Conference on Learning Representations (ICI.R), 2024.
[35] Diederik Kingma and Max Welling. Auto-encoding variational Bayes. International Confer- ence on Learning Representations (ICLR), 2014.
[36] Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit Q-learning. International Conference on Learning Representations (ICIRR) 2022.
[37] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. Neural Information Processing Systems (NeurIPS), 2020.
[38] Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-Garde Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of continuous generative flow networks. International Conference on Machine Learning (ICML), 2023.
[39] Sergey Levine, Aviral Kumar, G. Tucker, and Justin Fu. Offline reinforcement Icarmin Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.016433 2020.
[40] Xiang Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori Hashimoto. Diffusion-LM improves controllable text generation. Neural Information Processing Systems (NeurIPS), 2022.
[41] Nan Liu, Shuang Li, Yilun Du, Antonio Torralba, and Joshua B Tenenbaum. Compositional visual generation with composable diffusion models. European Conference on Computer Vision (ECCV), 2022.
[42] Yang Liu, Dan Iter, Yichong Xu, Shuohang Wang, Ruochen Xu, and Chenguang Zhu, G-Eval: NLG evaluation using GPT-4 with better human alignment. arXiv preprint 2rX1v:2303.166:: 2023.
[43] Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint 1xXv:23100.18344 2023.
[44] Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning. International Conference on Machine Learning (ICML), 2023.
[45] Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bose, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. International Conference on Machine Learning (ICML), 2022.
[46] Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment GFlowNets. Neural Information Processing Systems (NeurIPS), 2022.
[47] Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference. International Conference on Learning Representations (ICLR), 2023.
[48] Morteza Mardani, Jiaming Song, Jan Kautz, and Arash Vahdat. A variational perspective on solving inverse problems with diffusion models. International Conference on Learning Representations (ICLR), 2024.
13

---

[49] Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete score matching: Generalized score matching for discrete data. Neural Information Processing Systems (NeurIPS), 2022.
[50] Nasrin Mostafazadeh, Nathanael Chambers, Xiaodong He, Devi Parikh, Dhruv Batra, Lucy Vanderwende, Pushmeet Kohli, and James Allen. A corpus and cloze evaluation for deeper understanding of commonsense stories. North American Chapter of the Association for Compu- tational Linguistics: Human Language Technologies (NAACL-HLT), 2016.
[51] Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabili models. International Conference on Machine Learning (ICML), 2021.
[52] Nikolas Nüsken and Lorenz Richter. Solving nigh-dimensional Hamilton- Jacobi -ellmmn PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial Differential Equations and Applications, 2(4):48, 2021.
[53] Bernt Øksendal. Stochastic Differential Equations: An Introduction with Applications. Springer, 2003.
[54] Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: method for automatic evaluation of machine translation. Association for Computational Linguisties (ACL), 2002.
[55] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. advantage-weigghtn regression: Simple and scalable off-policy reinforcement learning. arXiv prepring arXiv:1910.00117, 2019.
[56] Ben Poole, Ajay Jain, Jonathan Barron, and Ben Mildenhall. DreamFusion: Text-to-3D using
2D diffusion. International Conference Learning Representations (ICLR), 2023.
[57] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 188)99
[58) Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. 2021.
[59] Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. Interna- tional Conference on Machine Learning (ICMI), 2015.
[60] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. International Conference Machine Learning (ICML), 2014.
[61] Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco J. R. Ruiz, and Ömer Deniz Aky- ildiz. VarGrad: A low-variance gradient estimator for variational inference. Neural Information Processing Systems (NeurIPS), 2020.
[62] Lorenz Richter, Julius Berner, and Guan-Horng Liu. Improved sampling via learned diffusions. International Conference on Learning Representations (ICLR), 2023.
[63] Robin Rombach, A. Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. Computer Vision and Pattern Recogni- tion (CVPR), 2022.
[64] Simo Särkkä and Arno Solin. Applied stochastic differential equations. Cambridge University Press, 2019.
[65] Marcin Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Rector-Brooks, Alexandre Adam, Yoshua Bengio, and Nikolay Malkin. On diffusion models for amortized inference: Benchmarking and improving stochastic control and sampling. arXiv preprint arXiv:2402.05098, 2024.
[66] Kiyoung Seong, Seonghyun Park, Seonghwan Kim, Woo Youn Kim, and Sungsoo Ahn. Col- lective variable free transition path sampling with generative flow network. arXiv preprint 1901 2024.
14

---

[67] Tiago Silva, Amauri H Souza, Luiz Max Carvalho, Samuel Kaski, and Diego Mesquita. Federated contrastive GFlowNets, 2024. URL https://openreeveer wwtp://nee/eetiiieee.em VJDFhkwQg6.
[68] Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning (ICML), 2015.
[69] Jiaming Song, Chenlin Meng. and Stefano Ermon. Denoising diffusion implicit models. International Conference on Learning Representations (ICLRR, 2021.
[70] 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. International Conference on Maching Learning (ICML), 2023.
[71] Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffuusion models. Neural Information Processing Systems (NeurIPS), 2021.
[72] Yang Song, Jascha Sohk-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations (ICLR), 2021.
[73] Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. International Conference on Learning Represen- tations (ICLR), 2022.
[74] Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. International Conference Learning Representations (ICLR), 2023.
[75] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT Press, 2018.
[76] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference Intelligent Robots and Systems, pages 5026-5033, 2012.
doi: 10.1 10.100//00522222381199999
[77] 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 as entropy-regularize control. arXiv preprint arXiv:2402.15 2024.
[78] Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. Interna- tional Conference on Learning Representations (ICLR), 2023.
[79] Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets variational inference: Controlled Monte Carlo diffusions. International Conference Learning Representations (ICLR), 2024.
[80] Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation, 23(7):1661-166) 2011.
[81] Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, and Shengyi Huang. Trl: Transformer reinforcement learning. attps://github com com//tuggingfcce/trl, 2020.
[82] Haochen Wang, Xiaodan Du, Jiahao Li, Raymond Yeh, and Greg Shakhnarovich. Score Jacobian chaining: Lifting pretrained 2D diffusion models for 3D generation. Conference on Computer Vision and Pattern Recognition (CVPR), 2023.
[83] Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies expres- sive policy class for offline reinforcement learning. International Conference on Learning Representations (ICLR), 2023.
15

---

[84] Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahem, Andrew Borst, Robert J Ragotte, Lukas Milles, al. De novo design of protein structure and function with rfdiffusion. Nature, 620(7976):0089-1100
2023.
[85] Sean Welleck, Ilia Kulikov, Jaedeok Kim, Richard Yuanzhe Pang, and Kyunghyun Cho. Consis- tency recurrent language model with respect to incomplete decoding. Empirical Methods in Natural Language Processing (EMNLP), 2020.
[86] Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, et al. Google's neural machine translation system: Bridging the gap between human and machine translation. arXiv preprini arXiv:1609.08; 144, 2016.
[87] 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. Neural Information Processing Systems (NeurIPS), 2023.
[88] Dinghuai Zhang, Ricky Tian Qi Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffusion generative flow samplers: Improving learning signals through partial trajectory optimization. International Conference Learning Representations (ICLR), 2024.
[89] Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sampling. International Conference on Learning Representations (ICLR), 2022.
[90] Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, and Yoav Artzi. BERTScore: Evaluating text generation with bert. International Conference on Learning Representations, 2020.
[91] Wanrong Zhu, Zhiting Hu, and Eric Xing. Text infilling. arXiv preprint arXiv: 1xXiv:1911.0015: 2019.
16

---

A Proofs
Proposition (Relative TB constraint). pe, ps and the scalar Z jointly satisfy the relative
trajectory balance (RTB) constraint
ZΦ
P& post (X0, Nar...., rxx) XAI. .AI....,XI)
(8)
for every denoising trajectory Xo XA 7...- X1, then po (x1) o po(x))r(x)), Le., the diffusion model pe samples the posterior distribution. Furthermore, if Po also satisfies the TB constraint (7) with respect to the noising process q and some target density ((x)), then p satisfies the TB
constraint with respect to the target density (x1) o q(x:)r(x1),
and Z q(x))r(x)) dxi-
post and Z jointly satisfy (8). Then necessarily Z # O, since the quantities on the right side are positive, We then have, using (4),
Proof of Prop. I. Suppose that Pa. P
post (x1) = ʃ
post (Xo, Nar..... X1) dNo dxar dx1-a p
((x)) f
Xse............... x1) dxo dxAr ... dxi-Ar
= Z (xi)po((x)
& (o(()))(xx),
desired.
Now suppose that pa also satisfies the TB constraint (7) with respect to q(xi). Then, for any denoising trajectory,
post (xo, NAr...., x1)
(Xo, Nar..... AI-Ar x1)= Po(Xo. Xar. xI)
Po
(14)
((x,)
((x))r(x;)/Z
showing that pr satisfies the TB constraint with respect the noising process q and the (not yet shown to be normalized) density (x)))(x)) We integrate out the variables No, XAr. X1- in (14), giving
P (x1)
(xx)r(x))
((x))r(x)) Zpanr post (x1).
Integrating x shows f (xx))r(xx) dx, Z.
B Relative TB as TB under the prior measure
The theoretical foundations for continuous generative flow networks [38] establish the correctness of enforcing constraints such as trajectory balance (7) for training sequential samplers, such as diffusion models, to match unnormalized target densities. While we have considered Gaussian transitions and identified transition kernels with their densities with respect to the Lebesgue mea- sure over Rdd these foundations generalize to more general reference measures. In application to diffusion samplers, suppose that x (x,) is collection of bebesgue-absolutely continuous den- sities over Rd for Ar..... 1 and that Fee((x) X,-a)), (X x,) are collections of aabessue-absolutter continuous transition kernels. If these densities jointly satisfy the detailed balance condition (x, ) iTr (N x,) Jeef (x-Arr 7 (X, x;-ar), then they satisfy the con- ditions to be reference measures. A main result of [38] is that if a pair of forward and backward processes satisfies the trajectory balance constraint (7) jointly with a reward density r. then the forward process samples from the distribution with density with all densities interpreted as relative to the reference measures Tref, x ref. π ref.
Recall that the relative density (or Radon-Nikode Nikodym derivative) of distribution with density under the Lebesgue measure relative to with density is simply the ratio of densities р/n.
17

---

If Po is a diffusion model that satisfies the TB constraint jointly with some reverse process q and target density (xx)), then one can take the reference transition kernels π ret. ref to be and respectively. In this case, the TB constraint for a target density (x1) and forward transition pe is
P post XX, XAI. x1)
2r(x1)9 ) No. XA
x1)
(15)
ref No. Xàr
X1) (Ka -r x1)
which is identical the RTB constraint (8). If (15) holds, then p samples from the distribution
with density *)(x1) relative to (x)), which is exactly ( We have thus recovered RTB as a case of TB for non-Lebesgue reference measures.
18

---

C Posterior inference on two-dimensional Gaussian mixture model
Setup We conduct toy experiments in low-dimensional spaces using samples from Gaussian mixture model with multiple modes to visually demonstrate its validity. The prior distribution p(x) trained on a Gaussian mixture model with 25 evenly weighted modes, while the target posterior plos (x)))p(x) uses a reward r(xi) to select and re-weight 9 modes from pxx)). More specifically, the resulting posterior is:
(1)) px))(xx))
n, - xx I).
1
(16)
(mi) = (000--5 (-5. (-5.0), (10, (1), (0.0). (5.0), (5, 10)) (f:) = (4, 4,5, 10, 15,4)
(17)
(18)
Our objective is to sample from the posterior (x1). We compare our method with several baselines, including policy gradient reinforcement learning (RL) with KL. constraint and classifier- guided diffusion models. For RL, we implemented the REINFORCE method with mean baseline and KL constraint, following recent work training diffusion models to optimize a reward function [8]. Sampling according to the RL policy leads to a distribution qa qa (x1), which is trained with the objective:
((0) Equ(x)) [r(xi)] + aDRL (qe(x))ll((x))
(19)
While the exact computation of KL (qo(x))llx(x)) is intractable, we follow the approximation method introduced by Fan et al. [16], which sums the divergence at every diffusion step. This approximation optimizes an upper bound of the marginal KL.
The other baseline is classifier (energy) guidance, which given a diffusion prior, samples using posterior score function estimate:
Vx, log (x,) Vx, log p(x,) + Vx, logr(x,)
(20)
Note that this a biased approximation of the true intractable score:
Xx, log prost (x,) xx, log p(x,) + vx, log bx [r(x)]
(21)
For our experiments, we follow the source code* provided in recent diffusion sampler benchmarks [65]. We utilize a batch size 500, with finetuning at 5,000 training iterations, a learning rate of 0.0001, diffusion time scale of 5.0, 100 steps, and log variance range of 4.0. The neural architecture employed is identical to that used in [65]. For pretraining the prior model, we use the same hyperpa- rameters above, but with 10,000 training iterations using maximum likelihood estimation with true samples.
Results. As we reported in the main text, in Fig. 1, we present illustrative results. The classifier- guided diffusion model shows biased posterior sampling (Fig. If), failing to provide accurate inference. RL with per step KL constraint cannot exactly optimize for the posterior distribution, making the tuning of the KL weight a crucial to achieving desirable output Fig. C.I. RTB asymptotically achieves the true posterior without introducing balancing hyperparameter a. Another advantage of our approach is off off-policy exploration for efficient mode coverage. RL methods for tie-ttnnng diffusion models (c.g., DPOK [16]. DDPO [8]) typically use policy gradient style methods that are on-policy. By using a simple off-policy trick introduced by [47, 38) and demonstrated by Sendera et al. [65], we can introduce randomness into the exploration process in diffusion by adding &. where is a noise hyperparameter and T is the diffusion timestep, into the variances and annealing it to zero over training iterations. We set 0.5 for off- off-policy exploration. As shown in Fig. C.2, RTB with off off-colicy exploration gives very close posterior inferences, whereas off-policy exploration in RL with 0.5 (which is carefully selected hyperparameter) does not improve performance due to its on- on-policy nature.
D Code
Code for all experiments is available at httts::/giithu/ http:///tihtttcco///////////aaaa hfrp://ggdddfffusosooooommuunnnn and will continue to be maintained and extended.
htttps:/ htattps://gtthubb wow.//POOg//gtnddiff.oooo.....ee
19

---

(a) Prior
(b) Posterior
(c) a 0.8
(e)==0)) a 0.5
(g) a =
Figure C.I: Tuning the KL weight in reinforcement learning: influences the balance between sticking to the prior distribution and moving towards the modes of the reward density. A higher a value maintains closer adherence to the prior, while lower allows gradual shift towards high values of rxx). Setting a below 0.3 tends to cause mode collapse, moving too far from the prior and focusing on maximizing rewards for single modes. a 0.5 gives samples that closest resembles the posterior.
(a) Prior
(b) Posterior
(c) RTB (on)
(d) RTB (off)
(c) RL (on)
(f) RL (off)
Figure C.2: Off-policy exploration benefits for RTB training. RTB, with simple off-policy exploration techniques that increase randomness in the diffusion process, significantly improves mode coverage.
On the other hand, policy gradient RL methods which are typically used to finetune diffusion models on-policy, and hence prone to mode collapse.
E
On classifier guidance and RTB posterior sampling
MNIST
CIFAR-10
Even
Ship
Car
Cat
Dog
Figure E.1: Samples from posterior model fine-tuned with RI. (no KI.). We observe early mode collapse, showcasing high-reward samples with minimal diversity.
E.1 Experimental Details
In our experiments, we fine-tune pretrained unconditional diffusion models with our RTB objective, to sample from a posterior distribution in the form prost(x (x y) p(y xx)(xx. In this section, detail the experimental settings for RTB as well as the compared baselines.
Experiments setting. For MNIST, we pretrain a noise-predictiin diffusion model on 28 x 28 (upscaled to 32 x 32) single channel images of digits from the MNIST datasets. We discretize the forward and backward processes into 200 steps and train our model until convergence. For CIFAR-10. we use pretrained model from [27], trained to generate 32x32 3-channel images from the CIFAR-10 dataset, while discretizing the noising/denoising processes into 1000 steps. For fine-tuning the prior, parametrize the posterior with LoRA weights [28], with the number of parameters equal to about 3% of the prior model's parameter count. We train our models on a single NVIDA V100 GPU.
We compute FID as a similarity score estimate of the true posterior distribution from the data. As such, the computation is limited to the total number of per-class-samples present the data, (between 5k and 6k for CIFAR- 10 and MNIST digits, and 30k for the even/odd task).
20

---

RTB. For RTB fine-tuning, we finetune a diffusion model following the objective in Equation 9. We impose the objective while sampling denoising paths following DDPM sampling scheme, with only 20% to 50% of the original trained steps. We employ loss clipping at 1. to account for imperfect constraints in the pretrained prior, and train each of our models for 1500 training iterations, well into convergence trends.
RL [16]. We implement two RL-based fine-tuning techniques derived from DPOK [16] and DDPO [8], respectively with and without KL regularization. These implementations use reinforce- ment learning baseline similar to the one in our experiments described in $3.2. By following the same sampling scheme as in our RTB experiments, we enable a direct comparison with RTB. To fine-tune the KL weight, we perform a search over a E {0.01, 0.1, 1.0).
DP [11]. We implement and adapt the Gaussian version of the posterior sampling scheme in [11], originally devised for noisy inverse problems. This method relaxes some of our experimental constraints, it requires differentiable reward r(x). We perform sweep over ten values of the suggest parameter range for the step size & E [.1, 1.J on MNIST single-digit sampling, and choose [ == 0.1 for our experiments.
LGD-MC [70]. We adapt the implementation of the algorithm in [70] to sample from the classifier- based posteriors in CIFAR-10 and MNIST. Similarly to the DP baseline, we use our pretrained classifier to perform measurements at each sampling step, and use Monte Carlo estimate of the gradient correction guide the denoising process. We choose & 0.1 following the DP experiments and default the number of particles to 10 as per the authors' guidelines.
E.2 Additional findings.
Mlasifier-guiddnt baselines. We find that the DP and LGD-MC classifier-guidanee based base- lines struggle to sample from the posterior distribution in our experimental settings. The baselines achieve the lowers classifier average rewards in all tested settings. Despite choosing & = 0.1 as the validate best performing hyperparameter, we also also observe the posterior samples from DP and LGD-MC to be close to the prior. As such, DP and LGD-MC score high in diversity, and low in FID for the Even/Odd experimental scenario, as expected from prior sampling benchmarks, but failing to appropriately model the posterior distribution.
RL and mode collapse. In the pure Reinforcement Learning objective imposed for the experiments in $3.1 (no KL), we observe significantly higher reward than other baseline methods, while showcasing increased FID and lower diversity. In Fig. E.1 we show a random set of 16 samples for posterior models trained on 4 different classes of the CIFAR-10 10 datasets, as well as the Even objective from the MNIST dataset, after 500 training iterations. and reward exploitation, visually evident from the little to no variation amongst samples for each class class, and single-digit collapse in the multi-modal even digits objective (see samples in Fig. for comparison with our RTB-finnttued models).
the figure, we observe early mode collapse
Conditional architectures. We repeat the even/odd posterior sampling experiment of $3.1 in a conditional setting, where the condition is an input the posterior model. For the posterior architecture, we use a naive modification of the prior with extra input channel, which is pop- ulated a full mask of 0 or for conditioning on the even and odd classes, respectively. The re- sults are shown in Table 1. We look forward to future work which develops more specialized ar- chitectures for handling conditional constraints.
Table E.1: Conditional experiment for MNIST even/odd posterior. Note that the posterior model in the conditional experiment is different from that in the baselines because it uses a different archi- tecture that includes additional input channel.
Dataset
MNIST even/odd
Algorithm Metric
Bxlogr(x)] (1)
FID (L)
DPS
-1.2270=amm2
1.1498-0.182
LGD-MC
-1.1720
1.144520.14
DDPO
-8.62123x10-11
1.8024..423
DPOK
-0,0783a0022
1.2536=0.206
RTB (unconditional)
-0.1816aa175
1.1794=0.171
F Infilling with discrete diffusion
RTB (conditional)
-0.1236
0.9112
Additional details. We illustrate some examples from the ROC Stories dataset used for training
in Table F.I. For the prior we use the sedd-small model, which uses an absorbing noising process [4] with a log linear noise schedule, as the diffusion prior p(z x). The posterior model is
'ntrps://hugginggace http:///........oooooooa//adddee small
21

---

Table F.1: Examples of training samples for the language infilling task.
Beginning (x)
Middle (z)
End (y)
I going a Halloween party. looked through my clothes but could not find costume. cut up my old clothes and constructed costume.
put my costume on and went to the party.
My friends loved my cos- tume.
Allen thought he poet. He attended college to study cre- ative writing. In college, he met boy named Carl.
very talented
Carl told him that he wasn't very good.
Because of this, Allen swore off poetry forever.
parameterized copy of the prior. To condition the diffusion model the beginning x set the tokens at the appropriate location in the state in the initial time step, i.e. = O. Our implementation is based on the original SEDD codebase 6. Training this model is computationally expensive (in terms of memory and speed) so we utilize the stochastic TB trick, only propagating the gradients through subset of the steps of the trajectory. We also use the loss clipping trick as discussed in $2.3. Specifically, we clip the loss below a certain threshold to O, resulting in updates only when the loss is larger. This threshold referred to as the loss clipping coefficient- - is hyperparameter. As this is a conditional problem we also use the relative VarGrad objective. We also use some tempering the reward likelihood which helps in learning (i.e., Preward inverse temperature parameter. We perform all experiments on an NVIDIA A A100-Large GPU. Note that we also tried baseline of simply fine-tuning the diffusion model on the data but encountered some training instabilities that we could not fix. The hyperparameters used for training RTB in our experiments are detailed in Table F.2.
z) where is the
Reward. For training Preward we follow the training procedure and implementation from [29]². Specifically, we fine-tune a GPT-2 Large model [57] on the stories dataset with full parameter fine-tuning using the trl library [81]. We trained for 20 epochs with batch size of 64 and 32 gradient accumulation steps and a learning rate of 0.0005.
Baselines. For the baselines, we adopt the implementations from [29]. critical difference in our experiments compared to [29] is that the posterior model is not initialized with base model that is fine-tuned on the stories dataset. To condition the model on and Y, as well as for the prompting baseline, we use the following prompt:
"Beginning: <X3Nn End: <Y>n Middle:
"
During training for the autoregressive GFlowNet fine-tuning, a (x. y) pair is sampled from the dataset and then sample (batch size) xs for every (X,Y), and Preward (eewarddXZi (XZY) used the reward. Both the GFlowNet fine-tuning and supervised fine-tuning baseline use LoRA fine-tuning. We use the default hyperparamteres from [29]. At test time, sample 100 infills for each example in the test set from all the models at temperature 0.9, and average over 5 such draws.
Additional results. Table F3, Table F.4 and Table F.5 illustrates some examples of the infills generated by the diffusion models. We note that the general quality of the samples is poor, due to a relatively weak prior. the same time we can observe that the prompting baselines often generate infills that are unrelated to the current story. We also note that the RTB fine-tuned model can sometimes generate repitions as the reward model tends to assign high likelihood to repititions [85]. We also attempted LLMEval [42] for evaluating the coherence of the stories but did not obtain statistically significant results.
22

---

Table F.2: Hyperparameters for the story infilling task.
Batch size
16
Gradient accumulation steps
8
Learning rate
le-5
Warmup Step
20
Optimizer
AdamW
Reward temperature start
1.2
Reward inverse temperature end
0.9
Reward inverse temperature horizon
5000
Number of training steps
1500
Loss clipping coefficient Discretization steps T
0.1
15
Table F.:: Examples of infills generated by the posterior trained with RTB along with reference infills for the stories infilling task.
Beginning (x)
Middle (z)
End (y)
David noticed he had put on a lot of weight recently. He examined his habits to try and figure out the reason. He realized he'd d been eating too much fast food lately.
He stopped going to burger places and started a vegefarian diet.
After few weeks, he started to feel much better.
He reviewed his habits to try to figure out how to change He asked he thought try to cut down on the amount amount.
He examined his habits to try and figure out the reason. He realized had been eating too much fast food recently.
Robbie was competing in a cross country meet. He was halfway through when his leg cramped up. Robbie wasn't sure he could go on.
He stopped for a minute and stretched his bad leg. Robbie was sure he could go on. Robbie was sure.
Robbic began to again and finished the race in second place.
He was floating and twisting his leg in half then.
His body just caught up with his legs. Robbie was.
He held his leg forward as he went through and his
G
Offline RL
G.1
Training details
Our method requires first training diffusion-based behavior policy By and a Q-function es. Once To and Qn are trained, The posterior policy πy is trained using RTB, with its weights initialized to the trained behavior policy weights
The behavior policy Tg is parametrized as state-conditioned noise-predicting denoising diffusion probabilistic model (DDPM) [27] with linear schedule, and 75 denoising steps. The diffusion model takes as input a state s, noised action a, and a noise level and predicts the source noise e. The state and noised action a, are concatenated with Fourier features computed on the noise level which are then fed through 3-layer MLP of hidden dimensionality 256, with layer normalization and GeLU activation after each hidden layer. The behavior policy trained using the Adam optimizer with batch size 512 and learning rate 5e-4 for 10000 epochs. The Q-function Qu is trained using IQL. We use the same IQL experimental configurations and training hyperparameters as [36]. That is, we set r 0.7. The architecture for Qv is 3-layer MIP with hidden dimensionality 256 and ReLU activations, which is trained using the Adam optimizer with learning rate 3e-4 and batch size 256 for 750000 gradient steps. The task rewards are normalized as in [36] and the target network is updated
23

---

Table F.4: Examples of infills generated by Prompt (x.y) along with reference infills for the stories infilling task.
Beginning (x)
Middle (z)
End (y)
David noticed he had put on lot of weight recently. He examined his habits to try and figure out the reason. He realized he'd been eating much fast food lately.
He stopped going to burger places and started a vegetarian diet. He'd d had less opportunities to eat prop- erly all of last.
After a few weeks, he started to feel much better.
Doctors made the note of the situation. He was treated.
He told him the guy for mic replacement. An In He felt empty for one reason and new fresh, too.
Robbie was competing in cross country meet. was halfway through when his leg cramped up. Robbic wasn't sure he could go
He stopped for a minute and stretched his bad leg.
Robbie began to run again and finished the race in second place.
Robbie wasn't sure Robbie's fuel tank was full.
Robbie took photograph with close friend. An
Only Stacey Ebers and Rand out there.
Robbie got bigger as the position got better. Än \n
Table F:5: Examples of infills generated by Prompt (x) along with reference infills for the stories infilling task.
Beginning (x)
Middle (z)
End (y)
David noticed he had put on lot of weight recently. He examined his habits to try and figure out the reason. He realized he'd d been eating too much fast food lately.
He stopped going to burger places and started a vegetarian diet. David, "All had told eat was problem. He got the backside what about that and he made the,
After few weeks, he started to feel much better.
He made just good of fast food and spliced it down.
He explained everything to them, reached them out, the problem.
Robbie competing in cross country meet. He was halfway through when his leg cramped up. Robbie wasn't sure he could go
He stopped for minute and stretched his bad leg.
Robbie began to run again and finished the race in second place.
Robbie and Robbie piling. Robbie and Robbie fistfight. said goodbye. Robbie dinner. Rob- bie agreed with.
He cut away a little Robbie's pace fleetingly.
He held off all the police and place. Rob- bie.
24

---

Table G.I: Mixed vs. online training on DR4L Tasks. We report meanasid over random seeds.
Task
RTB (Online)
RTB (Mixed)
puafhheetah-mmddiu--eeaaaa
46.88
48. 1 100.5
hopper- seteee-meiiim--ere -replay
99.23x3.22
100.40a0.21
walker2d-medium-ae
94.01
93.57=2.63
Table G.2: Temperature a
for D4RL tasks
Task
0.1
expert
0.5
walker2d-medium-eexe
0.1
halfcheetah-medium
0.05
hopper-mediie
0.1
fiiker2d-mediun
0.05
halfcheetah medium-replay
0.05
0.05
walker2d mallerrd-mediuu--replayy
0.1
with soft updates of m = 0.005. The posterior policy Ty is trained using the relative trajectory balance objective. Ty is also parametrized as caneiconddtioned noise- poise-predicting DDPM. initialized as a copy of the prior. We additionally use the Langevin dynamics inductive bias (11), and learn an additional MLP for the energy scaling network. The posterior noise prediction network also outputs an additive correction to the output of the prior noise prediction network. That is, the predicted noise of the posterior diffusion model defined as E(s, dr, t) == E(s, a,, f; 0) + E(s, drs Y), where e(:: the output of the prior noise prediction network and é(:; y) is the output of the posterior noise prediction network. We train all models single NVIDIA A A10ge GPU. The only hyperparameter tuned per task is the temperature a which we show in Table G.2.
Note that in these experiments both the prior and constraint are conditioned the state s. To prevent having to learn a neural network for log Z (s), we employ a variant of VarGrad objective [61]. For each state sampled in the minibatch, we further generate 64 on-policy trajectories (()) Ty. Each of these trajectories can be used to implicitly estimate log 7(s):
i=l with
log (s)) (i) = log Te ( (i) s) + Q (s, ) - Iog (T |s)
(22)
We then minimize the sample variance across the batch:
VarGrad (r) LTT
Σ log 2(s)
1 Σ log 2(s) (1)
(23)
RTB allows off-policy training so we are not restricted to train with samples generated on-policy. We thus also leverage the offline dataset, which samples from the prior and noise them with the DDPM noising process to generate off-policy trajectories with high density under the prior. Since there are actions in the replay buffer from high reward episodes in the tasks, this can help training efficiency compared to purely online training. We ran 5 seeds of training each with mixed training (off-policy and on-policy) and pure on-policy training the medium- replay tasks, with results shown in Table G.I, where mixed training outperforms pure online training on two of the three tasks.
G.2 Baseline details
As is standard in offline RL, we use the reported performance numbers from the previous papers. CQL. IQL are reported from the IQL paper. Diffuser (D), DD, D-QL and QGPO are reported from the QGPO paper. Their implementation improved the performance of and D-QL compared to their original papers. IDQL results are reported from the IDQL paper. We follow the evaluation protocol of previous work, and report the mean performance over 10 episodes, averaged across random seeds at the end of training (150k training steps).
25

---

H Fine-tuning text-to-image diffusion models
We build off the DPOK implementation", which fine-tunes stable-diffusion-vi- with ImageReward function. The posterior model to be fine-tuned is initialized as a copy of the prior model. We use LoRA [28] since it significantly more efficient than fine-tuning the entire model. Sampling of images is done with 50 steps of DDIM [69]. Even with LoRA, it is still difficult to fit gradients of all steps in the diffusion trajectory memory. To help with this, we use "stochastic subsampling" trick (SH.I).
We train all models single NVIDIA A A100-Large GPU. For the main experiments, default parameters for DPOK of reward weight & 10 and KI. weight = 0.01. For RTB we fix 1.0 for all prompts. We next perform ablation over different values of B.
the
We plot in Table H. the final average reward and diversity score for models trained with different values of reward weight B for the prompt "A green colored rabbit.". As expected, we find that increasing & increases reward at the cost of diversity for RTB and DPOK. The exception is B = 10 for RTB which has slightly lower final reward than 1, which could attribute to more difficult optimization due to the peaky distribution associated with higher reward weight.
Table H.I: Ablation of reward weights for "A green colored rabbit.".
Model J
B | Metric
Reward (f)
diversity (f)
Prior
-0.113
0.597
B= 0.01
-0.27
0.1488
DPOK (KL 0.0g
f= 0.1
-0.06
0,1486
f= 1.0
0.638
0.1362
# = 10.0
1.492
0.076
DDPO (KL weight=0.0)
A = 10.0
1.795
0,0493
A = 0.01
0.485
0.1431
#=0.1
1.525
0.0721
RTB
$= 1.0
1.756
0.0436
0 10.0
1.568
0.0689
H.1 Memory--fficient learning
We propose two methods to reduce the memory requirement of RTB fine-tuning.
Stochastic subsampling. The expected gradient of the RTB objective (9) is unaffected by propagat- ing gradient randomly sampled subset of the timesteps in trajectory and rescaling by the inverse proportion of timesteps sampled. Stochastically subsampling timesteps for gradient propagation in this way can significantly decrease memory consumption because computation graphs for the remaining timesteps do not need to be maintained; however, such subsampling increases gradient variance, it is preferable keep gradients for as many timesteps possible to fit in memory. For our text-to extttto-mmage experiments, we found sampling 8 timesteps out of 50 to keep gradients was sufficient.
Batched gradient computation. An important property of the RTB objective is that computing its gradient does not require storing the computation graph of all timesteps. The gradient of the RTB objective for single trajectory is just the sum of per-step likeliielddood gradients scaled by the RTB residual:
Z
T
p
post
(XIsr
X(i-1)xs))
V LKIB(TTT = log
+
log
post log p (xiar ((i-1)Ac).
rxx))
( X(i-1)A))
i=1l
Because the likelihood gradients can be accumulated during the forward pass, this allows for batched gradient accumulation version of the update. For trajectory length (number of diffusion steps) T and accumulation batch size (number of time steps receiving a gradient signal in each backward pass) B, the number batched forward passes required scales as B.
g
"nttps://githur http://ggttuuuuc.ggggggg//// htw.ggooggg--eeaarrcc//oooggoooo
26

---

Only the accumulation batch size B, not the trajectory length T. is constrained by the memory budget. This means we can easily scale training with large number of diffusion steps without increasing the variance of the gradient through stochastic subsampling, with training time growing linearly with number of time steps under a fixed memory budget. Although this method is not used in the main experiments presented here, preliminary experiments confirm these observations.
We highlight that both methods are not applicable to diffusion samplers based on differentiable simulation (e.g., PIS and DDS), which need to store the entire computation graph of SDE integration. For these methods, the memory requirement scales linearly with the trajectory length.
27

---

H.2
Generated images
H.2.1 A green-colored rabbit
Figure H.l: Prior
Figure H.2: DDPO
Figure H.3: DPOK
Figure H.4: RTB
H.2.2 Four roses
Figure H.5: Prior
Figure H.6: DDPO
Figure H.7: DPOK
Figure H.8: RTB
28

---

H.2.3 A cat and a dog
Figure H.9; Prior
Figure H.10: DDPO
Figure H.11: DPOK
Figure H.12: RTB
H.2.4
A half- masked rugged laboratory engineer with cybernetic enhancements from a distance, scifi character portrait by greg rutkowski, esuthio, craig mullins.
Figure H.13: Prior
Figure H.14: DDPO
Figure H.15: DPOK
Figure H.16: RTB
I Compute resources
For classifier guidance experiments $3.1 we use train on single NVIDIA V100 GPU. For text- conditional image generation $3.2, text infilling SF and offline SG, we use single NVIDIA A100 large GPU. The total estimated compute time our experiments is 3000 hours.
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We discuss theoretical claims about the RTB ojective in $2, and report experimental results in $3.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize other settings.
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the limitations of our proposed method in $5.
Guidelines:
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: See $A.
Guidelines:
The answer NA means that the paper does include theoretical results. the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We have provided code to reproduce our experiments in $D, and described training details in SC, $E, $F, &G and $H.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
30

---

If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Code is provided $D and can be used to reproduce our main results. We will be releasing our code publicly..
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
httee/uuiide//oeeeubniiiieee cassssonPooiyyy Policyy for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ mnips.ccppbbli. http://ggiieeeoeeeetttiiiiiiiiia /ot///ooodoSssseeeeeee www.bbbisssioooooooooooooo
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
31

---

Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification:
Training details are outlined in &C, SE, SF, and &G.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All experimental results have error bars, except the Stable Diffusion finetuning experiment $3.2 which was only trained on one seed per prompt due to compute constraints.
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative етror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resources used in experiments is outlined in SE, $F, $G, $H and summarised in ŞI.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
32

---

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https:/ httpa/////uuips...eeeee www.eripp..cpppuuiiieeeeee cctppblic//eteeecceececcc (etpllic/Ehhchhuieelliiiiie
Answer: [Yes]
Justification:
The paper conforms to the stated ethics guidelines.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss broader impact of the work $5.
Guidelines:
The answer NA means that there no socictal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release models that have risk of misuse.
Guidelines:
The answer NA means that the paper poses such risks.
33

---

Released models that have high risk for misuse or dual-use use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Creaters of code use are properly credited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's 'ss creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets. Researchers should communicate the details of the dataset/code/moddl part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: We not have any crowdsourcing experiments or research with human subjects.
34

---

Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not have any experiments with human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
35


Reviews:
[{"review_id": "jxIKQIgtkh", "reviewer": "Reviewer_4zDq", "summary": "This paper proposed a method to train a posterior $p^{post}(x)$ given a prior distribution $p(x)$ and some additional (possibly unnormalized) constraint function $r(x)$, when the prior is a pretrained diffusion model. Through the choice over $r(x)$, this setup can capture a wide range of tasks. The authors propose *relative trajectory balance* (RTB) constraint, which, if satisfied by the posterior under training, guarantees that $p^{post} \\sim p(x)r(x)$. By training $p^{post}(x)$ to satisfy the RTB constraint, one can recover the desired posterior. Importantly, this training objective allows for efficient computation of the loss and off-policy training. The paper provides convincing empirical results on a variety of tasks ranging from conditional generation to offline RL.", "questions": "* For the MNIST even-digit experiment, why was the reward function chosen to be $max_{i \\textrm{ even}} p(c=i \\mid x)$? Wouldn't $\\sum_{i \\textrm{ even}} p(c=i \\mid x)$ be a more reasonable reward? I'm curious if this particular choice of reward makes it more likely for certain methods to mode collapse (e.g. by maximizing the probability for a single even class), and whether the results in Table 2 would change for MNIST with a different $r(x)$.\n* Could you share some more details of how LoRA was used for MNIST/CIFAR setup (e.g. which weights were LoRAed)?  Also I'm curious what happens if you don't use LoRA and train the posterior on full set of weights (from prior weight or from scratch), as I LoRA could act as a regularizer.\n* How long is the ratio between prior vs. posterior training in terms of flops or wall-clock time?  How quickly the posterior can be fine-tuned could limit the applicability of this method on certain real-world tasks.", "limitations": "Yes. For concerns and potential limitations, see the comments and questions above.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "**Motivation and Theoretical Justification**\n* The paper provides a theoretical justification (Proposition 1) for the proposed method, which is a natural extension of the trajectory balance constraint for the setting where we cannot easily sample from $p^{post}$. This is a realistic constraint that one often encounters when trying to perform posterior inference given a pretrained prior. \n\n**Empirical Analysis**\n* While it's clear that the formulation does capture a wide range of tasks, it is nice to see an experimental evidence of that. The generality of the proposed method is emphasized through the results on four different tasks.\n* I appreciate that the authors took time to discuss some of the implementation details (Sec 2.3).\n\n**Writing Quality**\n* The presentation is very well-done. The problem being tackled is clear, and the technical aspects of the solution are intuitive and easy to follow.", "weaknesses": "**Scalability**\n* One main concern I have is how scalable the method is due to trajectory matching, especially as the prior becomes larger.  The main benefit of this technique relies on the assumption that one aims to repurpose an existing diffusion prior, which is presumable a large foundation model -- so it's important that the method scales desirably to the prior's size, which is not clear from the results in the paper.\n  * Section H discusses some techniques employed to reduce memory usage, but even this relatively small-scale experiment (50 steps on latent diffusion with LoRA), the authors report that they could fit only 8 timesteps on an A100.\n  * It'd be very useful to get some idea of how sensitive the method is to these hyperparameters. One extreme case to consider would be to use single-step gradient subsampling for a much larger T (e.g. 1000 steps).\n\n**Experimental Results**\n* Sec 2.3 claims that the proposed method should easily generalize to conditional constraints (which I believe), but there doesn't seem to be any experiment that tests this claim. For example, it'd strengthen the paper to have a result where RTB training successfully extracts a class-conditional model (that takes a conditioning input) from an unconditional one (perhaps for the CIFAR experiment).\n* The text-conditional image synthesis experiment seems to be done for only four hardcoded text prompts, which is pretty limited.\n* While the text in-filling results are great to have, a more practical and interesting task would be fine tuning a pretrained LLM. As hinted in Sec. 2.4, I'm curious if RTB can be used to efficiently customize an LLM, where we deviate from the Gaussian transition kernel.\n* Having a data-dependent reference as an upper bound on the performance of the proposed method would be useful (e.g. classifier guidance trained classifier trained on noisy data, or directly fine-tuning the prior itself on a specific task). This can tell us more about how much of the potential performance the model captures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback and good questions. We are glad you found the paper well-written and the empirical results compelling. We believe that the diversity of domains on which we showed the effectiveness of our approach makes this work both a valuable illustration of off-policy finetuning for intractable diffusion posteriors and a starting point for a number of possible domain-specific applications.\n\nWe have addressed each of your questions below and added three new experiments in response to your comments (please see the response to all reviewers).\n\n### Text infilling and LLM fine-tuning\n\nWe believe there may be a misunderstanding present. The infilling results are for the setting of *fine-tuning a pretrained discrete diffusion language model*, which does not use a Gaussian transition kernel. In discrete diffusion language models, the kernel is categorical at each position. \n\nThe text infilling experiment (Table 3 in the paper) thus supports two claims:\n- that the proposed diffusion fine-tuning algorithm is applicable outside the Gaussian setting (and even outside the continuous setting);\n- that fine-tuning of LLMs for intractable posteriors, such as infilling, using off-policy RL can be generalized beyond autoregressive language models (as done in [28]).\n\n### On scalability and cost\n\nWe answer a few questions related to scaling and fine-tuning cost comparisons.\n\n#### Scalability\n\nRTB actually scales surprisingly well to larger models and more diffusion steps (relative to simulation-based methods) despite being a trajectory-level objective.\n\nAn important observation about the RTB objective -- not included in our original submission -- is that computing the RTB gradient does not require storing the computation graph of all timesteps to be updated. We observe that the gradient of the RTB objective for a single trajectory is just the sum of per-step log-likelihood gradients scaled by the RTB residual: $$\n\\nabla_{\\phi}L_{RTB} = 2\\left(\\log\\frac{Z_{\\phi}}{r(x_1)} + \\sum_{i=1}^T\\log\\frac{p_{\\phi}^{post}(x_{i\\Delta t} \\mid x_{(i-1)\\Delta t})}{p_{\\theta}(x_{i\\Delta t} \\mid x_{(i-1)\\Delta t})}\\right) \\cdot \\nabla_{\\phi}\\sum_{i=1}^T\\log p_{\\phi}^{post}(x_{i\\Delta t} \\mid x_{(i-1)\\Delta t}).$$ Because the likelihood gradients can be accumulated during the forward pass, this allows for a batched gradient accumulation version of the update. For trajectory length (number of diffusion steps) $T$ and accumulation batch size (number of time steps receiving a gradient signal in each backward pass) $B$, the number of batched forward passes required scales as $T/B$. \n\nOnly the accumulation batch size $B$, not the trajectory length $T$, is constrained by the memory budget. This means we can easily scale training with large number of diffusion steps $T$ without increasing the variance of the gradient through stochastic subsampling, with training time growing only linearly with number of time steps under a fixed memory budget.\n\nThe batched update implementation was not included in our original codebase, but we have performed experiments to validate it and will include them in the revision.\n\nFinally, we highlight that this is in contrast to other diffusion samplers (e.g., PIS/DDS/DIS) which differentiate through the sampling (SDE integration) process and therefore need to store the entire computation graph. For these methods, memory necessarily scales with the length of the trajectory.\n\n#### Cost of training posterior vs. prior\n\nFor our image experiments (Table 2 and Figure 2), a very small number of finetuning steps is necessary for convergence to finetune a posterior with RTB (usually between 100 and 300 are sufficient, although we run for 1500), where each step performs a gradient update on a batch of trajectories (from 16 to 32, depending on the experiment). In contrast, training a prior diffusion model even on MNIST requires tens of thousands of training iterations.\n\n#### LoRA in MNIST/CIFAR setup\n\nThe experiments on discrete diffusion and offline RL (Sections 3.3 and 3.4) do not use LoRA, while the image experiments (Sections 3.1 and 3.2) do. Note that parameter-efficient fine-tuning is standard for large models such as Stable Diffusion and also adopted by baseline methods (e.g., DPOK).\n\nFor example, for MNIST and CIFAR we use LoRA (rank=32) on the key, query and value vectors of the denoiser’s attention mechanism of a UNet model. RTB can be applied without LoRA parameterization on these tasks too; please see the reponse to all reviewers and Table 1 in the attached pdf.\n\n### Conditional constraints\n\nUsing conditional constraints requires a more detailed study of architectures and conditioning strategies, since the unconditional prior cannot directly be finetuned. However, we have run a preliminary experiment on the MNIST even/odd digits task (Table 2); please see the response to all reviewers.\n\n### Data-dependent upper bound\n\nPlease see the response to all reviewers for such an upper bound on the text infilling task, where we have a larger ground truth dataset.\n\n### Max or sum in classifier reward\n\nThank you for the qustion. The choice of the reward function for the MNIST finetuning task (L220) was to discourage generation of ambiguous digits: for example,  $\\sum_{i\\text{ even}} p(c=i\\mid x)$ is high when $p(c=0\\mid x)\\approx p(c=2\\mid x)\\approx\\dots\\approx p(c=8\\mid x)$. In the current formulation, high reward can only be achieved by generating an image whose likelihood of at least one class is high. We observed that replacing max by sum resulted in more such ambiguous digits, which would appear to be a weakness of the pretrained prior model.\n\n### Prompts for ImageReward task\n\nFor the results in Figures 3 and 4, we use a similar setup to DPOK [15]. The prompts are copied verbatim from that paper.\n\n**We hope we have answered your questions satisfactorily; if so, please consider updating your score. Let us know if there are any other questions we can answer during the discussion phase.**", "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": "NqPYCmxXis", "reviewer": "Reviewer_VfXT", "summary": "The paper looks at the problem of finetuning/training a generative model to sample from a desired posterior distribution when given access to a diffusion model prior. The experiments validated the generation capability across different tasks that diffusion models (and conditional generation) can be applied to - text infilling, text-to-image generation and training policies through offline RL. The work proposes a fine-tuning objective to train the model that will sample from the posterior (which in this case is instantiated from the prior’s checkpoint), which can be computed using off policy data/trajectories, ie. trajectories that are not all sampled from the posterior model.", "questions": "1. On the text infilling task, am I correct in assuming that none of the other baselines had access to both x and y (the first 3 and the fifth sentences) since they are autoregressive and can’t condition on y? If yes, is the comparison to diffusion models that can condition on y fair? Should the comparison instead be with encoder-decoder or non-causal language models?\n2. Is there a typo in Line 112? It says “... would sample the product distribution 𝑝post(x1) ∝ 𝑝𝜃 (x1)𝑟(x1). If 𝑟(x1) = 𝑝(y | x1) is a conditional distribution over another variable y, then 𝑝post is the Bayesian posterior 𝑝𝜃 (x1 | y).” Shouldn’t the product distribution be equal to the probably of the intersection of x and y and not y given x? Or is there a term which has implicitly been assumed to be constant/irrelevant?", "limitations": "The paper includes a very brief discussion on limitations which could be improved. I can provide more suggestions for this discussion after engaging with the authors and understanding their assumptions better.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Conditional generation is an active field of study and being able to use a foundational diffusion/generative model as a prior that can be employed/modified to get a posterior distribution of interest, is of interest in many applications.\nThe paper is mostly written clearly, and presents experiments on three different generation problems. Developing an off-policy objective for RL training/finetuning of a generative prior is very useful, however it’s somewhat unclear to me how off-policy the data actually is - I believe the data used is mostly on-policy with small amounts of noise added - could the authors add a small discussion making it very clear how the training data is related to the on/off-policy distributions of the prior/posterior models and how it is different from the data used to train the baselines, as well as how realistic these data collection assumptions are?", "weaknesses": "1. Speaking very qualitatively, the provided generations in Fig 4 and H.1 don’t look very different from each other and DPOK (and often even DDPO) seems to be generating samples of the same quality. This leads me to think that marginal improvements in the metrics provided (log likelihood etc) are not strongly tied to the generation quality. Also how were these samples chosen? To avoid cherry-picking and encourage reproducibility, i’d suggest mentioning the seed they were sampled with in the corresponding code readme files.\n2. Why isn’t there a std dev/error reported for baseline numbers in Table 4? The numbers look close enough for many methods and aggregating metrics over seeds can significantly change results in RL.\n3. A major concern is that the authors should ablate each stabilization/implementation trick they used when it is applicable to other baselines that they compare to (Clipping, using the pretrained ckpt from the prior for the finetuning, using parameter efficient training, any tweaks made to noising schedules, any special tricks used in data collection  etc) since there are so many that have been used.\n\nOverall, I'm only moderately familiar with this area, and therefore will assign a lower confidence to my review. However, my current belief is that the paper has not demonstrated the advantages of the proposed method satisfactorily(or maybe this just has not been presented clearly) so far (either in terms of showing that it can really work well with very off-policy data that is easily available, or showing impressive gains in generation quality/diversity).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your comments. We are happy to see that you appreciated the breadth of applications presented in our paper and the value of the contribution.\n\nBelow we have attempted to answer your questions and to clarify a few misunderstandings in the review. We note that most of the listed weaknesses and questions pertain to specifics of individual experiments, not to the main contribution of the paper: a general-purpose method that can be applied to a wide range of tasks *simultaneously*, achieving results comparable to specialized methods that had been developed for each problem. We hope that our reponses will help to answer your questions and concerns.\n\n### Improvements in image quality \n\nMeasuring posterior sampling quality is difficult for text-to-image generation models, since we do not have ground truth samples from the posterior. We acknowledge that differences in visual quality compared to DPOK are not apparent, and the visual results (Figure 4) should be seen as complementing Figure 3 in showing comparable quality and diversity to baselines.\n\nRTB training facilitates asymptotically unbiased posterior sampling. This may not necessarily translate to improvements in auxiliary fidelity metrics such as those we must resort to when ground truth posterior samples are unavailable (as in the text-to-image experiments). However, in the experiments where we do have access to unbiased samples from the target classes (MNIST/CIFAR-10), it does translate to better FID scores (Table 2).\n\nIn practice, text-to-image generation systems may use a combination of techniques to improve samples (e.g., collecting data from human raters, imposing auxiliary aesthetic constraints, etc.). Our work adds a novel, general, and probabilistically principled technique to this toolbox.\n\n\n### Fair comparisons\n\nWe answer a few questions related to comparisons between other methods and hours, hoping to clarify why we believe they are fair and informative.\n\n#### Infilling baselines seeing the right context\n\nGood question; no, it is **not** the case that baselines see only $x$ (left context) and not $y$ (right context). The fact that this was not clear suggests we should write it explicitly in the paper; thank you for directing our attention to it.\n\nIn fact, all the baselines with the exception of \"Prompt (**x**)\" do condition on both $x$ and $y$, making the methods comparable. For example, for the prompting and SFT autoregressive baselines, the model receives a prompt of the form \"Beginning: {x}, End: {y}, Middle: \", following prior work (see Appendix F). The difference between \"Prompt (**x**)\" and \"Prompt (**x**,**y**)\" shows the importance of conditioning on $y$.\n\n#### Algorithm components\n\nTo clarify, we apply the same collection of techniques with RTB and with methods from prior work wherever applicable. For example, all methods used the same noising schedules. Gradient clipping, fine-tuning the pretrained prior checkpoint, and parameter-efficient training were done for RL baselines as well. \n\nOn the other hand, some techniques used with RTB training are *not applicable* to other methods. Data collection methods used for off-policy training cannot be applied to compared to manifestly on-policy finetuning algorithms such as DPOK. The inference-time classifier guidance baselines, such as those in Table 2, do not require any training, so training hyperparameter choices do not apply to them. \n\nAltogether, while we believe the comparisons we make are fair, we agree that ablations on the techniques used in the paper (while not being applicable to baseline methods) would add to the understanding of the proposed algorithm. We will include ablation studies about this in the final version.\n\n#### Prior/posterior training data\n\nIt is important to note that the newly proposed algorithms **and** most baselines are \"data-free\" -- that is, the prior model is pretrained and we receive no ground truth samples from the posterior, but rather must discover its modes through exploration and can query for the reward of a generated sample.\n\n#### How were samples for figures chosen?\n\nIn the image generation figures (Figures 3 and 4 and Appendix H.1), the samples are not cherry-picked. In Appendix H.1, in each row the images are generated from seeds 0-9.\n\n#### RL baselines\n\nThe reporting of baseline numbers and highlighting in Table 4 exactly matches the form in a line of previous offline RL papers (see, e.g., [35,43,36]), where std is reported only for the newly proposed algorithm alongside baseline means. However, we will update the paper to include std for comparisons, and have attached an updated table with std bars from baselines which report this metric in the global rebuttal document (Table 4).\n\n### Line 112\n\nIn that line, $p^{\\rm post}$ is a distribution over $x_1$, where $y$ is fixed. For a fixed $y$, the product is proportional *as a function of $x_1$* to the posterior $p(x_1\\mid y)$.\n\n**Thank you again for your feedback. We hope these answers and clarifications help; if they do, please consider updating your score. Let us know if there are any other questions we can answer during the discussion phase.**", "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": "8puWqSWfOu", "reviewer": "Reviewer_Z1Kk", "summary": "This paper studies the problem of training diffusion models to sample from an intractable posterior distribution, defined by a prior diffusion model and an arbitrary likelihood function. The contributions are summarized as follows:\n\n- This paper proposes relative trajectory balance (RTB) for training diffusion-based models to sample from a posterior distribution. RTB is derived from the perspective of continuous GFlowNets, which thus enables off-policy training. In contrast to related literature, the proposed approach performs posterior inference by fine-tuning a prior diffusion model in a parameter-efficient way.\n\n- The effectiveness of the proposed approach is validated through experiments in vision, language modeling, and continuous control benchmarks.", "questions": "Some typos:\n\n- Line 158: $x_{1} \\leftarrow x_{0}$ --> $x_{1} \\rightarrow x_{0}$?\n\n- Eq11: $NN_{2}(x_{t}, t)$ --> $NN_{2}(t)$?", "limitations": "Yes. Limitations were included.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- This paper proposes to train diffusion-based models to sample from an intractable posterior distribution. In particular, the proposed approach performs posterior inference by fine-tuning a prior diffusion model $p_{\\theta}$ in a parameter-efficient way.\n\n- The paper is well-written and well-organized.", "weaknesses": "We use a posterior sampler, parameterized by a prior diffusion model with a learnable drift term $u^{post}$, to simulate trajectories and obtain samples $x_{1}$. Then, we compute the likelihood function $r(x_{1})$ and the prior $p_{\\theta}$. Using these as an unnormalized density, the model parameters can be updated using GFlowNet-based loss. For fine-tuning a diffusion model, do we need to ensure that the initial state is an empty set (discrete) or a point mass (continuous)? In general, the initial state for DDPM follows a Gaussian distribution. Does this violate the GFlowNet assumption?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and positive assessment of the paper! We answer your question below:\n\n### Does Gaussian initial state violate assumptions?\n\nAssuming the initial distribution is Gaussian does not violate the GFlowNet assumption, as it amounts to assuming the first generation step transitions from an abstract initial state to a sample from that Gaussian distribution. When training, e.g., with the trajectory balance objective, the likelihood of this first transition -- from the abstract initial state to the Gaussian noise sample -- appears in the decomposition of the sampling trajectory's forward likelihood. With RTB, the likelihood of this initial-to-noise step is the same for the prior and posterior models and thus cancels in the loss.\n\n### Typos\n\n- The direction of the arrows in $x_1\\leftarrow\\dots\\leftarrow x_0$ is intentional: the arrows always show the transitions that exist in the MDP, but the sampling of the noising trajectory happens in reverse (starting from data $x_1$ and ending at noise $x_0$). We can revise if you think this causes confusion.\n- Equation (11): Yes, ${\\rm NN}_2$ isn't taking $x_t$ as input. We will fix this.\n\nThanks for reading carefully!\n\n**Please let us know if there are any other questions we can answer during the discussion period.**", "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:
**We would like to thank all the reviewers for their comments. The reviewers all noted that the paper is well-written (4zDq, VfXT, Z1kk) and pointed out the the broad utility of the proposed approach (4zDq, VfXT, Z1kk), the theoretical justification for the proposed method (4zDq), its efficiency (Z1kk), and the quality of the empirical analysis (4zDq).**

The attached PDF file contains the following:
- **Table 1:** RTB on CIFAR without LoRA (suggestion by Reviewer 4zDq, ablation for Reviewer VfXT).
  - While LoRA was used in the oriignal experiments, full model finetuning can be achieved effectively, using a very small (~1e-5) learning rate to avoid instabilities. We have included results for CIFAR-10 in Table 1. LoRA, however, provides significant memory and speed advantages, including a reduced sensitivity to the learning rate (with LoRA we can train effectively with lr=5e-4 or even lr=1e-3). We will include all details and results in the final version.
- **Table 2:** Conditional constraints experiment for even/odd MNIST (suggestion by Reviewer 4zDq).
  - We use a naive parametrization of training the prior with an extra input channel, which we populate with 0 or 1 for even or odd conditioning during RTB posterior fine-tuning. We look forward to future work which develops more specialized architectures for handling conditional constraints.
- **Table 3:** Upper bound for text-infilling task (suggestion by Reviewer 4zDq).
  - We (supervised-)fine-tune the prior diffusion language model on the entire Stories dataset (which consists of 50k examples compared to the 1k pairs of left and right context used for training with RTB and other baselines). This serves as an upper bound on the performance of all the data-free methods, including RTB.
- **Table 4:** Standard deviation bars for offline RL baselines that report them in their papers (suggestion by Reviewer VfXT).

Meta-review: This paper proposes Relative Trajectory Balance (RTB), a novel method for training diffusion models to sample from intractable posterior distributions by fine-tuning a pre-trained prior. Reviewers found the paper well-written (4zDq, VfXT, Z1kk) with a clear motivation and theoretical justification (4zDq). They appreciated the broad applicability of the method demonstrated through experiments across various domains (4zDq, VfXT, Z1kk).

Reviewers raised concerns regarding the scalability of RTB (4zDq) and the lack of certain experimental results, such as conditional constraints and data-dependent upper bounds (4zDq). They also questioned the significance of the empirical improvements, particularly in image generation quality (VfXT). The authors addressed these concerns in their rebuttal by providing additional experiments (conditional MNIST, text infilling upper bound) and clarifying the scalability of RTB through a batched gradient accumulation implementation. They also addressed the fairness of comparisons with baselines (VfXT) and provided further details on experimental setups (4zDq, VfXT). Although the question regarding the significance of empirical improvement still remains (VfXT).

While some reviewers initially expressed reservations about the empirical results (VfXT), the authors' clarifications and additional experiments paritally alleviated most of these concerns.  Reviewers ultimately found the method to be technically sound (4zDq, Z1kk) and offering a valuable contribution to the field (4zDq, VfXT, Z1kk). Leading to the decision of a weak acceptance.
Decision: Accept (poster)