Amortizing intractable inference in diffusion models for vision, language, and control

Authors: siddarth  venkatraman
moksh  jain
minsu  kim
marcin  sendera
mohsin  hasan
luke rowe mila
emmanuel  bengio
yoshua  bengio
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-10 14:28:33

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

Main paper: 
 introduction: 
 Diffusion models [67, 26, 71] are a powerful class of hierarchical generative models, used to model complex distributions over images [50, 11, 62] , text [4, 12, 39, 23, 22, 42] , and actions in reinforcement learning [29, 82, 31] 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 [54, 43] requires sampling from the normalized product of 𝜇(𝑎 | 𝑠) with a Boltzmann distribution defined by a 𝑄-function, 𝜋 * (𝑎 | 𝑠) ∝ 𝜇(𝑎 | 𝑠) exp(𝛽𝑄(𝑠, 𝑎)). In language modeling, various conditional generation problems [42, 22, 28] 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 [72, 32, 30, 10] or stochastic optimization [21, 47] . 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 [11] , but when such data is not available, one must resort to approximations or Monte Carlo estimates [69, 13, 9] , which are challenging to scale to high-dimensional problems. Reinforcement learning methods that have recently been proposed for this problem [7, 15] 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 [88, 61, 77, 64] , 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 [37] . 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). 
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). The prior is a diffusion model sampling a mixture of 25 Gaussians (a) and the posterior is the product of the prior with a constraint that masks all but 9 of the modes (b). Our method (RTB) samples close to the true posterior (c). RL methods with tuned KL regularization yield inaccurate inference (d), while without KL regularization, they mode-collapse (e). A classifier guidance (CG) approximation (f) results in biased outcomes. For details, see §C.A denoising diffusion model generates data x 1 by a Markovian generative process:EQUATIONwhere Δ𝑡 =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:EQUATIONThe 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 𝑢 𝑡 [52, 63] . 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 𝑝(x 0 , x Δ𝑡 , . . . , x 1 ) = 𝑝(x 0 ) 𝑇 𝑖=1 𝑝(x 𝑖Δ𝑡 | x (𝑖-1)Δ𝑡 )and defines a marginal density over the data space:𝑝(x 1 ) = ∫ 𝑝(x 0 , x Δ𝑡 , . . . , x 1 ) 𝑑x 0 𝑑x Δ𝑡 . . . 𝑑x 1-Δ𝑡 .A reverse-time process, x 1 → x 1-Δ𝑡 → • • • → x 0 , with densities 𝑞, can be defined analogously, and similarly defines a conditional density over trajectories:EQUATIONIn 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 𝑝 𝜃 :𝐷 KL (𝑞(x 0 , x Δ𝑡 , . . . , x 1 ) ∥ 𝑝 𝜃 (x 0 , x Δ𝑡 , . . . , x 1 )) (6) = 𝐷 KL (𝑞(x 1 ) ∥ 𝑝 𝜃 (x 1 )) + E x 1 ∼𝑞 (x 1 ) 𝐷 KL (𝑞(x 0 , x Δ𝑡 , . . . , x 1-Δ𝑡 | x 1 ) ∥ 𝑝 𝜃 (x 0 , x Δ𝑡 , . . . , x 1-Δ𝑡 | x 1 )) ≥ 𝐷 KL (𝑞(x 1 ) ∥ 𝑝 𝜃 (x 1 )) = E x 1 ∼𝑞 (x 1 ) [-log 𝑝 𝜃 (x 1 )] + const.where the inequality -an instance of the data processing inequality for the KL divergence -shows that minimizing the divergence between distributions over trajectories is equivalent to maximizing a lower bound on the data log-likelihood under the model 𝑝 𝜃 .As shown in [70] , minimization of the KL in (Equation 6) is essentially equivalent to the traditional approach to training diffusion models via denoising score matching [79, 67, 26] . 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,EQUATIONThis was named the trajectory balance (TB) constraint by [37] -by analogy with a constraint for discrete-space iterative sampling [45] -and is a time-discretized version of a constraint used for enforcing equality of continuous-time path space measures in [51] .In [60, 37] , 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 trajectoriesx 0 → x Δ𝑡 → • • • → x 1 usedfor 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 [64] . Such objectives contrast with on-policy, simulationbased approaches that require differentiating through the sampling process [e.g., 88, 77, 6, 78] .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) constraintEQUATIONfor every denoising trajectoryx 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 [45, 37] , 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):EQUATION)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 toEQUATION)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 [11] 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.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 [64] .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 [28] .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 [88] and shown to be useful for off-policy methods in [64] . namely, we writeEQUATIONwhere 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 [64] . This amounts to a relative variant of the VarGrad objective [60] (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).RTB as TB under the prior measure. The theoretical foundations for continuous generative flow networks [37] 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, 37] . 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 [28] to amortize intractable inference in autoregressive language models. 
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.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 [26] 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. on-policy from the current posterior model. We compare RTB with two RL-based fine-tuning techniques derived from DPOK [15] and DDPO [7] and with two classifier guidance baselines, namely DPS [10] , and LGD-MC [69] . 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.Diffusion models for text-conditional image generation [e.g. 62] 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 [7, 15, 76] .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 p to be close to the prior 𝑝:EQUATIONThe optimal p for (12) is p(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 [15] optimize an upper bound on the marginal KL in the form of a per-step KL penalty -𝛾 𝑇 𝑖=1 𝐷 KL [ p(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 p(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 [62] as a prior over 512 × 512 images. Following DPOK [15] , we use ImageReward [86] , 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 [7] .We measure the final average reward and the diversity of the generated image, as measured by the average pairwise cosine distance between CLIP embeddings [57] 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 [15, 86] . 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.To evaluate our approach on discrete diffusion models, we consider the problem of text infilling [90] , which involves filling in missing tokens given some context tokens. While discrete diffusion modelsunlike their continuous counterparts -can be challenging to train [4, 8, 48, 73] , score entropy discrete diffusion [SEDD; 42] 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 [49] , a dataset of short stories containing 5 sentences each. We adopt the task setup from [28] , 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 [28] , 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 [28] , we use three standard metrics to measure the similarity of the generated infills with the reference infills from the dataset: BERTScore [89] (with De-BERTa [25] ), BLEU-4 [53] , and GLEU-4 [85] . 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.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 [74] . In the offline RL setting [38] , 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 [54, 24, 43] . 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:EQUATIONwhere 𝑑 𝜇 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 [35] .Setup. We test on continuous control tasks in the D4RL suite [17] , which consists of offline datasets collected using a mixture of SAC policies of varying performance. We evaluate on the halfcheetah, hopper and walker2d MuJoCo [75] 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 [36] , and IQL [35] ) and diffusion-based offline RL methods which are currently state-ofthe-art: Diffuser [D; 29], Decision Diffuser [DD; 2], D-QL [82] , IDQL [24] , and QGPO [43] . 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 [43] , which learns intermediate energy densities for diffusion posterior sampling. 
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 [40, 14] . Amortized Bayesian inference [34, 59, 58, 19] is another domain of sampling from product distributions where diffusion models are now being used [20] . Beyond product models, [18] studies ways to amortize other kinds of compositions of hierarchical processes, including diffusion models, while [66] 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., 16 ]. 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 [28] .Diffusion samplers. Several prior works seek to amortize MCMC sampling from unnormalized densities by training diffusion models for efficient mode-mixing [6, 88, 77, 61, 78, 3] . Our work is most closely related to continuous GFlowNets [37] , which offer an alternative perspective on training diffusion samplers using off-policy flow consistency objectives [37, 87, 64] . 
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 [33, 64] to improve sample efficiency and mode coverage. Simulation-based objectives in the style of [88] 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 [44] . 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. [65] 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.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.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. 
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. 
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. 
 
 Appendix: Proposition 1 (Relative TB constraint). If 𝑝 𝜃 , 𝑝 post 𝜙 , and the scalar 𝑍 𝜙 jointly satisfy the relative trajectory balance (RTB) constraintEQUATIONfor every denoising trajectoryx 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),𝑝 post 𝜙 (x 1 ) = ∫ 𝑝 post 𝜙 (x 0 , x Δ𝑡 , . . . , x 1 ) 𝑑x 0 𝑑x Δ𝑡 . . . 𝑑x 1-Δ𝑡 = 1 𝑍 𝑟 (x 1 ) ∫ 𝑝 𝜃 (x 0 , x Δ𝑡 , . . . , x 1 ) 𝑑x 0 𝑑x Δ𝑡 . . . 𝑑x 1-Δ𝑡 = 1 𝑍 𝑟 (x 1 ) 𝑝 𝜃 (x 1 ) ∝ 𝑝 𝜃 (x 1 )𝑟 (x 1 ),as desired.Now suppose that 𝑝 𝜃 also satisfies the TB constraint (7) with respect to 𝑞(x 1 ). Then, for any denoising trajectory,EQUATION)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), giving1 = 𝑝 post 𝜙 (x 1 ) 𝑞(x 1 )𝑟 (x 1 )/𝑍 𝑞(x 1 )𝑟 (x 1 ) = 𝑍 𝑝 post 𝜙 (x 1 ). Integrating over x 1 shows ∫ 𝑞(x 1 )𝑟 (x 1 ) 𝑑x 1 = 𝑍. □The theoretical foundations for continuous generative flow networks [37] 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 [37] 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 . 3If 𝑝 𝜃 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 EQUATIONwhich 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.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:EQUATION{𝜇 𝑖 } = {(-10, -5), (-5, -10), (-5, 0), (10, -5), (0, 0), (0, 5), (5, -5), (5, 0), (5, 10)} (17) { π𝑖 } = {4, 10, 4, 5, 10, 5, 4, 15, 4}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 [7] .Sampling according to the RL policy leads to a distribution 𝑞 𝜃 (x 1 ), which is trained with the objective:EQUATIONWhile the exact computation of 𝐾 𝐿 (𝑞 𝜃 (x 1 ) ∥ 𝑝(x 1 )) is intractable, we follow the approximation method introduced by Fan et al. [15] , 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:EQUATIONNote that this is a biased approximation of the true intractable score:EQUATIONFor our experiments, we follow the source code 4 provided in recent diffusion sampler benchmarks [64] .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 [64] . 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 [15] , DDPO [7] ) typically use policy gradient style methods that are on-policy. By using a simple off-policy trick introduced by [46, 37] and demonstrated by Sendera et al. [64] , 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. On the other hand, policy gradient RL methods which are typically used to finetune diffusion models are on-policy, and hence prone to mode collapse. 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 [26] , 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 [27] , 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 [15] . We implement two RL-based fine-tuning techniques derived from DPOK [15] and DDPO [7] , 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 [10] . We implement and adapt the Gaussian version of the posterior sampling scheme in [10] , 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 [69] . We adapt the implementation of the algorithm in [69] 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.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.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-small5 model, which uses an absorbing noising process [4] with a log linear noise schedule, as the diffusion prior 𝑝(z | x). The posterior model is I put my costume on and went to the party.My friends loved my costume.Allen thought he was a very talented poet. He attended college to study creative writing. In college, he met a boy named Carl.Carl told him that he wasn't very good.Because of this, Allen swore off poetry forever.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.For training 𝑝 reward we follow the training procedure and implementation from [28] 7 . Specifically, we fine-tune a GPT-2 Large model [56] on the stories dataset with full parameter fine-tuning using the trl library [80] . 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 [28] . A critical difference in our experiments compared to [28] 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) xs 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 [28] . 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 [84] .We also attempted a LLMEval [41] for evaluating the coherence of the stories but did not obtain statistically significant results. 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 [60] . 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):EQUATIONWe then minimize the sample variance across the batch:EQUATIONRTB 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.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).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 [27] since it is significantly more efficient than fine-tuning the entire model. Sampling of images is done with 50 steps of DDIM [68] . 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. 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:∇ 𝜙 L RTB (𝜏; 𝜙) = 2 log 𝑍 𝜙 𝑟 (x 1 ) + 𝑇 ∑︁ 𝑖=1 log 𝑝 𝑝𝑜𝑠𝑡 𝜙 (x 𝑖Δ𝑡 | x (𝑖-1)Δ𝑡 ) 𝑝 𝜃 (x 𝑖Δ𝑡 | x (𝑖-1)Δ𝑡 ) • ∇ 𝜙 𝑇 ∑︁ 𝑖=1 log 𝑝 post 𝜙 (x 𝑖Δ𝑡 | x (𝑖-1)Δ𝑡 ).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. • 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. , 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.• 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.• 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. • 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.The time indexing suggestive of an SDE discretization is used for consistency with the diffusion samplers literature[88,64]. The indexing x 𝑇 → x 𝑇 -1 → • • • → x 0 is often used for diffusion models trained from data.Full prompt for "Cybernetic man": "A half -masked rugged laboratory engineer man with cybernetic enhancements as seen from a distance, scifi character portrait by greg rutkowski, esuthio, craig mullins."Recall that the relative density (or Radon-Nikodym derivative) of a distribution with density 𝑝 under the Lebesgue measure relative to one with density 𝜋 is simply the ratio of densities 𝑝/𝜋.https://github.com/GFNOrg/gfn-diffusionhttps://huggingface.co/louaaron/sedd-smallhttps://github.com/louaaron/Score-Entropy-Discrete-Diffusionhttps://github.com/GFNOrg/gfn-lm-tuninghttps://github.com/google-research/google-research/tree/master/dpok 


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

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

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

Reviews: review: paper_summary: paper_summary: This paper proposed a method to train a posterior $p^{post}(x)$ given a prior distribution $p(x)$ and some additional (possibly unnormalized) constraint function $r(x)$, when the prior is a pretrained diffusion model. Through the choice over $r(x)$, this setup can capture a wide range of tasks. The authors propose *relative trajectory balance* (RTB) constraint, which, if satisfied by the posterior under training, guarantees that $p^{post} \sim p(x)r(x)$. By training $p^{post}(x)$ to satisfy the RTB constraint, one can recover the desired posterior. Importantly, this training objective allows for efficient computation of the loss and off-policy training. The paper provides convincing empirical results on a variety of tasks ranging from conditional generation to offline RL.
strength_weakness: strength_weakness: **Scalability**
* 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.
  * 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.
  * 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).

**Experimental Results**
* 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).
* The text-conditional image synthesis experiment seems to be done for only four hardcoded text prompts, which is pretty limited.
* 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.
* 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.
questions: questions: * For the MNIST even-digit experiment, why was the reward function chosen to be $max_{i \textrm{ even}} p(c=i \mid x)$? Wouldn't $\sum_{i \textrm{ even}} p(c=i \mid x)$ be a more reasonable reward? I'm curious if this particular choice of reward makes it more likely for certain methods to mode collapse (e.g. by maximizing the probability for a single even class), and whether the results in Table 2 would change for MNIST with a different $r(x)$.
* 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.
* How long is the ratio between prior vs. posterior training in terms of flops or wall-clock time?  How quickly the posterior can be fine-tuned could limit the applicability of this method on certain real-world tasks.
limitations: limitations: Yes. For concerns and potential limitations, see the comments and questions above.
paper_summary: paper_summary: The paper looks at the problem of finetuning/training a generative model to sample from a desired posterior distribution when given access to a diffusion model prior. The experiments validated the generation capability across different tasks that diffusion models (and conditional generation) can be applied to - text infilling, text-to-image generation and training policies through offline RL. The work proposes a fine-tuning objective to train the model that will sample from the posterior (which in this case is instantiated from the prior’s checkpoint), which can be computed using off policy data/trajectories, ie. trajectories that are not all sampled from the posterior model.
strength_weakness: strength_weakness: 1. Speaking very qualitatively, the provided generations in Fig 4 and H.1 don’t look very different from each other and DPOK (and often even DDPO) seems to be generating samples of the same quality. This leads me to think that marginal improvements in the metrics provided (log likelihood etc) are not strongly tied to the generation quality. Also how were these samples chosen? To avoid cherry-picking and encourage reproducibility, i’d suggest mentioning the seed they were sampled with in the corresponding code readme files.
2. 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.
3. 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.

Overall, 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).
questions: questions: 1. On the text infilling task, am I correct in assuming that none of the other baselines had access to both x and y (the first 3 and the fifth sentences) since they are autoregressive and can’t condition on y? If yes, is the comparison to diffusion models that can condition on y fair? Should the comparison instead be with encoder-decoder or non-causal language models?
2. Is there a typo in Line 112? It says “... would sample the product distribution 𝑝post(x1) ∝ 𝑝𝜃 (x1)𝑟(x1). If 𝑟(x1) = 𝑝(y | x1) is a conditional distribution over another variable y, then 𝑝post is the Bayesian posterior 𝑝𝜃 (x1 | y).” Shouldn’t the product distribution be equal to the probably of the intersection of x and y and not y given x? Or is there a term which has implicitly been assumed to be constant/irrelevant?
limitations: limitations: The paper includes a very brief discussion on limitations which could be improved. I can provide more suggestions for this discussion after engaging with the authors and understanding their assumptions better.
paper_summary: paper_summary: This paper studies the problem of training diffusion models to sample from an intractable posterior distribution, defined by a prior diffusion model and an arbitrary likelihood function. The contributions are summarized as follows:

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

- The effectiveness of the proposed approach is validated through experiments in vision, language modeling, and continuous control benchmarks.
strength_weakness: strength_weakness: We use a posterior sampler, parameterized by a prior diffusion model with a learnable drift term $u^{post}$, to simulate trajectories and obtain samples $x_{1}$. Then, we compute the likelihood function $r(x_{1})$ and the prior $p_{\theta}$. Using these as an unnormalized density, the model parameters can be updated using GFlowNet-based loss. For fine-tuning a diffusion model, do we need to ensure that the initial state is an empty set (discrete) or a point mass (continuous)? In general, the initial state for DDPM follows a Gaussian distribution. Does this violate the GFlowNet assumption?
questions: questions: Some typos:

- Line 158: $x_{1} \leftarrow x_{0}$ --> $x_{1} \rightarrow x_{0}$?

- Eq11: $NN_{2}(x_{t}, t)$ --> $NN_{2}(t)$?
limitations: limitations: Yes. Limitations were included.

Comments: title: Any more questions?
comment: Dear Authors,

Thank you for your response. 

From what I understand, empirically, the results of RTB look similar to those of DPOK (or DPOK-like methods, ie, D-QL, like the authors said) on 2 of the 3 tasks: offline RL and text-to-image generation, and it was not tried out on the text-infilling task since that would require changes beyond what the DPOK paper had. Additionally, the qualitative examples/infills are not shown in the appendix for for the GFN baseline in the autoregressive prior category of models which seems to get the next best scores in table 3 so it is hard to say if the generated infills are much better than the next best performing method or not.

The experimental results do not demonstrate an advantage to using this approach, but the method is more principled compared to alternative methods (I have not validated the proofs in the appendix rigorously and will consult the other reviewers in the discussion phase to follow about this) and I will update my score to reflect this.
Dear Reviewer 4zDq,

Thank you again for your initial feedback. We have addressed your concerns in the above comment and added three new experiments based on your suggestions in the global response. 

Would you kindly let us know if these have affected your evaluation of the paper and if there is any further clarification we can provide?

The authors.
Thank you for your reply, and I will maintain my score.
Dear reviewer,

Thanks for your response. We are happy to clarify these points:
1. The training free baselines (DPS, LGD-MC) cannot be used with discrete diffusion models for the text infilling task and are only derived for the particular case of inverse problems in continuous space. DPOK/DDPO are RL finetuning baselines which can in principle be used for all the tasks, but do require significant changes to be adopted for the case of discrete diffusion (which was not considered in the original work). Note that **all the baselines are also biased samplers**, and hence expected to give worse performance for tasks requiring unbiased posterior inference such as classifier guidance (as we show in the paper) and other important scientific tasks (for example [1]). In the context of offline RL, the D-QL baseline trains the policy with an objective very similar to DPOK - The policy is trained to maximize the Q function with a per-step KL regularization term with behavior policy. We can add some notes in the appendix to make these connections clearer in the final version.
2. We agree that the claim "matches state-of-the-art" is more accurate, and will update the abstract in the revision.

**We hope we've addressed your questions and concerns. Feel free to reach out with any additional questions.**

[1] Adam, Alexandre, et al. "Posterior samples of source galaxies in strong gravitational lenses with score-based priors." arXiv preprint arXiv:2211.03812 (2022).
Dear Authors,

Thank you for your response.

I had the following clarification questions/thoughts:
1. " a general-purpose method that can be applied to a wide range of tasks simultaneously, achieving results comparable to specialized methods that had been developed for each problem." - is this true for all baselines considered in the work (that they are developed for specific problems)? Could the DPOK/DDPO baseline not have have been used on all the tasks? Could you explain why it was not applicable to the other two setups (Offline RL, text infilling)
2. On offline RL the method is slightly worse than both QGPO and D-QL (if we count the number of setups where a method is not within the top 5%) - and so I would suggest removing the claim from the abstract that it achieves state-of-the-art results - instead saying that it almost matches the state-of-the-art (if the other two methods are that).
Dear reviewers,

As we approach the end of the discussion period, we’d like to ask if we can give any final clarifications that would help you assess the paper.

Otherwise, if the responses and new results we shared earlier have resolved any doubts, please consider updating your ratings.

Thank you again for your time.

The authors