Interaction-Force Transport Gradient Flows
Authors: Egor Gladin
Pavel Dvurechensky
Alexander Mielke
Jia-Jie Zhu
Venue: NeurIPS 2024 poster (2024)
Keywords: kernel methods
gradient flow
optimal transport
Wasserstein
Fisher-Rao
Hellinger
unbalanced optimal transport
partial differential equation
optimization
calculus of variations
variational inference
MCMC
MMD
Primary area: probabilistic_methods
TL;DR: We propose the spherical interaction-force transport (IFT) gradient flows with the global exponential convergence analysis for both the MMD and the KL-energy gradient flows. Numerical experiments show stable inference performance.

Abstract: This paper presents a new gradient flow dissipation geometry over non-negative and probability measures.
This is motivated by a principled construction that combines the unbalanced optimal transport and interaction forces modeled by reproducing kernels. Using a precise connection between the Hellinger geometry and the maximum mean discrepancy (MMD), we propose the interaction-force transport (IFT) gradient flows and its spherical variant via an infimal convolution of the Wasserstein and spherical MMD tensors. We then develop a particle-based optimization algorithm based on the JKO-splitting scheme of the mass-preserving spherical IFT gradient flows. Finally, we provide both theoretical global exponential convergence guarantees and improved empirical simulation results for applying the IFT gradient flows to the sampling task of MMD-minimization. Furthermore, we prove that the spherical IFT  gradient flow enjoys the best of both worlds by providing the global exponential convergence guarantee for both the MMD and KL energy.

Interaction-Force Transport Gradient Flows
Egor Gladin
Pavel Dvurechensky Weierstrass Institute for Applied Analysis and Stochastics Berlin, Germany pavel. wwwel..dwreechenstyywii--s-tse
Humboldt University of Berlin
Berlin, Germany
& HSE University
egorgl egorgladiinyyandex..o
Alexander Mielke
Jia-Jie Zhu
Humboldt University of Berlin
Weierstrass Institute for Applied Analysis and Stochastics
& WIAS
Berlin, Germany
Berlin, Germany
al exander www.eeee..eeeeeeeeeeeeeeeeeeeeee miw.eeeeereiiaeerrrrrr.i.. de
jia-jie zhu@wias-berling
Abstract
This paper presents a new gradient flow dissipation geometry over non-negative and probability measures. This is motivated by principled construction that combines the unbalanced optimal transport and interaction forces modeled by reproducing kernels. Using precise connection between the Hellinger geometry and the maximum mean discrepancy (MMD), we propose the interaction-fon transport (IFT) gradient flows and its spherical variant via infimal convolution of the Wasserstein and spherical MMD tensors. We then develop a particle- basee optimization algorithm based the JKO-splitting scheme of the paass-prssrrving spherical IFT gradient flows. Finally, we provide both theoretical global exponential convergence guarantees and improved empirical simulation results for applying the IFT gradient flows to the sampling task of MiD-minimization. Furthermore, we prove that the spherical IFT gradient flow enjoys the best of both worlds by providing the global exponential convergence guarantee for both the MMD and KL energy.
Introduction
Optimal transport (OT) distances between probability measures, including the earth mover's distance
[Werman et al. 1985 Rubner et al., 2000) and Kantorvvato or Wasserstein distance (Villani
2008], are one the cornerstones of modern machine learning as they allow performing a variety
of machine learning tasks, e.g., unsupervised learning [Arjovsky et al., 2017
Bigot et al. 2017
semi-supervised learning [Solomon et al.) 2014], clustering Ho et al. 2017 text classification
Kusner al.) 2015), image retrieval, clustering and classification [Rubner et al. 2000 Cuturi 2013]|Sandler and Lindenbaum 2011], and distributionally robust optimization (Sinha et al. 2020
Mohajerin Esfahani and Kuhn 2018 Many recent works in machine learning adopted the techniques
from PDE gradient flows over optimal transport geometries and interacting particle systems for inference and sampling tasks. Those tools not only add new interpretations to the existing algorithms, but also provide a new perspective on designing new algorithms.
"Corresponding author: Jia-Jie Zhu
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

For example, the classical Bayesian inference framework minimizes the Kulback- divergence towards a target distribution π. From the optimization perspective, this can be viewed as solving
min #EACP
{F(e) == (x((T)))
(1)
where A is subset of the space of probability measures P. e.g., the Gaussian family. The Wasser- stein gradient flow of the KL gives the Fokker- Planck equation, which can be simulated using the Langevin SDE for MCMC. Beyond the KI., many researchers following Arbel et al. 2019) advocated using the squared MMD instead the driving ergy for the Wasserstein gradient flows for sam- pling. However, in contrast to the KL setting, there is little sound convergence analysis for the MMD- minimization scheme like the celebrated Bakry- Emery Theorem. Furthermore, it was shown, e.g-, in [Korba et al., 202T]. that Arbel et al.] (2019]ss ss algorithm suffers a few practical drawbacks. For example, their particles tend to collapse around the mode or get stuck local minima, and the algo- rithm requires a heuristic noise injection strategy that is tuned over the iterations; see Figure I aand $4 for illustrations. Subsequently, many such as Car-
Figure 1: (Left) Wasserstein flow of the MMD energy [Arbel et al. (009). Some particles get stuck at points away from the target. (Right) IFT gradient flow (this paper) of the MMD en- ergy. Particle mass teleported to close to the target, avoiding local minima. Hollow circles indicate particles with zero mass. The red dots are the initial particles, and the green dots are the target distribution. See $4 for more details.
rillo et al. [
Chewi et al. [2020] Korba al.
[2021], Glaser et al. 2021] Craig et al. 2023 Her-
trich et al. 2023]. J. Neumayer et al. 2024] proposed
modified energies to be used in the Wasserstein gra- dient flows. In contrast, this paper does not propose
new energy objectives. Instead, we propose a new gradient flow geometry the IFT gradient flows. To summarize, our main contributions are:
1.
We propose the interaction-foree transport (IFT)) gradient flow geometry over non-negative measures and spherical IFT over probability measures, constructed from the first principles of the reaction-diffusiee type equations, previously studied in the context of the Hellinger- Kantorovich (Wasserstein-Fiist
groups including Chizat et al. [2018] Rao) distance and gradient flows. It was first studied by three [iero et al. [2018). Kondratyev et al. (2016), Callouut and Monsaingeon| 2017]. Our IFT gradient flow is based on the inf-convolution of the Wasserstein and the newly constructed spherical MMD Riemannian metric tensors. This new unbalanced gradient flow geometry allows teleporting particle mass in addition to transportation, which avoids the flow getting stuck at local minima; see Figure T for illustration.
2. We provide theoretical analysis such as the global exponential decay of energy functionals via the type functional inequalities. As an application, we provide the first global exponential convergence analysis of IFT for both the MMD and KL energy functionals. That is, the IFT gradient flow enjoys the best of both worlds.
3. We provide a new algorithm for the implementation of the IFT gradient flow. We then empirically demonstrate the use of the IFT gradient flow for the MMD inference task. Compared to the algorithm of Arbel et al. [2019] IFT flow does not suffer issues such as the collapsing-to-nee mo-e issue. Leveraging the first-principled spherical IFT gradient flow, our method does not require a heuristic noise injection that is commonly tuned over the iterations in practice; see (Korba al., 2021] for discussion. Our method can also be viewed addressing
original MMDD-enrgyy-eee
stag-itanding issue of the kernel- embedding methods (Smola et al.) 2007 Muandet et al.
2017 Lacoste-Julien et al.
2015) for optimizing the support of distributions.
Notation We use the notation P(X). M* (X) to denote the space of probability and non non-negative
measures on the closed, bounded, (convex) set X C Idd. The base space symbol X is often dropped if there is no ambiguity in the context. We note also that many of our results hold for X = Rd In this paper, the first variation of functional F at μ E M is defined as function 6F [A] such that d F(u + e v)l==o = S AF [a](x) do(x) for any valid perturbation in measure o such that 8y
de + E.V M when working with gradient flows over M and μ + --- Pover P. We often

---

omit the time index t to lessen the notational burden, e.g., the measure at time t, μ(t, 1), is written as μ.. The infimal convolution (inf-convolution) of two functions f.g on Banach spaces is defined as (z) = infy y {f(y) + g(x y. In formal calculation, we often use measures and their density interchangeably. i.e., J f l means the integral w.r.t. the measure μ.. For a rigorous generalization of flows over continuous measures to discrete measures, see [Ambrosio et al. 2005]. Vgk(.,)) denotes the gradient w.r.t the second argument of the kernel.
Background
2.1 Gradient flows of probability measures for learning and inference
Gradient flows are powerful tools originated from the field of PDE. The intuition can be easily seen from the perspective of optimization solving the variational problem
min #EAC
F(u)
using a "continuous-time version" of gradient descent, over suitable metric space and, in particular, Riemannian manifold. Since the seminal works Otto 1996) and colleagues, one can view many PDEs as gradient flows over the aforementioned Wasserstein metric space, canonically denoted as (P2(X), W2); see [Villani 2008 Santambrogio, 2015) for comprehensive introduction.
Different from a standard OT problem, a gradient flow solution traverses along the path of the fastest dissipation of the energy F allowed by the corresponding geometry. In this paper, only concerned with the geometries with (pseudo-)Riemannian structure, such as the Wasserstein, (spherical) Hellinger Fisher-Rao geometries. In such cases, formal Otto calculus can be developed to greatly simplify the calculations. For example, the Wasserstein Onsager operator (which the inverse of the Riemannian metric tensor) Kw T
TpMM & div(pVe), where Ty.M is the tangent plane of MM at p and T Mr the cotangent plane. Using this notation, Wasserstein gradient flow equation of some energy F can be written as
8F
div(u 8F 5yc
i = Kw (1)
8p
In essence, many machine learning applications are about making different choices of the energy F in (2). e.g.. the KI., x² diveergge,,,, or MMD. However, Wasserstein and its flow equation (2) are by no means the only meaningful geometry for gradient flows. One major development in the field is the Dellinger-Kantorovee
a.k.a. the Wasserstein-Fisher-Raa (WFR) gradient flow. The WFR gradient flow equation is given by the reaction-diffusion equation, for some scaling coefficients a, B > O.
5F
5F
ù = a div(uV
Bu
du
5u
A few recent works have applied WFR to sampling and inference Yan et al. by choosing the energy functional to be the KL divergence.
2024
Lu et al.
2019
2.2 Reproducing kernel Hilbert space and MMD
In this paper, we refer to a bi-variate function l: X X as a symmetric positive definite kernel if k is symmetric and, for all E N, ai...... .,, On E and Ii...., X, we have Ei Σj=1 agayk (æ,, xI) 2 0. k is reproducing kernel it satisfies the reproducing property, i.c. for all x X and all functions in a Hilbert space / H. we have f(x) = (f,, k(-,x))H Furthermore, an RKHS if the Dirac functional ó, : H + R, 8,( == f(x) is continuous. can be shown that there is a one-to-one correspondence between the RKHS H. and the reproducing kernel k. Suppose the kernel square-integrable |A<][ == k(a, x))(())
the space
oo w.r. P E P. The integral operator Tr.p : 1 > is defined by Th,pg(a) == f (a, x') g (a') dp (a') for g E T.Z. With an abuse of terminology, we refer to the following composition also as the integral operator
K, = IdoTR.p L²(g) L²)
Kpp is compact, positive, self-adjoint, and nuclear; ef. [Steinwart and Christmann 2008). simplify the notation, we simply write K. when p is the Lebesgue measure.
3

---

The kernel maximum mean discrepancy (MMD) [rrtton et al. 2012 emerged as easy-to- compute alternative to optimal transport for computing the distance between probability measures.
i..., MMDD(M, (H, == || (N- 2) 11 = / k(x, x²) d(u (z) d(p - (x)) where H. is the
RKHS associated with the -effnite) kernel k. While the MMD enjoys many favorable properties, such as a closed-form estimator and favorable statistical properties (Tolstikhin et al. [017 2016], its mathematical theory is less developed compared to the Wasserstein space especially the geodesic structure and gradient flow geometries. It has been shown by ZZuu and Mielke] (2024) that MMD is a (de-)kernelized Hellinger or Fisher-Rao distance by using dynamic formulation
MMD² (m, v) min (o 11 dt - Et, u(0) μ. u(1) - v, & &
(4)
Mathematically, we can obtain the MMD geodesic structure if we kernelize the Hellinger (Fisher-Rao) Riemannian metric tensor,
GMD K2 (H)) KMMD = KHHe (1) o K
(5)
noting that the Onsager operator K is the inverse of the Riemannian metric tensor K = G The MMD suffers from some shortcomings in practice, such as the vanishing gradients and kernel choices that require careful tuning; see e.g., (Feydy et al. [019). Furthermore, a theoretical downside of the MMD as a tool for optimizing distributions, and kernel-mean embedding (Smola et al., 2007] Muandet et al., 2017] in general, is that they do not allow transport dynamics. This limitation is manifested practice, e.g., is intractable to optimize the location of particle distributions; see e.g. Lacoste-Julien et al. 2015 In this paper, we address all those issues.
3 IFT gradient flows over non-negative and probability measures
In this section, we propose the IFT gradi- ent flows over non-negative and probability measures. Note that our methodology is fundamentally different from a few related works in kernel methods and gradient flows such as AArel al.
gradient
2019
Korba et al.
2021 Hertrich et al.
2023 Glaser et al.
2021 Neumayer al.
2024 in that are
not concerned with the Wasserstein flows of a different energy, but new gradient flow dissipation geometry.
transport
interact. force
3.1 (Spherical)
IFT gradient flow equations over non-negative and probability measures
Figure 2: Illustration the IFT gradient flow. Atoms are subject to both the transport (Kantorovich) potential and the interaction (repulsive) force from other atoms.
The construction of the Measerrstiin-issher Rao gradient flows crucially relies on the inf-convolution from convex analy- sis [Liero et al. 2018 Chizat 2022). There, the WFR metric tensor is defined using inf-convolution of the Wasserstein tensor and the Hellinger (Fisher-Rao) tensor GwR(p)) (w) By Legendre transform, its inverse, the Onsager operator, is given by the sum KWER (1) = Kw (R) + (H). Therefore, we construct the IFT gradient flow by replacing the Hellinger (Fisher-Rao) tensor with the MMD tensor, as in
(p) Gw (GMD)))))))) KIF (p) Kw (p) + KMMD (p).
(6)
The MMD gradient flow equation is derived by Zhu and Mielke [2024) using the Onsager operator (5)
j -MMMDD(H)
6F
[m]
-K-1
5F [n]]
8p
5μ

---

Hence, we obtained the desired IF'T gradient flow equation using
6
i =akw (1)) [a] (KMMD(() dy
5F
åF δμ [u]] =a div(u
oF dp
[w]] -
---
5F
ōu [e]]
(8)
Formally, the IFT gradient flow equation can also be viewed as kemel-approximmtion to the Wasserstein Sesserstein-Fihhr-RReee Fisher-Rao gradient flow equation, i.e., the -iiffuusnn equation (3).
Corollary 3.1. Suppose f ko (z.-) du and the kernel-weighted-maann measure Ki (z.-) du + c6 as the bandwidth a + Then, the IFT gradient flow equation + 0, i.e., the reaction-difussie equation
converges to the Dirac
lends towards the WFR gradient flow equation
Like the WFR gradient flow over non-negative measures, the gradient flow equation (8) and ( are not guaranteed to stay within the probability measure space, i.e., total mass 1. This is useful in many applications such as chemical reaction systems. However, probability measures are often required for machine learning applications. We now provide mass-presnrving gradient flow equation that term the spherical IFT grudient flow. The term spherical used to emphasize that the flow stays within the probability measure, in the spherical Hellinger distance [Lscchos and Mielke 2019
To this end, we must first study spherical MMD flows over probability measures. Recall that is Hilbert space gradient flow (see (Ambrosio et al. 2 and does not stay within the probability space, Closely related, many works using kernel- mean embedding (Smola et al., 2007 Muandet et al. 2017) also suffer from this issue of not respecting the probability space. To produce restricted (or projected) flow in P, our starting point the MMD mmnimizing-movemee scheme restricted to the probability space
DF
11
1+1 +- argmin F(u) + (M,
(9)
Hk
PEP
21j
Mk11
We now derive the following mass-preserving spherical gradient flows for the MMD and IFT. Proposition 3.2 (Spherical MMD and spherical IFT gradient flow equations). The spherical MMD gradient flow equation is given by (where denotes the constant scalar)
8F [+]
KC
6F
[M]
i -
(10)
8p
JK-11
Consequently, the spherical IFT gradient flow equation is given by
5F (u]) - 3.- 5y
5F
1 KK-1
3p [M]
8F
i == div
Ōu [m]
(11)
SK-11
Furthermore, those equations mass-preserving. i.e., p==0 0.
So far, we have identified the gradient flow equations of interest. Now, we are ready to present our main theoretical results the convergence of the IFT gradient flow via functional inequalities. For example, the logarithmic Sobolev inequality (LSI)
dp
log
2 Cr.t ( for some qust >> 0
(LSII
dm
(P(
is sufficient to guarantee the convergence of the pure Wasserstein gradient flow of the KL divergence energy, which governs the same dynamics as the Langevin equation. The celebrated Bakry- Emery Theorem [Bakry and Emery| 1985), is a cornerstone convergence analysis for dynamical systems as it provides explicit sufficient condition: suppose the target measure π is A-log concave for some 0, then the global convergence is guaranteed, i.e.,
dx and V 2 X Id LSSI) with Gıst 2A glob. exp. convergence.
The question we answer below is whether the IFT gradient flow enjoys such favorable properties.
Our starting point is the (oolyak))oojsiievee type functional inequality.
5

---

Theorem 3.3. Suppose the following Lojasiewicz type inequality holds for some >> 0,
6F
5F
[4]
B.
[4
2 c.
(((u(t)) - inn F(H))
(IF'T-Loj)
8p
for the IFT gradient flow, or
8F 8p
5F 8p
8F K
[1a]
- inf F(w) (SIFT-Loj) cc. (F(u(())
a.
+
[u]
JK-11
1
for the spherical IFT gradient flow. Then, the energy F decays exponentially along the corresponding gradient flow, i.e., F(u(()) - inf, F(u) Seaa (F(u(0)) - inf, F(u)).
To understand a specific gradient flow, one must delve into the detailed analysis of the conditions under which the functional inequalities hold instead of assuming them to hold by default. We provide such analysis for the IF'T gradient flows next.
3.2 Global exponential convergence analysis
MMD energy functional As discussed in the introduction, the MMD energy has been proposed as an alternative to the KL divergence energy for sampling by Arbel et al. 2019 where they assume the access samples from the target measure Y/ π. However, the theoretical convergence guarantees under the MMD energy is a less-exploited topic. Those authors characterized a local decay behavior under the assumption that the [4 must already be close to the target measure π for all / >> 0. The assumptions they made are not only restrictive, but also difficult to check. There has Garbel also been global convergence analysis. For example, Arbel al.] [2099]'s 's Proposition 2 states that the MMD is non-increasing, which is not equivalent to convergence and is easily satisfied by other flows. The mathematical limitation is that the MMD is in general not guaranteed along the Wasserstein geodesics. In addition, our analysis also does not require the heuristic noise injection step was required in their implementation. MMouuh and Rigotti (2020) also used the MMD energy but with different gradient flow, which has been shown by Zhu and Mielke [2024] be kegull-regularized inf oonvoluootti of the Allen-Call and Cahn-Hillian type of dissipation. However, Mroueh and Rigotti] [2020)*s 's convergence analysis is not sufficient for establishing (global) exponential convergence as no functional inequality has been established there. contrast, we now provide full global exponential convergence guarantees.
be convex
We first provide interesting property that will become useful for our analysis.
Theorem 3.4. Suppose the driving energy is the squared MMD, F(u) MMMD X) and initial datum pr E P is a probability measure. Then, the spherical MMD gradient flow equation (10) coincides with the MMD) gradient flow equation
A== (1-x),
(MMD-MMI-GF)
whose solution is a linear interpolation between the initial measure μo and the target measure म, i.e.,
Hi e "o + (1
Furthermore, same coincidence holds for the spherical IFT and IFT gradient flow equation
a div (1 V2k(z,)) d (n= T) (x) 3 sx)-
(MMD-IFT-GFF
The explicit solution to the ODE (MMD-MMDD-G shows an exponential convergence to the target measure π along the (spherical) MMD gradient flow. The (spherical) IFT gradient flow equation (MMD-IFFT
differs from the Wasserstein gradient flow equation of [Arbel et al. 2019 by a linear term. This explains the intuition of why we can expect good convergence properties for the IFT gradient flow of the squared MMD energy. We exploit this feature of the 1FTT gradient flow to show global convergence guarantees for inference with the MMD energy. This has not been possible previously when confined to the pure Wasserstein gradient flow.
We do not use the terminology "MMD gradient flow" from [Arbel et al. 2019 since it is inconsistent with the naming convention of "Wasserstein gradient flow" as Wasserstein refers the geometry, not the functional.

---

Theorem 3.5 (Global exponential convergence of the IF'T flow of the MMD energy). Suppose the energy F is the squared MMD energy F(J)= = (p, v). Then, the (I -t.jj holds globally with constant 23 0.
Consequently, for any initialization within the non-negative measure cone pe M*, the squared MMD energy decays exponentially along the IFT gradient flow of non-negative measures, i.e.,
MMMD? 5) 1 (Mo"
(12)
Furthermore, the initial datum μo and the target measure x are probability measures Pa, ग E P. then the squared MMD energy decays exponentially globally along the spherical IFT gradient flow, i.e. the decay estimate (12) holds along the spherical 1FT gradient flow of probability measures.
We emphasize that no Bakry-Emery type kernel conditions are required the Lojasiewicz inequality holds globally when using the IF'T flow. In contrast, the Paasserstiin-Fisher. Rao gradient flow
Vak(u, :) d (u x) (4) - 3. f k(xx.) (u =x) (x) (MMD-WFR-GG)
j == a div
does not enjoy such global convergence guarantees.
Global exponential convergence under the KL divergence energy For variational inference and MCMC, common choice of the energy is the KL divergence energy, i... F(() = -KL (44 T This has already been studied by a large body of literature, including the case of Wasserstein- Fisher- Rao [Liero et al.. 2023 Lu et al. [019]. Not surprisingly. (LSI) is still sufficient for the exponential convergence the WFR type of gradient flows since the dissipation of the Wasserstein part alone is sufficient for driving the system to equilibrium. For the IFT gradient flows under the KI. divergence energy functional, the convergence can still be established. This showcases the strength of the IFT geometry - it enjoys the best of both worlds. The IFT gradient flow equation of the KI. divergence dy
dy
Unlike the MMD-energy flow case, the spherical IFT gradient flow of the KL. over probability measures P no longer coincides with that of the (non-spherical) IF' and is given by
energy reads jز = a div log
B KC-
di
dı
du du - 3 1- log
du dx
I log JK-11
dp du
j = a - div log
(13)
Proposition 3.6 (Exponential convergence of the SIFT gradient flow of the KI. divergence energy).
Suppose the (LSI) holds with CLSI 2X the target measure æ is X-log concave for some À >>> Then, the KL divergence energy decays exponentially globally along the spherical IFT gradient flow i.e. DKL(HITT) I -2adt (HoooT))))
The intuition behind the above result is that the SIFT gradient flow converges whenever the pure Wasserstein gradient flow, i.e., its convergence is least as fast as the Wasserstein gradient flow. However, we emphasize that the decay estimate of the KI divergence energy only holds along the spherical IFT flow over probability measures P, but not the full IFT flow over nen-negative measures M
3.3 Minimizing movement, JKO-splitting, and practical particle-based algorithm
In applications to machine learning and computation, continuous-time flow be discretized via the JKO scheme [Jordan et al. 1998), which based on the minimizing movement scheme (MMS) [De Giorgi 1993). For the reaction-diffusion type gradient flow equation in the Wasserstein- Fisher-Rao setting, the JKO-splitting a.k.a. time-splitting scheme has been studied by Gallouët and
Monsaingeon [017], Mielke et al. 22023] This amounts splitting the diffusion (Wasserstein) and
reaction (MMD) step in (8) i.e., at time step l 2
1 W2 (H, n²),
+ argmin F(u) + PEP
(Wasserstein step)
2T
(14)
1 ++ + argmin F(H) + MMD2 (H, μ b)).
1
(MMD step)
PEP
2m
7

---

A similar JKO -pllitiiting scheme can also be constructed via the WFR gradient flow, which amounts to replacing the MMD step in (14) with a proximal step the KL (as an approximation the Hellinger),
ie., f f p F(u) + (ulu ( )) which is well-studied in the optimization literature as the entropic mirror descent [Nemirovskij and Yudin 1983]. Our MMD step can also viewed as mirror descent step with the mirror map 2 | - However, for the task of MMD inference of [Arbel et al., 2019]. WFR flow does not possess convergence guarantees such as our Theorem 3.4 The MMD step can also be easily implemented in our simulation.
We summarize the resulting overall IFT particle gradient descent from the JKO splitting scheme in Algorithm [ in the appendix. We now look at those two steps respectively. For concreteness, we consider a flexible particle approximation to the probability measures, with possibly non-uniform weights allocated to the particles, i.e., J -i=1 aqda, a E A", I; X.
Wasserstein step: particle position update. (14) is a standard JKO scheme; see, e.g., [Santambro+ gio 2015]. The optimality condition of the Wasserstein proximal step can be implemented using particle gradient descent algorithm
OF [A (at), 1, T 5u
[+1
(15)
which is essentially the algorithm proposed by Arbel et al. [2019] when F(u) MMD m).
MMD step: particle weight update. The MMD step in (4) is a discretization step of the spherical MMD gradient flow, as shown in ( and Proposition 3.2 We propose to use the updated particle
location 2; from the Wasserstein step
(15 and update the weights 33 by solving
inf F SEA"
3,5 (11)
MMD Σ
3,
Σ
2m
ie., the MMD step only updates the weights. Alternatively, as in the classical mirror descent optimization, one can use linear approximation F(u) F(uE) + Vaa [ , -M- uf) 12. We also provide specialized discussion the MMD-energy minimization task of [Arbel et al. 2019]. Let the energy objective be the squared MMD F(u) := MMD(H, T)?. In this setting, we are given the particles sampled from the target measure yf ~ т. For the MMD step in (14), the computation is drastically simplified to MMD barycenter problem, which was also studied in [Cohen et al.) 2021).
8E
This amounts to solving a convex quadratic program with a simplex constraint; see the appendix for the detailed expression.
4 Numerical Example
The overall goal of the numerical experiments is to approximate the target measure π by minimizing the squared MMD energy, i.e., min MMD?(H, T). In all the experiments, we have access to the
#EACP
target measure in the form of samples yi т. This setting was studied in [Arbel al., 2019 as well as in many deep generative model applications. In the following experiments, we compare the performance of proposed algorithm of IFT gradient flow, which implements the JKO-splitting scheme in (14) and is detailed in Algorithm] I to that of (1 Arbel et al. [2019]'s s the "MMD flow" (see our discussion in () used their algorithm both with and without heuristic noise injection suggested by those authors; (2) the Wasserstein-Fistee
Rao flow of the MMD /MMD-WFR-GIT The
WFR flow was also used by Yan al. [2024. Lu et al. (2023) but for minimizing the KI. divergence function. As discussed in $3.2 the MMD flow of [Arbel et al.] [019] does not possess global convergence guarantees while IFT does. Furthermore, in the Gaussian mixture target experiment, the target measure π is not log-concave. We emphasize that our convergence guarantee still holds for the IFT flow while there no decay guarantee for the WFR flow. We provide the code for the implementation at /
com wwt.///.///eeeeeeec egorgl wwiggglainnnai /ittt/ffffffo//////
Gaussian target in 2D experiment Figures and 4 showcase the performance of the algorithms in a setting where nr and æ are both Gaussians in 2D. Specifically, nf2 ~ (5 1, and न ~
N (o. 1 1/2) )² The number of samples drawn from μº and x was set to n = 100. A Gaussian 1/2

---

MMD flow
MMD flow noise
IFT GD
12"1
10
12"
12"4
MMD flow noise
TFT GD
1000
2000
4200
0000
1000
2200
2000
4200
5200
Stpps
(a) Gaussian target experiment
(b) Gaussian mixture target experiment
Figure 3: Mean loss and standard deviation computed over 50 runs
(a) Vanilla MMD flow
(b) MMD flow noise
(c) IFT (Algorithm)
Figure 4: Trajectory of randomly selected subsample produced by different algorithms in the Gaussian target experiment. Color intensity indicates points' weights. The hollow dots indicate the particles that have already vanished.
kernel with bandwidth o = 10 was used. For all three algorithms, we chose the largest stepsize that didn't cause unstable behavior, observed from the trajectories produced by MMD flow (Figure 4 fa)), most points collapse into small clusters near the target mode. Some points drift far away from the target distribution and get stuck; the resulting samples represent the target distribution poorly, which is sign of suboptimal solution. MMD flow with the heuristic noise injection produces much better results. We suspect the noise
50. The parameter v (23) was 0.1. As can be
(a) Vanilla MMD flow
(b) MMD flow noise
(c) IFT (Algorithm
Figure 5: Trajectory of randomly selected subsample produced by different algorithms in the Gaussian mixture experiment. Color intensity indicates points' weights. The hollow dots indicate the particles that have already vanished.

---

helps to escape local minima; however, the injection needs to be heuristically tuned. However, it takes a large number iterations for points to get close locations with high density of the target distribution. Similarly to the previous research on noisy MMD flow, we use relatively large noise level (10) the beginning and "turn off" the noise after a sufficient number of iterations (3000 in our case). A drawback of this approach is that the right time for noise deactivation depends on the particular problem instance, which makes the algorithm behavior less predictable. Algorithm I achieves a similar accuracy to that of the noise-injected MMD flow, but much faster already after 1000 steps without any heuristic noise injection. For the few particles that did not make it close to the target Gaussian's 's mean, their mass is teleported those particles that close to the target. Hence, the resulting performance of the IFT algorithm does not deteriorate. The hollow dots in the trajectory plot in- dicate the particles whose mass has been teleported and hence their weights are zero. This a major advantage of unbalanced transport for dealing with local minima. For faster implementation, in the implemen- tation of the MMD step in (14). we only perform single step of projected gradient descent instead of computing solution to the auxiliary optimization problem (23). To be fair in comparison, we count each iteration as two steps. Thus, 6000 steps of the algorithm in Figure 3aa) correspond to only 3000 iterations, i.e., the results for IFT algorithm have already been handicapped by factor of 2. We would also like to note that Algorithm was executed with constant hy pyrrpaaraeetess without further tuning over the iterations, in contrast to the noisy MMD flow. In practical implementations, it is possible to further improve the performance by sampling new locations (particle rejuvenation) in the MMD step (14) as done similarly in (Dai et al., 2016). Since this paper is theoretical one and not about competitive benchmarking, leave this for future work.
MMD flow MMD flow noisc
IFT GD
WFR GD
10-1
MNC
10-1
2000
<002
6000
Stape
Figure Comparison with the WFR flow of the MMD in 100 dimensions
Gaussian mixture in 2D experiment The second experiment has a similar setup. However, this time the target is a mixture of equally weighted Gaussian distributions,
NV (0.(1/
N ((2).)). (1)22)))
(1/2
((2..(1211)
Figures ()) and [ showcase loss curves and trajectories produced by the considered algorithms.
WER flow for Gaussian mixture target in 100D We conducted an experiment dimension d' = 100. comparing the IF'T flow with the WFR flow of the MMD energy. The initial distribution is / I). and the target ı is a mixture of distributions ((m,, Σx), = 1, 1,2, 3, where m, and Σ, are randomly generated such that | /m:l22 ||| 20 and the smallest eigenvalue of Σ, is greater than 0.5. For fairesss each iteration of IFT particle GD (as well as its version with KI. step) counts as two steps, i.e., these methods only performed 4500 iterations. In the noisy MMD flow, the noise is disabled after 4000 steps. All methods are used with equal stepsize.
Discussion
In summary, the (spherical) IFT gradient flows are suitable choice for energy minimization of both the MMD and KL divergence energies with sound global exponential convergence guarantees. There is also an orthogonal line of works studying the Stein gradient flow and descent (Liu and Wang 2019 Duncan et al., 2 which also has the mechanics interpretation of repulsive forces. can be related to our work in that the IF'T gradient flow has a (de- deeeeleeee reaction part, while the Stein flow has a kernelized diffusion part. Furthermore, there is work Manupriya et al., 2024) that proposes static regurregulriieed Wasserstein distance, which should not be confused with our IFT gradient flow geometry. Another future direction is sampling and inference when we not have access the samples from the target distribution but can only evaluate its score function V log n.
10

---

Acknowledgments and Disclosure of Funding
We thank Gabriel Peyré for the helpful comments regarding the practical algorithms for the JKO- splitting scheme. This project has received funding from the Deutsche Forschungseemeinccatt (DFG, German Research Foundation) under Germany's Excellence Strategy The Berlin Mathematics Research Center MATH+ (EXC-2046/1, project ID: 390685689) and from the priority programme "Theoretical Foundations Deep Learning" (SPP 2298, project number: 543963649). During part of the project, the research of E. Gladin was prepared within the framework of the HSE University Basic Research Program.
References
L. Ambrosio, N. Gigli, and G. Savare. Gradient Flows: Probability Measures. Springer Science & Business Media, 2005.
In Metric Spaces and in the Space of
M. Arbel, A. Korba, A. SALIM, and A. Gretton. Maximum mean discrepancy gradient flow, In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, ed- itors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL htttps:/prroceedingss neurips. cttpiiis..ccppppppp faapper_flls/ppppeeeereeeeeeeeee 94455ac088aaeaaa0000cc0777222690 pdf
M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pages 21-2233. PMLR, 2017.
D. Bakry and M. Emery. Diffusions hypercontractiivss. In J. Azéma and M. Yor, editors, Séminaire de Probabilités XIX 1983/84, volume 1123, pages 177-206. Springer Berlin Heidelberg, Berlin, Heidelberg, 1985. ISBN 978-3-540-11 15200-9 978-3-540-39397-9) doi: 10. 10.1007/BF60075847.
J. Bigot, R. Gouet, T. Klein, and A. López. Geodesic PCA the Wasserstein space by convex PCAA. Ann. Inst. H. Poincaré Probab. Statist., 53(1):1-26, 02 2017. doi: 10.1214/1-AIHHP766 URL
https://doi. org/10. 1214/1111111111 11-14/11111176666
J. A. Carrillo, K. Craig, and F. S. Patacchini. A blob method for diffusion. Calculus of Variations and Partial Differential Equations, 58:1-53, 2019.
Chewi, T. Le Gouic, C. Lu, T. Maunu, and P. Rigollet. SVGD as kernelized Wasserstein gradient flow of the chi-squared divergence. Advances in Neural Information Processing Systems, 33: 2098-28 2009, 2020.
L. Chizat. Sparse optimization on measures with over-parameterized gradient descent. Mathematical Programming, 194(1-2)4487-332 2022.
L. Chizat, G. Peyré, B. Schmitzer, and F.-X. Vialard. An interpolating distance between optimal transport and Fisher-Rao metries. Foundations of Computational Mathematics, 18(1):1-44, Feb. 2018. ISSN 1615-3375, 1615-3383. doi: 10.100//s10008-000 016-933--1
L. Chizat, G. Peyré, B. Schmitzer, and F.-X. Vialard. Unbalanced optimal transport: Dynamic and Kantorovich formulation. arXiv:1508.05211 [math], Feb, 2019. URL http:://trrrr http::://xxxivorr/abb/ 1508.080 052166
arXiv: 1508.05216.
S. Cohen, M. Arbel, and M. P. Deisenroth. Estimating barycenters of measures in high dimensions. arXiv:2007.07105 {es, stat], Feb. 2021.
K. Craig, K. Elamvazhuthi, M. Haberland, and O. Turanova. A blob method for inhomogeneous diffusion with applications to multi-agent control and sampling. Mathematics of Computation, 92 (344):2575-2654, Nov. 2023. ISSN 0025-5718, 1088-6842. doi: 00.1090/mcom/3841.
M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
B. Dai, N. He, H. Dai, and L. Song. Provable Bayesian inference via particle mirror descent. In Proceedings of the 19th International Conference on Artificial Intelligence and Statistics, pages 985-994. PMLR, May 2016.

---

E. De Giorgi. New problems minimizing movements, Ennio de Giorgi: Selected Papers, pages 699-713, 1993.
A. Duncan, N. Nüsken, and L. Szpruch. On the geometry of Stein variational gradient descent. arXiv preprint arXiv:1912.008%/ 2019.
M. A. Efendiev and A. Mielke. On the anteiindeeendent limit of systems with dry friction and small viscosity. Journal of Convex Analysis. 13(1):151, 2006.
J. Feydy, T. Séjouréé, F.-X. Vialard, S.-i. Amari, A. Trouve, and G. Peyré. Interpolating between optimal transport and MMD using Sinkhorn divergences. In Proceedings of the Eventy-Second International Conference on Artificial Intelligence and Statistics, pages 1681-2690. PMLR, Apr. 2019.
T. O. Gallouët and L. Monsaingeon. A JKO splitting scheme for Kantorovich Sinntorrvich-iishe---ee gradient flows. SIAM Journal on Mathematical Analysis, 49(2):1 2017.
P. Glaser, M. Arbel, and A. Gretton. KALE flow: A relaxed KL gradient flow for probabilities with disjoint support. In Neural Information Processing Systems, June 2021.
A. Gretton, K. M. Borgwardt, M. J. Rasch, B. Schölkopf, and A. Smola. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012.
J. Hertrich, C. Wald. F. Altekrüger, and P. Hagemann. Generative sliced MMD flows with Riesz kemels.. arXiv preprint arXiv:2305. 1463, 2023.
N. Ho, X. Nguyen, M. Yurochkin, H. H. Bui, V. Huynh, and D. Phung. Multilevel clustering via Wasserstein means. In D. Precup and Y. W. Teh, editors, Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1501-1509, International Convention Centre, Sydney, Australia, 06-11 Aug 2017. PMLR. URL http://proceediing mlr, Tiress/v70/hho 17a. html
R. Jordan, D. Kinderlehrer, and F. Otto. The variational formulation of the Fokker Plaker equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998. Publisher: SIAM.
Kondratyev, L. Monsaingeon, and D. Vorotnikov. A new optimal transport distance on the space of finite Radon measures, Advances in Differential Equations, 2 21(11/12: 1164, 2016. doi: 10.57266//dee147766669
URL https: //doi. 10. 57262 ade/ 1476369298
A. Korba, P.-C. Aubin Frann-Frankowski, S. Majewski, and P. Ablin. Kernel Stein discrepancy descent. In Proceedings of the 38th International Conference on Machine Learning, pages 5719-5730. PMLR, July 2021.
M. J. Kusner, Y. Sun, N. I. Kolkin, and K. Q. Weinberger. From word embeddings to document distances. In Proceedings of the 32nd International Conference on International Conference on Machine Learning Volume 37, ICML' 15, pages 957-966. JMLR.org, 2015. URL http 7/dl. acm. org/ citation.... ffm?idd 3045111801111.11 11112211..04211111
Lacoste-Julien F. Lindsten, and F. Bach. Sequential kernel herding: Frank- Wolfe optimization for particle filtering. In G. Lebanon and S. V. Vishwanathan, editors, Proceedings of the Eighteenth International Conference on Artificial Intelligence and Statisties, volume 38 of Proceedings of Machine Learning Research, pages 544-552, San Diego, California, USA, 09-12 May 2015. PMLR. URL httpst :/pooeeed////ceeeegssssee mlr. press//338 wwt/////laaseeeeeeeeeeeeeeaeeeee julien15..... html
V. Laschos and A. Mielke. Geometric properties of cones with applications on the Hellinger- Kantorovich space, and a new distance on the space of probability measures. Funci. Analysis, 276(1 (76(11):35293377 2019. doi: 10. (0.1 8. 12.001..
M. Liero, A. Mielke, and Savaré. Optimal entropy-transport problems and a new Hellinger- Kantorrvich distance between positive measures. Inventiones mathematicae, 211 (3):969-1117, Mar. 2018. ISSN 0020-9910, 1432- doi: 10. 10.11077/00022-011--773333 URL httt:://inkk.. springer com http:/000100//000222011000055550
12

---

M. Liero, A. Mielke, and G. Savaré. Fine properties of geodesics and geodesic X-convexity for the Gellinger-Kantrrovee distance. Arch. Rat. Mech. Analysis, 247(112):1-73) 2023. doi: 10. 00.10077600225-023-009441111
Q. Liu and D. Wang. Stein variational gradient descent: A general purpose Bayesian inference algorithm. 1608.066 [cs, stat), Sept. 2019. URL BtRphttp://ariin http:///axxv..v.o wttpa/awwww....//////////// 04471 arXiv:
Y. Lu, J. Lu, and Nolen. Accelerating Langevin sampling with birth-death. ArXiv, May 2019.
Y. Lu, Slepčev, and L.. Wang. Birth-deate dynamics for sampling: global convergence, approxima- tions and their asymptotics. Nonlinearity, 36(11):5733 1):5731, 2023.
Manupriya, S. N. Jagarlapudi, and P. Jawanpuria. MMD-regullarize unbalanced optimal transport. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. URL https://opeerrviiee htt.://wwee..toormmmmcoeeeeiiiio
A. Mielke, R. Rossi, and Stephan. On time-splitting methods for gradient flows with two dissipation mechanisms. arXiv preprint 1711
P. Mohajerin Esfahani and D. Kuhn. Data-driven distributionally robust optimization using the Wasser- stein metric: Performance guarantees and tractable reformulations. Mathematicaal Programmingg 171(1-2): Sept. 2018. ISSN 0025-5610, 1436-4646. doi: 10.1007/410107-011--177..
Y. Mroueh and M. Rigotti. Unbalanced Sobolev descent, Sept. 2020. URL httt:tttta/xi...o/ abs/2008 14148 arXiv:2009. 14148 [es, stat].
K. Muandet, K. Fukumizu, B. Sriperumbudur, and B. Schölkopf. Kernel mean embedding of distributions: review and beyond. Foundations and Trends@ Machine Learning, 10(1-2): 1-141, 2017. ISSN 1935-8237, 1935-8245, doi: 10. 10.1561/2200000066
A. S. Nemirovskij and D. B. Yudin. Problem Complexity and Method Efficiency in Optimization. John Wiley, New York, 1983.
S. Neumayer, V. Stein, and G. Steidl. Wasserstein gradient flows for Morcau envelopes of f- divergences in reproducing kernel Hilbert spaces. arXiv preprint arXiv:2402.04613 2024.
Otto. Double degenerate diffusion equations as steepest descent. Bonn University, 1996. Preprint.
Y. Rubner, C. Tomasi, and L. J. Guibas. The earth mover's distance as metric for image retrieval. International journal of computer vision, 40(2):99-121, 2000.
R. Sandler and M. Lindenbaum. Nonnegative matrix factorization with earth mover's distance metric for image analysis. IEEE Transactions on Pattern Analysis and Machine Intelligence, 33(8): 1590-1602, 2011. ISSN 0162-8828. doi: 10. 10.100///A
Santambrogio. Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs, and Modeling. Springer International Publishing, 2015. ISBN 9783319208288 doi: 10.1007/ 978-3-3 3--1---1122222 htRh:/tp://ds wwo.o.oog/1..0 11t.:////0111100000-0000----10 20828-2
A. Sinha, H. Namkoong, R. Volpi, and J. Duchi. Certifying some distributional robustness with principled adversarial training. Xrxiv: 1057. [es, stat), May 2020.
A. Smola, A. Gretton, L. Song. and B. Schölkopf. A Hilbert space embedding for distributions. In M. Hutter, R. A. Servedio, and E. Takimoto, editors, Algorithmic Learning Theory, pages 13-31, Berlin, Heidelberg, 2007. Springer. ISBN 978-3-540-75225-7, doi: 10. 10.10079778-3-540-755 5.
Solomon, R. M. Rustamov. Guibas, and A. Butscher. Wasserstein propagation for semi- supervised learning. In Proceedings of the 31st International Conference on International Confer- ence on Machine Learning Volume 32, JCML' 14, pages 1-306-1-314, JMLR.org, 2014. URL httpp://ll acm. org/ httg/iitaioo.n.c cfm? id=344488 30000444880004441111 3044841
I. Steinwart and A. Christmann. Support vector machines. Springer Science & Business Media, 2008.
13

---

I. Tolstikhin, B. Sriperumbudur, and K. Muandet. Minimax estimation of kernel mean embeddings. 16X.:3662.04361 [math, stat], July 2017.
O. Tolstikhin,
Sriperumbudur, and B. Schölkopf. Minimax estimation of maximum mean discrepancy with radial kernels. In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Gar- nett, editors, Advances in Neural Information Processing Systems, volume 29. Curran Asso- ciates, Inc., 2016. URL http:::: f
neurips htte:/pppr_ffeee/eepeeeee www.lles/appp//00/// 2016/ Paper Paper. pdf
A. B. Tsybakov. Introduction to Nomparametric Estimation. Springer Series in Statistics. Springer, New York, NY, 2009. ISBN 978-0-387-790511--1 978---377777777-----7
doi: 10.1007/b13794, URL
https:///ina
C. Villani. Optimal transport: old and new, volume 338. Springer Science & Business Media, 2008.
M. Werman, S. Peleg, and A. Rosenfeld. A distance metric for multidimensional histograms. Computer Vision, Graphics, and Image Processing, 32(3):328 336, 1985. ISSN 1734-189X. doi: http://b..orr/10.000000/70444419
URL http://wwww... www.sccinnedirecc.oommm com/
Y. Yan, K. Wang, and P. Rigollet. Learning Gaussian mixtures using the Wasserstein-Fisher gradient flow. The Annals of Statistics, 52(4):1774-1795, 2024.
J.-J. Zhu and A. Mielke. Kernel approximation of Fisher-Rao gradient flows, 2024. URL https //axivv. ferxiv.org//e
20622
14

---

A Appendix: proofs and additional technical details
Proof of Proposition 3.2 We first consider the MMS step
9). The Lagrangian of the MMS step is,
for À E R.
/ p p-
((,,)) F())
MMD SA,DD²(,,, n²) + (u, 2
2n
The Euler-Lagraage equation gives
8F
|n'] + KK(u v²)) + +\
(16)
8j
f 1 1-=1
(17)
Rewriting the first equation,
,FF
P. = n¹ nIK
|1t nAK--
Integrating both sides, we obtain
== f K
8F 6μ
[n']
nn
K 1F OF 8f
[u]]
f
[K--1
Let the time step > 0 in the first equation in the Euger-Lagrange equation
we obtain
SF
5F
I K
5F 5p
[p|
A K
[)]
-
5y [m]
+
(18)
8p
SK-
which is the desired spherical MMD gradient flow equation.
Spherical IFT gradient flow equation is obtained by inf-convolution [Gallouët and Monsaingeon 2017 Liero et al. 2018 Chizat et al.) 2019] of the above spherical MMD and Wasserstein part. The verification of the mass- waas-prreeervig property is by straightforward integration of (18)
0 f j
/
K 8F 8µ
JK-
5F du
[p]
fx-11=0
[n]
JK-1
Hence, the theorem proved.
Proof of Theorem 3.3 The proof amounts to identifying the correct left-hand side of the Lojasiewicz type inequality. We take the time derivative of the energy
d 5F 8p
8F [4[) B K 8p
8F
F()) dt
[u] div(yV
[u])
8F
5F 8p
=
|
[1]]
11:g - -B.
[4]
which is the desired left-hand side of the Lojasiewicz type inequality. For the spherical IFT gradient flow, we have
8F d F()) = ( (oy 6F [m],0 a div(uV (u]) - 3. k-
5F
5F K ou
[R]]
6f [R]]
)
dt 8u
8y
S [K-11
8F ( 8y [u] a.div
8F
8F
1 JK- KK-1 1 8F
(p]
-B----1 ( 8FF [u]
f IC 6F
ou [x]
[ni))
+
[m]]
8p
[ KC-1
1 K-
8F
6F
f KC
[pi]
V 8h
[4]
[m]]
8yc
6p
JK-11
13
7
where the second equality follows from the fact that the spherical IFT gradient flow is mass- preserving. Hence, the left-hand side of the Lojasiewicz type inequality is obtained.
15

---

Proof of Corollary 3.7 The formal proof is by using well-known results for kemel smoothing in 2009 Note that the gradient flow equation can be rewritten as
non-parametric statisties (Tsybakov
SF
OF
OF [n]] + B uu 6f
O1F
- ----iiivvv
div (u]) + B K
[n] = aa - div(uV
8u [n].
5j
Recall the definition of the integral operator
Kf ko(-)))(y)) dy, Kuf= = No (o fhof=/holW))t
du(y).
(19)
Formally, as Kor (ar, :) dj - dgg. we have Ku + & for any & E (1)) Then, the IFT gradient flow equation tends towards the PDE
5F
i = -a div(uV ōu
[A]] + 8p
8u [n].
(20)
Furthermore, we also have
65 8y
|p]11 + a [M] ||: Hence, the conclusion follows.
Proof of Theorem 3.4
We first recall that the first variation of the squared MMD energy given by
5
(1 2
[w] [w]= (MD
k(a,, )u 2) dz).
(21)
5fl
Plugging the first variation of the squared MMD energy into the gradient flow equation (10)
and (TI, obtain the desired flow equations in the theorem. The ODE solution is obtained by elementary argument.
Proof of Theorem 3.3 We take the time derivative of the energy and apply the chain rule formally and noting the gradient flow equation (
d dt F(u) = /
8F [4], a. div ( 8p
SF 8y
5F
[u]) B B-K-
[n]]
L2
5yc
8F
8F 5j
6F 8p
= V
[+] 11:; 3-
[+]] 1 -
-3.
[A]
like
ōμc
Plugging in F(1)) MMDD (1, v). an elementary calculation shows that
dt ( (MD v)) s - k(x,,) - dz) = 2, MMD v),
d
(22)
which establishes the desired Łojasiewicz inequality specialized to the squared MMD energy. which reads
Ma.-M-M v)( dx) (dxx))
+
| M(x, v) dz)
2 (MND((,,)) (Loj)
By Grönwall's lemma, exponential decay established.
Furthermore, plugging the first variation of the squared MMD energy into the gradient flow equa- tion 10p, the extra term in the spherical flow equation becomes
K.
8F
[p]
- -)) 0.
bu
SK--1
JK-
Hence, the coincidence is proved.
By the Bakry-Emery Theorem, we have the LSI (LST) hold with cLsI 2X. Taking the time derivative of the KL divergence energy along the SIFT gradient flow, we have
Proof of Proposition 3.6
d
dip dn
DKL (MelTT) = di
V log
)
122
djp 2 dr 12
du
f KC
log
du dr
V log
-
log
dr
7
2ad (KL(H( +0.
By Grönwall's Jemma, exponential convergence is established.
16

---

Note that this result does not hold for the full IFT flows over non-negative measures as there exists no LSI globally on M
Remark A.1 (Regularized inverse of the integral operator). Strictly speaking, the integral operator is compuct and hence its inverse is unbounded. Using a viscosity- regulaizaegtlaiiatttinn techniques by Efendiev and Mielke (2066), we can obtain the flow equation where the inverse is always well well-defined, ie.. ft = a .div( V
5μ [n]. This corresponds to an additive regularization of the kernel Gram matrix in practical computation.
OF [4] - 3 (KC + -)- &f by
A particle gradient descent algorithm for IFT gradient flows We use the notation Xxx to denote the kernel Gram matrix Kxx = [ " ) Kxx for the cross kernel matrix
f+1 (+1))
K Kxx- = [M(a
etc.
Algorithm A JKO spliitiig for IFT particle gradient descent
Require:
1: for 6o to T 1 do
2: Compute the first variation of the energy at μ²: gf = 6F dy
[n]]. Then,
1+1 + o (x()
i 1.....,
12
(Wasserstein step)
1
[a"
Kxx -xx8
-xxx a
a 1+1 argmin F()
0r ++1 +
20p
at
Kxx a²
(MMD step)
"EA"
3: end for
4: Output the particle measure Ta' -E= of
Implementing the MMD step The MMD step in the JKO-splittng scheme is a convex quadratic program with simplex constraint, which can be formulated as
1 MMD? Σ SEA" 2
a
1 MMD
n [
a
inf
3,8_,tm (,, 21
B,8 E-1,
(23)
We further expand the optimization objective (multiplied by a factor of 2r for convenience)
Σ
3,0(11t1)
Σ
((U))
+
Σ
3,0(0
Σ
a
(aa
i=1
j=1
1=1
=
B8
3xx
3B TT Kxv
2
+ m
1 km1) + (3 Kxx - 23 Kxxa + a Kxxa)
m
=11+T) BT K 3
2r
2T
3T Kxyl 23 -
Kxxa. (24)
m
Therefore, the MMD step in Algorithm simplex constraint.
I can be implemented as a convex quadratic program with a
particle gradient descent algorithm for the WER flow of the MMD energy We provide the implementation details of the WFR flow of the MMD energy. The goal is simulate the PDE MMDD-WFR-GF WFR-GF). To the best our knowledge, there has been no prior implementation
this flow.
Nor there convergence guarantee. Similar to the JKO-splitting scheme of the IFT flow in (14) (MMD-WFR-GF) can be discretized using the two-step scheme
1 ++] +- arg min F(n) + W3 (Wasserstein step) PEP
2r
2+ argmin F(R) 1 -L( (KL step) 7)
MEP
where the energy function F is the squared MMD energy, F(n) = y MMD .). Use the explicit Euler scheme, the KL step amounts to the entropic mirror descent. In the optimization literature, this
17

---

step can be implemented as multiplicative update of the weights (or density), i.e., suppose x, is the new particle location after the Wasserstein step, then we update the weights vector a via
f+1
s+1 a. exp (-m óy
8F
wn(oftt))
ai
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, our claims in the abstract and introduction do accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
Justification: In various places in the paper, we discuss the limitations of our work. Especially in the discussion/conclusion section, we discuss the limitations of our approach.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sellsspeciiiationn asymptotic approximations only holding locally). ). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
19

---

Answer: [Yes]
Justification: Yes. For example, one focus of our work is precisely to study under what assumptions, functional inequalities such the Lojasiewicz inequality and the log-Sobolev inequality hold. We do not avoid discussing such assumptions. We provide a complete proof for each theoretical result.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Our implementation is standard and also similar to that of [Arell et al. 2019 and [Korba et al. 2021].
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibilityy In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
20

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: Similar to the answer to the previous question. We have submitted the code and added to the text link to github repo with the details the experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
wublic//giide//CdeStttimiss wit.ssammmsioooooooooooooo.o. for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https for more details.
//nips //nips.c/ppul
guides PoodeSuumissiinooo......ooo
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Ours is a theory paper with small-scale experiments. We provide all the necessary details in the main text, Sect. 4
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Yes, we have reported error bars in our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
21

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Our experiments are small-scale and be reproduced standard laptop. Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Question:
Answer: [Yes]
Justification: Our work is more on the theoretical side and, thus, we do not see any direct societal impact or potential harmful consequences. The research process did not use any personal data and did not involve human subjects or participants.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work is more on the theoretical side and, thus, we do not see any direct societal impact.
Guidelines:
22

---

The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deploymentt of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (c.g., gated release models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not release any data or models that have risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: No specific assets used in work. We have cited the related work with code and data.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
23

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms use in the datasets.paperssiii paperswi thcode. www.com//aaaa datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not introduce any new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
We did not use crowdsourcing and did not make research with Human Subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We did make research with Human Subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
24

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
25


Reviews:
[{"review_id": "tcgTugrXtQ", "reviewer": "Reviewer_rKJj", "summary": "This paper proposes a novel gradient flow geometry (IFT), based on the infimal convolution of the Wasserstein tensor with the MMD tensor. For this geometry, the authors show global exponential convergence guarantees for both MMD and KL energies. They then develop an algorithm for the IFT gradient flow and test it on an MMD inference task, showing empirically that it avoids mode collapse as Arbel et al. does.", "questions": "1. Looking at the numerical example, have you tried any heuristic approaches, e.g. some sort of branching, for eliminating and repopulating particles when weights get very low on certain particles? It seems such behavior might improve performance empirically.", "limitations": "Limitations are well-acknowledged by the authors.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. The exposition clarity is excellent. The authors do a great job of positioning their work relative to existing gradient flow works.\n2. In introducing a novel gradient flow geometry and showing favorable convergence characteristics, the work has good potential to inspire follow-on works.", "weaknesses": "1. The experiments run are relatively simple and low-dimensional, and it is not clear how practical the method would be for more realistic application scenarios. \n2. There is no example comparison of behavior for KL, which would have been nice to see.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the fair assessment and constructive suggestions. We are glad that the exposition of the paper is easy to follow -- thank you for this feedback. We have incorporated some of your suggestions and left more challenging ones for future work.\n\n> experiments run are relatively simple and low-dimensional, and it is not clear how practical\n\nYour assessment is correct -- indeed our experiments are only a proof-of-concept. We have now included results in higher dimensions in the attached PDF. We also plan to investigate more realistic applications in the future.\n\n> numerical example, have you tried any heuristic approaches,\n\nWe appreciate your constructive suggestions. \n\nFirst, we stated at the end of Sec.5 (L317-319) that there are simple improvements that can be made to improve the algorithm, and we deliberately left them out to keep the paper focused; almost playing handicapped. The IFT still outperforms the pure Wasserstein flow.\n\nSecond, thank you for the valuable suggestions. We believe they can definitely help with practical performance. As we reported in the paper, we actually let IFT run with a handicap (i.e. every IFT iteration counts as two iterations). We did so to show that IFT already has an advantage without further heuristic tricks. We have focused on implementing the WFR flow during the rebuttal to provide another new algorithm for comparison. We will further investigate your suggestions in the near future and comment in our Sec.5 if any of the tricks help.\n\n> comparison of behavior for KL\n\nWe agree that more experiments can strengthen the paper. A direct comparison between energy MMD vs KL is difficult, as we explain below. However, we did run the newly implemented WFR flow and reported them in the attached PDF. \n\nFirst, We must note that KL minimization is NOT the topic of this paper: our focus is a type of (unbalanced-transport) gradient flows and optimization with applications to the MMD-minimization tasks. IFT is tailored to address the gap in the literature on \"MMD flow\". We also do **not** claim that MMD inference is superior to KL inference, while some other papers might do so. Furthermore, we look forward to applying our theory to generative models using MMD, e.g., [Galashov et al].\n\nNonetheless, we appreciate the reviewer raising this. We can think of two interpretations of the question:\n\n1. KL energy with the proposed IFT flow: \nRecall that the MMD step is a discretization of the differential equation\n$$\n\\dot \\mu =- \\beta\\cdot \\operatorname{\\mathcal{K}}^{-1}  \n\\left(\\log \\frac{\\mathrm{d}\\mu}{\\mathrm{d}\\pi} - \\frac{\\int \\operatorname{\\mathcal{K}}^{-1}\\log \\frac{\\mathrm{d}\\mu}{\\mathrm{d}\\pi}}{\\int \\operatorname{\\mathcal{K}}^{-1}1}\\right)\n$$\nOne could view the update step as a step in the \"kernel-mean-embedding space\": let $e(\\mu):=\\operatorname{\\mathcal{K}} \\mu$  be the kernel mean embedding of $\\mu$, then the update rule is\n$$\ne(\\mu)^{\\ell+1}\\gets e(\\mu)^{\\ell}  - \\eta \\beta \\cdot \\left(\\log \\frac{\\mathrm{d}\\mu^{\\ell}}{\\mathrm{d}\\pi} -...\\right)\n$$\nHowever, as there is no guarantee that the velocity $\\log \\frac{\\mathrm{d}\\mu^{\\ell}}{\\mathrm{d}\\pi}-...$ to be in the RKHS. Hence, this step is theoretically interesting, but it's unclear how the infinite-dimensional update can be implemented in a principled way. \n\n2. KL energy with the WFR flow: this case belongs to the \"KL-inference\" where one only has access to the score (or potentials) of the target e.g. $\\nabla \\log \\pi$, instead of the \"MMD-inference\" (our paper, and Arbel et al., etc.) where we have access to the target $\\pi$ via its samples but not the score $\\nabla \\log \\pi$. Therefore, how to perform a \"fair comparison\" is unclear to us at the moment (e.g. how many samples from $\\pi$ ? Noisy evaluation of the score?).\n\nIn summary, we agree that the reviewer's proposal is interesting and can see room for future work. At this moment, we can only leave comparing MMD vs KL as future work as they apply in different tasks and settings (score-based vs sample-based).\n\nWe have now implemented the **WFR flow** of MMD. It works only slightly worse than IFT but does not come with guarantees. See our overall rebuttal summary and the PDF.\n\n\n### Reference\nGalashov A, de Bortoli V, Gretton A. Deep MMD Gradient Flow without adversarial training. arXiv; 2024", "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": "pZOWcZZMtm", "reviewer": "Reviewer_2Ho7", "summary": "This paper proposes a novel gradient flow geometry – interaction-force transport (IFT). It is theoretically shown that IFT gradient flow has global exponential convergence guarantees both for MMD and KL energies. The authors propose an algorithm based on the JKO-splitting scheme and test it on examples with 2D Gaussians and Gaussian mixture.", "questions": "Does your approach have some practical use cases?\n\nI suggest the authors to improve the experimental part of their paper by\n- including experiments in dimensions larger than $d=2$\n- performing comparisons with flows using KL divergence as the energy, e.g., with (Yan et al., 2023, Lu et al., 2019) which are cited in the paper\n\nI am open to adjusting my score if the authors address these suggestions.\n\n**References.** \n\nM. Arbel, A. Korba, A. Salim, and A. Gretton. Maximum Mean Discrepancy Gradient Flow. arXiv:1906.04370, Dec. 2019.\n\nA. Korba, P.-C. Aubin-Frankowski, S. Majewski, and P. Ablin. Kernel Stein Discrepancy Descent. In Proceedings of the 38th International Conference on Machine Learning, pages 5719–5730. PMLR, July 2021.\n\nY. Yan, K. Wang, and P. Rigollet. Learning Gaussian Mixtures Using the Wasserstein-Fisher-Rao Gradient Flow. arXiv:2301.01766, Jan. 2023. \n\nY. Lu, J. Lu, and J. Nolen. Accelerating Langevin Sampling with Birth-death. ArXiv, May 2019.", "limitations": "The authors have addressed the limitations of their approach in the discussion section.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper provides the proof of the exponential convergence guarantees for IFT gradient flow with MMD energy.", "weaknesses": "I am not convinced that the established theoretical results and provided experimental justifications are sufficiently significant. The main theoretical contribution of the current paper is the proof of exponential convergence guarantees for their IFT gradient flow both with MMD and KL divergence energy. (Actually, the proofs of these results do not seem to be very impressive, e.g., the proof of Proposition 3.8 immediately follows from two well-known facts. Anyway, it is not my major concern.) The authors state that the established convergence guarantees (especially, for the MMD energy) are the main motivation for considering the IFT gradient flow.  However, for the KL case, it is not a surprising property since even ordinary Wasserstein flow with KL divergence energy has the same exponential convergence guarantees. For the MMD case, the results are quite novel, although there exist several other works which prove some convergence properties of flows with MMD energy (Arbel et al., 2019). Thus, I am wondering, are the provided proofs of exponential convergence rates actually important for the practical use of the designed algorithm? The empirical evaluation of the algorithms seems to be very limited. The algorithm is tested only in low-dimensional experiments using 2D Gaussians and Gaussian Mixtures which immediately raises questions regarding the scalability of the approach. Besides, the authors compare their approach only with Wasserstein flows with MMD energy (with or w/o noise injection) (Arbel et al., 2019, Korba et al., 2021). However, it is important to see how the algorithm behaves in comparison to flows with KL divergence energy as well. \n\nOverall, my main concerns are related to the questionable significance of paper results. The proof of exponential convergence rates for their IFT gradient flow solely does not seem to be a significant contribution. Meanwhile, the experimental evaluation needs to be considerably enhanced.\n\n*Minor*:\n\n- line 218: 'expontial' - typo", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the constructive suggestions and the critical review. We have incorporated most of your suggestions, e.g., new experiments. We did find the one comparison you suggested with KL inference difficult, for which we explained the reason.\n\n> including experiments in dimensions larger than 𝑑=2\n\nDone. The results (attached PDF) are consistent with the other experiments.\n\n> performing comparisons...  (Yan et al., 2023, Lu et al., 2019) \n\nWe agree that more experiments can strengthen the paper. We have run new experiments and reported them in the attached PDF, such as WFR flows used by Yan et al. and Lu et al.\n\nFirst, We must note that KL minimization is NOT the topic of this paper: our focus is a type of (unbalanced-transport) gradient flows and optimization with applications to the MMD-minimization tasks. IFT is tailored to address the gap in the literature on \"MMD flow\". We also do not claim that MMD inference is superior to KL inference, while some other papers might do so. Furthermore, we look forward to applying our theory to generative models using MMD, e.g., [Galashov et al].\n\nNonetheless, we appreciate the reviewer raising this. We can think of two interpretations of the question:\n\n1. KL energy with the proposed IFT flow: \nRecall that the MMD step is a discretization of the differential equation\n$$\n\\dot \\mu =- \\beta\\cdot \\operatorname{\\mathcal{K}}^{-1}  \n    \\left(\\log \\frac{\\mathrm{d}\\mu}{\\mathrm{d}\\pi} - ... \\right)\n$$\nOne could view the update step as a step in the \"kernel-mean-embedding space\": let $e(\\mu):=\\operatorname{\\mathcal{K}} \\mu$  be the kernel mean embedding\n$$\ne(\\mu)^{\\ell+1}\n\\gets\ne(\\mu)^{\\ell}  - \\eta \\beta \\cdot \\left(\\log \\frac{\\mathrm{d}\\mu^{\\ell}}{\\mathrm{d}\\pi} -...\\right)\n$$\nHowever, as there is no guarantee that the velocity $\\log \\frac{\\mathrm{d}\\mu^{\\ell}}{\\mathrm{d}\\pi}-...$ to be in the RKHS. Hence, this step is theoretically interesting, but it's unclear how the infinite-dimensional update can be implemented in a principled way. \n\n2. KL energy with the WFR flow [Yan & Lu et al.]: \nTheir methods are adaptations of the original Wasserstein-Fisher-Rao flow. Their cases belong to the \"KL-inference\" where one only has access to the score (or potentials) of the target e.g. $\\nabla \\log \\pi$, instead of the \"MMD-inference\" (our paper, and Arbel et al., etc.) where we have access to the target $\\pi$ via its samples but not the score $\\nabla \\log \\pi$. Therefore, how to perform a \"fair comparison\" is unclear to us at the moment (e.g. how many samples from $\\pi$ ? Noisy evaluation of the score?).\n\nIn summary, we agree that the reviewer's proposal is interesting and can see room for future work. At this moment, we can only leave comparing MMD vs KL as future work as they apply in different tasks and settings (score-based vs sample-based).\n\nWe have now implemented the **WFR flow** of MMD. It works only slightly worse than IFT but does not come with guarantees. See our overall rebuttal summary and the PDF.\n\n> Proposition 3.8 immediately follows from two well-known facts... for the KL case, it is not a surprising property ...\n\nWe have already clearly stated that it is standard (L229-230). Neither did we claim it is surprising. It is also immediately evident that KL is not the focus of this paper, as can be seen from the shortness of Sec.3.3. Hence, we believe our current presentation is not likely to cause confusion. We emphasize that the discussion of KL is for completeness and to show that IFT enjoys the *best of both worlds*. \n\nBonus: Prop. 3.8 is not as trivial as it seems -- we just fixed an error in the current revision: the LSI only holds along the mass-preserving spherical flows (SIFT) over $\\cal P$, but the flow over $\\cal M ^+$. This technicality does not change ML applications, especially with MMD. But anyway, thanks for helping us notice this.\n\n> For the MMD case, the results are quite novel... other works... (Arbel et al., 2019). Thus, I am wondering, are the provided proofs of exponential convergence rates actually important...?\n\nThank you for the positive assessment. We will add a small discussion in the revision as per your comment. A sketch:  (Arbel et al., 2019) does not contain (global) convergence analysis. For example, their Prop.2 states energy is non-increasing. However, this is not equivalent to convergence and is easily satisfied by many flows. The mathematical limitation to their \"MMD flow\" is that the MMD is in general not guaranteed to be convex along the Wasserstein geodesics. [Arbel et al.] also contains a heuristic noise injection procedure, albeit without good analysis characterization. In contrast, we believe establishing the first global (exponential) convergence in our paper is, needless to say, important. Not to mention our theorems/proof are clean and do not contain un-verifiable assumptions. Therefore, in our case, we have both practical performance and theoretical guarantees\n\n> Does your approach have some practical use cases?\n\nThis paper is mainly methodological and analytical, but we also value practical application. MMD flow applications in the literature mainly include, e.g.,\n\n- image process, as done by [Hertrich et al.] and that research group at TU Berlin\n- generative models (MMD GAN)\n- MMD two-sample test\n\nWe also see some new directions of new deep generative models that are based on gradient flows of the MMD e.g. in [Galashov et al.]\n\n> Strengths: The paper provides the proof of...\n\nWe emphasize that the proofs are a part of but not the whole of our contributions, which also include the discovery of the IFT **gradient structure**, the implementable algorithms, PDE theory of the dissipation mechanism of the MMD and inf-convolution with Wasserstein (e.g. Prop 3.2, Thm 3.6). The work also advances the state-of-the-art understanding of \"MMD gradient flow\", which already has sizable literature.\n\n## Reference\nGalashov A, de Bortoli V, Gretton A. Deep MMD Gradient Flow without adversarial training. arXiv; 2024", "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": "joBO9Aa9vt", "reviewer": "Reviewer_f71K", "summary": "This manuscript proposes a new gradient flow over probability and non-negative measures termed the interaction-force transport (IFT). The flow is based on the inf-convolution of the Wasserstein Riemannian metric tensor and the spherical maximum mean discrepancy (MMD) Riemannian metric tensor. The authors provide a number of theoretical results related to their proposal: global exponential convergence guarantees for the MMD and Kullback-Leibler divergence energies. While convergence results are available for the KL divergence energy, not much is known theoretically in the context of the MMD energy. The authors then introduce a particle gradient descent algorithm for IFT, composed of two steps gradient flows (one to update particle locations via the Wasserstein step and one to update particle weights via the MMD step), and show two proof-of-concept examples that validate the developed theory.", "questions": "Could the authors comment on the roughness of the standard deviation bands in Figure 2 for the proposed method? Also, while this limitation is already mentioned briefly in the discussion, I feel that a bit more could be said about scalability of the presented algorithm. This does not necessarily have to be included in the main body of the paper, but rather in the appendix where the algorithm is presented.", "limitations": "Limitations were adequately addressed.", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper is well-written and well-organized. Despite heavy notation throughout, the authors do a nice job introducing the notation and staying consistent throughout the manuscript. The literature review on related method is also quite thorough.\n2. In my opinion, this is a significant theoretical contribution to the machine/statistical learning literature. The topic is timely and will appeal to a broad NeurIPS audience.\n3. While the paper is theoretical in nature, I appreciate that the authors provide an implementation of the proposed gradient flows. The presented algorithm is easy to understand ensuring reproducibility.\n4. The contribution of this manuscript is original in many aspects: the new definition of gradient flows, the proofs of global convergence for the MMD and KL energies, and comparison to previously defined methods for the MMD flow that required a heuristic noise injection step.", "weaknesses": "I did not identify many weaknesses in this work. While I understand that the contributions are theoretical in nature, I wish the authors would have presented more examples/comparisons to the work of Arbel et al. [2019] in an appendix. The presented examples are sufficient as proof of concept.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to carefully read and review the paper. We appreciate that you have noticed the many subtle features we put into the manuscript. Thank you for the kind summary in the \"Strengths\" section. Below, we respond to a few concerns and suggestions.\n\n> the roughness of the standard deviation bands in Figure 2 for the proposed method?\n\nThank you for the careful read and for noticing the details.\n\nFirst, we believe the comment refers to the std band in Fig.2, which occurs around and below 10^{-3} magnitude. We suspect that could be due to that we reported in the y-axis in log scale. Then, it is expected the bands get rougher the lower they get. We observe that the same happens to the \"MMD flow+noise\" when its performance improves eventually (over more iterations); so it does not pertain to IFT.\n\n\n> this limitation is already mentioned briefly in the discussion, I feel that a bit more could be said about scalability of the presented algorithm. This does not necessarily have to be included in the main body of the paper, but rather in the appendix where the algorithm is presented.  \n\nIndeed, we agree that scalability is important and deserves a more detailed discussion, at least in the technical details in the appendix.\n\nFirst, we note that our current implementation only looks at the batch case. Here, IFT performs very well as in, e.g., Fig.2. As we have larger data sets in higher dimensions, the vanilla MMD might be limiting. \n\nAdditionally, to speed up the implementation and hence scalability, we have reported that we perform a single step of projected gradient descent instead of computing a full solution to the optimization problem (18). This is effective as (18) is a convex program. \n\nWe note that we have not yet explored stochastic algorithms; we have focused on full-batch so far. Hence stochastic algorithms and parallel implementation could be future directions to improve scalability.\n\nSince we are also interested in applications to generative models, there are possibilities such as deep-net features inside the kernels, as used in MMD-GAN works.\n\n\n> more examples/comparisons to the work of Arbel et al. [2019] in an appendix.\n\nThank you for the suggestion -- we agree. We have now added new experiments, whose results we report in the PDF:\n- experiments in higher dimensions, vs Arbel et al. 2019;\n- new experiments that uses the Wasserstein-Fiher-Rao flow of the MMD energy. To the best of our knowledge, this is the first implementation of the MMD energy in this flow;\n- many miscellaneous improvements in terms of scalability and implementation, as suggested by the review. We will provide the details in the appendix of the revised manuscript as you suggested.", "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": "hXPcJ3jyNl", "reviewer": "Reviewer_qhY4", "summary": "The paper proposes a gradient flow in combined Wasserstein-MMD geometry w.r.t. certain functionals. The authors primarily consider MMD squared functional, but also have some theoretical results regarding KL divergence functional. The work is more about theory: the authors are concerned about some mathematical properties of their proposed flows and convergence analysis, and have only toy 2D illustrative experiments.", "questions": "* (a) What are the properties of inverse operator $\\mathcal{K}^1$. In particular, why is it linear on its arguments?\n* (b) The MMD minimizing movement step (eq. 16 and eq. 18) is being solved inexactly in practice, e.g., only weights of particles are optimized, while the exact solution is the mixture of source and target distributions. Moreover, even eq. 18 is substituted with a single step of project GD. What is the reason? How do such approximations affect theoretical and practical properties of the proposed method?\n* (c) In the appendix, proofs. Why does the equality hold: $\\langle \\frac{\\delta F}{\\delta \\mu}[\\mu], \\mathcal{K}^{-1}\\frac{\\delta F}{\\delta \\mu}[\\mu] \\rangle_{L_\\mu^2} = \\Vert \\frac{\\delta F}{\\delta \\mu}[\\mu] \\Vert^2_{\\mathcal{H}}$?", "limitations": "The limitations were addressed correctly", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "* Overall, the considered topic is quite interesting. The theory of Gradient Flows in different geometries is an emergent field which is at the intersection of Machine (Deep) learning and mathematics. This theory is full of remarkable, non-trivial, theoretical results. Transmitting all of this mathematical beauty into practical algorithms is praiseworthy.\n* The paper has some interesting theoretical results and statements.", "weaknesses": "* (A) At first, I found the manuscript to be a bit difficult to read, especially section 2. A lot of specific mathematical terms were used, e.g., “Onsager operator”; tangent/cotangent spaces and metric tensor of (probability) measures space. A lot of relationships between these specific objects were mentioned, e.g., $\\mathbb{K} = \\mathbb{G}^{-1}$; formulation of gradient flow through the Onsager operator (eq. 2). I think that in order to make the text more accessible for those who are not a specialist in geometry of (probability) measure spaces, it should be either simplified, or all necessary theoretical introductions should be done, e.g., in the appendix.\n* (B) I am not fully satisfied with the structure of the text. In particular:\n    * Why Remark 3.4 and Corollary 3.5. (some properties of pure IFT gradflow which was introduced much earlier) are located right after technical Theorem 3.3 (Lojasiewicz inequalities)? I think it is better to place these statements right after Remark 3.1.\n    * For me, it is a bit strange that the paper develops theory of spherical IFT gradflows (Section 3.1.), while the only practically considered case (where the driving functional for the gradflow is MMD) does not require this theory (Theorem 3.6) because spherical MMD (IFT) coincides with conventional MMD (IFT) flow. May be more emphasize (inculuding practical evaluations) should be put on KL-driving gradflows, where the sphericity matters.\n* (C) (lines 98-99). Machine learning applications of Wasserstein gradient flows: some missed links: [1-6]\n* (D) To be honest, I am a bit skeptical about the pure MMD gradient flow (in the MMD geometry) - which is denoted as (MMD-MMD-GF) in Theorem 3.6, and, correspondingly, my skepticism extends to MMD gradflow in the joined Wasserstein-MMD geometry - (MMD-IFT-GF) in Theorem 3.6. At first, (MMD-MMD-GF) was considered in literature, e.g., [7] (not cited!) - see Section 3.1, Case 1 of their paper. And it was noted that such pure MMD-MMD flow is undesirable in practice, exactly because it “teleports” mass between initial and target distribution (note that the solution to MMD-MMD is just interpolation between distributions, as noted by Theorem 3.6). As I understand, the idea of the paper under consideration is that by considering joined Wasserstein-MMD geometry (MMD-IFT flow) one can leverage this problem. However, in the paper, I didn’t find sufficient evidence that it is the case. In particular, the proposed practical optimization procedure includes solving proximal MMD minimizing-movement step, eq. 16. For $F = \\text{MMD}$ it boils down to eq. (18), which is MMD barycenter problem. It is known that MMD barycenter problem has solution, see [8, proposition 2], - it is just a mixture of input distributions. Therefore, if we fairly solve MMD minimizing movement step, the resulting $\\mu^{\\ell + 1}$ will mix $\\mu^{\\ell + \\frac{1}{2}}$ and target $\\pi$, i.e., the teleporting of mass will occur.\n* (E) The practical validation of the method is rather weak. Only a couple of 2D experiments with Gaussians/Mixture of Gaussians. Moreover, I didn’t find that the proposed method performs better than the alternatives. May be, according to some metrics it is indeed the case, but the visual performance of the method is somewhat disappointing. As I understand from Figures 3, 4 and gifs provided in the supplementary, the method leaves a considerable number of points far from the support of target distribution. The alternatives, even vanilla MMD flow, are better in terms of this characteristic.\n\n[1] Gao et. al., Deep Generative Learning via Variational Gradient Flow, ICML’2019\n\n[2] Gao et. al, Deep Generative Learning via Euler Particle Transport, MSML’2021\n\n[3] Mokrov et. al., Large-Scale Wasserstein Gradient Flows, NeurIPS’2021\n\n[4] Alvarez-Melis et. al., Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks, TMLR’2022\n\n[5] Bunne et. al., Proximal Optimal Transport Modeling of Population Dynamics, AISTATS’2022\n\n[6] Fan et. al., Variational Wasserstein gradient flow, ICML’2022\n\n[7] Mroueh et. al., Sobolev Descent, AISTATS’2019\n\n[8] Cohen et. al., Estimating Barycenters of Measures in High Dimensions", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for writing a long and critical review. We must point out that the review is filled with misunderstandings, which we do our best to clarify below. We would appreciate it if the reviewer could please consider our clarification. \n\n>  This theory is full of remarkable, non-trivial, theoretical results. Transmitting all of this mathematical beauty into practical algorithms is praiseworthy.\n\nThank you. But such a positive assessment does **not** match your score.\n\n> (E) The practical validation is rather weak... ... visual performance of the method is somewhat disappointing... the method leaves a considerable number of points far from the support of target distribution... vanilla MMD flow, are better\n\nWe are deeply confused by this assessment. The only explanation we can think of is that the reviewer completely misunderstood the point of our experiment and intuition of unbalanced transport. We have clearly stated **twice** in the manuscript (Fig.3 & 4 captions),  that **color intensity indicates points’ weights** and mention that **weights vanish**. The \"points far from the support\" the reviewer was referring to are likely (since we are confused by the review) the points of zero mass, hence the reviewer's comment is likely due to a major misunderstanding. Nonetheless, we will improve the presentation further to avoid misunderstandings.\nWe will explicitly say: **The hollow/white circles indicate the particles that have already vanished**. \n\nPlease note: as the paper indicates, the performance of the IFT is overwhelmingly good. Hence the review comment does not make sense. \n\n> only a couple of 2D experiments \n\nWe have now included results in higher dimensions in the attached PDF. \n\n> strange that the paper develops theory of spherical IFT gradflows.. does not require this theory (Theorem 3.6) because spherical MMD (IFT) coincides with conventional MMD (IFT) flow\n\nThe reasoning seems flawed here. Our results such as Thm 3.6, precisely establish this special-case equivalence of SIFT and allow an easy implementation. Without our theory, there is no literature on the PDE of a mass-preserving flow in the MMD/IFT, then no implementation is possible. \n\nThe reviewer claimed our theory is \"not needed\" after she/he read our theory. Without this result, one cannot implement the mass-preserving flow. Furthermore, rigorously speaking, it is also incorrect to say two gradient flows are equivalent. Hence, the logic of the claim that our result is ”not needed” is not sound\n\n> [7] (not cited!) \n\nWe have now cited a subsequent/expanded work by the same group [Y. Mroueh and M. Rigotti] in the revised preprint. This paper covers and overlaps with the older paper you mentioned, we believe the coverage is now enough.\n\nFurthermore, the paper you mentioned did not contain the PDE theory and principled gradient structure we uncovered. Hence, although related, it should not be used as an argument to undermine our contribution.\n\n> my skepticism extends to MMD gradflow in the joined Wasserstein-MMD geometry\n\nUnfortunately, we failed to find sound reasoning for this \"extension\" of skepticism. We never claimed MMD-Mmd-Gf (which is not our focus) is superior. IFT is what we advocate and the Wasserstein contribution made the difference. The review simply cited an old paper related to but didn't theoretically study **MMD-MMD-GF**, let alone IFT, and claimed this \"extends\" their \"skepticism\" to IFT. We do not see any sound reasoning here.\n\n> Therefore, if we fairly solve MMD minimizing movement step, the resulting μℓ+1 will mix μℓ+12 and target π, i.e., the teleporting of mass will occur.\n\nSorry, we did not get what is the question or concern here. The statement is straightforward from our PDE/ODE characterization in Thm 3.6.\n\n> (a) What are the properties of inverse operator 𝐾1.  In particular, why is it linear \n\nDo we understand it correctly, that the reviewer is asking why is the inverse of a linear operator, linear? Please see also the first paragraph in Sec 2.2 for properties of the integral operator and references that can provide more basics of kernel methods and the integral operator. Or we didn't understand your question correctly?\n\n\n> (c) In the appendix, proofs. Why does the equality hold:\n\nFirst, there is a typo that the inner products should be in unweighted L2 (thank you for helping us fix this), i.e.,\n$$\n\\langle\n       \\frac{\\delta F}{\\delta \\mu}\\left[\\mu\\right]\n       ,\\mathcal K^{-1}\n       \\frac{\\delta F}{\\delta \\mu}\\left[\\mu\\right]\n            \\rangle_{L^2}\n            \\\\\n            =\n \\| \\frac{\\delta F}{\\delta \\mu}\\left[\\mu\\right]\\|^2_{\\cal H}\n$$\nThis correct relation follows from the textbook definition of the RKHS (and its norm). See standard text such as\nCucker, F., & Zhou, D. X. (2007). _Learning theory: an approximation theory viewpoint_ (Vol. 24). \n\n> substituted with a single step of project GD. What is the reason? How do such approximations affect theoretical and practical properties of the proposed method?\n \n\nIn optimization, it is often desirable and significantly faster to perform inexact iterations to be more computationally efficient. That is the reason of our implementation. As our paper does not focus on the theory of the exact-inexact relation, hence it does not affect the existing analysis of the paper. In practice, we have observed such changes do not affect the performance significantly but improve the speed. Since eq (18) is a convex program, inexact iterations are preferred.\n\n\n> difficult to read...A lot of specific mathematical terms were used\n\nFirst, we will consider your suggestions.\n\nOur paper is a mathematically rigorous treatment, hence some mathematical terms are necessary. Plus, we have already opted for a minimal set of common terminologies for ML researchers on the topic. It also seems the other reviewers did not experience the same difficulty. Nonetheless, we will try to make the paper more accessible to non-experts (than it already is).", "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:
Dear reviewers, dear AC,

We would like to thank all reviewers for their constructive feedback. We are glad that the majority of the reviewers found our paper easy to read and contains non-trivial contributions.

The major concern expressed by some reviewers is that more experiments would strengthen the paper -- which we have now done. 

Based on the reviewers' suggestions, we have included new experiments. The figures are reported in the attached PDF. For example,

1. experiments in higher dimensions $d>2$;
2. a new Wasserstein-Fisher-Rao gradient flow of the MMD energy; see the appendix below for details. This is the first implementation of the MMD energy in this flow to the best of our knowledge. Note that this is the same algorithm used in [Yan et al. & Lu et al.], as requested by one reviewer. As we are not in the score-based sampling regime, we have to use the MMD energy instead of the KL energy (mentioned in some reviews);
3. many other improvements in terms of scalability and practicality, as suggested by the reviewers.

Therefore, we believe the reviewer's concerns have been addressed. We hope the reviewers will take our new results and improvement into consideration. If there are any further concerns, we are happy to discuss them during the discussion phase.

Thank you for your time and consideration,
Authors



## Appendix: Details on the newly implemented WFR flows for comparison (i.e., KL step for reweighting)

In the second paragraph in Sec.3.4 (L248-256), we discussed the comparison with the Wasserstein-Fisher-Rao (WFR) flow of the MMD. Note that there are no sound convergence guarantees for this scheme yet. To put things in perspective, let us make it clear: Yan et al. and Lu et al. use variants of WFR flows, but with KL energy. We can not directly compare with them as we are in the sample-based regime (only have access to samples of $\pi$) and they are in the score-based regime (only have access to the score $\nabla \log \pi$).

Nonetheless, per the reviewer's suggestion to compare with them, we have implemented the **WFR flow of the MMD energy**. Note that this is the first implementation of the MMD energy in this flow, to the best of our knowledge. We believe this has added value to the paper and to some extent addresses the reviewer's demand for more empirical comparisons. This amounts to the JKO splitting steps
$$
            \mu^{\ell+\frac12}
            \gets\arg\min_{\mu\in\cal P} F(\mu ) + \frac1{2\tau}W_2^2(\mu, \mu^\ell)
            \textrm{(Wasserstein step)}
  $$
$$
            \mu^{\ell+1}
            \gets\arg\min_{\mu\in\cal P} F(\mu ) + \frac1{\eta}{\mathrm{KL}}(\mu, \mu^{\ell+\frac12})
             \textrm{(KL step)}
$$
As is in the Wasserstein step, in practice, we use the
explicit Euler scheme which boils down to the **entropic mirror descent**.
As is well-known, especially in the optimization literature, the entropic mirror descent step can be implemented as multiplicative update of the weights (or density), i.e.,
suppose $x_i^{\ell+1}$ is the new particle location after the Wasserstein step, then we update the weights vector $\alpha$ via
$$
            \alpha_i ^{\ell+1}
            \gets \alpha_i  ^{\ell} \cdot \exp
            \left(
                -\eta \cdot 
                \frac{\delta F}{\delta \mu}[\mu^\ell] (x^{\ell+1}_i)
            \right)
$$
where the Riemannian velocity $\frac{\delta F}{\delta \mu}[\mu^\ell]$ is the same as used in the Wasserstein step above. For the MMD, it is given by (already provided in the manuscript)
$$\frac{\delta F}{\delta \mu}[\mu^\ell]=\int 
k(x, \cdot )
(\mu^\ell -\pi  )(\mathrm d x)$$

We have added more details to the appendix of the revised manuscript. Numerical results are reported in the PDF below.

Meta-review: This manuscript proposes a new gradient flow over probability and non-negative measures termed the interaction-force transport (IFT). The flow is based on the inf-convolution of the Wasserstein Riemannian metric tensor and the spherical maximum mean discrepancy (MMD) Riemannian metric tensor. The authors provide a number of theoretical results related to their proposal: global exponential convergence guarantees for the MMD and Kullback-Leibler divergence energies. While convergence results are available for the KL divergence energy, not much is known theoretically in the context of the MMD energy. The authors then introduce a particle gradient descent algorithm for IFT, composed of two steps gradient flows (one to update particle locations via the Wasserstein step and one to update particle weights via the MMD step), and show two proof-of-concept examples with 2D Gaussians and Gaussian mixture that validate the developed theory. 
In the rebuttal, the authors conducted higher dimensional experiments and implemented WFT flow for comparison, which enhance the work.

The paper is well-written and well-organized. Despite heavy notation throughout, the authors do a nice job introducing the notation and staying consistent throughout the manuscript. The literature review on related method is also quite thorough. This work has novelty and significant theoretical contributions to the machine/statistical learning literature:  the new definition of gradient flows, the proofs of global convergence for the MMD and KL energies, and comparison to previously defined methods for the MMD flow that required a heuristic noise injection step. While the paper is theoretical in nature, the authors also provide an implementation of the proposed gradient flows. The presented algorithm is easy to understand ensuring reproducibility. The topic is timely and will appeal to a broad NeurIPS audience.

The work has great theoretic contribution and has the potential to make impacts in the field. I recommend the acceptance.
Decision: Accept (poster)