Distributional Diffusion Models with Scoring Rules
Authors: Valentin De Bortoli
Alexandre Galashov
J Swaroop Guntupalli
Guangyao Zhou
Kevin Patrick Murphy
Arthur Gretton
Arnaud Doucet
Venue: ICML 2025 poster (2025)
Keywords: Diffusion models
energy distance
maximum mean discrepancy
scoring rules
accelerated sampling of diffusion models
Primary area: general_machine_learning
TL;DR: New losses for diffusion based on scoring rules; Better quality sampling in few steps regime

Abstract: Diffusion models generate high-quality synthetic data. They operate by defining a continuous-time forward process which gradually adds Gaussian noise to  data until fully corrupted. The corresponding reverse process progressively ``denoises" a Gaussian sample into a sample from the data distribution. However, generating high-quality outputs requires many discretization steps to obtain a faithful approximation of the reverse process. This is expensive and has motivated the development of many acceleration methods. We propose to speed up sample generation by learning the posterior distribution of  clean data samples  given their noisy versions, instead of only the mean of this distribution. This allows us to sample from the probability transitions of the reverse process on a coarse time scale, significantly accelerating inference with minimal degradation of the quality of the output. This is  accomplished by replacing the standard regression loss used to estimate conditional means with a  scoring rule. We validate our method on  image and robot trajectory generation, where we  consistently outperform standard diffusion models at few discretization steps.

Distributional Diffusion Models with Scoring Rules
Valentin De Bortoli" 1 Alexandre Galashov 112 J. Swaroop Guntupalli' Guangyao Zhou'
Kevin Murphy
Arthur Gretton 12 Arnaud Doucet
Abstract
ples, we sample from the time-reversed diffusion process, leveraging the learned denoiser. Despite their impressive capabilities, diffusion models often require a large number of steps faithfully approximate the time-reversal so as to obtain high-fidelity samples. To overcome this limita- tion, different solutions have been explored. One approach focuses on the development of improved numerical inte- grators (c.g. (Karras et al., 2022; Lu et al., 2022: Zheng et al., 2023)). Another common strategy involves the use of distillation techniques (e.g. (Luhman & Luhman, 2021; Song al., 2023; Luo, 2023; Salimans et al., 2024)), which requires training a smaller, more computationally efficient model to emulate the behavior of larger, pre-trained diffu- sion model. Finally, parallel simulation methods have also been explored (c.g. (Shih et al., 2023; Chen et al., 2024a)) and use more memory to avoid slow sequential processing.
Diffusion models generate high-quality synthetic data. They operate by defining continuous-time forward process which gradually adds Gaussian noise data until fully corrupted. The corre- sponding reverse process progressively "denoises" a Gaussian sample into a sample from the data dis- tribution. However, generating high-quality out- puts requires many discretization steps to obtain a faithful approximation of the reverse process. This expensive and has motivated the develop- ment of many acceleration methods. We propose to speed up sample generation by learning the posterior distribution of clean data samples given their noisy versions, instead of only the mean of this distribution. This allows us to sample from the probability transitions of the reverse process on a coarse time scale, significantly accelerating inference with minimal degradation of the quality of the output. This is accomplished by replacing the standard regression loss used to estimate con- ditional means with scoring rule. We validate our method on image and robot trajectory genera- tion, where we consistently outperform standard di diffusion models few discretization steps.
We depart from these earlier approaches in simply propos- ing to sample from the generative process coarser time scale. However, naively using denoisers obtained from the training of diffusion models significantly degrades the qual- ity of the outputs, as these denoisers do not capture the full posterior distribution of the clean data given its noisy version, but rather its conditional mean. We review the relevant background on diffusions and their limitations in Section 2.1. In earlier works, Xiao et al. (2022); Xu et al. (2024) proposed to use Generative Adversarial Networks (GANs) (Goodfellow et al., 2014) to learn such conditional distributions. We show here how to bypass adversarial train- ing by relying on generalized scoring rules (Gneiting & Raftery, 2007), which we describe in Section 2.2. This yields a simple loss which interpolates between the classical regression loss used in diffusion models and distributional losses. We present our new class of Distributional Diffu- sion Models in Section 3, and provide theoretical grounding for our choice of distributional loss, as well as details of the interpolation to the regression loss, in Section 4. We review related work in Section 5.
1. Introduction
Diffusion models (Ho et al., 2020; Song & Ermon, 2019; Sohl-Dickstein et al., 2015) have demonstrated remarkable success in synthesizing high-quality data across various do- mains, including images (Saharia et al., 2022b), videos (Ho et al., 2022), and 3D (Poole et al., 2023). These models proceed as follows: First, a forward diffusion process is detined where Gaussian noise progressively introduced to corrupt the data. This allows to learn denoiser at continuum of noise levels. Next, to generate new sam-
Experiments in Section 6 demonstrate that our approach produces high-quality samples with significantly fewer de- noising steps. We observe substantial benefits across a range of tasks, including latent spaces, and in robotics applications. All proofs are in the supplementary and a table of notation is in Appendix A.
Equual contribution 'Google DeepMind "Gatsby UCL. Corre- spondence to: Valentin De Bortoli <vdebortoli @rodebotool@ggoogl.coomm google.com>, Alexandre Galashov <agalashov googlc.com>>
tasks in both pixel and
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Distributional Diffusion Models
2. Background & Motivation
for a weighting function we (Kingma et al., 2021). Approx- imating (2))(2)(22)) with Dirac mass located (t, seems crude at first. However, as N + with (241 + fg. the resulting discrete time process converges to a contin- uous time process which recovers the data distribution if be (t, X,) = [[ X.]i see c.g. Song et al. (2021b).
2.1. Diffusion models
We follow the Denoising Diffusion Implicit Models (DDIM) framework of Song et al. (2021a). Let po be target data distribution on IE Consider X Po and the process Xtytx = (Xt, Xxx) distributed according
When performing a coarse time discretization, however, this result no longer holds, and the Dirac approximation is poor. We propose here to learn generative network to sample approximately from Poje Poje (2))
N-1
p(Zt [z(a = =p(0r, (1w))
II
place [Ta Fu-
with 0 to <.... < IN = 1. For 0 S < <s<s 1, we
2.2. Scoring rules
define p(z,, /ro, xz)= = (assi (zo, Iz), 2s,c) to ensure
that for any t E [0,
We next recall the framework of scoring rules, reviewed by Gneiting& & Raftery (2007), and then describe the specific scoring rules used in this work. Consider two probability distributions p. q. A scoring rule 5(p, indicates the qual- ity of a prediction p when the event = y is observed. The expected score is its expectation under y,
p( =
ofd,,
(2)
for some schedule Oi,, or chosen such that (o 01 = 1 and 01 To- = O. This ensures in particular that p( (xo) = (m1;; 0, Id). One possible popular schedule is given by the flow matching noise schedule (Lipman et al., 2023; Albergo et al., 2023; Gao et al., 2024), i.e.,
(p, = E,ls(p,
02 = 1-1 4,
Or =1
(3)
A scoring rule is proper when S(1 2 S().g)- It is called strictly proper with respect to a class of distributions P when the equality holds if and only if p = g, for all P.q. E P. In this work we focus on the class of proper scoring rules called kernel scores introduced by Gneiting & Raftery (2007, Section 5.1) which take the form
The mean and covariance of p( (ro s1) are given by
(r,, mr) = (e 1) + >² rr), )
+ a1 - (22 1) - (1- .,) Do)
Eg, = 01 (62r,, + (1 - =
(4)
1 Sp(p. v) ( - Eple(X,,)) v)l. -
with ts, (ar/as)' (0./o3); see (Song et al., 2021a) and Appendix F for a discussion. In (4), & E 10, 1] a churn parameter which interpolates between a determinis- tic process (E = 0) and a stochastic one (E = 1).
where p is a continuous negative definite kernel. Following Bouchacourt et al. (2016, Section 3.2), we will also consider generalized kernel scores of the form
We generate data Xto ~ Po by sampling Xen NO Id) and X1, p(11)) for k = N 1,... O, where for 0 < <stt<1
\ Sxp(p.8) = ( =xsp - Ex(p((X,)) v),,
(8)
p( = [ p(a.sloo x: (p dr)
(5)
with x E [0, 1]. One score of particular interest is the energy score (Gneiting & Raftery, 2007, eq. 22) denoted Sg, which uses p(x, || 2'' with 3 E (0, 2). We denote Sx.p. its generalized version, called the generalized energy score. We will also employ = - with k a continuous positive definite kernel, see e.g. (Berlinet & Thomas-Agnan, 2004) and Appendix A for a definition. Similarly to (7), we define the (generalized) expected energy score and (generalized) expected kernel score.
It can be easily checked that this process has the same marginal distributions, denoted pta (zt,), as the process de- fined by po( (r )p(ee, (1x))) (see (1)) so that in particular X ~ Po; see e.g. (Shi et al., 2024, Appendix E).
Subsequently, to avoid ambiguity, we write (a) (xr) for p(x.(xx) ( and (z. 24) for xe) for any 0 < < t t 1. Usually Poja (ar)) in (5) is approximated for any by δān(t,ra), where to (t, xx) x E[Xo] X; xx) is neural network denoiser trained using a regression loss; i.c.
Positive definite kernels are said to be characteristic to P when S, yields strictly proper scoring rule on P (Sripe- rumbudur et al., 2010; 2011). While the exponentiated quadratic (rbf) kernel satisfies this property, we will also investigate other kernels such as the inverse multiquadratic
(0) = =f,' w, - (o, Jdd,
(6)
2

---

Distributional Diffusion Models
(imq) kernel and the exponential (exp) kernel.
generalized energy score with & E (0, and A E [0, 1]
1/2, kimg (1r, x²) = (1|æ 2|2 + -
(9a)
X
Sx.o (1)c()))((
2 poje (1 1100)))))))))))))))))))))) (1I - x1
pgg
Rebe (21,2') = exp [ || x' 1|//20]]
(9b)
- E pg ( (1 | - =ollo].
(12)
kexp(z, x') = exp |- || - x' 11/o].
(9c)
Any proper scoring rule Sp defines a divergence
The conditional energy score was first proposed for train- ing conditional generative models by Bouchacourt et al. (2016), and its use in extrapolation and model estimation was established by Shen & Meinshausen (2024) in the case À = B = 1. We emphasize that (12) is conditional, since it compares the conditional distribution P (1(ara) to Fp. We recall that this scoring rule is strictly proper when À = 1 and 3 E (0, 20. The hyperparameter À allows us to control the trade-off between diversity (first "interaction" term) and accuracy (second "confinement" term).
Dplp. = 5(g,4) - S(p,g).
(10)
When S Sp and p(x,x²) -k( x²) with k a pos- itive definite kernel, we recover the squared Maximum Mean Discrepancy (MMD²) (Gretton et al., 2012). For p(a, x²) = ||0 - 2/11,, we obtain the energy distance (Rizzo & Székely, 2016). The relation between these discrepancies was established by Sejdinovic et al. (2013) where the spe- cific positive definite kernel family for which the MMD and energy distance are equivalent is described.
Energy diffusion loss. Our final expected loss integrates the conditional generalized energy score over both the noise level t E [0, 1], and the samples from the dataset, Xo ~ Po, together with noisy samples X, ~ Prjo(- so that (No, Xt) ~ Po,t- This gives the energy diffusion loss
Letting 3 + 2 for the energy distance, we obtain (p(p. q) = [X] - E. XX] resembling the integrand of (6). Indeed, up to a constant C' independent of f, letting p = - and q = P we get
[(0) -00=-/o- wr Wp
dt, (13)
Ep, (D,( g)] x (t. x
(11)
see Appendix B for proof. the case of p k with k given by (9a) or (9b), the relation between D, and the diffu- sion loss (6) is less immediate than for the energy distance. However the connection becomes apparent by keeping the leading terms in the Taylor expansions of the kernels, as shown in Proposition 4.3.
where we is a user-defined weighting function. The loss (13) has the following remarkable property: when we set / + 2 and À + O, we recover the classical regression diffusion loss (6). When 8 E (0,2) and X = 1 we obtain an (inte- grated) strictly proper expected conditional energy score. By selecting À E (0, 1) and ,3 E (0, 2), we can interpo- late between these two cases. Our approach shares with (Bouchacourt et al., 2016; Gritsenko al., 2020; Chen et al., 2024b; Pacchiardi et al., 2024; Shen & Meinshausen, 2024) the conditional energy score (12), but differs in that these earlier works learned only a single conditional distribution, whereas we integrate over many different noise levels.
3. Distributional Diffusion Models
We now introduce Distributional Diffusion Models. At a high level, we replace the regression loss (6), used to learn approximation of the conditional mean B[Xo]_X] in diffusion models, with a loss based on scoring rules, to learn an approximation ( (r)) of Po)e(Ho)zz)) This is achieved by learning generative network (t)
Empirical energy diffusion loss. For samples (Xi)= Lid. Po from the training set and correspond- ing noise levels {t. L.d. 4'0, 1]. sample noisy data points X; Xg,t t ~ (-)))- from the forward diffusion process (2) with noise schedule (3). Then, for each pair ples & ~ (t,, Xx,, g.). We then obtain the following empirical en- ergy diffusion loss
which aims to produce samples Xo ~ P This model takes as input time t [0, 1], a noisy sample x, and Gaussian noise sample & ~ (V(0,,d) Sampling from the noise & allows us to sample from this model (as in a GAN). The model's objective is to approximate the full distribution Poje (*) i.e., (t(1, X,, {) d Xo! I and & ~ (N(O,dd)) Whenever the model (t, Xx, <) is used, we refer to such method as distributional. Using this model allows us to sample from diffusion models with fewer steps than classical diffusion models, leveraging (5).
we sample m Gaussian noise sam- j E [m] which allows us to compute
1 n.m
Wr, || zo(lit X1, {t) {3)l (14)
Cn,m (0)
v,j=1
Conditional Generalized Energy Score. We base our training loss on generalized energy scores as introduced in Section 2.2. For each 1, No, It, we consider conditional
Co Xe, 2o(ti, 2/ }11?
2(m
1)
}'
3

---

Distributional Diffusion Models
Sig now outputs an approximate sample from (o)o()))2)) instead of an approximation of ( m4].
10
18
Algorithm 2 Distributional Diffusion Model (sampling)
Require: {4x ) with to 0 < ... < IN 1, churn
parameter
Sample Xty ~ (N Id)
for & E {N - -,...,, 0} do
Sample { ~ ~N(,
Sample (o,,
Figure 1. Left: reduction factor f(, a function of B E [0, 2] and À= E 00, 1]. Right: evolution of the covariance during the sampling distributional diffusion model for a Gaussian target po = (0, 4ld) for 20 of N E (5, 10, 95, 1000, À = 0.5. As increases, the output variance gets closer to 4.
Set Xo = fto(tt+1 x
Compute Meu (Xo, Xeats) ) 1//2
thitk using (4)
Set Xr, Max (Xo, X (ex11) + 2 2
1-1/2
end for Return: Xo
In Section 4, we discuss an alternative training loss using the joint generalized energy score which would focus on pl.: and show that its empirical version suffers from much higher variance. See also Appendix D for longer discussion.
4. Theoretical analysis
In this section, we provide some theoretical understanding of the generalized scoring rules introduced in Section 3.
Kernel diffusion loss. Similarly, we can define kernel diffusion loss by replacing | - æ' z²11 with a characteristic kernel -k(x z²), such as (9a) or (9b), in (12) and (13). As for the energy score, we can recover the diffusion loss (6) for kinq and krot by identifying the leading terms in the Taylor expansion of these kernels. We refer the reader to Section 4.3 for more details.
4.1. Gaussian analysis
We shed more light on the two main hyperparameters of the generalized kernel score introduced in Section 2.2: (i) the kernel parameter (8 in the case of the energy score); (ii) the trade-off term X. We focus here on the case of generalized energy score defined in (8).
Architecture and training. Unlike standard diffusion mod- els, the neural network to (f, xr, 5) takes not only and x; as input, but also & ~ (N0 Id) to obtain an approximate sam- ple from Poj e (z)) For simplicity, in image generation, we concatenate [sx, <] along the channel dimension without modifying the rest of the architecture, see Appendix I for more details, and Algorithm for the full training algorithm.
Proposition 4.1: Assume that p NNH, ( for μ E
IRd and a > O. Then, for any X [0. and B E 10, 2). we have that 42,5 = N(u,, oidd) maximizes Sx,3 g) defined by (8) among all Gaussian distributions g. with parameters
a. = (2X-2/12-0) -) -
Algorithm 1 Distributional Diffusion Model (training)
P. = H4
Require: M training steps, schedule (ay, ot), distribution Po. weights 0, batch size n, population size m
A notable outcome of maximizing the generalized energy score is that μ¹ = f no matter the value of À and B. With 2, we recover expected the correct variance, since Sx,8 is strictly proper energy score. As soon as X < 1, the variance a² is underestimated by factor f(A, B) = (2X () 1)- see Figure 1.
for A: 1: M do
X = and
Sample 1; ~ [ni 1J) for E [m]
iid. Sample X3 ~ Po for i E [n]]
Sample X; ~ P- X)) for E [n] using (2)
Sample Er,3 ~ (O Id) for i E [m].] E [m]
Set 0% 0x-11 Ln (0)lez-, using (14)
We now consider the effect of these hyperparameters when sampling from the corresponding generative model using Algorithm 2. For simplicity, we let the churn parameter E but our analysis also applies to 5 [0, 1). We consider a Gaussian target po -N(,, o Id). In that case, Poje (fol2;) is also Gaussian for any / E [0. 1j. Therefore if this density were estimated using the generalized energy score, then the results of Proposition 4.1 apply. In particular, this implies that the induced estimate of the Gaussian distribution paje obtained by plugging our approximation of Poje into (5)
end for
Return: 0as
Once again, we emphasize that setting À = 0 and B = 2 in Algorithm recovers the classical diffusion model loss.
Sampling. Once we have trained our model f.. we generate samples using Algorithm 2. This procedure is very similar to DDIM (Song et al., 2021a). The only difference is that
4

---

Distributional Diffusion Models
would have the correct mean but incorrect variance for
As expected, differences between classical diffusion models and distributional diffusion models arise when considering larger discretization stepsizes, where the Dirac approxima- tion of Puje (zo (rt)) is no longer valid, see Figure
4.2. Joint or conditional scoring rules
Given two unbiased estimators and B of A and B, we say that A is easier to approximate than B if SNR(Á) SNR (B) where, for an unbiased estimate Č. SNR(C) B[C] /Var(C); i.e. the inverse relative variance of C. We investigate here one alternative to the loss (13) and show that its empirical version exhibits lower SNR than the empirical version of (13) under the simplifying assumption that 0 is chosen so that Poje = Poje- While the loss (13) leverages generalized energy score on the conditional distribution Poje. we can instead leverage the generalized energy score on the joint distribution P define a new loss
Figure 2. SNR((Z,) (red) and SNR((In.join)) (blue) w.r.t. o² (x- axis) for possible bias levels b E {0, 1, (from dark to light). Vertically, we plot (ag)² such that for the bias b, we have SNR(I.) SNR(In.jont) for oz.
Hoogeboom et al., 2023), where b E R is some bias. The SNR of U, V can easily be computed and we observe that, across a large range of values of a and bias values we indeed have SNR((n,m)) 2 (N.(iz, in the large m limit, see Figure 2.
L.¹ Ep.. (S (Xu, (Xu, Xii)] dt,
Cjoint (0)
(15)
4.3. From kernel scores to diffusion losses
where Po denotes 161, (2o) (p In the case of the conditional loss (13), an empirical interaction estimate is
We have shown in eq. (11) of Section 2.2 that the energy diffusion loss (14) recovers the diffusion loss (6) when / 2. In that case, we say that the scoring rule is diffussion compatible. A natural question to ask is whether other scoring rules are diffusion compatible: i.e., given P. with c E RR a hyperparameter of the kernel, that there exists c* E [-00, +20| and f - R. such that lim [(c) (p. 1) = [[Ep[X]] E21
Inmm
2n(m 1) Aij
where Dr.j (tt &,,) Lo Xe.r el elp where X ~ Pt.. g) ~ (1, Id) for all j. the case of the joint loss (15), an empirical interaction estimate is
c
A
|IXI
Xx +1 A..)
The following result shows that kimg and keot can also recover the diffusion loss.
Joimiooint
2m(m- 1)
i=1 jt'
Assume that p= k with k given by (9a) or (9b). Then the scoring rule diffusion compatible.
Proposition 4.3:
where M, | X1, (/) Eo x2' 11 Here
we have X ~ pt.,, c/- ~ NO, Id) for all j, We then ask Unfortunately, we cannot answer this in the general case. However, if Po (0, I'Id), and leveraging results from U-statistics, have the following result.
whether SNR(In,m) 2 SNR(In.m.joaa)
This proposition justifies using scoring rules other than the energy one, which also allow recovering diffusion models in the limit. We compare the performance of these different kemnels in Section 6.2. Our experiments suggest that dif- fusion compatibility, as demonstrated by the energy score, kingg and kebe, is sufficient for defining a loss function that leads to high-quality sample generation when used to train model. However, we have identified other kernels, such as kexp that do not satisfy the diffusion compatibility require- ment but still result in models with desirable properties.
Proposition 4.2: Let U pu and V ~ py where pu a²0² wr/(1 + of )1/2 and pv & wela + (aza² + 2/) are two distributions on [0, 1]. Then, we have that =-0
SNR(I.) := lim SNR(In.m) nSNR(U),
SNR(In.joinn = lim SNR(In.m.join) moinn)) = nSNR(V). +00
5. ReRated Work
In practice, we consider the sigmoid weighting scheme
Accelerated diffusion models. Many strategies acceler- ate diffusion models, broadly classifiable as distillation-
ut = (1 + exp[b log(af/o?)]) (Kingma et al., 2021;
5

---

Distributional Diffusion Models
Squared [generated target]
Average (on sld of Aall.
hebe level C- 6.99
Nusse leve bas
Marie lowel t- D.45,
Nose level
No revett 0.01
6651
1.23
6631
6054
1.00
0.0
Num sampling steps
Timestept
Left: Squared MMD between target distribution and sampled data according to different models, z-axis denotes the number of sampling steps. Right: average standard deviation of samples Xol.X, produced by either true posterior distribution Poje or the models (t, Xx, E), a-axis is the timestep f.
Figure 3.
Postertor pull
Sytt. 1. 0.1.A
Aglt. m. 2=
Xalo
Figure 4. Samples from true posterior Poje(-)az,) (light blue) for a
sample xi (black cross) from (e)((2s) (2a) for specific zo (orange cross) and samples Sig (f, ze. E) ~ d Poju(- (green dots) for 3 = 0.1, 1. Top/bottom row: zo is the mean of the left/right
or sampling-based. Distillation methods (Luhman & Luh- man, 2021; Salimans & Ho, 2022; Luo, 2023; Salimans et al., 2024; Dieleman, 2024; Liu et al., 2023; Meng et al., 2023; Song et al., 2023: Franceschi et al., 2024; Huang et al., 2024; Sauer al., 2025) train student model to mimic teacher diffusion model, leveraging consistency losses (Song et al., 2023), adversarial losses (Franceschi et al.. 2024; Xu et al., 2024; Sauer et al., 2025), or noise coupling (Huang et al., 2024). See (Dieleman, 2024) for a detailed discussion. While distillation is prevalent, other work focuses on improved samplers for larger step sizes (ooliceeu-Martincauu et al., 2021; Lu et al., 2022; Zheng et al., 2023). Our work differs from both, neither training student nor proposing new samplers.
Gaussian, Samples from (t ze, E) (red) with & = 2, A = 0 concentrate around [[ as expected.
Binkowski et al., 2018) and energy distances (Liu, 2017: Bellemare et al., 2017; Salimans al., 2018) have also been defined on adversarially trained discriminative neu- ral net features. The conditional generalized energy score of (12) has been used in learning conditional distributions by the DISCONet approach of Bouchacourt et al. (2016), and by the engression approach of Shen & Meinshausen (2024) (which corresponds to the special case of X, B 1). Recently, energy distances have been applied to speech syn- thesis (Gritsenko et al., 2020), normalizing flows (Si et al., 2023), neural SDEs (Issa et al., 2024), and other generative models (Chen et al., 2024b; Pacchiardi et al., 2024).
Discrepancy and diffusion. We modify the diffusion model training loss to learn the conditional distribution (o 22.- The importance of approximating the covariance of this distribution was noted by Nichol & Dhariwal (2021) and exploited in (Ho et al., 2020; Nichol & Dhariwal, 2021; Bao et al., 2022a;b; Ou et al., 2024). Closest to our ap- proach is (Xiao et al., 2022), which uses a GAN to learn (=)o()))))) for fewer sampling steps. While sharing the same motivation, our method avoids adversarial training and discriminator training, offering a simple loss modifi- cation that encompasses standard diffusion models. Other loss modifications, like using the & loss (Chen al., 2021; Saharia al., 2022a), aim to improve output quality, not reduce sampling steps. Galashov et al. (2025) can be seen as dual to our work: they learn discriminative kernel features with a flow at different noise levels, while we focus solely on learning a generator.
To our knowledge, such losses have not previously been used to train diffusion models. It is possible, however, to use an adaptation of (Galashov et al., 2025) in combina- tion with the approach of Bouchacourt al. (2016); Shen & Meinshausen (2024), in order to incorporate conditional GAN-style generation into the reverse process of a diffu- sion model. This idea was proposed by X. Shen (personal communication, 6th November 2024) in an open discussion following a presentation at Google Deepmind, as a future research direction of interest. The main idea of Galashov et al. (2025) is retained, namely use standard forward diffusion process, and distributional loss for the reverse process: however the adaptive-kereel MMD loss for the reverse process is replaced with fixed-kernel energy score; and this score is used to train a sequence of conditional GAN generators to approximate p( I.t,,), rather than using particle diffusion directly. The approach indeed rep- resents a promising line of work (since developed in Shen et al., 2025) for discretizing the reverse diffusion process, as distinct from the DDIM-style approach adopted in Section 3. The approach can further be understood as non-adversarial formulation of Cheng et al. (2024), which uses an adaptive conditional GAN critic in place of the fixed-kernel energy score.
Energy distances, MMD, and generative modeling. The MMD (Gretton et al., 2012), of which energy distances (Székely & Rizzo, 2013; Rizzo & Székely, 2016) are spe- cial case (Sejdinovic et al., 2013), has been widely used as distributional loss in generative modeling. GANs have used MMDs with fixed kerness as critics to distinguish generated from reference samples (Li et al., 2015; Dziugaite et al., 2015; Unterthiner et al., 2018). MMDs (Li et al., 2017:
6

---

Distributional Diffusion Models
6. Experiments
CIFAR-10
CelebA
In this section, we validate the performance of our approach in 2D, image generation, and robotics settings. The main ex- periments are presented here, while additional experiments and ablations given in Appendix K. All the experimen- tal details are described in Appendix I, while Appendix J analyzes computational complexity.
FID
Num diffusion steps
Num diffusion steps
- 2=2.-- 3=2200 R-2A-SE
6.1. 2D experiments
Consider target given by mixture of two Gaussians,
Figure 5. Conditional image generation. a-axis denotes number of denoising steps, y-axis represents FID. Black dashed line denotes the performance of diffusion model at 100 steps.
i.c. Po = (1.. Id) + 0.5. (p2, o2Id), where μl1 = (3.3). U2 = (-3,3) 3) and Ø = 0.5. We train an unconditional distributional model žg with , = 0.1 and X = using Al- gorithm 1. Additionally, we train a baseline unconditional diffusion model by optimizing (6) and the model že with 3 2 and X= = 0 using Algorithm 1, i.e. classical diffu- sion model using the same architecture as the distributional models. More details are provided in Appendix 1.1.
LSUN Bedrooms
Latent CelebA HQ
FID
To measure the quality of the samples produced, we use the MMD squared given by Dp (10) with p(z,z') = -kene( x²) (9b) for = 1. Figure (left), we observe that with few denoising steps, the distributional variant has smaller Dp compared to other models, and achieves similar performance the diffusion model using large number of steps. In Figure 3 (right), for each X;, we produce 8 sam- ples & ~ (O, Id) and we compute the standard deviation (std) of (t, zij, <) over g. We then average the std over all the r/. While classical diffusion models cannot model the variance of the posterior, distributional diffusion models have std close that of the true posterior.
Num diffusion steps
Num diffusion steps
---21-0 142000 BC2A-02
F42.1-35 -12. Dihaainn
Figure 6. Unconditional image generation. zeaxis denotes num- ber of denoising steps, y-axis represents FID. Black dashed line denotes the performance of the diffusion model at 100 steps.
pendix 1.2 for details.
In Figure 5, we show results of conditional image genera- tion on CIFAR- 10 and CelebA, and in Figure results of unconditional image generation on LSUN Bedrooms and latent CelebA-HQ. Whenever we report performance of dis- tributional model and write 3 < 2, we select parameter / given one fixed À which minimizes the FID. More detailed results are given in Appendix K.2, which include FIDs for different numbers of steps for every combination of (B, A).
To further our analysis, we visualize the samples from the distributional model in Figure 4. For given 2, values, we sample from the true posterior (ooe( ar)) which is available in closed form. We also sample (t, Ze. {) ~ Poj (-Ix;) where { (O, Id) for a model trained using 3 = 0.1, 1, showing that our generating network is able to learn good approximation of the posterior. We also display samples from (t, .e. <) trained with 3 = 2, A = 0 which are concentrated around E[Xo|ara) as expected.
Main takeaways. The results suggest that whenever the number of diffusion steps is low, then distributional models with X = 1, < 2 achieve better performance than classical diffusion models. This confirms that the diffusion model approximation is poor in the "few steps" regime, and can be improved with better modeling of Poje as explained in Sec- tion 2. As the number of diffusion steps increases, the Dirac approximation becomes more accurate and diffusion models achieve the best performance. In that scenario, we observe that distributional models with À = 1, 3 < 2 yield worse performance. Our hypothesis is that since distribu- tional models ie only approximate the posterior distribu- tion Po this could lead to an increased amount of noise at sampling time and consequently to error accumulation. However, we observe that distributional model variants with
6.2. Image experiments
Main results. We train conditional pixel-space models on CIFAR-10 (32x32x3) and on CelebA (64x64x3), as well as unconditional pixel-space models on LSUN Bed- rooms (64x64x3). We further use an autoencoder trained on CelebA-HQ (256x256x3) producing latent codes of shape 64x64x3, which are then used to build latent unconditional models. For each dataset, we train a diffusion model by op- timizing (6) as well as a distributional model (Algorithm 1) for different >= E [0, 1] and B E [0.001, 2]. Models are evaluated using the FID score (Heusel et al., 2017). See Ap-
7

---

Distributional Diffusion Models
B
Best P
Imvase NG
Exponential
204-
1.5
100
024
Number of affls Staps
Diffision
4
45
Num diffusion steps
Num diffusion steps
Figure 8. Conditional generation CIFAR-I0 10 with different ker- nels. z-axis is the number of diffusion steps. y-axis is the FID.
CIFAR- (cand)
LSUN Bedrooms (uncond)
CelebA (cond)
Latent CelebA HQ (uncond)
Figure? 7. Parameters À and 8 minimizing the FID for given num- ber diffusion steps. z-axii indicates the number diffusion steps and legend color indicates the dataset.
À = 1 for few diffusion steps and better performance with X < when number of steps increases; see Appendix K.2 for detailed results.
Comparisons to distillation and numerical solvers. In Appendix K.3, we compare our approach with the DPM- solver++ method Lu et al. (2022) and the state-of-the-art, multi-step distillation method, moment- moment-matching matching distilla- tion (Salimans et al., 2024). Overall, our approach achieves results competitive with multistep distillation (where we compared different numbers of steps) and outperforms DPM-solver++ for more than 8 NFEs. The performance of DPM-solve+++ degrades as the number of steps increases, which is expected due to its numerical instability.
E (0,1) still yield good performance as the number of diffusion steps increases. This highlights the trade-off be- tween distributional and diffusion models controlled by the parameter X. As argued in Section 4.1, À < leads to underestimation of the target variance which, for denoising diffusions, implies that the variance of âg is more concen- trated around Poj: than would occur for posterior samples. Remarkably, we also found that using 3 < 2, À = led to comparable or even slightly better performance compared to classical diffusion models. This suggests that could be valuable to train diffusion models using a different loss than Mean Square Error (MSE) regression.
Our approach offers two compelling benefits over multistep distillation. First, does not rely on distillation, eliminating the need to train large (and slow) teacher model. Second, does not require specifying additional sampler hyperparame- ters during training. Combining our distributional approach with modern distillation method is an interesting direction for future research.
Conditional vs unconditional generation. We observe slightly different behavior of distributional models in the conditional image generation case (Figure 5), and in the un- conditional case (Figure 6). For a small number of diffusion steps, we notice that the gap in performance between the model trained with À= = 3 < 2 and a diffusion model, is much smaller in the unconditional case (1.5x) compared to the conditional one (3x). In Figure 7 we visualize the best parameters X, 3 selected to achieve the lowest FID, as a function of diffusion steps. More detailed results are presented in Appendix K.2. We observe that overall À fol- lows downward trend, decreasing as the number of steps increases. However, / behaves differently in each setting, increasing monotonically as a function of diffusion steps in the conditional case, while remaining at a minimal value in the unconditional case. We hypothesise that this is because learning unconditional distributions is a much harder task since the data is spread out more, leading to much higher variance in the distribution of žg.
6.3. Robotics experiments
Experimental setup. We demonstrate here the effectiveness of our distributional diffusion models for robotics applica- tions. We experiment with diffusion policies on the Libero (Liu et al., 2024) benchmark, a life-long learning benchmark that consists of 130 lation tasks. In our experiments, we used Libero suites with 10 tasks in each: Libero-Long, Libero-Goal. Libero- Object, and Libero-Spatial. Our multi-task diffusion policy is conditioned on the encoded visual and proprioceptive ob- servations, and the task descriptions and generates a chunk of 8 7-dimensional actions to execute; see Appendix I.3 for additional details on the architecture.
conditioned robotic manipu-
Distributional diffusion model. We train the distributional diffusion model with Algorithm using the energy diffusion loss with population size m 16 and varying / and X. We use diffusion model to model a sequence of 8 actions a = (a)...., ag). where a; E IR.. We use the noise & E 8,2 which we concatenate with a on the last dimension. At sampling time, we follow Algorithm to sample sequence
Using different kernels. As discussed in Section 2.2 and in Section 3, we can also use the kernel diffusion loss. We train de with kernels (9a), (9b). (9c), varying kernel parame- ters and À E [o. 1]; see Figure& 8 for results on conditional im- age generation for CIFAR-10. 10. We observe similar behavior to the energy diffusion loss, achieving the best results when
8

---

Distributional Diffusion Models
Libero10
(Salimans al., 2024) is further avenue for future work.
Max success rate
Finally kernel scores allow for more general domains than R by way of characteristic kernels on groups (e.g. the group SO(3) of 3d rotations) and semigroups (d- dimensional histograms) (Fukumizu et al., 2008); and graphs (Vishwanathan et al., 2010). thus possible to extend our methodology these scenarios.
0.70
0.65
0.60
0.55
0.50
Impact Statement
16
50
Num diffusion steps
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences our work, none which we feel must be specifically highlighted here.
==10A=00
=
8-1.0=00.1
Diffusion
Figure 9. Performance in Liberol0 as function of diffusion steps and A. We plot the maximum success rate from best checkpoint during training. À 0 works best when more diffusion steps are used, and > 0 gives best results when fewest diffusion steps used. À > 0 also reduces the performance by a littlc for 50 diffusion steps (see 0.5 vs 0.0).
References
Aiello, E., Valsesia, D., and Magli, E. Fast inference in de- noising diffuusin models via MMD finetuning. IEEE Ac- cess, 12:106912-106977 2024. doi: 10.1109/ACCES. 2024.3436698.
Albergo, M. S., Boffi, N. M.. and Vanden Dijnen-, Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:3303.08797. 2023.
of actions from diffussin policy and we execute it. We sweep over number (2, 16, 50} of diffusion steps.
Results. For the main result, focus on Libero-Long. the most challenging suite with 10 tasks that features long- horizon tasks with diverse object interactions. Similar to the results on image experiments, À 0 works best when more diffusion steps (50) are used during evaluation. Using > gives best results when fewer diffusion steps are used (2, 16). We note that the performance slightly deteriorates with À > 0 when the number of diffusion steps is large, see Figure 9. See Figure 19 for similar results other Libero suites.
Bao, F. Li, C., Sun, J., Zhu, J., and Zhang, B. Estimating the optimal covariance with imperfect mean in diffusion probabilistic models. In International Conference Machine Learning, 2022a.
Bao, F., Li, C.,, Zhu, J., and Zhang, B. Analytic-DPMr an analytic estimate of the optimal reverse variance in diffu- sion probabilistic models. In International Conference on Learning Representations, 2022b.
Bellemare, M. G., Danihelka, I., Dabney, W., Mohamed, S., Lakshminarayanan, B., Hoyer, S., and Munos, R. The Cramer distance as a solution to biased Wasserstein gra- dients. arXiv preprint arXiv: I70. 2017.
7. Discussion
We present a novel approach for training diffusion mod- els, by learning the full conditional distributions Poia using scoring rules. This approach achieves good performance in image generation and robotics tasks, outperforming stan- dard diffusion models in the "few-step" regime with mini- mal degradation of output quality. Our experimental results further suggest that standard diffusion models can also be trained with objectives other than classical MSE regression.
Berlinet, A. and Thomas-Agnas C. Reproducing Kernel Hilbert Spaces in Probability and Statistics. Kluwer, 2004.
Binkowski, M., Sutherland, D. J., Arbel, M., and Gretton, A. Demystifying MMD GANs. In International Conference on Learning Representations, 2018.
So far, our generative networks approximating Poje incor- porate the noise& & only by concatenating it to the x; along the channel dimension. We believe that other architectures could significantly improve sample quality, and will ex- plore this in future work. Another promising direction is in learning the kernel and the parameter À to maximize generative performance in the few-step regime. Combining our approach with powerful distillation techniques such as
Bouchacourt, D.,, Mudigonda, P. K., and Nowozin, S. Disco nets: Dissimilarity coefficients networks. In Advances in Neural Information Processing Systems, 2016.
Chen, H., Ren, Y., Ying, L., and Rotskoff, G. M. Acceler- ating diffusion models with parallel sampling: Inference at sub-linear time complexity. In Advances in Neural Information Processing Systems, 2024a.

---

Distributional Diffusion Models
Chen, J., Janke, T., Steinke, F., and Lerch, S. Genera- tive machine learning methods for multivariate ensemble postprocessing. The Annals of Applied Statistics, 18(1): 159-183, 2024b.
Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. Advances in Neural Informa- tion Processing Systems, 2014.
Chen, N., Zhang, Y., Zen, H., Weiss, R. J., Norouzi, M., and Chan, W. WaveGrad: Estimating gradients for waveform generation. In International Conference on Learning Representations, 2021.
Gretton, A. Introduction to RKHS, and some simple ker- nel algorithms. Advanced Topies in Machine Learning lecture, University College London, 2013.
Gretton, A., Borgwardt, M., Rasch, M. J., Schölkopf, B., and Smola, A. A kernel two-sample test. Journal of Machine Learning Research, 13(1):723-773, 2012.
Cheng. Y., Liang, M., Huang, S., Han, G., Ning, J., and Liu, W. Conditional GAN for enhancing diffusion models in efficient and authentic global gesture generation from audios. arXiv preprint arXiv:2410.20359 2024.
Gritsenko, A., Salimans, T., van den Berg, R., Snoek, J., and Kalchbrenner, N. A spectral energy distance for parallel speech synthesis. In Advances in Neural Information Processing Systems, 2020.
Chi, C.,, Xu, Z., Feng, S., Cousineau, E., Du, Y., Burchfiel, B., Tedrake, R., and Song, S. Diffusion policy: Visuo- motor policy learning via action diffusion. International Journal of Robotics Research, 2023.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. In Proceedings of the IEEE conference Computer Vision and Pattern Recognition, 2016.
Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. preaaingg of deep bidirectional trans- formers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, S. GANs trained by a two time-scale update rule converge to a local Nash equilibrium. In Advances Neural Information Processing Systems, 2017.
Dieleman, S. The paradox of diffusion distillation, 2024. URL https: http:///anerrrreeeeeeeeeee att.dderr.a//222/2///2// paradox.. html.
Ho, J., Jain, A., and Abbeel, Denoising diffusion prob- abilistic models. In Advances in Neural Information Processing Systems, 2020.
Dziugaite, G. K., Roy, D. M., and Ghahramani, Z. Training generative neural networks via maximum mean discrep- ancy optimization. In Uncertainty Artificial Intelli- gence, 2015.
Ho, J., Salimans, T., Gritsenko, A., Chan, W., Norouzi, M., and Fleet, J. Video diffusion models. In Advances in Neural Information Processing Systems, 2022.
Franceschi, J.-Y., Gartrell, M., Dos Santos, L., Issenhuth, T., de Bézenac, E., Chen, M., and Rakotomamonjy, A. Unifying GANs and score-based diffusion as generative particle models. Advances in Neural Information Pro- cessing Systems, 2024.
Hoogeboom, E., Heek, J., and Salimans, T. simple diffu- sion: End-to-end diffusion for high resolution images. In International Conference on Machine Learning, 2023.
Fukumizu, K., Gretton, A., Schölkopf. B., and Sriperum- budur, B. K. Characteristic kernels groups and semi- groups. In Advances in Neural Information Processing Systems, 2008.
Huang, Z., Geng, Z., Luo, W., and Qi, G.-j. Flow generator matching. arXiv preprint 2410v 19310, 2024.
Hutchinson, M. F. A stochastic estimator of the trace of the influence matrix for Laplacian smoothing splines. Com- munications Statistics- Staristics-Simulation and Computation, 183)::059-1076, 1989.
Galashov, A., de Bortoli, V., and Gretton, A. Deep MMD gradient flow without adversarial training. In Interna- tional Conference on Learning Representations, 2025.
Issa. Z., Horvath, B., Lemercier, M., and Salvi, C. Non- adversarial training of neural SDEs with signature kernel scores. In Advances in Neural Information Processing Systems, 2024.
Gao, R., Hoogeboom, E., Heek, J., Bortoli, V. D., Murphy, K. P., and Salimans, T. Diffusion meets flow matching: Two sides of the same coin. 2024. URL https:// diffusionflos github.ic..
Saiicoeur-aattneea A., Li, K., Seihe-aail,,,er, R., Kach- man, T., and Mitliagkas, I. Gotta go fast when gen- erating data with score-based models. arXiv preprint arXiv:2105 /4080 2021.
Gneiting, T. and Raftery, A. E. Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association, 102(477):359-374 2007.
10

---

Distributional Diffusion Models
Karras, T., Aittala, M., Aila, T., and Laine, S. Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems, 2022.
Ou, Z.., Zhang, M., Zhang, A., Xiao, Z., Li, Y., and Barber,
D. Diffusion model with optimal covariance matching. arXiv preprint arXiv:2406. 10808, 2024.
Pacchiardi, L., Adewoyin, R. A., Dueben, P., and Dutta, R. Probabilistic forecasting with generative networks via scoring rule minimization. Journal of Machine Learning Research, 25(45): 1-64, 2024.
Kingma, D., Salimans, T., Poole, B., and Ho, J. Varia- tional diffusion models. Advances in Neural Information Processing Systems, 2021.
Li, C.-L., Chang, W.-C., Cheng, Y., Yang. Y., and Poczos, B. Mmd gan: Towards deeper understanding of moment matching network. In Advances in Neural Information Processing Systems, 2017.
Poole, B., Jain, A., Barron, J. T., and Mildenhall, B. Dream- fusion: Text-to-3d using 2d diffusion. In International Conference on Learning Representations, 2023.
Rizzo, M. L. and Székely, G. Energy distance. Wiley Interdisciplinary Reviews: Computational Statistics, 8(1):
Li, Y., Swersky, K., and Zemel, R. Generative moment matching networks. In International Conference on Ma- chine Learning, 2015.
27-38, 2016.
Rombach, R., Blattmann, A., Lorenz, D., Esser, P., and Ommer, B. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF con- ference Computer Vision and Pattern Recognition, 2022.
Lipman, Y., Chen, R. T., Ben-Hamul H., Nickel, M., and Le, M. Flow matching for generative modeling. In Interna- tional Conference on Learning Representations, 2023.
Liu, B., Zhu, Y., Gao, C., Feng, Y., Liu, Q. Zhu, Y., and Stone, P. Libero: Benchmarking knowledge transfer for lifelong robot learning. In Advances in Neural Informa- tion Processing Systems, 2024.
Ronneberger, O., Fischer, P., and Brox, T. U-net: Con- volutional networks for biomedical image segmenta- tion. In Medical image computing and computer-assisted Ciaveentto-MMIICCA 2015: 18th international confer- ence, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pp. 234-241, Springer, 2015.
Liu, L. On the tao-sa statistic approach to generative adversarial networks. Master's the- sis, University of Princeton Senior Thesis, April 2017. URL http:: //ack/////kk// www.rs.....inttoo....... edu/ark: /88435 (86435/dsp017440555
Rozet, F., Andry, G., Lanusse, F., and Louppe, G. Learning diffusion priors from observations by expectation maxi- mization. In Advances in Neural Information Processing Systems, 2024.
Liu, X., Gong, C., et al. Flow straight and fast: Learn- ing to generate and transfer data with rectified flow. In International Conference on Learning Representations, 2023.
Russakovsky, O., Deng, J., Su, H.,, Krause, J., Satheesh, S., Ma, S., Huang, Z.,, Karpathy, A., Khosla, A., Bernstein, M., Berg, A. C., and Fei-Fei, L. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3):211-52,, 2015. doi: 1.1007/s11263-015-083
Lu, C., Zhou, Y., Bao, F., Chen, J., Li, C., and Zhu, J. DPM- solver++: Fast solver for guided sampling of diffusion probabilistic models. arXiv preprint arXiv:2211.0109) 2022.
Saharia, C., Chan, W., Chang, H., Lee, C., Ho, J., Salimans, T., Fleet, D., and Norouzi, M. Palette: mmage-to-image diffusion models. In ACM SIGGRAPH, 2022a.
Luhman, E. and Luhman, T. Knowledge distillation in iterative generative models for improved sampling speed. arXiv preprini arXiv:2 2021.
Saharia, C., Chan, W., Saxena, S., Li, L., Whang. J., Denton, E. L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., and Salimans. T. Photorealistic text-to- text-to-image diffu- sion models with deep language understanding. Advances in Neural Information Processing Systems, 2022b.
Luo, W. A comprehensive survey on knowledge distillation of diffusion models. arXiv preprini arXiv:2304.04200 2023.
Meng, C., Rombach, R., Gao, R., Kingma, D., Ermon, S., Ho, J., and Salimans, T. On distillation of guided diffu- sion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.
Salimans, T. and Ho, J. Progressive distillation for fast sam- pling of diffusion models. In International Conference on Learning Representations, 2022.
Nichol, A. Q. and Dhariwal, P. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning. 2021.
Salimans, T., Zhang, H., Radford, A., and Metaxas, D. Im- proving GANs using optimal transport. In International Conference on Learning Representations, 2018

---

Distributional Diffusion Models
Salimans, T., Mensink, T., Heek, J., and Hoogeboom, E. Multistep distillation of diffusion models via moment matching. In Advances in Neural Information Processing Systems, 2024.
Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. Consis- tency models. In International Conference Machine Learning, 2023.
Sriperumbudur, B. K., Gretton, A., Fukumizu, K., Lanckriet, G. R. G., and Schölkopf, B. Hilbert space embeddings and metrics on probability measures. Journal of Machine Learning Research, 11:1517-1561, 2010.
Sauer, A., Lorenz, D.. Blattmann, A., and Rombach, R. Ad- versarial diffusion distillation. In European Conference on Computer Vision, pp. 87-103. Springer, 2025.
Sriperumbudur, B. K., Fukumizu, K., and Lanckriet, G. R. G. Universality, characteristic kernels and RKHS embedding of measures. Journal of Machine Learning Research, 12: 2389-224 2410, 2011.
Sejdinovic, D., Sriperumbudur, B., Gretton, A., and Fuku- mizu, K. Equivalence of distance-based and RKHS-basee statistics in hypothesis testing. The Annals of Statistics, pp. 2263-2291. 2013.
Steinwart, I. and Christmann, A. Support Vector Machines. Information Science and Statistics. Springer, 2008.
Serfling, R. J. Approximation Theorems of Mathematical Statistics. John Wiley & Sons, 2009. 2009.
Székely, J. and Rizzo, M. L. Energy statistics: A class of statistics based on distances. Journal of Statistical Planning and Inference, 14498):224-1127, 2013.
Shen, X. and Meinshausen, N. Engression: extrapolation through the lens of distributional regression. Journal of the Royal Statistical Society Series B: Statistical Method- ology, 2024.
Unterthiner, Nessler, Seward, Klambauer, Heusel, Ram- sauer, and Hochreiter. Coulomb GANs: Provably optimal Nash equilibria via potential fields. In International Con- ference on Learning Representations, 2018.
Shen, X., Meinshausen, N., and Zhang, T. Reverse Markov learning: Multi-step generative models for complex dis- tributions. arXiv preprint 1rX1v:2502.13747, 2025.
van den Oord, A., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. In Advances in Neural Information Processing Systems, 2017.
Shi, Y., De Bortoli, V., Campbell, A., and Doucet, A. Diffu- sion Schrödinger bridge matching. In Advances Neural Information Processing Systems, 2024.
Vishwanathan, S., Schraudolph, N. N., Kondor, R., and Borgwardt, K. M. Graph kernels. Journal of Machine Learning Research, 11(40):1201-222 2010.
Shih, A., Belkhale, S., Ermon, S., Sadigh, D., and Anari, N. Parallel sampling of diffusion models. In Advances in Neural Information Processing Systems, 2023.
Xiao, Z., Kreis, K., and Vahdat, A. Tackling the generative learning trilemma with denoising diffusion GANs. In International Conference on Learning Representations, 2022.
Si, P., Chen, Z., Sahoo, S., Schiff, Y., and Kuleshov, V. ceni-aautoeeresssvve energy flows: exploring likelihood- free training of normalizing flows. In International Con- ference Machine Learning, 2023.
Xu, Y., Tong, S., and Jaakkola, T. Stable target field for reduced variance score estimation in diffusion models. In International Conference on Learning Representations, 2023.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequi- librium thermodynamics. In International Conference on Machine Learning, 2015.
Xu, Y., Zhao, Y., Xiao, Z., and Hou, T. UFOgen: You forward once large scale text-to-image generation via diffusion GANs. In Proceedings of the IEEE/CVF Con- ference on Computer Vision and Pattern Recognition, pp. 8196-8206, 2024.
Song, J., Meng, C., and Ermon, S. Denoising diffusion implicit models. In International Conference on Learning Representations, 2021a
Zhao, T. Z., Tompson, J., Driess, D., Florence, P., Ghasemipour, K., Finn, C., and Wahid, A. ALOHA unleashed: A simple recipe for robot dexterity. arXiv preprint arXiv:2410..3126 2024.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems, 2019.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Er- mon, S., and Poole, B. Score- generative modeling through stochastic differrnntial equations. In International Conference on Learning Representations, 2021b.
Zheng, K., Lu, C., Chen, J., and Zhu, J. DPM-solver-v3: Improved diffusion ODE solver with empirical model statistics. In Advances in Neural Information Processing Systems, 2023.
12

---

Distributional Diffusion Models
Organization of the Appendix
The appendix organized as follows. Appendix A, we present a notation table, to help the reader. Appendix B, we prove that diffusion compatible scoring rules can be linked to diffusion model losses as claimed in Section 2.2. In Appendix we prove the results regarding the minimization of generalized expected energy score Gaussian setting presented in Section 4.1. I. In Appendix D, expand on the discussion of Section 4.2 regarding joint conditional scoring rules. In Appendix E, we present the results of Section 4.3 regarding diffusion-compatiibe kernels. In Appendix F, we present a new perspective on the DDIM updates derived in (Song et al., 2021a). More precisely, we show that these updates can be recovered using a Stochastic Differential Equation perspective. This justifies the form of the DDIM mean and covariance updates presented in (4). We present some pseudocode to compute our loss function in Appendix G. Experimental details are provided Appendix We present a computational complexity analysis in Appendix We present an extended related work Appendix H. Appendix K, we present additional experimental results including ablation studies.
A. Notation
We recall that a kernel defined on space X is a symmetric function k: XxX + R, i.e. for any 2, E X, k( y) - k(,, kernel & is said to be positive semi-definite if for any n E N, fai,..., .n)) E E and (cI....,Gr) Cn E R. we have
Geyk(1,, 20.
(16)
The kernel & is positive definite if equality in (16) occurs if and only = 0 for any i E (1....,n) (Berlinet & Thomas-Agnan, 2004; Steinwart & Christmann, 2008; Gretton, 2013) for an overview on kernel theory.
n). We refer the reader to
Name
Notation
Definition
Appeared 1st
Continuous negative definite kernel
Sec 2.2
k
Continuous positive definite kernel
Sec 2.2
Scoring rule
spp.y)
Sec 2.2
Expected score
S(p.q) 4)
S(p, E. (Sp, Y)]
See 2.2
Kernel score
Sp(p.v)
Sp(p.)) v)
B(p(X,1) (x,
Sec 2.2
Energy score
Sa(p. v)
Sp(p,y) S,(p,)) 3) with p(z,, x') ) = | 2'11
Sec 2.2
Expected Kernel score
Sp(p.9)
Sp(p, Byy (p, V)
Sec 2.2
Expected Energy score
Sg(p. 4)
Sp(p.9) = Sp(p, 4) with plxx, x²) = || 1'113
Sec 2.2
Generalized (Gen.) Kernel score
Sxp(p.v)
Sxp(p.v) v) = 2 Epsplp(X X')] E[p(X, .)1
Sec 2.2
Generalized Energy score
(p.)))
Sx u) = Sx v) with p(r, x²) = || 2'11
See 2.2
Expected Gen. Kernel score
(x,(P).))
Sx 9) = Ey S Y)]
Sec 2.2
Expected Gen. Energy score
Sx,(lp,9)
Sx,(pp,) 9) (p 9) with plx,, x') ||--
See 2.2
Squared MMD
Dp(p.4)
Dp 9) Spl9.9) 9) Sp(p. 9)
See 2.2
Energy Distance
Dz(p.g)
Dzlp, 9) D Dplp, 9) with plx, x²) ||0: 2'1
Sec 2.2
Diffusion Loss
(a)
Can(A) So Wr..p. 11 (t, ) 12| dt
Sec 2.1
Energy Diffusion Loss
L())
[(() Sa "e W...
SA.0(P6 (-1 Xx) dt
Sec
Kernel Diffusion Loss
C())
(()) - So' Wylppp..
(- dt
Sec
Joint Diffusion Energy Loss
Cjoin (0)
Cjoint (0) = - -So w/ Eyo.
S2,0(18 (Xo, Xx)) di
Sec 4.2
13

---

Distributional Diffusion Models
B. Correspondence between Discrepancy and Diffusion Loss
In this section, we demonstrate the connection between diffusion-compatible scoring rules and diffusion model losses, as stated in Section 2.2. More precisely, we prove the following result.
Proposition B.1: Assume that Dp(p, 9) = | X] [x]2 and that = (t.(t.X.) and g = --o)e(-(Xe). Then, we have
Ep,(D((p,))] = Epo. ||| Do x,) Epo. (1.xo [x]11)
where the second term the rh.ss. independent of 0.
Proof. First we have
Ep, (D,(p, x)) (1 Xt) B[X Xx]1
ppo, [1 Xt) Xo + X, - R. XX [x]1
E (t, Xx) - -Xol + E ||| X, E|Xo] X, [| + po, (t, X,) - No. Xo -E xxII
Epo, (1,:0 X2) - xoo1²] + Po, [o B[Xo| X |1]1
+ p0. ( (t), Xo | [x]) 2E [.. [(oo, Xo - E|Xo/X,111
Epo, /l Xee) -ooll]] + po. [l 15 /
p.. Ioo..:l(ooo... Xo - E[Xo [.
(17)
In addition, we have
po 11 E|Xp/ XxI1 pp Xull Ep [1 /XAeLxX
(18)
Similarly, we have Combining (17), (18) and (19), we get
Epo 1(No No - FE X,II xx(ooXx,))]
= Ep ||| X - Epp | X X [1]
(19)
Ep, (Dp = Epo, 11 Bo xx)1 Ep. [l [x
which concludes the proof.
C. Generalized Energy Score for Gaussian Targets
We first investigate the properties of the generalized energy score for a single Gaussian target in Appendix C.I (Proposi- tion 4.1) and then apply these results to diffusion models in Appendix C.2
C.1. Learning a Gaussian distribution
For any 8 E 20 and À [0, 1] we recall that Sx.8 is the generalized expected energy score given for any distributions p, on Id with 3 moments by
X Sp.( -ex |||X Y11 + A (xxx))))= ||| x
We refer the reader Appendix A for remainder the notation used throughout the paper. For any l E Idd and a 0. let Py.er (N,, Id)) and denote by & the space of Gaussian distributions with scalar covariance, i.e.,
c = (Id, v Rd 7 0}.
We consider the following minimization problem: for any μ Rd and solve
Pp.... = a S Pp.a)
When = 1 and 3 E (0, 2). Sp.s q) is a proper scoring rule (Rizzo & Székely, 2016; Gneiting & Raftery, 2007) so we have Pp..o, = Pp.o- For [0, 1). we have the following result. This a simple restatement of Proposition 4.1.
14

---

Distributional Diffusion Models
Proposition C.1: For any 8 E (0 and X E (0, 1), we have that op..o. E argmax Sx Pu.o) with
o, = (2X 222-21(2-)) 1)
p. = ,4,
We start with the following lemma.
Lemma C.2: Assume that X ~ (10, I > 0. Then for any 3 > 0, we have 0 E org El|x cll
Proof. For any B O, let [s(c) E cll"). For any 3 > O, we have that Sa is continuous, using the dominated convergence theorem, and coercive, using Fatou's 's lemma, hence admits at least one minimizer. If 2 1, then fa is convex. Assume that c* is minimizer. Then -c* is also a minimizer and by convexity, we have that 0 is also a minimizer. If 3 E (0, 1) then Sa is strictly concave. Assume that e* is minimizer. Then -c* is also a minimizer. If O is not a minimizer of fa. then f0) > fa(c*). Hence, by strict concavity, we have that lim fa(tc* = -oo which absurd. Hence, 0 also a minimizer.
Proof. First, we have that for any q E & with mean j and covariance 52Ld
Sx Pp.c) -B[lI jı + (a² + 0/))//221 E 2
2
where Z ~ Id) Maximizing with respect to jı first, we get that μr = μ thanks to Lemma C.2. In addition. if ğ = we have
Sx,8(a = -B]ll j + (o² + 03/1/221111 + Rill B[Il225211)
- 02) 2 2 1
== -0000000000000000 0/3
20/200)
Maximizing with respect to o, get that
0.
So the effect of the introduction of < is that a, underestimates a, i.e. the model is too concentrated around its mean. The effect is stronger as 3 >> 2. In what follows, we denote
f(x, B) = 222-2
(20)
C.2. Application to diffusion sampling
Consider a Gaussian target distribution po- We investigate the effect of using the generalized energy score learn the Gaussian conditionals Poje the sampling updates used inference time, ie. when substituting approximation of Paja within (5) to obtain approximation of Pajt-
We first recall a few useful lemmas. Using (2), the forward model is given by
X, 04 Xu + o.Z.
(21)
with Z ~ (O, Id) and Xn ~ Pp.o. the target density. The diffusion (Xt)sejo,1) given by
dx,- ftX, + g.
(22)
where (B:)eeoo.jj is d-dimensional Brownian motion has the same marginal distributions as (21) for
ft 0, log(or),
gt (0//a)
15

---

Distributional Diffusion Models
We refer Appendix F for derivation of this fact, also (Song et al., 2021b).
Lemma C.3: For any
[0, 1] we have that
((1te) ano
(H) (
(afo² of)ld )
(Xu, X1) ~ N
Proof. This a direct consequence (21).
In particular, we also have the following lemma.
Lemma C.4: For any 1 E [0, 1] we have that
Xo ~ N(r(()x, (1-r(())), 0 (2,((t))ld),
where
010²
a²'² 0, + o
r(t) a² + of
T2,((t)
Proof. This is combination of Lemma C.3 and the formula for computing Gaussian posteriors.
Lemma C.5: For s, t [0, with 2 s we have that
of - ( 0g0t;
X, ~ N
as
X,,,,
Id
as
as
Proof. This is a direct consequence of (22).
Lemma C.6: For s, E [0, 1] with 2 s we have that
Xo X, ~ (r(s, 1)x, + ass( - (2) o1 - 122(8, 1))ld)
where
a,o²
a7o
r(s,t)
r22(3, 0?
Proof. Using that X. is the same as X, X,, X, by the Markov property, we get the distribution of (X,, X.) We conclude upon using the formula for computing Gaussian posteriors.
Note that in Appendix F, we will establish Lemma C.6 in a more general setting, i.c., we will introduce a churn parameter and draw further connection with DDIM (Song et al., 2021a). At this stage, Lemma C.6 recovers (Ho et al., 2020) which corresponds to setting the churn parameter to 1. Now, combining (20), Lemma C.4 and Lemma C.6 we get the following proposition.
16

---

Distributional Diffusion Models
Proposition C.7: For s, E [0, 1] with 2 we have that
X, X =
with
Aalt (r) + as(1 72,2 (5))(()))) t 12,( t)) r(t)p,
8/ of (,,2(s t)) + f(,, 3) a (1
In particular, denoting X, ~ Ni4 Id) we get that
p. --((,,)) (1 - 2 (() + (1,(1 - (,, ()) -ar())) -2 (2, ()) /A, 3)0 ag (2,2(,, (1) r2,) + ot (r(, as(1 (22((,(()))(1)
(23)
In particular, the update on the mean in (23) is not dependent on and B. Therefore we estimate correctly the mean of Paje but incorrectly the variance for À <
The curves presented in Figure are obtained using Proposition C.7.
D. Different learning objectives
First in Appendix D.1 we compare the energy diffusion loss (13) with a similar loss based on Maximum Mean Discrepancy and show that they only differ by a term which does not depend on 8. In Appendix D.2, we compare the conditional loss and the joint loss, as discussed Section 4.2. In Appendix D.3, we consider a third possible option, i.e., we compare the conditional loss with a marginal loss. We highlight the limitations of the marginal loss in that case. We derive the empirical versions of the conditional and joint losses in Appendix D.4. Finally in Appendix D.5, we prove the main result of Section 4.2, i.e., we prove Proposition 4.2 and compare the SNR of the interaction terms of the joint and conditional losses.
D.1. Energy diffusion loss and Maximum Mean Discrepancy diffusion loss
In this section, first compare the energy diffusion loss (13) and another loss based Maximum Mean Diserepancy (MMD), denoted MMD diffusion loss. We show that these two losses are equal up to a constant which does not depend on the network parameters 0. However, this additional term appearing in the MMD diffusion loss is intractable, even though it could be potentially estimated using important sampling techniques. We consider the case where the kernel & is given by k(x,, =²) = - || - Il While our discussion can be extended to other settings, we also focus on the case where = 1, = 1 and wt for simplicity.
First, we consider the following MMD diffusion loss
LMMD (0)= S.' J (x(x:) (pole: Poj)da,,ds
(24)
We recall that the squared MMD is given by (p(p. Sp(g,9) sp(p.9) with p -k and
1.'' / J (2. xz
[(0) ((0)=- =
(25)
First, we highlight the following result which draws connection between kernel scores and squared MMD.
Proposition D.1: Let x E IRd be such that k(xx x) = 0, then we have Dp Sx)) x))
Proof. For any distribution we have Sp(g, q) 1/2 1 (k(,,,lk((XXX)) (XXXX))) and therefore Sp 8z) = kk(( x) = 0. Similarly, we have that
-2 Epsp 1xxX Epss, 16(X,Y)) = 1 [k [XX, E,,[X,,, Sp(, -
Sp(p,8,) -
17

---

Distributional Diffusion Models
Therefore, have that
Dp(p, = Sp(08,, (1p)).pp.b..
0.) -Sp(p, x).
Using Proposition D.1, (25) can be rewritten as
[()) -
I JESXRe po,e(0o, ai Szo ddrddxddl.
(26)
Therefore, our energy diffusion loss can be expressed terms of squared MMD. We observe that both (24) and (26) involve an integrated version of the squared MMD. However, in (26) we target Sx, while in (24) we target Poje- In the rest this section, we show that these two losses actually only differ by term which does not depend on the network parameters 6.
Let us start with the energy diffusion loss L given in (25).. We have
1 1/2 //' J.eyy po, xe P (+ a Daje (xg) = =d
((0)
2
L/ /
21) (2 20 ddlddndd//dddd
1
(20 =dlldz/deddrddtt I'' /eep (r(
2
/ (Rays
pr (o
dzd
For the MMD diffusion loss (24), we have
LMMD (0) -- J (Rays p(or))one(((lleeeoooeed ddlldz)dzzddr.dt
1 L'' Seeep J p((x))p ( ( x 11a ال dedddd
2
+ L."
(eeays / pt (x()) (2o (20 (nd xx |
To summarize, in this section have shown that the kernel score losses that have defined in Section 3 can be expressed in terms of squared MMD. More precisely, we have shown the following result in the specific case where k(x, = - || x'| (but the result remains true for every positive definite kernel).
Proposition D.2: We have that
(()) = (MMD( + C,
with C E R constant independent of 0.
This means that P also minimizes C(0). Note that we could also have derived this result from the properties of the strictly proper scoring rules.
When optimizing with respect to the parameters of the kernel, here the kernel x(xxx') || x'|l does not depend on any parameter, the losses C(0) and (0) differ since now we cannot neglect the additional term appearing in LMMD (0). Unfortunately, this term recalled below is computationally problematic
S. I (aays
ddlldzadz/drddd..
18

---

Distributional Diffusion Models
Indeed estimate of this integral requires sampling =i from the conditional distribution (oje(zo//z)) However this be approximated by using the following self-normalized importance sampling approximation of this distribution
(o)o(2)) Pri[(2 (o)
w u' (2o); w' x Pro Xi)
w =1
(27)
Pr(
where X Lid. diffusion models. However, for small, the distribution of the weights (w,) ) very degenerate and the approximation of (2))( will be poor.
Po- Such approximation was considered by Xu et al. (2023) to derive low variance regression targets in
D.2. Joint Diffusion Energy Loss
Now that we have established the connection between MMD diffusion losses and energy diffusion losses in Proposition D.2, we are going to discuss other possible learning objectives. We recall the joint diffusion energy loss introduced (15)
Ljoint (0) = -.. ut Boo. (S (Xo, X.)]] dt.
and the LMMD. joint loss which is defined using joint distributions as follows
LMMD joint (0) = S. (p Po.
We emphasize again that D, appearing in LonaDjooint is defined over Rd x Rd and not R2 as in LAMD-
Similarly to Proposition D.2, we can show the following result. Recall that we consider = 1 to simplify presentation. Proposition D.3: Let p((20o æx), (1), a()) = p( a'6) + ((11;2)) We have that
Cjoine (0) LMMD (0) + C,
with C P. constant independent 0. Developing foin (0), obtain
1 //' /
LMMD joint (0) = 2 J J(RO)E
x + + - dzdd
L'' S (r(2))) p((c)))
+3ل1 + | - 2 221l dæ
+ 1.' I polx))
=61 + || x da
We recall that the term which does not depend on 0 in LMMD is difficult approximate, the proposed approximation using importance sampling (27) is indeed expected to perform poorly. On the contrary, the term which does not depend on 0 in LMMD.joiit is very easy to approximate. This suggests that if one would optimize the loss with respect to the parameters of the kernel, then the choice LMMDDjoiit is more suitable. However, in the case where the parameters are fixed, as it is the case in our framework, then we will show in Appendix D.5 that LMMD is more sample efficient.
D.3. Marginal Diffusion Energy Loss
Finally, we introduce a last discrepancy L marginal given by
LMMD SMMDDmarrgina = -[( Dolpo po)dt,
19

---

Distributional Diffusion Models
where
9.t (xro)
pe(xr)) (
Po
After few simplifications, we obtain
LNMD. marginal (0) L. / (auy pe plchhmeeddddehhhhhhaahdeeedd)))d - deldddeendee da
1
1
/ [Baj4 pc p (
|
zilldrddz@dee de
2
L Jeey Pe(x;) po (a )p (,) Poj (~ =2
Note that while the minimizers of the conditional and joint losses unique and given by Pole- this is not the case for the marginal loss. Indeed, simply have that minimizer of the marginal loss should satisfy
po(xo) = Jrr Poje (zol (x )p (2t, ddz)
This simply says that we should transport Pt to Po and there is infinite number of such transports.
D.4. Empirical versions of MMD and MMD Joint diffusion losses
In this section we justify the form of (14), i.e., the empirical version L given by (13). For completeness, here we present the empirical version of CMDD which can be expressed as (M = L()) + C, where term which does not depend on 0, according Proposition D.2
We recall that have
1'' I pla)maalddleeeeeedddeal.ll =glldzfdzgdr.dd
LMMD
Jo (saass
(1) - = da 1 L'' Sapp po(x)))) (xb bn (
2
+ L/
S pr
( da
Confinement term. We first focus on the last term of the loss, known as confinement term
C(0) S/¹ J po. (1) df dt.
This can be estimated unbiasedly by
Cn.m 1
1x! - x2 11.
where 4 (n, 1J), (x .tt)) pa., for i E [n] and x2 p (-1x for j E [m].
Prediction Interaction term. We now focus on the prediction interaction term
I I= f(' T-/'/
/ Juee pe po (n (1) =.lddrgderddeed..
This can be approximated unbiasedly by
In.m
1 1
xx x-o 1|.
m(m 1)
'-'==
where 4, Uni 1)), X3, ~ Pt, for E [n] and x pg. (-1x) for j E [m].
20

---

Distribational Diffusion Models
Target interaction term. We finally focus on the larget interaction lerm. This term is harder describe and implement. However, emphasize that it is not needed in order to optimize L given by (13). It is only needed when computing (24). We have
I' = J
(z
- zdlldzjdzjdzzdd
As explained before it can also be written as
po. x/)) m
(o ( (x po (+ 10 - sflldrddefddr dt.
I' =
Jac (ar |
One estimate of this objective given by
Pt. Pt (xg)
It R,m
X
n
Σ k4 Pt, Xi X(x
i=1
j=1
where t Uni 1)), Xb. Xi, Pa,te
D.5. Sample efficiency
Previous subsections assume ut for ease of presentation. We consider general weighting function wy here. Finally, in this section, discuss the sample efficiency of the empirical version of the energy diffusion loss (13) and the joint energy diffusion loss (15). In particular, we prove Proposition 4.2.
U-statistics. We need first to recall few basic results about U-statistics of order 2. Consider symmetric function h: (IR))2 + IR and let
Un =
Σ
X,,,
iid X₂
p.
2
IS Ki
We recall Hoeffding's theorem on the variance of U-statistics, see (Serfling, 2009) for instance.
Proposition D.4 (Hoeffding's theorem): We have
2 Var(U,) n(n 1) 2((n Var [x]) + Var(h(X;,
This implies that
lim Var(U),) (V) (1)).
Another classical result the law of total variance.
Lemma D.5: We have
Var(Y) = Var(E[V X]) + E[Var(Y XX)]
Combining those two results, we obtain the following proposition which central to the rest of the study.
21

---

Distributional Diffusion Models
Proposition D.6: Let X; ltd. for E {l...., n and N. Define also {UiN=1 such that for each i {l...., Us is independent from X, for j # i, [X. are i.d. and U, is a U USstatissics of order 2 with samples. We define
U = Σ U4
Then, we have that
Var(U) = Var(E
X1l) + B[[VarUU [x)]]
n
In particular, we have that
lim Var(U) - Var(BiU1 Xi]).
m - 00
Conditional and joint interaction terms. We consider two interaction terms. Recall that we assume that we are working here under the strong assumption that we are at parameter 0 such that PỘ = Poja- The first one is given by the empirical interaction term in the joint energy diffusion loss
1
Wt 11x Xo || |
X II:
moiutjoiut
mmm ==1==1
lid with t ~ Un, 1I), {x n.m Lid. Jij=1 ~ and x Pt.jo-l1Xx The second one is given by the empirical interaction
term the conditional loss:
n
m
I Inmm
wi, (11 -x- II},
m(m 1)
4-'==
where lid. Uni 1J) and for any i E [n]. X1, Pr. In addition, (x) j conditionally independent, i.e. X Poja, (-I.,)).
We start with the following proposition giving the mean In,m and In.m.joint-
Proposition D.7: We have that
- -(0 J
E a joinnt
wi wallloo x'ol + ||= (00
Similarly, we have that
J/' S wellzo (1)):((())))))))))))))))))))) (2 ddgdd/ddddddd
E E[I
Next, we can obtain the asymptotic variance of those random variables.
Proposition D.8: We have that
lim Var(II.mmjoinn)
1 -Var--nit((o..)) (Ley) wi fll0 -'6l + ||+. -ll.p//o.choeepp.oo/.///daaaaaad
n
Similarly, have that
lim Var(In.m) 1 Vart-Unl((001))xx-rp t000
(/rey welloro 2/ (olloo((<oo (x0ool (ord)ddd
22

---

Distributional Diffusion Models
Proof. This is application of Proposition D.6, with X ~~ Uni 1J) in the of I jointoinn and X (1, Unif([O, 1]) Pt in the case of I, -_
In order to compare the sample efficiency of these different terms, we compute their Signal-to-Noise Ratio (SNR). In particular, we compare
T [[ m. joint 12 my -+++0 Var [I, mn
SNR (nnnnoint)
lim SNR(IInm..jinn
lim
E[In,ml
ST,( = lim SNR(In, lim +-0
m-+000
Var (In,
The larger the SNR the better. In particular, if (T,( Z SNR then claim that the conditional version is more sample efficient than the joint version. Note that both SNR(I,) and SNR(Znn,oimt) are linear functions of N, Therefore if SNR(I,)) 2 SNR(In.join) then we can interpret this as one sample used in the conditional version is worth used in the joint version.
SNR(Z.)/SNR(Z,,
While the general expression for these SNR can be hard to obtain, we are going to analyze the specific case where the target Gaussian, In that case the posterior known and explicit expressions can be derived.
Gaussian case. shown that
We assume that pu N/ IdId) and consider Prjo(za/zo) Naxx a,20 Id). In that case it can be
Po (x) = Nwoo an.o + +o of),,, o²
a (a/o of))ld).
This can be rewritten
q aue + u
Po)e(2o(115) =N
20:
Tt
Td
+ us
where u, = 0./a/oo* Note that lim ue = and lim Uy = too. First, we compute the mean in both cases.
Proposition D.9: We have that
[' wilo - (af0² + 07) )2.
E [mI.io..oine]] = V2I[|1211 =
In addition, we have that
E[In,m] = -20 211 //'
1/2
1 Ut
dl.
Then, we compute the variance terms.
Proposition D.10: We have that
lim Var(In.m.joint)) 2 [1Z]] too
J w o + (a,00 + 07)1/2222 - ( m/ ( (a30² o0)02+0))1/2)a4
Similarly, we have that
[| 2
( Wt
1/2 27
lim Var(In.m) 40000000
wi
2 dt
dt
n
1 + ut
1 us
Combining Proposition D.9 and Proposition D.10, we get the following result.
23

---

Distributional Diffusion Models
Proposition D.11: We have that
(Jo' wi + (afo² + of)1/2)d1)
SNR(In..oln)) = n2 w6 (af0² + 07) (Jo wilo + (afo² o2)) 2'
Similarly, we have
1/2
So We
lit [tuu
dl
SNR(In)
So wi Itu+
dl
So Wt
us I+aa
dt
This concludes the proof of Proposition 4.2. In Figure 2, show that in setting where the weighting is chosen follow the sigmoid one of (Kingma et al., 2021; Hoogeboom et al., 2023), obtain that SNR(In,join) S SNR(I,) for larger range of a.
E. Diffusion-comppatible kernels
In this section, we prove the results of Section 4.3. In particular, we prove Proposition 4.3. We recall that for a continuous negative kernel Pe with E R a hyperparameter of the scoring rule, we say that Pe is diffusion compatible, if there exist ex [-00, +ool and f IR + R such that
lim fcc) = [ Eg [x]] c-
We also recall some of the kernels we use.
kimg (z, x²) (| 22|1 + 1)2
(28a)
kibe (z, x') = exp ||| - x' 1|2 112/20²|
(28b)
kaxp (x. x3) = exp |- x' 11/0].
(28c)
In particular, we prove the following proposition.
Proposition E.1: Assume that - k with lo given by (28a) or (28b). Then the scoring rule is diffusion compatible. In particular, if k is given by (28a), we have that +00 and [(c) = 2c. If given by (28b), we have that +000 and f(c) 2c.
Proof. First, we have that for any continuous positive kernel k with p p= -k, we get
Dp(p.9) = S(g, g) S(p. q)
= Egool (-k(k,
-
(xX,))-((X,)))
- -veol-
2
1
Epssplht (k(,, X')] 3 [k
EEoon Sxnoolk(XX))) (XXXY)
In addition, we have that for any u, u E R
uD,(p. g) Epoy ((((,,))) w)] Epop (uXXXX b)] [ - --)]-
Let be given by (28a), i.e., for any x, x' Rdd, we have that kim (I,, x²) (| 2'12 + 11/2 1/2 Let 0, and
24

---

Distributional Diffusion Models
c-//2 and consider (Ca)nEN such that Gp +00. We have that for any x; x' and
lu(k(z,xx") w) = 20²/2 - - x'1
c)
20112211z 2²1 + c) 2c
(110 x2 + c)1/2
2c 20.1/11100 x'
(1|x 2212 + 1/
2 x'
I (1|æ 2/12//c 1)1/2 S <2 - x'1l.
In addition, we have that for any x; x' and n E N
2 2-2(11x 22113/ / Cn. + 1)
20 2cn (1 r'112 +c,)
(1|= x x²112 2/112/cm + 1)1/2
Therefore, letting + to0 we have that
lim 2/2/22 2c, (1|5r r112 6n) 12/2 = || x 112
1 +-++00
Hence, setting f(c) = 2c, and using the dominated convergence theorem, we get that
1 Epeso IlIx - =1 1 LEpoolll - x-1 Exsoolll - Y'11
lim ((c)),(p.g) +
|||
Eg [X]1?
which concludes the first part of the proof. Next, we consider i be given by (28b), i.e., for any as a' Rdd we have that 1. Then, we have that
kexp(z expl-1|x 22112/00) Let 20 and
lu(k(z,22) v)1 20 expl- æ'1 2=11"/20²) < 1| - x'11
Now consider (cn kne such that Gn - +00. We have that
lim 20 (1 expl- 2' 2/12/20111 =| 2/12? 2-+00
Hence, setting (c) = 2c, and using the dominated convergence theorem, we get that
f(c) = Ep ||| - Y1 マ E x'11 } R [|| - y'
lim ++0
2
|| [X]]
which concludes the second part of the proof.
Finally, we highlight that while the exponential kernel (28c) is not diffusion compatible it satisfies another compatibility rule. More precisely, we can show the following result. The proof is similar to the proof of Proposition E. in the case of kybr-
Proposition E.2: Assume that = -k with k given by (28c), ie.. M(x,x') = exp|-I|x x x'11/o). In that case, upon setting f(c) = e and e" = +00 we get that
lim (c)) = (p.9),
where p² (x, x!) || I'Il, iee. recover the energy distance with 1.
F. DDIM updates from a Stochastic Differential Equation perspective
In this section, we show that one can exactly recover the DDIM updates from the stochastic process point of view with careful choice of forward process.
We first show that different choices of forward process yields different interpolation densities P=jo.; in Appendix F.1. With a careful choice of the forward process and appropriate parameters, we show that we can recover DDIM updates (Song et al., 2021a) Appendix F.2.
25

---

Distributional Diffusion Models
F.1. General forward processes
In this section, we consider general forward processes. The main property of all those forward processes (X₂ tE)000011 is that they will satisfy
X, = axXu + t
X ~ po:
(29)
with Z' ~ (O, Id). First, we define
=
We also consider Z ~ (O, Id) and introduce the dynamics (Xt)re(o.1)
dx, 0, + 9 (1 5²)Z dt + eg. 20,
(30)
where (Be)tejo.1) is a d-dimensional Brownian motion. In (30), we identify different forward processes which are all non-Markov except in the case where 1. The remarkable property of these forward processes is that they all recover the interpolation (29). This means that the forward trajectories might be very different but they all admit the same marginal distributions. More formally, get the following result.
Proposition F1: For any E [0, 1). let ex E [0, 1]. Ler Z NN Id) and Xo po- Additionally, assume that the following SDE admits solution
g6 (1 =2) dt + Eg.
dX.
i log(a))) +
20,
where (Br)eejo.1 is d-dimensional Brownian motion, p(z) = (z; Id) and pr (-12) is the density of Xt conditionally to Z = 2. Denote Pr(x.) Jg (((22)))p(2)dd. We have that for any E [0, 11. a, + o,Z ~ pt.
Explicit integration.
We can solve exactly (30) and its solution is given for any 2 s by
a 0²
1/2
X,
x, + -1-
oo ax
Z
+
0,
z
as
a
with 2 ~ Id) This can also be rewritten as
1/2
X, an = at (X, - a,X +(1 +1- 62/1/2 01
at 0s
0i 2
Z +
0i
2
as
02
1/2
d+ [=1 - 12200220022 + +00 + -
E1 2 1/2
0t
dtt as
Z + 8
00
a
2
0
a,
1/2
o,Z + 80 = (1 11 1/2
at 7 + 8 0, ai o
7,
(31)
a2
a
with X, = ax + (1 - 1) + 0, with Z, 2 and Ż independent Gaussian random variables with zero mean and
identity covariance. We can now compute the probability distribution Ps.sju- This will allow us to compute Pajo,e using Gaussian posteriors. Using (31), we have
(X,, X,) ~ ((xxX))
(o?ld C (axXo)
(
C2 ofld)
with
C.
2 02
0'
(1
52) 01
22
2 04 o² + (1 =?)o,o, ld.
0+
a2
26

---

Distributional Diffusion Models
Now, let compute the obtained posterior.
2 at 0
E X Xo. X = ax
+
+ - =
(X, - -x
Os 11
0z
22 ar o
(1 - c2) 00
X, + 0,
1
_ a o²
+ (1 5²)
0x at
Xo.
a₂ oi
ag di
Finally, we have that
2
C,, Xo) obId
E² at 0 + (1 E² )0)00
o bld
ld
ax
02
_² ar 0s + - =2) (#)
Ogg Ft
Note that in the special case where 0, we recover deterministic sampler. To summarize, we have the following result which justifies (4).
Proposition F.2: Let [0, 1] and (Xe)ceju,1) given by Proposition F.1. Then, we have that for any s, E [0, 1] with
S
X, = (62 +11
+ + r2,2((8,)) 1) - (1 - (1-62)m,, (s,,)),,() +o - (E - (1 -1- (1)
(32)
with 2 ~ ((0, and where
(8, a o
a of
Limiting behavior. Looking at (32) one can wonder if we recover the SDE framework when we let s -> t. This is indeed the case as established below. We will make use of the following result
lim (t,s(s,) t) - i0, log(as) - jo: log(or).
t t- 8
Therefore, get that taking the limit in (32) have
dx, [[e? (O: log(ar) 20, log(or)) (1 - log(o))))
(E0) log(ar) 20, log(o,)) + (1 - (8, log(as) - a log(o:))) dt
+ + (0, log(as) 0 dBB
can easily be shown that In addition, we have
20, log(as) - 0, log(o.)) 20,0)0((0r/00)) = 96-
(33)
dx, =
-0) -0
- O log(o:)))X,
- a log(o,))) X)o] dt -arls" log(az) - 20, log(o:)) (1 (2)
- 0, ( log(ar) 20, log(a;)) = (
+
+(1 - =3 (0, log(ar) O log(o,)))))))))) l axXo)] dt + g.
- log(a:) Xx + (1+5 += log(ar) 8 log( (X, a:Xo)] dt sg.
27

---

Distributional Diffusion Models
Now, combining this result and (33) get that
1 1+++* axxo X dX,- = -ù log(ar) ) + gi
dl + eg.
of
Taking the conditional expectation in this expression, we recover the score logp: and therefore
+ s² dx, - log(ar) + g6 2 Vlogp: di + eg.dB.
Therefore, we recover the usual sampler derived from the stochastic process point of view by taking the limit s + f in (32).
F.2. Connection with DDIM
We recall that DDIM (Song et al., 2021a), the authors assume that the schedule (ate, σz) satisfies až + of 1. this case, they consider
p(z,|00. xx) = N(xxs (E. Iz), En,);
with
- a² n
a -sat
(0,0 =
Dt
+
0ts - 001
300>
Da. ld.
(34)
Note that in (Song et al., 2021a), as is replaced with or and Na,t is denoted de Recall that in our introduction of diffusion models in Section 2.1, see (4) we have
(s, 21) (=9ri,2(s, -1-
+ (1 (2,2 (1=,))).,(s,t)) Eet = o1 - (E (e2 + (1 - =
(35)
The following result proves that the updates DDIM (34) and the ones we present (35) are identical upon identification of one parameter.
Proposition F3: Assume that v 101--632 (1 Then have that E,,e E.,, and u x.)) x1).
Proof. First, have that
a
1/,2
- a3
- (s2r1, -02(1
+ +1 (1-=)))33 +²))²)
of
of
1 a3 - (1- (1-ag)(1 - (s) (1 (1 =?)))
1 1- o
/ruu(s,)( -1 (1- +1-
=ro.1ss,))637 -11 (1-
- = 1) (I -1-=²)m (s,
28

---

Distributional Diffusion Models
Hence, we get that
0
2
1
1 og
1
2
It +
02 /
V
To
a -
2
02 -
Nal
2
It + a
1 -,,,),(,,t)
T0
on
ai
(e + (1- (o, + +a [1 (222)(8. (11---1---1)))0 (1-=2)ma((stt)) 20 Mo,e Iz).
which concludes the proof.
In our work, we call E the churn parameter. In (Song et al., 2021a, Equation (16)), we have that
0 of Nstt
1 -
2
1
(ro,,,2(s 22)
-22.(s)))))
In that case 7j is another churn parameter. However, we write the following relation between those parameters. In particular, in both cases we have that we recover if DDPM if e (or equivalently if 1) and DDIM if e 0 (or equivalently if = 0).
Proposition F::: We have that
1/2 o2(1 (s, (1- ro,2(8, 1) - (2,2(8,1)
Similarly, we have that
17 20.2(()))-222((11)) 1/2 03
8² =
(8,1
Remarkably, we have that = and only in = 0 and e = 1 if and only in 1.
Proof. Proving that 0 if and only if 1 0 is straightforward and left to the reader. Now assume that e 1, we are going to prove that 1.
n/² 03 (s (1-
- 03.22,22081)
(o, - (s)
(1 (--a3)(1-ff )a)aa (11
a²(1-a;
To,2(s, t) 2
(1-a3)(a²-a) a?(1-af)
ro,2(,, r2,2(8,) 1)
To,2(s, 1)(1 - 12.0(8,)) 1))
ro,2(8,1) - (22((,t) 1.
Similarly, we get that implies that 1.
G. Pseudo-code for loss function
G.I. Some replication utils
Given batch of [xr1, x2 24, ". assuming that n = 4 and m = 2, this function will output (before the reshape operation)
29

---

Distributional Diffusion Models
x
x
f
def replicate_ff (n: int, m: int, x: chex.Arraye -> chex.Array:
batch_size, data_shape = x. sh)p x. shape [1:]
x[in]
x jnp. reshape (x, (n, 1, sdata_shape))
x jnp hap.tilelin (x, (1, m) + (1,) * 1en shape_
x jnp. reshape (x, (n +data_shape))
return x
We also assume that we are provided function split_fn such that given æ with shape (nm,...) the output of split_fn has shape (n, m....
G.2. Loss function
We assume that we are given two functions cn and and y with shape (n, m.. the output of compute.rro..daggnaa..
such that given æ E Ey= P(Zi.) Vk.) and the output =jj' p (1 V.y Function is consistent with the notation defined of compute e rho fn is E mmm I] > m in Section 2.2.
In addition, we assume that we have access to a function add.noise se_fn such that given of shape (n,) and zo of shape (n....) it outputs 22 0y + 0₂ where 2 has the same shape as x and has independent Gaussian entries WV 1).
Finally, we assume that we have access function apply- such that given of shape (n,) and it of shape (n,...) outputs 2t).
def Loss_fn(t: Anp.Arryy, x0: jnp.Array) -> jnp.Array:
# add noise and replicate
add_noise_fn (key=key, t=t, x0=x0)
key, = random www.eddm.....llioooooeeeeeeeoite pplit(key) x0 x0_population = replicate_fn (n-n, m-m, x=x0) t_population feplicate_fn (n-n, st__oopuatiinn replicate in (n=n, I= x=xt]
x=t)
10
# compute prediction
11
eps_population - jax. random. normal (key-key, shapexx__ppplaationnsheee shape)
key, jax. random. split (key)
13
output_population apply fn (t=t_population, xt-xt_population, eps=epa_populationn
14
15
# split the populations
x0 populaticn = split_ fn (n=n, m=m, x=x0 poo_pooplaaion
output_population = split_fn (n=n, m=n, xroutput_populationn)
18
19
# compute confinement term confinemen
20
crmputeerth_ddaagnnll fn (aapoaalfffxxx0__ppp__
y=output_popuiation
compute interaction interaction_prediieet prediction compute_rie shattee__ro____e www.ouuttuutpppuuttttteeep proutput_populatiice
for predictions
23
24
# Generalized kernel score
score 0.5 1bda interaction confinement
27
28
I Our aim to maximize the score, so the loss is negative score
29
loss scooe
return loss
30

---

Distributional Diffusion Models
In this function, have focused the preprcccition and have omitted the weighting for simplicity.
H. Extended related work
Diffusion and GANs. In (Xiao et al., 2022), the authors propose to replace the conditional mean estimator of diffusion models with a GAN. This allows to model bigger steps in the denoising process. Recall that, seeing diffusion models as hierarchical VAEs, the ELBO given by
log pg(p)()
Σ 1 (pp + C,
where C is a constant. In (Xiao et al., 2022), the relative entropy (p0 replaced by an adversarial loss. Namely, the authors consider
NV. (Fprre 1 1 (ep (Da(t Ita-1 1 Ftt))] ( + Epo (wr2 (1r || log(1 (ts ((m))))) awwll)
()
k=1
In that case, D& discriminator and pa is generator. The main differences with our setting is that in our case, we do not train discriminator. In our setting the discriminative parameters only given by the kernel parameters, i.c., in the case of the energy distance, the parameters are given by À and 9 in (13) and are fixed. In this regards our parameterization is much more lightweight. In addition, while the transition po ) is estimated in similar way as in our paper, i.e.,
pa(Ste | - | po (o )p(re, J Ft,, )dro;
the discriminator is not applied on the so variable conditionally to Tt, but on the couple (Iea Tte .). This is similar to the differences we highlighted Appendix D.5 when comparing joint loss and conditional loss. Note that our approach could be extended to consider kernels not on Pojs but instead leveraging the characterisation
p(=,(=))) = [ plzulll (( )p(x) ( Ie dd..
for u S mi)n(,)). In that case, choosing 0 is simply special case.
Modeling the covariance. The importance of modeling Poje in the low step regime has been highlighted in numerous papers, see (Ho et al., 2020; Nichol & Dhariwal, 2021; Bao et al., 2022a;b; Ou et al., 2024). A number of approaches aim at leaminng/pproxxmatingg the covariance of Pajt or Pojr. see (Ho et al., 2020; Nichol & Dhariwal, 2021; Bao et al., 2022a;b; Ou al., 2024; Rozet al., 2024). Most of the time in order to obtain tractable approximation of this covariance they consider diagonal approximations leveraging the Hutchinson trace estimator (Hutchinson, 1989). In contrast, we do not assume Gaussian form for Poj: or Pstt- Doing so we model more complex distributions than Gaussian transitions (ii) avoid the additional complexity of having to model covariance matrix.
Relationship to works using MMD. Scoring rules have corresponding MMD divergences, and approaches exist that use the MMD in diffusion and particle flow models. Galashov et al. (2025) generates sequence of distributions from a forward diffusion process. It then generates noise-dependent neural MMDs between clean and noisy data, and performs gradient flow to move particles from noise level / to lower noise level 8 < f. Compared to ours, Galashov et al. (2025) does not use generator since a particle flow. Their best CIFAR-I0 FID 7.7. Aiello al. (2024) first trains a diffusion model, and then refines/distills by coarsening the reverse timesteps, and uses MMD on CLIP features to finetune the DDM denoiser. Their CIFARIO FID for NFE- 10 improves from 13.6 to 3.8, while ours is 3.19.
I. Experimental details
L.1. 2D experiments
We consider a target distribution po given by a mixture of two Gaussians, Po = 0.5N o21d) + (,5, olld), where #1 = (3,3), //2 (-3,3) 3) and 0.5. We create a dataset from this distribution by sampling 102400 points.
31

---

Distributional Diffusson Models
We train diffusion model and distributional diffusion models for 100k steps with batch size 128 with learning rate le 3 using Adam optimizer with a cosine warmup for first 100 iterations. We use b 0.9, b2 0.999, 1e 8 in Adam optimizer. On top of that, we clip the updates by their global norm (with the maximal norm being 1). We use EMA decay of 0.99. We use the flow matching noise schedule (3) and we use safety epsilon 1e - 2. When training the model, we use sigmoid loss weighting as in (Kingma et al., 2021). The time is encoded via sinusoidal embedding into the dimension 2048 followed by layer MLP where hidden dimension is 2048 and the output dimension is 2048, using gelu activation. As a backbone architecture, we use a 9-layers MLP with gelu lu activation and with hidden dimension of 64. The MLP is applied as follows. First, we apply 4 preprocessing MLP layers to the embedding of time and separately to the zt. After that, we concatenate these two on the last dimension and pass through 4 MLP layers which is then followed by one output MLP layer of dimension 2. When we use the distributional diffusion model, the noise & has the same dimensionality Tt and we concatenate the two along the last dimension. After passing it through the MIP, which produces a vector of dimension 4, we ignore first 2 dimension to get the output žg (t, Tt. <) of dimension 2. We use the population size m = 32 (see Algorithm I) and we use 3 = 0.1, À = for distributional model and / = 2, À = for diffusion-lis variant.
For evaluation, we sample 4096 samples X3 ~ Po as well as samples X3 from the forward process (3). For Figure 2, right, we use the MMD squared given by D, (10) with x() = - (z,z²) 2') (see (9b)) for = 1.
For each X7, we produce 8 samples & ~ (V, 0, Id) and we compute standard deviation of (t,t, over Ş. Then, we average the standard deviation over all the s... On top of that, since we know the posterior pl ((Xo|X)), Xx), which is given by
p(uo xx) = w/N(woii vI, Σ) + w2N v2, E), wk x N (aj0² + of)ld),
0+
1k = Σ
Xt
+ oz 1k)
of
-1
ai
Σ =
+ oi
ld.
02
where a with k = 1, 2 are the means of po and 0² the variance. The values of a and or are given by the schedule (3). We produce 8 samples X, ~ p(-| and compute their standard deviation and then average over X We report this metric in Figure 3, right. In order to produce Figure 4, we fix zo to be equal to either μl to μ2z We then produce Xz/20 for each I (one sample for each l and uo). After that, we visualize io (t, X3, E) for n 4096 samples of g ~ (O( Id).
1.2. Image space experiments
Architecture. In our image space experiments consider a U-net architecture (Ronneberger et al., 2015). We consider channel of size 256 and apply the following multiplier for the channel at each level of the Uenet, 2,, 2). At each level we consider 2 residual blocks. We apply an attention layer at the second level of the U-net (resolution 16 in the case of CIFAR-10 and resolution 32 in the case of CelebA and LSUN Bedrooms). We use an attention block in the bottleneck block of the Untt Each residual block consists of normalizing the input, applying a 3 x 3 convolution block and applying a non- non-linearity. This is then followed by dropout layer and a x 3 convolution block. Finally we add the input (processed through 1 x1 1 convolutional block) of the residual block to this output (this is the residual connection). We use the RMSnorm for the normalization layer and GELU for the non-linearity. When using an attention layer, this is done after the convolutional residual block just described. We use multi-head attention mechanism with 8 heads. We process the time with a sinusoidal embedding followed by MLP layer. In the case of CelebA, we use a linear embedding layer to process the conditioning vector of shape (40,) ) while in the case of CIFAR10 we consider an embedding matrix to process the class information. Once the time conditioning and the (optional) other conditioning have been obtained we sum them. The conditioning is used in the model by replacing the RMSnorm with an adaptive RMMnorm layer, i.e., the scale and bias of the normalization layers are obtained by passing the conditioning vector through a MLP layer. Similarly to the two dimensional setting, we concatenate 22 of shape (b, h, w., c) and & of shape (b, h, w. c) along the last dimension in order to get an input of shape of shape (b, h, w, 2c). The output of the U-net is of shape of shape (b, h, w, 2c) and we drop half of the channels to get a sample of shape (b, h, wr, c).
For latent CelebA-HQ, we follow the recipe of (Rombach et al., 2022) to train the autoencoder. We train the autoencoder with a similar architecture as the autoencoder used in LDM-4 (see (Rombach et al., 2022)). It encodes CelebA-HO -HQ of shape (256, 256, 3) into latent code of shape (64, 64, 3). Then, for latent diffusion model and distributional latent model, we use a
32

---

Distributional Diffusion Models
similar architecture as LDM-4 except that we concatenate the latent dy and & over the last dimension. More precisely, we consider a channel of size 256 and apply the following multiplier for the channel at each level of the U-net, 2, At each level we consider 2 residual blocks. We apply an attention layer at the foruth level of the U-net. The structure of the residual block is identical as in image space. Similarly, we use the same type of normalization layer and non-linearity. We use a multi-head attention mechanism with 8 heads.
Training and evaluation details. We train all the pixel-space models for 1e6 steps with batch size 256 for CIFAR-10. batch size of 64 for CelebA, LSUN Bedrooms. The latent space model consists of two-stage training regime, where we first train latent autoencoder for 2e6 steps with batch size 16. It is then followed by training a latent diffusion model for 2e6 steps with batch size 16.
For all the experiments, we use the Adam optimizer with additional norm clipping (maximal norm equal to 1) with the learning rate equal to le-4 and an additional cosine warmup for 100 first iterations.
Our autoencoder on CelebA-HQ trained with 0.1 coefficient on adversarial loss, 100 coefficient on generator loss and 100 coefficient codebook, from the VQ-VAE (van den Oord et al., 2017) loss. We use 3- 3-VAE with | = 1e - 6.
For all the models, we sweep over weighting w/ (see (6), (13)), where we either do not use any weighting or we use sigmoid weighting, see (Kingma et al., 2021), weighting with bias parameters {-2, -1, 0, 1, 2}. When we train classical diffusion model, we always use velocity as prediction target, i.e., the target is Xn X,. When we train distributional model, we use Xo as prediction target. For distributional models, we additionally sweep over À E 0.1, 0.5, 1.0), E (0.0001, 0.001, 0.01,, 0.1, 0.5, 1.0, 1,,, 1.99. 2.0). We use population size m = 4. We also tried m but it led to slightly worse results (the resulting FID was 7% worse). We did not use higher population sizes since it led to higher computational and memory cost (see Appendix J for more details). When and Section 3, for (imq) kernel (9a), we sweep over E {0.01, 0.1, 0.3, 0.5, 1.0, 1.5, 2.0, 10.0, 50.0, 100.03. For for (rbf) kernel (9b) and for (exp) kernel (9c), we firstly compute median M = ||2 - z' z'112 for (rbf) kernel and M = || x' 12 for (exp) kernel, on a subset of 100 batches of size 512 from the dataset. Then, for (rbf) kernel, we define x² = M²y and we sweep over y E (0.01, 0.1, 0.3, 1.0, 1,,, 10.03. For (exp) kernel, we define o = My and we sweep over y E {0.01, 0.1, 0.3, 0.5, 1.0, 1.5, 2.0, 10.03. Since we only use kernel diffusion losses on CIFAR-10, we only report the value of the medians in this case. We have that 2 1357. 8127 and M 36.9.
kernel diffusion losses, see Section 2.2
In order to produce samples, we follow Algorithm 2 for distributional models and SDE sampler for classical diffusion models. To evaluate performance, we use FID (Heusel et al., 2017) metric. As we train the models, every 10000 steps we produce samples with 10 denoising steps for distributional models and with 100 denoising steps for classical diffusion model. We then evaluate the performance of the models by computing FID on a subset of 2048 datapoints from the original dataset. Then, we use this metric in order to select the best hyperparameters and the best checkpoints for each best hyperparameter. After training is done, for every combination of X, 3 (or c, a, o² in case of other kernels), we select the best checkpoint based on their FID on this small subset of data. The corresponding checkpoints are then used for evaluation when sampling with 10, 30, 50, 70, 90, 100 denoising steps. To compute the final FIDs, we use 50000 samples for CIFAR-10, LSUN Bedrooms and CelebA. We use 30000 samples for latent CelebA-HQ.
I.3. Robotics experiments
We used the Libero benchmark (Liu al., 2024), a life-long learning benchmark that consists of 130 language-conditionee robotic manipulation tasks. There are 5 suites in the Libero benchmark, Libero-Spatial, Libero-Object, Libero-Goal, Libero-Long and Libero-90. Each suite contains 10 tasks, except Libero-90 which has 90 tasks. In our experiments we focus on Libero-Long, the most challenging suite with 10 tasks that features long-horizon tasks with diverse object interactions, and reproduce the main experiment settings on three other suites: Libero-Spatial, Libero-Object, Libero-Goal, with 10 tasks. We train language-conditionee multitask diffusion policies for the 10 tasks of Libero 10 suite and evaluate the success rates of the policies in simulation. We encode the visual observations using a ResNet (He et al., 2016). and encode the language instruction using pretrained and frozen Bert (Devlin et al., 2018) encoder. Our multitask diffusion policy does action chunking (Chi et al., 2023), and predicts and executes chunk of 8 7-dimensional actions conditioned on the language instruction, the current visual observations and proprioceptive states. For the denoiser, we follow Aloha unleashed (Zhao et al., 2024) and use a cross-attention based transformer denoiser.
We follow the original Libero paper for training and evaluating our multitask diffusion policy. As a loss function for training the diffusion policy, we use the Mean Squared Error (MSE). Libero suite has 138090 steps of training data and we use a
33

---

Distributional Diffusion Models
batch size of 32 train policies for 250k steps (roughly 50 epochs). Evaluation on simulated environment for 50 episodes per task following the Libero benchmark. We evaluate policy every 25k steps (roughly 5 epochs). We take the median success rate of the policy across all tasks in the suite and repeat the experiment over 3 seeds. We report the median success rates over all tasks and best checkpoint during training. We use the same set for other three suites and train for 50 epochs and evaluate every 5 epochs.
Energy distance difflusion model. For the energy distance diffusion variant, we concatenated noise dimensions, sampled from (0; Id) to the action dimension of the chunks. For training we used population size m 16 per data sample, and computed the loss with varying Bs and Äs. For evaluation, use a single sample and execute the predicted action chunk from the diffusion policy using a varying number of diffusion steps (2, 50).
J. Computational complexity
Our method has the same computational complexity during sampling as ordinary diffusion models. However, it has an increased computational complexity during training which detailed below.
Assume that the computational complexity of the forward pass in our diffusion model O(F) and that the dimensionality of x is D.
Diffusion loss case. The loss function in (14) with m = 1 and À 0, can be thought of as a standard diffusion model loss function. It requires OnnF) evaluations to compute (44,,, Xx,) for every element of a batch. Then, to evaluate the loss (the norm), the complexity is OnnD). Therefore the total cost of computing the loss is (nn(F' + D)). The backwards pass is proportional to the forward pass and has computational complexity OnE)).
Distribution loss case. For À + 0 and m 1, first need (nmmF) function evaluations compute 2 X b Er.s) need (mm. operations compute the first (diffusion-
for i 1..... and for Er.j. j ....., m. Additionally, like) terms of the loss. Then, in order to compute the interaction terms, need Onnm(m 1)D) time. Therefore, the total cost is of computing the loss is (mm ( + D) + nm(m D).). Naively, the backwards pass will take Ommm++ nm - P) = O(mm² F) time. However, since the 2nd term uses the gradients of x (0, Xt,, Er.j), one could precompute these for all = I,..., n and j = I,..., m and therefore decrease the total backwards cost to OnmFF. We found that practice, the XLA compiler in JAX performs this optimization without explicit coding, see our results below.
vary A GPU (40Gb of memory) with batch size 16, H100 GPU (80Gb of memory) with batch size 64 and TPUv5p (95Gb of memory) with batch size 64 (per device with 4 devices in total). The results below indicate that steps per second decreases proportionally to m.
Runs on real hardware. We compared the training times real hardware. We study the training of CIFAR-10 full forward+backward step. As hardware,
m. We report steps per second metric where step corresponds
Hardware
Diffusion
Distributional (m=2)
Distributional (m=4)
Distributional (m=8)
A100 (batch size =16) H100 (batch size =64)
9.05
6.78
4.6
2.77
14.3
7.85
4. 15
2.06
TPUv5 (batch size =64)
11.2
8.5
4.3
2.22
Table Steps per seconds on real hardware for different models.
K. Additional experiments
K.1. Additional 2D experiments
We trained an unconditional standard diffusion model and our distributional models on a more complex 2D distribution checkerboard, see Galashov al. (2025). We report MMD between sampled and target distributions with different NFEs, smaller MMD in bold, in Table 2 Our approach outperforms standard diffusion for small NFEs.
K.2. Additional image experiments
In this section, we provide additional results for image space experiments.
34

---

Distributional Diffusion Models
NFEs
Diffusion
Distributional
5
7.00e-3
1.57e-4
10
2.10e-3
1.96e-4
50
2.70e-4
1.55e-4
100
1.96e-4
2.00e-4
1000
1.71e-4
1.83c-4
Table 2. MMD between target and sampled distribution for checkerboard dataset.
steps
Num steps
Num steps
steps -100
25
To
14
24
12
55
F4
2:4
14
TE
83
24
18
2B
24
13
24
24
2.
64
M
-45
24
43
18
98
34
13
23
50
3.7
21
001
11
18
12
14
13
-35
21
81
64
61 14
85
83
15
24
1.4
Figure 10. Heatmap of FIDs for conditional image generation CIFAR-10. a--xis corresponds X, y-axis corresponds to B, the color corresponds to FID going from minimal (blue) to maximal (red) values. Different columns denote number diffusion steps. Black squares correspond the cases where performance reaches FID higher than specified threshold.
Impact of parameters 33 and À on FID. We report detailed performance of distributional models èg based on chosen and X for different number of diffusion steps. We present results for conditional image generation on CIFAR-10 in Figure 10 and on CelebA Figure 11. The unconditional image generation for LSUN Bedrooms are given in Figure 12 and for latent CelebA-HQ in Figure 13. In practice, we noticed that performance of models for some À and led to very high FIDs. In order to better visually inspect the results, mask the corresponding blocks if FID reaches certain threshold. We used following thresholds: 9.8 for CIFAR-10, 20 for CelebA and LSUN Bedrooms, 300 for latent CelebA HQ. The results on CIFAR-10 and CelebA suggest that for a small number of steps, the best values of parameters are located in the bottom right corner, i.e. A close to 1.0 and /3 small. As we increase the number of steps, better values appear in the top left corner, i.e., À close or equal to 0 and / r 2. In case of unconditional generation, we noticed generally that using the minimal possible led the best results, as can be seen from Figure and Figure 13, where the bottom row leads to the best performance.
Detailed results with different kernels. We present results on conditional image generation on CIFAR-10 with different kernels. First, in Figure 14, we show the value of the X for every kernel which achieves the lowest FID for the given number of diffusion steps. Overall, we observe that À follows a downward trend as we increase the number of diffusion steps, which consistent with out experiments in the main paper. Furthermore, we show the corresponding value of the kernel parameter which achieves the lowest FID for given number of diffusion steps. Below, we also show the detailed figures of FID for every À and kernel parameter combination.
Num steps
10
Num steps 30
Num steps 50
Num steps 100
53
24
12
253
157
5A
33
43
11
355
23
31
37
FID
FID
05
245
x-3
40
43
43
39
39
062
12.
85
42
46
32
37
39
05
31 81
45
14
Figure 11. Hcatmap of FIDs for conditional image generation on CelebA, z-axis corresponds to X, y-axis corresponds to 8, the color corresponds FID going from minimal (blue) to maximal (red) values. Different columns denote number of diffusion steps. Black squares correspond to the cases where performance reaches FID higher than specified threshold.
35

---

Distributional Diffusion Models
Num steps 10
Num steps
30
Num steps
50
Num steps
100
13
11:2
313
231
14.7
61
59
3L9
5.7
E5
119
49
F4
384
10
13:4
18.3
13.8
9:4
14
83
0.00
L0
Figure Heatmap FIDs unconditional image generation on LSUN Bedrooms. æ-axís corresponds X, y-axis corresponds the color corresponds to FID going from minimal (blue) to maximal (red) values. Different columns denote number of diffusion steps. Black squares correspond to cases where performance reaches FID higher than specified threshold.
Num steps
10
Num steps 30
Num steps 50
Num steps
100
224
16.4
25.
213
11.4
SA
113
1L.
LAS
AA
11.2
x7
253
10.3
36.
21.0
11-2
112
31.4
11.1
A
À
Figure Heatmap of FIDs for unconditional image generation Latent CelebA-HQ. z-axis corresponds X, y-axis corresponds the color corresponds to FID going from minimal (blue) to maximal (red) values. Different columns denote number of diffusion steps. Black squares correspond to cases where performance reaches FID higher than specified threshold.
Best A
Inverse MQ
Exponential
RBF
Boite
Best o²
58-
06-
100
100
166
106
Num diffusion steps
Num diffusion steps
Num diffusion steps
Num diffusion steps
Figure 14. Best parameters À and corresponding kernel parameters for each kernel, achieving the lowest FID for the given number of diffusion steps.
36

---

Distributional Diffusion Models
Best
1
A- 0,1
A-
A=10
Num diffusiin steps
N.m diftusionn staps
kur diffucen steps
hur difiusden steps
Nam diftuston steps
Nam diftusion steps
Inverse MQ
Exponential
RBF
Diffusion
Figure 15. Performance of different kernels conditional image generation CIFAR-IO, 10, for different values of À (different columns). The kernel parameters are chosen to minimize FID for the given À condition. We also report performance of the diffusion model (red) and show in the black dashed line the performance of difffusion model with 100 steps.
Num steps
10
Num
30
Num steps 50
Num steps 100
22
36.
31
83
44
37
3#
33
X5
34
25
23
32
39
23
22
32
44
28
FID
99
35
85
48
97
27
23
2.5
44
65
57
0
64
54
81-
A22
10:1
16.
À
À
Figure 16. Heatmap of FIDs for conditional image generation CIFAR-10 for Inverse MQ (9a) kernel. æ-axis corresponds to X, y-axis corresponds to c, the color corresponds to FID going from minimal (blue) to maximal (red) values. Different columns denote number of diffusion steps. Black squares correspond the cases where performance reaches FID higher than specified threshold.
On top of that, we compare performance of different kernels for fixed value of X. We also report the performance of energy kernel, which is defined by p(z, x²) || x'11 see also Section 2.2. The results are given in Figure 15. We see that overall using different (to energy) kernels leads similar results overall.
Finally, we present detailed results of FIDs for different kernels with different parameters À and corresponding kernel parameters. We report the corresponding heatmaps. Whenever the value of FID is larger than a threshold, we put black rectangles. For Inverse MQ, we use threshold 13.3. For RBF kernel, we use threshold 10. For Exponential kernel, we use threshold 10. For Inverse MQ, the results are given in Figure 16. We observe that overall the best FIDs are achieved for larger values of parameter e for all three kernels we investigate. Moreover, when number of diffusion steps increases, the region with the best FIDs moves from À R 0 to À 1. The results for RBF kernel are given in Figure 17. We observe that in case of small number of diffusion steps, the region of the best values is around X = 1 and smaller values of o². As we increase the number of diffusion steps, we generally observe that many values of À and Ø² yield similar reasonable results. Results for Exponential kernel are given in Figure 18. When the number of steps is low, the best values are achieved for À close to 1. As we increase the value of X, the best region moves closer to smaller values of X. This is coherent with our theoretical and methodological motivation, i.c., confirms our experiments when using the energy kernel p(x, x') = x' I| see Figure 10.
is beneficial in the low number of diffusion steps regime, and
steps
Num steps
30
steps
Num 100
13518.13
459
425
2V
13
18
12
13
3
27
3.3
41
24
43
61
66
54
29
14
231
23
24
26
2A4
34
22
32
24
32
40.34
135
65
66
Figure 17. Hcatmap of FIDs for conditional image generation on CIFAR-10 for RBF (9b) kernel. æ-axis corresponds to X, y-axis corresponds to o², the color corresponds to FID going from minimal (blue) maximal (red) values. Different columns denote number diffusion steps. Black squares correspond to the cases where performance reaches FID higher than specified threshold.
37

---

Distributional Diffusion Models
Num steps
Num steps
30
Num steps
50
Num steps
100
24
24
27
24
32
23
24
23
44
2.4
42
24 3.9
244
30
5.0
27 2. 32
54
24
25
35
61
64 64 Lo
65 L8
L8
Figure 18. Heatmap of FIDs for conditional image generation on CIFAR-10 for Exponential (9c) kernel. z-axis corresponds to A, y-axis corresponds o, the color corresponds to FID going from minimal (blue) maximal (red) values. Different columns denote number of diffusion steps. Black squares correspond to the cases where performance reaches FID higher than specified threshold.
ImageNet. We ran initial experiments on ImageNet (Russakovsky et al., 2015) with resolution 64 x 64 x 3, but so far were unable to obtain satisfactory results. We will focus on scaling our method to large datasets in the follow-up work.
K.3. Comparisons to distillation and numerical solvers
In this section, compare approach with the DPM-solver++ method Lu et al. (2022) and the state-of-the-arr multi-step distillation method, moment-matching distillation (Salimans et al., 2024).
For all the sampling methods, we sweep over the "safety" parameter which specifies the time interval to be [7, - v)- We considered values {0.1, 0.001, 0.00013. Moreover, for all the methods we also sweep over the "churn" parameter e which controls the stochasticity, see Equation (4), we tried {0.0, 0.25, 0.5, 0.75, 1.0]. We found the stochastic version of DPM-solver++ (with churn parameter close to overall worked the best. For ordinary diffusion and our method, we use "uniform" time schedule. For DPM-solver++ best results (we also tried EDM and uniform).
the "logsnr" time schedule (see Lu al. (2022)), which led the
For our distillation comparisons, we distill pre-trained diffusion models into student models. These student models are trained for 100,000 iterations, using batch size of 256 for pixel-space models and 32 for latent-space models. During training, we sweep over the learning rate and Exponential Moving Average (EMA) decay.
The student model utilizes stochastic DDIM sampler with churn parameter, E. We sweep over e values in the set {0.0, 0.5, 1.0]. We found that sampling time, setting the churn parameter 0 for the distilled student yielded the best results.
For all the methods, we present the results with the best hyperparameters in the tables below (Table 3,Table 4, 4,Table 5, 5,Table 6).
Table 3. CIFAR-10 (Conditional) results. The metric is FID averaged across random seeds. We report results with different NFEs for haseline diffusion model, DPM++, Distributional (ours) and moment matching distillation. In bold, we highlight the method with the lowest FID and in italic, we highlight the method with the second lowest FID.
NFEs
Diffusion
DPM++
Distributional
Distillation
80.77
19.41
29.33
5.19
23.31
12.60
4.67
3.84
7.53
3.95
3.21
3.13
10
5.67
3.55
3.19
2.99
12
4.83
3.69
3.03
2.93
15
3.61
3.91
2.87
2.76
The results indicate that approach is competitive with multistep distillation (across varied number of steps) and outperforms + at 8 or more Number of Function Evaluations (NFEs). The performance of DPM-Solver++ degrades as the number of steps increases, which is an expected consequence of its numerical instability.
Our approach has two compelling advantages over multistep distillation. First, it does not rely on distillation, eliminating the need to train a large, slow teacher model. Second, does not require specifying additional sampler hyperparameters during training.
38

---

Distributional Diffusion Models
Table CelebA (Conditional) results. The metric averaged across random seeds. We report results different NFEs for baseline diffusion model, DPM++, Distributional (ours) and moment matching distillation. In bold, we highlight the method with the lowest FID and in italic, we highlight the method with the second lowest FID.
NFEs
Diffusion
DPM++
Distributional
Distillation
2
62.85
58.39
81.79
7.02
4
26.25
20.39
22.98
6.47
8
13.07
17.28
4.94
4.43
10
11.67
22.48
3.60
4.43
12
10.00
23.23
3.37
4.25
15
8.51
23.53
3,45
4.27
Table 5. LSUN (Unconditional) results. The metric is FID averaged across random seeds. We report results with different NFEs for baseline diffusion model, DPM++, Distributional (ours) and moment matching distillation. bold, we highlight the method with the lowest FID and in italic, we highlight the method with the second lowest FID.
NFBs
Diffusion
+PM++
Distributional
Distillation
238.85
120.99
233.24
16.04
41.57
43.60
69.90
5.54
10.82
12.95
9.79
4.52
10
9.08
17.85
7.14
4.64
12
7.04
16.07
6.12
3.81
15
6.51
23.91
5.78
4.39
Table Latent CelcbA-HQ (Unconditional) results. metric averaged across random seeds. We report results with different
NFEs for baseline diffusion model, DPM++, Distributional (ours) and moment matching distillation. In bold, we highlight the method with the lowest FID and in italic, we highlight the method with the second lowest FID.
NFEs
Diffusion
DPM++
Distributional
Distillation
101.47
114.44
67.60
37.88
4
53.83
26.05
23.80
14.47
8
15.60
20.70
9.84
8.22
10
11.89
21.84
8.15
6.76
12
9.97
24.61
7.04
6.05
15
8.39
27.31
6.21
5.46
Furthermore, combining our distributional approach with modern distillation method a potentially interesting topic for future research.
K.4. Additional robotics experiments
In addition to Libero10 suite, we repeated our main experiment on three other suites in Libero benchmark: Liber-Goal, Libero-Spatial, and Libero-Object. We present the performance on these three additional suites in Figure 19. all three suites, distributional diffusion helps the performance when fewer number of diffusion steps are used, reproducing our results on Liberol0.
In the setting of Section 6.3, we only consider the case 3 = 1 and compared our results for several values of À [0, 1]. namely À E {0.0, 0.1, 0.5}. In Figure 20, we present similar study as Figure 9 but in the case where B = 2.
We also ran an ablation over m with fixed / 1 and À = 0.5. The results are reported in Table 7.
K.S. Additional samples
For CIFAR-10, 10, samples from the distributional model è trained with 3 0.1, and with 10 sampling steps are shown in Figure 21, left. Samples from the one trained with = 2, A = 0 and with 100 sampling steps are shown in Figure 22, left. The samples from the diffusion model with steps are shown in Figure 21, right and with 100 steps are shown in
39

---

Distributional Diffusion Models
Libero-Goal
Libero-Spatial
Jibero-Objest
0.80
0.90
Max success rate
0.75
Max success rate
0.85
Max success rate
0.98
0.70
0.80
0.96
0.75
0.94
0.60
0.92
0.70
0.65
0.88
0.50
0.60
0.86
0.40
0.55
0.84
2
16
50
2
16
50
2
16
50
Num diffusion steps
Num diffusion steps
Num diffusion steps
--10A-0.0
A-1.0
8-10A-00
==10A-0.5 0.5
A-1.0A-0.0
B-10A-05
2=10A-011
Diffusion
6-10A-01
Diffusion
A-10A-01
Figure 19. Performance in three other Libero suites Libero-Goal, Libero-Spatial, and Libero-Object as function of diffusion steps and with same settings as Figure 9. We plot the maximum success rate from best checkpoint during training. Distributional models consistently performed the best when fewest diffusion steps used. >> 0 also performs better than À
Liberol0
Libero10
0.90
Max success rate
0.80
0.85
Max success rate
0.75
0.80
0.70
0.65
0.65
0.60
0.60
0.55
0.55
0.50
16 Num diffusion steps
50
Num diffusion steps
8 2.0 A 0.0
2-0.0
8- 0
B 8-1 A 0.5 Diffusion
= A 0.1
Diffusion
8-10 - 0.1
Figure Left. Performance in Libero10 function diffusion steps and A, same as Figure for Right. we reproduce
Figure for ease of read.
40

---

Distributional Diffusion Models
m
Success rate
2
0.783
4
0.73
0.817
0.800
32
0.840
64
0.797
Table 7. Ablation over for robotics dataset.
Figure 21. Left, samples from distributional model trained CIFAR-10 and sampled with 10 steps. = 0.1, Right, samples from diffusion model trained on CIFAR- 10 and sampled with 10 steps.
Figure 22, right.
For CelebA, samples from the distributional model trained with 3 = 0.001, A and with 10 sampling steps, are shown in Figure 23, left. Samples from the one trained with 3 = 2, À = O and with 100 sampling steps, are shown in Figure 24, left. Samples from diffusion model with 10 steps are shown Figure 23, right and with 100 steps are shown in Figure 24, right.
For LSUN Bedrooms, samples from the distributional model trained with 0.0001, shown in Figure 25, left. Samples from the trained with 0.0001, 0 and with 100 sampling steps, shown in Figure 26, left. Samples from the diffusion model with 10 steps shown in Figure 25, right and with 100 steps shown in Figure 26, right.
and with 10 sampling steps,
For latent CelebA HQ, samples from the distributional model trained with 3 = 0.01, = 0.5 and with 10 sampling steps, are shown in Figure 27, left. The samples from the one trained with 3 = 0.01, 0.1 and with 100 sampling steps, are shown in Figure 28, left. Samples from the diffusion model with 10 steps are shown Figure 27, right and with 100 steps shown Figure 28, right.

---

Distributional Diffusion Models
Figure 22. Left, samples from distributional model trained on CIFAR-I0 and sampled 100 steps. 2, À 0. Right, samples
from diffusion model trained on CIFAR-10 10 and sampled with 100 steps.
Figure 23. Left, samples from distributional model trained on CelebA and sampled with 10 steps. 0.001, Right, samples
from diffusion model trained on CelebA and sampled with 10 steps.
42

---

Distributional Diffusion Models
Figure Left, samples from distributional model trained CelebA and sampled with 100 steps. , 2, A Right, samples from diffusion model trained on CelebA and sampled with 100 steps.
Figure Left, samples from distributional model trained LSUN Bedrooms and sampled with steps. , 0.0001, Right.
samples from diffusion model trained on LSUN Bedrooms and sampled with 10 steps
43

---

Distributional Diffusion Models
Figure 26. Left, samples from distributional model trained on LSUN Bedrooms and sampled with 100 steps. 3 0.0001, Right. samples from diffusion model trained on LSUN Bedrooms and sampled with 100 steps
Figure Left, samples from distributional model trained latent CelebA HQ sampled with 10 steps. 0.01, À 0.5. Right. samples from diffusion model trained on latent CelebA HQ and sampled with 10 steps.
44

---

Distributional Diffusion Models
Figure Left, samples from distributional model trained on latent CelebA HQ sampled with 100 steps. 0.01, À 0.1. Right, samples from diffusion model trained on latent CelebA HQ and sampled with 100 steps.
45


Reviews:
[{"review_id": "XjIlLf0ax2", "reviewer": "Reviewer_waLv", "summary": "This paper proposes Distributional Diffusion Models (DDM), an acceleration method of diffusion model. Specifically, the authors follow the  Denoising Diffusion Implicit Models (DDIM) and replace the regression loss with a loss based on scoring rules to accomplish sample generation. This method learns the posterior distribution of clean data samples given their noisy versions, instead of only the mean of this distribution.  Additionally, the authors provide some theoretical analysis of the generalized scoring rules and validate DDM on image and robot trajectory generation.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank reviewer waLv for their comments and positive feedback.\n\nAs requested, we re-implemented and have added comparisons to a popular numerical solver DPM-solver++ [1] (as suggested by reviewer r8Za) and to the recommended distillation method [2]. We did implement the parallel sampling of diffusion models but haven't had enough time to sweep over hyperparameters to obtain competitive results. We will add this comparison as well as the comparisons below, in the final version of the paper.\n\n**Experimental details**: For all the sampling methods, we sweep over the “safety” parameter $\\eta$, which specifies the time interval to be “[$\\eta$,$1-\\eta$]”. We consider values $[1e-1,1e-2,1e-3,1e-4]$. Moreover, for all the methods we sweep over the “churn” parameter which controls the stochasticity, we tried $[0.0,0.25,0.5,0.75,1.0]$. We found the stochastic version of DPM-solver++ (with a churn parameter close to $1$) overall worked the best. For ordinary diffusion and our method, we use a “uniform” time schedule. For DPM-solver++ we use the “logsnr” time schedule (see [1]), which led to the best results (we also tried EDM and uniform).\n\nFor distillation comparisons, we distill pretrained diffusion models into a student model with a given number of student sampling steps. We train a distilled model for 100k iterations with batch size 256 for pixel-space and with batch size 32 for latent space models. We sweep over EMA decay and learning rate. For the student, we use a stochastic DDIM sampler with a churn parameter $\\epsilon$ and we sweep over it with considered values $[0,0.5,1.0]$. At sampling time, we found that using a different churn parameter $\\epsilon=0$ for the distilled student, led to the best results.\n\nFor all the methods, we present the results with the best hyperparameters.\n\n**Results**: The results are presented in the table below. They indicate that our approach is competitive with multistep distillation (we compared different amounts of steps) and outperforms DPM-solver++ over 8 NFEs (whose performance degrades as expected as the number of steps increases due to numerical instability).\n Two compelling benefits of our approach over multistep distillation are that it does not rely on distillation, so there is no need to train a large (slow) teacher model; and additionally it does not require specifying additional sampler hyperparameters during training. It is certainly possible for our distributional approach to be combined with a modern distillation method. This is an interesting topic for future research.\n\nThe tables below contain FID scores (averaged across 3 random seeds) for different number of function evaluations (NFEs) and for different methods. In **bold**, we highlight the method with the lowest FID and in *italic*, we highlight the method with the second lowest FID.\n\n## Dataset name = CIFAR-10 (Conditional)\n|NFEs| Diffusion | DPM++ | Distributional | Distillation |\n|------|-----------|-------|----------------|--------------|\n| 2|80.77|*19.41*|29.33|**5.19**|\n| 4|23.31|12.60|*4.67*|**3.84**|\n| 8|7.53|3.95|*3.21*|**3.13**|\n| 10|5.67|3.55|*3.19*|**2.99**|\n| 12|4.83|3.69|*3.03*|**2.93**|\n| 15|3.61|3.91|*2.87*|**2.76**|\n\n## Dataset name = CelebA (Conditional)\n|NFEs| Diffusion | DPM++ | Distributional | Distillation |\n|------|-----------|-------|----------------|--------------|\n| 2|62.85|*58.39*|81.79|**7.02**|\n| 4|26.25|*20.39*|22.98|**6.47**|\n| 8|13.07|17.28|*4.94*|**4.43**|\n| 10|11.67|22.48|**3.60**|*4.43*|\n| 12|10.00|23.23|**3.37**|*4.25*|\n| 15|8.51|23.53|**3.45**|*4.27*|\n\n## Dataset name = LSUN (Unconditional)\n|NFEs| Diffusion | DPM++ | Distributional | Distillation |\n|------|-----------|-------|----------------|--------------|\n| 2|238.85|*120.99*|233.24|**16.04**|\n| 4|*41.57*|43.60|69.90|**5.54**|\n| 8|10.82|12.95|*9.79*|**4.52**|\n| 10|9.08|17.85|*7.14*|**4.64**|\n| 12|7.04|16.07|*6.12*|**3.81**|\n| 15|6.51|23.91|*5.78*|**4.39**|\n\n## Dataset name = Latent CelebA-HQ (Unconditional)\n|NFEs| Diffusion | DPM++ | Distributional | Distillation |\n|------|-----------|-------|----------------|--------------|\n| 2|101.47|114.44|*67.60*|**37.88**|\n| 4|53.83|26.05|*23.80*|**14.47**|\n| 8|15.60|20.70|*9.84*|**8.22**|\n| 10|11.89|21.84|*8.15*|**6.76**|\n| 12|9.97|24.61|*7.04*|**6.05**|\n| 15|8.39|27.31|*6.21*|**5.46**|\n\nWe hope that our updated experiments have addressed your concerns. If so, we would be grateful if you could consider increasing your score.\n\n**References**:\n\n[1] DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models, Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, Jun Zhu, 2022\n\n[2] Multistep Distillation of Diffusion Models via Moment Matching, Tim Salimans et al., 2024.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "The proposed method and evaluation criteria in this paper are both reasonable for diffusion acceleration.", "theoretical_claims": "Yes, I have checked the gaussian analysis, joint or conditional scoring rules analysis, and from kernel scores to diffusion losses analysis in the theoretical analysis section. I did not find any apparent issues.", "experimental_designs_or_analyses": "Yes, I have checked the experimental designs and analyses, including 2D experiments, image experiments and robotics experiments.", "relation_to_broader_scientific_literature": "1. This paper  provides a unique research perspective by learning the posterior distribution of clean data samples given their noisy counterparts, rather than only the mean of this distribution.\n\n2. Diffusion models are combined with traditional scoring rules to improve the loss function, enabling the model to learn the full distribution rather than the mean.\n\n3. This paper provides the gaussian analysis, joint or conditional scoring rules analysis, and from kernel scores to diffusion losses analysis in the theoretical analysis section, which collectively explain why the use of scoring rules can improve the performance of diffusion models.", "essential_references_not_discussed": "In DDM, the authors take DDIM as a starting point to introduce research on scoring rules. The main contribution is to replace the regression loss with a loss based on scoring rules and to provide the theoretical analysis. No significant related work appears to have been omitted.", "other_strengths_and_weaknesses": "Strengths:\n\n1. The paper provides an acceleration method of diffusion models, learns the posterior distribution of clean data samples given their noisy versions, instead of only the mean of this distribution.\n\n2. The paper replace the regression loss with a loss based on scoring rules, and provides the theoretical analysis.\n\n3. The authors conduct so many experiments on image and robot trajectory generation.\n\n4. The theoretical analysis is detailed. \n\nWeaknesses:\n\n1. In the experimental section, this study only trains a baseline unconditional diffusion model by optimizing (6), but it lacks comparisons with other acceleration methods. As mentioned in the introduction, improved numerical integrators[1], distillation techniques[2], and parallel simulation methods[3], among others, are commonly employed strategies. Comparisons with these methods would make the experiment more comprehensive.\n\n[1] DPM-Solver-v3: Improved Diffusion ODE Solver with Empirical Model Statistics\n\n[2] Multistep Distillation of Diffusion Models via Moment Matching\n\n[3] Parallel Sampling of Diffusion Models\n\nIt would be better to compare with other baselines.", "other_comments_or_suggestions": "Double quotation marks are wrong in Section 7 (line 419).", "questions_for_authors": "I have no questions for the 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": "U9XEZ2NaYq", "reviewer": "Reviewer_r8Za", "summary": "The paper proposes to learn the conditional distribution $p(x_0 | x_t)$ instead of learning the mean of this conditional distribution. To do this, the authors propose the use of scoring rules. In practice, this resorts to balancing two losses: the standard denoising loss, and a diversity loss on the learned conditional distribution $p(x_0 | x_t)$. Theoretical justifications of the proposed method are provided in the Gaussian case. Also, some theoretical comparison is done between two configurations of the loss (conditional vs joint distribution learning). Empirical evidence demonstrates that the proposed method improves sampling quality of diffusion models, mostly in the few-step setting.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer r8Za for their overall positive assessment of our paper. Below, we address your comments.\n\n| Weakness 1\n\nAs suggested, we have now compared our method to the DPM-solver++ [1]. We have also compared it to a SOTA (multistep) distillation method [2]. Please see our response to the reviewer waLv above, where we added these results.\n\nOur results show that our approach is competitive with multistep distillation (we compared different amounts of steps) and outperforms DPM-solver++ over 8 NFEs (whose performance degrades as expected as the number of steps increases, due to numerical instability).\n Two compelling benefits of our approach over multistep distillation are that it does not rely on distillation, so there is no need to train a large (slow) teacher model; and that it does not require specifying additional sampler hyperparameters during training. It is certainly possible for our distributional approach to be combined with a modern distillation method. This is an interesting topic for future research.\n\n| Weakness 2\n\nWe thank you for bringing up this point regarding the training time complexity. Indeed, as you suggested, our paper will benefit from adding a complexity analysis of our method. We will add the following section to the appendix of our paper.\n\n## Computational complexity\nOur method has the same computational complexity during sampling as ordinary diffusion models, and only has an increased computational complexity during training which is detailed below.\n\nAssume that the computational complexity of the forward pass in our diffusion model is $O(F)$ and that the dimensionality of $x$ is $D$.\n\n**Diffusion loss case.** The loss function in eq. (14) with $m=1$ and $\\lambda=0$, can be thought of as a standard diffusion model loss function. It requires $O(nF)$ evaluations to compute $x_{\\theta}(t_i,X_{t_i})$ for every element of a batch. Then, to evaluate the loss (the norm), the complexity is $O(nD)$. Therefore the total cost of computing the loss is $O(n (F+D))$. The backwards pass is proportional to the forward pass and has computational complexity $O(nF)$.\n\n**Distribution loss case.** For $\\lambda \\neq 0$ and $m>1$, we first need $O(nmF)$ function evaluations to compute $x_{\\theta}(t_i, X_{t_i}, \\xi_{i,j})$ for $i=1,\\ldots,n$ and for $\\xi_{i,j}$, $j=1,\\ldots,m$. After that, we need $O(nmD)$ operations to compute the first (diffusion-like) terms of the loss. Then, in order to compute the interaction terms, we need $O(nm(m-1)D)$ time. Therefore, the total cost is of computing the loss $O(nm(F+D) + nm(m-1)D)$. Naively, backwards pass will take $O(nmF + nm(m-1)F) = O(nm^2F)$ time. However, since the 2nd term uses the gradients of $x_{\\theta}(t_i, X_{t_i}, \\xi_{i,j})$, one could precompute these for all $i=1,\\ldots,n$ and $j=1,\\ldots,m$ and therefore decrease the total backwards cost to $O(nmF)$. We found that in practice, the XLA compiler in jax performs this optimization without explicit coding (see below).\n\n**Runs on real hardware.** We compared the training times on real hardware. We study the training of CIFAR-10 and we vary $m$. We report **steps per second** metric where a step corresponds to a full forward+backward step. As hardware, we use A100 GPU (40 Gb of memory) with batch size = 16, H100 GPU (80Gb of memory) with batch size = 64 and TPUv5p (with 95Gb of memory) with batch size = 64 (per device with 4 devices in total). The results below indicate that **steps per second** decreases proportionally to $m$.\n\nHardware| Diffusion|Distributional (m=2)|Distributional (m=4)| Distributional (m=8)\n-|-|-|-|-\nA100 (bs=16)|  9.05 | 6.78 | 4.6 | 2.77\nH100 (bs=64) |  14.3  | 7.85 | 4.15 | 2.06\nTPUv5 (bs=64)|  11.2   | 8.5 | 4.3  | 2.22\n\n| Could this framework be extended to score-based models with ODE sampling?\n\nIf the churn parameter $\\epsilon$ is taken to be zero in the DDIM equations, one recovers a deterministic sampler which can be thought of as the discretization of an ODE. This ODE is a special case of the standard probability flow ODE; see Section 4.3 of [3] for a discussion. In particular, the DDIM update coincides with the exponential integrator for the probability flow ODE.\n\nWe emphasize that our method can also be extended directly to sample a transport map between any two distributions $p_0$ and $p_1$ and does not require one of the distributions to be Gaussians. This is a framework for which flow matching/stochastic interpolants are typically used.  We will add a section in our paper to detail this extension. \n\nWe hope that our updated experiments have addressed your concerns. If so, we would be grateful if you could consider increasing your score.\n\n**References**:\n\n[1] DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic Models, Cheng Lu et al., 2022\n\n[2] Multistep Distillation of Diffusion Models via Moment Matching, Tim Salimans et al., 2024.\n\n[3] Denoising diffusion implicit models, Jiaming Song et al., 2020.", "claims_and_evidence": "Claims made in the submission are supported by clear and convincing evidence. The main claim is that learning conditional distribution $p(x_0 | x_t)$ allows to sample with a coarser discretization grid. This is validated empirically on image benchmarks with FID metric, and on robotic dataset.", "methods_and_evaluation_criteria": "Method: The method is sound, elegant, and simple to implement. \n\nEvaluation criteria: Evaluations criteria is well done. It consists in: theoretical proofs in a simplified setting (Gaussian target distribution), experiments on toy datasets with squared MMD metric, image datasets with FID, and robotic datasets with max success rate on a robotic manipulation task. \nSomething that is missing is a comparison with other methods that improve the sampling quality when using few number of steps. In the current approach, the authors only compare to vanilla DDIM. However, many approaches have been proposed to reduce the neural function evaluations (NFEs) required for achieving high sampling quality. For example, it would be interesting to compare to works such as DPM-solver (Lu et al.).", "theoretical_claims": "I checked the proof of Proposition 4.1 (proof of the loss in the Gaussian setting) and it is correct. I did not check the other proofs in detail. Also note that there is no theoretical proofs of the main claim, which is that the proposed method allows to reduce the number of sampling steps compared to a standard diffusion model.", "experimental_designs_or_analyses": "I checked the design of all experiments, and there is no strong issue. The main weakness is the lack of comparison with other methods aiming high performance of diffusion models with low number of sampling steps.", "relation_to_broader_scientific_literature": "This paper proposes a novel view to improve the sampling quality of diffusion models with low number of steps. In my knowledge, the approach of learning $p(x_0 | x_t)$ with scoring rules is novel and original. Other approaches to learn $p(x_0 | x_t)$ made use of GANs and are discussed by the authors.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Weakness 1: Although this paper improves sampling quality at low number of sampling steps, it lacks comparisons with other methods that share the same purpose. Many different methods exist and share the same goal, e.g. improved samplers for diffusion models. How does this compare to some of these other methods? \n\nWeakness 2: The proposed method leads to a computational overhead during training (generating m points for each training point). I did not see any study on the training time, but this would be important to understand the practical feasibility of the method.", "other_comments_or_suggestions": "No.", "questions_for_authors": "Could this framework be extended to score-based models with ODE sampling?", "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": "2AFZT66tIu", "reviewer": "Reviewer_FCZ5", "summary": "This paper introduces a approach to accelerate diffusion models by learning the full conditional posterior distribution of clean data given noisy samples, rather than just the conditional mean, as done in standard diffusion models. The authors propose replacing the traditional regression loss (used to estimate conditional means) with a generalized energy score based on scoring rules, which allows sampling from the posterior distribution on a coarser time scale, thus reducing the number of denoising steps required.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer FCZ5 for their feedback on our work. Please see our response below\n\n|$p_{0|t}(x_0|x_t)$,Dirac…not true\n\nUsing $p_{0|t}(x_0|x_t)\\approx\\delta_{\\hat{x}_\\theta(t,x_t)}$ in $p(x_s|x_t)$ (see eq.5) leads to\n\n$p(x_s|x_t)\\approx p(x_s|x_0=\\hat{x}_\\theta(t,x_t),x_t),$\n\nwhere $p(x_s|x_0,x_t)$ is Gaussian (see eq.4). Sampling from these approximations of $p(x_{t_k}|x_{t_{k+1}})$ for some $t_N=1>t_{N-1}>....>t_1>t_0=0$ gives exactly the DDIM algorithm (see DDIM paper [1],eq.10)\n\n|Other metrics,diversity\n\nFID [2] already quantifies the diversity of the sampled images as it is Wasserstein-2 distance between Gaussian approximations to the distributions in the space of Inception V3 coding layer features, and it takes into account the covariances of these distributions. See [3] regarding image metrics for evaluation (i.e., “[the FID score] unlike IS [..] is able to detect intra-class mode dropping”)\n\nA good metric for evaluating diffusion policy in robotics is an open challenge and so far success rates during independent evaluation is the only measure that has proven useful for robotics. Systematic evaluations of various metrics for policies yielded no reliable metrics [4]\n\n|more 2D\n\nWe trained unconditional standard diffusion and our distributional models on a more complex 2D distribution – checkerboard, see [5]. We report MMD between sampled and target distributions with different NFEs, smaller MMD is in **bold**. Our approach outperforms standard diffusion for small NFEs. We will add these results to the paper\n\nNFEs|Diffusion|Distributional\n-|-|-\n5|7.00e-3|**1.57e-4**\n10|2.10e-3|**1.96e-4**\n50|2.70e-4|**1.55e-4**\n100|**1.96e-4**|2.00e-4\n1000|**1.71e-4**|1.83e-4\n\n|[5-6]\n\nThank you for pointing out these references. We’ll discuss as follows:\n\n“Scoring rules have corresponding MMD divergences, and approaches exist that use the MMD in diffusion and particle flow models. [5] generates a sequence of distributions from a forward diffusion process. It then generates noise-dependent neural MMDs between clean and noisy data, and performs gradient flow to move particles from noise level $t$ to a lower noise level $s<t$. Compared to ours, [5] does not use a generator since it is a particle flow. Their best CIFAR10 FID is 7.7. [6] first trains a diffusion model, and then refines/distills by coarsening the reverse timesteps, and uses MMD on CLIP features to finetune the DDM denoiser. Their CIFAR10 FID for NFE=10 improves from 13.6 to 3.8, while ours is 3.19.\n\n|m sensitivity\n\nPlease see our response to reviewer r8Za about m and computational complexity. For image experiments $m=2,4$ already worked well. We will add an ablation over $m$ for CIFAR10 in the revised version of the paper. For robotics, we ran an ablation over $m$ ($\\beta=1,\\lambda=0.5$), see results below\n\nm|Success rate\n-|-\n2|0.783\n4|0.73\n8|0.817\n16|0.800\n32|0.840\n64|0.797\n\n|IMQ,RBF,kexp\n\nProper scoring rules define divergences. For kernel scores, this is the squared MMD. The energy distance kernel is appealing – its loss (13) is similar to standard diffusion loss, and other kernels can recover this loss in certain limits. While this recovery is good for interpretability, any characteristic kernel (including the $kexp$) gives a valid scoring rule. A benefit of (13) is its numerical stability without exponentiation\n\n|additional experiments\n\nPlease see our response to reviewer waLv for comparisons with DPM-solver++ and SOTA multi-step distillation.\n\nWe started an experiment on ImageNet (64x64), but due to time constraints, will only add results after rebuttal. As for robotics, we studied 10 most challenging tasks from Libero benchmark, which is already a diverse set of language-conditioned robotics tasks. We started an experiment on all tasks from Libero and will report results after rebuttal\n\n|$\\Sigma_{s,t}$...why learning necessary?\n\nWe need to approximately sample from $p(x_s|x_t)=\\int p(x_s|x_0,x_t)p(x_0|x_t)dx_0$ where $p(x_s|x_0,x_t)=N(x_s;\\mu_{s,t}(x_0),\\Sigma_{s,t})$ is fixed. We obtain a better approximation by learning $p(x_0|x_t)$ with scoring rules instead of Dirac mass as denoiser. Minimizing loss (13) recovers the true $p(x_0|x_t)$, if $p^\\theta(x_0|x_t)$ is rich enough. The resulting $p(x_s|x_t)$ is indeed non-Gaussian and its form is fully problem dependent.\n\nWe hope that our explanations and updated experiments have addressed your concerns. If so, we would be grateful if you would consider raising your score.\n\n[1]Song J et al.,Denoising Diffusion Implicit Models,ICLR2021\n\n[2]Heusel M et al.,GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium,NeurIPS2017\n\n[3]Borji A,Pros and Cons of GAN Evaluation Measures,2018\n\n[4]Hussenot L, et al.,Hyperparameter Selection for Imitation Learning,2021\n\n[5]Galashov A et al.,Deep MMD gradient flow without adversarial training,ICLR2025\n\n[6]Aiello E et al.,Fast inference in denoising diffusion models via MMD finetuning,IEEE2024\n\n[7]Arbel M et al.,MMD Gradient Flow,NeurIPS2019", "claims_and_evidence": "The author claim that $p_{0|t}(x_0|x_t)$ is approximated by $\\delta_{\\hat{x}_\\theta}(t,x_t)$, which is not true. In ddpm, $p\\_{0|t}(x_0|x_t)=\\mathcal{N}(x_0;\\tilde{\\mu}(x_t), \\tilde{\\beta}_t)$, where $\\tilde{\\mu}(x_t)=\\frac{1}{\\sqrt{\\bar{\\alpha}_t}}(x_t-\\sqrt{1-\\bar{\\alpha_t}}\\epsilon_t)$ and $\\tilde{\\beta}_t=\\frac{(1-\\bar{\\alpha}_t)}{\\bar{\\alpha}_t}I$, $\\epsilon_t = \\mathbb{E}[\\epsilon]$ denotes noise added in the forward process. Hence, it's strange to claim that  posterior distribution is  Dirac. I think the author should explain this thing, maybe i am wrong, but author should give the reason.", "methods_and_evaluation_criteria": "The proposed method replaces the regression loss in standard diffusion models with a generalized energy score (Equation 12) or kernel score, training a generative network to approximate the full posterior. The evaluation criteria include:\n\n - 2D Experiments: The authors use a mixture of Gaussians as the target distribution and measure performance via squared Maximum Mean Discrepancy (MMD) with an RBF kernel.\n - Image Generation: Evaluated on CIFAR-10, CelebA, LSUN Bedrooms, and CelebA-HQ (latent space) using the FID.\n - Robotics: Evaluated on the Libero-Long benchmark, measuring success rates for multi-task robotic manipulation.", "theoretical_claims": "The paper makes several theoretical claims, primarily in Section 4. The theoretical claims are generally sound for the specific cases analyzed (Gaussian distributions, specific kernels)", "experimental_designs_or_analyses": "The experimental designs are sound and validate the core hypothesis of improved performance with fewer steps. However, the reliance on a single metric (FID for images, success rate for robotics) limits the depth of the evaluation. Including diversity metrics and more complex 2D distributions would strengthen the analysis.", "relation_to_broader_scientific_literature": "They leveraged the idea of diffusion distillation to fit the posterior distribution across different step sizes. Due to the accumulation of multiple states, an excessively large step size can lead to significant approximation errors in the posterior estimation, causing it to deviate from a Gaussian distribution. In this case, the predicted \\(\\epsilon\\) can no longer serve as the score, and the variance becomes uncertain in its evolution, requiring a neural network for estimation.", "essential_references_not_discussed": "[1] Galashov A, de Bortoli V, Gretton A. Deep MMD gradient flow without adversarial training[J]. arXiv preprint arXiv:2405.06780, 2024.\n\n[2] Aiello E, Valsesia D, Magli E. Fast inference in denoising diffusion models via mmd finetuning[J]. IEEE Access, 2024.", "other_strengths_and_weaknesses": "Strengths:\n\n- The use of scoring rules to learn the full posterior in diffusion models is a novel and elegant solution to the problem of slow sampling, avoiding the complexity of GANs or distillation.\n\n - The method’s ability to reduce the number of denoising steps while maintaining quality has practical implications for real-time applications (e.g., robotics).\n\n - The paper is well-written, with a clear structure and detailed explanations of the method and experiments.\n\nWeaknesses:\n\n - The author do not discuss the relationship between this paper and [1, 2].\n\n - The author claim that $p_{0|t}(x_0|x_t)$ is approximated by $\\delta_{\\hat{x}_\\theta}(t,x_t)$ may be wrong.\n\n - The proposed method is similar to Maximum Mean Discrepancy (MMD), but use a different story.\n\n[1] Galashov A, de Bortoli V, Gretton A. Deep MMD gradient flow without adversarial training[J]. arXiv preprint arXiv:2405.06780, 2024.\n\n[2] Aiello E, Valsesia D, Magli E. Fast inference in denoising diffusion models via mmd finetuning[J]. IEEE Access, 2024.\n- The experiments focus heavily on FID and success rates, neglecting metrics like sample diversity or robustness, which are crucial for generative models.", "other_comments_or_suggestions": "Additional Experiments: Including experiments on larger model (e.g., sdxl, sd3.5) or more diverse robotics tasks would strengthen the empirical validation.", "questions_for_authors": "- In the empirical energy diffusion loss (Equation 14), you use m=8 (2D experiments) and m=16 (robotics). How sensitive is the performance to the choice of m, and what motivated these specific values? A larger m might better approximate the posterior but increase computational cost—how does this trade-off affect your method?\n\n - Proposition 4.3 shows that IMQ and RBF kernels are diffusion-compatible, but the exponential kernel (kexp) is not. However, you note that kexp still results in desirable properties (Section 4.3). Can you elaborate on these properties and why non-diffusion-compatible kernels might still be useful? This could clarify the practical implications of kernel choice.\n\n - Perhaps the author needs to explain that if $\\Sigma_{s,t} = \\sigma_s^2 \\left( 1 - \\left( \\varepsilon^2 r_{1,1}(s,t) + (1 - \\varepsilon^2) \\right)^2 \\right) \\text{Id},$ is defined as such, then why is learning still necessary? Wouldn't it be possible to directly use this $\\Sigma_{s,t}$? (In the algorithm, the author also directly uses this $\\Sigma_{s,t}$, with only the mean being modified.)  \n If, as I mentioned earlier, the transition probability distribution no longer follows a Gaussian distribution due to cross-step state superposition, then what should $\\Sigma_{s,t}$ look like? Is the learned result from the method proposed by the author correct?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The manuscript proposes, in the context of diffusion models, learning the posterior distribution of clean samples given their noisy counterparts using a scoring rule–based loss. The reviewers found that the manuscript provides a novel view on improving the sampling quality of diffusion models and that the theoretical contributions are sound, elegant and analytically detailed. The reviewers had concerns regarding limitations in the empirical evaluation, particularly the lack of comparison to other methods and baselines, analysis of the computation overhead during training, and relation to other literature. However, the authors addressed all these concerns during rebuttal, and all reviewers recommend accepting the paper. Therefore, I also recommend accepting the paper. Since the rebuttal contains many additional details and experience, it is crucial that the authors incorporate all promised changes into the final version of the paper.
Decision: Accept (poster)