Title: Amortizing intractable inference in diffusion models for vision, language, and control
Venue: NeurIPS 2024
Accepted: yes
Keywords: diffusion

inverse problems

conditional generation

language models

infilling

discrete diffusion

offline RL

planning

GFlowNet
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

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, x ∼ 𝑝 post ( x ) ∝ 𝑝 ( x ) 𝑟 ( x ) , in a model that consists of a diffusion generative model prior 𝑝 ( x ) and a black-box constraint or likelihood function 𝑟 ( 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. 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 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 .

1 Introduction

Diffusion models [68, 27, 72] are a 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 𝑝 ( x ) that is multiplied by an auxiliary constraint 𝑟 ( x ) . For example, if 𝑝 ( x ) is a prior over images defined by a diffusion model, and 𝑟 ( x ) = 𝑝 ( 𝑐 | x ) is the likelihood that an image x belongs to class 𝑐 , then class-conditional image generation requires sampling from the Bayesian posterior 𝑝 ( x | 𝑐 ) ∝ 𝑝 ( x ) 𝑝 ( 𝑐 | x ) . In offline reinforcement learning, if 𝜇 ( 𝑎 | 𝑠 ) 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 𝜇 ( 𝑎 | 𝑠 ) with a Boltzmann distribution defined by a 𝑄 -function, 𝜋 ∗ ( 𝑎 | 𝑠 ) ∝ 𝜇 ( 𝑎 | 𝑠 ) exp ( 𝛽𝑄 ( 𝑠, 𝑎 )) . In language modeling, various conditional generation problems [43, 23, 29] amount to posterior sampling under a discrete diffusion model prior. Table 1 summarizes four such problems that the proposed method improves upon prior work.

The hierarchical nature of the generative process in diffusion models, which generate samples from 𝑝 ( x ) by a deep chain of stochastic transformations, makes exact sampling from posteriors 𝑝 ( x ) 𝑟 ( x ) under a black-box function 𝑟 ( 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 in 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], we propose an asymptotically unbiased training objective, called relative trajectory balance (RTB), for training diffusion models that sample from posterior distributions under a 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 to 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-guided image generation for unconditional diffusion vision priors (§3.1) and can be used to improve caption-conditioned generation under text-to-image foundation model priors (§3.2).

In 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 a hierarchical generative model. In this section, we present the mathematical setting (§2.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 R 𝑑 , and later discuss how the methods generalize to arbitrary hierarchical priors (§2.4).

(a) Prior

(b) Posterior

(c) RTB (ours)

(d) RL (KL reg.)

(e) RL (no reg.)

(f) CG

2.1 Background and setting: Diffusion models as hierarchical generative models

A denoising diffusion model generates data x 1 by a Markovian generative process:

where Δ 𝑡 = 1 𝑇 and 𝑇 is the number of discretization steps. 1 The initial distribution 𝑝 ( x 0 ) is fixed (typically to N( 0 , I ) ) and the transition from x 𝑡 -1 to x 𝑡 is modeled as a Gaussian perturbation with time-dependent variance:

The scaling of the mean and variance by Δ 𝑡 is insubstantial for fixed 𝑇 , but ensures that the diffusion process is well-defined in the limit 𝑇 →∞ assuming regularity conditions on 𝑢 𝑡 [53, 64]. The process given by (1, 2) is then identical to Euler-Maruyama integration of the stochastic differential equation (SDE) 𝑑 x 𝑡 = 𝑢 𝑡 ( x 𝑡 ) 𝑑𝑡 + 𝜎 𝑡 𝑑 w 𝑡 .

The likelihood of a denoising trajectory x 0 → x Δ 𝑡 →··· → x 1 factors as

and defines a marginal density over the data space:

A reverse-time process, x 1 → x 1 -Δ 𝑡 →··· → x 0 , with densities 𝑞 , can be defined analogously, and similarly defines a conditional density over trajectories:

In the training of diffusion models, as discussed below, the process 𝑞 is typically fixed to a simple distribution (usually a discretized Ornstein-Uhlenbeck process), and the result of training is that 𝑝 and 𝑞 are close as distributions over trajectories.

Diffusion model training as divergence minimization. Diffusion models parametrize the drift 𝑢 𝑡 ( x 𝑡 ) in (Equation 2) as a neural network 𝑢 ( x 𝑡 , 𝑡 ; 𝜃 ) with parameters 𝜃 and taking x 𝑡 and 𝑡 as input. We denote the distributions over trajectories induced by (Equation 3, Equation 4) by 𝑝 𝜃 to show their dependence on the parameter.

In the most common setting, diffusion models are trained to maximize the likelihood of a dataset. In the notation above, this corresponds to assuming 𝑞 ( x 1 ) is fixed to an empirical measure (with the

points of a training dataset D assumed to be i.i.d. samples from 𝑞 ( x 1 ) ). Training minimizes with respect to 𝜃 the divergence between the processes 𝑞 and 𝑝 𝜃 :

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 a lower bound on the data log-likelihood under the model 𝑝 𝜃 .

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, 27]. Such training exploits that for typical choices of the noising process 𝑞 , the optimal 𝑢 𝑡 ( x 𝑡 ) can be expressed in terms of the Stein score of 𝑞 ( x 1 ) convolved with a Gaussian, allowing an efficient stochastic regression objective for 𝑢 𝑡 . For full generality of our exposition for arbitrary iterative generative processes, we 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-matching training. From (Equation 6) we also see that the bound is tight if the conditionals of 𝑝 𝜃 and 𝑞 on x 1 coincide, i.e. , 𝑞 is equal to the posterior distribution of 𝑝 conditioned on x 1 . Indeed, the model 𝑝 𝜃 minimizes (Equation 6) for a distribution with continuous density 𝑞 ( x 1 ) if and only if, for all denoising trajectories,

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 continuoustime 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.i.d. samples from 𝑞 ( x 1 ) one has access to a (possibly unnormalized) density 𝑞 ( x 1 ) = 𝑒 -E( x 1 ) / 𝑍 from which one wishes to sample. These objectives minimize the squared log-ratio between the two sides of (7), which allows the trajectories x 0 → x Δ 𝑡 → · · · → x 1 used for training to be sampled from any training distribution, such as 'exploratory' modifications of 𝑝 𝜃 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, simulationbased approaches that require differentiating through the sampling process [ e.g. , 89, 78, 6, 79].

2.2 Intractable inference under diffusion priors

Consider a diffusion model 𝑝 𝜃 , defining a marginal density 𝑝 𝜃 ( x 1 ) , and a positive constraint function 𝑟 : R 𝑑 → R > 0 . We are interested in training a diffusion model 𝑝 post 𝜙 , with drift function 𝑢 post 𝜙 , that would sample the product distribution 𝑝 post ( x 1 ) ∝ 𝑝 𝜃 ( x 1 ) 𝑟 ( x 1 ) . If 𝑟 ( x 1 ) = 𝑝 ( y | x 1 ) is a conditional distribution over another variable y , then 𝑝 post is the Bayesian posterior 𝑝 𝜃 ( x 1 | y ) .

Because samples from 𝑝 post ( x 1 ) are not assumed to be available, one cannot directly train 𝑝 using the objective (6). Nor can one directly apply objectives for distribution-matching training, such as those that enforce (7), since the marginal 𝑝 𝜃 ( x 1 ) is not available. However, we make the following observation (proof in §A).

Proposition 1 (Relative TB constraint) . If 𝑝 𝜃 , 𝑝 post 𝜙 , and the scalar 𝑍 𝜙 jointly satisfy the relative trajectory balance (RTB) constraint

for every denoising trajectory x 0 → x Δ 𝑡 →··· → x 1 , then 𝑝 post 𝜙 ( x 1 ) ∝ 𝑝 𝜃 ( x 1 ) 𝑟 ( x 1 ) , i.e., the diffusion model 𝑝 post 𝜙 samples the posterior distribution. Furthermore, if 𝑝 𝜃 also satisfies the TB constraint (7) with respect to the noising process 𝑞 and some target density 𝑞 ( x 1 ) , then 𝑝 post 𝜙 satisfies the TB constraint with respect to the target density 𝑞 post ( x 1 ) ∝ 𝑞 ( x 1 ) 𝑟 ( x 1 ) , and 𝑍 = ∫ 𝑞 ( x 1 ) 𝑟 ( x 1 ) 𝑑 x 1 .

Note that the two joints appearing in (8) are defined as products over transitions, via (3).

Relative trajectory balance as a 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 as a function of the vector 𝜙 that parametrizes the posterior diffusion model and the scalar 𝑍 𝜙 (parametrized via log 𝑍 𝜙 for numerical stability):

Optimizing this objective to 0 for all trajectories ensures that (8) is satisfied. While the RTB constraint (8) has a similar form to TB (7), RTB involves the ratio of two denoising processes, while TB involves the ratio of a 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 we assume 𝑝 𝜃 ( x 0 ) = 𝑝 post 𝜙 ( x 0 ) are fixed ( e.g. , to a standard normal), then (9) reduces to

Notably, the gradient of this objective with respect to 𝜙 does not require differentiation (backpropagation) into the sampling process that produced a trajectory x 0 →··· → x 1 . This offers two advantages over on-policy simulation-based methods: (1) the ability to optimize L RTB as an off-policy objective, i.e. , sampling trajectories for training from a distribution different from 𝑝 post 𝜙 itself, as 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 𝑟 ( x 1 ) = 𝑝 ( y | x 1 ) is a conditional likelihood, classifier guidance relies upon writing 𝑢 𝑡 ( x 𝑡 ) -𝑢 post 𝑡 ( x 𝑡 ) explicitly in terms of ∇ x 𝑡 log 𝑝 ( y | x 𝑡 ) , by combining the expression of the optimal drift 𝑢 𝑡 in terms of the score of the target distribution convolved with a Gaussian (cf. §2.1), with the 'Bayes' rule' for the Stein score: ∇ x 𝑡 log 𝑝 ( x 𝑡 | y ) = ∇ x 𝑡 log 𝑝 ( x 𝑡 ) + ∇ x 𝑡 log 𝑝 ( y | x 𝑡 ) .

Classifier guidance gives the exact solution for the posterior drift when a differentiable classifier on noisy data, 𝑝 ( y | x 𝑡 ) = ∫ 𝑝 ( y | x 1 ) 𝑝 ( x 1 | x 𝑡 ) 𝑑 x 1 , is available. Unfortunately, such a classifier is not, in general, tractable to derive from the classifier on noiseless data, 𝑝 ( y | x 1 ) , 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 a large impact on sample efficiency. In on-policy training, we use the current policy 𝑝 post 𝜙 to generate trajectories 𝜏 = ( x 0 → . . . → x 1 ) , evaluate the reward log 𝑟 ( x 1 ) and the likelihood of 𝜏 under 𝑝 𝜃 , and a gradient updates on 𝜙 to minimize L RTB ( 𝜏 ; 𝜙 ) .

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 x 1 that have high density under the target distribution, we can sample noising trajectories x 1 ← x 1 -Δ 𝑡 ← . . . ← x 0 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 onpolicy training with off-policy exploration to be a strength of RTB over on-policy RL methods, as was shown for distribution-matching 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 𝑟 ( x 1 ) . If the constraints depend on other variables y - for example, 𝑟 ( x 1; y ) = 𝑝 ( y | x 1 ) - then the posterior drift 𝑢 post 𝜙 can be conditioned on y and the learned scalar log 𝑍 𝜙 replaced by a 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 easily generalizes to priors that are conditioned on some context variable.

Efficient parametrization and Langevin inductive bias. Because the deep features learned by the prior model 𝑢 𝜃 are expected to be useful in expressing the posterior drift 𝑢 post 𝜙 , we can choose to initialize 𝑢 post 𝜙 as a copy of 𝑢 𝜃 and to fine-tune it, possibly in a parameter-efficient way (as described in each section of §3). This choice is inspired by the method of amortizing inference in large language models by fine-tuning a prior model to sample an intractable posterior [29].

Furthermore, if the constraint 𝑟 ( x 1 ) is differentiable, we can impose an inductive bias on the posterior drift similar to the one introduced for diffusion samplers of unnormalized target densities in [89] and shown to be useful for off-policy methods in [65]. namely, we write

where NN 1 and NN 2 are neural networks outputting a vector and a scalar, respectively. This parametrization allows the constraint to provide a 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 a minibatch of the quantity inside the square, which removes dependence on log 𝑍 𝜙 and is especially useful in conditional settings, consistent with the findings of [65]. This amounts to a relative variant of the VarGrad objective [61] (see (23) in §G). 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 to 0 (see §E,§F).

2.4 Generative flow networks and extension to 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 R 𝑑 , 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 a 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 in R 𝑑 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 𝑞 becomes trivial, and the RTB constraint for a model 𝑝 post 𝜙 to sample a sequence x from a distribution with density 𝑟 ( x ) 𝑝 𝜃 ( x ) is simply 𝑍 𝜙 𝑝 post 𝜙 ( x ) = 𝑟 ( x ) 𝑝 𝜃 ( x ) for all sequences x . We note that a sub-trajectory generalization of this objective was used in [29] to amortize intractable inference in autoregressive language models.

3 Experiments

In this section, we present empirical results to 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 a classifier-guided visual task where we wish to learn a diffusion posterior 𝑝 post 𝜙 ( x | 𝑐 ) ∝ 𝑝 𝜃 ( x ) 𝑝 ( 𝑐 | x ) given a pretrained diffusion prior 𝑝 𝜃 ( x ) and a classifier 𝑟 ( x ) = 𝑝 ( 𝑐 | 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 𝑝 ( 𝑐 | x ) for both datasets. We perform parameter-efficient fine-tuning of 𝑝 post 𝜙 , initialized as a copy of the prior 𝑝 𝜃 , using the RTB objective (see §E.1 for details). The RTB objective is optimized on trajectories sampled

Table 2: Classifier-guided posterior sampling with pretrained unconditional diffusion priors. We report the mean ± std of each metric computed across all relevant classes for each experiment set, and highlight ± 5% from highest/lower experimental value. The FID 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 log 𝑟 ( 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.1 for conditional variants.

Dataset

→

MNIST

MNIST even/odd

CIFAR-10

Algorithm

↓

Metric

→

E

[

log

𝑟

(

x

)]

(

↑

)

FID (

↓

)

Diversity (

↑

)

E

[

log

𝑟

(

x

)]

(

↑

)

FID (

↓

)

Diversity (

↑

)

E

[

log

𝑟

(

x

)]

(

↑

)

FID (

↓

)

Diversity (

↑

)

DPS

-

2

.

1597

±

0.423

1

.

2913

±

0.410

0

.

1609

±

0.000

-

1

.

2270

±

0.202

1

.

1498

±

0.182

0

.

1713

±

0.000

-

3

.

6025

±

0.503

0

.

7371

±

0.216

0

.

2738

±

0.000

LGD

-

MC

-

2

.

1389

±

0.480

1

.

2873

±

0.412

0

.

1600

±

0.000

-

1

.

1720

±

0.199

1

.

1445

±

0.184

0

.

1600

±

0.000

-

3

.

0988

±

0.359

0

.

7402

±

0.214

0

.

2743

±

0.000

DDPO

-

1

.

5

±

4.7

×

10

-

3

1

.

5822

±

0.583

0

.

1350

±

0.005

-

8

.

6

±

12.3

×

10

-

11

1

.

8024

±

0.423

0

.

1314

±

0.002

-

2

.

7

±

8.5

×

10

-

4

1

.

7686

±

0.589

0

.

1575

±

0.015

DPOK

-

0

.

1379

±

0.225

1

.

2063

±

0.316

0

.

1442

±

0.004

-

0

.

0783

±

0.082

1

.

2536

±

0.206

0

.

1631

±

0.007

-

2

.

4414

±

3.266

0

.

5316

±

0.157

0

.

2415

±

0.024

RTB (ours)

-

0

.

1734

±

0.194

1

.

1823

±

0.288

0

.

1474

±

0.003

-

0

.

1816

±

0.175

1

.

1794

±

0.171

0

.

1679

±

0.004

-

2

.

1625

±

0.879

0

.

4717

±

0.138

0

.

2440

±

0.011

Prior

Posterior(seven)

Posterior (even)

Prior

Posterior (dog)

3

0

S

7

7

8

8

9

B

7

22

7

8

4

3

8

7

7

t

7

7

0

8

2

8

S

y

0

7

7

7

7

6

MNIST

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 singledigit posterior (learning to sample images of each digit class 𝑐 ), CIFAR-10 single-class posterior (analogous to the previous), and MNIST multi-digit posterior. The latter is a multimodal posterior, for which we set 𝑟 ( x ) = max 𝑖 ∈{ 0 , 2 , 4 , 6 , 8 } 𝑝 ( 𝑐 = 𝑖 | x ) to generate even digits, and similarly for odd digits.

Results. Samples from the RTB-fine-tuned posterior models are shown in Fig. 2. In Table 2 we report mean ± std 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 𝑟 ( x ) . Pure RL fine-tuning (no KL regularization) displays mode collapse characteristics, achieving high rewards in exchange for significantly poorer diversity and FID scores (see also Fig. E.1). Classifier-guidance-based methods, like DP and LGD-MC, exhibit high diversity, but fail to appropriately model the posterior distribution (lowest log 𝑟 ( x ) ). Additional results can be found in §E.2.

3.2 Fine-tuning a 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'). Finetuning pretrained text-to-image diffusion models 𝑝 𝜃 ( x 1 | z ) as RL policies to maximize some reward 𝑟 ( x 1 , z ) based on 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 ˜ 𝑝 to be close to the prior 𝑝 :

The optimal ˜ 𝑝 for (12) is ˜ 𝑝 ( x 1 | z ) ∝ 𝑝 ( x 1 | z ) exp ( 𝛽𝑟 ( x 1 , z )) for some inverse temperature 𝛽 . The marginal KL 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 -𝛾 ˝ 𝑇 𝑖 = 1 𝐷 KL [ ˜ 𝑝 ( x 𝑖 Δ 𝑡 | x ( 𝑖 -1 ) Δ 𝑡 , z )|| 𝑝 ( x 𝑖 Δ 𝑡 | x ( 𝑖 -1 ) Δ 𝑡 , z )] 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 ˜ 𝑝 ( x 1 | 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 × 512

(a) Four roses.

(b) A green rabbit.

(c) A cat and a dog.

(d) Cybernetic man.

ROOSS

images. Following DPOK [16], we use 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, as the reward log 𝑟 ( x 1 , z ) . As reference, we present comparisons against DPOK with the default KL regularization 𝛾 = 0 . 01 and DPOK with 𝛾 = 0 . 0 , which is equivalent to DDPO [8]. We measure the final average reward and the diversity of the generated image, as 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 are discussed in §H.

Results. Fig. 3 plots the diversity versus log reward on a set of prompts from [16, 87]. In terms of average log 𝑟 ( x 1 , 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 is 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-art methods for fine-tuning text-toimage diffusion models.

0.05

0.10

Diversity

0.0

0.5

1.0

1.5

Avg. reward

A green colored rabbit.

None (prior)

DDPO

DPOK

RTB (ours)

0.10

0.12

0.14

Diversity

A cat and a dog.

0.15

0.20

0.25

Diversity

Four roses.

0.10

0.12

0.14

0.16

Diversity

Cybernetic man

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 of autoregressive language models of similar scale. Non-autoregressive generation in diffusion language models can provide useful inductive biases for infilling, such as the ability to attend to context on both sides of a target token.

Setup. We use the ROCStories corpus [50], a dataset of short stories containing 5 sentences each. We adopt the task setup from [29], where the first 3 sentences of a story x and the last sentence y are given, and the goal is 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 𝑝 post ( z | x , y ) ∝ 𝑝 ( z | x ) 𝑝 reward ( y | x , z ) where 𝑝 is a SEDD language model prior (a conditional model over z given x ) and 𝑝 reward is an autoregressive language model fine-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 x (Prompt ( x ) ) and x , y (Prompt ( x , y ) ). Additionally, to contextualize the performance, we also consider autoregressive language model baselines from [29], which studied this problem under an autoregressive prior 𝑝 ( 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.

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 DeBERTa [26]), BLEU-4 [54], and GLEU-4 [86]. Table 3 summarizes the results. We observe that the diffusion language model performs significantly 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 GFlowNet fine-tuning. We provide some examples of generated text in §F.

Model

Algorithm

↓

Metric

→

BLEU-4

GLEU-4

BERTScore

Autoreg.

Prompting

0

.

010

±

0

.

002

0

.

022

±

0

.

001

0

.

005

±

0

.

001

Supervised fine-tuning

0

.

012

±

0

.

001

0

.

023

±

0

.

001

0

.

013

±

0

.

002

GFN fine-tuning [29]

0

.

019

±

0

.

001

0

.

031

±

0

.

002

0

.

102

±

0

.

005

Discrete

diffusion

Prompt

(

x

)

0

.

011

±

0

.

002

0

.

023

±

0

.

002

0

.

014

±

0

.

003

Prompt

(

x

,

y

)

0

.

014

±

0

.

003

0

.

027

±

0

.

003

0

.

092

±

0

.

004

RTB (ours)

0

.

025

±

0

.

002

0

.

045

±

0

.

002

0

.

156

±

0

.

003

SFT (upper bound)

0

.

031

±

0

.

002

0

.

057

±

0

.

004

0

.

182

±

0

.

005

3.4 KL-constrained policy search in offline reinforcement learning

The goal of RL algorithms is to learn a policy 𝜋 ( 𝑎 | 𝑠 ) , i.e. , a mapping from states 𝑠 to actions 𝑎 in an environment, that maximizes the expected cumulative discounted reward [75]. In the offline RL setting [39], the agent has access to a dataset D = {( 𝑠 𝑖 𝑡 , 𝑎 𝑖 𝑡 , 𝑠 𝑖 𝑡 + 1 , 𝑟 𝑖 𝑡 )} 𝑁 𝑖 = 1 of transitions (where each sample ( 𝑠 𝑡 , 𝑎 𝑡 , 𝑠 𝑡 + 1 , 𝑟 𝑡 ) indicates that an agent taking action 𝑎 𝑡 at state 𝑠 𝑡 transitioned to the next state 𝑠 𝑡 + 1 and received reward 𝑟 𝑡 ). This dataset is assumed to be generated by a behavior policy 𝜇 ( 𝑎 | 𝑠 ) , 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 in offline RL is policy extraction from trained 𝑄 -functions [55, 25, 44]. For reliable extrapolation, one wants the policy to predict actions that have high 𝑄 -values, but also have high density under the behavior policy 𝜇 , as naive maximization can result in choosing actions with low probability under 𝜇 and thus unreliable predictions from the 𝑄 -function. This is formulated as a KL-constrained policy search problem:

where 𝑑 𝜇 is the distribution over states induced by following the policy 𝜇 . The optimal policy 𝜋 in (13) is the product distribution 𝜋 ∗ ( 𝑎 | 𝑠 ) ∝ 𝜇 ( 𝑎 | 𝑠 ) exp ( 𝛽𝑄 ( 𝑠, 𝑎 )) for some inverse temperature 𝛽 . If 𝜇 ( 𝑎 | 𝑠 ) is a conditional diffusion model over continuous actions 𝑎 conditioned on state 𝑠 , we use RTB to fine-tune a diffusion behavior policy to sample from 𝜋 ∗ , using 𝜇 as the prior and exp ( 𝛽𝑄 ( 𝑠, 𝑎 )) as the target constraint. We use a 𝑄 -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), 'medium-expert' (collected from 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-ofthe-art: Diffuser [D; 30], Decision Diffuser [DD; 2], D-QL [83], IDQL [25], and QGPO [44]. For algorithm implementation details, hyperparameters, and a report of baselines, see §G.

Results. Table 4 shows that RTB matches state-of-the-art 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.

4 Other related work

Composing iterative generative processes. Beyond the approximate posterior sampling algorithms and application-specific 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] is another domain of sampling from product distributions where diffusion models are now being used [21]. Beyond product models, [19]

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-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.

5 Conclusions and future work

Relative trajectory balance provides a new approach to training diffusion models to generate unbiased posterior samples given a diffusion prior and an arbitrary reward function. Through experiments on a 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, as should simulation-free extensions, e.g. , through objectives that are local in time [45]. The ability to handle arbitrary black-box likelihoods also makes RTB a 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. , 1], medical imaging [ e.g. , 73], and molecular structure prediction [ e.g. , 84].

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, Seong et al. [66] have already explored a 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 memory-intensive. Additionally, the RTB objective is computed on complete trajectories without any local credit-assignment 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.

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 AI4Discovery, Samsung, and Recursion.

The research was enabled in part by computational resources provided by the Digital Research Alliance of Canada ( https://alliancecan.ca ), Mila ( https://mila.quebec ), and NVIDIA.

References

Alexandre Adam, Adam Coogan, Nikolay Malkin, Ronan Legin, Laurence Perreault-Levasseur, Yashar Hezaveh, and Yoshua Bengio. Posterior samples of source galaxies in strong gravitational lenses with score-based priors. arXiv preprint arXiv:2211.03812 , 2022.

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.

Tara Akhound-Sadegh, Jarrid Rector-Brooks, Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, and Alexander Tong. Iterated denoising energy matching for sampling from Boltzmann densities. International Conference on Machine Learning (ICML) , 2024.

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.

Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet foundations. Journal of Machine Learning Research , 24(210):1-55, 2023.

Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusionbased generative modeling. Transactions on Machine Learning Research (TMLR) , 2024.

Julius Berner, Lorenz Richter, Marcin Sendera, Jarrid Rector-Brooks, and Nikolay Malkin. From discrete-time policies to continuous-time diffusion samplers: Asymptotic equivalences and faster training. arXiv preprint arXiv:2501.06148 , 2025.

Kevin Black, Michael Janner, Yilun Du, Ilya Kostrikov, and Sergey Levine. Training diffusion models with reinforcement learning. International Conference on Learning Representations (ICLR) , 2024.

Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Neural Information Processing Systems (NeurIPS) , 2022.

Gabriel Cardoso, Yazid Janati el idrissi, Sylvain Le Corff, and Eric Moulines. Monte carlo guided denoising diffusion models for bayesian linear inverse problems. International Conference on Learning Representations (ICLR) , 2024.

Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. International Conference on Learning Representations (ICLR) , 2023.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat GANs on image synthesis. Neural Information Processing Systems (NeurIPS) , 2021.

Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Continuous diffusion for categorical data. arXiv preprint arXiv:2211.15089 , 2022.

Zehao Dou and Yang Song. Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. International Conference on Learning Representations (ICLR) , 2024.

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 on Machine Learning (ICML) , 2023.

Ying Fan, Olivia Watkins, Yuqing Du, Hao Liu, Moonkyung Ryu, Craig Boutilier, Pieter Abbeel, Mohammad Ghavamzadeh, Kangwook Lee, and Kimin Lee. Reinforcement learning for finetuning text-to-image diffusion models. Neural Information Processing Systems (NeurIPS) , 2023.

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.

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.

Timur Garipov, Sebastiaan De Peuter, Ge Yang, Vikas Garg, Samuel Kaski, and Tommi Jaakkola. Compositional sculpting of iterative generative processes. Neural Information Processing Systems (NeurIPS) , 2023.

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.

Tomas Geffner, George Papamakarios, and Andriy Mnih. Compositional score modeling for simulation-based inference. International Conference on Machine Learning (ICML) , 2023.

Alexandros Graikos, Nikolay Malkin, Nebojsa Jojic, and Dimitris Samaras. Diffusion models as plug-and-play priors. Neural Information Processing Systems (NeurIPS) , 2022.

Ishaan Gulrajani and Tatsunori B Hashimoto. Likelihood-based diffusion language models. Neural Information Processing Systems (NeurIPS) , 2023.

Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov. SSD-LM: Semi-autoregressive simplexbased diffusion language model for text generation and modular control. Association for Computational Linguistics (ACL) , 2023.

Philippe Hansen-Estruch, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. IDQL: Implicit Q-learning as an actor-critic method with diffusion policies. arXiv preprint arXiv:2304.10573 , 2023.

Pengcheng He, Xiaodong Liu, Jianfeng Gao, and Weizhu Chen. Deberta: Decoding-enhanced bert with disentangled attention. International Conference on Learning Representations (ICLR) , 2021.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Neural Information Processing Systems (NeurIPS) , 2020.

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. International Conference on Learning Representations (ICLR) , 2022.

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.

Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. International Conference on Machine Learning (ICML) , 2022.

Zahra Kadkhodaie and Eero P. Simoncelli. Solving linear inverse problems using the prior implicit in a denoiser. Neural Information Processing Systems (NeurIPS) , 2021.

Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. Neural Information Processing Systems (NeurIPS) , 2024.

Bahjat Kawar, Gregory Vaksman, and Michael Elad. SNIPS: Solving noisy inverse problems stochastically. Neural Information Processing Systems (NeurIPS) , 2021.

Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio, Sungsoo Ahn, and Jinkyoo Park. Local search GFlowNets. International Conference on Learning Representations (ICLR) , 2024.

Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. International Conference on Learning Representations (ICLR) , 2014.

Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit Q-learning. International Conference on Learning Representations (ICLR) , 2022.

Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. Neural Information Processing Systems (NeurIPS) , 2020.

Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hernández-Garcıa, Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of continuous generative flow networks. International Conference on Machine Learning (ICML) , 2023.

Sergey Levine, Aviral Kumar, G. Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020.

Xiang Li, John Thickstun, Ishaan Gulrajani, Percy S Liang, and Tatsunori B Hashimoto. Diffusion-LM improves controllable text generation. Neural Information Processing Systems (NeurIPS) , 2022.

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.

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 arXiv:2303.16634 , 2023.

Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. arXiv preprint arXiv:2310.16834 , 2023.

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.

Kanika Madan, Jarrid Rector-Brooks, Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. International Conference on Machine Learning (ICML) , 2022.

Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in GFlowNets. Neural Information Processing Systems (NeurIPS) , 2022.

Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference. International Conference on Learning Representations (ICLR) , 2023.

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.

Chenlin Meng, Kristy Choi, Jiaming Song, and Stefano Ermon. Concrete score matching: Generalized score matching for discrete data. Neural Information Processing Systems (NeurIPS) , 2022.

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 Computational Linguistics: Human Language Technologies (NAACL-HLT) , 2016.

Alex Nichol and Prafulla Dhariwal. Improved denoising diffusion probabili1stic models. International Conference on Machine Learning (ICML) , 2021.

Nikolas Nüsken and Lorenz Richter. Solving high-dimensional Hamilton-Jacobi-Bellman 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.

Bernt Øksendal. Stochastic Differential Equations: An Introduction with Applications . Springer, 2003.

Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu. Bleu: a method for automatic evaluation of machine translation. Association for Computational Linguistics (ACL) , 2002.

Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv prepring arXiv:1910.00177 , 2019.

Ben Poole, Ajay Jain, Jonathan T. Barron, and Ben Mildenhall. DreamFusion: Text-to-3D using 2D diffusion. International Conference on Learning Representations (ICLR) , 2023.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

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.

Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. International Conference on Machine Learning (ICML) , 2015.

Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. International Conference on Machine Learning (ICML) , 2014.

Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco J. R. Ruiz, and Ömer Deniz Akyildiz. VarGrad: A low-variance gradient estimator for variational inference. Neural Information Processing Systems (NeurIPS) , 2020.

Lorenz Richter, Julius Berner, and Guan-Horng Liu. Improved sampling via learned diffusions. International Conference on Learning Representations (ICLR) , 2023.

Robin Rombach, A. Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. Computer Vision and Pattern Recognition (CVPR) , 2022.

Simo Särkkä and Arno Solin. Applied stochastic differential equations . Cambridge University Press, 2019.

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.

Kiyoung Seong, Seonghyun Park, Seonghwan Kim, Woo Youn Kim, and Sungsoo Ahn. Collective variable free transition path sampling with generative flow network. arXiv preprint arXiv:2405.19961 , 2024.

Tiago Silva, Amauri H Souza, Luiz Max Carvalho, Samuel Kaski, and Diego Mesquita. Federated contrastive GFlowNets, 2024. URL https://openreview.net/forum?id= VJDFhkwQg6 .

Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning (ICML) , 2015.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. International Conference on Learning Representations (ICLR) , 2021.

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.

Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-based diffusion models. Neural Information Processing Systems (NeurIPS) , 2021.

Yang Song, Jascha Sohl-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.

Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with score-based generative models. International Conference on Learning Representations (ICLR) , 2022.

Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. International Conference on Learning Representations (ICLR) , 2023.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT Press, 2018.

Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems , pages 5026-5033, 2012. doi: 10.1109/IROS.2012.6386109.

Masatoshi Uehara, Yulai Zhao, Kevin Black, Ehsan Hajiramezanali, Gabriele Scalia, Nathaniel Lee Diamant, Alex M Tseng, Tommaso Biancalani, and Sergey Levine. Finetuning of continuous-time diffusion models as entropy-regularized control. arXiv preprint arXiv:2402.15194 , 2024.

Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. International Conference on Learning Representations (ICLR) , 2023.

Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets variational inference: Controlled Monte Carlo diffusions. International Conference on Learning Representations (ICLR) , 2024.

Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computation , 23(7):1661-1674, 2011.

Leandro von Werra, Younes Belkada, Lewis Tunstall, Edward Beeching, Tristan Thrush, Nathan Lambert, and Shengyi Huang. Trl: Transformer reinforcement learning. https://github. com/huggingface/trl , 2020.

Haochen Wang, Xiaodan Du, Jiahao Li, Raymond A. Yeh, and Greg Shakhnarovich. Score Jacobian chaining: Lifting pretrained 2D diffusion models for 3D generation. Conference on Computer Vision and Pattern Recognition (CVPR) , 2023.

Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. International Conference on Learning Representations (ICLR) , 2023.

Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature , 620(7976):1089-1100, 2023.

Sean Welleck, Ilia Kulikov, Jaedeok Kim, Richard Yuanzhe Pang, and Kyunghyun Cho. Consistency of a recurrent language model with respect to incomplete decoding. Empirical Methods in Natural Language Processing (EMNLP) , 2020.

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 preprint arXiv:1609.08144 , 2016.

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.

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 on Learning Representations (ICLR) , 2024.

Qinsheng Zhang and Yongxin Chen. Path integral sampler: a stochastic control approach for sampling. International Conference on Learning Representations (ICLR) , 2022.

Tianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger, and Yoav Artzi. BERTScore: Evaluating text generation with bert. International Conference on Learning Representations , 2020.

Wanrong Zhu, Zhiting Hu, and Eric Xing. Text infilling. arXiv preprint arXiv:1901.00158 , 2019.

A Proofs

Proposition 1 (Relative TB constraint) . If 𝑝 𝜃 , 𝑝 post 𝜙 , and the scalar 𝑍 𝜙 jointly satisfy the relative trajectory balance (RTB) constraint

for every denoising trajectory x 0 → x Δ 𝑡 →··· → x 1 , then 𝑝 post 𝜙 ( x 1 ) ∝ 𝑝 𝜃 ( x 1 ) 𝑟 ( x 1 ) , i.e., the diffusion model 𝑝 post 𝜙 samples the posterior distribution. Furthermore, if 𝑝 𝜃 also satisfies the TB constraint (7) with respect to the noising process 𝑞 and some target density 𝑞 ( x 1 ) , then 𝑝 post 𝜙 satisfies the TB constraint with respect to the target density 𝑞 post ( x 1 ) ∝ 𝑞 ( x 1 ) 𝑟 ( x 1 ) , and 𝑍 = ∫ 𝑞 ( x 1 ) 𝑟 ( x 1 ) 𝑑 x 1 .

Proof of Prop. 1. Suppose that 𝑝 𝜃 , 𝑝 post 𝜙 , and 𝑍 jointly satisfy (8). Then necessarily 𝑍 ≠ 0 , since the quantities on the right side are positive. We then have, using (4),

as desired.

Now suppose that 𝑝 𝜃 also satisfies the TB constraint (7) with respect to 𝑞 ( x 1 ) . Then, for any denoising trajectory,

showing that 𝑝 post 𝜙 satisfies the TB constraint with respect to the noising process 𝑞 and the (not yet shown to be normalized) density 1 𝑍 𝑞 ( x 1 ) 𝑟 ( x 1 ) . We integrate out the variables x 0 , x Δ 𝑡 , . . . , x 1 -Δ 𝑡 in (14), giving

Integrating over x 1 shows ∫ 𝑞 ( x 1 ) 𝑟 ( x 1 ) 𝑑 x 1 = 𝑍 .

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 measure over R 𝑑 , these foundations generalize to more general reference measures . In application to diffusion samplers, suppose that 𝜋 ref ( x 𝑡 ) is a collection of Lebesgue-absolutely continuous densities over R 𝑑 for 𝑡 = 0 , Δ 𝑡, . . . , 1 and that - → 𝜋 ref ( x 𝑡 | x 𝑡 -Δ 𝑡 ) , ← -𝜋 ref ( x 𝑡 -Δ 𝑡 | x 𝑡 ) are collections of Lebesgue-absolutely continuous transition kernels. If these densities jointly satisfy the detailed balance condition 𝜋 ref ( x 𝑡 ) ← -𝜋 ref ( x 𝑡 -Δ 𝑡 | x 𝑡 ) = 𝜋 ref ( x 𝑡 -Δ 𝑡 ) - → 𝜋 ref ( x 𝑡 | x 𝑡 -Δ 𝑡 ) , then they satisfy the conditions 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 𝑟 , then the forward process 𝑝 samples from the distribution with density 𝑟 , with all densities interpreted as relative to the reference measures 𝜋 ref , ← -𝜋 ref , - → 𝜋 ref . 3

□

If 𝑝 𝜃 is a diffusion model that satisfies the TB constraint jointly with some reverse process 𝑞 and target density 𝑞 ( x 1 ) , then one can take the reference transition kernels - → 𝜋 ref , ← -𝜋 ref to be 𝑝 and 𝑞 , respectively. In this case, the TB constraint for a target density 1 𝑍 𝑟 ( x 1 ) and forward transition 𝑝 post 𝜙 is

which is identical to the RTB constraint (8). If (15) holds, then 𝑝 post 𝜙 samples from the distribution with density 1 𝑍 𝑟 ( x 1 ) relative to 𝜋 ref ( x 1 ) , which is exactly 1 𝑍 𝑝 𝜃 ( x 1 ) 𝑟 ( x 1 ) . We have thus recovered RTB as a case of TB for non-Lebesgue reference measures.

C Posterior inference on two-dimensional Gaussian mixture model

Setup We conduct toy experiments in low-dimensional spaces using samples from a Gaussian mixture model with multiple modes to visually demonstrate its validity. The prior distribution 𝑝 ( x 1 ) is trained on a Gaussian mixture model with 25 evenly weighted modes, while the target posterior 𝑝 post ( x 1 ) = 𝑟 ( x 1 ) 𝑝 ( x 1 ) uses a reward 𝑟 ( x 1 ) to select and re-weight 9 modes from 𝑝 ( x 1 ) . More specifically, the resulting posterior is:

Our objective is to sample from the posterior 𝑝 post ( x 1 ) . We compare our method with several baselines, including policy gradient reinforcement learning (RL) with KL constraint and classifierguided diffusion models. For RL, we implemented the REINFORCE method with a mean baseline and a KL constraint, following recent work training diffusion models to optimize a reward function [8]. Sampling according to the RL policy leads to a distribution 𝑞 𝜃 ( x 1 ) , which is trained with the objective:

While the exact computation of 𝐾𝐿 ( 𝑞 𝜃 ( x 1 )∥ 𝑝 ( x 1 )) 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 a posterior score function estimate:

Note that this is a biased approximation of the true intractable score:

For our experiments, we follow the source code 4 provided in recent diffusion sampler benchmarks [65]. We utilize a batch size of 500, with finetuning at 5,000 training iterations, a learning rate of 0.0001, a diffusion time scale of 5.0, 100 steps, and a 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 hyperparameters as 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 classifierguided diffusion model shows biased posterior sampling (Fig. 1f), failing to provide accurate inference. RL with a per step KL constraint cannot exactly optimize for the posterior distribution, making the tuning of the KL weight 𝛼 crucial to achieving desirable output Fig. C.1. RTB asymptotically achieves the true posterior without introducing a balancing hyperparameter 𝛼 . Another advantage of our approach is off-policy exploration for efficient mode coverage. RL methods for fine-tuning diffusion models (e.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 𝜖 2 𝑇 , where 𝜖 is a noise hyperparameter and 𝑇 is the diffusion timestep, into the variances and annealing it to zero over training iterations. We set 𝜖 = 0 . 5 for off-policy exploration. As shown in Fig. C.2, RTB with off-policy exploration gives very close posterior inferences, whereas off-policy exploration in RL with 𝛼 = 0 . 5 (which is a carefully selected hyperparameter) does not improve performance due to its on-policy nature.

D Code

Code for all experiments is available at https://github.com/GFNOrg/diffusion-finetuning and will continue to be maintained and extended.

(a) Prior

(b) Posterior

(c)

𝛼

=

1

.

0

(d)

𝛼

=

0

.

8

(e)

𝛼

=

0

.

7

(f)

𝛼

=

0

.

5

(g)

𝛼

=

0

.

3

O

O

O

(a) Prior

(b) Posterior

(c) RTB (on)

(d) RTB (off)

(e) RL (on)

(f) RL (off)

E On classifier guidance and RTB posterior sampling

MNIST

CIFAR-10

0000

0000

o00o

0000

Even

Ship

Car

Cat

Dog

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 𝑝 post ( 𝑥 | 𝑦 ) = 𝑝 ( 𝑦 | 𝑥 ) 𝑝 ( 𝑥 ) . In this section, we detail the experimental settings for RTB as well as the compared baselines.

Experiments setting. For MNIST, we pretrain a noise-predicting diffusion model on 28 × 28 (upscaled to 32 × 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 a pretrained model from [27], trained to generate 32 × 32 3-channel images from the CIFAR-10 dataset, while discretizing the noising/denoising processes into 1000 steps. For fine-tuning the prior, we 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 in the data, (between 5k and 6k for CIFAR-10 and MNIST digits, and 30k for the even/odd task).

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 a DDPM sampling scheme, with only 20% to 50% of the original trained steps. We employ loss clipping at 0 . 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 a reinforcement 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 𝛼 ∈ { 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, as it requires a differentiable reward 𝑟 ( x ) . We perform a sweep over ten values of the suggest parameter range for the step size 𝜁 ∈ [ . 1 , 1 . ] 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 classifierbased 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 a Monte Carlo estimate of the gradient correction to 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.

Classifier-guidance baselines. We find that the DP and LGD-MC classifier-guidance based baselines struggle to sample from the true 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 a 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 datasets, as well as the Even objective from the MNIST dataset, after 500 training iterations. In the figure, we observe early mode collapse 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. 2 for comparison with our RTB-finetuned models).

Conditional architectures. We repeat the even/odd posterior sampling experiment of §3.1 in a conditional setting, where the condition is an input to the posterior model. For the posterior architecture, we use a naive modification of the prior with an extra input channel, which is populated a full mask of 0 or 1 for conditioning on the even and odd classes, respectively. The results are shown in Table E.1. We look forward to future work which develops more specialized architectures for handling conditional constraints.

F Infilling with discrete diffusion

Additional details. We illustrate some examples from the ROC Stories dataset used for training in Table F.1. For the prior we use the sedd-small 5 model, which uses an absorbing noising process [4] with a log linear noise schedule, as the diffusion prior 𝑝 ( z | x ) . The posterior model is

parameterized as a copy of the prior. To condition the diffusion model on the beginning x we set the tokens at the appropriate location in the state in the initial time step, i.e. 𝑡 = 0 . 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 a 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 0, resulting in updates only when the loss is larger. This threshold - referred to as the loss clipping coefficient - is a hyperparameter. As this is a conditional problem we also use the relative VarGrad objective. We also use some tempering on the reward likelihood which helps in learning ( i.e. , 𝑝 reward ( y | x , z ) 𝛽 ) where 𝛽 is the inverse temperature parameter. We perform all experiments on an NVIDIA A100-Large GPU. Note that we also tried a 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.

Reward. For training 𝑝 reward we follow the training procedure and implementation from [29] 7 . 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 a 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]. A critical difference in our experiments compared to [29] is that the posterior model is not initialized with a base model that is fine-tuned on the stories dataset. To condition the model on 𝑋 and 𝑌 , as well as for the prompting baseline, we use the following prompt:

"Beginning: {X}\n 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) x s for every ( 𝑋,𝑌 ) , and 𝑝 reward ( 𝑋𝑍𝑌 ) is used as 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, we 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 F.3, 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. At 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 a LLMEval [42] for evaluating the coherence of the stories but did not obtain statistically significant results.

G Offline RL

G.1 Training details

Our method requires first training a diffusion-based behavior policy 𝜋 𝜃 and a Q-function 𝑄 𝜓 . Once 𝜋 𝜃 and 𝑄 𝜓 are trained, The posterior policy 𝜋 𝛾 is trained using RTB, with its weights initialized to the trained behavior policy weights 𝜃 .

The behavior policy 𝜋 𝜃 is parametrized as a state-conditioned noise-predicting denoising diffusion probabilistic model (DDPM) [27] with a linear schedule, and 75 denoising steps. The diffusion model takes as input a state 𝑠 , a noised action 𝑎 𝑡 and a noise level 𝑡 and predicts the source noise 𝜖 . The state 𝑠 and noised action 𝑎 𝑡 are concatenated with Fourier features computed on the noise level 𝑡 , which are then fed through a 3-layer MLP of hidden dimensionality 256, with layer normalization and a GeLU activation after each hidden layer. The behavior policy is trained using the Adam optimizer with batch size 512 and learning rate 5e-4 for 10000 epochs. The Q-function 𝑄 𝜓 is trained using IQL. We use the same IQL experimental configurations and training hyperparameters as in [36]. That is, we set 𝜏 = 0 . 7 . The architecture for 𝑄 𝜓 is a 3-layer MLP with hidden dimensionality 256 and ReLU activations, which is trained using the Adam optimizer with a 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

with soft updates of 𝑚 = 0 . 005 . The posterior policy 𝜋 𝛾 is trained using the relative trajectory balance objective. 𝜋 𝛾 is also parametrized as a state-conditioned noise-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 is defined as 𝜖 ( 𝑠, 𝑎 𝑡 , 𝑡 ) : = 𝜖 ( 𝑠, 𝑎 𝑡 , 𝑡 ; 𝜃 ) + 𝜖 ( 𝑠, 𝑎 𝑡 , 𝑡 ; 𝛾 ) , where 𝜖 (· ; 𝜃 ) is the output of the prior noise prediction network and 𝜖 (· ; 𝛾 ) is the output of the posterior noise prediction network. We train all models on a single NVIDIA A100-Large GPU. The only hyperparameter tuned per task is the temperature 𝛼 which we show in Table G.2.

Note that in these experiments both the prior and constraint are conditioned on the state s . To prevent having to learn a neural network for log 𝑍 𝜙 ( s ) , we employ a variant of VarGrad objective [61]. For each state s sampled in the minibatch, we further generate 𝑘 = 64 on-policy trajectories 𝜏 ( 𝑖 ) 𝑘 𝑖 = 1 with 𝜋 𝛾 . Each of these trajectories can be used to implicitly estimate log 𝑍 ( s ) :

We then minimize the sample variance across the batch:

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 are 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 on the medium-replay tasks, with results shown in Table G.1, 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 D 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 5 random seeds at the end of training (150k training steps).

H Fine-tuning text-to-image diffusion models

We build off the DPOK implementation 8 , which fine-tunes stable-diffusion-v1-5 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 is 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 in memory. To help with this, we use a 'stochastic subsampling' trick (§H.1).

We train all models on a single NVIDIA A100-Large GPU. For the main experiments, we use the default parameters for DPOK of reward weight 𝛽 = 10 and KL weight = 0 . 01 . For RTB we fix 𝛽 = 1 . 0 for all prompts. We next perform an ablation over different values of 𝛽 .

We plot in Table H.1 the final average reward and diversity score for models trained with different values of reward weight 𝛽 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 𝛽 = 10 for RTB which has slightly lower final reward than 𝛽 = 1 , which we could attribute to more difficult optimization due to the peaky distribution associated with higher reward weight.

H.1 Memory-efficient 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 propagating gradient to a randomly sampled subset of the timesteps in a 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, so it is preferable to keep gradients for as many timesteps as possible to fit in memory. For our text-to-image 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 a single trajectory is just the sum of per-step log-likelihood gradients scaled by the RTB residual:

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) 𝑇 and accumulation batch size (number of time steps receiving a gradient signal in each backward pass) 𝐵 , the number of batched forward passes required scales as 𝑇 𝐵 .

Only the accumulation batch size 𝐵 , not the trajectory length 𝑇 , 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.

H.2 Generated images

H.2.1 A green-colored rabbit

H.2.2 Four roses

天

R

FOKSO

FOUR

ROORSS

FOUAS

Fouss

Fouss

FORUI

ToROss

FORM!

FOURS

FOSLSS

Fous

880

FouRRoFouRRest

urRown

FourRoss

FazrRos

ourcRoss

Fourlvm

FOMUNS

FR

Rolsos

FoUENs

OYS

Couss

H.2.3 A cat and a dog

高

H.2.4 A half - masked rugged laboratory engineer man with cybernetic enhancements as seen from a distance, scifi character portrait by greg rutkowski, esuthio, craig mullins.

福

省

I Compute resources

For classifier guidance experiments §3.1 we use train on a single NVIDIA V100 GPU. For textconditional image generation §3.2, text infilling §F and offline §G, we use a single NVIDIA A100 large GPU. The total estimated compute time for all our experiments is 3000 hours.

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 our 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 in 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 or 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 to other settings.

It is fine to include aspirational goals as motivation as long as it 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 not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or 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 a short proof sketch to 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of 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 §C, §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 are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to 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 is 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 to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions 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 a 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 be a way to access this model for reproducing the results or a 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 to have some path to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: Code is provided in §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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a 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

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Training details are outlined in §C, §E, §F, and §G.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or 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 not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at 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, or 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.

It 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 a 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 of range (e.g. negative error rates).

If error bars are reported in tables or 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Compute resources used in experiments is outlined in §E, §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 as well as estimate the total compute.

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://neurips.cc/public/EthicsGuidelines ?

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 a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 in §5.

Guidelines:

The answer NA means that there is no societal 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 particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to 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 as 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: We do not release models that have risk of misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-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 a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Creaters of code we 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 a URL.

The name of the license (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's 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/model as 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 as details about compensation (if any)?

Answer: [NA]

Justification: We do not have any crowdsourcing experiments or research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 to 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 we 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 as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nWe 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?\n\nQUESTIONS:\nSome 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)$?\n\nLIMITATIONS:\nYes. Limitations were included.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nConditional 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?\n\nWEAKNESSES:\n1. 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).\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n**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.\n\nWEAKNESSES:\n**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.\n\nQUESTIONS:\n* 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.\n\nLIMITATIONS:\nYes. For concerns and potential limitations, see the comments and questions above.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-30"}