Deep MMD Gradient Flow without adversarial training
Authors: Alexandre Galashov
Valentin De Bortoli
Arthur Gretton
Venue: ICLR 2025 Poster (2025)
Keywords: Generative modeling
diffusion models
Wasserstein gradient flows
generative adversarial networks
discriminator flow
Primary area: generative models
TL;DR: A method to train noise conditional GAN-like discriminator without adversarial training with a forward diffusion and sample from the noise adaptive Wasserstein gradient flow

Abstract: We propose a gradient flow procedure for generative modeling by transporting particles from an initial source distribution to a target distribution, where the gradient field on the particles is given by a noise-adaptive Wasserstein Gradient of the Maximum Mean Discrepancy (MMD). The noise adaptive MMD is trained on data distributions corrupted by increasing levels of noise, obtained via a forward diffusion process, as commonly used in denoising diffusion probabilistic models. The result is a generalization of MMD Gradient Flow, which we call Diffusion-MMD-Gradient Flow or DMMD. The divergence training procedure is related to discriminator training in Generative Adversarial Networks (GAN), but does not require adversarial training. We obtain competitive empirical performance in unconditional image generation on CIFAR10, MNIST, CELEB-A (64 x64) and LSUN Church (64 x 64). Furthermore, we demonstrate the validity of the approach when MMD is replaced by a lower bound on the KL divergence.

Published as a conference paper at ICLR 2025
DEEP MMD GRADIENT FLOW WITHOUT ADVERSARIAL TRAINING
Alexandre Galashov UCL Gatsby Google DeepMind aga lashaaaegooggle.com com
Valentin De Bortoli
Arthur Gretton UCL Gatsby Google DeepMind
Google DeepMind vdeborteli loolollgoooogleooooo com com
gretton com
ABSTRACT
We propose gradient flow procedure for generative modeling by transporting particles from an initial source distribution a target distribution, where the gradient field on the particles is given by noise-adaptive Wasserstein Gradient of the Maximum Mean Discrepancy (MMD). The noise-adaptive MMD is trained on data distributions corrupted by increasing levels of noise, obtained via a forward diffusion process, as commonly used in denoising diffusion probabilistic models. The result is a generalization of MMD Gradient Flow, which call Diffusion- Flow or DMMD. The divergence training procedure is related discriminator training in Generative Adversarial Networks (GAN), but does not require adversarial training. We obtain competitive empirical performance unconditional image generation on CIFARIO, MNIST, CELEB-A (64 x64) and LSUN Church (64 x 64). Furthermore, we demonstrate the validity of the approach when MMD replaced by a lower bound on the KL divergence.
INTRODUCTION
In recent years, generative models have achieved impressive capabilities on image Saharia et al. (2022), audio Le et al. (2023) and video generation Ho et al. (2022) tasks but also protein modeling Watson et al. (2023) and 3d generation Poole et al. (2023). Diffusion models (Sohl-Dickstein et al.
2015 Ho et al. 2020 Song et al. 2021 Rombach et al. (022)) underpin these new methods. In these models, we learn a backward denoising diffusion process via denoising score matching (Hyvärinen
2005| Vincent 201T). This backward process corresponds to the time-reversal a forward noising process. At sampling time, starting from random Gaussian noise, diffusion models produce samples by discretizing the backward process.
One challenge that arises when applying these models in practice is that the Stein score (that is, the gradient log of the current noisy density) becomes ill-behaved near the data distribution (Yang et al. 2024) the diffusion process needs to be slowed down at this point, which incurs a large number of sampling steps near the data distribution. Indeed, if the manifold hypothesis holds TTnenbaum et al. al. (2016); Brown et al. (2022) and the data is supported on lower dimensional space, it is expected that the score will explode for noise levels close to zero, to ensure that the backward process concentrates on this lower dimensional manifold[Bortoli (2022) Pidstrigach] (2022); Chen et al. (202)). While strategies exist to mitigate these issues, they trade-off the quality of the output against inference speed, see for instance (Song et al. 2023
(0000):eefferman
Xu al. 2024 Sauer et al. 2023).
Generative Adversarial Networks (GANs) (Goodfellow et al. 2014) represent an alternative popular generative modelling framework (Brock et produced by generator: a neural net mapping low dimensional noise to high dimensional images. The generator is trained in alternation with a discriminator, which is a measure of discrepancy between the generator and target images. An advantage of GANs is that image generation is fast once the GAN is trained (Xiao et al. (002). although image samples are of lower quality than for the best diffusion models (Ho et al. 2020 Rombach et al., (022). When learning a GAN model, the main challenge arises due to the presence of the generator, which must be trained adversarially alongside the discriminator. This requires careful hyperparameter tuning (Brock et al. 2019 Karras
Karras et al. 2020a). Candidate samples are

---

Published as a conference paper at ICLR 2025
et al. 2020b Liu et al. (021). without which GANs may suffer from training instability and mode collapse (Arora et al. 2017; Kodali et al. 2017 Salimans et al. 2016).
Nonetheless, the process of GAN design has given rise to strong understanding of discriminator functions, and a wide variety of different divergence measures have been applicd. These fall broadly
into two categories: the integral probability metrics (among which, the Wasserstein distance (Arjovsky
eet al.
2017:
Gulrajani et al., 22177] Genevay et al. 2018) and the Maximum Mean Discrepancy (Li
al.
2017 Binkowski et [0118] Arbel et al. 2018 and the f-divergences (Goodfellow al.
2014 Nowozin et al. 2016]] Mescheder el al. 2018 Brock et al. 2019). While it would appear that
f-divergences ought to suffer from the same shortcomings as diffusions when the target distribution is supported on a submanifold Arjovsky et al. (2017), the divergences used in GANs are in practice variational lower bounds on their corresponding f-divergences (Nowozin et al., 2016), and in fact behave closer to IPMs in that they do not require overlapping support of the target and generator samples, and can metrize weak convergence (Arbel et al. [202T] Proposition 14) and (Zhang et al., 2018) (there remain important differences, however: notably, f-divergences and their variational lower bounds need not be symmetric in their arguments).
A natural question then arises: is it possible to define a Wasserstein gradient flow (Ambrosio et al. 2008; Santammboobooooo 2015) using a GAN discriminator as a divergence measure? In this setting, the divergence (discriminator) provides gradient field directly onto set of particles (rather than to generator), transporting them to the target distribution. Contributions in this direction include the MMD flow Arbel et al. (2019): Heertrich al. (2024), which defines Wasserstein Gradient Flow on the Maximum Mean Discrepancy (Gretton et al. 2012); and the KALE (KL approximate lower-bound estimator) flow Glaser et al. (2021). which defines a Wasserstein gradient flow on a KL lower bound of the kind used as a GAN discriminator based on an f-divergence (Nowozin et al.) 2016). We describe the MMD and its corresponding Wasserstein gradient flow in Section [ These approaches employ fixed function classes (namely, reproducing kernel Hilbert spaces) for the divergence, and are thus not suited to high dimensional settings such as images. Moreover, we show in this work that even for simple examples in low dimensions, an adaptive discriminator ensures faster convergence of a source distribution to the target, see Section
A number of more recent approaches employ trained neural net features in divergences for a subse- quent gradient flow (e.g. Fan et al., 2022 Franceschi et al. 2023). Broadly speaking, these works used adversarial means train series of discriminator functions, which are then applied in sequence to a population of particles. While more successful on images than kernel divergences, the approaches retain two shortcomings: they still require adversarial training (on their own prior output), with all the challenges that this entails; and their empirical performance falls short in comparison with modern diffusions and GANs (see related work in Section for details).
In the present work, we propose a novel Wasserstein Gradient flow on a noise-adaptive MMD diver- gence measure, leveraging insights from both GANs and diffusion models. To train the discriminator, we start with clean data, and use a forward diffusion process from (Ho et al., 2020) to produce noisy versions of the data with given levels of noise (data with high levels of noise are analogous to the output of a poorly trained generator, whereas low noise is analogous to well trained generator). The added noise always Gaussian. For a given level of noise, we train noise conditional MMD) discriminator to distinguish between the clean and the noisy data, using a single network across all noise levels. This allows us to have better control over the discriminator training procedure than would be achievable with GAN generator at different levels of refinement, where this control is implicit and hard to characterize.
To draw new samples, we propose novel noise-adaptive version of MMD gradient flow (Arbel et al. 2019). Starting from Gaussian distribution, we move them in the direction of the target distribution by following MMD Gradient flow (Arbel et al.) 2019). adapting corresponding level of noise. See Section 4 for details. This allows us to have a fine grained control over the sampling process. As final challenge, MMD gradient flows have previously required large populations of interacting particles for the generation of novel samples, which is expensive (quadratic in the number of particles) and impractical. In Section 5 we propose a scalable approximate sampling procedure for a case of a linear base kernel, which allows single samples to be generated with a very little loss in quality, at cost independent of the number of particles used in training. The MMD is an instance of an integral probability metric, however many GANs have been designed using discriminators derived from f-divergences. Section demonstrates how our approach can be applied
MMD discriminator to the
2

---

Published as a conference paper at ICLR 2025
to such divergences. using lower bound on the KL divergence an illustration. Section contains a review of alternative approaches to using GAN discriminators for sample generation. Finally, in Section 7 we show that method, Diffusion-MMDD -Ma-ggddieettt flow (DMMD), yields competitive performance in generative modeling on 2-D datasets as well as in unconditional image generation on CIFAR1O (Krizhevsky & Hinton 2009), MNIST, CELEB-A, A, LSUN Church.
BACKGROUND
In this section, we define the MMD GAN discriminator, then describe Wasserstein gradient flow as it applies for this divergence measure.
MMD GAN. Let X C ED and P(X) be the set of probability distributions on X. Let P E ((X)) be the target (data) distribution and Qu E P(X) be a distribution associated with generator parameterized by vb E R². Let H be Reproducing Kernel Hilbert Space (RKHS), see (Schölkopf & Smola 2018) for details, for some kernel
X - E. The Maximum Mean Discrepancy (MMD) (Gretton et al.) 2012) between Q+ and P is defined as MMD(Q P) = supl]s] w11 Q [](X] - Ep
We refer to the function fQ..P that
attains the supremum as the wimeess function,
(() x k(x 2)QQ(() kly, 2))Py))
which will be essential in defining our gradient flow. Given X = {x:) Q and yM ( ~ p&M, the empirical witness function known in closed form, fq.. p(x) x N E k) xr Σラ ()), and unbiased estimate of MMD (Gretton et al. 2012 is
likewise straightforward. In the MMD GAN (Binkowski et al. 2018
Li et al. 2017)) the kernel
k( y) Keass (o(z; 0), o1;; 0)),
where lipase is base kernel and ((:: : X - K are neural networks discriminator features with parameters E " We use the modified notation [X Y to highlight the functional dependence the discriminator parameters. The MMD is Integral Probability Metric (IPM) (Muller 1997), ) and thus well defined on distributions with disjoint support: this argument was made in favor of IPMs by Arrovsky al. (2017). Note further that the Wasserstein GAN discriminators of Arjjvssyy et al. (2017): Gulrajani et al. (2017) can be understood in the MMD framework, when the base kernel linear. Indeed, was observed by Genevay al. (2018) that requiring closer approximation to a true Wasserstein distance resulted in decreased performance in GAN image generation, likely due to the the exponential dependence of sample complexity on dimension for the exact computation of the Wasserstein distance; this motivates an interpretation of these discriminators simply as IPMs using class of linear functions of learned features. We further note that the variational lower bounds used in approximating f-divergences for GANs share with IPMs the property of being well defined on distribtions with disjoint support[Nowozl Nowozin et al. (2016); Arbel et al. (202I). although they need not be symmetric in their arguments. Finally, while Qs and 0 are trained adversarially in GANs, our setting will only require us to learn the discriminator parameter A.
Wasserstein gradient flows. Instead of GAN generator, can move sample of particles along the Wasserstein Gradient flow associated with the discriminator (Ambrosio et al. (2008). Let P2())) be a set of probability distributions
X' with finite second moment equipped with the 2- Wasserstein distance. Let F(v) P2()') + R be functional defined over P2())) with a property that arg inf, Flv)= P. We consider the problem of transporting mass from an initial distribution
Po = Q to target distribution μ = P. finding a continuous path (M))>0 starting from v0 that converges to μ. This problem studied in Optimal Transport theory (Villani) 2008 Santambrogio 2015). This path can be discretized as sequence of random variables (X,) such that X,, Pn.
Xn+1 X,, - VF (p,) Xo ~ Q,
(3)
where 77 > and ((,,,, is the first variation of F associated with the Wasserstein gradient, (Ambrosio et al. 2008 Arbel et al.) 2019) for precise definitions. As oo and 0. depending on the conditions on J, the process (3) will convergence to the gradient flow as a
continuous time limit (Ambrosio al. 2008).
3

---

Published as a conference paper at ICLR 2025
MMD gradient flow. For choice F(v) MMD" Pl and fixed kernel, conditions for convergence of the process in (3 to are given by Arbel et al. (2019). Moreover, the first variation
of (v)) = fu, E H is the witness function defined earher U Usng T-3) gradient flow for any n E N is given by
3 the discretized MMD
Xn+1 = Xn V Su,.
Xo ~ Q.
This provides an algorithm to (approximately) sample from the target distribution P. We remark that Arbel et al. (2019): Hertrich et al. (2024) used a kernel with fixed hyperparameters. In the next section, we will argue that even for RBF kernels (where only the bandwidth chosen), faster convergence will be attained using kernels that adapt during the gradient flow. Details of kernel choice for alternative approaches are given in related work (Section 6
A MOTIVATION FOR ADAPTIVE KERNELS
In this section, we demonstrate the benefit of using an adaptive kernel when performing MMD gradient flow. We show that even in the simple setting of Gaussian sources and targets, an adaptive kernel improves the convergence of the flow. Let kig (1r, y) = a-d exp[-I|ar y (2(2)2(2) be the normalized Gaussian kernel. For any μ E H and > 0 we denote by TH, the Gaussian distribution with mean μ and covariance matrix o Id. We denote MMD, the MMD associated with hia
a="exp(-le
Proposition 3.1. For any Po E IRd and a > 0, let a* be given by
a= = argmax IV no MMD (20,0:
Then, we have that
a"
= ReLLU([lolll///(/ /(d++///+++))) - 20221/2
(5)
The result is proved in Appendix
K The quantity II po MMD (FO,a represents how much the mean of the Gaussian " is displaced by flow w.r.t. MMD2, We want (MD' (To ol large as possible as it denotes the maximum displacement possible.
We show that a² maximizing this displacement is given by (5). Assuming > 0 is fixed, it is notable that this quantity depends on ||ol||, i.e. the distance between the two distributions. This observation justifies our approach of following an adaptive MMD flow at inference time. We further highlight the phase transition behaviour of Proposition] 3.1 once the Gaussians are sufficiently close, the optimal kernel width is zero (note that this phase transition would not be observed in the simpler Dirae GAN example Mescheder al, (2018), where the source and target distributions are Dirac masses with no variance). This phase transition suggests that the flow associated with MMD benefits less from adaptivity as the supports of the distributions overlap. We exploit this observation by introducing an optional denoising stage to our procedure; see the end of Section
In practice, is not desirable to approximate the distributions of interest by Gaussians, and richer neu- ral network kernel features o(z; ()) are used (see Section 7). Arbel et al. (2018) describe approaches to optimize the MMD parameters for GAN training. which serve as proxies for convergence speed: not sufficient simply maximize the MMD, since the witness function should remain Lipschitz to ensure convergence (Arbel et al.) 20188] Proposition 2). Regularization of the witness is achieved in practice by controlling the gradient of the witness function; take a similar approach in Section
DIFFUSION MAXIMUM MEAN DISCREPANCY GRADIENT FLOW
In this section, present Diffasion Maximum Mean Discrepancy gradient flow (DMMD), new generative model with a training procedure of MMD discriminator which does not rely on adversarial training, and leverages ideas from diffusion models. The sampling part of DMMD consists in following noise adaptive variant of MMD gradient flow.
In the case of variational lower bounds on f-divergences, the witness function is still well defined, and the first variation takes the same form in respect of this witness function: Glaser al. (02T for the case of the KL divergence.
4

---

Published as a conference paper at ICLR 2025
Ad Adversarial-fre training of noise conditional discriminators. In order to train discriminator without adversarial training, we propose to use insights from GANs training. In a GAN setting, at the beginning of the training, the generator is randomly initialized and therefore produces samples close to random noise. This would produce a coarse discriminator since it is trained to distinguish clean data from random noise. As the training progresses and the generator improves, so too does the discriminative power of the discriminator. This discriminator behavior is central in the training of GANs (Goodfellow et 2014). We propose a way to replicate this gradually improving behavior without adversarial training and instead relying on principles from diffusion models (Ho et al. 2020).
The forward process in diffusion models allows us to generate a probability path Pt, t [0. 1], such that Po- = P. where P is our target distribution and P1 N(0, Id) is a Gaussian noise. Given samples 2o ~ Po = P. the samples 2o are given by
Tt 20 + Bee, E N(, Id),
(6)
with On = B, = and on = 8o noise level the samples Fe are very close the original data zo. whereas for the large values of Ie they are close to a unit Gaussian random variable. Using the GANs terminology, thought as the output of generator such that for high/low noise level f, it would correspond to undertrained well-trained generator. Using this insight, for each noise level t E [o, 1], we define f, discriminator MMD² P; 0) using the kernel of type (2) with noise-conditional discriminator features <(;; I; 0) parameterized by Neural Network with learned parameters 0. We consider the following noise-conditional loss function
From the form of the 2e we observe that for low
could be
((O,t) MM.
where the minus sign comes from the fact that our aim is to maximize the squared MMD. In addition, regularize this loss with f2-penalty (Binkowski et al. 2018) denoted Ltg (0. L) as well as with the gradient penalty (Binkowski et al. 2018 Gulrajani et al. 2017) denoted L (0, t). see Appendix C.2
for the precise definition of these two losses. The total conditonnl loss is then given
(0, t) = C(8, t) + Xe, Ci, (0. (0,
(8)
for a suitable choice of hyperparameters Xe, 2 0, 2 O. Finally, the total loss is given as (0) U0,11 (0, f)]. where U10, is uniform distribution. In practice,
sampled- based unbiased estimator of MMD, see Appendix C.2 The procedure is described in Algorithm
Adaptive gradient flow sampling. In order to produce samples from P. we use the adaptive MMD gradient flow with noise conditional discriminators MMD [P,, P; 0" where O the discriminator parameters obtained using Algorithm
[ Let f,2 = mmin + i = 0,.... . T be the noise discretisation, where At = (tmax - (min) such that to mmin tT fmax for some min = e and Imax = 1 - 6, where e < 1. We sample Np initial particles {Z ~ N(O, N(), {- For each f, we
Nr
follow MMD gradient flow 4 for N₂ steps with learning rate 1)
z zin+1 =Z,
Z, - V P. ((z 11)0).
Here k,.. I/N, 8zz, is the empirical distribution of particles (2₂ 1 No at the noise level and the iteration n, 5 is a Dirae mass measure. The function Sukg... .p .(2, e, t; (*) is adapted from equation (T) where is replaced by this empirical distribution. After following the gradient flow (9) for N, steps, we initialize a new gradient flow with initial particles Z = Zi... for each i= 1,..., No- with the decreased level of noise i At. The recurrence is initialized with z Z'
where {Z+ 1 are the initial particles. This procedure corresponds to running T 1 consecutive MMD gradient flows for N, iterations each, gradually decreasing the noise level from tmax to m
The resulting particles (Z (Zi.N...0 No
Mt=1 are used as samples from P. See Algorithm 2
In practice we sample (once) a large batch N. of (X)),% P&N, from the data distribution and denote by PN, (Xo) the corresponding empirical distribution. Then we use the empirical witness
Different schedules (ar, 3e) are available in the literature. We focus on Variance Preserving SDE ones Song 202T) here
5

---

Published as a conference paper at ICLR 2025
function
fuky in (X.) Px 0*) given by
N, Li=1 Knaw ko (o(2, t: )) (22,t;0))
N. Σ
Khase
(o(2), t:60 ((2,t;8))
(10)
Algorithm I Train noise-conditional MMD dis- Algorithm 2 Noise-adaptive MMD gradient flow
criminator
Inputs: - number of noise levels
Input:
{xr Discriminator features (xz, 1; 0) with parame-
Dataset D
mmax miin maximum/minimum noise levels N₂ number of gradient flow steps per noise level
ters 0 E R
Av 2 A, 2 gradient and l penalty
7 > 0 0- gradient flow learning rate
coefficients
Np number noisy particles
> 0 - learning rate
Batch of clean particles Xo ~ Po-
Niter number of iterations, B batch size
Steps: Sample initial particles N(0,1d)
Nooise number of noise levels per batch for i to Niter do
Set At = (fmux fmmin) //min//T
for
to 0 do
Sample a batch B of clean particles Xp P(Xg)
Set the noise level mmi i
Set Z Z
for = 1 to No do
for to N₂ 1 do
Sample noise level ln ~ ((0 1]
Use
100
to 2 t; 0*)
compute
Sample Xt,, ~ "p( (X(,,,) Xo, (n)
fug Nx,,NN. (Xo)
Let the clean and noisy features be
Z Zi"
V. fV
16 = =((oo,tnn 0n
(Xo)
Xin = =((,,, tn:
end for
Pt
Set Z Z,N
For linear base kernel (1] use opti-
end for Output Z
mized (19) to compute MMD loss
Compute the loss Lix (f, fn) using end for
Compute total loss
(i)(() Nugue Ever Lum (0,
Update discriminator features < ADAM(6, ())
end for
Final denoising. In diffusion models Hoo et al. 2020), it is common to use denoising step at the end to improve sample quality. We found empirically that few MMD gradient flow steps at the end of the sampling with higher learning rate 1) allowed to reduce noise and improve performance.
5 SCALABLE DMMD WITH LINEAR KERNEL
The computational complexity of the MMD estimate two sets of N samples is O(N2). This is likewise the cost of computing the witness function (10) at particles, when computed using N clean and noisy particles. However, using linear base kernel (see (2)
kbase (I, y) = (x. y)-
(11)
allows to reduce the computation complexity both quantities down to O(N), see Appendix We consider the average noise conditional discriminator features on the whole dataset
C.3
5(Xo, t; 0*) = + ΣWW. o(Xi, 0*).
(12)
With the linear kernel (II) we can use average features (12) in the second term of (T).. practice, we can precompute these K- K-dimensional features for T timesteps and store them in memory for later use for sampling purposes, with storage cost of OTT)). This removes the need to store store the N data sample {X)1 since we only retain the average feature vector of this sample.
Approximate sampling procedure. MMD gradient flow (9) requires us use multiple interacting particles Z to produce samples, where the interaction is captured by the first term in (10). In practice this means that the performance will depend on the number of these particles. In this

---

Published as a conference paper at ICLR 2025
Clein esta
DNAB (OUn))
NVD GAN
Lnsariint
0.05 e.1.D
Figure I: Samples from MMD Gradient flow with different parameters for the RBF kernel.
section, we propose an approximation to MMD gradient flow with a linear base kernel [TT) which allows to sample partieles individually and independently, therefore removing the need for multiple particles. For a linear kernel, the interaction term in (1O) for a particle Z, equals to NN, Ee (ZZ,"", 4 0)), of (2, 1; 0*))- For a large number of particles Np. the contribution of each particle Z on the interaction term with Z will be small. For a sufficiently large Np. we hypothesize 11
that No E 2lZ ,1;0") R + Σラ= o(x;, (1))), 0*), where N the size of the dataset and X are produced by the forward diffusion process (6 applied to each Xg. Thus, we consider an approximate witness function
(()) = (d(2, 1; 07), (x, - -000-~((00) (xx,,1;0))
(13)
with /0 ((x,, t;0') 0*) precomputed using 12 Once again, crucially, we have no need to store the sample
{X/) fj=1, since we need only retain its feature mean. We may then sample single particle Z
N(0, Id) and follow noise-adaptive MMD gradient flow with (13) i.. 2, 2 - - pp.
The corresponding algorithm is described in Appendix C.4
RELATED WORK
Ad Adversarial training and MMD-GAN. Integral Probability Metries (IPMs) are good candidates to define discriminators in the context of generative modeling, since they are well defined even in the case of distributions with non-overlapping support (Muller 1997). Moreover, implementations of f-divergence discriminators in GANs rely on variational lower bounds (Nowozin et al., 2016): as noted earlier, these share useful properties of IPMs in theory and in practice (notably, they remain well defined for distributions with disjoint support, and may metrize weak convergence for sufficiently rich witness function classes (Arbel et al. [202T] Proposition 14) and (Zhang et al. (018)), Several works (Arjovsky et al. 2017 Gulrajani al. 2017 Genevay al. 2018 Li et al. 2017 2017|Binkows eet al. (2018) have exploited IPMs as discriminators for training GANs, where the IPMs are MMDs using (linear nonlinear) kernels defined on learned neural net features, making them suited to high dimensional settings such as image generation. Interpreting the IPM-based GAN discriminator as a squared MMD yields an interesting theoretical insight: |Fancceshhi et al. (2022) show that training GAN with an IPM objective implicitly optimizes MMDD" in the Neural Tangent Kernel (NTK) limit (Jacot et al.) 2018). IPM GAN discriminators are trained jointly with the generator in a min-max game. Adversarial training is challenging, and can suffer from instability, mode collapse, and misconvergence (Xiao et al. 2022; Binkowski et al. 2018 Li et al. 2017 Arora et al. 2017] Kodali et al. 2017 Salimans et al. (016)) Note that once a GAN has been trained, the samples can be refined via MCMC sampling in the generator latent space (e.g., using kinetic Langevin dynamics;
see Ansari et al. 2021 Che et al. 2020 Arbel et
2021.
Discriminator flows for generative modeling. Wasserstein Gradient flows (Ambrosio et al. 2008
Santambrogio 2015) applied to a GAN discriminator are informally called discriminator flowss see (Franceschi et al. 2023). A number of recent works have focused replacing GAN generator by discriminator flow. Fan et al. (0022) propose a discretisation of JKO [Jordan et al. 19 1998) scheme to define Kullback-Leible (KL.) divergence gradient flow. Other approaches have used a discretized
interactive particle-hased approach instead of JKO, similar to (3). HHng et al. (2023) ; Franceschi et al.)
(2023) build such flow based on f-divergences, whereas Franceschi al. (2023) focuses on MMD gradient flow. In all these works, an explicit generator is replaced by corresponding discriminator flow. The sampling process during training is as follows: Let Y& be the samples produced training iteration & by the gradient flow F& induced by the discriminator D; applied to samples YK-1 from the previous iteration. We denote this by Y& < FD, 11).. Then, the discriminator at iteration l +1 1 is trained on samples Y₁. A challenge of this process is that the training sample for the next discriminator will be determined by the previous discriminators, and thus the generation process is
7

---

Published as a conference paper at ICLR 2025
still adversarial: particle transport minimizes the previous discriminator value, and the subsequent discriminator is maximized on these particles. Consequently, it is difficult to control or predict the overall sample trajectory from the initial distribution to the target, which might explain the performance shortfall of these methods in image generation settings. By contrast, we have explicit control over the training particle trajectory via the forward noising diffusion process.
Furthermore, these approaches (except for [Hnn et al., 2023) require to store all intermediate discriminators D,,..., Dy throughout training (N is the total number of training iterations). These diseriminators are then used to produce new samples by applying the sequence of gradient flows FN ..... F -)) to Y, sampled from the initial distribution. This creates large memory overhead. An alternative is to use pretrained features obtained elsewhere or a fixed kernel with empirically selected hyperparameters (see Hertrich al., 2024, Hagemann al., 2024 Altekrüger et al. 2023), however this limits the applicability of the method. To the best of our know ledge, our approach is the first to demonstrate the possibility to train a discriminator without adversarial training. such that this discriminator can then be used to produce samples with a gradient flow. Unlike the alternatives, our approach does not require to store intermediate discriminators.
MMD for diffusion refnemmen/regularization. MMD has been used to either regularize training of diffusion models (Li & van der Schaar, (024) or to finetune them (Aiello et al.) 2024) for fast sampling. The MMD kernel in these works has the form ( with Inception features (Szegedy et al.. 2015). Our method removes the need to use pretrained features by training the MMD discriminator.
Diffusion models.
Diffusion models (Sohi-Dickstein al. 2015 Ho et al. 2020 Song al. 2021
represent a powerful new family of generative models due to their strong empirical performance in many domains (Saharia et al., 2022, Le al. 2023 Ho et al. 2022 Watson al.) 2023 Poole et al. 2023). Unlike GANs, diffusion models do not require adversarial training. At training time, denoiser learned for multiple noise levels. As noted above, our work borrows from the training of diffusion models, we train a discriminator on multiple noise levels of the forward diffusion process (Ho et al., 2020). This gives better control of the training samples for the (noise adapted) discriminator than using an incompletely trained GAN generator. We may also consider the setting of flow matching, which is related to the diffusion setting. The potential advantages of our approach to flow matching are discussed in Appendix G
Predictor-corrector: Backwards diffusion might produce samples at time which do not correspond to the forward process the same time. To fix this discrepancy, one can leverage the predictor- corrector scheme (Song et al. (021). At sampling time I, we run a Langevin algorithm for a few iterations targeting P. (corrector) before performing the jump to another noise level (predictor). Our Algorithm 2 can be interpreted as a "corrector only" scheme, where instead of using the Langevin algorithm, we use MMD flows.
EXPERIMENTS
Understanding DMMD behavior in 2-D. Our aim is to get understanding of the behavior of DMMD described in Section We expect DMMD to mimic GAN discriminator training via noise conditional discriminator learning. To see whether this manifests in practice, we design an experiment with a Radial Basis Function (RBF) kernel for MMD, (2 3) = expl || y /(200 (t; 0))], where the noise dependent kernel width function af; 0) [0. + [0, bo) is parameterized by K This parameter controls the coarseness of the MMD discriminator. We consider 2-D checkerboard dataset, see Figure [ left. We learn noise-conditional kernel widths (t; 6) using a neural network. As baselines, we train MMD-GAN where the discriminator learns a, as well as MMD gradient flow with fixed values of a and with a manually selected noise-dependent ar (1) = 0. - 1) + 0.54, i.e. linear interpolation. All experimental details are provided in Appendix H
We report the learned RBF kernel widths for DMMD in Figure 2 left. As expected, as noise level goes from high to low, the kernel width o(t) decreases. In Figure 2 center, we show the learned MMD-GAN kernel width parameter a as a function of training iterations. As the training progresses, this parameter decreases, since the corresponding generator produces samples close to the target distribution. The behaviors of DMMD and MMD-GAN are quite similar and so as the range of values for the kernel widths is also similar. This highlights our point that DMMD mimies the training
8

---

Published as a conference paper at ICLR 2025
of a GAN discriminator. The exact dynamics for (t)) in DMMD depends the parameters of the forward diffusion process (6) The sharp phase transition is consistent with the phase transition highlighted in Section 3. In addition, we report MMD P:t) for different methods in Figure 2 right. We see that DMMD behaves similarly to linear interpolation, but is more nuanced for higher noise levels. The samples are reported in Figure I DMMD produces samples which are visually better than the other baselines. For the RBF kernel, noticed the presence of outliers. The amount of outliers generally depends on the kernel, see Appendix of (Hertrich et al. 2024) for more details.
Image generation We study the performance of DMMD on unconditional image generation of CIFARIO (Krizhevsky & Hinton 2009). We use the same forward diffusion process as in (Ho et al.
2020) to produce noisy images. We use a U-Net (Ronneberger et al.
2015) backbone for discriminator
feature network (zx, t; 0), with slightly different architecture from (Ho et al. 2020), see Appendix [
For all image-based experiments, we use a linear base kernel
[1] We explored using other kernels
such RBF and Rational Quadratic (RQ), but did not find an improvement in performance. We use FID (Hessel et al. 2017) and Inception Score (Salimans et al., (016)) for evaluation, see Appendix [ Unless specified otherwise,
the number N₂- 200 of particles for Algorithm [ We perform
ablation over the number of particles in Appenidx L.3 The total number of iterations for DMMD equals T x N,, where T is the number of noise levels and N, is the number of steps per noise level. For consistency with diffusion models, we call this number of function evaluations (NFE) and we show performance of DMMD with different NFEs. As we show in Appendix P (see Table 8 there is an improvement on FID as we increase NFEs, but only up to a point (NFE=250).
ANY Ww
Table 1: Unconditional generation on CIFAR- 10. For MMD GAN (orig.), we used mixed-RQ kerned (see (Binkowski et al., 2018)). "Orig." original paper, "impl." our implementation. For JKO-Flow (Fan et al. 2022), the NFE is taken from their Figure 12.
c.3
12
24
10
-12
Figure 2: Toy experiment. Left, learned RBF kernel widths (t) for DMMD. Center, for MMD-GAN as function of training iterations.
Method
FID
IS
NFE
MMD GAN (orig.)
39.90
6.51
MMD GAN (impl.)
13.62
8.93
Right, MMD"(P,, P;t) for different methods.
DDPM (orig.)
3.17
9.46
1000
DDPM (impl.)
5.19
8.90
100
Discriminator flow baselines
Table 2: Approximate sampling performance on CIFARIO. IS stands for Inception score
DGGF-KL
28.80
110
JKO-Flow
23.10
7.48
150
Method
FID
IS
NFE
MMD flow baselines
DMMD
8.31
9.09
100
MMD-GAN-FIN Flow
450
1.21
100
DMMD-e
8.21
8.99
102
GS M-M-MDDR RK
55.00
86
a-DMMD
24.86
9.10
50
DMMD (ours)
8,31
9.09
100
a-DMMD-e
9.185
8.70
52
DMMD (ours)
7.74
9.12
250
a-DMMD-a
11.22
9.00
52
As baselines, we reimplement MMD-GAN (Binkowski et al. 2018) with linear base kernel and DDPM (Ho et al. 2020) using the same NN backbones as for DMMD. We also report results from the original papers. We further consider baselines based on discriminator flows: JKO-Flow (Fan et al., 2022), which uses JKO (Jordan et al. 1998) scheme for the KL gradient flow; and Deep Generative Wasserstein Gradient Flows (DGGF-K) KL) (Heng et al. 2023), which uses particle-
approach (as in (3) for the KI. gradient flow. These methods adversarial training to train discriminators (see Section 6). We further compare against Generative Sliced MMD Flows with Riesz Kernels (GS-MMD-RR RK) (Hertrich et al., (024) which uses similar particle-based approach to DGGF-KI KL to construct MMD flow, but with a fixed (kernel) discriminator. We also report results using a discriminator flow defined on trained MMD--AV GAN discriminator, which we call MMD-GAN- Experimental details are given in Appendix ! The results are provided in Table [
We see that DMMD performs better than MMD GAN. As expected, LIMD-GAN-FIL does not work. This is because the MMD--AAN discriminator at convergence was trained on samples close to the target distribution. For the RBF kernel experiment, this means that the gradient of MMD will be very small on samples far away from the target distribution. This highlights the benefit of

---

Published as a conference paper at ICLR 2025
adaptive MMD discriminators. Moreover, DMMD performs better than GS-MMD-RO which uses fixed kernel. This highlights the advantage of learning discriminator features in DMMD. DMMD achieves superior performance compared to other discriminator flow baselines. We believe that one of the reasons for the under-performance of these methods is adversarial training, which makes the hyperparameter choice tricky. DMMD, on the other hand, relies simple non-adversarial training procedure from Algortihm[] [ Finally, we see that DDPM performs better than DMMD. This not surprising, since both, U-Net architecture and forward diffusion process 6 were optimized for DDPM performance. Nevertheless, DMMD demonstrates strong empirical performance as a discriminator flow method trained without adversarial training. Samples from our method are L.I We provide results on CELEB-A, LSUN Church and MNIST below.
provided in Appendix
Approximate sampling. We run approximate MMD gradient flow (see Section [ with the same discriminator as for DMMD. We call this variant e-DMMD, where stands for approximale. We then refine the solution of a-DMMD in a "denoising" step by taking two additional MMD gradient flow steps with a higher learning rate, using either the approximate gradient flow, which we call a- 1-DMM-D,, or the exact gradient flow (9) applied to single particle, which we call o-DMMD-e, where e stands for exact. Finally, for reference, we add a final denoising step to the original DMMD flow, which we call DMMD-e. Results are provided in Table 2 We observe that a- a-DMMD performs worse than DMMD, as expected. Applying a denoising step improves performance of 4-DMMD, bringing it closer to DMMD. This suggests that the approximation (13) moves the particles close to the target distribution; but once close to the target, a more refined procedure is required. By contrast, we see that denoising helps DMMD only marginally. This suggests that the exact noisecconditional witness function (10) accurately captures fine detail close to the target distribution.
Results on MNIST, CELEB-A (64x64) and LSUN-Church (64x64) Besides CIFAR-10, we study the performance of DMMD on MNIST (Lecun et al., 1998), CELEB-A (64x64 (Liu et al.) 2015) and LSUN-Church (64x64) Yu et al. 2015). For MNIST and CELEB-A, we consider the same splits and evaluation regime in (Franceschi ct al. 2023). For LSUN Church, the splits and the evaluation regime are taken from (Ho et al. (020). For more details, see Appendix [.] where we also provide results for the high resolution dataset CELEB-A-HOQ (128x128), and results CELEB- A-HQ (256x256) using DMMD in latent space. As baselines, we consider our implementations of DDPM (Ho et al.) 2020), MMD-GAN (Binkowski et 2018). In addition to DMMD, we report the performance of Discriminator flow baseline from (Franceschi et all 2023) with numbers taken from the corresponding paper. This baseline uses adversarial training together with MMD gradient flow to produce samples. The results are provided in Table We see that DMMD performance is better compared to the discriminator flow and MMD-GAN, which is consistent with our findings on CIFAR-10. 10. It also underperforms compared to DDPM. The corresponding samples are provided in Appendix [.22
Table 3: Unconditional generation on additional datasets. The metric used is FID.
Dataset
MMD-GAN
DDPM
DMMD
Disc. flow
MNIST
7.0
1.94
3.0
4.0
CELEB-A
12.1
6.72
8.3
41.0
LSUN
8.4
3.84
6.1
8 CONCLUSION
In this paper we have presented a method to train a noise conditional discriminator without adversarial training, using a forward diffusion process. We use this noise conditional discriminator to generate samples using a noise adaptive MMD gradient flow. We provide theoretical insight into why an adaptive gradient flow can provide faster convergence than the non-adaptive variant. We demonstrate strong empirical performance of our method on uncoditional image generation of CIFAR 10, as well as on additional, similar image datasets. We propose scalable approximation of our approach which has close to the original empirical performance.
number of questions remain open for future work. The empirical performance of DMMD will be of interest in regimes where diffusion models could be ill-behaved, such as in generative modeling on Riemannian manifolds; as well as on larger datasets such as ImageNet. DMMD provides a way of training a discriminator, which may be applicable in other areas where a domain discriminator might be required. Finally, it will be of interest to establish theoretical foundations for DMMD in general settings, and to derive convergence results for the associated flow.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Aiello, E.. Valsesia, D., and Magli, E. Fast inference in denoising diffusion models via MMD finetwning. IEEE Access, 12: 121106922-00993, 2024, doi: ATL.:///AACESS22224.44333 URL
https: //d/://doi.o ww..orggooo. 10. 1 109/ 2024 3436698
Albergo, M. S., Boffi, N. M., and Vanden-Fijnden, E. Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint arXiv:2303.08797 2023.
Altekrüger, F., Hertrich, J., and Steidl, G. Neural wasserstein gradient flows for maximum mean discrepancies with riesz kernels, 2023.
Ambrosio, L., Gigli, N., and Savaré, G. Gradient Flows in Metric Spaces and in the Space of Probability Measures. Lectures in Mathematics ETH Zürich. Birkhäuser, 2. ed edition, 2008. ISBN 978-3-7643-8722
OCLC: 254181287.
Ansari, A. F., Ang. M. L., and Soh, H. Refining deep generative models via discriminator gradient flow. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
net forr/forumm
id=2bc-ue9p ZE
Arbel, M., Sutherland, D. J., Binkowski, M., and Gretton, A. On gradient regularizers for MMD gans. In Bengio. S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 6701-6711, 2018. URL http::ttttt:oceeeecaaacc -c html
neurips cc wwtpppppee//.///aaaaeeeeeee
Arbel, M., Korba, A.. Salim, A., and Gretton, A. Maximum mean discrepancy gradient flow. In Wallach, H. M., Larochelle, H., Beygelzimer, A., d' Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.),
Advances in Neural Information Processing Systems 32: Anmual Conference
Neural Information Processing Systems 2019, NewrlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp.
nwwrinu.pppp..m wttps.cc/pae wtt.cc/pppp//21.01 22009//hase 9/hash/ 6481-6492 2019. URL /tttpp://proceeeee Abstract
at html
Arbel, M.. Zhou, L., and Gretton, A. Generalized energy based models. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
http:/oorrmmm.dddPPtPPPP.3P.. d=OPtUPR ew www.://oo///e/e./
Arjovsky, M., Chintala, S., and Bottou, L. Wasserstein generative adversarial networks. In Precup, D.
and Teh, Y. W. (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 214 223. PMLR, 06-11 Aug 2017. URL httpppp//pproceediicc
ml prwss/ wttppppsssssp/apppp.eaa /ttp//////////////////
html
Arora, S., Ge, R., Liang, Y., Ma, T., and Zhang, Y. Generalization and equilibrium in generative adversarial nets (gans). In Precup, and Teh, Y. W. (eds.). Proceedings of the 34th Interna- tional Conference on Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, PMLR, 2017. URL
volume 70 of Proceedings of Machine Learning Research, pp.
http://proceedi ngs. mlr, www.sressseee.... www00///rraaaa. html
Bansal, A., Borgnia, E., Chu, H.-M., Li, J., Kazemi, H., Huang, F., Goldblum, M., Geiping, J., and Goldstein, T. Cold diffusion: Inverting arbitrary image transforms without noise. Advances in Neural Information Processing Systems, 36:41259 (6:41259-41222 41282, 2023.
Binkowski, M., Sutherland, D. J., Arbel, M., and Gretton, A. Demystifying MMD gans. In 6th International Conference on Learning Representations, ICLR 2018, Vanconver, BC, Canada, April 30- May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL https
Bortoli, D. Convergence of denoising diffusion models under the manifold hypothesis. Transactions on Machine Learning Research, 2022. ISSN 2835-8855 URL https: httpp:://ppereeview.. Forum? Marrmmidd-MKS5XXooo Expert Certification.

---

Published as a conference paper at ICLR 2025
Brock, A., Donahue, J., and Simonyan, K. Large scale GAN training for high fidelity natural image synthesis. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL hatta::/ net cest/forumt forum? dd= / X
Brown, B. C., Caterini, A. L., Ross, B. L.,, Cresswell, J. C., and Loaiza-Ganem, G. The union of manifolds hypothesis and its implications for deep generative modelling. arXiv preprint arXiv:2207.02862, 2022.
Che, T., Zhang, R., Sohl-Dickstein, J., Larochelle, H., Paull, L., Cao, Y., and Bengio, Y. Your gan is secretly an energy-based model and you should use discriminator driven latent sampling. In Proceedings of the 34th International Conference Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Chen, S., Chewi, S., Li, J., Li, Y., Salim, A., and Zhang, A. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In The Eleventh International Conference on Learning Representations, 2023. URL htLhhtss: https:://ppenreview net/forum? id=zyLVMgs20U
Daras, G., Delbracio, M., Talebi, H., Dimakis, A., and Milanfar, P.
Soft diffusion: Score matching with general corruptions. Transactions Machine Learning Research, 2023. ISSN 2835-8856.
URL htttps:/
net forum? wit.ummdddd-99Bexxx....
Fan, J., Zhang, Q.. Taghvaei, A., and Chen, Y. Variational Wasserstein gradient flow. In Chaudhuri, K., Jegelka, S., Song. L., Szepesvari, C., Niu, G., and Sabato, S. (eds.), Proceedings of the
39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pp. 6185 6215. PMLR, 17-23 Jul 2022. URL httttttt/ppcceeeantttt mlr. press vl www.2//fan2nd... html
Fefferman, C., Mitter, S., and Narayanan, H. Testing the manifold hypothesis. Journal of the American Mathmatical Society, 29(4))983-1049 2016
Franceschi, J., de Bézenac, E.,, Ayed, I., Chen, M., Lamprier, S., and Gallinari, P. A neural tan- gent kernel perspective of gans. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvári, C., Niu, G., and Sabato. S. (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Re- search, pp. 6660-6704. PMLR, 2022. URL htttttsssspppppeeeiiiiete mm press/v htt.:/s...s...../////e....... 62/
Franceschi22e html
Franceschi, J.-Y., Gartrell, M., Santos, L. D., Issenhuth, T., de Bézenac, E., Chen, M., and Rakotoma- monjy, A. Unifying gans and score-based diffusion as generative particle models. In Proceedings of the 37th International Conference on Neurul Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2023. Curran Associates Inc.
Gao, R., Hoogeboom, E., Heek, J., Bortoli, D., Murphy, K. P., and Salimans, T. Diffusion meets low. g thub
flow matching: Two sides of the same coin. 2024. URL httphhtts///ee
io/
Genevay, A., Peyre, G., and Cuturi, M. Learning generative models with sinkhorn divergences. In Proceedings of the twenty-First International Conference on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pp. 1608-1617. PMLR, 2018.
Glaser, P., Arbel, M., and Gretton, A. KALE flow: A relaxed KL gradient flow for probabilities with disjoint support. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, 2021. URL
id=7BeCVTCs
Goodfellow, I., Pouget-Abadie, Mirza, M., Xu, B., Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Generative adversarial nets. In Ghahramani. Z., Welling, M., Cortes, C., Lawrence, N., and Weinberger. K. (eds.). Advances in Neural Information Processing Systems, volume 27. Cur- cc ran Associates, Inc., 2014. URL https: pts.etpppe///procciiinsssssss www.sinnnnneuuuurr httplpps.ccpppppp cctpaper_feesss files/
12

---

Published as a conference paper at ICLR 2025
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. A kerne two-sample test. Journal of Machine Learning Research, 13(25):723-77) 773, 2012. URL http:// /mte://mmrrr... org/ ptterr//v3//xeetooooo
html
Gulrajani, I., Ahmed, F., Arjovsky, M., Dumoulin, V., and Courville, A. Improved training of wasserstein gans. In Proceedings of the 31st International Conference on Neural Information Processing Systems. NIPS' 17, pp. 5769-5779, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Hagemann, P.,
Hertrich, J., Altekrügger, F., Beinert, R., Chemseddine, J., and Steidl, G. Posterior sampling based on gradient flows of the MMD with negative distance kernel. The Twelfth International Conference Learning Representations, 2024. http::ttttpppp/aaeeeeett
Heng, A., Ansari, A. E,, and Soh, H. Deep generative wasserstein gradient flows, 2023, URL https://openrevient h
Hertrich, J., Wald, C., Altekrüger, F., and Hagemann, Generative sliced MMD flows with riesz kernlss. In The Twelfth International Conference on Learning Representations, 2024. URL
btttp:://operreveee wwweenreieew.ne/fooruiii wttp:/ooooommmmdd/.
Heusel, M., Ramsauer, H., Unterthiner, T., Nessler, B., and Hochreiter, Gans trained by a two time- scale update rule converge local nash equilibrium. In Proceedings of the 31 International Conference on Neural Information Processing Systems, NIPS' 17, pp. 6629-6640, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 9781510860964.
Ho, J., Jain, A., and Abbeel, Denoising diffusion probabilistic models. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020. Curran Associates Inc. ISBN 9781713829546.
Ho, J., Chan, W., Saharia, C., Whang, J.. Gao, R., Gritsenko, A., Kingma, D. P., Poole, B., Norouzi, M., Fleet, D. J., et al. Imagen video: High definition video generation with diffusion models. arXiv preprint ar 22X1:22200.02303, 2022.
Hyvärinen, A. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research, 6(24):695-709, 2005. URL httt::/ttp/mittttt http:://mm...rr//ppppess///////s
hyvarinenO5a.c ntm
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent kernel: convergence and generalization in neural networks. Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS' 18, pp. 8580-8589, Red Hook, NY, USA, 2018. Curran Associates Inc.
Jordan, R., Kinderlehrer, D., and Otto. F. The variational formulation of the fokker- equation. SIAM Journal on Mathematical Analysis, 1911):1-17, 1998. doi: 10. 1 101137770333111119003333339 URL https: // .org 0. 111111100031111030300303
Karras, T.. Aittala, M., Hellsten, J.. Laine, S., Lehtinen, J., and Aila, Training generative adversarial networks with limited data. In Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20, Red Hook, NY, USA, 2020a. Curran Associates Inc. ISBN 9781
Karras, T., Laine, S., Aittala, M., Hellsten, J., Lehtinen, J., and Aila, T. Analyzing and improving the image quality of stylegan. In 2020 JEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 8107-8116, 2020b. doi: 00.1009CCVP42260.2200.000
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In Bengio, Y. and LeCun, Y. (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL
www./orrg/aaa abs/
1412. 6980]
Kodali, N., Abernethy, J., Hays, J., and Kira, Z. On convergence and stability of gans, 2017. URL https: fttps://errivve
http:///wwwwwvo./oo/////// www.ooogg////ooooooo 07215
13

---

Published as a conference paper at ICLR 2025
Krizhevsky, A. and Hinton, G. Learning multiple layers of features from tiny images. Technical Report 0, University of Toronto, Toronto, Ontario, 2009. URL https: /ww. cs toronto eddu/krizz Sewriizzeeannnngg-eeeeee
2009-TR-------- pdf
Le. M., Vyas, A., Shi, B., Karrer, B., Sari, L., Moritz, R., Williamson, M., Manohar, V., Adi, Y., Mahadeokar, J., and Hsu, W.-N. Voicebox: text-guided multilingual universal speech generation at scale. Proceedings of the 37th International Conference on Neural Information Processing Systems, NIPS '23, Red Hook, NY, USA, 2023, Curran Associates Inc.
Lecun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied to document recognition. Proceedings of the IEEE, (6(11))2278-223 23244, 1998. doi: 10.1109/5.72671
Li, C.-L.,, Chang, W.-C., Cheng. Y., Yang. Y., and Póczos, B. Mmd gan: towards deeper understanding of moment matching network. In Proceedings the 31st International Conference on Neural Information Processing Systems, NIPS' 17, pp. 2200-2210, Red Hook, NY, USA, 2017. Curran
Associates Inc. ISBN
Li, Y. and van der Schaar, M. error propagation of diffusion models. In The Twelfth International Conference on Learning Representations, 2024. htlpttttps///pppeeeeee..e... forum?
SIERAAct1E223
Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., and Le, M. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2023. URL wwt.a//.eeeeeeeeeeee.eeeee forum? 1=
https: https:///penrest pppnnevieee net.
Liu, M.-Y., Huang, X., Yu, J., Wang, T.-C., and Mallya, A. Generative adversarial networks for image and video synthesis: Algorithms and applications. Proceedings of the IEEE, 10965):839-862 2021. doi: 10.1109//PROC.202 20233.0999.... 1.349999..
I Liu, Z., Luo, P., Wang, X., and Tang, X. Deep learning face attributes in the wild. In 2015 IEEE International Conference on Computer Vision (ICCV), pp. 3730-3738, 2015, doi: 10.1 109/CV. 2015.425.
Mescheder, L.., Geiger, A., and Nowozin, S. Which training methods for gans do actually converge? In International conference on machine learning, pp. 3481-3490. PMLR, 2018.
Muller, A. Integral probability metrics and their generating classes of functions. volume 29, pp. 429-444 Advances in Applied Probability, 1997.
Nowozin, S., Cseke, B., and Tomioka, R. f-gan: training generative neural samplers using variational divergence minimization. In Proceedings of the 30th International Conference on Neural Informa- tion Processing Systems, NIPS' 16, pp. 271-279, Red Hook, NY, USA, 2016. Curran Associates Inc. ISBN 9781510888 10838819.
Pidstrigach, J. Score-based generative models detect manifolds. Advances in Neural Information Processing Systems, 35:35852-35865, 2022.
Poole, B., Jain, A., Barron, J. T., and Mildenhall, B. Dreamfusion: Text-to-3d using 2d diffusion. In The Eleventh International Conference on Learning Representations, 2023. URL https hooeenevvew.net//eee wewtt//ooommiddee NNys577y
Rombach, R., Blattmann, A., Lorenz, D.,, Esser, P., and Ommer, B. High-Resolution Image Synthesis with Latent Diffusson Models In 2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10674-10685. Los Alamitos, CA, USA, June 2022. IEEE Computer Soci- URL htthhtt::://d... www.::ewwweccccccccc.......
ety. doi: 10. 11199C/VPR5288822222000000 org/10. 1Tt.0011999/99222280
01042
Ronneberger, O., Fischer, P., and Brox, T. U-net: Convolutional networks for biomedical image segmentation. In Navab, N., Hornegger, J., Wells, W. M., and Frangi, A. F. (eds.), Medical Image Computing and Computer-Assisted Intervention MICCAI 2015, pp. 234 241,,, Cham, 2015. Springer International Publishing. ISBN 978-3-319-2454444
14

---

Published as a conference paper at ICLR 2025
Saharia, C., Chan, W., Saxena, S., Li, L., Whang, J., Denton, E. L., Ghasemipour, K., Gontijo Lopes, R., Karagol Ayan, B., Salimans, T., et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances in Neural Information Processing Systems, 35:36479-36494 2022.
Salimans, T., Goodfellow, I., Zaremba, W., Cheung, V., Radford, A., and Chen, X. Improved techniques for training gans. In Proceedings of the 30th International Conference Neural Information Processing Systems. NIPS' 16, pp. 2242,, Red Hook, NY, USA, 2016. Curran Associates Inc. ISBN 9781510838819.
Santambrogio, F. Optimal transport for applied mathematicians. Birküuser, NY, 55(58-63):944 2015.
Sauer, A., Lorenz, D., Blattmann,, A., and Rombach, R. Adversarial diffusion distillation, 2023. URL
https://arxiv/ wtt.
17042
Schölkopf, B. and Smola, A. J. Learning with Kernels: Support Vector Machines, Regularization, Optimization, and Beyond. The MIT Press, 06 2018. ISBN 9780262256933, doi: 10.755 Imitpes/mitpress// 4175.001.0000 URL httpt:ttt///ooot http://ti..ogg////// 10. /tt./mmtprrss/////120
0001
J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In Bach, F. and Blei, D. (eds.). Proceedings of the 32nd International Conference on Machine Learning. volume 37 of Proceedings of Machine Learning Research, pp. 2256 2266-2265, Lille, France, 07-09 Jul 2015. PMLR. URL https:// proceedings mlr www.s/eessv.v ptw.ss.s...ppppsppe.......s / sohl dickstein15. html
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based genera- tive modeling through stochastic differential equations. In International Conference on Learning Representations, 2021. /tp.tts://peeeneeeeee.eee Ltw.peeeweeeeeeeeeeeooooooooaaaa Pxt.oooummmiiiIIIIIII PXTIGGRRRSSS
Song. Y., Dhariwal, P., Chen, M., and Sutskever, Consistency models. Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett, J. (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceediiggs of Machine Learning Research, pp. 32211-32252. PMLR, 23-29 Jul 2023. URL httpt: son.g2a.httm
ing.e.nniinn....ooo...... mwt.pperss///S. v202/
Szegedy, C., Liu, W., Jia, Y., Sermanet, P., Reed, S., Anguelov, D., Erhan, D., Vanhoucke, V., and Rabinovich, A. Going deeper with convolutions. In 2015 IEEE Conference Computer Vision and Pattern Recognition (CVPR), pp. 1-9, Los Alamitos, CA, USA, June 2015. IEEE Computer ety
Society. doi: 10.1 orr/100 1109/111119//// 2015. 2015..299999
201. URL nitpss://do http::////dddieeeeccmpoperrrrre ieseecoppttersscc...
Tenenbaum, B., de Silva, V., and Langford, J. C. A global geometric framework for nonlinear dimensionality reduction. Science, 29065500)2211-2222. 2000. doi: 10. 126/scienncee229.550000 2319.
290. 5500
2319
van den Oord, A., Vinyals, O., and Kavukcuoglu, K. Neural discrete representation learning. In Pro- ceedings of the 31 International Conference on Neural Information Processing Systems, NIPS'17, 17, pp. 6309-6318, Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 97815 151660964,
Villani, C.
Optimal Transport: Old and New. Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 2008 ISBN 9783540710509. URL http:::ttt/ooooo google google
co. htkk/ooks?id
5tkC
Vincent, P. A connection between score matching and denoising autoencoders. Neural Computation,
23(7):166111674 2011.
Watson, J. L.,, Juergens, D., Bennett, N. R., Trippe, B. L., Yim, J., Eisenach, H. E., Ahern, W. Borst, A. J., Ragotte, R. J., Milles, L. F., Wicky, B. M., Hanikel. N., Pellock, S. J., Courbet, A., Sheffler, W., Wang, J., Venkatesh, P., Sappington, I., Torres, S. V., Lauko, A., De Bortoli, V.,
Mathieu, E., Ovchinnikov, S., Barzilay, R., Jaakkola, T. S., DiMaio, F., Back, M. and Baker, D. De novo design of protein structure and function with rfdiffusion. Nature. 620(7976) http:://xxxxdd.o.oo/
Jul 2023. ISSN 1476-4687. doi: 10. 10.1008//41186-1 0288860-23-0--5--- URL http: hx doi.
10. 1 00000/ 418/5/11888661 023 004 06415-8
15

---

Published as a conference paper at ICLR 2025
Xiao, Z., Kreis, K., and Vahdat, A. Tackling the generative learning trilemma with denoising diffusion GANs. International Conference on Learning Representations, 2022. URL https
Xu, Y., Zhao, Y., Xiao, Z., and Hou, T. Ufogen: You forward once large scale text-to-image generation via diffusion gans. In 2024 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 8196-8206, 2024. doi: 10. 11009/CVR55233332220..07777
Yang. Z., Feng. R., Zhang, H., Shen, Y. Zhu, K., Huang, L.,, Zhang, Y., Liu, Y., Zhao, D., Zhou, J., and Cheng, F. Lipschitz singularities in diffusion models. In The Twelfth International Conference on Learning Representations, 2024. URL
net/ forum?
id=WNkkOcOwiz
Yu, F., Zhang, Y., Song, S., Seff, A., and Xiao, J. Lsun: Construction of large-scale image dataset using deep learning with humans in the loop (v2017), arXiv preprint 1xXiv:5506.03666 2015.
Zhang. Liu, Q. Zhou, D., Xu, T., and He, X. On the discrimination -eneralizainnnnn disciiiinatoonggenerliizziiooono tradcoff in gans. In 6th International Conference on Learning Representations, 2018
16

---

Published as a conference paper at ICLR 2025
A ORGANIZATION OF THE SUPPLEMENTARY MATERIAL
In Appendix B we briefly describe the adversarial training in the context of Generative Adversarial Networks (GANs). In Appendix C we describe in details the training and sampling procedures for DMMD. In Appendix [ we provide more details about the computational complexity of DMMD during trianing and sampling. In Appendix E. we explain how our approach could be applied to f-divergences resulting in DKALE-Flow and in Appendix [F. we provide more details about this method. In Appendix G we discuss the connection Flow Matching. In Appendix H we describe more details for the 2d experiments, In Appendix II we provide experimental details for the image datasets as well as the additional results. In Appendix K we provide proof for the theoretical results described in Appendix 3 from the main section of the paper. Finally, in Appendix L we present the samples from DMMD on different image datasets.
B ADVERSARIAL TRAINING
We briefly describe the notion of the adversarial training in the context of Generative Adversarial Networks (GANs). The objective function for GANs is
F(B, v) = 1.- D(X
where G(Z, ()) is a generator with parameters v and D[x; G(Z; (b); is discriminator divergence
with parameters 0. The objective for the generator is
arg min max F(0, v/)
6
The objective for the discriminator is
0" f- arg DIaX min F(O, (+)
In practice, we alternate the updates the objective F for the generator and the discriminator. Let N be the iteration number. Then, the generator parameters are updated as
5h ax V F
2 =²)
and the discriminator parameters are updated as
gn+1 + tt + aa = 8t, z)=+1).
where Ogg,, og are learning rates.
When we train DMMD, we could interpret it as having a fixed generator given by the forward diffusion process, and the training is no longer adversarial since there are no min/max optimization going on.
C DMMD TRAINING AND SAMPLING
C. MMD DISCRIMINATOR
Let X' C C and P(X') be the set of probability distributions defined on X. Let P E P(X) be the
target or data distribution and Q2 E ((x)) be a distribution associated with a generator parameterized by V E R. Let H. be Reproducing Kernel Hilbert Space (RKHS), see (Schölkopf & Smola 2018) for details, for some kernel k X' X + R. Maximum Mean Discrepancy (MMD) (Gretton et al.
2012) between Qs and P defined as MMD(Qu. P) sup (/ - E [x)]]. Given
X
(xx) 1= 1= QaN and YM A psM, an unbiased estimate of MMD Gretton
et al. 2012 is given by
MMD M N(N-1) ΣN M(1,,15))
(14)
Σ Σijj k(Ve, Uj)
2 ΣN Σ M ())
MGM
NM
In MMD GAN (Binkowski et al. k (r, y) kiusse (bbx:; (1r; 0)),
2018 Li et al.
2017). the kernel in the objective 14 is given as
(15)
17

---

Published as a conference paper at ICLR 2025
where Kaase is base kernel and ((:;0) X - RK neural networks discriminator features with parameters 0 E R" We use the modified notation of MMD Y 0] for equation (14 to highlight the functional dependence on the discriminator parameters. MMD am instance Integral Probability Metric (IPM) (see (Arjovsky et al. 2017)) which is well defined distributions with disjoint support unlike f-divergences (Nowozin et al. (016). An advantage of using MMD over other IPMs (see for example, Wasserstein GAN (Arjovsky et al. 2017)) is the flexibility choose a kernel k. Another form of MMD is expressed as a norm of a witness function
MMD(Q P) = sup (Eq [f(x)] [x[f(x)]) ll plln: SCH
where the witness function fQ..P is given as
lQ.,p(=)= / Max,))dQu, - M k().
Given of samples X {ai Qoo and y = (wi p&M, and the kernel 15 the empirical witness function is given
fQ..
Khase A (6(1;; ()
Khase Chase(d(153) (0(1;;0),(z;;0))
N
M
The l2 penalty (Binkowski et al. 2018) is defined as
N
Ce (0)
Σ
| 1| +
1
( 0)
N
N
Assuming that M and following (Binkowski et al. 2018
Gulrajani et al. 2017). for oş ~
U10,, 1], where U10 1] is uniform distribution [0, I|. construct 2; 2g0;; + a)9 for all
i = -,,..,, N. Then, the gradient penalty (Binkowski et al.
2018 Guirajani et al. 2017) is defined as
N (())) 1)2
1 N
L (0)
We denote by C()) the MMD discriminator loss given
[()) - MMD yM;0 N(N
Aiuuse (0(2;; (((s;))
Arase (ol:: ((,y;;)
N.
2
Σ
i=
N
M Σ
j=1
( 0))
Abase
Then, the total loss for the discriminator on the two samples of data assuming that N M is given as
(a) L()) (o) De,
for some constants A 2 0 and \r₃ 2 0.
C.2 NOISE-DEPENDENT MMD
In Section we describe the approach to train MMD discriminator from forward diffusion using noise-dependent diseriminators. For that, assume that we given noise level L where U|0, 1] is a uniform distribution on [0, 1J. and a set of clean data X [ p@N Then we produce a of noisy samples z using forward diffusion process 66). We denote these samples by X.N = [ We define noise conditional kernel
U
k(x, y: t, 0) = Anuse (0(x, t; 0), o))
18

---

Published as a conference paper at ICLR 2025
with noise conditional features 6(x, t; 0). This allows to define the noise conditional discriminator loss
C(0,t) [ xx ,t,0
Kbase (oo
(16)
N(N-1)
N
Aibase (b((=;; (,)
N(N
No
Aiaaso (0x",
The noise conditional l2 penalty is given as
(0,
1 N E 01
1
Σ
N
The noise conditional gradient penalty is given
N
(v(0,t) Σ (IVFr,(()))ll 1)?
where 2; = a + (1 - oz)xx* for oi ~ UO0 1] and the noise conditional witness function
1 N
1
fpa(2) =
Aibase (0(x)) 0), (1=;0))
Aibaso (hn:
(17)
N
Therefore, the total noise conditional loss given
(,) C(e, (v + Le, for some constants A 2 0 and λo₃ 2 O.
(18)
C.3 LINEAR KERNEL FOR SCALABLE MMD
Computational complexity of (18) is O(N)). Here, we assume that the base kernel is linear, i.c.
Ribusse (x. y) (r.y) v)
This allows to simplify the MMD computation (16)
[X X 444 0]
N(N - 1) (xx(X:") (xxx") - 1106113" 1106113(xx.) + -
N(N 1) -xxx (1)))) 2
NN (5 (19)
where
5(x?²) Σ 0(s):0))
(x(x)) Σ )(x;;6))
j=1
N
( Σ
|1 0,)11
N
10 = 11o(x";0.)) 0.)11
==1
19

---

Published as a conference paper at ICLR 2025
Algorithm 3 Approximate noise-adaptive MMD gradient flow for single particle
Inputs: T is the number noise levels
Emax: Emin are maximum and minimum noise levels
N₂ is the number of gradient flow steps per noise level
O is the gradient flow learning rate
<(Xo, 1, 0*) precomputed clean features for all 1,.... , T with (0))
((Xe, (, ) precomputed noisy features for all / - 1...., 7 with 20
Steps: Sample initial noisy particle N(0, Id)
for i T to 0 do
Set the noise level = itt and Z Z
for 0 to N, 1 do
2 2 7 VV 1; 0) ,x,,, - o(Xo,t,
end for Set Z N
end for Output Z
Therefore can pre-compute quantities (. (oxll((())))))) (x11)((X which takes
O(N) and compute MMD X X²" , ,1,0) in ((1) time. This also leads ((1) computation complex- ity for Lt, and ((N) complexity for Lv. This means that we simplify the computational complexity to O(N) from (N
At sampling, following D requires to compute the witness function 17 for each particle, which for general kernel takes O(N2) in total. Using the linear kernel above, simplifies the complexity of the witness as follows
fF.((z) (Zx (x 0(; 0)),
where ZN is set of N noisy particles. We can precompute (Z*) in O(N) time. Therefore one iteration of a witness function will take O(1) time and for N noisy particles it makes O(N).
C.4 APPROXIMATE SAMPLING PROCEDURE
In this section we provide an algorithm for the approximate sampling procedure. The only change with the original Algorithm 2 is the approximate witness function
(z) = (d(2, t; 0*), (,,, - o(xo,1,0)).
where
N
2(Xo,
pl 1; l*)
(20)
N
==1
1 N
o(X,, 1,
Σ
( 1; 0*)
Here 26, 1,..., correspond to the whole training set of clean samples and æf, I,.... correspond to the noisy version these clean samples produced by the forward diffusion process (6) for a given noise level These features can be precomputed once for every noise level The resulting algorithm given in Algorithm (B). A second crucial difference with the original algorithm is that it generates single particles Z independently.
D COMPUTATIONAL COMPLEXITY
D.1 TRAINING-TIME COMPLEXITY
Computing total loss in training iteration (see Algorithm I on a batch B samples is O(B²) for an arbitrary kernel and O(B) for a linear kernel (see Appendix C.3) We compute it for N,, noise levels meaning that the cost scales as ON,, for an arbitrary kernel and O(N, B) for a linear kernel. This is more expensive than DDPM (Ho et al. 2020) which cost scales O(B).
20

---

The approach described in Section 4 can be applied to any divergence which has well defined Wasserstein Gradient Flow described by gradient of the associated witness function. Such diver- gences include the variational lower bounds f-divergences, which are popular in GAN training, and were indeed the basis of the original GAN discriminator (Goodfellow et al., 2014). One such f-divergence is the KL Approximate Lower bound Estimator (KALE, Glaser et al. 2021). Unlike the original KI. divergence, which requires density ratio, the KALE remains well defined for distributions with non-overlapping support. Similarly MMD, the Wasserstein Gradient of KALE is given by the gradient of a learned witness function. Thus, we train noise-conditional KALE discriminator and use corresponding noise-conditional Wasserstein gradient This approach Published as a conference paper at ICLR 2025
D.2 SAMMPLING-TINE COMPLEXITY
During sampling time (see Algorithm 2) we use 7 noise levens and do N, steps per noise level with Nr noisy particles. Let ((F) denote the cost of forward pass for the backbone network. The gradient takes ((2F) to compute. Each sampling step for DDPM (Ho et al. 2020) costs O(F), while one sampling step for DMMD costs O(FN)) for an arbitrary kernel, 0 FFF))) for linear kernel and OFF) when using approximate sampling procedure.
E F-DIVERGENCES
described by (Nowozin et al. (2066),
flow, as with DMMD. We call this method Diffiusion KALE flow (D-KALE-FFI
described in Appendix [F] We found this approach to lead to reasonable empirical results, however unlike with DMMD, it achieved worse performance than corresponding GAN, see Appendix
F D-KALE-FLOW
In this section, describe the DKALE-flow algorithm mentioned in Section Let X' C R and ((t) be the set of probability distributions defined on X. Let P E P(()) be the targer or data distribution and Q E ((X) be some distribution. The KALE objective (see (Glaser et al. 2021 is defined as
/ hdQ - -e
K KALE = (1 max
(21)
NEH
where À 2 O is a positive constant and H. is the RKHS with a kernel k. In practice, KALE divergence can be replaced by a corresponding parametric objective
A KALE(Q, = (1+ ( ((xx8,0))Q(()) - / Ch - 11a113) (22)
where
h(X;0,a) (xx;0aa)=(xx00))) o)
with o( ((;;00 0) E D and E RDD The objective function (22) can then be maximized with respect to and a for given Q and P. Similar to DMMD, we consider a coise-conditional witness function
l[;;; 1, 0, a, vb) = <2;; 1, ()
From here, the noise-conditional KALE objective is given
L(8 sh, t/A) = K KALE P)),0,a),
where P, is the distribution corresponding total coise-conditional objective is given as
forward diffusion process, see Section Then, the
Lua (0, w. (A) = C(O, v. f/A) Lo Xe, Le, (0,t),
where gradient penalty has similar form to WGAN-GP (Gurajjani et al. 2017
1. = =
a, e) 1)2,
where 7 -3X 1--8888,8 UO. P(X) and Y P(Y). The 12 penalty is given
1 - (Ex-pix)l 2 Ey-
Leg(0 (0.
21

---

Published as a conference paper at ICLR 2025
Algorithm 4 Noise-adaptive KALE flow for single particle
Inputs: T is the number of noise levels
Emax: Emin are maximum and minimum noise levels
N₂ is the number of gradient flow steps per noise level
> O is the gradient flow learning rate
Trained witness function h(:; I, 0*,, xb*)
Steps: Sample initial noisy particle Z ~ N(O, Id) Set A/ mmmxx m/mi)//T
for i T to 0 do
Set the noise level = tmin + it and Zó =
for 0 to N, 1 do
77 v/ ; ,
end for
Set Z = N
end for Output Z
Therefore, the final objective function to train the discriminator is
Lox(0, (,, 3 (- = 10 [oor (0, sb, t/A]]
This objective function depends on RKHS regularization A, on gradient penalty regularization Äv and on 12-penalty regularization At. Unlike in DMMD, we do not use an explicit form for the witness function and do not use the RKHS parameterisation. On one hand, this allows to have more scalable approach, since we can compute KALE and the witness function for each individual particle. On the other hand, the explicit form of the witness function DMMD introduces beneficial inductive bias. In DMMD, when we train the discriminator, we only learn the kemnel features, i.e. corresponding RKHS. In D-KALE, we need to learn both, the kernel features Φ(x; I, 0) as well as the RKHS projections att; ).) This makes the learning problem for D-KALE more complex. The corresponding noise adaptive gradient flow for KALE divergence is described Algorithm| 4 An advantage over original DMMD gradient flow is the ability to this flow individually for each particle.
G CONNECTION TO FLOW MATCHING
Diffusion models (Ho et al. 2020) try to reverse a Gaussian noising dynamic while flow match- ing (Lipman et al. 2023) methods try to match interpolation dynamics between the data distribution and a Gaussian distribution. Both methods can be shown to be equivalent given the choice of the schedule (Gao et al. 2024).
Our approach is agnostic to the destruction process and does not rely on an underlying Gaussian assumption of the prior distribution 2 In particular, one could design destruction processes which are more suitable to the problem at hand. While extensions have been proposed for such destruction processes in the case of Gaussian diffusion (see (Bansal al., 2023 Daras al. [2023 for instance) they either rely on soft Gaussian assumption for the noising process (Daras et al. (2023) or do not enjoy the theoretical properties of our framework (Bansal al. 2023): optimality we recover the data distribution, up to the discretization error, regardless of destruction process or prior distribution.
H TOY 2-D DATASETS EXPERIMENTS
For the 2-D experiments, we train DMMD using Algorithm (T) for Nuer = 50000 steps with a batch
size of B 256 and number of noise levels per batch equal to Nnoise 128. The Gradient penalty constant \ = 0.1 whereas the l penalty is not used. To learn noise-conditional MMD for DMMD, we use 4-layers MLP g(t; 0) with ReLU activation to encode (t; 0) amin ReLU 0)) with O 0.001, which ensures a (t: 0) > O. The MLP layers have the architecture of [64, 32, 16, 1].
Note that this might also not be necessary in the case of flow matching and stochastic interpolant, Albergo al. 2023 for instance
22

---

Published as a conference paper at ICLR 2025
Before passing the noise level E [0 to the MLP, used in (Ho et al. 2020), which produces a feature vector of size 1024. The forward diffusion process from (Ho et al. (020) have modified parameters such that corresponding 31 10 4, Br 0.0002. On top of that, we discretize the corresponding process using only 1000 possible noise levels, with
sinusoidal embedding similar to the one
min 0.05 and max 1.0. At sampling time for the algorithm [ we use mmin 0.05, mma 1.0, N₂ = 10 and T = 100. The learning rate 1.0. As basleines, we consider MMD-GAN with a generator parameterised by a 3-layer MLP with ELU activations. The architecture of the MLP is [256, 256, 2]. The initial noise for the generator is produced from a uniform distribution - with dimensionality of 128. The gradient penalty coefficient equals to 0.1. As for the discriminator, the only learnable parameter is a. We train MMD-GAN for 250000 iterations with batch size of B = 256. Other variants of MMD gradient flow use the same sampling parameters as DMMD.
We used 1 a100 GPU with 40GB of memory to run these experiments. In total, all the experiments took less than 2 hours.
I IMAGE GENERATION EXPERIMENTS
dataset. We use the same forward diffusion process in (Ho et al. 2020). As Neural Network backbone. we use U-Net (Ronneberger et al. 2013) with a slightly modified architecture from (Ho et al. 2020). Our neural network architecture follows the backbone used in (Ho et al.) 2020). On top of that output the intermediate features at four levels- - before down sampling, after womp-samm,ing, before upsampling and a final layer. Each of these feature vectors is processed by group normalization, the activation function and a linear layer producing an output vector of size 32, To produce the output of discriminator features, these four feature vectors are concatenated to produce final output feature vector of size 128. The noise level time is processed via sinusoidal time embedding similar to (Ho) al. 2020). We use dropout of 0.2. DMMD is trained for Niter = 250000 iterations with a batch size B = 64 with number N = 16 of noise levels per batch. We use a gradient penalty Av 1.0 and E regularisation strength AG, = 0.1. As evaluation metries, we use FID (Heusel et al., 2017 and
For the image experiments, we use CIFARIO (Krizhevsky & Hinton 2009)
Inception Score (Salimans et al. 2016) using the same evaluation regime as in (Ho et al. 2020). To select hyperparameters and track performance during training, we use FID evaluated on a subset of 1024 images from training set of CIFARIO.
For CIFARIO, we use random flip data augmentation.
In DMMD we have two sets of hyperparameters, one is used for training in Algorithm D and one used for sampling in Algorithm 2 During training, we fix the sampling parameters and always use these to select the best set of training time hyperparameters. We use n 0.1 gra- dient flow learning rate, T = 10 number of noise levels, N₂ = 200 number of noisy particles, N, 5 number of gradient flow steps per noise level, Eini 0.001 and tma 1 - 0.001. We use a batch of 400 clean particles during training. For hyperparameters, we do a grid search for Ay {0, 0.001, 0.01, 0.1, 1.0, 10.0). for Xe, {0, 0.00, 0.01, 0.1, 1.0, 10.0), for dropout rate {0, 0.1,0.2, 0.3). for batch size {16, 32, 64). To train the model, we use the same optimization procedure as in (Ho al. 2020), notably Adam (Kingma & Bal 2013) optimizer with a learning rate 0.0001. We also swept over the the dimensionality of the output layer 32, 64, 128, such that each of four feature vectors got the equal dimension. Moreover, we swept over the number of channels for U-Net (32, 64, 128} (the original one was 32) and we found that 128 gave us the best empirical results.
After having selected the training-time hyperparameters and having trained the model, we run sweep for the sampling time hyperparameters over 7) E {1, 0.5, 0.1, 0.04, 0.013, E {1, 5, 10, 50). N₂ E 51,5, 10, 50}, mmin E {0.001, 0.01, 0.1, 0.2). fmax {0.9. 1 - 0.001 J. We found that the best
hyperparameters for DMMD were 7) = 0.1, N₂ = 10, T 10, 4min 0.1 and mmax = 0.9. On top of that, we a variant for DMMD with T 50 and N₂- = 5.
For a- D-DMMD method, we used the same pretrained discriminator as for DMMD but we did an additional sweep over sampling time hyperparameters, because in principle these could be different. We found that the best hyperparameters for a-DMMD are 1) = 0.04, mmin = 0.2, faux = 0.9, T = 5.
N, 10.
23

---

Published as a conference paper at ICLR 2025
For the denoising step, see Table for DMMD-e, used steps of DMMD gradient flow with higher learning rate 17 = 0.5 with mmax = 0.1 and Emin = 0.001. For a-DMMD-e, we used 2 steps of DMMD gradient flow with higher learning rate of 17" 0.5 with max 0.2 and min 0.001. For a- a-DMMD-e, we used 2 steps of DMMD gradient flow with a higher learning rate of 1y¹ = 0.1 with
mmax 0.2 and miin 0.001. The only parameter we swept over in this experiment this higher learning rate '7"
After having found the best hyperparameters for sampling, we run the evaluation compute FID on the whole CIFARIO dataset using the same regime as described in (Ho et al. 2020).
For MMD-GAN experiment, we use the same discriminator as for DMMD but top of that train& a generator using the same U-net architecture as for DMMD with an exception that we do not use the 4 levels of features. We use a higher gradient penalty of Xy = 10 and the same lz penalty Xe 0.1. We use batch size of B = 64 and the same learning rate as for DMMD. We use dropout of 0.2. We train MMD-GAN for 250000 iterations. For each generator update, we do discriminator updates, following (Brock et al. 2019).
For MMD-GAN-FFON experiment, we take the pretrained discriminator from MMD-GAN and run a gradient flow of type 4 it, starting from a random noise sampled from a Gaussian. We swept over different parameters such as learning rate 1) and number of iterations Niar- We found that none of our parameters led to any reasonable performance. The results in Table Jare reported using 1 0.1 and Niter 100.
1.1 ADDITIONAL DATASETS
1.1.1 RESULTS ON CELEB-A,
MNIST
We study performance of DMMD on additional datasets, MNIST (Lecun et al. 1998), on CELEB-A (64x64 (Liu et al. 2015). LSUN-Churde (64x64) (Yu et al. 2015). For MNIST and CELEB-A we use the same training/test split as well as the evaluation protocol as in (Fanceechhi et al. 2023). For LSUN-Church, For LSUN Church, et al. 2020 For MNIST, we used the same hyperparameters during training and sampling as for CIFAR-10 with NFE=100, see Appendix |1 For CELEB-A and LSUN, we ran sweep over dr₃ E {0, 0.001,0.00000000.000 0.01, 0.0, 10.0} and found that &g 0.001 led to the best results. For sampling, we used the same parameters as for CIFAR-10 with NFE=100. We use the same architecture for DMMD as for CIFAR- 10 experiments. The reported results in Table are given with NFE=100.
compute FID on 50000 samples similar to DDPM (Ho
The results are provided in Table 4 In addition to DMMD, we report the performance of the Discriminator flow baseline from (Franceschi et al. (2033] with numbers taken from the corresponding paper. We see that DMMD performance is significantly better compared to the discriminator flow, which is consistent with our findings on 100. The corresponding samples are provided in Appendix L.2
Table 4: Unconditional image generation on additional datasets. The metric used is FID. The number of gradient flow steps for DMMD is 100 and NFE for DDPM 100.
Dataset
MMD--AN
DDPM
DMMD
Disc. flow (Franceschi 2023
CELEB-A
12.1
6.72
8.3
41.0
SSUN-Church
3.84
MNIST
7.0
1.94
3.0
4.0
I.1.2 RESULTS ON CELEB-A-HQ
We study performance of DMMD on a high-resolution dataset CELEB-A-HQ (128x128) which contains 30000 samples. During training, we track performance on a subset of 2048 samples. We use the same backbone architecture as in DDPM paper (Ho et al. 2020) applied to CELEB-A-HO (256x256) dataset. For DMMD. we ran a sweep over de, E (0, 0.001, 0.01, 0.1, 1.0, 10.0) and found that & = 0.001 led to the best results. At sampling time, we use T = 10 noise levels with N, = 10 and 17 = 0.001 gradient flow learning rate and with N₂- = 50 noisy particles (see Algorithm 2 for
24

---

Published as a conference paper at ICLR 2025
more details). We the same architecture for DMMD for CIFAR-10 experiments. For both methods, NFE=100.
The results are reported in Table In addition to DMMD, also train DDPM (Ho et al. 2020 and report the corresponding performance. The results suggest that DMMD still has a performance gap compared to DDPM observed for lower resolution datasets. The corresponding samples are provided in Appendix
Table 5: Unconditional image generation on CELEB-A-HQ (128x128). The metric used is FID. The number of gradient flow steps for DMMD is 100 and we use NFE=100 for DDPM.
DDPM
DMMD
11.67
14.09
1.1.3 RESULTS ON LATENT-SPACE CELEB-A-HE
For memory efficient sample generation in higher dimensions, study performance of DMMD in the latent space of a VQ-VAE (van den Oord et al., (017)) trained on CELEB-A-HQ (256x256). We follow the approach of the latent diffusion paper (Rombach al. 2022), and train a VQ-VAE of the corresponding LDM-4 architecture which produces a latent space of the shape (64 x 64 x 3). Our baseline for comparison a latent diffusion model with the LDM-4 architecture. We trained DMMD on the same latent space, where we swept over A₃ E {0.05, 0.1, 0.2), and found that &g = 0.2 led to the best results. We used a gradient penalty & 1, and the same architecture for DMMD as for CIFAR- 10 experiments. At sampling time, we used T 10 noise levels with N, = 10 steps per noise level, a gradient flow learning rate 7) = 0.1, and N, = 200 noisy particles (see Algorithm for details). The results are given in 6 with NFE=100 for both methods. The corresponding samples are provided in Appendix
Table Unconditional image generation on latent space (64x64x3) CELEB-A-HQ (256x256).
The metric used is FID. The number of gradient flow steps for DMMD is 100 and NFE=100 for DDPM.
Latent Diffusion
Latent DMMD)
5.6
8.9
I.2 D-KALE-FLOW DETAILS
We study performance of D-KALE-flow on CIFARIO. We use the same architectural setting as in DMMD with the only difference of adding an additional mapping a(t: bb) from noise level to D dimensional feature vector, which is represented by 2 layer MLP with hidden dimensionality of 64 and GELU activation function. We use batch size B 256, dropout
rate equal to 0.3. For sampling time parameters during training, we use 1 0.5, total num- ber of noise levels T = 20, and number of steps per noise level N, = 5. At training, we sweep over RKHS regularization À {0, 1, 10, 100, 500, 1000, 2000}, gradient penalty x {0, 0.1, 1.0, 10.0, 50.0, 100.0, 250.0, 500.0, 1000.0). 12 penalty in {0, 0.1, 0.01, 0.001
I.3 NUMBER OF PARTICLES ABLATION
Number of particles. In Table 7 report performance of DMMD depending the number of particles Np at sampling time. As expected as the number of particles increases, the FID score decreases, but the overall performance is sensitive to the number of particles. This motivates the approximate sampling procedure from Section 5
J PERFORMANCE VS. NUMBER OF GRADIENT FLOW STEPS TRADE-OFF
Here, we provide a table showing the dependence of the performance of DMMD on number of total DMMD gradient flow steps, which we call NFE. The NFE is the total number of gradient flow
25

---

Published as a conference paper at ICLR 2025
Table 7: Number of particles ablation, FIDs on CIFARIO.
N, 50
N, = 100
N, 200
9.76
8.55
8.31
iterations, which equals to N,T, where N₂ is the number of steps per noise level and T is the number of noise levels. By default, we use the gradient flow learning rate rj = 0.1, see (9. We also found that as we increase the number of total gradient flow steps, it was sometimes beneficial to use smaller learning rate, 7j = 0.05. Results are given in Table 8 We see that as we increase NEE, the FID improves up to a point (NFE 250). After NFE=250, we do not see further improvement. Moreover, as we noticed in our experiments, increasing the total compute at sampling time might require readjusting the gradient flow learning rate.
Table 8: Dependence of the FID on CIFAR-10 on the total number of gradient flow steps (NFE). is the gradient flow learning rate, see
Total number of steps (NFE)
FID
10(1 0.1)
377.5
50(n = 0.1)
36.4
100(1) 0.1)
8.5
250(n 0.1)
12.1
250(n 0.05)
7.74
500(1) 0.05)
8.6
1000(n 0.05)
9.1
J.1 RESULTS WITH F-DIVERGENCE
We study performance of D-KALE-FUE described in Section B and Appendix F in the setting of unconditional image generation for CIFAR-10. We compare against a GAN baseline which uses the KALE divergence in the discriminator, but has adversarially trained generator. More details described in Appendix and Appendix [.2] The results are given in Table We see that unlike with DMMD, D-KALE-FlOw achieves worse performance than corresponding KALE-GAN indicating that the inductive bias provided by the generator may be more helpful in this case this is a topic for future study.
Table Unconditional image generation on CIFAR-10 with KALE-divergence, The number of gradient flow steps is 100.
Method
FID
Inception score
D-KALE-Flow KALE-GAN
15.8
8.5 8.7
12.7
J.2 COMPUTE RESOURCES FOR IMAGE EXPERIMENTS
For all the experiments, we used A100 GPUs with 40 GB of memory. To train DMMD for 250k: steps, we needed run training for around 24 hours. The total hyperparameter sweep for DMMD required 36 runs to figure out regularization constants, 12 runs figure out batch size and dropout rate and then runs to figure out the dimensionality of the U-Net and the same 3 runs where the features of the U-Net were coming only from the last layer. This required 54 runs in total.
Running inference on small subset of CIFAR- 10 required around minutes of GPU time, and we ran full grid search to select best sampling time parameters, which is around 1080 values. We did this sweep for DMMD and a DMMD. For DMMD e, we additionally swept over higher learning rate at the second stage which required more runs. For DMMD and DMMD a,
26

---

Published as a conference paper at ICLR 2025
swept over learning rates at second stage which required 10 more runs. After having found the best parameters, we run sampling with the best parameters on full CIFAR-10 which takes about hour for NFE 100.
For additional datasets, for MNIST we used the same best parameters as for CIFAR-10, which required one run only since we ran an additional sweep over regularization strength which required 6 training runs per dataset and additional runs for sampling the whole datasets.
very good performance out of the box. For CELEB-A and LSUN,
For MMD - GAN, the training runs were faster, by around 2-x factor. We did a grid search over the regularization strengths which took 36 training runs and 12 runs figure out batch size and drop-out rate.
For DKALE-flow, the experiment was as fast as MMD GAN and we ran grid search with 67 runs for regularization and 4 runs for dropout. The same was done for DKALE - GAN.
K OPTIMAL KERNEL WITH GAUSSIAN MMD
In this section, we prove the results of Section 3 We consider the following unnormalized Gaussian kernel
ka v) = a cxxp[|1== - y1l2/(2a²))
For any E Fd and a > 0 we denote TH.o the Gaussian distribution with mean μ and covariance matrix o² Idd. We denote MMD the MMD2 associated with k.. More precisely for any 14, A22 Rd
and 01, 02 >> O we have MMD (T p2 =Emgio 0am [Ro (X,
2E 85 (X, Y)] +
(23)
E Emp 8 (aa (Y,Y)))
In this section we prove the following result. Proposition K.1. For any po IR and > O, let a" be given by
a= = argmax a20 |V no MMD2 (Fu,or pl..
Then, we have that
202)1/2
(24)
let us provide some insights on the result. The quantity |[ MMD3 (T al) represents how much the mean of the Gaussian Tyopo displaced when considering a flow on the mean of the Gaussian w.r.t. MMD?. Intuitively, we aim for
Before proving Proposition
"V,,,MD' (To al to be as large as possible as this represents the maximum displacement possible. Hence, this justifies our goal of maximizing V V,, MMD" (To Il with respect to the width parameter
We show that the optimal width a* has a closed form given by 24 is notable that, assuming that 0 is fixed, this quantity depends [| i.e. how far the modes of the two distributions are. This observation justifies our approach of following an adaptive MMD flow at inference time. Finally, we observe that there exists a threshold, i.e. || / //d 2) 2o for which lower values of [| still yield a* = 0. This phase transition behavior is also observed in our experiments.
We define ((a, o. po, #1) for any a, o > 0 and no: #1 Rd given by D(a, a. vo p1) Spexxd ki ka, po ( (z))dTm)
Proposition K.2. For any a, > 0 and μoo p1 E Rdd have
D(a,o, Ho. = [a² (022="(1/22 1/ 4/2 expll (ool 11 /22x)) + 1/0.11/(221)
(fo. p2) a² || ( - -1mm11 /200)
with
A (a² μt 1² (2240)/(( (n²² + a²),
Au = (a? po + R²2 h1) (1² + a²))
where K = (1/02 + 1/02)-1/22
27

---

Published as a conference paper at ICLR 2025
Proof. In what follows, start by computing D(a, o, no: (1) for any a, 0 and Mo. R given by
D(a, o, No, (1)
Sanxad ka (2, yod (x)d
=1/(270aa)) SRaxRa expl=llxpl u / po
=1/(270 a) JEAXE expl // /22 -1|z
-00ll1 (200) - lly paall/(20""]ddddy
In what follows, we denote K = (1/0² + 1/12)/ We have
|lul)//2xx)) + (2,, u) / + (a,, (o),wa)/² + (v, mi) mi) with C(po,ui) ui) 1/(2no02a))) expl- /20²) - || (2 In what follows, we denote
D(a, o, No, (1) C(uo, M1) fRA expl //rpp||/lll//(/////1)/
P(a,y) v) the second-order polynomial given by
(z, P(r,y) = 1==11 (lul) - (x, w3/a2 -x, No) - m1)/
Note that we have
D(a, o, no, M1) - C(po, H1) J expl- (x,, v)ldrdy.
(25)
Next, for any fo. Pu E RRd, we consider Q(2, 3) given by
Q(2, v) || - ioo13 /2xx) Ily in (2x²) (x - Po,y m2)
|1x]|//(2x²) + | /2xx)) + |/lll/ -111111/ (1. /22 -200/12211100 -ly. -22-2z
(x -po. A22/2²
= ||=||2/(2x²) + [1/aol + |lvll - -11x111/(2x3) - -(r, -wy.i2)/12
(x, v) /02 (jo, in) (v, /o)/ /2² (er,i)//0²
P(z, v) + |0oll2((xx)) + +1A111/((2x)) (z, /O) (u,Pe)/n² + ( (o)/ + 13.110/02
(jo: im) / + (v, (o) + (z., /2)/a²
P(z,y) + 10 12 12xx) + || / - /2)/0a²
+ (x. po/ 00 2 po/k² + /a) (y-12/00² - jin/1² 10/002)
In what follows, we set Po. Ar such that
11/0² = f/2/n² m
Ha/0² = .u/// i22022
We get that
Ar (mI/(02 + Ho/ a (1/ - 1/a²).
A0 (No/(0² + 11/
- -1/a")
We have that
01/1/x-- 1/at) = 0210/ 22((0'0²)) = 1/0" +2/0 = 1/16² + 1/a2.
(26)
Therefore, we get that
11 (HI/K² + po/ 2o/oo)//1/////2 1L/ +1/a2),
A (HO/122 + 11/02)//11// +1/02).
Finally, we get that
Å (a²p) x² ho) = a²),
Au (a² po N² un ( (122 +a?)
With this choice, we get that
P(z.v) Q(2 v) |1 ///0ol1/(22)) 122x2) 11 1(2x3) + (Ao, R1)/a²
(27)
We also have that for any x, E R.d
x (=-An) Pi)
T
(
Id/n²
-Id/a²)
x Ao
Q(r,))= = (1/2)
y f11)
Id/a²
Id/k²
28

---

Published as a conference paper at ICLR 2025
Using this result we have that
SRaxRa exp[-Q(x, y) (2TT) det((=-1)-))
(28)
with
Dー (
Id/k² -Id/a²
-1d/a² Td/n²
Using 26 we get that
det(E = [11/0 (1/m² + +1/02)1
Combining this result and 28 we get that
fge xxp-Q((x:) y)] = ((10) +1/3///)31//2)-13 -
Combining this result, and (25 we get that
D(a, o, Po, p1) C(po, M1) exp 2oll Therefore, we get that
Ho. D(a, No. M1) = = [0202(1/k²² (1/=0( -1/02)) /22)) |1 11/211//(22²))
(Ao, (o)//1)/a22 || /2202)
We investigate special cases of Proposition K.2
First, we show that if (Au = Al then D(a, a, Po: (b) does not depend on Po-
Proposition K.3. For any >> O and fi I& we have D(a, o, Po, Mu) = (a? + 200)
Proof. We have that fūo = فu = μ1 = μ4 in Proposition K.2 In addition, we have that
(1/2x²) + (1/2x²) 1/a² 1/(20²) 11/(22)=
Therefore, we have that
cxplll@oll//(xx²) + || / (2x (Po, pn) (11oo1/((207))
=
which concludes the proof upon using that 1/x² 1/a2 1/02.
Proposition K.3 might seem surprising at first but in fact it simply highlights the fact that when trying to differentiate a Gaussian measure with itself, the result is independent of the location the Gaussian and only depends on its scale. Then, study the case where μ4 0.
Proposition K.4. For any > O and fr E IRd we have
D(a, M4, 0) = (a² + 20220=2000 d/2
+ 20
Proof. First, we have that
Noo a² a2/ 222 +22)
À m + t2
Therefore, get that
D(a, No, [02(1/n² +1/3)) exp (xp(())) [|0 -2/0
Using 26 we get that
a²/1²² 2²= a" N²)/0²
Therefore, get that
(at/k² -3) - -
which concludes the proof.
Using Proposition
K.3 Proposition
K.4 and definition
23
have the following result.
29

---

Published as a conference paper at ICLR 2025
Eioksor of to
200.00
Figure 3: Evolution of the norm of the mean br of the Gaussian distribution " according to a gradient flow the mean le w.r.t. MMD In the adaptive case a is given by Proposition 3.1 while in the non adaptive case, 04 = 00 = 1. In our experiment we consider d = 1 and a = for illustration purposes.
Proposition K.5. For any a, a > and μo E IRd we have
MMD Too.o) = 2(a²² + 20²2) 2/2(1 expl-lluol3 + (22((=++2)2]
In addition, we have
(0/2)-d/2-20 -d/2-1 V Pg MMD Toooa) 22(a2 20 -
expl-11mol1
+ +20
Finally, we have the following proposition. Proposition K.6. For any po E TRd and a > 0 let a be given by
xx = argmax 2200 VV MMD (n (Fo.a: pp..
Then, we have that
a = ReLUU[] /(d+ - 2//2
Proof. Let a > and μo Rdd First, using Proposition
K.3 we have that for
V MMD
(TT, o a)) = 4a 24 || (a² 20²² d-2 expl I1loll //lll//((/2/// + 20²)
Next, we study the function f: [0, to]
IR given for any E [0, tol by
f(tt) = 2ーه expl-t|||aol|].
with to 1/(20²). We have that
f(t)= = gd+ exp 10010|10001000 [ 2) 11mol1
We then consider two cases. First, if to S (d + 2) 01 i.c. o² S || 12 (2 2)), then f is
increasing on [0. to] and we have that f is maximum if = tg. Hence, if o² < |oo| //(2 2)), we have that 0. Second, if to S (d 2)/ i.e. 0² s25 (22 2)) then is increasing on [0, ("), non-increasing on [t [o] with L" (d + 2)/ 2)/l|//| and we have that f is maximum if == 1" Hence, if x² 2 || (2/d + 2)). we have that a = 21)/ol|/// concludes the proof.
//(l/////dd/// (d= 2) 202)1/2, which
K.1 PHASE TRANSITION BEHAVIOUR
L IMAGE GENERATION SAMPLES
L.1 CIFARIO SAMPLES
Samples from DMMD with NFF=100 and NFE=250 are given in Figure 4 Samples from DMMD with NFF=100 and from a-DMMD with NFE=50 are given in Figure
L.2 ADDITIONAL DATASETS SAMPLES
Samples for MNIST are given in Figure 6 for CELEB-A (64x64) are given in Figure 7 and for LSUN Church (64x64) are given in Figure
30

---

Published as a conference paper at ICLR 2025
Figure 4: CIFAR-10 samples from DMMD with NFE=250 on the left and with NFE=100 on the right
Figure 5: CIFAR-I0 samples from DMMD with NFE=100 on the left and samples from the a- DMMD-e with NFE=50 the right
31

---

Published as a conference paper at ICLR 2025
55
5
6
6
5
9
3
9
7
777800
5
92794
9
099560
5426 4
Figure DMMD samples for MNIST.
m
Figure 7: DMMD samples for CELEB-A (64x64).
32

---

Published as a conference paper at ICLR 2025
00
Figure DMMD samples for LSUN Church (64x64).
Figure 9: CELEB-A-HQ (128x128) samples. Left, the samples from DMMD, Right, the samples from DDPM.
33

---

Published as a conference paper at ICLR 2025
Figure 10: Latent (64x64x3) CELEB-A-HQ (256x256x3) samples for latent DMMD.
L.3 ADDITIONAL DATASETS SAMPLES ON CELEB-A-HQ
The samples for CELEB-A-HQ (128x128) 128) are given in Figure
L.4 ADDITIONAL DATASETS SAMPLES ON LATENT (64x64x3) CELEB-A-HO (256x256x3)
The samples for latent DMMD are given in Figure 10 and the samples for latent diffusion are given in Figure 11
34

---

Published as a conference paper at ICLR 2025
Figure 11: Latent (64x64x3) CELEB-A-HO (256x256x3) samples for latent diffusion.
35


Reviews:
[{"review_id": "Sdwc6b5SJK", "reviewer": "Reviewer_QpwN", "summary": "The paper presents a novel approach, \"Diffusion Maximum Mean Discrepancy Gradient Flow\" (DMMD), to improve generative modeling by combining Maximum Mean Discrepancy (MMD) with a noise-adaptive gradient flow mechanism. Unlike GANs, DMMD eliminates adversarial training by utilizing a noise-conditional MMD discriminator. DMMD introduces a particle transport technique, adapting MMD as the divergence metric to transport particles from a source distribution to a target distribution.", "questions": "1. As mentioned above, what challenges do you consider in scaling DMMD to larger datasets like ImageNet?\n2. What factors limit DMMD's FID performance relative to state-of-the-art diffusion models?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper is clearly written and well-organized, tackling a genuine problem and effectively presenting its contributions and findings.\n\n2. The paper establishes a solid mathematical foundation, rigorously linking each section to prior research.\n\n3. The results presented are sufficient to validate the theoretical findings and showcase the effectiveness of the proposed approach.", "weaknesses": "1. While the paper shows promising results, it is still outperformed by standard diffusion models, especially in terms of FID scores. Further work might be necessary to reach SOTA performance on larger datasets like ImageNet.\n\n2. Related to the previous point, the experimental results are primarily limited to smaller datasets (CIFAR10, MNIST, CELEB-A, and LSUN Church), which may not reflect the potential scalability of DMMD to more complex, high-resolution datasets.\n\n3. Although the method avoids adversarial training, the noise-adaptive MMD flow still introduce complexity, which may limit reproducibility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "78m3sT6BwX", "reviewer": "Reviewer_KT31", "summary": "The paper integrates MMD GAN with the diffusion forward process, introducing a novel generative model framework called DMMD. Through the design of a noise-adaptive MMD gradient flow, this framework aims to reduce the challenges of adversarial training and address the singularity issues found in score-based methods.", "questions": "Q1 - The need for dataset features during sampling seems counterintuitive. Do the authors have any insights into potential solutions for addressing this limitation?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "S1 - The paper provides a well-formulated background and problem statement, with a theoretically motivated and well-grounded DMMD framework.\n\nS2 - The idea of using an adversarial training-free discriminator based on the diffusion forward process could offer valuable insights to the community.", "weaknesses": "W1 - The framework's absolute performance is a concern, as DMMD shows a significant performance gap compared to DDPM and more modern methods on the selected image generation benchmarks.\n\nW2 - Its broader application potential is limited, with empirical evaluation restricted to small datasets like MNIST and CIFAR.\n\nW3 - The sampling method appears restrictive, requiring reference features from the ground truth dataset to formulate the witness function.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5y9xpMEVaN", "reviewer": "Reviewer_xQpY", "summary": "The article proposes a gradient flow procedure for generative modeling by transporting particles from an initial source distribution to a target distribution, where the gradient field on the particles is given by a noise-adaptive Maximum Mean Discrepancy (MMD) divergence. The noise adaptive MMD is trained on a data distribution corrupted by forward diffusion process. The divergence training procedure is related to discriminator training in Generative Adversarial Networks (GAN), but does not require adversarial training. The article demonstrates competitive empirical performance of the method in unconditional image generation on CIFAR10 dataset.", "questions": "1. I recommend that the author conduct additional experiments using a variety of datasets to demonstrate the effectiveness of the DMMD. The results on CIFAR10 are not as impressive as those achieved by state-of-the-art generative models. I also suggest that the author include more comparative experiments with additional relevant works.\n\n2. I am wondering if using vector Z in Eq.(15) to train a generator will result in better outcomes than using it as a sampling process.\n\n3. I am curious about the time required to train the MMD Gradient Flow on the CIFAR-10 dataset. In my experience, training a diffusion-based discriminator takes significantly longer than training the original models. I find that incorporating a diffusion process in gradient flow-based models can reduce the number of accumulated gradient steps. I would like to know if the sampling steps of the DMMD model are influenced by the diffusion process.\n\n4. I wonder how the sampling time changes with increasing data dimensions, such as in 256x256 high-dimensional datasets.\n\n5. Can the MMD Gradient Flow be applied to other tasks such as Super-Resolution or Inpainting? These tasks also involve particle transportation from an initial source to a target distribution.\n\n6. There are writing errors in the article, particularly in Table 4 and Algorithm 2.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The article explores the integration of a diffusion process into models based on Maximum Mean Discrepancy (MMD) gradient flow, offering a theoretical foundation for both the discriminator and the sampling process used in MMD-GAN with this diffusion approach. Additionally, it employs a linear kernel for the scalable MMD-GAN to reduce computational complexity. The authors also conduct experiments using other forms of KL divergence, such as KALE divergence, to demonstrate its effectiveness.", "weaknesses": "1. I believe the contribution of this article is inadequate. Previous research has utilized the diffusion process in the discriminator, as noted in this work [1]. However, this article does not provide theoretical proof demonstrating that the MMD GAN can converge to more optimal points when using the diffusion process.\n\n2. Additionally, the effectiveness of MMD Gradient Flow has only been tested on low-resolution datasets, which does not provide sufficient evidence to confirm its overall efficacy. I recommend that the author conduct additional experiments using high-resolution datasets at a resolution of 256x256, specifically on the LSUN and CelebA datasets. These experiments should include evaluations based on the number of metric sampling steps (NFE) and diversity (FID).\n\n[1] Xiao Z, Kreis K, Vahdat A. Tackling the generative learning trilemma with denoising diffusion", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HBSwl5xrwx", "reviewer": "Reviewer_vKuM", "summary": "This paper proposes a generative model based on the Wasserstein gradient flow of the Maximum Mean Discrepancy (MMD), called DMMD. Given the forward process in diffusion models, the proposed method learns a discriminator between clean data and noisy data. The Wasserstein gradient flow is represented via this discriminator. Therefore, DMMD generates samples by simulating particle trajectories through this flow. DMMD achieves superior performance compared to other discriminator flow and MMD flow baselines in image generation across various datasets.", "questions": "- In Table 1, do all the other baselines use the same backbone network?\n- Could you provide the wall-clock time comparison to demonstrate the computational benefits of approximate sampling in Table 2?\n- Is this method applicable to higher-dimensional datasets, such as CelebA-HQ (256x256)?\n- Could you provide the potential advantages of this approach compared to other non-adversarial dynamic generative models, such as Flow Matching [1]?\n\n[1] Lipman, Yaron, et al. \"Flow matching for generative modeling.\" ICLR 2023.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- This paper is easy to follow.\n- This paper suggests an efficient approximate sampling procedure for the linear kernel.\n- DMMD is evaluated on various image datasets, such as CIFAR-10, MNIST, CelebA (64x64), and LSUN-Church (64x64).", "weaknesses": "- My main concern is whether the trajectories of the probability distributions for the forward process $\\\\{ p^{1}\\_{t} \\\\}\\_{t \\geq 0}$ in the diffusion model and the MMD gradient flow $\\\\{ p^{2}\\_{t} \\\\}\\_{t \\geq 0}$ coincide. If these trajectories are different, DMMD learns the Wasserstein gradient flow for minimizing $MMD (p^{1}\\_{t}, p_{data})$. However, during the generation process, the particles follow $\\\\{ p^{2}\\_{t} \\\\}$ at $ t - \\triangle t$. This gradient flow mismatch cannot guarantee that the particles correctly generate the target distribution. Could you clarify this concern?\n- The generation process is computationally expensive. How many $N_{s}$ steps are required for each time $t$ during sample generation (Eq 9)?\n- For a general kernel, the sampling from DMMD requires access to the training data (Eq 10). Only for linear kernel, this issue can be avoided by saving the average features for each time $t$ (Eq. 12).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)