Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Authors: Siddarth Venkatraman
Mohsin Hasan
Minsu Kim
Luca Scimeca
Marcin Sendera
Yoshua Bengio
Glen Berseth
Nikolay Malkin
Venue: ICML 2025 poster (2025)
Keywords: diffusion
amortized inference
inverse problems
fine-tuning
VAEs
GANs
CNFs
flow matching
generative models
Primary area: probabilistic_methods->bayesian_models_and_methods
TL;DR: To sample Bayesian posteriors (for constrained generation, RLHF, and more) under a pretrained prior, we train diffusion models to sample the noise space, applicable to any generative model.

Abstract: Any well-behaved generative model over a variable $\mathbf{x}$ can be expressed as a deterministic transformation of an exogenous (‘*outsourced'*) Gaussian noise variable $\mathbf{z}$: $\mathbf{x}=f_\theta(\mathbf{z})$. 
In such a model (*eg*, a VAE, GAN, or continuous-time flow-based model), sampling of the target variable $\mathbf{x} \sim p_\theta(\mathbf{x})$ is straightforward, but sampling from a posterior distribution of the form $p(\mathbf{x}\mid\mathbf{y}) \propto p_\theta(\mathbf{x})r(\mathbf{x},\mathbf{y})$, where $r$ is a constraint function depending on an auxiliary variable $\mathbf{y}$, is generally intractable.
We propose to amortize the cost of sampling from such posterior distributions with diffusion models that sample a distribution in the noise space ($\mathbf{z}$). These diffusion samplers are trained by reinforcement learning algorithms to enforce that the transformed samples $f_\theta(\mathbf{z})$ are distributed according to the posterior in the data space ($\mathbf{x}$). 
For many models and constraints, the posterior in noise space is smoother than in data space, making it more suitable for amortized inference. Our method enables conditional sampling under unconditional GAN, (H)VAE, and flow-based priors, comparing favorably with other inference methods. We demonstrate the proposed ___outsourced diffusion sampling___ in several experiments with large pretrained prior models: conditional image generation, reinforcement learning with human feedback, and protein structure generation.

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Siddarth Venkatraman 12 Mohsin Hasan 12 Minsu Kim 123 Luca Scimeca 12 Marcin Sendera 4
Yoshua Bengio 125 Glen Berseth 125
Nikolay Malkin
Abstract
Table 1. Common families of generative models can be expressed as deterministic transformations of noisc. Here, distent refers to the dimension of the latent variable typically understood in the model; typically, diatent & ddata ($2). Posteriors under these generative model priors generally intractable, but some models can be fine-tuned by asymptotically unbiased variational objec- tives to sample intractable posteriors ($3.1). Outsourced diffusion sampling approximates posteriors in the noise space ($3.2 and $4). Invertible?
Any well-behaved generative model over vari- able x can be expressed as a deterministic transfor- mation of an exogenous ('outsourced") Gaussian noise variable z: x = fz(z). In such model (e.g.. VAE, GAN, continuous-ting model), sampling of the target variable x (x(x) is straightforward, but sampling from posterior distribution of the form p(x y) & pe (x)r(xX y). where r is a constraint function depending on an auxiliary variable y. is generally intractable. We propose to amortize the cost of sampling from such posterior distributions with diffusion models that sample distribution in the noise space (z). These diffusion samplers are trained by reinforcement learning algorithms to enforce that the transformed samples (z are distributed according to the posterior in the data space (x). For many models and constraints, the posterior in noise space is smoother than in data space, making more suitable for amortized inference. Our method enables conditional sampling under unconditional GAN, (H)VAE, and flow-based priors, comparing favorably with other inference methods. We demonstrate the proposed out- sourced diffusson sampling in several experiments with large pretrained prior models: conditional image generation, reinforcement learning with human feedback, and protein structure generation.
flow-based
Model
Noise dim. daiio
Variational tuning?
(H)VAE
diacn N + dani
x
GAN
duse
/
dam dam
Diffusion
Latent diffusion
duac (T daas
inference problems. The aim is to approximate or sam- ple from the product of the modeled distribution over a data space with an observation likelihood or other constraint function. While such diverse applications as conditional gen- eration (Dhariwal & Nichol, 2021; Ho & Salimans, 2022), inverse problems (Song et al., 2022: Chung et al., 2023; Venkatraman al., 2024), and constrained improvement from human feedback (Korbak et al., 2022; Fan et al., 2023) can be cast as posterior inference tasks, sampling from such posterior distributions when no unbiased target data is avail- able is generally intractable. For some model families, ap- proximate solutions, such as MCMC, approximate guidance, and variational inference, may be possible. Each of those methods has limitations, such high cost reach conver- gence for multimodal posteriors, intractability of accurate density estimation, and reliance on techniques specialized to the model and constraint.
1. Introduction
Fundamentally, generative models are probabilistic pro- grams that produce samples from the distributions they de- fine by combination of deterministic computation and injection of random noise.' This paper argues that the noise space of generative models - the space where the noise in- jected during generation resides - is an effective target for posterior inference. To be precise, we consider a generative
Generative models, trained on a dataset to maximize likeli- hood or related quantities, can become priors for Bayesian
"qqual contribution 'ila Québec Al Institute Univerrsité de Montréal kKAIST "Jagiellonian University SCIFAR University of Edinburgh. Correspondence to: <(siddarth, venkatraman, massin } @mila..uebec>>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
We refer here to models that produce samples in bounded number of operations, not objects such deep energy- models, which sampling is intractable.

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Figure Left:: Top row: Marginal densities CNF that transforms Gaussian distribution = 0) to Swiss roll = 1). Middle row: The constraint function mixture of two Gaussians centered an an observation y (+) and its reflection through the origin pulled back X,. Bottom
2-
0.00
r- 0.50
t- 0.75
t- 1,00
Posterior densities Nr, propor- tional to the product of the first two rows, The rightmost column shows samples the data space. Right: The same objects shown in noise and data space for GAN that transforms noise to data (x). Outsourced dif-
p0r1 xi
ply X
pur
pie 19)
giz
fusion samplers approximate P(No y) or p(z y), which are smoother than p y) (see Fig. 2).
model that expresses data as a deterministic transformation of noise, = (Z), where the noise variable z follows known distribution
Target
(z(z) and are the model parame- ters. This model defines distribution over x the pushfor- ward [/o]-P, of the noise distribution by the deterministic transformation with density pxx). A constraint function yxx.y) y) in the data space, depending on an auxiliary variable y. defines a posterior distribution p(x y) & pxx)r(x,) y). This posterior can be sampled by inferring a distribution over the noise variable that, when transformed by fe. aligns with the posterior in the data space. The posterior in the noise space is often smoother (and lower-dimensionall than the corresponding posterior in the data space (Fig. 1), making it more amenable to efficient sampling (Fig. 2).
Outsourced
piks y)
Figure 2. Marginal densities diffusion sampler the posteriors from the CNF example in Fig. I in data space and noise space. The data space posterior (top row) has well-separated modes and harder to sample from than the outsourced posterior (hottom row).
complex posteriors ($5).
While posterior sampling noise space still intractable, can be addressed by methods of black-box variational inference. Recent advances in diffusion samplers diffusion dataset, but to match given
2. Outsourcing Noise in Generative Models
Consider a probabilistic model over variable taking val- ues in Rdum, with auxiliary latent variables valued in Ruum The model is joint distribution induces distribution over its marginalization over w. In terms of densities (if they exist), if p(w, x) is the joint den- sity, then the marginal density of x is p(x) = p(w, x) dw.
models trained not unnormalized density (Zhang & Chen, 2022; Vargas et al., 2023; Richter & Berner, 2024; Sendera et al., 2024) open an opportunity to model complex posteriors in noise space. We call such amortized posterior inference in the noise space outsourced diffusion sampling.
and w, and it
A form of the noise outsourcing lemma (see, e.g., Austin, 2015) states that, under basic assumptions, any such model is equivalent one augmented with additional latent vari- ables w', independent of w and following a fixed distribu- tion, such that x is a (deterministic) function of w and w'. In particular, if w and w' are both standard Gaussian, then x is a deterministic function of Gaussian noise variable z (the concatentation of w and w'), called the 'outsourced' noise:
Our exposition and experiments support three claims:
(1) Outsourced diffusion sampling is agnostic the form of the mapping fe and applicable to a wide range of prior models, including VAES, GANs, normalizing flows, and continuous-time flow-based models (Table l).
(2) Outsourced diffusion sampling an effective poste- rior inference method under large pretrained genera- tive model priors in variety of domains: conditional image generation, reinforcement learning with human feedback, discriminator- adjusted GAN sampling, and protein structure generation (Table 2).
Proposition 2.1 (Noise outsourcing lemma for Gaussians). Let w and be Borel-measurable random variables val- ued in Rdwa and Rdaan, respectively, with w marginally standard Gaussian, and let duoise distent- There exists random variable z in Rduace such thut:
(3) Outsourced diffusion sampling is more efficient than amortized inference methods that fit model to sample the data space posterior directly and than non-amortized methods like MCMC, illustrating the flexibility of diffu- sion sampling in outsourced noise spaces for sampling
(1) z is standard Gaussian;
(2) w the projection of onto its first diatent coordinates;
(3) there exists a measurable function f R.dme - R Rdm. such that (w, x) (w, fz)) almost surely.
2

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Prop. 2.1 ensures that any generative model with marginally Gaussian latent variables can be rewritten as a deterministic function of higher-dimeensonal Gaussian noise variable, but does not specify the form of the function f (which is very non-unique). We will be interested in modeling Bayesian posteriors over x given observations by pulling them back to the noise variable z, using methods agnostic to the form of f. which is seen as a black-box transformation.
differential equation (ODE)
dx, = v (x,, dt.
This includes ODEs derived from diffusion models, e.g-, DDIMs (Song et al., 2021a), and those trained with flow matching, the family of methods introduced by Lipman et al. (2023); Alhergo & Vanden-Fijnden (2023); Liu et al. (2023).
Under regularity conditions on va. a distribution over initial conditions No induces marginal distributions over X, for > O, and in particular over the data variable x1- The CNF is generative model with outsourced noise variable z = xo-
Common families of generative models have natural form for f. obtainable explicitly from their latent variable struc- ture, that we will exploit. We now explain how several model families, in their basic form, can be expressed as de- terministic transformations of noise z following a Gaussian distribution over R Rdmine See Table 1 for a summary.
Diffusion models. Diffusion models -Sohh-DDckktein et al., 2015; Ho al., 2020) and latent diffusion models (Rom- bach et al., 2021) can also be expressed as deterministic transformations of noise; see Appendix A.1 for discussion and connections.
Variational autoencoders (VAEs; Kingma & Welling, 2014). The generative model in simple VAE may have the form x ~ (W) ( ( (dm)))) where w follows Gaussian distribution in R Rduen and μ4 and dge are neural networks outputting a vector and scalar, respectively. This model may be reparametrized as
3. Posteriors under Generative Model Priors
For a generative model p(x) of any of the types described in $2, and a positive constraint function r(x.y) such that Ex-p(x [r(x. is finite, we are interested in sampling the posterior distribution p(x y) xe (x))r(x,y). Various sources of constraints will be described in $5 (see Table 2).
x (w(W) (W)))
where & ~ (O, Idam). Thus x is a deterministic transfor- mation of the concatenation of and É. which follows Gaussian distribution in tdaun (The encoder, an aux- iliary object used in training the VAE, does not form part of the generative model.) Hierarchical VAFs EsV Rezende et al., 2014), generalize VAES, using a Marko- vian chain of latent variables in the generative process i.e., a graphical model structure of WN - - Wi - x with each transition conditional Gaussian distribution. If -valued, then can be similarly reparametrized function of the N dideensiom Gaussian noises injected on each transition W;+1 - W; and the dimenniins noise on the last step, as in (1).
3.1. Posterior Sampling and Approximation
In this section, we describe existing methods for sampling approximately from such intractable posteriors.
Model-agnostic methods. The most general methods for sampling from distributions defined by unnormalized densities are Markov chain Monte Carlo (MCMC) methods. These methods may not require fitting parametric models, although hybrid methods such as adaptive importance sampling (Bugallo et al., 2017), twisted SMC variants (Lawson et al., 2018; 2022), and neural boostrap algorithms (Midgley et al., 2023) - can accelerate their convergence.
these variables are all -dluuea
Generative adversarial networks (GANs; Goodfellow et al., 2014). In a GAN, a generator Go maps Gaussian- distributed noise ~ NO, Idma.)) deterministically to data, x (z). Thus GAN is naturally model with out- sourced noise in Rduu,, (The discriminator is an auxiliary object used in training, not part of the generative model.)
MCMC methods are agnostic to the form of the target dis- tribution and are guaranteed converge to it under mild conditions in the limit of infinite time (or memory in the case of particle filtering methods like SMC (Del Moral et al., 2006; Doucet et al.. 2009)), making them anytime algorithms that can trade off computation cost for accuracy. However, MCMCs assume access to the target density and possibly to its gradient, limiting their applicability:
Normalizing flows (NFs; Rezende & Mohamed, 2015). In a NF - also naturally a generative model with out- sourced noise- the generator So maps Gaussian noise z NO, Ida..) to data x deterministically, x fz)... Un- like GAN generator, the function fo is constrained to be in- vertible and necessarily (in order to model a full-support dis- tribution) must have noise of the same dimension as the data.
For VH)S and their special case diffusion models, the density cannot be computed exactly; only variational bounds are available.
For GANs, the density cannot be computed because the generator is not injective (invertible) and may not even define a full-support distribution over the target space. For CNFs, the density can be approximated using the
Continuous normalizing flows (CNFs; Chen et al., 2018) A CNF is an invertible transformation from noise z = Xo
to data x X1 that the solution of neural ordinary
3

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Hutchinson trace estimator (Hutchinson, 1989; Grath- wohl al., 2019), but accurately computing the gradient is expensive, as it requires backpropagating through the computation graph of neural ODE integrator.
with v < fau, then fo (( (f. o f) 1) =v.
In particular, if y is a probability measure and h : X Rzo is f-u-integrable, then A == [hdlsin) (h f) 4 is probability measure on Z, and fad= = Jhdlfau) (h saj) is
Monte Carlo methods in latent space. MCMC sampling can be performed at intermediate time points of continuous- time flow-based models (Cabezas et al., 2024). MCMC techniques are also used for sampling in GAN latent spaces for discriminatoo-guidded sampling (Che et al., 2020; Hou et al., 2025) and for conditional generation by sampling in intermediate activation spaces (Nguyen et al., 2017). Simi- lar approaches are applied to normalizing flows (Coeurdoux et al., 2024). Monte Carlo techniques are also used in diffu- sion models (Appendix A.
a probability measure on X.
(See Appendix A.2 for the proof.) In common terms, in terms of densities, the relevance of the proposition of our setting is as follows. Let f R - Rdam be function from the noise space to the data space. A prior density p(z) in the noise space, transformed via x = f(z), defines a prior p(x) (a density with respect to some reference measure, such as the volume measure on the image of f). If h(x) = =(xx y) is a constraint function with f y)x dx < co, then if z' is a variable in the noise space distributed with density proportional to p(z')r(f'')) and x' = f(z',, then x' is distributed with density proportional to p(x')r(x', y) in the data space. This means that to sample the posterior in latent space, we can sample from the posterior in noise space (with density (f y)) and transform the sample by f. (Note that such sampling does not require computation of the push- forward density. Indeed, f need not be injective or smooth, as would typically be required for such computations.)
Amortized inference and fine-tuning, For some families of models, it is possible to train model that, at convergence to the global optimum, samples from the posterior distribution exactly. This is a problem of variational inference: the model is trained - or perhaps fine-tuned using the prior model as initialization - to be close to the target distribution in some measure of divergence.
For CNFs, method for fine-tuning the drift function vo to yield CNF that samples from the posterior distribution, known as adjoint matching, has recently been proposed by Domingo-Enrich et al. (2024). While this method is asymptotically unbiused, it requires access to the gradient of the likelihood function. Furthermore, it is only applicable to a narrow class of fow-based models, namely, those that are trained from certain marginal couplings and interpolants and closely related to the probability flow ODEs of diffuusin models. This restrictiveness is due to adjoint matching converting the neural ODE to an equivalent neural SDE?, which is not possible in general (e.g., for flow-based models trained using minibatch optimal transport couplings (Tong et al., 2024: Pooladian et al., 2023) or with Con-Gaussian source distributions), Naively applying adjoint matching to such CNFs gives biased results (Fig. 3).
Amortizing outsourced posterior sampling. Although noise space posteriors might be simpler than the distribution in target space (Figs. 1 and 2), they can still be multimodal and high-dimensional. MCMC methods have been used to sample from noise spaces of NFs and GANs (Che al., 2020; Cannella et al., 2021), but suffer from long mixing times. In addition, many MCMC methods assume that the target density p(z)r((()), (z), y) is (efficiently) differentiable. which is not the case when f is CNF.
Instead, it can be desirable to use amortized variational inference to fit a fast sampler to the latent posterior, that is, to approximate by a parametric model. We have no samples from this posterior, but have access to its unnormalized density R(z y) = pf
3.2. Bayesian Posterior in Noise Space
We call such a model an outsourced sampler, a name moti- vated by the fact that the factors of variation in the posterior are 'outsourced' to the noise space via the pullback opera- tion. We shall use diffusion models as the variational family, as will be discussed in $4.2.
We describe how posterior distributions can be pulled back to the noise space of generative model expressed as de- terministic transformation of an outsourced variable. This relies on a basic measure- measure-theoretie theoretic fact regarding the trans- formation of density functions under pushforward measures:
4. Outsourced Diffusion Sampling
Proposition 3.1. Suppose that (Z, Ez) and (X, Ex) are measurable spaces and f Z. + X is measurable. If y is a o--finite measure on Z and v is a o-finite measure on X
4.1. Diffusion Samplers for Amortized Inference
Diffuusion sampling is the variational inference problem of approximating a distribution over Rd, with a given unnor- malized density R Rd + Rso, by a diffusion model. Samples from the target distribution, which has density
In diffusion models (SDEs), the prior model can also he fine- tuned sample from the posterior using objectives closely related to those proposed here; see Appendix A.
4

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Ptarger (2) 1 R(),, are not available, nor do we have access to the normalizing constant Z; however, we have the ability to query for the unnormalized density R(z) at any point The goal is to train a neural stochastic differential equation
y) == p())r((((),))), The sampler can be conditioned on the auxiliary variable y (taking it as input, resulting in amortization over y and the possibility of generalizing to new y) or can be trained for single, fixed value of y to sample the corresponding posterior.
Zo (a) dz, u (Z, () dt o, dB,, (2)
Training objective. To train a diffusion model to sample from the target density R(z y), we use the trajectory bal- ance objective (TB; Malkin et al., 2022). TB was first intro- duced in the context of (discrete) generative flow networks (GFlowNets; Bengio et al., 2021; 2023) and generalized to the continuous setting in Lahlou et al. (2023); it is also a close relative of the VarGrad objective (Richter et al., 2020). It was comprehensively evaluated for diffusion samplers in Sendera et al. (2024) and its asymptotic consistency in the continuous-time limit was established in Berner et al. (2025). We briefly review the TB objective and refer to those works for further details.
where a is a neural network, of a scalar function of time, and B, is standard Brownian motion, that the induced distribution over
z is close to the target distribution PZrger(Z) Z) in some measure of divergence. (Note that, unlike for diffusion models trained from data, it is standard for generation to proceed in increasing time (from noise at == 0 to the target at = 1.) The model can be sampled by simulating (2) in a time discretization, e.g., using the Euler- Maler-Maruyama method.
Training diffusion samplers is more difficult than training typical diffusion models, i.e., maximizing a variational bound on log-likelihood of a dataset. Various objectives have been proposed, including: (1) ones that rely on differ- entiable simulation of the generative process during training (Li et al., 2020; Kidger et al., 2021; Zhang & Chen, 2022; Vargas et al., 2023) and are linked with optimal control (Berner et al., 2024; Vargas et al., 2024), (2) ones using biased but asymptotically consistent Monte Carlo estimates of the score function (Vargas et al. (2022); Huang et al. (2024); Akhound-Sadegh et al. (2024)), and (3) 'off-policy" divergences that can be optimized on arbitrary generative trajectories not necessarily sampled from the current itera- tion of the model (Richter et al., 2020; Nüsken & Richter, 2021: Sendera et al., 2024, inter alia). A unifying perspec- tive on these methods and analysis in the continuous-tin limit were recently given by Berner et al. (2025).
A difftusinn model - neural network with parameters & and possibly conditioned on defines a Markovian distribution over denoising trajectories T = (Zo - Zar + - 21). where At- = + the time step of the discretization of the SDE (2), via
T [1
PT =))=p(z))
(f I((-1))..)) i=1
(3)
Here p(zo) is the density of a fixed distribution over the initial noise (recall that generation goes forward in time) & is the density of the transition kernel defined by the and Pp model, i.e., the probability of transitioning from a sample at given noise level to a sample at the next-lowest noise level. Similarly, the (fixed) noising process defines a distribution over noising trajectories conditioned their terminal end- point: Pa(T 21) = IT P= Ziar). The TB ob- jective associated with trajectory r is a squared log-ratio:
In this work, we adopt training methods of the third kind, using off-policy divergences, as they have have two notable advantages. First, they treat the target as black-box re- ward function and do not require access to the score V log R, as differentiable simulation methods and some Monte Carlo methods do. Second, they can be trained off-policy. on trajectories obtained through exploration, allowing the flexi- ble use of exploration strategies and thus promoting mode discovery, as demonstrated for continuous-space samplers by Malkin al. (2023); Sendera et al. (2024); Phillips & Cipcigan (2024); Kim et al. (2025). The exact form of the off off-policy divergence will be described in $4.2.
Zo(r) ly)) ) LTB(T,y, (108 log R(z: 1))B(T |z))
(4)
where Za is learned model that, at optimality, estimates the partition function f R(z y) dz. This objective aims to match two distributions over trajectories: the one defined by the denoising model and that defined by the target dis- tribution and the noising kernel. If the two distributions are equal, then their marginal densities at = also coincide.
4.2. Diffusion Sampling in Noise Space
For training, one draws trajectories T from some training distribution (which is not necessarily the current model pp) and optimizes (4) with respect to the parameters $. If the TB loss is optimized to 0 for every trajectory T in the continuous- time limit, the model p asymptotically samples from the target density R(z y) (Berner et al., 2025). The training loop described in Algorithm 1.
Diffusion models are an attractive choice of variational family due to their ability to sample from complex high- dimensional distributions. We train outsourced diffusion samplers by using the methods introduced in $4.1 to approx- imate posteriors in outsourced noise spaces ($3.2).
The target density we wish to sample takes the form R(z|
5

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Figure (a) Flow paths of a CNF trained with OT-CFM (Tong et al., 2024) from a source distribution (+) to the '2 moons' distribution (+). The source is Gaussian (top row) or mixture of 8 Gaussians (bottom row). (b) The constraint constructed such that posterior is the lower moon. CNF flow paths from the lower moon (target posterior) to the source latents (outsourced posterior). (c) Flow paths from naive application of Adjoint Matching (Domingo- Enrich et al., 2024). which is is biased for OT flows and non-Gaussian sources. (d) Flows starting samples from an outsourced diffusion model, which samples the latent posterior, give target samples close the ground truth.
(a) OT-CFM prior (b) True Posterior (c) Adj. Matching
(d) Ours
Algorithm Training loop for Outsourced Diffusion Sam- pling
5.1. Class-Conditional Sampling
Setup. The prior model pe(x) is an off-the-shelf uncon- ditional image generator trained on the CIFAR-10 10 dataset (Krizhevsky, 2009). Using a CIFAR-10 classifier p(y x), train posterior cass-connitioned generative model p(x y) or (x x). For our experiments we work with two priors which achieve high fidelity unconditional generation: a flow matching (CNF) model trained with inde- pendent coupling and linear interpolants (I-CFM; Tong et al., 2024), and spectrally normalized GAN (SN-GAN; Miy- ato et al., 2018). We use a 13-layer VGG-net model as the classifier (Simonyan & Zisserman, 2015). We compare sourced diffusion sampling against powerful MCMC base- line and recent method for fine-tuming flow-based models:
1: Initialize: deterministic prior function f, randomly initialized noise posterior model p randomly initial- ized (y), VP-SDE backward policy Pa. log reward function logr(x, y), on-policy update fraction p.
2: for each step ı = 1, ,,...,N N do
3: Sample a batch trajectories: (T -Pp(T 1y)
4: for I..... do
5: Compute log density:
log R(1) - log (( (....D++(((a)) +
6: Store experience (7 log R()) in buffer D
7:
end for
8:
Draw Uniform(O, 1)
9:
if u I then
Hamiltonian Monte Carlo (HMC; Brooks et al., 2011) applied in the noise spaces of both the SN-GAN and CFM priors sample the outsourced posterior. We highlight that HMC with CNF priors can be quite slow, since requires differentiating through the ODE integrator.
10:
Keep on-policy batch ((T(0), log R )
11: else
12:
Sample off-policy batch ((T(0, log R (
end if
14:
Compute LTB(T;Y, Φ)) for batch using Equation 4.
Adjoint Matching EDorih et al., 2024): see $3.1 and Fig. 3 for discussion. Note that adjoint matching requires access to the gradient of the constraint function.
15:
à Update pp. - using V VTLTB(T:y.d)
16:
end for
For the SN-GAN prior, we train an outsourced diffusion model sample generator noise z E R512, For the CFM prior, we sample the initial noise latent E R 3x32x32 See details in Appendix B.2.
O, where
Exploration and credit assignment techniques. We bor- row number of off-policy exploration techniques (such as replay buffers), as well as methods to make training more stable (such as temperature annealing) from the diffusion samplers literature. For details, see Appendix B.
We report average log-reward (classifier log- likelihood) of samples and FID scores (computed with the dataset images of the given class) in Table 3. The CFM posteriors consistently outperform SN-GAN posteriors, re- flecting the superior quality of the CFM prior. Among the methods evaluated, adjoint matching with the CFM prior achieves the best performance as expected, since it is specifically designed for fine-tuning I-CFM models with a Gaussian source (samples in Fig. 8). Outsourced diffusion, a more general approach, also delivers strong conditional gen- eration (visualized in Fig. 4). Unlike adjoint matching and HMC, outsourced diffusion does not rely on gradient infor- mation from either the classifier or the prior. Additionally,
Results.
5. Experiments
The goal of our experiments is demonstrate the general applicability of outsourced sampling, highlighting tasks which lack specialized techniques for posterior inference. We list the different tasks, alongside the sources of priors, constraints, and the dimension noise sampled by the outsourced diffusion model in Table 2.
6

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Table 2. The priors and constraints studied in $5. Outsourced diffusion sampling works in noise spaces of wide range of generative models and is agnostic to their specific form. Task
Constraint
Prior
Prior type
dgoisc
daata
SN-GAN
GAN
128 3 x 32
3 3 32 x 32 3 32 x 32
CIFAR-10 classifer guidance CIFAR-I0 classifer
1-CFM
CNF
StyleGAN3 NVAE
GAN
512
3 x 256
FFHQ text conditioning
ImageReward
Hierarchical VAB
x x 8 8
3 x 256 X 256
Text-to-Image model RLHF
ImageReward
Stable Diffusion 3
Latent-CNF
16 64 x 64
3 x 512 x 512
Protein structure
Structure Diversity
FoldFlow 2
Riemannian CNF
7 64
7 64
Table 3. CIFAR-10 posterior sampling results for GAN and CNF priors. We report expected classifier log probability and FID scores for the class posteriors, averaged over all 10 classes.
Prior
Sampler
B llog p(y xT)
FID (J)
Prior
-5.37
97.14
SN-GAN
Latent HMC
33.26
75.33
(a) SN-GAN Prior
(b) Posterior (Dog)
(c) Posterior (Car)
Outsourced Diff.
-3.84
68.12
Prior
-5.88
84.79
Latent HMC
-2.80
46.69
I-CFM
Adj. Matching
-3.09
19.45
Outsourced Diff.
-3.35
34.28
(d) I-CFM Prior
(e) Posterior (Dog) (I) Posterior (Car)
& Kautz, 2020) and StyleGAN3 (Karras et al., 2021a), both of which achieve high-fideliyy unconditional generation.
Figure CIFAR-IO samples generated using SN-GAN (top row) and CFM (bottom row) priors and posterior samples from trained outsourced diffusion models for the 'Dog' and 'Car' classes.
NVAE is a deep hierarchical VAE with a large number of latents of different scales. Vahdat & Kautz (2020, Appendix B.6) notes that almost all the feature variance is captured by the first 4 levels of the latent hierarchy. We train the outsourced diffusion model to sample noise z R4x20x8x8 for these levels, which turns out to be sufficient for condi- tional generation. For StyleGAN3, we sample the generator noise space z E 522, Due to the absence of specialized variational techniques for posterior inference with GANs and VABs, we use HMC sampling of the outsourced noise, targeting the same distribution the outsourced diffusion sampler, as the baseline. More details in Appendix B.3.
the outsourced diffusion model offers significant advantages in training efficiency, requiring approximately 5 hours on an A100 GPU compared to 12 hours for adjoint matching (see Appendix C.2). Moreover, we show how we can distill one-step outsourced diffusion samplers in Appendix C. with no performance degradation and significant sampling speed advantage.
5.2. Conditional -esh-Rnsn Face Generation
Setup. Given generator of high-resolutin (256 x 256) human face images trained on the FFHQ dataset (Karras et al., 2021b) as the prior (x), we aim to generate faces aligned with a specified text caption y. To achieve this, we use constraint function given by ImageReward (Xu al., 2023), text-image reward model built on the BLIP backbone (Li et al., 2022) that scores images based their alignment with the provided text prompt and aesthetic quality. The ImageReward score serves as the log-constraint function log (x, y) in our formulation, enabling us to frame the text-conditional face generation problem as posterior inference. For the prior models, we employ NVAE (Vahdat
Results. We report average ImageReward score and diver- sity, measured as average cosine distance of CLIP (Radford et al., 2021) embeddings for 100 generated images from the posterior, for 4 different prompts in Table 4. We find HMC for StyleGAN3 can get stuck in bad reward modes, but some- times obtains high reward. HMC is consistently poor with the NVAE prior, which we attribute to combination of di- mensionality and high energy barriers. Outsourced diffusion samplers consistently generate prompt-accurate posterior samples. Illustrative samples are displayed in Fig. 5 and more uncurated samples in Appendix D. StyleGAN3 pos- teriors are of higher quality than NVAE posteriors, likely
7

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Table 4. FFHQ text-conditioning results for NVAR and StyleGAN3 priors. We report average log-reward and CLIP cosine distance (diversity) for posteriors, averaged over the prompts shown in Fig. 5.
NVAE
Prior
Sampler
CLIP Diversiry (T)
StyleGAN3
Prior
1.94
0.30
NVAE
Latent HMC:
-1.2
0.30
Outsourced Diff.
0.98
0,26
Prior
-152
0.36
StylGAN3
Latent HMC
-0622
0.31
Outsourced Diff.
1.23
0.26
Asian girl with glasses
Bald man, black beard
Brown haired child
Prior
Old
Figure 5. Sampled images from the FFHQ priors and outsourced diffusion posteriors for different prompts. More examples in Appendix D.
because the prior is also stronger.
Table 5, RLHF finetuning results for SD3 prior. We report ex- pected log reward and CLIP cosine distance (diversity) for posteri- ors, averaged over the prompts listed in Fig.
5.3. Text-to-Image RLHF
Setup. Diffusion and flow matching models that gener- ate images conditioned on textual prompts often struggle with complex prompts that involve compositional relation- ships. A promising strategy to address this limitation is to fine-tune such models using reward functions that quan- tify image-caption alignment. In previous work, Fan et al. (2023); Venkatraman et al. (2024) used Stable-Diffusion- 1.5 (Rombach et al., 2021) as a caption-conditionee prior p(x y) and ImageReward as the unnormalized log- likelihood logr(x,,)
Sampler
E[log rxx. ())
CLIP diversity (1)
Prior
0.791
0.19
CFG
0.84
0.17
Outsourced Diff.
1.27
0.16
A and dog.
A cat riding llama.
In both of those works, the prior diffusion model was tuned sample approximately from the posterior p plligned (x y) p(x| y)r(x,)) y). In our experi- ments, we instead align Stable Diffusion 3 (SD3; Esser et al., 2024), which is CNF. not diffusion model, and is thus unsuitable for fine-tuning using the mentioned techniques.
Prior
Posterior
Prior
Posterior
We train outsourced diffusion samplers of the CNF's 16 x 64 64-dimensional noise space. Gradient-basse posterior inference techniques, such adjoint matching and HMC, prohibitive for flow model at the scale of SD3, since they involve differentiating through the reward model, high resolution multiscale decoder, and ODE integrator. Instead, as a baseline, we tune the classifier-free wwee--rriddnneee weight (Ho & Salimans, 2022) individually for each prompt and report the best performance. See Appendix B.4 for details.
Quiet village disrupted by meteor strike.
A human with horse face and human with wolf face.
Figure 6. Sampled images from the SD3 prior and outsourced dif- fusion posterior for different prompts. More examples in Ap- pendix E.
5.4. Protein Secondary Structure Diversity
Results. We report the ImwgeReward score and the average CLIP cosine distance averaged over 4 prompts in Table 5. We present illustrative examples in Fig. 6, with further un- curated samples provided in Appendix E. We find that latent sampling greatly improves reward (and our qualitative as- sessments) compared to the prior. These results, along with the analysis in $5.2, demonstrate the effectiveness of our pro- posed method fine-tune high-dimensional image priors.
Setup. Many protein generative models have been proposed to tackle the problem of designing novel yet realistic protein structures (Watson et al., 2023; Bose et al., 2024). They learn to produce proteins of N residues by sampling ro- tations and translations applied to each residue backbone (the space SE)3) As our prior pa(x), we use the recent FoldFlow 2 model, which is a Riemannian CNF (on the man- ifold SE(3) embedded in RxNN) ) trained with minibatch OT coupling (Huguet et al., 2024).
Protein residues fold into patterns called secondary struc-
8

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Table Results for protein structure experiments. We report the average log-reward, as well as the pairwise TM-Score (as mea- sure of sample diversity), averaged over 64 samples. Standard deviation over seeds is reported.
l
(a) CFM Prior
Method
B[lograv rav (x)] (1)
Pairwise TM-Score (L)
Prior
-1.325 0.014
0,4480 0.0044
RW MCMC
-0.640 0.073
0,4181 * 0.0057
(b) Outsourced Diffusion Posterior
Outsourced Diff.
0.422 0.225
0.4407 0.0706
Figure 7. Protein samples with pertinent secondary structures high-
lighted: a- a-helices (blue), B-shects (red), and coils (green).
tures, which include a-helices, B-sheets, and coils. Many protein generative models have issues producing proteins with diverse secondary structures, typically under-sampling proteins with ß-sheets. A natural problem is to produce samples which are both probable under the prior model and exhibit this structural diversity. This can be framed as sampling from Po(X)ras (x)x))ivxx,))))))))))))))>)))( (x), where the constraint function raiv (x) assigns high-values to proteins with diverse sec- ondary structures -arpartiiular, the presence of B-sheets).
strated the effectiveness of this method in a variety of set- tings, there are many questions for future work. One natural direction is to extend the method to discrete problems, where the noise space is discrete or the transformation involves discretization. Another is to adapt outsourced diffusion sampling to general probabilistic programs, where the gen- erative model includes both stochasticity and nontrivial con- trol flow, and where current inference methods use MCMC sampling in outsourced noise spaces (Dash et al., 2023).
Let p LP Pp. Pe] be vector representing the proportion of residues that are in a-helices, B-sheets or coils. The particular constraint function that we use, adapted from
Acknowledgments
e-12w"p
Huguet et al. (2024), is raiv H.2-Hlp]' where w [1. [.5) is a weight vector, and H[p) is the entropy of p.
The authors thank Moksh Jain for comments and sugges- tions in the course of this project, Kirill Tamogashev for comments on a draft of the manuscript, and Prudencio Tossou and Dominique Beaini for helpful discussions and suggestions about the protein conformation experiments.
This problem poses two challenges. First, since the model is a Riemannian CNF, the flow ODE cannot be converted into a diffusion SDE, so that adjoint matching and diffusion fine- tuning techniques are not applicable. Second, the constraint function is not differentiable with respect to the generative model's output, ruling out methods such as HMC. We there- fore compare our method gradient-free MCMC method in the noise space We fix the protein length 64 residues, and evaluate the model achieving the highest diversity score during training. See Appendix B.5 for details.
L.S. acknowledges support from Recursion. The research of M.S. was funded by National Science Centre, Poland, 2022/45/N/ST6/03374 Y.B. acknowledges the support from CIFAR and the CIFAR AI Chair program as well as NSERC funding for the Herzberg Canada Gold medal. G.B. ac- knowledges funding support from Samsung Al Lab and CIFAR.
We report the average log raiv rdiv(x), (x), as well as a di- versity metric (the pairwise TM-Score) in Table 6. The latter calculates the similarity between pairs of proteins, averaged across a set of generated samples (Zhang & Skolnick, 2004). Uncurated protein samples are shown in Fig. 7. We find that our proposed method samples diverse protein structures rich in B-sheets more frequently than the baselines, while maintaining a T'M-Score comparable to the prior.
Results.
The research was enabled in part by computational re- sources provided by the Digital Research Alliance of Canada (https: /aal://///aaanannnnnmmaaaaa ca), Mila (https: //mile quebec), and NVIDIA.
Impact Statement
Improving reward finetuning strategies for generative mod- els enhance their usefulness but also carries risks asso- ciated with misuse, reward misspecification and unintended generalization. Careful evaluation and responsible deploy- ment are essential as these methods scale.
6. Conclusion
We have proposed outsourced diffusion samplers for effi- cient posterior inference in the noise spaces of generative models. These samplers take advantage of the expressive- ness of diffusion models and the flexibility of off-policy training algorithms for blx target distributions and can be applied to any model that can be written as deter- ministic transformation of noise. While we have demon-
References
Akhound-Sadegh, T., Rector-Brooks, J., Bose, A. J., Mittal, S., Lemos, P., Liu, C.-H., Sendera, M., Ravanbakhsh, S., Gidel, G., Bengio, Y., Malkin, N., and Tong, A. Iterated

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
denoising energy matching for sampling from Boltzmann densities. International Conference Machine Learning (ICML), 2024.
Cardoso, G., Janati el idrissi, Y., Le Corff, S., and Moulines, E. Monte Carlo guided denoising diffusion models for Bayesian linear inverse problems. International Confer- ence on Learning Representations (ICLR), 2024.
Albergo, M. S. and Vanden-Fijnden, E. Building normal- izing flows with stochastic interpolants. International Conference on Learning Representations (ICLR), 2023.
Che, T., Zhang, R., Sohl-Dickstein, J., Larochelle, H., Paull, L., Cao, Y., and Bengio, Y. Your GAN is secretly an energy-based model and you should use discriminator driven latent sampling. Neural Information Processing Systems (NeurIPS), 2020.
Austin, T. Exchangeable random measures. Annales de I'Institul Henri Poincaré, Probabilités et Statistiques, (3):842 861, 2015.
Chen, J., Richter, L., Berner, J., Blessing, D., Neumann, G., and Anandkumar, A. Sequential controlled Langevin diffusions. International Conference on Learning Repre- sentations (ICLR), 2025.
Bengio, E., Jain, M., Korablyov, M., Precup, D., and Ben- gio, Y. Flow network based generative models for non- iterative diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. Neural Information Processing Systems (NIPS), 31. 2018.
Bengio, Y., Lahlou, S., Deleu, T., Hu, E. J., Tiwari, M., and Bengio, E. GelowNet foundations. Journal of Machine Learning Research, 24(210):1-55 2022..
Chung, H., Kim, J., Mccann, M. T., Klasky, M. L., and
Berner, J., Richter, L., and Ullrich, K. An optimal con- trol perspective on diffusion- based generative model- ing. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https: attps:///ppereet cpenrevie.....e net/ www.uorrrddeoor 3w
Ye, J. C. Diffusion posterior sampling for general noisy inverse problems. International Conference on Learning Representations (ICLR), 2023.
Cobb, A. D. and Jalaian, Scaling Hamiltonian Monte
Carlo inference for Bayesian neural networks with sym- metric splitting. Uncertainty in Artificial Intelligence (UAI), 2021.
Berner, J., Richter, L., Sendera, M., Rector- cector-Brooks,) J., and Malkin, N. From discrete-time policies to continuous- time diffusion samplers: Asymptotic equivalences and faster training. arXiv preprint arXiv:2501.061448 2025.
Coeurdoux, F,, Dobigeon, N., and Chainais, P. Normalizing flow sampling with Langevin dynamics in the latent space. Machine Learning. 113(1 1):11):::
Bose, J., Akhound-Sadeghh T., Huguet, G., FATRAS, K., Sector-Brookss J.. Liu, C.-H., Nica, A. C., Korablyov, M., Bronstein, M. M., and Tong, A. SE(3)-stochastic flow matching for protein backbone generation. International Conference on Learning Representations (ICLR), 2024.
2024.
Dash, S., Kaddar, Y., Paquet, H., and Staton, S. Affine monads and lazy structures for Bayesian programming. Proc. ACM Program. Lang., 7(POPL), January 2023.
Del Moral, P., Doucet, A., and Jasra, A. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, (8(3):41-4333 2006.
Brooks, S., Gelman, A., Jones, G., and Meng, X.-L. Hand- book of Markov Chain Monte Carlo. Chapman and Hall/CRC, May 2011. ISBN 9780429138508. doi: 10.1201/b10905. 1201/b10900
URL http:/ dx doi http://ddxxdd....oggo... 000
Dhariwal, P. and Nichol, A. Q. Diffusion models beat GANs on image synthesis. Neural Information Processing Sys- tems (NeurIPS), 2021.
Bugallo, M. F. Elvira, V., Martino, L., Luengo, D., Miguez, J., and Djuric, P. M. Adaptive importance sampling: The past, the present, and the future. IEEE Signal Processing Magazine, 34(4):60-7) 79, 2017.
Domingo-Enrich, C., Drozdzal, M., Karrer, B., and Chen, R. T. Q. Adjoint matching: Fine-tuning flow and diffusion generative models with memoryless stochastic optimal control. arXiv preprint 2rXiv:2409.08861 2024.
Cabezas, A., Sharrock, L., and Nemeth, C. Markovian flow matching: Accelerating MCMC with continuous normalizing flows. Neural Information Processing Sys- tems (NeurIPS), 2024.
Dou, Z. and Song. Y. Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. International Conference on Learning Representations (ICLR), 2024.
Cannella, C., Soltani, M., and Tarokh, V. Projected latent Markov chain Monte Carlo: Conditional sampling of normalizing flows. International Conference on Learning Representations (ICLR), 2021.
Doucet, A., Johansen, A. M., et al. A tutorial on particle filtering and smoothing: Fifteen years later. Handbook of nonlinear filtering, 12(656-704)33, 2009.
10

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Doucet, A., Grathwohl, W. S., Matthews, A. G., and Strath- mann, H. Score-based diffusion meets annealed impor- tance sampling. Neural Information Processing Systems (NeurIPS), 2022.
Hutchinson, M. F. A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Com- munications in Statistics Simulation and Computation, 18:1059-1076, 1989.
Esser, P. Kulal, S., Blattmann, A., Entezari, R., Müller, J., Saini, H., Levi, Y., Lorenz, D., Sauer, A., Boesel, F., Podell, D., Dockhorn, T., English, Z.. Lacey, K., Good- win, A., Marek, Y., and Rombach, R. Scaling rectified flow transformers for high-resolution image synthesis. International Conference on Machine Learning (ICML), 2024.
Kadkhodaie, Z. and Simoncelli, E. P. Solving linear inverse problems using the prior implicit in a denoiser. Neural Information Processing Systems (NeurIPS), 2021.
Karras, T., Aittala, M., Laine, S., Härkönen, E., Hellsten, J., Lehtinen, J., and Aila, T. Alias-free generative adver- sarial networks. Neural Information Processing Systems (NeurIPS), 2021a.
Fan, Y., Watkins, O., Du, Y., Liu, H., Ryu, M., Boutilier, C., Abbeel, P., Ghavamzadeh, M., Lee, K., and Lee, K. Reinforcement learning for fine-tuning text-to-image dif- fusion models. Neural Information Processing Systems (NeurIPS), 2023.
Karras, T., Laine, S., and Aila, T. A Style- Generator Architecture for Generative Adversarial Networks IEEE Transactions on Pattern Analysis & Machine Intelligence, 43(12)42217-423 December 2021b. ISSN 1939-3539. doi: 10. 10.10099TAAM...22022290000
URL
https: //tpss://doi: www.ccompteeeeocccccccc..... org/ 10. N9.1109/TPAAN 2020 .00.0022700000 2970919.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. Neural Information Process- ing Systems (NIPS), 2014.
Kawar, B., Vaksman, G., and Elad, M. SNIPS: Solving noisy inverse problems stochastically. Neural Information Processing Systems (NeurIPS), 2021.
Graikos, A., Malkin, N., Jojic, N., and Samaras, D. Diffu- sion models as plug-and-pla priors. Neural Information Processing Systems (NeurIPS), 2022.
Kidger, P., Foster, J., Li, X. C., and Lyons, T. Efficient and accurate gradients for neural SDEs. Neural Information Processing Systems (NeurIPS), 2021.
Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., and Duvenaud, D. FFJORD: Free-form continuous dy- namics for scalable reversible generative models. Interna- tional Conference on Learning Representations (ICLR), 2019.
Kim, M., Choi, S.. Yun, T., Bengio, E., Feng, L., Rector Brooks, J., Ahn, S., Park, J., Malkin, N., and Bengio, Adaptive teachers for amortized samplers. International Conference on Learning Representations (ICLR), 2025.
Ho, J. and Salimans, T. Classifier-free diffusion guidance. arXiv preprint 2tt.:::2227.25999 2022.
Kingma, D. P. and Welling, M. Auto-encoding variational Bayes. International Conference on Learning Represen- tations (ICLR), 2014.
Ho, J., Jain, A., and Abbeel, P. Denoising diffusion proba- bilistic models. Neural Information Processing Systems (NeurIPS), 2020.
Korbak, T., Perez, E., and Buckley, C. L. RL with KL
penalties is better viewed as bayesian inference. Findings Empirical Methods in Natural Language Procedding (EMNLP), 2022.
Hou, Z., Lang, N., and Zhou, X. WL-GAN: Learning to sample in generative latent space. Information Sciences, 700:121834, 2025.
Krizhevsky, A. Learning multiple layers of features from tiny images. Master's thesis, University of Tront, 2009.
Huang, X., Dong, H., Hao, Y., Ma, Y.-A., and Zhang, T. Reverse diffusion Monte Carlo. International Conference on Learning Representations (ICLR), 2024.
Lahlou, S., Deleu, T., Lemos, P., Zhang, D., Volokhova, A., Hernández-Garcia. A., Ezzine, N., Bengio, Y., and Malkin, N. theory of continuous generative flow networks. International Conference on Machine Learning (ICML), 2023.
Huguet, G., Vuckovic, J., Fatras, K., Sibbodeau-Laufer E., Lemos, P., Islam, R., Liu, C.-H., Rector-Brooks J., Sahound-Sadegll T., Bronstein, M., et al. Sequence- augmented SE(3)-flow matching for conditional protein backbone generation. Neural Information Processing Systems (NeurIPS), 2024.
Lawson, D., Tucker, G., Naesseth, C. A., Maddison, C. J., Adams, R. P., and Teh, Y. W. Twisted variational se- quential Monte Carlo. 2018, URL https: wwww.w/w.oommmmmoooooooo.// cs. utoronto. caaonto.cc//aae comcaadddda wttpdddds/p/////eeeeeeeeeeeeeeee ttabbb//tssc....o pdf.

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Lawson, D., Raventós, A., Warrington, A., and Linderman, S. SIXO: Smoothing inference with twisted objectives. Neural Information Processing Systems (NeurIPS), 2022.
Phillips, D. and Cipcigan, MetaGFN: Exploring dis- tant modes with adapted metadynamics for continuous GFlowNets. arXiv preprint arXiv:2408. 2024.
Li, Li, D., Xiong, C., and Hoi, S. BLIP: Bootstrapping language image pre-training for unified vision- vision-language understanding and generation. International Conference on Machine Learning (ICML), 2022.
Pooladian, A.-A. Ben-Hamu, H., Domingo-Earich, C., Amos, B., Lipman, Y., and Chen, R. T. Multisample flow matching: Straightening flows with minibatch couplings. International Conference on Learning Representations (ICLR), 2023.
Li, X., Wong, T.-K. I.., Chen, R. T., and Duvenaud, D. Scalable gradients for stochastic differential equations. Artificial Intelligence and Statistics (AISTATS), 2020.
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. Learning transferable visual models from natu- ral language supervision. International Conference on Machine Learning (ICML), 2021.
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. Interna- tional Conference on Learning Representations (ICLR), 2023.
Rezende, D. and Mohamed, S. Variational inference with normalizing flows. International Conference on Machine Learning (ICML), 2015.
Liu, X., Gong, C., and Liu, Q. Flow straight and fast: Learning to generate and transfer data with rectified flow. International Conference on Learning Representations (ICLR), 2023.
Rezende, D. J., Mohamed, S., and Wierstra, D. Stochas- tic backpropagation and approximate inference in deep generative models. International Conference on Machine Learning (ICML). 2014.
Malkin, N., Jain, M., Bengio, E.,, Sun, C., and Bengio, Y. Trajectory balance: Improved credit assignment in gflownets. Neurul Information Processing Systems (NeurIPS), 2022.
Richter, L. and Berner, J. Improved sampling via learned diffusions. International Conference Learning Repre- sentations (ICLR), 2024.
Malkin, N., Lahlou, S., Deleu, T., Ji, X., Hu, E., Everett, K., Zhang, D., and Bengio, Y. GFlowNets and varia- tional inference. International Conference on Learning Representations (ICLR), 2023.
Richter, L., Boustati, A., Nüsken, N., Ruiz, F. J. R., and Ömer Deniz Akyildiz. VarGrad: A low-variance gradient estimator for variational inference. Neural Information Processing Systems (NeurIPS), 2020.
Mardani, M., Song, J., Kautz, J., and Vahdat, A. A varia- tional perspective on solving inverse problems with dif- fusion models. International Conference on Learning Representations (ICLR), 2024.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
Midgley, L. I., Stimper, V., Simm, G. N., Schölkopf, B., and Hernández-Looene Lobato, J. M. Flow annealed importance sam- pling bootstrap. International Conference on Learning Representations (ICLR), 2023.
Ronneberger, O., Fischer, P., and Brox, T. U-Net: Con- volutional networks for biomedical image segmentation. Medical Image Computing and Computer-Assisted Inter- vention (MICCAI), 2015.
Miyato, T., Kataoka, T., Koyama, M., and Yoshida, Y. Spec- tral normalization for generative adversarial networks. International Conference on Learning Representations (ICLR), 2018
Sendera, M., Kim, M., Mittal, S., Lemos, P., Scimeca, L., Rector-Brooks, J., Adam, A., Bengio, Y., and Malkin, N. Improved off-policy training of diffusion samplers. Neural Information Processing Systems (NeurIPS), 2024.
Nguyen, A., Clune, J., Bengio, Y., Dosovitskiy, A., and Yosinski, J. Plug & play generative networks: Conditional iterative generation of images in latent space. Computer Vision and Puttern Recognition (CVPR), 2017.
Simonyan, K. and Zisserman, A. Very deep convolutional networks for large-scale image recognition. International Conference on Learning Representations (ICLR), 2015.
Nüsken, N. and Richter, L. Solving high- -ige-sinensioaal Mamillon-lacoii-ellnaa 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.
Sohl-Dickstein, J., Weiss, E. A., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. International Conference on Machine Learning (ICML., 2015.
12

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. International Conference Learning Representations (ICLR), 2021a.
Watson, J. L.,, Juergens, D., Bennett, N. R., Trippe, B. L.,
Yim, J., Eisenach, H. E.., Ahern, W., Borst, A. J., Ragotte, R. J., Milles, L. F., Wicky, B. I. M., Hanikel, N., Pellock, S. J., Courbet, A., Sheffler, W., Wang, J., Venkatesh, P., Sappington, I., Torres, S. V., Lauko, A., Bortoli, V. D., Mathicu, E., Ovchinnikov, S., Barzilay, R., Jaakkola, T., DiMaio, F,, Baek, M., and Baker, D. De novo design of protein structure and function with rfdiffusion. Nature, 620: 100:1089-1100 2023.
Song. J., Zhang, Q.. Yin, H., Mardani, M., Liu, M.-Y., Kautz, J., Chen, Y., and Vahdat, A. Loss- diffusion models for plug-and-play controllable generation. International Conference Maching Learning (ICML), 2023.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. International Conference Learning Representations (ICLR), 20211,.
Xu, J., Liu, X., Wu, Y., Tong, Y., Li, Q. Ding, M., Tang, J., and Dong, Y. ImageReward: learning and evaluating human preferences for text-to-image generation. Neural Information Processing Systems (NeurIPS), 2023.
Song, Y., Shen, L.., Xing. L.., and Ermon, S. Solving inverse problems in medical imaging with score-based generative models. International Conference on Learning Represen- tations (ICLR), 2022.
Zhang, Q. and Chen, Y. Path integral sampler: a stochastic control approach for sampling. International Conference on Learning Representations (ICLR), 2022.
Tong, A., Fatras, K., Malkin, N., Huguet, G., Zhang. Y., Rector-Brooks, J., Wolf, G., and Bengio, Y. Improv- ing and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research (TMLR), 2024.
Zhang, Y. and Skolnick, J. Scoring function for automated assessment of protein structure template quality. Proteins: Structure, 57, 2004.
Tzen, B. and Raginsky, M. Neural stochastic differential equations: Deep latent Gaussian models in the diffusion limit. arXiv preprint arXiv: www..9955..9... 2019.
Uehara, M., Zhao, Y., Black, K., Hajiramezanali, E., Scalia, G., Diamant, N. L., Tseng, A. M., Biancalani, T., and Levine, S. Fine-tuning of continuous-time diffusion models eetropy-rreularrized control. arXiv preprint 151444 2024.
Vahdat, A. and Kautz, NVAE: A deep hierarchical varia- tional autoencoder. Neural Information Processing Sys- tems (NeurIPS), 2020.
Vargas, F., Ovsianas, A., Fernandes, D., Girolami, M., Lawrence, N. D., and Nüsken, N. Bayesian learning via neural Schrödinger-Föllmee flows. Statisties and Computing, 33(1):3, 2022.
Vargas, F., Grathwohl, W., and Doucet, A. Denoising dif- fusion samplers. International Conference on Learning Representations (ICLR), 2023.
Vargas, F., Padhy, S., Blessing, D., and Nüsken, N. Trans- port meets variational inference: Controlled Monte Carlo diffusions. International Conference on Learning Repre- sentations (ICLR), 2024.
Venkatraman, S., Jain, M., Scimeca, I.., Kim, M., Sendera, M., Hasan, M., Rowe, L., Mittal, S., Lemos, P., Bengio, E., et al. Amortizing intractable inference in diffusion models for vision, language, and control. Neural Infor- mation Processing Systems (NeurIPS), 2024.
13

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
A. Theory and Method Details
A.I. Outsourced Sampling for Diffusion Priors
We describe posterior sampling, outsourced posterior inference and diffusion sampling when the prior generative model is itself diffusion model.
Monte Carlo methods. The 'guidance" term in diffusion posteriors the difference between the scores of the noised prior and posterior distributions can be estimated by Monte Carlo integration (Song et al., 2023; Cardoso et al., 2024) or using approximations specialized for constraints arising from linear inverse problems (Kawar et al., 2021; Kadkhodaie & Simoncelli, 2021; Song et al., 2022: Chung et al., 2023). Posterior estimation can also be achieved through stochastic optimization (Graikos et al., 2022; Mardani et al., 2024)
Methods related to sequential Monte Carlo, which treat a modification to the denoising transition kernel as proposal, have also been proposed (Doucet al., 2022: Dou & Song. 2024; Chen et al., 2025).
Amortized methods. Because generation proceeds in long sequence of sampling steps, and the modes of the posterior are not known priori, these methods use reinforcement learning techniques to discover regions of high posterior density, Asymptotically unbiased methods include ELEGANT (Uchara et al., 2024) and relative trajectory balance (Venkatraman et al., 2024).
Outsourcing noise in diffusion models. Generation of data x Xp is modeled as a Markov process Xy ..... where Xy N(0. Idam) and the transition from X, X,-1 is conditionally spherical-Gaussee trick, the trajectory of latent variables can be expressed as function of the initial sample xy and the T standard Gaussian noises injected at each step of sampling, just as in a VAE. Thus diffusion model is generative model with outsourced noise in (T+1) (see (6) below).
Xo. Via the reparametrization
Generalizing this setting, typical latent diffusion model chains diffusion model in latent space Rdlser WT W1 - Wo- with a Gaussian decoder Wo - x of the same form as a VAE decoder. Combining the two, the data x is a deterministic transformation of the concatenation of the initial latent variable wy, the T standard Gaussian noises injected at each step of sampling, and the noise in the final decoder. Thus a latent diffusion model is generative model with outsourced noise in Rua (T+1)+dan
(Note the similarity to the outsourced interpretation of HVAEs above: diffusion model indeed be understood as deep hierarchical VAE. However, a diffusion model is also neural stochastic differential equation (Tzen & Raginsky, 2019; Song et al., 2021b) integrated in discrete time. In this view, in the continuous-time limit, the outsourced noise is a sample of Brownian motion, and indeed an Itō integral is a deterministic transformation of a Brownian noise random variable.)
Outsourced autoregressive sampling under diffusion priors recovers relative trajectory balance. Venkatraman et al. (2024) studied the problem of fine-tuning a diffusion model po seen as a transition policy po (Xax, X(--)ar) to yield diffusion model p post that samples the product of the distribution p*(x) defined by the prior model with a constraint r(x)), where the prior and posterior diffusion model share the noising process and standard Gaussian noise distribution p(xo). The relative trajectory balance (RTB) objective was proposed; for trajectory T.
pr (xias I x(i-1)80)) 2
Z
LRTB(T: Φ) log
log 1-1 Px ((i-1)Ar)
(5)
r(x1)
where Z is learned scalar (note the resemblance to (4)).
Let & be the standard Gaussian noise injected in sampling xia conditionally on x(i-1)d so that the prior model can be rewritten as a deterministic function of the noises fo. .f. via
Xo Eo. Xiar = X(I-1)Ar (o(X (i ()(1) VTANVVE
(6)
where μe outputs the drift of the generative SDE. Similarly, let f be the drift the posterior generative SDE. For a
trajectory r sampled using a sequence of noises &o. .i.....E7 under the prior model, the RTB loss (5) can then be rewritten
14

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
in terms of the noises:
Z
(KG-1)A (i 1)Ar)Ar, Tis Ar) ((xarr X(G-1)MHH post
LRIB [RIB(T;))
log f(x))
log
N(xiN - X(i-I) HO (KC-))A (i 1)Ar)Ar, T Ar)
i=l
4
N(EG IG- (1- VA///I/ d
log
log
r(xi)
N(E:O, Idam
i=1
(Eo O, Idam 11 NEE pain (a (Xii-))ar (i 1) VAr/ Idam r(xi) I N(E: Idam
log
where (xin((.1) = post (x.t) - (x,
Consider now an amortized sampler of the outsourced noises 51.51..... that generates the variables autoregressively (Eo from a standard Gaussian, then subsequently each & from a Gaussian with unit variance and mean conditioned on previously sampled noises). The transition policy density of this sampler can be written
pr(E: Ei-1) NE: Moutsourced Ei-1). Idam).
Under a model parametrization in which the policy takes as input the intermediate state X(i-1)Ar- computed as a function of the noises Eo... Ei-1 using the prior model,
Moutsuuced (60.....))..)) = uai (X(i-1) (i 1)Ar) I)
the numerator in (7) is precisely ZePr(T), where is the sampling trajectory of the autoregressive sampler generating Eo. The denominator is rxi) multiplied with the prior density of the noise, which is the target density for the sampler (the unnormalized density of the outsourced posterior).
Thus we see that (7) exactly recovers the trajectory balance objective (Malkin et al., 2022) for an autoregressive sampler of outsourced noise.
A.2. Proof of Prop. 3.1
Proposition A.1. Suppose that (Z, Ez) and (X, Ex) are measurable spaces and f Z + X is measurable. If j is a or-finite measure on Z and is a o-finite measure on X with f.p, then f. ((dir f) -4 =
In particular, y is probability measure and h: X - R2o f-p-integrable, then d Thdlfau (ho f) .)..p. probability measure Z, and fel [hd(sw) (h f-4) is probability measure X.
Proof of Prop. 3.1. Let h = dv Zí.H we must show that f-((h f) 11) = h- f.H. Let E E Σx and D = (E) E Σz. By definitions and properties of pushforward measures,
fa(( A) m)E) ((h S..o).p (D)
o filD du
(h1g) of du
[ hig df-u
= (h (a))
as required.
For the second part of the proposition, yf h- f.p,, then df. h f-u-almost everywhere, and the result follows easily from the first part by linearity of the pushforward.
15

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
B. Experiment Details
Code for our experiments is available at httpa//wwwo.o........ooooccoctta Sampling.
B.1. Diffusion model
For all experiments, we use a convolutional UNet architecture (Ronneberger et al., 2015) for the diffusion model. This architecture is naturally well-suited for latent spaces structured as image feature maps. However, for other types of latent representations, we found that simply reshaping them into H W x C feature maps and treating them as images yielded surprisingly effective results.
We use the variance preserving SDE (Song et al., 202 as backward policy PB for the trajectory balance loss in (4), with a discretisation of 25 steps. For all of the experiments, we use off off-polliy training with a replay buffer similar to Sendera et al. (2024). Specifically, for every training update we randomly sample either from the replay buffer, or generate new on-policy samples which are added to the replay buffer according to the buffer probability a - whose value specific to experiments is provided in the following sections.
B.2. Caad-Conniiional Sampling
We set the classifier inverse temperature B = for all our baselines, otherwise the soft logits of the classifier results in a fuzzy posterior. We provide some additional information regarding implementation of the baselines below:
Outsoureed Diffusion. We learn to sample the 128 dimensional generator noise with the outsourced sampler. To keep architectures consistent across experiments, we used the convolutional UNet to sample this noise by reshaping them to 2 x 88 x 8 feature maps. For sampling from the CNF prior, we use 45 Euler steps to discretise the ODE. The outsourced sampler operates on the noise space which the same dimensionality as the output image - 3 x 32 x 32.
Adjoint Matching. To fine-tune flow matching models with a classifier, we use adjoint matching following the approach described in Cemingo-Fnnriee et al. (2024). Due to the unavailability of the open source code at this time, we made our own implementation which we found works extremely well. First, we convert the ODE inference of conditional flow matching into a memoryless SDE, ensuring that both have the same marginal distribution as outlined in the referenced paper. Next, we apply the adjoint matching method for the stochastic optimal controls of the memoryless SDE. This process requires gradient information of the log-reward (i.e., the classifier's 's oogllike,hoood,, which we compute using Py PyTorch's autograd functionality.
All typerparameters remain identical to those of the outsourced diffusion sampler, including the learning rate, except for the number of training iterations and the temperature annealing schedule. Training is significantly slower because it requires computing the log-reward gradients and simulating the full trajectory for adjoint matching. Therefore, we set a total of 2,500 iterations with 1, 1,000 linear temperature annealing steps. This configuration results in similar reduced wall time compared to the outsourced diffusion sampler while achieving stable convergence and high performance on 10..R-10. Figure 8 shows the qualitative results of applying adjoint matching to the CIFAR-10 10 task.
Note that for Stable Diffusion 3, the actual implementation of the inference pipeline- which involves a diffusson transformer with various combinations of language embedding fusion for multimodality-is is not fully open-sourced. Implementing a memoryless SDE on top of this pipeline is non- toivial and requires careful tuning and integration with the internal processes of Stable Diffusion's inference mechanisms.
Temperature Annealing. We anneal the inverse temperature from B = to the final B = 4. We tune the schedule linearly over the first 2, 000 steps of training for outsourced diffusion, and 1, 000 steps for adjoint matching.
HMC baselines. We use hamiltorch (Cobb & Jalaian, 2021) to implement HMC. We use step size of 10-2, with 5 leapfrog integration steps. We use a burn-in chain of length 100 before starting to collecting samples, spaced out by 10 samples. Chains are run for 1, 000 samples, after which we reset the seed to help diversity. The runtime of HMC is quite slow with the CNF prior due to gradient computation. A chain of 1, 000 samples takes close to hour A 100 GPUs. We only keep 90 samples from this chain to preserve diversity, but these are still correlated. It would take close to 10 hours to generate 1, 000 samples for FID computation, but using extra resources we can run parallel chains. For comparison, outsourced diffusion takes close to 5 hours for full training, after which sample generation is extremely cheap due to amortization. Despite the
16

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
(a) Adj. (Dog)
(b) Adj. (Car)
Figure CIFAR-10 samples generated Adjoint Matching (Adj.) the Dog and Car classes.
need train a model, outsourced diffusion actually more memory-efficient than even single HMC chain, as it eliminates the need compute gradients through the ODE integrator. We visualize samples from a latent HMC chain in Fig. 40.
Our results for each class listed in Table 7 and Table 8. We additionally evaluate the ELBO for I-CFM posteriors in Table 9.
Table 7. SN-GAN results CIFAR-10
Class
Prior
HMC
Outsourced Diff.
Reward
FID
Reward
FID
Reward
FID
Airplane
-5.82
100.28
-2.53
82.56
-3.89
79.55
Car
-5.36
125.16
-3.91
94.81
-4.22
85.88
Bird
-5.40
71.30
-4.86
52.14
-3.91
59.76
Cat
-5.40
68.30
-2.21
67.55
-3.97
55.05
Deer
-5.40
70.10
-3.67
51.10
-3.38
43,04
Dog
-5.32
95.62
-4.00
82.32
-4.23
68.52
Frog
-4.98
93.00
-4.20
68.12
-3.91
74.70
Horse
-5.23
97.53
-3.21
64.60
-3.70
52.87
Ship
-5.19
116.89
-2.11
94.45
-3.76
79.73
Truck
-5.59
133.23
-2.25
97.88
-3.50
82.07
AVG
-5.37
97.14
-3.26
75.33
-3.847
68.117
B.3. Conditional High-RRsolutton Face Generation
For all prompts we use fixed inverse temperature / 100, which we found to be suitable reward scale.
Outsourced Diffusion. StyleGAN3 uses 512 dimensional generator noise. Similar our SN-GAN experiment for CIFAR, reshape this into 2 x 16 x 16 feature map be passed the UNet model. The NVAE prior for FFHQ is very deep latent variable model, having 36 total latent groups starting from 8 x 8 scale all the way up 128 x 128. The joint dimensionality of this latent space is extremely large, making joint posterior inference very challenging. Luckily most variable features of interest are captured in the first 4 latent groups of size 20 x 8 x 8 each. We stack the noise groups to create 80 x x feature maps jointly diffused by the UNet. We use 25 steps for diffusion sampling.
HMC baseline. We use 5. 10-3 as the step size for StyleGAN3, and 1 as the step size for NVAE, We use 5 leapfrog integration steps. We only require 100 samples for evaluation in this task (unlike 1000 needed for FID in CIFAR), and we afford to only collect 2 samples for every 1000 length chain. We collect the samples at 500 and 1000. This takes around 3 hours for StyleGAN3 and 5 hours for NVAE without parallelism. We only sample the first 4 latent groups with NVAE prior.
We find that unlike the CIFAR-10 priors, HMC struggles to obtain high reward with these priors, but performs better with StyleGAN3 than NVAE. We attribute the particularly poor performance of NVAE to the high dimensionality and high energy barriers. Interestingly, outsourced diffusion performs significantly better for sampling these posteriors. We suspect that
17

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Table 8 I-CFM results CIFAR-10
Class
Prior
HMC
Adj. Matching
Outsourced Diff.
Reward
FID
Reward
FID
Reward
FID
Reward
FID
Airplane
-5.81
73.09
-2.45
62.24
-3.30
27.64
-3.36
47.61
Car
-6.22
92.07
-2.12
24.85
-2.93
17.51
-3.02
19.12
Bird
-5.94
73.48
-2.93
60.75
-3.16
23.09
-3.11
41.04
Cat
-5.53
70.33
-3.60
54.32
-3.33
20.11
-3.57
43.66
Deer
-5.59
72.19
-2.72
49,46
-3.08
15.52
-3.78
31.87
Dog
-5.83
89.38
-3.89
41.13
-3.21
23.11
-3.57
34.33
Frog
-6.66
93.08
-3.36
55.64
-3,27
20.86
-3.98
34.74
Horse
-6.13
82.88
-2.10
46.22
-2.79
16.59
-2.95
31.85
Ship
-5.85
102.33
-2.28
37.80
-3.09
18.00
-2.92
32.13
Truck
-5,22
99.10
-2.55
34.48
-2.55
12.34
-3.20
26.46
AVG
-5.88
84.79
-2.80
46.69
-3.09
19.45
-3.35
34.28
Table Additional comparisons against the training-free DPS (Chung et al., 2023) and diffusion finc-tuning with RTB (Venkatraman
al., 2024). For methods where it is applicable, we also report the estimated ELBO.
Model
Sampler
Bllog ply x)] (1)
FID (1)
ELBO (1)
Prior
-5.88
84.79
224.04
DPS
-2.2
84.96
I-CFM
RTB
-4.20
90.77
-147.69
Latent HMC
-2.80
46.69
Adj. Matching
-3.09
19.45
-17.23
Outsourced Diff.
-3.35
34.28
-20.36
primarily the nice mode mixing properties of the diffusion annealing path that facilitate this. However, an additional factor might be the benefits of amortization, which an interesting direction for future work to investigate. We present the results for different prompts in Table 10, and showcase the first 10 samples for each prompt, generated from fixed random seed, in Appendix D.
Table 10. StyleGAN3 results on FFHQ
Prompt
Prior
HMC
Outsourced Diff.
Reward
Diversity
Reward
Diversity
Reward
Diversity
An old man
-1.02
0.35
-0.60
0.3
1.62
0.32
young Asian girl with glasses
-1.97
0.36
-0.81
0.32
1.13
0.21
bald man with black beard
-1.89
0.38
-0.71
0.32
1.02
0.27
brown-haired child
-1.2
0.36
-0.35
0.28
1.14
0.23
AVG
-1.52
0.36
-0.62
0.31
1.23
0.26
B.4. Text-To-Imare RLHF
We intentionally choose prompts that pose challenge for Stable Diffusion 3 while still receiving reliable feedback from ImageReward. Since SD3 is generally a more powerful model than ImageReward, this approach is not applicable to most prompts. However, in real-world scenario, we anticipate the use of better preference model trained with human feedback, which would offer more reliable guidance for improving the generative model. For training, we use a fixed inverse temperature B == 30.
Outsourced Diffusion. We learn sample from the noise space SD3, which is latent CNE This means we can
18

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
Table 11. StyleGAN3 results on FFHQ
Prompt
Prior
HMC
Outsourced Diff.
Reward
Diversity
Reward
Diversity
Reward
Diversity
An old man.
-1.74
0.30
-1.24
0.32
1.38
0.31
A young Asian girl with glasses.
-2112
0.30
-0.95
0.31
0.70
0.25
A bald man with a black beard.
-2.16
0.29
-1.40
0.29
1.32
0.26
A brown-haired child.
-1776
0.30
-1.21
0.31
0.53
0.23
AVG
-1.94
0.30
-1.20
0.30
0.98
0.26
sample a reduced dimensionality from the full image (of size 3 x 512 x 512), however the latent space still fairly high dimensional with shape 16 64 64.
Classifier-Free Guidance. Since SD3 is trained both unconditional and conditional model, CFG to approximately sample from lowered temperature conditional distribution:
(x, y) = (1 w) - (v)(x,)
(8)
Increasing the guidance scale w generally guides the model to be more prompt accurate at the cost of diversity and increased too much, image fidelity. We find tuning the CFG weight slightly improves score with ImageReward. Increasing resulted in degraded performance ("A cat riding a Ilama."), so we report score with the default guidance scale 5.0 (same as prior).
We present the results for different prompts in Table 12, and showcase the first 10 samples for each prompt, generated from fixed random seed, in Appendix E. We set w = (6.0, 5.0, 5.5) for the prompts ordered as in the table.
Additionally, we train Outsourced Diffusion posteriors using Stable Diffusion 1.5 with the same prompts as in (Venkatraman ct al., 2024). We compare our results against those reported in their paper in Table 13.
Table 12, SD3 RLHF results for each prompt
Prompt
Prior
CFG
Outsourced Diff.
Reward
Diversity
Reward
Diversity
Reward
Diversity
cat and a dog.
1.24
27.25
riding
quiet village disrupted strike.
0.65
0.24
0.71
0.20
0.94
23.20
human with horse face and human with wolf face.
1.22
0.20
19.32
1.26
0.22
1.36
0.18
26.10
AVG
1.27
Table 13. SD 1.5 RLHF results, averaged across prompts
Sampler
E [lgrr(x.)][[)
CLIP diversity (f)
Prior
-0.17
0.18
DDPO
1.37
0.09
DPOK
1.23
0.13
RTB
1.4
0.11
Outsoureed Diff.
1.26
0.14
B.5. Protein Structure Prediction
The specific form of the constraint function (x) (which can also be thought of as reward) is adapted from (Huguet et al., 2024). It a monotonic function of the entropy over secondary structure proportions (restricted to a-helices, B-sheets and coils). The non-linear transformation increases the discrepancy between the rewards of samples with and without sheets.. The expression is adjusted so that is positive, to allow us to take the logarithm.
19

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
The e-1.2 multiplier ensures the log-constraint function roughly centered around 0. We found this improve numerical stability in the diffusion training.
Outsourced Diffusion. In the outsourced diffusion experiments, an inverse temperature of ß 400.0 was used to allow for improvement in this sparse reward setting. The diffusion sampler used 20 sampling steps.
The SE(3)* N elements representing protein structure were parameterized as R vectors, with the first coordinates being a quaternion representation for rotation matrix and the 3 other numbers representing the translation vector. For the UNet model, the protein coordinates were shaped into 7x8x8 vectors. Additionally, the diffusion sampler was first pre-trained (for 200 epochs) using the denoising score matching objective (Ho et al., 2020), with samples from (SO(3)) x N(O, 10 13) (where U is the Haar measure, i.e., the unique invariant probability measure SO(3)*), ), the initial distribution used for training the CNF prior model. All inference parameters for the flow model were based on the default configuration from the FoldFlow paper, including 50 steps for integrating the flow ODE (Huguet et al., 2024).
For TB training, we used a replay buffer, where 1/4 of the samples are drawn proportional to their reward, and the rest are sampled uniformly. The buffer used each iteration with probability of 0.2 (ar 0.2). The reason for the modification from the standard uniform buffer used other experiments to make the best use of high reward samples, which are rare especially early on in training. learning rate of 10-5, and batch size of 16 was used. Gradient 12 norms were clipped to 0.05.
Due to policy collapse and training instability, models were saved every 100 training iterations, and the model with highest reward was selected for evaluation. The diffusion model was trained for A100 GPU hours.
Random Walk MCMC baseline. Gaussian proposal, p(z' z) (z'' 0.013 (i.e., step-size of 0.01). The quaternion dimensions normalized to have unit norm (projected back to )O(3)
Note this proposal is symmetric. The MCMC chain run for 1000 iterations, with 32 chains in parallel. Metrics were evaluated on samples from the iterations 900 and 1000 (to reduce sample correlation). This was run for 8 A100 GPU hours.
20

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
CIFAR-10
Stable Diffusion
StyleGAN3
Original
Distilled
Figure 9. Posterior samples via the original outsourced diffusion sampler and the distilled one-step sampler in three experiments. We use the prompt "A cat and dog." for the stable diffusion experiments, and "An old man" for the SN-GAN experiments.
C. Additional Experiments
C.1. Distillation and One-Step Outsourced Samplers
Table 14. FID scores of standard and distilled outsourced diffusion samplers for a CFM
Training outsourced samplers of the kind described in this paper may have the downside of an increase in the number of sampling steps necessary to produce posterior samples of interest. We perform further experiments to show how we can mitigate these side effects via distillation. We distill outsourced models for the CMF posterior on the 'car' class in CIFAR-10 10, the SN-GAN posterior for "A dog and cat", and stable diffusion posterior for "A green car". We train original diffusion samplers.
prior on CIFAR- 10.
Sampler
FID
distilled one-step samplers equivalent in architecture and size We use the trained outsourced samplers as teacher models and employ a simple training regime whereby we sample 2 NO,, Ia), use an ODE sampling scheme to sample 21, and then learn a one-step mapping from zo to zι with the student model. We use Mean Square Error (MSE) loss and a simple variance agnostic regularizer to train the student model and encourage diversity.
SDE
38.65
ODE
36.71
Distilled
33.85
We report in Fig. 9 original and posterior samples for all experiments. We observe high-quality distilled samples, undistin- guishable from the original outsourced diffusion sampler. Furthemore, report in Table 14 the FID score computed from samples from the original and distilled sampler, with respect to the original CIFAR-10 10 samples of the class of interest. We observe comparable, if not improved, FID for our distilled model. In line with the results in the main paper, we posit that the ability to fit high fidelity one-step samplers is due to the simplicity of the properties of the target distribution in latent space, often smoother and lower dimensional than the distribution in data space, leading to casy-to-lcarn transforms from Zo to Z
C.2. Efficiency Analysis
The results in Appendix C.1 show that the diffusion sampler can be distilled into single-step generator, enabling inference of the fine-tuned model without incurring additional sampling time. As a result, the inference efficiency remains comparable to that of direct finetuning methods.
On the CIFAR-I0 dataset, our method achieves approximately twice the training speed of adjoint matching when using an NVIDIA 100 GPU. This performance advantage is expected to grow for digees-aimeniional outputs. It is important to note that the requirement for gradient computation not only raises training memory costs but also limits the flexibility of the method. For tasks where the reward gradient is unavailable, such as many protein or molecule tasks, Adjoint Matching is not applicable. In contrast, the outsourced diffusion sampler can be efficiently applied to arbitrary black-box tasks where the reward (or prior) is not differentiable.
21

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
D. FFHQ Samples
D.I. Prior.
Figure 10. NVAE prior
Figure 11. StyleGAN3 prior
D.2. An old man.
Figure 12. NVAE HMC
Figure 13. NVAE Outsourced Diffusion
Figure 14. StyleGAN3 HMC
Figure 15. StyleGAN3 Outsourced Diffusion
22

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
D.3. An asian girl with glasses.
Figure 16. NVAE HMC
Figure 17. NVAE Outsourced Diffusion
Figure 18. StyleGAN3 HMC
Figure 19. StyleGAN3 Outsourced Diffusion
23

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
D.4. Bald man with black beard.
Figure 20. NVAE HMC
Figure 21. NVAE Outsoureed Diffusion
Figure 22, StyleGAN3 HMC
Figure 23. StyleGAN3 Outsourced Diffusion
24

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
D.S. Brown haired child.
Figure 24. NVAE HMC
Figure 25. NVAE Outsourced Diffusion
Figure 26. StyleGAN3 HMC
Figure 27. StyleGAN3 Outsourced Diffusion
25

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
E. Stable Diffusion 3 Samples
E.1. cat and dog.
Figure 28. Prior
Figure 29. Classifer- Free Guidance
Figure 30. Outsourced Diffusion
E.2. A cat riding llama.
Figure 31, Prior
Figure 32, Classifer-Free Guidance
Figure 33. Outsourced Diffussio
26

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
E.3. A quiet village is disrupted by meteor strike.
Figure 34, Prior
Figure 35. Classifer-Free Guidance
Figure 36. Outsourced Diffusion
E.4. A human with a horse face and a human with a wolf face.
Figure 37. Prior
Figure 38. Classifer-Free Guidance
Figure 39, Outsourced Diffusion
27

---

Outsourced Diffusion Sampling: Efficient Posterior Inference in Latent Spaces of Generative Models
F. CIFAR-10 HMC Chain
(a) 1=100
(b) 151
(c) t=200
(d) t=250
(e) t=300
(1) 1=350
(g) 40400
(h) 1=450
ii) (=500
ii 1=550
(k) 1=00
()) t=650
(m) t=700
(n) 1=750
(o) t=800
(p) t=80
(q) 1000
(r) 1950
(s) t= 10100
Figure 40. Samples from HMC chain of length 1000 after 100 steps of burn-in, for CIFAR class 'Airplane" with CNF prior. We find that in the latent space, MCMC smoothly traverses through different modes. The samples (d,o.g.rs) are distinctly identifiable airplanes and (c) is partially identifiable as an airplane. The samples approach the correct class after long mixing time.
28


Reviews:
[{"review_id": "BhVBfKNI9w", "reviewer": "Reviewer_qmBQ", "summary": "## Summary\n* This paper proposed a more general approach towards posterior inference for generative model with Gaussian prior (e.g. GAN, Flow, Diffusion models). More specifically, it proposed to learn a non-Gaussian prior (noise space z).\n* The paper is evaluated on various models and various scale from toy size to Stable Diffusion 3. Solid empirical results support", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the helpful review. We address the weaknesses and questions raised by you in our response below.\n\n### __Comparison with model-specific approaches__\n\nWe thank the reviewer for pointing out the additional model-specific baselines with which to compare out general approach. \n- We note that for GANs, we compare to a latent space exploration baseline similar to PULSE, Hamiltonian Monte Carlo (HMC). \n- We have added the DPS baseline to the CIFAR-10 experiments.\n- We have added RTB as a diffusion-specific baseline for CIFAR-10 and Stable Diffusion-1.5 experiments.\n\n__The additional results are included in Tables 1 and 2 below, and also in the [linked pdf](https://anonymous.4open.science/r/results-FED2/results.pdf).__\n\nTable 1: CIFAR-10 posterior sampling results, averaged over 10 classes\n| Model | Sampler              | $\\mathbb{E}[\\log p(\\textbf{y} \\mid \\textbf{x})]$ $(\\uparrow)$ | FID $(\\downarrow)$ | ELBO $(\\uparrow)$ |\n|-------|----------------------|------------------------------------------------------------|--------------------|--------------------------------|\n| I-CFM | Prior                | -5.88                                                      | 84.79              |      -24.04                    |\n|       | DPS                  | -2.22                                                      | 84.96              |        -                       |\n|       | RTB                  | -4.20                                                      | 90.77              |    -147.69                     |\n|       | Latent HMC           | -2.80                                                      | 46.69              |        -                       |\n|       | Adj. Matching        | -3.09                                                      | 19.45              |      -17.23                    |\n|       | **Outsourced Diff.** | -3.35                                                      | 34.28              |      -20.36                    |\n\nIn the CIFAR-10 results we note that DPS obtains a higher reward than Outsourced Diff., with a worse FID score (and ELBO -- see the response to Reviewer nZ5P). This points to possible reward hacking. Additionally, the added RTB baseline for this experiment proved much more unstable during training than our method, reflected in its worse FID (and ELBO) scores, while improving the reward relative to the prior. Some RTB runs underwent policy collapse, requiring early stopping for the reported results.\n\nTable 2: SD 1.5 fine-tuning results. DDPO, DPOK and RTB results taken from [Venkatraman et al.]\n\n| Sampler              | $\\mathbb{E}[\\log r(\\textbf{x}, \\textbf{y})] (\\uparrow)$ | CLIP diversity $(\\uparrow)$ |\n|----------------------|--------------------------------------------------------|-----------------------------|\n| Prior                |                -0.17                                   |          0.18               |\n| DDPO                 |                 1.37                                   |          0.09               |\n| DPOK                 |                 1.23                                   |          0.13               |\n| RTB                  |                 1.4                                    |          0.11               |\n| **Outsourced Diff.** |                 1.26                                   |          0.14               |\n\n\nFor the SD1.5 experiment, we see that our method achieves a slightly lower reward than RTB, while obtaining a higher diversity over all baselines (other than the prior).\n\n### __Relationship to non-amortized methods (golden noise and test-time scaling)__\n\nWe agree that incorporating a discussion comparing our proposed amortized sampler with these inference-time optimization methods would strengthen the paper. Notably, the latent space HMC baseline used in our CIFAR-10 and FFHQ experiments shares conceptual similarities with these inference-time noise optimization approaches. While the methods in prior work do not asymptotically sample from the true Bayesian posterior over the noise space, they effectively bias the sampling process toward high-reward regions, often producing qualitatively similar results. In contrast, our method frames prior-regularized reward fine-tuning as posterior inference, whereas these inference-time techniques are more aligned with direct reward maximization. We will integrate this discussion into the appendix of the final version.\n\n### __Organization of results__\n\nWe thank the reviewer for their feedback on the organization of results. We will merge some of the results that share similar metrics and baselines (such as CIFAR and FFHQ), and reorganize the tables to make the presentation more clear for the final submission.\n\n**Thank you again for your review, and please do not hesitate to let us know if there is anything more we can clarify in the second response phase.**", "claims_and_evidence": "Yes. The authors claim a general approach to achieve conditional sample by learning the prior p(z) of various generative model. They verify their approach in different kinds of models such as flow and diffusion.", "methods_and_evaluation_criteria": "The evaluation spans from toy CIFAR dataset to practical scale image generation. The metrics can be organized in a better way, but common metrics such as FID and imagereward are shown.", "theoretical_claims": "There is no theoretical claim in this paper.", "experimental_designs_or_analyses": "## Experimental designs\n* The flow matching model in SD3 is a special type of flow matching model that start from a Gaussian noise. This type of flow is not much different from diffusion models \n* For some models such as GAN and diffusion, there exists tailored method to sample from posterior, such as DPS [diffusion posterior sampling for general noisy inverse problems] and GAN inversion [PULSE: Self-Supervised Photo Upsampling via Latent Space Exploration of Generative Models]. As the authors propose a general approach for different types of generative model, it would strengthen the paper if the authors show the performance of their approach against model specific approaches.", "relation_to_broader_scientific_literature": "* The authors propose a general approach to learn amortized conditional generative model by learning the prior. This paper is also related to a recent hot topic: golden noise [Not All Noises Are Created Equally:Diffusion Noise Selection and Optimization] and test time scaling in text to image diffusion model [Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps]. Despite this paper is amortized, it is better to discuss the relationships to those non-amortized works.", "essential_references_not_discussed": "See Relation To Broader Scientific Literature.", "other_strengths_and_weaknesses": "## Weakness\n* The authors are encouraged to organize the results in a better way. For now too many results are shown from image to portein generation. Different datasets, methods and metrics are scattered across the main text and appendix.", "other_comments_or_suggestions": "-", "questions_for_authors": "-", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oFo0OKMsqL", "reviewer": "Reviewer_nz5P", "summary": "This work targets the problem of generating samples from the posterior $p(x | y) \\propto p(x) r(x, y)$, where the prior $p(x)$ is a (pre-trained) generative model and $r(x, y)$ is a reward function. The authors argue that (most) generative models can be formulated as the application of a pushforward $f(z) = x$ to a simple latent distribution $p(z)$ (e.g. Gaussian). In this case, the initial problem is equivalent to generating from the posterior $p(z | y) \\propto p(z) r(f(z), y)$, which the authors argue is easier to tackle. The authors train a *diffusion sampler* $p^\\phi(z | y)$ using the trajectory balance (TB) objective (Malkin et al., 2022), which only requires access to the unnormalized density $p(z) r(f(z), y)$. The authors demonstrate on 3 text-to-image benchmarks and 1 protein design benchmark that their method is competitive against a few alternatives.\n\n### update after rebuttal\n\nI thank the authors for their rebuttal and taking my comments into account. I appreciate the additional discussion, and hope that they will be able to add it to the manuscript. My concerns regarding the evaluation of the posterior (calibration) somewhat remains, but I agree that proper evaluation remains a challenge in high dimension, especially in the lack of data pairs $(x, y)$. In this light, I will be raising my score to 4 (accept).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and constructive feedback, as well as pointing out that the paper is very well written. \n\nYou raise some important concerns, which have helped us improve the paper and which we hope to address:\n\n### __Posterior evaluation__\n\nWe agree that more comprehensive posterior evaluations are important for a thorough assessment. \n\nFirst, we emphasize that the FID scores reported in our original CIFAR-10 experiments capture both reward and coverage/diversity, and serves as a proxy metric for closeness to the true posterior. FID remains a widely accepted and robust metric for evaluating the sample quality of image generative models.\n\nFor the CIFAR-10 experiments (comparing our amortized sampler to ground truth posterior samples), we tried posterior evaluation with TARP [Lemos et al., 2023] and PQMass [Lemos et al., 2024], an unconditional two-sample variant of TARP, but were not able to obtain results that seemed meaningful. For example, running TARP frequently showed prior samples as perfectly calibrated and unbiased relative the conditional dataset, which is obviously false -- but the discrepancy is captured correctly in the FID scores. (Published code from those papers was used; to ensure the results do not contain an error, we checked that our code could reproduce values in line with those reported on CIFAR-10 in [Lemos et al., 2024].)\n\nFor other experiments, ground truth samples from the posterior are not available, complicating evaluation (for instance, FID cannot be computed), which is the reason we resort to other metrics. For prompt-conditioned sampling we instead report the average CLIP feature distance, following the approach used by [Venkatraman et al.]. In our Stable Diffusion 1.5 experiments, the CLIP feature diversity scores exhibit a consistent and interpretable trend: the prior achieves the highest diversity, while DDPO, which performs greedy reward maximization without KL regularization, yields the lowest diversity. These trends support the reliability of our evaluation metrics in capturing meaningful aspects of sample diversity.\n\nAt the reviewer's suggestion, we also report the ELBO (i.e., the lower bound on the log-normalizing constant $\\log Z$, where the ELBO-to-log-likelihood gap is the reverse KL divergence). This is a widely used metric in the diffusion samplers literature, including [Sendera et al.], and is maximized by a perfect sampler of the target density. We note however, that this metric cannot be used to evaluate inference-time baselines or MCMC.\n\n__Addition of the ELBO metrics for CIFAR-10 and SD3 are included in Tables 1 and 3 repectively, within the [linked PDF](https://anonymous.4open.science/r/results-FED2/)__.\n\n### __Additional baselines__\n\nWe include experimental results for the RTB baseline proposed by [Venkatraman et al.] for flow models trained with independent coupling, which can subsequently be fine-tuned as diffusion models. This evaluation is done for RTB using the I-CFM prior on CIFAR-10; we also include new results comparing our method against RTB, DPOK [Fan et al.] and DDPO [Black et al.] with Stable Diffusion 1.5 as a prior with the same setup used by [Venkatraman et al.]. \n__These are included in Tables 1 and 2 of the response to Reviewer qmBQ, and also in the [linked pdf](https://anonymous.4open.science/r/results-FED2/results.pdf)__.\n\nWe note that for CIFAR-10, the RTB baseline is significantly more unstable to train: requiring LoRA fine-tuning [Hu et al.], as well as lower learning rates to avoid a quick policy collapse. This aligns with design choices in the RTB paper [Venkatraman et al.]. \nEven with this, the training remains somewhat unstable for the flow model architecture, being very sensitive to training hyperparameters. Within the training time allocated to ODS, the model achieves a modest improvement in reward; however, the generated samples often fail to consistently belong to the target class and frequently exhibit visual artifacts. This degradation in quality is further reflected in poorer posterior metrics (FID, ELBO). We argue the training instability is caused by the flow prior, which uses a different noise schedule as well as much fewer inference steps compared to experiments in [Venkatraman et al.]. This instability could lead to \"reward-hacking\" behavior where reward improves at the cost of sample quality. \n\nBy contrast, ODS improves largely within the first 2-3 GPU hours, and only marginally afterwards. The general difficulty of training RTB with the flow prior reinforces the claim that the latent posterior $p(z|y)$ is a simpler target than the image posterior $p(x|y)$.\n\nFinally, **thank you for your helpful suggestions on improving the notation and writing clarity**, as well as for pointing out the typos. We will implement this feedback for our final submission.\n\n**Thank you again for your review, and please do not hesitate to let us know if there is anything more we can clarify in the second response phase.**", "claims_and_evidence": "/", "methods_and_evaluation_criteria": "Yes, the method is sound, the evaluation tasks are relevant, and several baselines are considered. However,\n\n1. In the case where $p(x)$ is a flow-matching/diffusion model, an important baseline should be to fine-tune the generative model using the trajectory balance (TB) objective directly, which would validate the authors' claim that $p(z | y)$ is an easier target than $p(x | y)$ (line 74.5). This approach is taken by Venkatraman et al. [1].\n\n2. The assessment of the quality of the inferred posterior distributions is not sound. $\\mathbb{E}[\\log r(f(z), y)]$ is maximized by a collapsed distribution around $z^* = \\arg\\max_z r(f(z), y)$. The diversity is maximized by a maximal entropy distribution. Although these quantities are relevant for evaluating the inferred posterior $p^\\phi(z | y)$, they are not sufficient. In fact, I suspect the method presented in this paper to be subject to mode collapses, and the current evaluation does not rule out this hypothesis. For example, you can observe in Figure 4.e that most dogs are white and facing the camera, which is not the case in the CIFAR-10 dataset. Similarly, in Figure 6, there is a collapse of the image composition (tabby cat on the llama, meteor falling bottom left, horse on the left).\n\n    There is an extensive literature on the evaluation of posterior distributions, notably emerging from the SBI [2] community. For example, in scientific applications, the calibration of the posterior distribution is extremely important [3-5]. Another sensible metric could be the (reverse) KL divergence between $p^\\phi(z | y)$ and $p(z | y)$. This KL can be computed up-to the normalizing constant $Z(y) = \\int p(z) r(f(z), y) dz$, and therefore can be used to compare different approximations of $p(z | y)$.\n\n    In my opinion, a proper evaluation of the inferred posteriors is a *sine qua non* for a paper claiming to perform posterior inference. I would be happy to raise my score if the authors address this concern.\n\n[1] Amortizing intractable inference in diffusion models for vision, language, and control (Venkatraman et al., 2024)\n\n[2] The frontier of simulation-based inference (Cranmer et al., 2020)\n\n[3] Validating Bayesian Inference Algorithms with Simulation-Based Calibration (Talts et al., 2018)\n\n[4] A Trust Crisis In Simulation-Based Inference? Your Posterior Approximations Can Be Unfaithful (Delaunoy et al., 2021)\n\n[5] Sampling-Based Accuracy Testing of Posterior Estimators for General Inference (L3mos et al., 2023)", "theoretical_claims": "Yes, I quickly checked the proof of Prop. 3.1 and it seemed valid. The work is mainly applicative, as it consists in applying already existing methods (mainly TB) to a widespread problem.", "experimental_designs_or_analyses": "See Methods and Evaluation criteria.", "relation_to_broader_scientific_literature": "The literature review in this work is complete and extensive, without being overwhelming. Kudos to the authors.", "essential_references_not_discussed": "There is an extensive literature on the evaluation of posterior distributions [3-5] (this is a very limited selection), which this work does not consider/discuss.", "other_strengths_and_weaknesses": "1. The paper is very well written. The figures are clear. The experiments are relevant.\n2. The evaluation is lacking a proper assessment of the posteriors' quality.\n3. In my opinion, the methodological contribution of this work is modest with respect to previous works. The impact, however, could be significant. Unfortunately, this cannot be assessed without a proper evaluation.", "other_comments_or_suggestions": "1. In Table 1, diffusion models should have $d_{data}$ as noise dimension, if they are to be considered as deterministic functions. The decoder of latent diffusion models is typically considered deterministic (even by the authors), leading to a noise dimension $d_{latent}$. This is consistent with Table 2.\n1. In Figure 1, \"two Gaussians **and** an observation\" \n1. In Figure 1, the bottom right plot should read $p(x | y)$\n1. Line 240, \"**the one** defined by the target\"\n1. Line 241, \"noising kernel\" is undefined\n1. Line 246, TB \"loss\" was introduced as \"objective\" in Eq. (4)\n1. In Eq. (2), I would avoid the notation $\\sigma_t$ as it has another meaning in the diffusion model literature\n1. Line 312.5 $z \\in \\mathbb{R}^{128}$", "questions_for_authors": "none", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lqrnX10pXb", "reviewer": "Reviewer_Ge1S", "summary": "This paper addresses the posterior inference problem using diffusion sampling. By comparing their approach with existing MCMC methods and amortized inference methods, the authors demonstrate that their proposed outsourced diffusion sampling method, optimized through the trajectory balance objective, is both efficient and effective. They evaluate its performance across three application domains: conditional image generation, text-to-image generation, and protein structure generation.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your good questions and constructive feedback. We have responded to your points below.\n\n### __Algorithm details__\n\nWe agree that the addition of an algorithm block would help readers understand how to implement the core training loop, and added this to __Algorithm 1 in page 2 of the [linked pdf](https://anonymous.4open.science/r/results-FED2/results.pdf).__\n\n### __FID scores not state-of-the-art__\n\nWhile FID scores are indeed worse than adjoint matching on CIFAR-10, we highlight that ODS is a general approach applicable to any generative model, while adjoint matching is specific to flow models (and not all flow models, see end of §3.1) and requires differentiating the reward. \n\nThe FID scores reported in [Venkatraman et al.] are an order of magnitude lower than those presented in our paper. On examining the underlying codebases we discovered that the discrepancy arises from the FID packages used. Specifically, [Venkatraman et al.] uses [this codebase](https://github.com/marcojira/fld), while we use [this one](https://github.com/GaParmar/clean-fid). The distinction lies in the handling of image resolution: CIFAR images are smaller than ImageNet images, and only the latter implementation (which we use) accounts for this difference appropriately.\n\n### __Efficiency__\nOur efficiency claim refers to lower training time of ODS relative to finetuning methods, while having much cheaper inference than MCMC. There is a small (25 step) inference overhead relative to Adj. matching (for reducing this, see Appendix C.1).\n\n### __Which part of the proposed method is new?__\n\nTrajectory balance has previously been employed for training diffusion samplers from unnormalized densities (notably in [Lahlou et al.] and [Sendera et al.]), and a variant of it, relative TB, is a diffusion-specific fine-tuning method in [Venkatraman et al.]. Our novel contribution is the general capability to train a sampler (using TB) of the Bayesian posterior in noise space (§3.2) for any generative model, providing a general-purpose framework for amortized posterior inference. The exact way in which ODS generalizes relative TB is detailed in Appendix A.1.\n\n### __Training-free baselines and classifier guidance__\n\nUnbiased **classifier guidance** requires a *time-dependent* likelihood gradient: $\\nabla \\log p(y \\mid x_t) = \\nabla \\log \\mathbb{E}_{p(x_0 \\mid x_t)}[p(y \\mid x_0)]$\n\nThis necessitates training a classifier on noised inputs, as using one trained on clean data renders the gradient intractable. Such methods therefore make stronger assumptions than our setting, which assumes a black-box classifier on clean data.\n\nTraining-free methods like DPS are biased approximations to classifier guidance but work well in practice. DPS uses the approximation $\\nabla \\log p(y \\mid x_t) \\approx \\nabla \\log p(y \\mid \\mathbb{E}[x_0 \\mid x_t])$, and requires the same assumptions as our proposed algorithm making it an appropriate baseline.\n\nWe’ve added results for DPS (Diffusion Posterior Sampling, [Chung et al.]), a strong training-free baseline for approximate posterior sampling with diffusion and flow models using the I-CFM prior on CIFAR-10. __Results are in Table 1 of the response to Reviewer qmBQ and also [here](https://anonymous.4open.science/r/results-FED2/results.pdf).__ DPS achieves high reward but poor FID, reflecting highly biased sampling.\n\n### __Additional experiments with SD1.5__\n\nWe conducted additional experiments with Stable Diffusion 1.5 to compare against DDPO [Black et al.], DPOK [Fan et al.], and RTB [Venkatraman et al.], as requested. We also included RTB with the I-CFM prior as a baseline for the CIFAR-10 class-conditional sampling task. __Results are shown in Tables 2 and 1 of our response to reviewer qmBQ and are available [here](https://anonymous.4open.science/r/results-FED2/results.pdf).__ On SD1.5, ODS achieves the a strong balance between reward and sample diversity.\n\n### __Stability of training__\n\nODS uses TB, an off-policy RL objective that does **not** require backpropagating through the full sampling chain. Instead, it only needs local gradients of log-likelihoods for each sampling step (see [Venkatraman et al.], Appendix H.1)\n\nSuch off-policy optimization has three benefits, according to prior work (see, e.g., [Nüsken & Richter], [Sendera et al.]):\n- The reward function need not be differentiable, since TB avoids using log-reward gradients; \n- TB avoids instability and mode collapse seen in methods that backprop through the full SDE trajectory (e.g., PIS [Zhang & Chen]);\n- Memory usage is minimal thanks to a gradient accumulation strategy (included in our code), avoiding storage of full computation graph\n \nIn our experiments, we found training to be very stable for all tasks except for the protein diversity task, where policy collapse sometimes occurs.\n\n**Thank you again for your review, and please do not hesitate to let us know if there is anything more we can clarify in the second response phase.**", "claims_and_evidence": "This paper makes the following three claims. First, ODS is agnostic to the form of priors, which is demonstrated through the experiments. Second, ODS is an effective posterior inference method that can be applied across multiple domains. While ODS shows improved results in certain application domains given the baseline authors provided, some important baselines have not been considered, which I will discuss in detail in the Experiments & Designs section. \n\nIn addition, certain results are a bit far from the state-of-art methods, i.e. the FID of ODS with I-CFM in the conditional image generation task. The experiments cover three application domains, which I consider sufficient.\n\nThe third claim concerns efficiency, where the authors claim that ODS is more efficient than amortized inference methods and MCMC methods.  I suggest the authors make this more explicit, is it the training time efficiency or sampling time efficiency? \nTraining time ODS is better than Adjoint matching (Domingo-Enrich et al., 2024). But the NFEs at sampling time are not mentioned between ODS and Adjoint matching in the conditional CIFAR-10 data.", "methods_and_evaluation_criteria": "The proposed methods is a proper fit to the problem and also to the applications that the author is trying to target. \nIn the Conditional High-Resolution Face Generation and Text-to-Image RLHF experiments, the generated image quality metric like FID is not present.", "theoretical_claims": "N/A", "experimental_designs_or_analyses": "Class-Conditional Sampling with Cifar-10, Text-to-Image RLHF is carefully checked. \n\nFor the first one, a simple training-free baseline such as classifier-guidance is missing from Table 3. Additionally, I'm unsure if mentioning distillation on ODS is relevant.\n\nFor the second one, is the choice of SD3 instead of SD1.5 as the prior for not comparing with Fan et al. (2023) and Venkatraman et al. (2024)? Both methods are closely related to the proposed method. Although the paper mentions that CNF is not a diffusion model, recent work [1] demonstrates that flow-based models can be viewed as diffusion models with different noise schedules and parameterizations. Therefore, comparisons should also include methods closely related to this work in the diffusion land.\n\n\n[1] Gao, R., Hoogeboom, E., Heek, J., De Bortoli, V., Murphy, K. P., & Salimans, T. (2024). Diffusion meets flow matching: Two sides of the same coin.", "relation_to_broader_scientific_literature": "Posterior inference is an important problem to the broader scientific field. Traditional MCMC methods may yield better results, but they suffer from slow sampling times. Amortized inference offers an alternative approach, however, it requires obtaining massive numbers of samples from the posterior for training. Therefore, developing an efficient solution to this problem is important, which is what this paper targets.", "essential_references_not_discussed": "The discussion on the relevant works is sufficient.", "other_strengths_and_weaknesses": "- Strength: \n\nThe proposed method offers a promising solution to the posterior inference problem at scale and is universal to the prior distribution. In addition, the proposed method does not require the reward or constraint function to be differentiable. The paper also includes applications from different domains. \n\n- Weakness\n\nMy main concerns in this work are the writing and the experiment. The details of the proposed method in Section 4.2 are missing, and there is no algorithm block describing the overall method. This leaves an audience with no prior knowledge about the trajectory balance objective line of work struggling to differentiate between what already exists in the literature and what is newly proposed in this work. So I suggest instead of introducing what VAE/GAN is, set up the audience with more details like “off-policy” divergences approach and TB objective. \n\nSome of the weaknesses are mentioned in the Experimental Designs or Analyses. Some of the results, such as the FID scores on conditional CIFAR-10, are far from those of state-of-the-art methods, which makes hard to judge whether ODS can obtain high quality posterior samples.", "other_comments_or_suggestions": "Figure 1 comment: “The constraint function – a mixture of two Gaussians centered **an an** observation”", "questions_for_authors": "In section 4.2, which part of the proposed method exists in the literature and which part is new in this work? \n\nDid authors consider training-free posterior sampling methods to be a proper baseline? \n\nWhy is classifier guidance not appear in the condition Cifar10 experiment? \n\nCould the author compare ODS with Fan et al. (2023) and Venkatraman et al. (2024) on SD1.5 to demonstrate the effectiveness?\n\nIs the training stable given the ODS need to back-propogate through the entire sampling chain according to Eq(4)?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "n5QQeyTyNX", "reviewer": "Reviewer_RBkw", "summary": "This paper presents a new method for posterior sampling from a wide variety of generative models (GANs, Flow, VAE) that can be expressed as a deterministic transformation x ~ f(z) for z sampled from a simpler distribution like Gaussian noise. The key idea is to train an auxiliary diffusion model to produce an initial point z' which when passed through f(z') results in samples x' from the the desired posterior in the data space p(x | y).", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed review and constructive feedback, as well as the positive comments on the exposition and evaluations.\n\nTo address the questions:\n\n### __Efficiency of training a model for each task__\n\nIt is true that the outsourced diffusion sampler must be trained separately for each likelihood constraint $p(y\\mid x)$. However, for the diffusion sampler, the inference-time cost per sample is much lower than for MCMC (i.e., the sampling cost is amortized). In fact, the total cost for generating samples for evaluation was much lower using our method (**including** training time) compared to MCMC, for instance:\n- The CIFAR-10 evaluation required generating 1000 samples, which took 10 hours with Hamiltonian Monte Carlo (using multiple chains in parallel), while the outsourced diffusion sampler took 5 hours for both training and sampling (of which sampling was a neglible fraction). \n- A similar pattern held for MCMC in the protein experiments: the entire experiment took 8 hours for MCMC and 4 total hours for our method. \n\nAdditional details on the timing comparisons between MCMC baselines and our method are available in Appendix B.\n\nRegarding the **need to retrain for new constraints**, we point out that the regular functionality of generative models remains intact when using our method to sample the latent space. Namely, we can still apply *approximate* training-free methods for guidance to solve intractable sampling problems (e.g., for inpainting or linear inverse problems), with the only difference being that the noise components are obtained from our diffusion sampler. Our method offers a computational trade-off whereby additional compute can be spent to sample more accurately according to a given posterior distribution, compared to the generative model with standard noise.\n\n### __Size of example images: Can we use larger models?__\n\nThe largest latent space we explore in our experiments is with Stable Diffusion 3, featuring a latent resolution of $16 \\times 64 \\times 64$, corresponding to an image resolution of $3 \\times 512 \\times 512$. We consider this to be a reasonably high resolution for practical applications: most image generative models, save those with excessive hardware and sampling time requirements, use similar latent dimensions. (Although diffusion models trained in pixel space with higher dimension exist, they typically require far more sampling steps than ones that work in a latent space.) Generally, larger latent spaces do increase training times, but our experiments demonstrate that our method scales effectively even with high-fidelity models like SD3. \n\nWe note that for many generative models like GANs, the noise space is of much lower dimension than the data space (e.g., 512 vs. $3\\times256\\times256$ for the FFHQ StyleGAN3 we consider), which is further motivation to perform outsourced sampling.\n\nFinally, thank you for your suggestions on reorganizing the figures. We will incorporate this feedback into our final submission.\n\n**Thank you again for your review, and please do not hesitate to let us know if there is anything more we can clarify in the second response phase.**", "claims_and_evidence": "There are 3 main claims: 1) It's applicable to a wide range of prior models 2) It's effective under a variety of domains 3) It's more efficient than MCMC type methods.\n\nI think the claims are well supported. Results are shown on a variety of models and a lot of different tasks. One thing that wasn't clear to me is whether we need to train a different model for each posterior task. If so, then the efficiency is questionable, because MCMC posterior sampling methods can generally substitute in different likelihood constraints p(y | x) at inference without retraining. Since the R constraint is in the loss function (eq 4) it seems like the outsourced diffusion model is trained for each specific task.", "methods_and_evaluation_criteria": "In general the method is well explained and makes sense. I like some of the core evaluations on the ImageReward examples to show how this proposed method improves over the baselines. I wonder why the examples are mostly 256 x 256. Is it harder to train in the latent space of larger models? Given that the first couple of figures are on toy examples (the swiss roll pictures) and the experiments are on small images, it seems like there may be more work needed to make this method truly applicable.", "theoretical_claims": "I checked most of the math closely (excluding some of the proofs in the appendix). Everything seemed to follow logically, and I appreciate the authors interpreting many of the detailed equations in more intuitive/layman terms.", "experimental_designs_or_analyses": "The broad experiments were great. There were many different generative models, constraints, and datasets used. As mentioned earlier, it would be nice to see some higher resolution examples if it's possible.", "relation_to_broader_scientific_literature": "I think this paper ties together posterior sampling across many different types of generative models nicely. It also positions itself in contrast to MCMC posterior sampling methods, which have generally been extremely slow and required many steps to converge. I like that it is positioned well regarding all the posterior methods that the community uses.  It is also a novel idea from my understanding.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The paper is very clear. I found the method section easy to follow. I worry about the cost of training an auxiliary model for posterior sampling. A major advantage of using generative models for posterior sampling, is the ability to handle many tasks (inpainting, super-resolution etc) with a single pretrained generative model. If the outsourced diffusion sampler has to be retrained for different R functions then this method becomes much more constrained. In fact I think this factor should be discussed more, the (retraining of the outsourced sampler for different constraints)", "other_comments_or_suggestions": "I think there should be one nice teaser figure at the top. To start the paper by only seeing the toy swiss roll examples makes me think the practicality is already limited. The casual reader would greatly benefit from seeing a clear, compelling example up front, like maybe the cat on the llama example.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a method for sampling from data posteriors within the framework of large pretrained generative models, via diffusion samplers. The method is model-agnostic and can be applied to a variety of generative models that are deterministic push forward functions on Gaussian distribution. Experimental results have been shown on variety of tasks and models demonstrating the method's wide utility. 

3/4 reviewers have recommended acceptance post rebuttal and the major concern of the reject review is on organization and writing. 

I think this is a good paper worth accepting, albeit I urge the Authors to address the comments on better writing and evaluating the method against some state of the art models if possible.
Decision: Accept (poster)