End-to-end Learning of Gaussian Mixture Priors for Diffusion Sampler
Authors: Denis Blessing
Xiaogang Jia
Gerhard Neumann
Venue: ICLR 2025 Poster (2025)
Keywords: Variational Inference
Sampling
Diffusion Models
Mixture Models
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We focus on improving existing diffusion-based sampling methods via end-to-end learning of Gaussian mixture priors

Abstract: Diffusion models optimized via variational inference (VI) have emerged as a promising tool for generating samples from unnormalized target densities. These models create samples by simulating a stochastic differential equation, starting from a simple, tractable prior, typically a Gaussian distribution. However, when the support of this prior differs greatly from that of the target distribution, diffusion models often struggle to explore effectively or suffer from large discretization errors. Moreover, learning the prior distribution can lead to mode-collapse, exacerbated by the mode-seeking nature of reverse Kullback-Leibler divergence commonly used in VI.
To address these challenges, we propose end-to-end learnable Gaussian mixture priors (GMPs). GMPs offer improved control over exploration, adaptability to target support, and increased expressiveness to counteract mode collapse. We further leverage the structure of mixture models by proposing a strategy to iteratively refine the model through the addition of mixture components during training. Our experimental results demonstrate significant performance improvements across a diverse range of real-world and synthetic benchmark problems when using GMPs without requiring additional target evaluations.

Published as a conference paper at ICLR 2025
END-TO-END LEARNING OF GAUSSIAN MIXTURE PRIORS FOR DIFFUSION SAMPLER
Denis Blessing*. Xiaogang Jia', Gerhard Neumann
Autonomous Learning Robots, Karlsruhe Institute of Technology 2 FZI Research Center for Information Technology
ABSTRACT
Diffusion models optimized via variational inference (VI) have emerged as a promising tool for generating samples from unnormalized target densities. These models create samples by simulating a stochastic differential equation, starting from a simple, tractable prior, typically a Gaussian distribution. However, when the support of this prior differs greatly from that of the target distribution, diffu- sion models often struggle to explore effectively or suffer from large discretiza- tion errors. Moreover, learning the prior distribution can lead to mode-collapse, exacerbated by the mode-seeking nature of reverse Kullback Leibler divergence commonly used in VI. To address these challenges, we propose end-to-end learn- able Gaussian mixture priors (GMPs). GMPs offer improved control over explo- ration, adaptability to target support, and increased expressiveness to counteract mode collapse. We further leverage the structure of mixture models by proposing strategy to iteratively refine the model by adding mixture components during training. Our experimental results demonstrate significant performance improve- ments across a diverse range of real-world and synthetic benchmark problems when using GMPs without requiring additional target evaluations.
INTRODUCTION
Sampling methods are designed to address the challenge of generating approximate samples esti- mating the intractable normalization constant Z for a probability density æ on IR the form
p(x) Z 2 p(x)dæ,
x <(x)
Rd
where : Rdd - (0, oo) can be evaluated. This formulation has broad applications in fields such as Bayesian statistics, the natural sciences (Liu & Liu, 2001; Stoltz et al., 2010; Frenkel & Smit, 2023; Schopmans & Friederich, 2024), or robotics (Zhou et al., 2024; Celik et al., 2025).
Monte Carlo (MC) methods (Hammersley, 2013), Annealed Importance Sampling (AIS) (Neal, 2001), and their Sequential Monte Carlo (SMC) extensions (Del Moral et al., 2006; Arbel et al., 2021: Matthews et al., 2022; Midgley et al., 2022) have long been regarded as the gold standard for tackling complex sampling problems. An alternative approach is variational inference (VI) (Blei al., 2017), which approximates tractable distributions. Recently, there has been growing interest in diffusion models (Zhang & Chen, 2021; Berner et al., 2022; Richter al., 2023; Vargas et al., 2023a;b; Grenioux al., 2024), which employ stochastic processes to transport samples from a simple, tractable prior distribution the target distribution. While diffusion models have shown great success in generative modeling (Ho et al., 2020; Song et al., 2020), their application to sampling tasks introduces unique challenges.
intractable target distribution by parameterizing family of
We identify these challenges as follows (C1-C3): Unlike generative modeling, where the support of the target distribution is often known, in sampling tasks, the target's support is usually unknown. This makes it difficult to set the prior appropriately and requires the model explore the relevant regions of the space-an exploration that becomes exponentially harder dimensionality increases (CI). Additionally, large discrepancies between the support of the prior and the target distribution
"Correspondence den b.l esaii ig@kiit.e

---

Published as a conference paper at ICLR 2025
C1: Lack of exploracion capabilisies.
C2: Complex dynanies require mazy diffusion steps.
collague reverse
Learned mixture prices to
Figure Illustration of challenges (CI-C3) associated with diffusion-basee sampling methods and how learned Gaussian mixture priors address them (bottom right). Here, π denotes the target distribution.
can lead to highly non-linear dynamics, necessitating many diffusion steps to mitigate discretization errors (C2). Finally, while joint optimization of the prior and diffusion process is possible, using simple priors like Gaussians can result mode collapse due to mode-seeking behavior of the re- verse Kullback-Leibler (KL) divergence commonly used in VI (C3). These challenges are illustrated in Figure 1.
Outline. In Section 3, present overview of diffusion- based sampling methods within the framework of variational inference. Next, we discuss the necessary adaptations for supporting the learning of arbitrary prior distributions, illustrated through specific examples of diffusion models (Section 4). We then provide a rationale for our choice of Gaussian Mixture Priors (GMPs) and introduce novel training scheme designed iteratively refine diffusion models during training (Section 5). Finally, in Section 6, we assess our method through experiments on a range of real- world and synthetic benchmark problems, demonstrating consistent improvements in performance.
RELATED WORK
Sampling and Variational Inference. Numerous works have studied the problem of sampling from unnormalized densities to estimate the partition function Z, including Monte Carlo (MC) methods such as Markov Chain Monte Carlo (MCMC) (Kass al., 1998) and Sequential Importance Sam- pling (Liu et al., 2001). Seminal works include Annealed Importance Sampling (Neal, 2001; Guo al., 2025) and its Sequential Monte Carlo extensions (Del Moral al., 2006; Arbel et al., 2021; Wu et al., 2020; Matthews et al., 2022; Midgley al., 2022). Another line of work approaches the sampling problem by utilizing tools from optimization to fit a parametric family of distributions to the target density π. known Variational Inference (VI) (Blei et al.. 2017). To that end, one typi- cally uses the reverse Kullback-Leible divergence, although other discrepancies have been studied (Li & Turner, 2016; Midgley et al., 2022; Dieng et al., 2017; Richter al., 2020; Wan al., 2020; Naesseth et al., 2020).
Diffusion-based Sampling Methods. Recently, there has been growing interest in combining Monte Carlo methods with variational techniques by constructing a sequence of variational dis- tributions through the parameterization of Markov chains (Naesseth al., 2018: Geffner & Domke, 2021; Thin et al., 2021; Zhang et al., 2021; Chen et al., 2024b). In the limit of infinitely many steps, these Markov chains converge to stochastic differential equations (SDEs) (Särkkä Solin, 2019), which has led to further research on diffusion- based models for sampling, particularly in light of advances in generative modeling (Ho al., 2020; Song et al., 2020). One line of work con- sidered parameterized drift functions to improve annealed Langevin diffusions in the overdamped (Doucet et al., 2022a) or underdamped (Geffner & Domke, 2022) regime. Another line of work casts diffusion-based sampling as a stochastic optimal control problem (Dai Pra, 1991) including denois- ing diffuuson models (Berner et al., 2022; Vargas et al., 2023a), and Follmer sampling (Föllmer, 2005: Zhang& & Chen, 2021; Vargas al., 2023b). A unifying view was later provided by Vargas et al. (2024); Richter et al. (2023): Blessing et al. (2025). Further extensions to diffusion-basee sampling methods have been proposed such improved learning objectives (Zhang et al., 2023;
2

---

Published as a conference paper at ICLR 2025
Akhound-Sadepee et al., 2024; Noble et al., 2024) combinations with sequential importance sam- pling (Phillips et al., 2024; Chen et al., 2024a). Another study leverages physics- informed neural networks (PINNs, (Raissi et al., 2019)) to learn the Fokker-Plande equation governing the density evolution of the diffusion process (Sun et al., 2024; Shi et al., 2024).
PRELIMINARIES
In this section, we offer a concise overview of diffusion models within the context of variational inference. Our discussion draws primarily from the works of Richter et al. (2023); Vargas et al. (2024). While these studies emphasize the continuous-time perspective, we adopt an approach that largely emphasizes discrete time, aiming make this work more accessible readers without background in stochastic calculus.
3.1 CONTROLLED DIFFUSIONS, DISCRETIZATION, AND COUPLINGS
We consider two -value stochastic processes (X))) on the time-interval [o, T): One starts from prior distribution po and runs forward in time whereas the other starts from the target distribu- tion Pr = π and runs backward in time. These processes are governed by the stochastic differential equations (SDEs) given by controlled diffusions, that is,
d.X: tf(x,, + ou (Xc, (X,, dt + VoodB,,
Xo Po.
(2a)
dX;; = [f((),,) - (X,, d) dt + V2ddB,, XT ~ =
(2b)
T) Rd, respectively. Further, (B, (Be)rEjo,T is a d-dimensional Brownian motion and a E R+ diffusion coefficient. For integration, we consider the Euler-Maruyama (EM) method with constant discretization step size 2 such that N T is integer. To simplify notation, write In, instead of Fnôt. Integrating (2a) yields
with drift, and parameterized control functions f, u, IR 07;
In+1 = In + [f(z,, n) + ou (Zn, m)] öt + o 2öten,
To Po:
(3)
where en ~ I), The EM discretizations of (2a) and (2b) admit the following Markov Processes
N
n (<o:N) = Po(2))
II
Pnin- - (In
1n-1),
and
(4)
n=
N
p (To:N) (*(
II P. (In-1lxan)
(5)
in sense that 8 and pº converge to their law, respectively, as 5t + O. Here,
F1 (Tu+1 (xn) = N (In+ilæn + [f n) + ou (x,, m)] 8t, 20 2tt 1 (#n-1 l2n) // (In-1læn -If [f n) o (xn, n)] 8t, 20 20 Stl).
and
(6)
(7)
The goal of diffusion- sampling methods is to align these processes by learning control func- tions n and 7 such that
(To:N) (2uuN).
Assuming (8) holds, we have f (xo.N Jd = ff (Zo.N )ooo:w)ddoo-11- (x meaning that we can sample 2o ~ Po and integrate the 'forward' diffusion process (2a) obtain samples from æ. In contrast, the 'backward' process (2b) is not needed for generating samples from n, but required for estimating the normalization constant Z, and obtaining a tractable optimization objective which is discussed in the next section. Lastly, we want to highlight that this formulation of diffusion fiffusiio-bbsedd sampling is very generic and that most instances of samplers, such as denoising diffusion samplers (Berner et al., 2022; Vargas et al., 2023a), be recovered by choosing the drift and control functions in (2) appropriately. We refer the interested reader to Richter et al. (2023) for further details.
3

---

Published as a conference paper at ICLR 2025
# (((on)) po(20) II Pofn- (2al (0ull)))))))
ming,; DEL 187)
F (Mnen) = s(ZN) (1-111)))
(Ma-1(pa)
Diffusion Based Sampling: The goal align two parameterized Markov Processes 5 and ? The former starts at the prior po and runs forward in time while latter starts at target and runs backward.
Figure 2:
3.2 VARIATIONAL INFERENCE FOR DIFFUSION MODELS
Variational Inference (Blei et al., 2017) uses a parameterized tractable distribution pe and minimizes divergence the target distribution æ with respect its parameters 0, typically the Kullback- Leibler divergence, i.e.,
(x) p(x)
DKL (x))(T(()) = Exzg log
+ log 2 = ELBO log Z,
It directly follows that minimizing DKL, or equivalently, maximizing the ELBO' does not require access to the true normalization constant as it is independent of Moreover, using the fact that DKL. 2 O, it is straightforward to see that ELBO(() S log Z.
In the case of diffusion models, we are interested learning the parameters 0, of the control functions u", 07, Directly minimizing DKL between (Ix) (zo. and a is challenging. However, the droc-psocessing inequality (Cover, 1999)
DKL (p% (((rr))((((T))) < DKL (Zo
(10)
provides an auxiliary, tractable, objective for optimizing (0,7).
7), that is,
Pu( p(IN)
F Pain- - (Fxr (1)) log
DKL (5" (Io:N) = E EROIN
log
+
+ +log Z,
P (I [xn
-((,,))
(11)
where L(0,7) is often referred to as augmented or extended evidence lower bound, it has ad- ditional looseness due the latent variables To:N-1 (Geffner Domke, 2021). Note that the VI setting for optimizing diffusion models is different from techniques used when samples from the target, i.e., NN π are available. The former requires simulations Zo.N ~ p for optimization, while the latter minimizes the forward KL DRL. (b ") allowing for simulation-free optimiza- tion techniques such as denoising score-matching matching (Liu et al., 2022; Shi al., 2024). Moreover, recent works consider minimizing other loss functions that the KI. divergence (11). ). A recent overview of possible alternatives can be found in Domingo-Eie -nrich (2024). For further details, the interested reader is referred to Berner et al. (2022); Vargas et al. (2024).
(Vincent, 2011; Song Ermon, 2019) or bridge
4 END-TO-ENT END LEARNING OF PRIOR DISTRIBUTIONS
We aim to learn a parametric prior Pd with parameters & end-to-end when maximizing the extended ELBO L (11)). To that end, consider two requirements:
1. We can compute gradients of L with respect to é.
Evidence Lower Bound. The terminology stems from Bayesian inference, where log Z is equivalent to the evidence of the data.
4

---

Published as a conference paper at ICLR 2025
2. There exists and y such that po (20) (
For the former, we assume that Pb (In) ) amendable to the reparameterization trick?, i.e., we can express sample from Pe as deterministic function of random variable with some fixed distribution and the parameters Φ, i.c., To = g(), We can then obtain gradients of
log p(aN) Pó (x0)
N
Pn (In- (x P (In (2n-1)
C(O, = Baa.N
log
(12)
n=1
with 57 (xTON)) = (x0) II Pa Pn/n-1 (n))) with respect to o, by differentiating through the stochastic process
In+ = Inn (In, n) ou (wn; n)] 8t oV28lln: in g(S,
(13)
The second requirement is necessary to obtain coupling between pf and æ,, i.c., to satisfy (8). This requirement is trivially fulfilled for a controlled backward process (2b), where we can learn a such that p transports π back to pg. However, this requirement can be more intricate for other processes and will be discussed in the next sections. In particular, we look two instances of (2), namely denoising diffusion models (Berner et al., 2022; Vargas et al., 2023a) and annealed Langevin diffusions (Doucet al., 2022a; Vargas al., 2024).
4.1 DENOISING DIFFUSION MODELS
Denoising diffusion models use Ornstein Uhlenbeck (OU) process for (2b), that is 3,
dX;; = o² Xtdt + 22ddB XT ~ PT = n,
(14)
and, hence, a special case of (2) with f(x,,t) o²X, and v² 0. Assuming a sufficiently large (or T), it holds that Po (xo) = J N )d R 7). In other words, the OU process transports the target π to a Gaussian distribution. We extend denoising diffusion models to support learning arbitrary priors based on Proposition 1, whose proof can be found in Appendix A.1.
Proposition 1. Let (2b) be a (uncontrolled) stochastic process as defined in (2) with v² = 0, starting from Pr w. For a fime-independent drift, i.e., f(x,t) f(x), the stationary distribution (x) for which Ops(2s) = O holds, is given by dt
1 z
1 ( f(o)dz)
(=(x)
exp
(15)
with normalization constant Z
Rewriting (15), yields f = -02 vn log resulting in the SDE
dX, -'V, log (X.)dl + 2odB. XT ~ pr=p-=- =
(16)
with stationary distribution (x)) Note that denoising diffusion models leverage this result by setting p² = (0, 1). resulting in the OU process ( (14)) since V, logg =z. Hence, we can
adapt existing denoising diffusion sampling methods (Vargas al., 2023a; Berner et al., 2022) to arbitrary priors dd using
dx, - log d (Xe)dt + V20dB XT ST~PT=R
(17)
However, contrary to the OU process, where the relaxation time, i.e., the time scale which the system loses memory of its initial conditions and approaches its stationary distribution, can be estimated analytically, it is unknown for general yd and is only guaranteed T (Roberts & Tweedie, 1996). We address this by additionally learning the time horizon T = Not by treating the discretization step size òt learnable parameter. As such, the parameters of the stationary
NNote that this requirement not necessary when minimizing loss function where the expectation is not computed with respect to samples from p For further details see e.g. (Richter et al., 2023). 'The sign differs here compared works as the SDE is integrated backward in time.
5

---

Published as a conference paper at ICLR 2025
distribution, i.e., the prior distribution and the discretization step size öt are optimized jointly by maximizing the extended ELBO
p(aN) Po (3r0)
Pn 1/22 (en- Pn|n-1 (2n((nn1))
C(0, Φ) = ExO.N ap. log
log
(18)
R=1
with additional parameters Φ, öt. Please note that we omit dependence ôt in Equation (18) to keep the notation uncluttered. Proposition thus suggests, that for any Ø, there exists a öt such that Po (2o) = f p 9 (# )d as N + co. Empirically, we observe substantial improvements for finite values of N, as demonstrated in Section 6
4.2 ANNEALED LANGEVIN DIFFUSIONS
Annealed Langevin Diffusions equation by constructing a sequence of distributions (Te)) 27 that anneal smoothly from the prior distribution To Po the target distribution रт π. One typically uses the geometric average, that is, (2)2) = Po (z) (z) (x)) 1-A, for 3r monotonically increasing in t with Bo = 0 and 3T = 1.
annealed version of the (overdamped) Langevin diffusion
When learning the prior,
parametric annealing, i.c., { (x) Po & (z)) The corresponding stochastic processes can be described as an instance of (2) given by
dX,-
0² V, x log कं (Xt) + ou (Xx: de + V20dBr, Xo po = p.
(19)
dX, V² log के (X,) σo (X dt + 2odB XT NT~PT = n,
(20)
when setting f Vz l.
n. Note that V log oh can be computed without knowing the normaliza-
tion constant Z of æ. Different variants be derived from using either controlled or uncontrolled processes: Monte Carlo Diffusions (MCD) (Doucet et al., 2022b) uses a controlled process (2b) but uncontrolled (2a) and Controlled Monte Carlo Diffusions (CMCD) (Vargas et al., 2024) control both processes. Since both methods use controlled backward processes (2b), the second require- ment is satisfied. Finally, while this work focuses on overdamped approaches, we want highlight that there exist methods that are based on the underdamped Langevin equation (Geffner & Domke, 2021; Geffner & Domke, 2022), however, the idea of learning prior end-to-end straightforwardly transfers to these approaches.
5 GAUSSIAN MIXTURE PRIORS AND ITERATIVE MODEL REFINEMENT
In this work, we focus on end-to-end learned Gaussian mixture priors (GMPs), that is,
K
(x))
O&PP (xo)
auN(soluu Ex), 0g 2 0,
Σ a 1,
(21)
with mixture weights og, Gaussian components p = Axo Mh; Σx) and parameters
U (ak, Φx) with du {ur, Σx}. Having established how the prior is learned in Section we discuss desirable properties to address the challenges outlined in Section 1 and how GMPs address them.
A key objective is to improve the exploration capabilities of iuffusion-based sampling methods to address C1. GMPs allow control over exploration by adjusting the initial variance of each Gaussian component. Additionally, the means of the Gaussian components can be initialized to incorporate prior knowledge of the target density, even if this knowledge is limited to a rough estimate of the target's support. This aspect will be claborated on later in this section.
Another important consideration to adjust the support of the prior such that it matches the target density, which reduces the complexity of the dynamics and, in turn, minimizes the number of dif- fusion steps required. GMPs demonstrate rapid adaptation capabilities, partially through their small parameter count, making them particularly suitable for addressing C2.
To prevent the model from focusing only on a subset of the target support (C3), which may oceur due to the optimization of the mode-seeking reverse KI. divergence, we require more expressive

---

Published as a conference paper at ICLR 2025
distribution than single Gaussian prior. GMPs provide solution by combining multiple Gaussian components, each of which can focus on different subsets of the target support.
Finally, efficient evaluation of P is crucial, as it must be performed at each discretization step of the stochastic differential equation (SDE) that governs the diffusion process. This requirement is satisfied by GMPs, particularly when using diagonal covariance matrices.
Iterative Model Refinement. Gradually increasing the model complexity during the optimization process has demonstrated promising results in previous studies (Guo et al., 2016; Miller et al., 2017; Arenz et al., 2018; Cranko & Nock. 2019), and is directly applicable our approach. We begin with an initial prior distribution P = pe with parameters Φ1. These parameters are optimized using (18). After a predefined criterion met, such as a fixed number of iterations, a second distribution p is added, forming new prior: Po A c110" d1 + a220 d2 with E RC and (1) +02 1. This process is repeated, resulting in a mixture model p (x) = -KK a&Po ot (z)).
We identify the benefits of this iterative scheme twofold: First, it can simplify optimization by focusing on learning a subset of parameters ΦR at time, rather than jointly optimizing all Φ (Bengio et al., 2009). Second, it enables the initialization of newly added components based on a partially trained model, potentially preventing mixture components to focus on similar parts of the target support. For GMPs, for instance, the mean of new component Vnew can be placed in promising region, potentially informed by prior knowledge of the task or by running a rriivariant Markov chain to obtain a set of promising samples. More generally, consider set of candidate samples = {
We propose initializing the mean of a new component Unew as follows:
log p(ENN) Po (mx)
N
2-2 ( log 0,0 p n - (.n|
Pnow = arg max E 2N EC
BOIN
(22)
where Po is the current model. This heuristic balances exploration and exploitation by favoring samples with high target likelihood and low prior likelihood, while also accounting for the diffusion process. See Appendix B for further details.
6 NUMERICAL EVALUATION
In this section, we test the impact of our proposed end-to- end learning scheme for prior distributions. Specifically, we consider three distinct settings: First, we evaluate these meth- ods with Gaussian prior that is fixed during training. Sec- ond and third, we consider learned Gaussian (GP) and Gaus- sian mixture priors (GMP). We indicate these different settings as X. X-GP, and X-GMP, respectively, where X is the corre- sponding acronym of the diffusion-based sampling methods. We consider four different methods: ReverRReersed Diffusion Sampler (DIS) (Berner et al., 2022), Monte Carlo Diffusions (MCD) (Doucet et al., 2022b), Controlled Monte Carlo Diffu-
METHOD (X)
MCD
V Plogs? x2
CMCD'
V log
DIS
V p
x
DBS
ANY
Table I: Diffusion-basee sampling methods considered this work based on (2). Crosses indicate that the con- trol is set to zero.
sions (CMCD) ( (Vargas et al., 2024) and Diffusion Bridge Sampler (DBS) (Richter et al., 2023). A summary is shown in Table It is worth noting that do not separately consider the Denoising Diffusion Sampler (DDS) (Vargas et al., 2023a), as it can be viewed as a special case of DIS. For reference, we consider Gaussian (GVI) and Gaussian mixture (GMVI) mean-field approximations (Wainwright & Jordan, 2008), both of which are special cases of the aforementioned methods for N diffusion steps with K 1, K 2 1, respectively (cf. Appendix B). Lastly, we consider three competing oofte--f-th-ar methods, namely, Sequential Monte Carlo (SMC) (Del Moral et al., 2006), Continual Repeated Annealed Flow Transport (CRAFT) (Matthews et al., 2022), and Flow Annealed Importance Sampling Bootstrap (FAB) (Midgley et al., 2022).
For evaluation, we consider the effective sample size (ESS) and the marginal or extended evidence lower bound as performance criteria. Both are denoted as 'ELBO' for convenience. Next. if the ground truth normalization constant Z is available, we use importance-weigheed estimate 2 to
Vargas et al. (2024) the same in control in (2a) and (2b) by leveraging Nelson's relation (Nelson, 2020).

---

Published as a conference paper at ICLR 2025
compute the estimation error A log 2 = log 2 - -log 2). Additionally, if samples from the target are available, we compute the Sinkhorn distance W (Cuturi, 2013).
To ensure fair comparison, all experiments conducted under identical settings. Our evaluation methodology adheres to the protocol by Blessing et al. (2024). For a comprehensive overview of
the experimental setup see Appendix C. Moreover, comprehensive set of ablation studies and additional experiments, are provided in Appendix D.
GVI: K-1 N-
GMVI: R 10, N
Method
ELBOT 1
A logg 4
ESS
WV? 1
GVI
-18841
0.000annee
178. 24:100
GMVI
0:056-1.014
0.764
102 8.6
0.201
0.207a
164,852
MCD-GP
0.173u8.245
a.206e
166.967.6.11
GMP
0.014-0.001
0.922at.e12
174..74
CMCD
-02200
0:020:9.:03
0.588atves
104.853 052.652
CMCD-GP
-0211a3002
0.020-9.001
108.644
DIS-GP: X 1, M 126
DIS-GMP: R- N 128
CMCD-GMP
-0.0
0.005=1.00
05000...00 aus
102 022.0.
0:04
0.483atveas
107, 858.8.
DIS GP
-0220
0.047:9.201
0.498.a
107, 607.458
GMP
0:019-3:001
0.029eatuezz
DBS
-0.180 10
0.019
0.600stvv44
102. 06166.6643
DBS
0:021-9.00
0.600
102.85 652.853
DBS-GMP
0:012-9:011
0.9
100 230 ses
SMC
140. 353
CRAFT
-0.07
0:09::9.015
136 335
FAB
0.014:0.02
1.001110.00
Num. mixtune compoteetts
N.. ditfusioa stopa
Figure 3: Left side: Results for Funnel target, averaged across four seeds. Evaluation criteria include evidence errors for estimating the oog-ooraaiziin constant log Z, effective sample size ESS, Sinkhorn distance W,"" The best overall results are highlighted in hold, with category-specifie best results underlined. Arrows (t, J) indicate whether higher lower values are preferable, respectively. Blue and green shading indicate that the method uses learned Gaussian (GP) and Gaussian mixture priors (GMP), respectively. Red shading indicate competing state-of-the-an methods. Note that ESS cannot be computed due to the use of resampling schemes. Right side: Visualization of the first two dimensions of the Funnel target. Colored cllipses and circles denote standard deviations and means of the Gaussian components, respectively. Red dots illustrate samples of the model.
lower bound ELBO, importance-weigheee
6.1 BENCHMARK PROBLEMS
We evaluate the different methods on various real-world and synthetic target densities.
Real-World Densities. We consider six real-world target densities: Four Bayesian inference tasks, where inference is performed over the parameters of a logistic regression model, namely Credit (d 25), Cancer (d 31), lonosphere (d = 35), and Sonar (d 61). Moreover, Seeds (d 26) and Brownian (d = 32). where the goal is to perform inference over the parameters of a random effect regression model, and the time discretization of Brownian motion, respectively. For these densities, we do not have access to the ground truth normalizer Z or samples from и preventing us from computing errors for log normalization estimation A log 2 and Sinkhorn distances W?. The resulting ELBO values are presented in Table 2.
Synthetic Densities. The Funnel density was introduced by Neal (2003) as has a shape that resem- bles funnel, where one part is tight and highly concentrated, while the other is spread out over wide region, making it challenging for sampling algorithms to explore the distribution effectively. Next, we consider the Fashion target which uses NICE (Dinh et al., 2014) to train normalizing flow on the high-dimensional d = 28 x 28 = 784 MNIST Fashion dataset. A recent study by Blessing al. (2024) showed that current state-of-the-ard methods were not able to generate samples with high quality from multiple modes.
6.2 RESULTS
Impact of Learned Gaussian (GP) and Gaussian Mixture (GMP) Priors. We evaluated the performance of our proposed methods both real-world tasks and the Funnel density, employing N = 128 diffusion steps across all methods and K = 10 mixture components for X-GMP. To ensure fair comparison, we initialized the priors of all diffusion-based methods with mean
8

---

Published as a conference paper at ICLR 2025
METHOD
CREDIT
CANCER
BROWNIAN
IONOSPHERE
SONAR
-605.561_0.10 _00.10
10.041
147.453...14
3.885g0.000
123.391 10 ora
137.096 o13
GMVI
-603.422
+0.222
1445.556 0..
2.250=0.001
122. 0.0
195.959++..o
MCD
1399..22
(899.211=49..11
75. (3,6999eo.uu 01s
1448. 171 sos
15. 15519580.
114 320000 007
112. 639 n2s
MCD-GP
554.3300_omm _..00
73.542.0 o00
89. 39.67661 188
0.771 g.o 10.00s
11.887 oo1
100.338 001
MCD-GMP
276-0.0
73. 33,466++..00
0.9813=0.00s
111 827+0.00
197.
CMCD
-586.956=a.oc
-74.033+0..0 010
80 076
1.3316
112 183+ oos
109. 330.332++.00 noe
CMCD-GP
73.456. o00
78.5761 now
1.043go oo5
111.687ash ous
108. 108.689 one
CMCD-GMP
-58.
162-0.002
73. 002
78, -0.007
1.087=0.001
111 682 one
589.636__... _.757
74.400 007
86.502 107
3.503 ona
112 112.525..o ous
110. 150.153...0 o2e
DIS-GP
545.247__0.o0 _0.00
73.540.. ous
85.005 286
0.588
111.847.. ooe
100.260.0 094
DIS-GMP
223-0.006
73. 231.492+00 to.
+2.117
0.885=0.00m
111 111.811+o0m 002
157.11.770000
DB$
587.366_n.dd _.663
73.918 ons
82.466 $2.466114.000
0.773 onn
112.070.o ous
109. 168 oos
DBS-GP
-585.524-0444
73. 437+0.001
833 -83.3955+4..16
1.081=0.00s
111.6734e..0 002
1188.995+0.00
DBS-GMP
-586.14820.00
73.418::00 001
78. 160 oss
1.118
111 600.0.000
-088. 20..5488
SMC
698.403_.4.16
74.699
194.059 11591...11
1.884_n..o
114.751 zas
1.1555511.11 177
CRAFT
5991.795-0.411
-73. 73..00
-5.7737=1007
0.886=0.05s
1122 112.38644..a 182
-155.618+1.84
FAB
773.418:00 002
1.031 20.010
1111.67840.m
108.593.1..0
Table 2: Evidence lower bound (ELBO) values for various woaldworld benchmark problems, averaged across four seeds. The best overall results are highlighted in bold, with category-specifee best results underlined. Blue and green shading indicate that the method uses learned Gaussian (GP) and Gaussian mixture priors (GMP), respectively. Red shading indicate competing state-of-the-art methods.
and unit variance. Table 2 and Figure 3 present our findings. The analysis demonstrates that GP consistently achieves tighter ELBO values compared to fixed priors, with GMP yielding further improvements over GP. Furthermore, Figure 3 illustrates both qualitatively and quantitatively that GMP effectively combines the strengths of Gaussian mixture and diffusion models, resulting in significant improvements. Specifically, we observed that the Gaussian components adapt well to the target's support, covering both the neck and opening of the funnel shape. This results in less complex dynamics and better target coverage for DIS-GMP compared to using a single Gaussian (DIS-GP). Notably, the combination of DBS and GMP outperforms sfate-of-the-art methods across the majority of tasks and evaluation metrics.
DIS-GMP
DIS-CMP IMR
Fadiion
Methed
ELB01
& kig 2
EMC
78.70800008
47 J08
1500
GMVI
-72.65 -72.614mmm
45 1597771100
1506
GMVE IMR
-57.00 lesuea
50.44
580
58.8888130 anc
0 0.76ss
32 300
SMC
CRAFT
FAB
DIS-CP
-24 aMz
10 501811m
167
DES GMP
-35875
180006.0.00
1700.000.1.00
-22 ma
37 07
513.776-11
SMU
-12.180arsee
13. 21.70a
0916
CRAFT
-520 ATeaoen
1413.M0
-99207415
596.54.900
367
0.34Du
Figure 4: Left side: Results for Fashion target, averaged across four seeds. Evaluation criteria include evi- dence lower bound ELBO, amportance-weighled emrors for estimating the log-normalizing constant A log Z, and Sinkhorn distance W/". The best overall results are highlighted in bold, with category-specific best results underlined. Arrows (t. 4) indicate whether higher or lower values are preferable, respectively. Blue and green shading indicate that the method uses learned Gaussian (GP) and Gaussian mixture priors (GMP), respectively. Orange shading indicates that the method uses iterative model refinement (IMR). Red shading indicate com- peting state-of the methods. Right side: Visualization of the samples. Top left corner visualizes samples from the target distribution. Colored frames indicate samples from different components of the Gaussian mixture.
x 784 dimensional Fashion
Ablation Study: Number of Mixture Components K and Diffusion Steps N. We further in- vestigated the effect of varying the number of diffusion steps N and mixture components K on subset of tasks for DIS. The results, shown in Figure 5, demonstrate consistent improvements in effective sample size (ESS) with increases in both K and N. Additionally, we consistently observed that the combination of a higher number of components and diffusion steps yields the best overall performance. These trends hold across other metries, as further detailed in Appendix D.

---

Published as a conference paper at ICLR 2025
Iterative Model Refinement (IMR). Lastly, investigated the impact of IMR, as detailed in Section 5, using DIS. For this analysis, we focused on the muult--mdall Fashion target, which neces- sitates exploration in a high-ding sish-dimessioal space (d 784). In addition to the performance criteria outlined in Section 6, we quantify how many of the modes the model discovered via the entropic mode coverage (EMC) introduced by Blessing et al. (2024). EMC evaluates the mode coverage of a sampler by leveraging prior knowledge of the target density. It holds that EMC E [0. 1] where EMC = indicates that the model achieves uniform coverage over all modes whereas EMC indicates that the model only produces samples from a single mode. We employed the Metropolis- adjusted Langevin algorithm (MALA) (Cheng et al., 2018) to generate a set of candidate samples, noting that the computational cost of this process is comparable to a single gradient step in most diffusion-based sampling methods. The initial candidate samples as well as the support of DIS without learned prior are initialized such that they roughly cover the target support. Additional de- tails are provided in Appendix C.2. We iteratively increased the number of components to K = 10, utilizing N 128 diffusion steps throughout. Figure 4 presents our findings, demonstrating that the absence of IMR leads to mode collapse across all methods, as evidenced by high Sinkhorn dis- tance values. The qualitative results highlight the role of candidate samples in facilitating mode discovery. Notably, the color-coding of DIS-GMP + IMR illustrates that each mixture component concentrates on distinct mode, validating the effectiveness of the initialization heuristic proposed in (22) in balancing exploration and exploitation. This finding is also quantitatively reflected by the high EMC and low Sinkhorn distance values. In contrast, the ELBO and A log Z values are slightly worse when using GMPs and IMR. This is attributed to the fact that these performance cri- teria are not well-suited for quantifying the model performance for multi-modal targets and tend to favor models that fit single mode perfectly (Blessing et al., 2024). Moreover, with higher K, the diffusion model has to learn more complex control functions, as it needs to operate over the support of the entire Gaussian Mixture Model (GMM) rather than single Gaussian. This added complexity can introduce more opportunities for approximation errors, which may negatively impact ELBO and A log Z values compared to using single learnable Gaussian. Nevertheless, the resulting samples from DIS- GMP are closer to the target distribution in terms of Sinkhorn distance. Importantly, these errors remain significantly smaller than those observed with non-learnable priors.
Scarl
Creck
0:13
A47
0.3
1:22
6:35
AT2
422
0152
600
0:50
1.25
02.
1:31
0.52
6.55
3.00
0.01
012
8:2
8:30
8.14
8:22
342
0144
0.21
0:22
0:28
4.77
0.54
014
Figure 5: Effective sample size (ESS) of DIS-GMP for various real-world benchmark problems, averaged across four seeds. Here, N denotes the number of discretization steps and K the number of components in den Gaussian mixture.
7 CONCLUSION AND FUTURE WORK
In this paper, we propose a novel approach for improving diffusion-based sampling techniques by introducing end-to-end learnable Gaussian Mixture Priors (GMPs). Our method addresses key chal- lenges in diffusion models-namely large discretization errors, mode collapse, and poor explo- ration-by providing more expressive and adaptable priors compared the conventional Gaussian priors. We conducted comprehensive experiments on both synthetic and real-world datasets, which consistently demonstrated the superior performance of our proposed method. The results underscore the effectiveness of GMPs in overcoming the limitations of traditional diffusion models while re- quiring little to no hyperparameter tuning. Furthermore, we developed a novel strategy for iterative model refinement, which involves progressively adding components to the mixture during training. and demonstrated its effectiveness on a challenging high-dimensional problem.
A promising direction for future research is the improvement of the iterative model refinement strat- egy. While we showed that progressively increasing the number of components in the Gaussian mixture improves performance, optimizing the selection criteria for adding new components, gener- ating better candidate samples, or dynamically adjusting the number of components during training, could lead to further gains in efficiency and accuracy.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS AND DISCLOSURE OF FUNDING
D.B. acknowledges support by funding from the pilot program Core Informatics of the Helmholtz through bwHPC, as well the HoreKa
Association (HGF) and the state of supercomputer funded by the Ministry of Science, Research and the Arts Baden-W Württemberg and by the German Federal Ministry of Education and Research.
REFERENCES
Tara Sahound-Sadeghh Jarrid Rector-Brooks. Avishek Joey Bose, Sarthak Mittal, Pablo Lemos, Cheng-Hao Liu, Marcin Sendera, Siamak Ravanbakhsh, Gauthier Gidel, Yoshua Bengio, et al. Iterated denoising energy matching for sampling from boltzmann densities. arXiv preprint 066112 2024.
Michael Arbel, Alex Matthews, and Armaud Doucet. Annealed flow transport monte carlo. In International Conference Machine Learning, pp. 318-330. PMLR, 2021.
Oleg Arenz, Gerhard Neumann, and Mingjun Zhong. Efficient gradient- free variational inference using policy search. In International conference on machine learning, pp. 234-243. PMLR, 2018.
Oleg Arenz, Philipp Dahlinger, Zihan Ye, Michael Volpp, and Gerhard Neumann. A unified per- spective on natural gradient variational inference with gaussian mixture models. arXiv preprint arXiv:2209.111 11533, 2022.
Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Proceedings of the 26th annual international conference on machine learning, pp. 41-48, 2009.
Julius Berner, Lorenz Richter, and Karen Ullrich. An optimal control perspective on diffusion-based generative modeling. arXiv preprint arXiv:2211.01366 2022.
David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational inference: review for statisti- cians. Journal of the American statistical Association, (12(518)8599-33 2017.
Denis Blessing, Xiaogang Jia, Johannes Esslinger, Francisco Vargas, and Gerhard Neumann. Be- yond elbos: A large-scale evaluation of variational methods for sampling. arXiv preprint arXiv:2406. 2444::2406.07423, 07423, 2024.
Denis Blessing, Julius Berner, Lorenz Richter, and Gerhard Neumann. Underdamped diffusion bridges with upplications sampling. In The Thirteenth International Conference on Learning Representations, 2025. URL https:/ /ttp//ppeeeeeeeeeeeeeeeeee
//peattss///peeeeeeiee. net forum? fttt///oooomm0m00000
James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Minne al. Jax: Autograd and xla. Astrophysics Source Code Library, pp. ascl 221-2111,, 2021.
Onur Celik, Zechu Li, Denis Blessing, Ge Li, Daniel Palanicek. Jan Peters, Georgia Chalvatzaki, and Gerhard Neumann. Dime: Diffusion-based maximum entropy reinforcement learning. arXiv preprint arXiv:2502.02316, 2025.
Junhua Chen, Lorenz Richter, Julius Berner, Denis Blessing, Gerhard Neumann, and Anima Anand- kumar. Sequential controlled langevin diffusions. arXiv preprint arXiv:2412.07088, 2024a.
Wenlin Chen, Mingtian Zhang, Brooks Paige, José Miguel Hernández-Lobato, and David Barber. Diffusive gibbs sampling. arXiv preprint arX1v:2402.03008, 2024b.
Xiang Cheng, Niladri Chatterji, Peter I. Bartlett, and Michael Jordan. Underdamped langevin meme: A non-asympiotic analysis. In Conference learning theory, pp. 300-323. PMLR, 2018.
Thomas M Cover. Elements of information theory. John Wiley Sons, 1999.
Zac Cranko and Richard Nock. Boosted density estimation remastered. In International Conference Machine Learning, pp. 1416-1425 PMLR, 2019.

---

Published as a conference paper at ICLR 2025
Marco Cuturi. Sinkhor distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Marco Cuturi, Laetitia Meng-Papaxane Panagaaaaxanthos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201. 122224 2022.
Paolo Dai Pra. A stochastic control approach to reciprocal diffusion processes. Applied mathematics and Optimization, 23(1):313-329 1991.
Pierre Del Moral, Araad Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. Journal the Royal Statistical Society: Series B (Statistical Methodology), 68(3)::11-433) 2006.
Adji Bousso Dieng, Dustin Tran, Rajesh Ranganath, John Paisley, and David Blei. Variational inference via X upper bound minimization. Advances Neural Information Processing Systems, 30, 2017.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components esti- mation. arXiv preprint arXiv:1410.85/6, 2014.
Laurent Dinh, Jascha Dick-Di,, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv: 2016.
Carles Domingg-Farich. A taxonomy of loss functions for stochastic optimal control. arXiv preprint arXiv:2410.000345, 2024.
Arnaud Doucet, Will Grathwohl, Alexander Matthews, and Heiko Strathmann. Score-based diffu- sion meets annealed importance sampling. Advances in Neural Information Processing Systems, 21114822214994 2022a.
Arnaud Doucet, Will Grathwohl, Alexander G de G Matthews, and Heiko Strathmann. Score-based diffusion meets annealed importance sampling. In Advances in Neural Information Processing Systems, 2022b.
Hans Föllmer. An entropy approach to the time reversal of diffusion processes. Stochastic Differential Systems Filtering and Control: Proceedings of the IFIP-WG 7/1 Working Conference Marseille-Luminy France, March 12-17, 1984, pp. 156-163. Springer, 2005.
Daan Frenkel and Berend Smit. Understanding molecular simulation: from algorithms to applica- tions. Elsevier, 2023.
Tomas Geffner and Justin Domke. MCMC variational inference via uncorrected Hamiltonian an- nealing. In Advances in Neural Information Processing Systems, 2021.
Tomas Geffner and Justin Domke. Langevin diffusion variational inference. arXiv preprint arXiv:2208.00 07743, 2022.
Louis Grenioux, Maxence Noble, Marylou Gabrié, and Alain Oliviero Durmus. Stochastic localiza- tion via iterative posterior sampling. arXiv preprint arXiv:2402.10758, 2024.
Fangjian Guo, Xiangyu Wang, Kai Fan, Tamara Broderick, and David B Dunson. Boosting varia- tional inference. arXiv preprint 161 1.05599, 2016.
Wei Guo, Molei Tao, and Yongxin Chen. Complexity analysis of normalizing constant estima- tion: from jarzynski equality to annealed importance sampling and beyond. arXiv preprint arXiv:2502.04577. 2025.
John Hammersley, Monte carlo methods. Springer Science & Business Media, 2013.
Carsten Hartmann, Omar Kebiri, Lara Neureither, and Lorenz Richter. Variational approach to rare event simulation using least-squares regression. Chaos; An Interdisciplinary Journal of Nonlinear Science, 29(6), 2019.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851. 2020.
12

---

Published as a conference paper at ICLR 2025
Robert E Kass, Bradley P Carlin, Andrew Gelman, and Radford M Neal. Markov chain monte carlo in practice: a roundtable discussion. The American Statistician, 52(2):93-100, 1998.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 14/2.6980, 2014.
Xuechen Li, Ting-K Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pp. 3870-3882. PMLR, 2020.
Yingzhen Li and Richard E Turner. Rényi divergence variational inference. Advances in neural information processing systems, 29, 2016.
Jun S Liu and Jun S Liu. Monte Carlo strategies in scientific computing, volume 75. Springer, 2001.
Jun Liu, Rong Chen, and Tanya Logvinenko. A theoretical framework for sequential importance sampling with resampling. In Sequential Monte Carlo methods in practice, pp. 225-246. Springer, 2001.
Xingchao Liu, Lemeng Wu, Mao Ye, and Qiang Liu. Let us build bridges: Understanding and extending diffusion generative models. arXiv preprint arXiv:2208. 14699, 2022.
Alex Matthews, Michael Arbel, Danilo Jimenez Rezende, and Arnaud Doucet. Continual repeated annealed flow transport monte carlo. International Conference on Machine Learning, pp- 15196-15219, PMLR, 2022.
Laurence Illing Midgley, Vincent Stimper, Gregor NC Simm, Bernhard Schölkopf, and José Miguel Hernáneez-Lobato Flow annealed importance sampling bootstrap. arXiv preprint 2rXiv:2208.01889 2022.
Andrew C Miller, Nicholas J Foti, and Ryan Adams. Variational boosting: Iteratively refining posterior approximations. In International Conference on Machine Learning, pp. 2420-24299 PMLR, 2017.
Christian Naesseth, Scott Linderman, Rajesh Ranganath, and David Blei. Variational sequential monte carlo. International conference on artificial intelligence and statistics, pp. 968-977. PMLR, 2018.
Christian Naesseth, Fredrík Lindsten, and David Blei. Markovian score climbing: Variational infer- ence with kl (p- q). Advances in Neural Information Processing Systems, 33:15499-15510 15510, 2020.
Radford M Neal. Annealed importance sampling. Statistics and Computing, 11(2):125-139, 2001.
Radford M Neal. Slice sampling. The annals of statistics, 31(3):705-767, 2003.
Edward Nelson. Dynamical theories of Brownian motion, volume 101. Princeton university press, 2020.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International conference on machine learning, pp. 8162-8171 8171. PMLR, 2021.
Kim A Nicoli, Christopher J Anders, Lena Funcke, Tobias Hartung, Karl Jansen, Pan Kessel, Shinichi Nakajima, and Paolo Stornati. Estimation of thermodynamic observables in lattice field theories with deep generative models. Physical review letters, 126(3):33001, 2021.
Maxence Noble, Louis Grenioux, Marylou Gabrié, and Alain Oliviero Durmus. Learned reference- based diffusion sampling for multi-modal distributions. arXiv preprint arXiv:2410.19449 2024.
George Papamakarios, Eric Nalisnick. Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lak- shminarayanan. Normalizing flows for probabilistic modeling and inference. The Journal of Machine Learning Research, 22(1):2617-266) 2021.
13

---

Published as a conference paper at ICLR 2025
Angus Phillips, Hai-Dang Dau, Michael John Hutchinson, Valentin De Bortoli, George Deligianni- dis, and Araad Doucet. Particle denoising diffusion sampler. arXiv preprint arXiv:2402.006320, 2024.
Maziar Raissi, Paris Perdikaris, and George Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. Var- grad: a low-variance gradient estimator variational inference. Advances Neural Information Processing Systems, 184481113422 2020.
Lorenz Richter, Julius Berner, and Guan Hoag Liu. Improved sampling via learned diffusions. arXiv preprint arXiv:2307.01198, 2023.
Gareth O Roberts and Richard L Tweedie. Exponential convergence of langevin distributions and their discrete approximations. 1996.
Simo Särkkä and Amo Solin. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Henrik Schopmans and Pascal Friederich. Conditional normalizing flows for active learning of coarse-grained molecular representations. arXiv preprint arXiv:2402.01 195, 2024.
Yuyang Shi, Valentin De Bortoli, Andrew Campbell, and Arnaud Doucet. Diffusion schrödinger bridge matching. Advances in Neural Information Processing Systems, 36, 2024.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems, 32, 2019.
Yang Song. Jascha Sohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011. 134466, 2020.
Gabriel Stoltz, Mathias Rousset, et al. Free energy computations: A mathematical perspective. World Scientific, 2010.
Jingtong Sun, Julius Berner, Lorenz Richter, Marius Zeinhofer, Johannes Müller, Kamyar Aziz- zadenesheli, and Anima Anandkumar. Dynamical measure transport and neural pde solvers for sampling. arXiv preprint arXiv:2407.07873, 2024.
Achille Thin, Nikita Kotelevskii, Alain Durmus, Eric Moulines, Maxim Panov, and Arnau Doucet. Monte Carlo variational auto-encoders. In International Conference on Machine Learning, 2021.
Francisco Vargas, Will Grathwohl, and Arnaud Doucet. Denoising diffusion samplers. arXiv preprint arXiv:2302. 13834, 2023a.
Francisco Vargas, Andrius Ovsianas, David Fernandes, Mark Girolami, Neil D Lawrence, and Niko- las Nüsken. Bayesian learning via neural cchrödinger-fölllne flows. Statistics and Computing, 33(1):3, 2023b.
Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nüsken. Transport meets variational inference: Controlled monte carlo diffusions. In The Twelfth International Conference on Learn- ing Representations, 2024.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural compu- tation, 23(7):1661-1674 2011.
Martin J. Wainwright and Michael I. Jordan. Graphical Models, Exponential Families, and Varia- tional Inference. Foundations and Trends@ in Machine Learning, 1(1-2):1-305, November 2008.
Neng Wan, Dapeng Li, and Naira Hovakimyan. F-divergence variational inference. Advances in neural information processing systems, 33:17370-17379, 2020.
14

---

Published as a conference paper at ICLR 2025
Hao Wu, Jonas Köhler, and Frank Noé. Stochastic normalizing flows. Advances in Neural Informa- tion Processing Systems, 33:5933-5944 5944,,9442222 2020.
Dinghuai Zhang, Ricky Tian Qi Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffu- sion generative flow samplers: Improving learning signals through partial trajectory optimization. arXiv preprint arXiv:2310.02679, 2023.
Guodong Zhang, Kyle Hsu, Jianing Li, Chelsea Finn, and Roger Grosse. Differentiable annealed im- portance sampling and the perils of gradient noise. In Advances in Neural Information Processing Systems, 2021.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sam- pling. arXiv preprint ar114 2021.
Hongyi Zhou, Denis Blessing, Ge Li, Onur Celik, Xiaogang Jia, Gerhard Neumann, and Rudolf Lioutikov. Variational distillation of diffusion policies into mixture of experts. arXiv preprint arXiv:2406. 12538, 2024.
A PROOFS
A. PROOF OF PROPOSITION
We will use the Fokker-Planck equation (FPE) and show that the given stationary distribution satis- fies when the time derivative is set to zero:
First, recall the FPE for the probability density p(x, I) of process described by the stochastic differential equation (SDE)
dr, (x( + odw.
(23)
is given
Öp(x, t)
V ([x))((e,
2 Vir(z,1)
(24)
Öt
For the stationary distribution (x)), we set
0p(x.1) 0t
0:
0 = V (.()))) + o²
vz)'(0)
(25)
2
Next, recall the proposed stationary distribution:
p" (x) x exp (-2 / f(z)dz)
2 0²
(26)
Next, we verify that this satisfies stationary FPE ( (25)). First, let's compute the gradient and Lapla- cian of (x))
(-22)
2
= (()))>"())
[(x)
(27)
02
(
2
<2p"(2) V (x)
f(x)]
(28)
σ²
-p²(x) (-
(-20
2
[fx)]2
p²))))
(-2
2 v f(x)
(29)
15

---

Published as a conference paper at ICLR 2025
Algorithm 1 Training of diffusion sampler with learnable prior
Require:
control functions uo, Vy with initial parameters Oo. Yo
prior distribution a with initial parameters @p
initial discretization step size 8l/
number of gradient steps G, number of diffusion steps N, step size !)
On {0o: To, Φ0 òto}
for i + O..... G do
Tn < g(5,00
E~ p(-) [ - Eo log (ro)
sample p. via reparameterization (batched in practice)
for t O....,,
do
In+1 In + [f( m) + ou (zn nt]] öt, + a o 28t,cn 01.00
t L Log Pntiin (#n+1 (xn) - log P 7i.0: (In (2++1)
C C log C- ((NN)
O f- O + VeL
maximize (extended) ELBO
return optimized parameters OG
Finally, we substitute these into the left side of (25), that is,
V [)()) + _h (x)
(30)
v-1[(x)p (x)) + d²
2 (x) (- [f(x)] + (-22) V f(x)
(31)
>"(x)V f() + +x(())
[f(x)] ->'(x/V f(x)
(32)
02
2 (-22) f(x) (x)
-/'(z>V f(z) + f(z)p
- 02 [f(z)) (x) f(x)
(33)
0,
(34)
which yields the desired result.
B ADDITIONAL DETAILS FOR DIFFUSION-BASED SAMPLER
We additionally provide pseudocode in Algorithm 1 for a generic diffusion sampler with learnable prior ps- For clarity, we present an update step for a single sample. In practice, however, one would use mini-batches for these updates.
Pseudocode:
Special Cases of X-GMP: Consider the generic (extended) ELBO for X-GMP, that is,
p(aN)
P Pn-1jn (
LGMP (0,7,4)
Exa.xx n.e
log Ek= O&PP (IT)
log 00, P (In)))n)))
(35)
with p dx (zo) = [Zool/u Στ). We obtain the following special cases:
GVI (K N,N 0): For single Gaussian mixture component and zero diffusion steps. Equation (35) reduces to the marginal ELBO objective in Equation (9) for a Gaussian distribution, that
(ovi(0) = E B20 Spoppi log p(za)
(36)
Pd (2o)
16

---

Published as a conference paper at ICLR 2025
GMMVI (K N, = 0): Similarly, if have zero diffusion steps, but multiple Gaus- sian mixture components we obtain the marginal ELBO for Gaussian mixture models, i.c.
p(
(oMvI = =B momro
log
(37)
K Σ -k=1 Ci (zr))
Please note that there are more sophisticated methods train Gaussian mixture models for VI, see Arenz et al. (2018; 2022).
X-GP (K = N, > 0): For single mixture component and multiple diffusion steps, we obtain the objective for X-GP, i... for diffusion- -oodll with learned Gaussian prior, given by
log p(æN) Po (xr)
Pn-1/n ( (25m))
Co 7.0) E Bf o.s
+
log -0.6 Pnin-1 (In, (n-1))
(38)
X-GMP (K > I,N 0): Having multiple multiple mixture components K and diffusion steps N results the full X-GMP objective, as in Equation (35).
Forward and Backward Transitions. We provide further information about the diffusion-based sampling methods considered in this work in Table 3. Specifically, we provide expressions for the forward and backward transitions.
Time complexity. Diffusion-based samplers that Gaussian prior have time complexity of O(N). whereas Gaussian Mixture Priors (GMPs) incur time complexity of O(NK). The additional factor K arises from the need to compute the likelihood of the GMP each diffusion step. However, in practice, the evaluation of the likelihood of the GMP can be parallelized across its components, which substantially reduces the computational overhead. This parallelization allows for efficient implementation despite the increased theoretical complexity.
Memory consumption. When using the standard "discrete-then-oppiiizee approach to minimize the KL divergence in (11), which requires differentiation through the SDE, memory consumption scales linearly with both K (number of components) and N (number of diffusion steps). In contrast, methods like the stochastic adjoint approach for KL optimization (Li et al., 2020) achieve constant memory consumption, making them more suitable for scenarios with a large number of components or diffusion steps.
In our experiments, we opted for the former approach due to its simplicity. However, for tasks involving extensive components or steps, the stochastic adjoint method or similar approaches may be more practical. Additionally, constant memory consumption can also be achieved by using al- ternative loss functions such as the log-variance loss (Richter al., 2020; 2023) or moment-loss (Hartmann et al., 2019).
Initialization heuristic for iterative model refinement. As an alternative to (22) one could con- sider the following heuristic for selecting the mean of new component:
Mnew arg max E log p(IN)
2 log
i| (In-1/2a))
(39)
ZO
Pe (m))
-0, Pnin
(2n (
==1
However, empirically we found that choosing (22) results in increased sample diversity. We leave an in-depth exploration of different heuristics as future work.
C EXPERIMENTAL DETAILS
C.I BENCHMARKING TARGETS
This section introduces the target densities considered in our experiments. Please note that the major- ity of tasks are taken from the recent benchmark study from Blessing et al. (2024). For convenience, provide brief explanation of the target densities.
17

---

Published as a conference paper at ICLR 2025
Method
Fi 2i (Suuaala)))
Fa dls (1e(
DIS
N Emss (-oP log (m.) ou (x...)) ct, 2096t))
N (f alxo opP log (1))) (6060))6
MCD
N (Iaeslis -o o log 20
N (ia-alite - Vo ((())))) (x..))) 20 20)
CMCD
(rassline - [ M. loy (2r)) (1mm)) fr, 207807)
N (lreaalr Vo log (2())) on (f.. M) 8t, 20 207
DBS
N (ineellas - ((..))) +ou (A.. 81, 267))
N (os-ilas () -oo (o." 81, 20,20
Table 3: Comparison of different forward and backward transitions pi "0 and j 7xP, respectively, for diffusion- based sampling methods based f. ma Po u2 and v" as defined in the text.
Bayesian Logistic Regression: We evaluate a Bayesian logistic regression model four standard- ized binary classification datasets:
lonosphere (d 35, 351 (2t 36) pairs)
Sonar (d = 61, 208 (zi, Ut) pairs)
German Credit (d = 25, 1000 (æ,, yb) pairs)
Breast Cancer (d 31, 569 (Ir, vi) pairs)
The model assumes:
V(0,0.1)
1/4 Bernouli((sgmoodde x;)),
where features are standardized for linear logistic regression. Here, we perform inference over the parameters of the (linear) logistic regression model. In Blessing al. (2024), the authors used an uninformative prior for the parameters of the Bayesian logistic regression models for the Credit and Cancer tasks, which frequently caused numerical instabilities. To maintain the challenge of the tasks while ensuring stability, we opted for a Gaussian prior with zero mean and variance of o? = 100.
Random Effect Regression: We apply random effect regression to the Seeds dataset (d 26):
Gamma 0.01),
ao, 01 02 012 ~ (O, 10),
b, NO 1
= 1,..,,2, 21,
logits; = + a1 + 42th + a12 + bj.
Ti Binomial(ooiiss, N:).
with inference conducted over model parameters given observed data.
Time Series Models: For time series analysis, we use the Brownian model (d = 32):
Qinn Log 2),
Globa LogNormal(0,22) 2),
NO,, ainn).
i= 2..,....2 20, 34 N( Oman,
Wiz aobs),
i == 1,...,30,. 30,
with inference focusing on parameters Oiin, Onbss and latent states (2 30 Funnel: (d = 10), a funnel-shaped distribution defined by:
x (x)) =
exp(xx))),
with o3 9.
Fashion and Digits. MNIST variants (DIGITS) and Fashion MNIST (Fashion) datasets using NICE (Dinh et al., 2014) to train normalizing flows, with resolutions 14 and DIGITS and 28 28 for Fashion.
18

---

Published as a conference paper at ICLR 2025
C.2 DIFFUSION-BASE METHODS: DETAILS AND TUNING
General setting: All experiments are conducted using the Jax library (Bradbury et al., 2021). Our default experimental setup, unless specified otherwise, is as follows: We use batch size of 2000 (halved if memory-constrainele and train for 140k gradient steps to ensure approximate conver- gence. We use the Adam optimizer (Kingma & Ba, 2014), gradient clipping with a value of and learning rate scheduler that starts at 8 x 10 and uses a cosine decay starting at 60k gradient steps. We utilized 128 discretization steps and the Huler-Maruyama method for integration. The control functions u and 2 were parameterized as two-layer neural networks with 128 neurons. For DBS, we set the drift to f = - log n.
Unlike Zhang& & Chen (2021), we did not include the gradient of the target density in the network architecture. Inspired by Nichol & Dhariwal (2021), we applied a cosine-square scheduler for the discretization step size: ôt a cos² (2 N where [0, oo) - (0, oo) is learned for all methods. We enforced non-negativity of a via an ddement-wise softplus transformation. The diffusion coeffi- cient was set to 1 for all experiments. Furthermore, we set the initial a to 0.1 for all experiments except Brownian, where we set 0.01. We did not perform any hyperparameter tuning since most parameters are learned end-to-end.
Gaussian Priors (GP) and Gaussian Mixture Priors (GMP): We learn diagonal Gaussian priors and ensure positive definiteness with an element-wise softplus transformation. We use separate learning rate of 10-2 for all experiments to allow for quick adaptation of the Gaussian components. Furthermore, the mean was initialized at 0 and the initial covariance matrix was set the identity except for Fashion where we set the initial variance to 5 which roughly covers the support of the target. The individual components in the Gaussian mixture follow the setup of Gaussian priors. The mixture weights are uniformly initialized and fixed during training. If not otherwise specified, we use K = 10 mixture components for X-GMP.
Iterative Model Refinement (IMR): For IMR, we add a new component after 500 training itera- tions starting with a single component. The initial means were selected with the heuristic presented in Equation (22). The variance of the newly added components was set to be 1. The candidate sam- ple set was generated using the Metropolis Adjusted Langevin Algorithm (MALA) (Cheng et al., 2018). For that, used 2000 random samples from Gaussian with zero mean and variance 5, which roughly covers the support of the Fashion target. Please note that competing methods also use this prior knowledge for initialization of the prior, see Table 4. We use 128 steps steps, that is,
Fz+1 = Is + à V (z:)) + ā. / c~N(-10,, T)
(40)
with ā = 5 and an additional Metropolis adjustment step. Here, ā was chosen such that the final set of samples yields high target log- lieliielssoods log p(z). The final samples are used as candidate set. We note that this procedure brings the new components close to different modes in the target distribution and therefore facilitates exploration. Moreover, the computation of such candidate set is very cheap, i.e., the equivalent of single gradient step for e.g. MCD or CMCD.
C.3 COMPETING METHODS: DETAILS AND TUNING
The results for competing methods presented in this work are primarily drawn from Blessing et al. (2024), where hyperparameters were carefully optimized. For convenience, we repeat the details. Since our experimental setup differs for the Credit and Cancer tasks (detailed in Section C. we adhered to the tuning recommendations provided by Blessing et al. (2024). Details about hyperpa- rameters can be found in Table 4.
Sequential Monte Carlo (SMC) and Continual Repeated Annealed Flow Transport (CRAFT):
The Sequential Monte Carlo (SMC) approach was implemented with 2000 particles and 128 an- nealing steps, matching the number of sequential steps used in diffusion- -baee sampling methods. Resampling was performed with a threshold of 0.3, and one Hamiltonian Monte Carlo (HMC) step was applied per temperature, using 5 leapfrog steps. The HMC step size was tuned according to Table 4, with different step sizes based on the annealing parameter 3, Additionally, the scale of the initial proposal distribution was tuned. As CRAFT builds on the SMC framework, it used the same SMC specifications, incorporating diagonal affine flows (Papamakarios et al., 2021) as transition models.
19

---

Published as a conference paper at ICLR 2025
Parmmeters
Funnel
Credit
Sotur
Seeds
lonesghere
Intiial Scale
10.1. I. 10)
I
51
0.1
1
1
57
HMC stepsize 0.5)
(0.005, 0,001, 0.01, 0,05, U.1, 0.2)
0.001
0.2
0.01
0.01
0,001
0,05
0.2
0.2
[.)
HMC stepsize
(0.05, 0,001, 0.01, 0,05, 0.2)
0.2
0.005
0,055
0.001
0.05
0.05
CRAFT
1, 10) [10 10"6,0 x .1001
Learning Rase
10
10
6.x 10
5 10
10"
10
10"5
10"
FAB
Initidl Scale
[O.1.
Leurning
[10
.10
Table 4: Hyperparameter selection for all different sampling algorithms. The "Grid" column indicates the values which performed grid search. The values in the column which are marked with experiment names indicate which values were chosen for the reported results. The values parameters indicated with are set by using prior knowledge about the task.
Flow Annealed Importance Sampling Bootstrap (FAB): Automatic step size tuning for the SMC sampler was applied on top of the normalizing flow (Papamakarios et al., 2021). The flow architec- ture utilized RealNVP (Dinh et al., 2016), with an 8-layer MLP serving as the conditioner. FAB's replay buffer was employed to accelerate computations. The learning rate and base distribution scale were adjusted for target specificity as outlined in Table 4. A batch size of 2000 was used, and FAB was trained until reaching approximate convergence, which sufficient achieve approximate convergence.
C.4 EVALUATION
Evaluation protocol and model selection We follow the evaluation protocol of prior work (Bless- ing al., 2024) and evaluate all performance criteria 100 times during training, using 2000 samples for each evaluation. To smooth out short-term fluctuations and obtain more robust results within single run, we apply running average with window of 5 evaluations. We conduct each experiment using four different random seeds and average the best results of each run.
Performance Criteria: In order to define the performance criteria, first define the unnormal- ized (extended) importance weights w. that is,
p(æN) 1. Py Pn- 1/m (In- 12xx) Po (200 II n= Pnin-1 ((m((n--1)
(41)
0800
We consider the following following performance criteria:
Evidence lower bound (ELBO): We compute the (extended) ELBO as
ELBO
[log w]
log io
(42)
Estimation error A log Z: When having access to the ground truth normalization constant log Z, we can compute the estimation error A log Z = |log 2 -loo Z using an importance weighted estimate, that is,
1 m
log 2 = log E Exg. Exg. W
log m
60
(43)
Effective sample size (ESS): Moreover, we compute the (normalized) ESS as
( 1 ()())2 2
ESS
(44)
171 Σ (17())²
Sinkhorn distance: We estimate the Sinkhorn distance W (Cuturi, 2013), i.e., an entropy regularized optimal transport distance between a set of samples from the model and target using the Jax ott library (Cuturi et al., 2022). Note that computing W.? requires samples from the target density which typically not available for real-world target densities.
20

---

Published as a conference paper at ICLR 2025
DIS GMP
DIS-GMP
CMCD- GMMP
K
REL.
ABS. [4]
REL. [%]]
0.103
0.128
24.27
10
0.155
Table 5: Wallclock time of DIS-GMP and CMCD-GMP for the Fashion target for N 128. Here. the number components in den Gaussian mixture, 'abs.' denotes absolute time per gradient step seconds, and 'rel.' denotes the relative increase in percent compared to K =
Louosphere
Browsian
0:33
0.25
11:06
0.I9
0.03
0.04
0.13
0.00
0.07
0:11
0.15
0:05
0.00
0.07
0.1
12
0100
0.04
0:04
0.06
0.08
0.11
12
0.:00
0.04
0:055
0.00
0.000
0:14
01:0
0.05
01:05
0:05
0.001
0.14
11.8
11.00
0.06
16
32
64
128
32
64
125
16
32
64
122
Sceds
Credis
0.05
0:05
0.08
6:15
0.21
0.05
0:01
C.05
D100
6.05
0.13
A
0.03
8.04
0.05
1L.6
0.08
0.1.1
0.03
01.4
0.04
01.00
0.07
0.200
0.05
0.05
0.00
0.07
0.11
0.05
0.01
0.05
0.06
0.06
0.08
64
125
16
52
64
128
Figure 6: Wallelock time per gradient step of DIS-GMP for various benchmark problems. Here, N denotes the number discretization steps and the number components Gaussian mixture.
Entropic mode coverage (EMC): EMC evaluates the mode coverage of sampler by leveraging prior knowledge of the target density. It holds that EMC E [0, 1] where EMC = 1 indicates that the model achieves uniform coverage over all modes whereas EMC indicates that the model only produces samples from single mode. Please note that EMC does not provide any information about the sample quality. For further details, we refer the interested reader to Blessing et al. (2024).
D FURTHER NUMERICAL RESULTS
Here, we provide further numerical results.
Wallclock time We further report the wallclock time per gradient step for DIS for different number of diffusion steps N and mixture components K. The results are shown in Figure 6. For
N < 64, the Gaussian mixture prior barely influences the wallclock time where using K 10 components roughly adds a 20 percent increase. Considering the performance improvements this is a good trade-off. For = 128, Using hood of the prior has to be evaluated in every diffusion step. However, since most diffusion-bassd methods apart from DIS additionally require evaluating the target density at every step, the relative costs of using GMPs reduce if the target is more expensive evaluate. We empirically validated this by additionally including a comparison between the wallclock time for DIS and CMCD on the Fashion target in Table 5. In this setting, the relative cost added by the GMP is minor.
10 roughly results 50 percent increase as the likeli-
21

---

Published as a conference paper at ICLR 2025
Hlnesttn
Sceth
Crefitt
1..00
0:28
0.77
-TUA)
-7.47
1231
1122
11345
101.0
32.0
0.14
80
007
8.9
0.99
-3.22
75.00
-7499
-7.00
-1.40
3014
co.r
121-3
112.0
11211
11:M1
0:96
6:53
0.30
0.55
878
0.30
5935
73.58
33.01
2000
7351
46.10 19221
Figure 7: Evidence Lower Bound (ELBO) of DIS-GMP for various real worldw benchmark problems, averaged across four seeds. Here, N denotes the number of discretization steps and K the number of components in den Gaussian mixture.
ELBO
E358
A k:
AX21
+033
4U35
-100
1.75
0.20
1.85
00:0
5182
B31
4122
355
13:11
3-111
1:4
140
0.20
061
4.43
0.30
0.2
6111
3.9
6.3
8:6
0.20
0.96
0:5
3.73
0.12
647
0.5
29:27
142
1240
61
08
61
185
Figure 8: Various performance criteria of DIS-GMP for the Funnel target, averaged across four seeds. Here, N denotes the number of discretization steps and K the number of components in den Gaussian mixture.
Additional results for DIS-GMP We present further details regarding the ablation study from Section 6. Specifically, we report ELBO values the real-world benchmark problems Figure and various metrics for the Funnel target in Figure 8. The results are consistent with the results in Figure 5, where the performance improves with higher number of mixture components and diffusion steps N.
DIS-GMP on Digits target We additionally investigate the performance of DIS-GMP on the syn- thetic digits target. The results reported in Figure 9. Here, observe that the ELBO get looser when using more mixture components K, while A log Z stays roughly constant. However, the sam- ple diversity improves significantly shown quantitatively from the Sinkhorn distance W and qualitatively in the Figure on the right-hand side.
Digits (d 196)
K
ELBO
A llog Z 4
W?
12.090
5..491
297.566200340
10
X 201
- -12.303
4.419
183.241 18.776
10
-13.820+0.58
4.658=0.260
164.827
20
-15.41340.31 5.66340.085
151.006::0.644
K: Num minturs compusatia
Figure 9: Left side: Results for Digits target, averaged across four seeds using DIS GMS-GMP+MMR. Evaluation criteria include evidence lower bound ELBO, importance-weighted errors for estimating the log-normalizing constant A log Z, and Sinkhorn distance W,". The best results are highlighted in bold. Arrows (t. 4) indicate whether higher or lower values are preferable, respectively. Right side: Visualization of the d = 14x 14 = 196 dimensional Digits samples for a different number of mixture components K.
22

---

Published as a conference paper at ICLR 2025
DIS
DIS-GP
DIS-GMP
DIS-MMP+IMI
Figure 10: Visualization two-dimensional Gaussian mixture target density for different variants DIS
with N = 128 diffusion steps and K = 10 components for GMP versions. Colored ellipses and circles denote standard deviations and means of the Gaussian components, respectively. Red dots illustrate samples of the learned model.
METROD
V lo
CREDIT
SEEDS
CANCER
BROWNIAN
TONOSPHERE
SONAR
DIS-GP
X
585, 247
-33544020.00
15.006
0.58820.015
-11.8877
-109 280 -109:280ga.co
DIS-GP
5992 20222.00
73. 73.197--0000
96. 180=10004
NIA
111.977
-109. 173
DIS CMP
X
585. o..00
-33.492:0.03
:4.161
0.885
oos
111.8110..002
-009 199.115:0... 00000
DIS-GMP
586.817.10.90
73.475.10.m 002
84.732.0-46
NIA
112.108
109.248_o.com
Table Evidence lower bound (ELBO) values for various real-world benchmark problems, averaged across four secds. Hcre, V log v indicates if the model architecture uses target score as described in Equation (46). The best overall results are highlighted in bold, with category- best results underlined. Blue and green shading indicate that the method uses learned Gaussian (GP) and Gaussian mixture priors (GMP), respectively.
Ablation on Gaussian Mixture Target We additionally experiment with using two-dimensional Gaussian mixture model (GMM) as the target density. The GMM has ten components where the means uniformly sampled in [-12, 12] and the covariance matrices distribution. In addition to the performance criteria outlined in Section 6, we quantify the variation of the dynamics time using the spectral norm of the Jacobian of the learned control, i.e.,
sampled from Wishart
[ Dou (x, t)
== EzorTwa
dt
(45)
dr
For DIS we initialized the prior with standard deviation of 12 such that the prior covers the support of the target. For DIS-GMP, we use K = 10 components that are initialized with a standard devia- tion of 1. We report qualitative and qualitative results in Figure 10 and Figure 11. We find that DIS without learned prior and sufficiently large prior support is able to cover all modes as indicated by EMC 1. While the sample quality is similar between DIS and GMP counterparts (see W plain DIS needs significantly more diffusion steps in order to achieve similar EI BO/ESS values compared the GMP counterparts which achieve good performance with few 8 diffusion steps. The quirement of plain DIS for more discretization steps is additionally reflected in the variation of the dynamics over time S. Lastly, the GP version is not able cover all modes due the mode-seeking nature of the reverse KL as indicated by the EMC and W values.
Ablation: Influence of the control architecture We further evaluate the performance using the architecture by Zhang & Chen (2021) which additionally incorporates the score of the target, i.c. V log n, into the architecture via
n' (x. t) (x. + f2" (L) log (stopp gradient(z:))-
(46)
where /i and f2 are parameterized function approximatior with parameters 0. Zhang & Chen (2021); Vargas et al. (2023a) found that detaching, that is, using a atop-gradient operator on x yields superior results due to the simplification of the computational graph. We adopt this change and report the results in Table We find that using the score of the target leads, in the majority of experiments, slightly worse results.
Ablation: Bullback-Leibler vs. Log- Variance divergence We further compare the KL diver- gence to the log-variance divergence introduced in Richter et al. (2020) and later extended to diffu-
23

---

Published as a conference paper at ICLR 2025
DIS
DIS-GP
DIS-GMP
DIS-GMP+ +MM
120
100
11,30
1385
11.0
60
0.4
40-
70-
0.2
222
16
0.3
0.9
0.2
DNCC
0.8
Figure 11: Results for the two-dimensional Gaussian mixture target, averaged across four seeds and reported across different numbers of difffusion steps N for different variants of DIS. Evaluation criteria include evidence lower bound ELBO. importance weighted enrors for estimating the log-normalizing constant A log Z, and Sinkhorn distance W2² 7 entropic mode coverage EMC, and the uneeintegrated spectral norm of the control (see Equation (45)).
METHOD
CANCER
BROWNIAN
TONOSPHERE
SONAR
DIS
KL
589 777
-4.4000...oor
36.5022...10r
3.503
112.525
-110.153-..ozz
DIS
LV
-5170 845es
-1..654-0.0222
98.3779=1.40
5.682=0.906
112. 02..609
-110.622e0.07
DIS-GP
KL.
-885 cm
-3.5440_0.00
-5.006 1996
0.588.0.01a
109.240 02d
DIS-GP
LV
5161.451 2.02
73.7033
549.071 : 400.900
0.729 00.24
111.83920.00
1199.498 20.0
DIS-GMP
KL
-585, 200
73 -90
-81.061-2.117
0.885-0.00
- 111.811=0.0c
109. 1109.157 157 60.00.000
DIS-GMP
I.
51552.228 cot
73.777_0.000
-6..45620..57
0.72220.006
109.443 o00
Table 7: Evidence lower bound (ELBO) values for various world benchmark problems, averaged across four seeds. Here, 'Div.' indicates if the model is trained using the Kullback-Leible (KL) or log- lorvvariance (LV)
divergence. The best overall results highlighted in bold, with category-speciff best results underlined. Blue and green shading indicate that the method learned Gaussian (GP) and Gaussian mixture priors (GMP), respectively.
sion models in Richter et al. (2023). The log-variance divergence is defined as
5 7.0 (xo:N) L(,, Φ) = V No log 15 5000 (To.N)
(47)
where R describes a reference process, e.g. Equation (2a) where d is replaced with an arbitrary control. In practice, one typically uses the generative process 80 with an additional stop gradi- ent operator on the parameters (Richter et al., 2023). Not computing the expectations with respect to samples from the generative process significantly reduces memory consumption and does not require the prior distribution to be amendable to the reparameterization trick. The results are re- ported in Table and follow the same experimental setting as outlined the main part of the paper. We find that the KL divergence typically performs better than the log-veriance divergence. Most significantly, the log-variance divergence seems to be numerically unstable for the Credit target.
Comparison to long-run Sequential Monte Carlo We additionally compare diffusion samplers with learned Gaussian mixture prior to a Sequential Monte Carlo with a high number of discretiza- tion steps N. The results shown in Table 8 and Table 9. While long-run SMC significantly
24

---

Published as a conference paper at ICLR 2025
Мктноо
N
CREDIT
CANCER
BROWNIAN
SONAR
MCD-GMP
128
01a
73
88.562
0.90322000
111.82720007
100.197.
CMCD-GMP
128
586. 162ee ouz
-33-42920.0
-8. 23.400200.00
1087
111.
111.6822000
-08. 63-00
DIS-GMP
128
585 223eu ove
-73. +30
-84.061=221 #2.117
-11.811+
-109.157 119.157=+0.00 20.000
DBS-GN BBS-GMP
128
585. -.000
73. -1.000000
78. 78.160+0o.oz
1.118-0.002
111.657+0.. +57000002
108. e.00...00.
128
(28.403.4.114 146
74.6999o.150
184.0399...ot 10.013
1.844...22
114.751 2o.23s
115,355 1.177
256
708.185, 115.188 385
73.97226-04
140.757
0.380
113. 10
100. 822
0.414
112.354
SMC
1021
11000011::2..m
-73.5825:..0x
128. 808
0.786
111. 101.902
-108 291
2018
-72.6022+1..22 222
-7..496e0.017
128.376=a.50
0.992
111. 786
108. 261-2611+
4096
121.950+a..15
1.088
10..6622...m
110.8770 71667766...22 20.
Table 8: Evidence lower bound (ELBO) values for various real-world benchmark problems and different numbers of discretization steps N, averaged across four seeds. The with category-spectoo best results underlined. green shading indicate that the method Gaussian mixture priors (GMP).
hest overall results are highlighted in hold,
FASHION (d 784)
METHOD
ELBO f
log 24
W +
EMC f
DIS-GMP4IMR
128
662.482-2 282222.722
27 97.645+3.111
513.776+13.B8
0.780=0.08p
128
- -12181.999 193221122....
11747.5 17747.582215 133.292
259
0.026 027
256
-10095. 076 1076.10 723
19901113::107
6018.423 195.:199....0 144
0.191
512
-9340.233 222 804
9254.499 108...00.222
5821 18921422__.n0 492
0.14110.:10
SMC
1024
24229.5577742.2222
9190.558=742.688
5610.51 -223. 885
0.075##0.22
2048
-8472.600 288
8454.062.2281 .475
5718.030 909
0.2571100088
4096
1199.465+1144
8390.302
1583.099-179.44
0.102++0.10
Table 9: Results Fashion target, averaged across four seeds reported across different numbers of discretization steps N. Evaluation criteria include evidence lower bound ELBO, seportance-weigheee errors for estimating the log-normalizing constant A log Z, and Sinkhom distance W2² and entropic mode coverage EMC. The best overall results are highlighted in bold, with category-specific best results underlined. Arrows (t. J) indicate whether higher lower values are preferable, respectively. Orange shading indicates that the method uses iterative model refinement (IMR).
METHOD
FUNNEL
SEEDS
BROWNIAN
IONOSPHERE
SONAR
DIS DIS-GMP
2.993+t0.022
4. 6888+0.055
6.266+t..29
4.394+0.065
4.840+0.0s1
1.898
2.367=0.000
2.445=0.004
2.736++0.004
3.861+0.030
Variability in the dynamics of the learned model via time-integrated spectral of the control x 102 (see Equation (45)) for various benchmark problems. Both DIS and DIS-GMP use = 128 diffusion steps. Here, DIS-GMP uses K = 10 components. Lower values indicate lower variability in the dynamics of the learned model.
Table 10:
increases ELBO values, GMP-based diffusion sampler yield superior results in most experiments. Moreover, the results on the Fashion target indicate that more discretization steps yield better ELBO, log Z and Sinkorn distances, but are not able to prevent mode collapse as indicated by the low EMC values.
Ablation: Variation of dynamics We additionally compare the variability in the dynamics of the learned model between DIS and DIS-GM DIS-GMP via time- inteeinteeetted spectral norm of the control S (see Equation (45)). The results are shown in Table 10 and show that DIS-GMP indeed has less variation in the dynamics. These findings are also in line with those in Figure 3 and Table 2 where DIS-GMP has significantly higher ELBO values compared to DIS without learned prior.
25

---

Published as a conference paper at ICLR 2025
E LATTICE o² THEORY
We apply our method to simulate statistical lattice field theory near and beyond the phase transi- tion. This phase transition marks the progression of the lattice from disordered to semi-ordered and ultimately to a fully ordered state, where neighboring sites exhibit strong correlations in sign and magnitude.
We study the lattice ot theory in D = 2 spacetime dimensions (distinct from the problem's 's di- mensionality as described below). The random variables in this setting are field configurations @ E RLXL where I. represents the lattice extent in space and time. The density of these configura- tions defined as
e-((
(())
Z
where the potential (Ø) is given by:
( = 2x ΣΣ SerZra + (1-2A) Σ 22
+ E 1
(48)
Here, the summation over runs over all lattice sites, and the summation over μ considers the neighbors of each site. The parameters À and & are referred to as the bare coupling constant and the hopping parameter, respectively. Following Nicoli et al. (2021),
set À = 0.022, identifying the critical threshold of the theory (the transition from ordered to disordered states) at K 2 0.3. Near this threshold, sampling becomes increasingly challenging due to the multimodality of the density, with modes becoming more separated for larger values of K.
We conduct experiments for E {0.2, 0.5} across various problem dimensions d L L. The methods compared include DIS, DIS-GP, and DIS-GMP, each with N = 128 diffusion steps, as well as long-run SMC sampler with N 4096. For all methods, the initial support is set to 5, approximately covering the target's support for all tested values of K. The tuned parameters of the HMC kernel for the SMC sampler are detailed in Table 4, while additional parameter settings are provided in Appendix C.2. Note that DIS (and its extensions) do not undergo hyperparameter tuning due to their end-to-end learning framework.
To compare the different methods, we utilize the negative variational free energy of the system, defined as:
1
-J
1 log Z
1 ELBO.
(49)
This bound follows from the inequality log Z > ELBO discussed in Section 3.2 and provides a means of comparison between sampling methods. However, as the ELBO (and thus F) is not sensitive mode collapse (Blessing et al., 2024), and since samples from the target distribution unavailable, we also qualitatively assess the methods by visualizing the (normalized) histogram of the average magnetization M (Ø) Σx Φx across lattice configurations ó.
Quantitative results are presented in Table 11, with qualitative findings illustrated in Figure 12. The results indicate that learning the prior (i.e. DIS-GP/DIS-GMPP significantly improves free energy estimates compared to DIS without a learned prior. Moreover, Figure 12 demonstrates that DIS- GMP avoids mode collapse while achieving comparable better free energy estimates than both DIS-GP and the long-run SMC sampler in the majority of settings. While SMC captures multi- modality at the phase transition (k = 0.3), it struggles with the multimodality in the fully ordered phase (N = 0.5). By contrast, DIS and DIS-GP are prone to mode collapse. Lastly, the performance of DIS degrades significantly with increasing problem dimension d, which mitigated when using learned Gaussian or Gaussian mixture prior.
26

---

Published as a conference paper at ICLR 2025
METROD
= 16
64
100
144
106
d = 206
DIS
0.0186s9.000
0.6100-..coos
0.61 45:zoor
0.5097=0.000
0.5719ec.om 2015
DIS-GP
0.6176
0.6158g0.om
0.2
DIS-GMP + IMR
0.6278
2.6292:0.uom
6231=v.uoe
0.6 0.6660200.00
0.6156
SMC
0.6167 o..00
0.61 7üge non6.
0.618610 2003
0.61642nnor
0.608720010
0.614 42g4 cosn
DIS
10633e
8277
1.007740.004
0.963
0.890S w0s
0.3
DIS-GMP IMR
1.0940:0.com
0-966e unss
100001i0.cmm
LD40620000
1.038020.c0u
1.004750..
SMC
00844880001
1,05144 4. 0007
0448__0.000
1.0437=0.000
10333g0004
1.0389 2. oons
DIS
12. 2.000
12 2007 00.0023
9.2810aa.seu
11.7179.0.0700
9.4392 1.000...04
10 10.85466s.omm 0006
0.5
DIS-GMP + IMR
12.3167.0
12. 28.2800
12.2761 20.611..000 0000
12:2730
12.2722:0.com
12.20884n
SMC
12. 12:0009
12 a007
12.2070sn..0m*
12.209020c0m
12.2116anuzu
12.2007as.m2
Table 1I: Lower bound values for negative variational free energy -F as defined in (49) for lattice d² theory prohlem with different values for the space-time extend vd T. averaged (i.c. the highest) overall results are highlighted in hold each configuration hopping parameter space-time extend.
seeds. The best
DIS-GP
DIS-GMP
SMC
Average Magnetisation
Figure 12: Normalized histogram the average magnetization ( Σ, dz for 2000 samples K LXL and space-time extend
14 for DIS. DIS-GP, DIS-GMP and long-run SMC different values hopping parameter s. The plots are generated using the same random seed 0.
27


Reviews:
[{"review_id": "3gthXZHuMK", "reviewer": "Reviewer_JR3v", "summary": "This paper introduces a novel approach to diffusion-based sampling that incorporates end-to-end learnable Gaussian Mixture Priors (GMP), offering a more expressive alternative to conventional Gaussian priors. The flexibility of GMP mitigates mode collapse, enhances exploration capabilities, and reduces the need for extensive diffusion steps by alleviating non-linear dynamics. The authors also propose an Iterative Model Refinement (IMR) strategy, progressively increasing model complexity to optimize performance on high-dimensional, multi-modal targets. Extensive experiments demonstrate that using GMP in place of Gaussian priors in existing diffusion-based sampling methods achieves improved sampling performance and outperforms previous methods across synthetic and real-world datasets.", "questions": "- In Section 5, it seems that the positioning of challenges C1 and C2 may have been reversed. Specifically, in lines 308–310, reducing the complexity of dynamics to minimize the required diffusion steps seems more directly related to the issue identified in the Introduction as C2 (lines 51–53). Conversely, enhancing the exploration capability of diffusion-based sampling methods to cover diverse regions seems more relevant to addressing C1.\n- While Figure 6 presents wall-clock time for relatively low-dimensional cases, it would be helpful to understand how computational demands scale in higher-dimensional scenarios. Could you provide an analysis of the computational cost associated with varying the number of mixture components $K$ and diffusion steps $N$ for a more formal comparison? (e.g. Big-O notation)\n- How does estimated memory consumption vary with changes in $K$ and $N$?\n- In Figure 4, DIS-GMP performs worse than DIS-GP across all metrics of sampling quality. This outcome seems contrary to what the paper suggests; how might this be explained? Furthermore, while DIS-GMP + IMR with progressively learned components shows better Sinkhorn distance and qualitative sampling results, its ELBO and $\\triangle$ log Z performance is lower than that of DIS-GP and DIS-GMP. This observation suggests that ELBO may have limitations in evaluating sampling performance for multi-modal target distributions (as also seen in Figure 9).\n- Table 2 presents only ELBO results; does the Sinkhorn distance follow a similar trend to what was observed in the Funnel case?\n- The IMR approach uses a criterion of adding components every 500 training iterations. Are there any ablation study results regarding this criterion for IMR?\n\nMinor:\n- In Equation 31 of Appendix A.1, the term $(-\\frac{2}{\\sigma^2})^2$ seems to have a sign inconsistency with Equation 32, where the third term should perhaps be positive.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-structured, first identifying the limitations of Gaussian priors in diffusion-based sampling and then logically progressing to propose end-to-end learnable Gaussian Mixture Priors (GMP) as a solution. Each step in this argument is clearly justified, supported by sufficient background information and helpful illustrations, making it easy for readers to follow.\n\n- The experimental design effectively supports the theoretical claims, providing empirical evidence of the advantages offered by GMP over traditional Gaussian priors. Additionally, the comprehensive review of related work and background offers readers a useful overview of recent developments in this field.\n\n- The proposed method enables end-to-end learning of the prior distribution, allowing existing diffusion-based sampling methods, which previously relied on fixed naive priors, to benefit from a more optimized starting prior. This adaptability enhances sampling performance and broadens the applicability of the approach across various diffusion-based sampling methods.", "weaknesses": "- Although GMP can achieve competitive performance with a single Gaussian prior using fewer steps, high-dimensional data often require more complex dynamics and an increased number of components to handle multiple modes effectively. As both the number of components $K$ and steps $N$ grow, the computational cost increases substantially, raising concerns about the scalability of this approach for even higher-dimensional applications in the future.\n- As acknowledged in the paper, there is no clear criterion for determining the appropriate number of mixture components needed for optimal performance. Furthermore, IMR relies on heuristic rules, such as a fixed number of iterations, to add new components, introducing additional hyperparameter tuning requirements for model optimization.\n- Some minor inconsistencies and certain claims in the paper are not fully supported by the experimental results, which would benefit from further clarification. For instance, additional explanation may be needed for the DIS results in Figure 4.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oeaxoV1FO8", "reviewer": "Reviewer_d8iD", "summary": "The paper proposes using learned Gaussian mixture priors in diffusion samplers instead of the standard fixed, unimodal Gaussian. The authors experimentally demonstrate the advantages of having a multi-modal prior; better exploration, fewer diffusion steps, and avoiding mode collapse. To effectively train the learned Gaussian mixture prior model, the authors introduce an iterative refinement approach that gradually adds new modes to the prior distribution.", "questions": "- For the exploration and mode collapse issues discussed, is it not possible to transform the target data distribution such that its support is covered by a single unit Gaussian, given that you use the knowledge of the support of the true distribution when adding new modes? How much worse is the base DIS compared to the learned and learned+mixture variants under this experiment setting?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- **Novelty**: The idea of learning a mixture prior in a diffusion process is underexplored and could have significant impact on future works. To the best of my knowledge, there have been few papers that explore a learned prior in diffusion models. Additionally, the idea of gradually adding complexity to the prior is very interesting.\n- **Evaluation**: The authors provide adequate comparisons between the proposed method and existing approaches that demonstrate the overall improvements when using the proposed model. The results show that their method better captures the different modes of the distributions learned.\n- **Presentation**: The writing is succinct and helps the reader understand the necessary preliminaries before diving into the details of the proposed method. Even with limited prior knowledge of the different diffusion samplers out there, the paper is easy to follow.", "weaknesses": "- Some more ablations could help with understanding the specific advantages of the proposed method better. One of the issues of the unimodal Gaussian prior that is mentioned is the increased complexity in the learned dynamics. In some examples of Fig. 5 the EES only marginally improves with more modes, making the number of diffusion steps the more impactful hyperparameter. This hints towards the dynamics not being simpler with more modes, as we would expect fewer diffusion steps to also yield good EES in those cases. A more in-depth experiment that demonstrates the differences in dynamics could help ground the overall argument for a mixture prior.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PzdwJiyvb3", "reviewer": "Reviewer_PAqT", "summary": "This paper proposes an end-to-end learnable Gaussian mixture priors (GMPs) to address challenges in diffusion-based sampling methods including poor exploration, non-linear dynamics and mode collapse. Generally speaking, the paper contains some interesting ideas, and some experimental results are provided to demonstrate the performance improvement of the proposed method.", "questions": "1. How can the proposed iterative model refinement (IMR) strategy be further improved to optimize the selection criteria for adding new components and dynamically adjust the number of components during training?\n\n2. As shown in the left side of Figure 4, the proposed IMR strategy seems not work well in all metrics (DIS-GMP, DIS-GMP-IMR), how could this be explained?\n\n3. As introduced in Section 5, the IMR strategy requires to optimize parameters until a predefined criterion is met, how to choose the criterion and how does the criterion affect the training result?\n\n4. The mean of a new added Gaussian component is chosen according to (22), how about the choice of variance?\n\n5. What are the potential limitations or drawbacks of using Gaussian mixture priors compared to other types of learnable priors, and how could these be addressed?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper presents an approach for improving diffusion-based sampling techniques by introducing end-to-end learnable Gaussian Mixture Priors (GMPs). The authors identify three key challenges in diffusion models: lack of exploration capabilities (C1), non-linear dynamics requiring many diffusion steps (C2), and mode collapse due to reverse KL minimization (C3). To address these challenges, the paper proposes GMPs, which offer improved control over exploration, adaptability to target support, and increased expressiveness to counteract mode collapse. Experimental results are provided to demonstrate performance improvement of the proposed methods.", "weaknesses": "The main concern of the paper is that the considered numerical experiments are oversimplified. The authors only considered logistic regression models and the fashion target on the MNIST fashion dataset in experiments. It is well known that diffusion models are widely used for generating samples with complicated priors, such as high-quality real-world images. Yet, from the numeral results provided in the paper, the reviewer cannot see any potential of the proposed method in this respect. As such, it is not convincing to claim that the proposed methodology can indeed solve the problems (C1)-(C3) of the existing diffusion models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ittFqEf4C7", "reviewer": "Reviewer_3sKi", "summary": "In this paper, the authors propose an innovative sampling approach using the framework of controlled diffusions. They address current limitations in diffusion samplers, such as poor exploration, mode collapse, and numerical instabilities, which hinder model effectiveness. To counter these issues, they suggest parametrizing the prior of the backward diffusion process as a Gaussian Mixture Model (GMM). The authors detail how this adjustment can be applied to both denoising diffusions and Langevin diffusions. They introduce an iterative learning procedure that jointly optimizes the control networks and the GMM prior parameters. Through extensive experiments, the authors demonstrate that their approach outperforms traditional diffusion samplers across a wide range of tasks.", "questions": "- Why not using the log-variance objective which as proven to be very effective ? [2] (Note that this objective wouldn't require the reparameterization trick on the prior)\n- Could you provide an ablation study in a controlled environment (typically sampling Mixture of Gaussian) to actually challenge the procedure and assess that the claims were addressed ?\n- How does your algorithm scale with the number of modes $K$ as GMM notoriously showcase numerical instabilities in this scenario ?\n- Fig 4 and Fig 9 clearly show that the model mode collapsed as most of the generated images are almost identical. How do you explain this behavior ?\n- According to Fig 6, the Funnel sampling procedure (for DIS) took over 5 hours. Is it actually competitive against MCMC samplers like MALA or HMC under the same computational budget ?\n- It seems like SMC and CRAFT have been unfairly implemented. Indeed, as those algorithms use a different annealing scheme, there are no reasons to keep the same number of intermediate distributions. Moreover, as they are cheaper, this number should be increased. Could you re-run the experiments with increased budget ?\n- Could you compare X-GMP against X-GP but adding the target informed parametrization of the controlled networks ? It seems like it could make a big difference.\n- The problem of overparameterisation of the GMM (when the GMM as more mode than the target distribution) has two different interpretations between L479-483 and Fig 9. Could you clarify it ?\n- Could you compute mode-collapse aware metrics such as the EUBO (and derivations) or the EMC from [1] ?\n\n[1] Blessing, D., Jia, X., Esslinger, J., Vargas, F., & Neumann, G. (2024). Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling. In Proceedings of the 41st International Conference on Machine Learning (pp. 4205–4229). PMLR.\n\n[2] Nüsken, N., Richter, L. Solving high-dimensional Hamilton–Jacobi–Bellman PDEs using neural networks: perspectives from the theory of controlled diffusions and measures on path space. Partial Differ. Equ. Appl. 2, 48 (2021). https://doi.org/10.1007/s42985-021-00102-x", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is well written\n- The idea of parametrizing the prior seems novel", "weaknesses": "- Lack of multimodal experiments (while it was the main motivation according to Fig. 1)\n- Competing algorithms seems to have been wrongly implemented as the control networks lack the $\\nabla \\log \\pi$ information (according to L927-928)\n- Sec 6 only considers metrics known to be unsuited for multimodal distributions (and especially to detect mode collapse C3) as per [1] (cited by the authors)\n- The computational procedure is very heavy", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)