Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Authors: Clément Bonet
Christophe Vauthier
Anna Korba
Venue: ICML 2025 oral (2025)
Keywords: Wasserstein gradient flows
optimal transport
datasets
Primary area: probabilistic_methods->everything_else
TL;DR: We flow datasets using Wasserstein over Wasserstein gradient flows.

Abstract: Many applications in machine learning involve data represented as probability distributions. The emergence of such data requires radically novel techniques to design tractable gradient flows on probability distributions over this type of (infinite-dimensional) objects. For instance, being able to flow labeled datasets is a core task for applications ranging from domain adaptation to transfer learning or dataset distillation. In this setting, we propose to represent each class by the associated conditional distribution of features, and to model the dataset as a mixture distribution supported on these classes (which are themselves probability distributions), meaning that labeled datasets can be seen as probability distributions over probability distributions. We endow this space with a metric structure from optimal transport, namely the Wasserstein over Wasserstein (WoW) distance, derive a differential structure on this space, and define WoW gradient flows. The latter enables to design dynamics over this space that decrease a given objective functional. We apply our framework to transfer learning and dataset distillation tasks, leveraging our gradient flow construction as well as novel tractable functionals that take the form of Maximum Mean Discrepancies with Sliced-Wasserstein based kernels between probability distributions.

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Clément Bonet Christophe Vauthier *2
Anna Korba
Abstract
2023; Haviv et al., 2024b). Remarkably, it has been shown that one can embed any finite dataset with little or distor- tion (Andoni et al., 2018; 88; Kratsios et al., 2023) in the Wasser- stein space, i.e., the space of probability distributions (e.g., over a Fuclidean space) equipped with the Wasserstein-22 dis- tance from Optimal Transport (OT). This has motivated the use of this space to embed many types of data ranging from words (Vilnis & McCallum, 2015) to knowledge graphs (He et al., 2015; Wang et al., 2022), graphs (Bojchevski & Günnemann, 2018; Petric Maretic et al., 2019), or neuro- science data (Bonet et al., 2023). Therefore, it is essential to develop tools to work on the space of probability measures over probability measures, also known as random measures. In particular, they provide a natural way to represent labeled datasets as mixtures (Alvarez-Melis & Fusi, 2020).
Many applications in machine learning involve data represented as probability distributions. The emergence of such data requires radically novel techniques to design tractable gradient flows on probability distributions over this type of (infinite- dimensional) objects. For instance, being able to flow labeled datasets is a core task for applica- tions ranging from domain adaptation to transfer learning or dataset distillation. In this setting, we propose to represent each class by the associated conditional distribution of features, and to model the dataset as mixture distribution supported on these classes (which are themselves probability distributions), meaning that labeled datasets can be seen as probability distributions over proba- bility distributions. We endow this space with a metric structure from optimal transport, namely the Wasserstein over Wasserstein (WoW) distance, derive differential structure on this space, and define WoW gradient flows. The latter enables to design dynamics over this space that decrease a given objective functional. We apply our frame- work transfer learning and dataset distillation tasks, leveraging our gradient flow construction well as novel tractable functionals that take the form of Maximum Mean Discrepancies with Sliced- Wasserstein based kernels between proba- bility distributions.
A natural distance on this space is the Wasserstein over Wasserstein distance (WoW) (Nguyen, 2016; Catalano & Lavenant, 2024), also known as the Hierarchical OT distance, which lifts the Wasserstein distance between prob- ability distributions as a ground cost, to define a Wasserstein distance between random measures. The latter has been used for generative modeling applications (Dukler al., 2019), domain adaptation tasks (El Hamri et al., 2022), comparing documents (Yurochkin et al., 2019) or multilevel clustering (Ho et al., 2017). It has also been used to compare Gaussian mixtures (Chen et al., 2018; Delon & Desolneux, 2020; Wilson et al., 2024) or generic mixtures (Dusson et al., 2023; Chen Zhang, 2024). However, its poor sample complexity has motivated the development of alternative distance mea- sures, such those based Integral Probability Metrics (Catalano & Lavenant, 2024). Nonetheless, this space pos- sesses rich Riemannian structure, enabling the definition of concepts like geodesics. This has been leveraged recently by Haviv et al. (2024a) to perform generative modeling over the space of probability distributions with Flow Matchings.
1. Introduction
Probability measures provide powerful way to represent many data types. For instance, they allow to naturally repre- sent documents (Kusner et al., 2015), genes (Bellazzi et al., 2021), point clouds (Qi et al., 2017; Geuter et al., 2025), images (Sodini et al., 2025), or single-cell data (Persad et al.,
While this space naturally supports a range of machine learn- ing tasks, optimization methods tailored to it have received limited attention. Yet, this is important for multiple ap- plications, including variational inference with a Gaussian mixture family (Lambert et al., 2022; Huix et al., 2024), computing barycenters (Delon & Desolneux, 2020), or flow- ing datasets (Alvarez-Melis & Fusi, 2021), e.g., for domain adaptation, transfer learning (Alvarez-Melis & Fusi, 2021; Hua et al., 2023) or dataset distillation (Wang et al., 2018).
Eqqual contribution 'ENSAE, CREST, IP Paris Uniierrsité Paris-Saclay, Laboratoire de Mathématique d' Orsay. Correspon- dence to: Clément Bonet <element. bonet ensae.fr>, .f>>, Chrisophe Vauthier <christophe, vauthier universite- -arit--aa-ssiieaee paris- saclay.fr>.
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
In this paper, we propose to leverage the Riemannian struc- ture of random measures equipped with the WoW distance, by defining and simulating gradient flows, i.e. paths of ran- dom measures that follow the steepest descent of a given objective functional.
mixtures of probability distributions corresponding the class-conditional distributions. We focus on image datasets, and show that the flow enables structured transitions of classes toward other classes, with applications to transfer learning and dataset distillation.
Notations. For a Riemannian manifold M, d .M x M - R4 is its geodesic distance. For x E M,, we denote by T.M the tangent space at x, and by || IIæ the Rieman- nian metrie. We define by T.M {(1,, u), 2 E M and T..M) the tangent bundle. We define for (2, E TM the projections T.M (z, v) = x and T T' v) = v. exp TM + M is the exponential map. For x E M. if exp, TrM- + .M is invertible, we note log,, its inverse. V and div refer to the Riemannian gradient and divergence on M. For metric space (X,d), P2(X) denotes the space of probability distributions ments, i.c., P2(X) {u E P(X), I d(x, dd((x( < oo} with E X some arbitrary origin. For any E P2(M), (M, TM) is the set of functions V M - TM < oo. For a measurable map : + M, we note by TH the pushforward measure. Id denotes the identity map on M. P2.sc C (2(M) is the space of measures absolutely continuous w.r.f. the vol- ume measure on M. For μ4, E P(X), ), we denote μ < v if μ is absolutely continuous w.r.f. V. [(H,D) = {7 P(X X), T μr, T 7 = n} with x' (#1, + x₁, is the set of couplings, and TI (12, v/ the set of optimal couplings.
Related works. An elegant and popular way to perform optimization over probability distributions (over manifold) is to leverage the Riemannian structure of the Wasserstein space (Otto, 2001), and to use Wasserstein gradient flows (Ambrosio et al., 2008; Santambrogio, 2017). Several time discretizations of these flows have been studied (Jordan et al., 1998; Salim et al., 2020; Bonet et al., 2024), and they have been applied simulate the flow dynamics of multiple objectives such as the Kullback-Leibler divergence (Wibisono, 2018; Salim et al., 2020; Diao et al., 2023), the Maximum Mean Discrepancy (MMD) (Arbel et al., 2019; Altekrüger et al., 2023; Hertrich et al., 2024a;b) and variants thereof (Glaser et al., 2021: Chen et al., 2024; Neumayer et al., 2024; Chazal et al., 2024) or the Sliced-Waa Wasserstein distance (Liutkus et al., 2019; Du et al., 2023; Bonet al., 2025). Yet, all these works focus on the case where the prob- ability distributions are defined over mnite-dimensional manifold, e.g. Rd In practice, simulating these flows often boils down to simulating a particle system in Rd Hence, these works do not address probability distributions defined on infiiit-dimensionn spaces, such as the space of probability measures, which is the focus of this work.
X with second finite mo-
such that f [[](()))
The closest works to ours are the ones of Alvarez-Melis & Fusi (2021) and Hua et al. (2023). These papers cast labeled datasets as measures over a product space of the features and the conditional distributions (i.e., the distributions of the features of a given class). However, they circumvent the issue of designing gradient flows on this space by modeling the conditional probabilities as Gaussian distributions, hence parametrized by mean and covariance, which are finite- dimensional objects. While this enables them to leverage standard Wasserstein gradient flows, this Gaussian modeling of mixture components is a strong assumption that may not capture the true shape of many labeled datasets in practice.
2. Background
We begin by introducing some background on Optimal Transport (OT) and on Wasserstein Gradient Flows. For the- oretical purposes, we provide background on the geometry of (P2(M), W2) with M a Riemannian manifold, as in the next section, we will rely on results which hold on compact Riemannian manifolds (without boundary). Nonetheless, the applications will be done for M = Rd. The reader may refer to Appendix A for more details.
Optimal Transport The Wasserstein distance between 1, E P2(M) is defined as
Contributions. In this work, we introduce a principled framework for optimizing functionals over the space of probability measures on probability measures, leveraging the Riemannian structure of this space to develop Wasser- stein over Wasserstein (WoW) gradient flows. We provide a theoretical construction of the flows, and then a practical implementation through time discretization using forward Fuler scheme. We also propose novel functional objective, that writes as an MMD with kernel between distributions based on the Sliced- Wasserstein distance, and whose gradi- ent flow simulation is tractable. We then apply this scheme to flow datasets viewed as random measures; specifically, as
/ | y)² dj(
Wz(m = inf gETT(a.))
(1)
The metric space (P2(M), W2) has Riemannian structure (Otto, 2001; Erbar, 2010). In particular, if the log map is well defined p-almost everywhere (a.c.), (constant-speed) geodesics between u,v v are defined as Pe = cxxp,, a log,1 )π²)) # 7 with ý E Io v) an op- timal coupling. If μ E (M), there is map T, namely the OT map, such that Tæ = V and 7 = (Id, T) n by McCann's 's theorem for a wide range of manifolds (McCann, 2001; Figalli, 2007). In particular, T exPra o( V
2

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
with Pu, Kantorovich potential between j and v, and geodesics become 14 = exPra o( o() At any μ E (M),, we can define a tangent space P2(2
If such a gradient exists, then we say that F is Wasser- stein differentiable at j.. There is unique gradient be- longing to T, P2MM) and we restrict to this gradient. In- formally, the Wasserstein gradient of F can be computed as F(r) = V& (n), with & öy (u) the first variation defined, when it exists, as the unique function (up to an additive constant) such that, for x satisfying I dx = O, di F(u + +xxx= t= = -J oz (1) dx (Ambrosio et al., 2008, Lemma 10.4.1). Examples of differentiable functionals in- clude potential energies V(n)= = f VVd and interaction en- ergies W(H) = ff (1x,, y))(, for V :M- + R and W : :Mx x M + IR twice differentiable with bounded Hessian, for which Vw, V(n)= = VV and Vw, (z)) = J W y) V.W(y, x))dd(()) Moreover, if the func- tional F.. P2(RR - IR has a closed-form over dis- crete measures, i.e., there exists F (RR)) + R such that F 6z.) F(z1,...., Zn), then backpropagation on F and find the Wasserstein gradi- ent of F using the relation F(, W oe,) = n FF, ..,) (see Proposition A.9).
[Ve, E CMMM) (M with (M) the space of smooth compactly supported functions on M (Erbar, 2010; Gigli, 2011). This is Hilbert space endowed with the (2,,TMM inner product. The exponential map on P2(M) is then defined as exp,(() = (exPra ov) # for μ E P2(M), u T (MM). For instance, when M Rd with d(z, ||2 21.. then exp,(u) = + y and (v) y x for all x. E Rd.
Let P2(TM) = ( E P(TM), [(d(z, o)² + (ly))))d)(zz v) < bo} where o E M is any reference point. Following (Gigli, 2011), we define for every u, v E (M(M.,
(x) = E Pa(MM)) 7 11, exp# 7 = 1/
use
/ [1u1Bd7(z = W3(4, })} (2)
the set of plans 7 E P2(TM) such that (T-M, 7x is an OT plan between j and v. This allows one to avoid using the logarithm map, which might not be well defined everywhere, e.g. being multivalued. This space carries more information than the set of optimal couplings as it precises which geodesic was chosen to move the mass, as M4 (exPaM o(tr")) 7 are constant speed geodesics between and i (Gigli, 201 Theorem 1.11). On P2(R)). this trans- lates as exp,, (v) = {(x) nT2 - x +), 5 E [o, ))} (Gigli, 2004; Hertrich et al., 2024a). We show in the next proposition, whose proof can be found in Appendix C.I, that can build surjective map from exp,-'(v) to [o(H,, v).
A Wasserstein gradient flow of a differentiable functional JF is a curve t + ir which continuity equation O.juc = div (pr F(4)) A possi- ble discretization is the Riemannian Wasserstein gradient descent (Bonnabel, 2013; Bonet et al., 2025), defined as Pk+1 = exP1a - T -V,,F(mx) u For discrete distri- butions us = n Ex_1 duy. it translates as, for all 2 O, For M = Rd. this is simply 2441 = af F(u)) which corresponds to Wasserstein gradient descent.
(weak) solution of the
E (1, (l,..., n), x
exp25 ( V-.
Proposition 2.1. Let p,v E (MM). A surjective map from exp, (v) to [u, is given by y ++ (TMM. (TM,)x In particular, expr (v) is not empty, and i E exp, (w), then d(zz, (v)
3. Wasserstein over Wasserstein Space
We introduce in this section the Wasserstein over Wasser- stein space (P)))) Ww.). i.e. the space of probability distributions over probability distributions P2 (P2(M)), endowed with the OT distance with the squared Wasserstein distance on P2(M) as groundcost. We first state some properties of this distance, and then introduce a differential structure on this space which will be used in the next sections to develop suitable optimization methods. In the following, M is compact and connected manifold. The proofs can be found in Appendix C.
[[lla y-a.e. (2, T.M.
Additionally, if M is compact and connected, and μ E (M), then there exists unique y E exxp (v), of the form 7 = (Id, p
Wasserstein Gradient Flows. Let :P(MM + R be a lower semi-continuous functional. We briefly introduce the differential structure on (P2(M), W2), i.e. probabil- ity measures on manifolds, inspired by (Erbar, 2010) and (Lanzetti et al., 2025).
3.1. OT Distance and Riemannian Structure
Let μ E (2) We say that F(,) E T,, TMM is a Wasserstein gradient of F at if for any v E (M( and any 7 E cxp,, (v), we have the Taylor expansion
The WoW distance is defined as the OT problem with the squared Wasserstein distance on P2(M) as groundcost, i.e. for P, Q E P2 (M)) inf
/ (W(LLp) dT(,))). (4)
Ww, (P, Q) rell(,.))
= [W)=)(m)> + (V.: 1) d)(a,))
This defines a distance (Nguyen, 2016). Analogously to P2(M) and Brenier-MeCann" 's theorem, it has been shown
+ v))
(3)
3

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
that there exists an OT map from P to under absolute continuity of IP with respect to suitable reference measure Po E Pa (P2()) (Emami & Pass, 2025), which has no atom and satisfies an integration by part formula (Dello Schi- avo, 2020). We refer to Appendix B for more details.
where the I in the sup are selected in expp (Q). Simi- larly, & 1. (PM))) belongs to the super-differential θ F(P) F at if - E (-F))().
If the functional admits a sub- and super-differential, which coincide, we can define a gradient.
Now, let us denote for any Y E P2(TMM, the projections d " (+) exp# y and (7)) = m 7-
Definition 3.3. F is Wasserstein differentiable P (2(P2(M)) if F+F(() n a ((P) o 0. In this case, is WoW gradient of IF at
For any P E P2 (Pa(M)) let us also define
say that E (P) P. Pndnd it satisfies for any Q E Pa (P2()), I E expp (Q),
expp (Q) = {T (2(PP(TTM)) ## TT= = P, @# T Q2
/f |111/d u) (7) ww, (P, Q)) (5)
F(Q) = F(P) + //
M (1), w) (4(m)
dr v)) + (W (P, Q))
Relying on Proposition 2.1, we can define for any P. E
(7)
(M)) a surjective map from expp (Q) to (P, Q)-
In the following, we note Vww, (PP) such a gradient.
Proposition 3.1. Ler P,Q E (P2( Then, [ +
(o-M, (T is a surjective map from (Q) to
We can also define notion of strong sub- and super- differential, as well as gradient, by allowing the coupling [ to be non-optimal, in contrast with the previous definitions. Definition 3.4. & E L2 (P2))) is a strong subd- ifferential of F at PP if for all Q E (P))) for all I Te (P s.t. MT P. o#
(IP, Q). In particular, expp'(Q) is not empty and if TE TExxpp (Q), 7 E exp_ -1,,,, pa bb (oxP* # 7) for Ta..e y.
Additionally, if P < Po. there exists unique U E expp (Q), (Q)) of the form (μ + (Id, - (Vp.T(a))=)) with T the unique transport map from P to Q and Pu,T(n) Kantorovich potential between μ. T(u) E (MM).
exp Q.
F(Q) 2 F(P)+ + l/ (2)), o) dr (7)
The proof of Proposition 3.1 can be found in Appendix C.2. The previous construction enables us to formalize the Rie- mannian structure of (P2 (P)(). Ww,). ), without having to define notion of logarithm map on P2(M), which might be ill-defined when the OT plan is not unique. Between P, E P2 (P2())). we can define for IT E (Q) a geodesic t + Pt- = expAM otto")+T, which satisfies for all 8, E [0, 11, Ww, (Pa, Pt) = -t- - 8 WWW, (P, Q), see Appendix B. For F < Po. using Proposition 3.1,
o
| (1
(8)
Strong superdifferentials and gradients are defined similarly.
The latter definition is particularly useful when perturbing a measure along a non-optimal direction, as in the case of the forward Euler schemes we will compute in the next section.
the curve simplifies Moreover, for M- Rd, this reads as Pt (# (Id (7))))))))))) P.
PP, exPra o(- +
We now turn to examples of functional on P2 (P2 that take the form of free energies. Given F: P2(M) - R, we define potential energy V : (P((P( as V(P) = f F(H)dP(p)) Analogously to classical Wasserstein gradients, its WoW gradient is obtained as Vww, ((P)(H) = F(,F(H). Given kernel W Pa(Rd)) x P2(Rd)) + R. we define interaction Sf v) dP(H)dP(M). and
R
3.2. Differential Structure
We now provide a differential structure to (P2 Ww,). following the one of (Ambro- sio et al., 2008; Erbar. 2010; Lanzetti et al., 2025) for (P2(M), W2). In this section, let F (P))) + R be lower semi-continuous functional. We define formally the Hilbert space
energies as W(P) their WoW gradients are obtained as Vww, ((P)()) = + Vw, (W(,,, dP(c).. We refer to Appendix B.4 for more details.
ʃ (Vw,, W(H,)) (V,,,,((,,,
of functions from (2(
to (MM in Appendix B. First, we define the notions of (extended) sub- and
Let us now define cylinder functions, which provide a class of Wasserstein differentiable functionals (von Renesse & Sturm, 2009; Dello Schiavo, 2020; Fornasier et al., 2023). Definition 3.5. functional F: :P2(M) - R is a cylinder if there exists k 2 O, F C. (R+) and Vi,..., Vvk C (M) such that, for all μ E (MM
Definition 3.2. & E (P( (P)( belongs to the sub- differential 0- (PP) of IF at if for all Q e (M)))
F(Q)= +(p) /| (6( (((()))(0)))))))) v)2 dr ())
F()) = F (Svidu.... /Vidu).
+ +o(Ww, Q)), (6)
(9)
4

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
In this case, note F E Cy)(P2(M)). Similarly, for an interval, we note E Cyl(T x Pa(M)). if F(t, p)) = F (t, Vidu..... VVddu) for every t E I and c E P2(M), this time for some F E C x R*).
this orthogonal part has null contribution in the Taylor expansion given in (7). Thus, in this case, we can restrict ourselves to the unique WoW gradient belonging to the tangent space, in particular to write optimization schemes.
Using the chain rule, any FE ())(P)(M) is Wasserstein differentiable and for all μ E P2(M).
4. WoW Gradient Flows
In this section, we aim at minimizing IF P22(P2 + some functional. We first show the existence of the WoW gradient flow of this functional as the limit of the JKO scheme (Jordan et al., 1998) for F convex along generalized geodesics (Ambrosio et al., 2008). Then, building on the differentiable structure of the space introduced earlier, we propose forward (explicit) scheme that is computationally more efficient in practice than the implicit JKO scheme, and tractable for relevant functionals.
Vw, F(A)= = Σ (f (/vida... dz,
I Vidue
VV.
(10)
This provides the main building block for defining tangent space, in which we will show that WoW gradients reside.
Definition 3.6. The tangent space P E (P2(MM is
TpP TpP2((P(M)) = WW & (y)()2(M)) (11)
4.1. Optimization Schemes on P2((P((R))
where the closure is taken in the space (P,,
JKO Scheme. Let Po E PP2(P a.c (R4)). The JKO sheme of F is defined, for all & 2 0 and T > 0, as
We now justify the definition of this tangent space. We show the existence of velocity fields (Ur)r belonging to the latter, associated to any absolutely continuous curves (Pt),. such that the pair (ct.P.), satisfy a continuity equation. We recall that a curve (Pt))e(0.1)
1
Pk+1 E argmin (P.
Ww, (P,PP)))) + =(()- (13) 2T
abso- lutely continuous if there exists g E ([0, 1]) such that (P,, Pt) < J g())du, and its metric derivative is [P) = lim x Ww, (Pr4n, PP), which exists a.e. (Ambrosio et al., 2008, Th. 1.1.2).
Its Wasserstein gradient flow is defined as the limit when T + O. Leveraging (Ambrosio et al., 2008, Theorem 4.0.4), we show in the next Proposition the existence of the flow for functionals F that are X-convex along gen- eralized geodesics P. (((1 t)Tzi + tT # T
Proposition 3.7. Let (Px)rrr r be an absolutely continuous curve on (PM) Then, for a.e. t E I, there ex- ists U2 E Pe (P2(M)) such that [|v (1(1123( x [P))(() for all & E Cyl(T x P2(M)).
#
between Q,O (2 where I [I( Q) 1.21 E (P, Q). 13T E [I, 0) with satisfies ##
1.2 (2) H (I, x1. (0, - (ar, and
P E P2 (2 (Rd)). Since P E Po (R2)))(( there is always an OT map starting from ~ P towards any v EE (Rd),, which we write K
Sf (Rope(u) (Vw,pi(H),
ec(e)) (eppp(l)).((p)))).) dP (p)ddt O.
(12)
Proposition 4.1. Let 12 O. Let Pe (P2(RR) + R be proper, coercive, lower-semi continuous and X-convex along generalized geodesics, i.c., satisfying for all E [0. 1].
The proof of Proposition 3.7 is deferred to Appendix C.3. We leave the investigation of the converse implication to fu- ture work, f.e. that satisfying the (weak) continuity equation (12) implies absolute continuity of the curve (PP)). We then have the following properties, which show that elements of the tangent space are strong gradients and are unique. Their proofs are deferred to Appendix C.4 and Appendix C.5.
Ae(1
((P.) S (1-t)
W2 (Po,
2
(14)
(-(,, ((11-))T t) + (TTT) T!) I. D Q) for P- = that satisfies TH 1.2 E (1, Q). Ty r E (P, and IP E (P (Rd)). Then, the gradient flow of exists and is unique.
Proposition 3.8. Let & E 0 F(P) n T+P2(P2(M)) Then & is strong subdifferential o at IP. Proposition 3.9. There is at most one element in θ F(P) (P) n (P2()(2
The proof of Proposition 4.1 can be found in Appendix C.6. Examples of X-convex FF on P2 (2(P((R))) include potential energies for any F X-convex along generalized geodesics on (2(R"), and interaction energies for X = and W jointly convex along generalized geodesies, see Appendix B.5.
As L"(P, (PM))) and the tangent space are Hilbert spaces, one can always decompose a WoW gradient with a part in P2 (PM) and another part orthogonal to it. We show in Appendix C.5 that under technical assumptions,
5

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Forward Scheme. Given the existence of the WoW gradi- ent of F, established in the previous section, we propose an alternative to the implicit JKO scheme: forward scheme, commonly referred to as Wasserstein gradient descent, de- fined as follows
on the Sliced- Wasseeetenssr (SW) (Rabin et al., 2012; Bonneel et al., 2015), defined between fr, E P2(RR))
SW = / 196-1
W # f4, P do
(19)
with Sth1 = 10 E Rd, ||0||= 1 the sphere, p (x) (a, 0) the coordinate of the projection of E Rd on the line AR for 0 E 5 and a the uniform measure on Sd-1 For instance, positive definite kernels include the Gaussian SW kernel K (p. 2) SW3(u.p))/ (Kolouri al., 2016; Carriere al., 2017; Meunier al., 2022). We also exper- iment with the Riesz SW kernel K (1r, v) = -SW (M, v) in analogy with the Riesz kernel (sometimes referred to as negative distance kernel), k( y) = - || - yll on Rd x Rd, which is not positive definite, but which has demonstrated very good results in practice (Hertrich et al., 2024b) and does not require tuning a bandwidth h.
(15) At the "distribution particle" level in P2(M), this means that for each distribution Pk ~ PR., we update it as
Ak 2 0, Pk+1 exPp2 (- (-rVww, F(P))).
Ma+1 = exPmu (~rVww, ((P))
(16)
In practice, we will mostly focus distributions of the form P = v >e= 5re with μº 1/2 Σm= 8ag, which notably include labeled datasets (assuming for simplicity now that all classes e = .,...: C contain n examples). Thus, we apply to each particle in M = Rd the update
21,6+1 = = (- TV VV F( = = TVww, TWw, ((P
Wow gradient of the MMD. Given v P2(RR), if K2 j + ((,,v) is a Wasserstein differentiable func- tional, then F is differentiable, and its WoW gradient at PE P2 (P2(R)) is of the form, for all μ E P2(R'),
(17)
We see that there are two levels of interactions for each particle in M: one "intra-class" through the dependence in the distribution Pe and one "inter-class" between the distributions μr ~ P through the dependence in P& in the gradient. Thus, expect to observe an interaction between particles of each distribution μ², but also between each distribution S.a
V Vww, V ((P)())) = / /VaaKolu d(P Q)().. (20)
For the Gaussian SW kernel K(,, e d SW3 de- noting F(u) = 2 1 v), its gradient can be obtained by the chain rule as
4.2. Examples of Discrepancies
Vw, Ka(k)
= - h
-
1
th SW3(HN)) VW, F(u),
(21)
Classical functionals in the study of Wasserstein gradient flows are obtained as linear combinations of potential en- ergies, interaction energies and internal energies (Santam- brogio, 2015). We focus here on potential energies and interaction energies. We leave the study of internal ener- gies on this space for future works. We refer to e.g. (von Renesse & Sturm, 2009; Sturm, 2024) for discussions of entropy functionals on this space.
where F(, = fga- 6ollE, 0))0 (0) with the the Kantorovich potential between P. lu and pgu (Bonnotte, 2013, Proposition 1.7.). In practice, the Sliced-Wassersie distance, involving an integral over the sphere, is approxi- mated through Monte Carlo. Morcover, for discrete mea- sures P = Ee 5pme and Q = Ee 8pan with f. 2 dag and Dan todifferentiation over x == (of) ee of F(x) rescale the Fuclidean gradient of F by X C to obtain the WoW gradient V Vww, ((P)( (x:) = nCV This is analogous to the Wasserstein gradient case, and coincides with the WoW gradient for functionals with a closed-form over discrete measures (see Proposition B.7).
im Oyft we use au-
F F(P), and
A classical discrepancy to compare probability distributions, which can be written as a sum of a potential energy and an interaction energy, the Maximum Mean Discrepancy (MMD) (Gretton et al., 2012). Given a positive definite kernel K (2d) P2(RR R, PP, Q E (P2()))... let (P) = MMMD (P, Q) be defined as
F F(P) = 2. /f K ((H,v) d(p ())) - Q)()) (18)
5. Applications
In this section, we minimize the MMD on Pe (P2 (R") to solve various tasks'. We represent labeled datasets with C classes as distributions P = Σe Spci, where each
= =V(P) + W(P) + cst,
where V(P) = /V())PP()) V(n) = K(u, (v))
WPP)= = 1/2 ʃf K(H, v) (P)))dP(v) and the constant only depends on Q that is fixed. For K, we will use kernels based
'oode available at https://gi com/ /t coneet/ P1 Plooiin_Datasetss. _ow Btcp_oooo__rrddddeeeeeeeee Gradient_ Flows.
6

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
iter
Iter 10
Iter 25
Iter 100
Gaussian
O
00000
QOO COO
O
Riesz
Figure 1:
Minimization of F(P) 2 MMD Q) with Q mixture of 3 rings, and with kernels either the Gaussian SW
kernel with bandwidth / = 0.05 or Riesz SW kernel, for a learning rate of r = 0.1. We observe that they first form a ring for each distribution, and then each ring converges to a target ring.
2 = = Ei dge is the distribution of samples belonging to class c. We emphasize that we are the first to represent labeled datasets this way. We first verify on synthetic data and datasets of images that minimizing such distance allows to transport classes between the source and target. Then, we dataset distillation and a transfer learning task. We focus here on learning target distributions k= of the form Q è Ei= dye.n where yan 1/2 Ei= dyg are empirical distributions, each =s has the same number of particles n and the number of class is supposed to be known. Similarly as (Hertrich et al., 2024b), we add a mo- mentum to accelerate the scheme for image-based datasets. We refer to Appendix D for more details about the experi- ments, as well additional experiments using other kernels and an ablation study for the number of projections to ap- proximate SW. Related works (Alvarez-Melis & Fusi, 2021; Hua et al., 2023) are described in detail in Appendix E.
O
/
>
>
>
3
3
3
3
3
3333
3
leverage this property
4
2
4
5
5
V
S
S
4
6
6
7
7
7
8
8
A
A
8
aoee
9
9
Figure 2: Samples along the flow from MNIST to FMNIST (Left), KMNIST (Middle) and USPS (Right).
ages with C classes. Thus, we have Po = t EE Spmin and Q Σ 8ye.n, and vam. μ.رn" represent the em- pirical distribution of images belonging to the class c. We consider the *NIST datasets, i.e., MNIST (LeCun & Cortes, 2010), Fashion-MNIS, (FMNIST) (Wang et al., 2018) 8), KM- NIST (Clanuwat et al., 2018) 8) and USPS (Hull, 1994). These datasets all have C = 10 classes and are of size 28 x 28 (except for USPS which is upscaled to 28 x 28). We also consider CIFAR10 (Krizhevsky al., 2009) and SVHN (Netzer et al., 2011) which are of size 32 x 32 x 3. We first show in Figure 2 examples of trajectories starting from MNIST to the other *NIST datasets (with step size 7 = 0.05, momentum m = 0.9 and n = 200). We see that samples from MNIST are sent to samples of the target dataset, i.e. that the flow converges well. We also observe that images from each class are mapped one-to-one to images within the same class (see Figure 11 in the Appendix), without overlap or collapse across classes
Synthetic Data. We illustrate on Figure/ the evolu- tion of particles when minimizing the MMD with kernels K K(A, v) = e Sw3 and K K(1, v) = SS, v). for a target being the three-ring dataset. Each ring repre- sents distribution se with n 80, and the target is thus mixture of three Dirac, i.e., Q = 3 >3 8yc.. We learn distribution P of the same form, with the same number of particles for each distribution. We observe for both kernels that the particles of each distribution na (i.e., the different point clouds) form a ring early in the gradient flow dynam- ics, and then move in a structured manner towards the target. This illustrates the two level of interactions at the intra and inter distributions levels. In Appendix D.2, we add compar- isons with other hyperparameters and other kernels. Overall, the kernel K (μ, v) SW2 (1, 1/) is the simplest to use, it does not require tuning a bandwidth, and converges well in general. Thus, in the following experiments, we restrict our- selves to this kernel, and name the resulting loss MMDSW.
To verify this quantitatively, we perform a domain adapta- tion task as in (Alvarez- & Fusi, 2021, Section 7.3). Here, we first train a classifier on 5000 samples of MNIST with 500 images by class. Then, we flow the other
Domain Adaptation. We now focus on the case where both the source Po and the target Q are distributions of im-
7

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Table I: Accuracy of the classifier on synthetic datasets. We compare Distribution Matching (DM) with the MMD with Riesz SW kernel on MNIST and Fashion MNIST, using p E {1, 10,1 50} synthetic images by class.
Dataset
A
160
AA
Id
A
Baselines
DM
MMDSW
DM
MMDSW
DM
MMDSW
DM
MMDSW
Random
Full data
61.1+6.5
66.5.5.5
26.825.3
87.8.0.6
60.3gg.4
87.7
60.923.5
55.8
MNIST
88,2+28
93.220.7
88.723.3
93.8.0.7
97.0-0.1
96.4202
g0.
96. 46.4
92.2+1.
99.4
50
95.920.9
97.0g0.2
95.3-114
97.5g0.
98.4
98.4g0.1
98.4
98.4
97.620.
54.4+3.
60.0=4.1
60.6+3.6
58.7+0.4
60.9=2.6
58.7=0.6
60.8
49.0g7.
FMNIST
10
74.6+1.0
76.7=1.0
74.7=0.8
76.6:1.1
81.2+2.3
78, 78.020.9
82.5+0.3
78.9+1.2
75.3+0.7
92.4
50
81.310.5
84.21o.1
81.4g1.
85.0g0.2
87.610.2
87.610.2
87.5_o.1
87.610.2
83.2 83.210.
FMNIST to MNIST
SVHN to CIFAR1
source dataset (Courty et al., 2016). Thus, to flow the data as we did just earlier, we would need first to find pseudo-labels on the source datasets, e.g. with clustering (Alvarez-Melis & Fusi, 2021; El Hamri et al., 2022). However, this is not the goal of the paper.
100
100
Accuracy
MMDSW
50
90
OTDO (final)
0
5
10
10" 1.5 2-
2-125.1-11
110
10"
1.5-1022-1002 22-1002---02
lterations
Kerations
Dataset Distillation. Dataset distillation or condensation (Wang et al., 2018) seeks to produce a compact synthetic dataset derived from a large training set, such that training neural network on the synthetic data yields performance close to that obtained with the full dataset. Zhao & Bilen (2023) proposed to learn the synthetic dataset by performing Distribution Matching, i.e., denoting o' the distribution of each class c of the target dataset, they minimize
Figure 3: Accuracy of the pretrained classifiers along the flow from FMNIST (Left) and SVHN (Right) towards MNIST and CIFAR10.
datasets to MNIST (with r = 0.1 and momentum m = 0.9), and measure the accuracy of the pretrained classifier on the flowed dataset. Note that while we use the class labels of the flowed dataset perform the gradient flow dynamies, we do not know priori which class in the flowed dataset cor- responds to which class in MNIST, yet it is needed for the evaluation of domain adaptation. To perform this alignment, we solve an OT problem with W3 as groundcost between P and QQ (i.e., the WoW OT problem) with IP the distribu- tions obtained at the end of the flow dynamic and Q the ones of the target dataset. Since these distributions have finite support of the same size (C), solving this OT problem provides such an alignment: we can associate prediction of the pretrained model to an image and "true class" of the flowed dataset. We also perform this experiment with a pretrained neural network on CIFAR1O, flowing SVHN toward CIFARIO 10, with n = 100 samples by class, step size 0.1 and momentum m 0.9.
Σ MMD (AAA w Ageo
F((())) = Foow
(22)
with is the linear kernel k(o, v) (z,, v). Ao Rd - Rd random data augmentation (e.g. rotation, cropping, see (Zhao & Bilen, 2021)) and v0 Rd + Rd with d' << d a randomly initialized neural network used to embed the data.
Let Q = d Le de be the target dataset, (M² = we A f and P à #
EE Syc. Note that Φ#
6lPP p
y Doni Oug uggggae We propose to minimize
F(P) Eow
[MMD]
0. (0#
(23)
On Figure 3, we report the accuracy of the pretrained clas- sifier on the data flowed starting from FMNIST towards MNIST and from SVHN towards CIFAR10, 10, over the iter- ations (averaged over flows started at different splits of the source data). We also report the value from (Alvarez- Melis & Fusi, 2021) using OTDD on the MNIST dataset. We observe that the classifier converges to 100% accuracy for sufficient number of iterations. This demonstrates that the flow is able to perfectly match one class from the source dataset with class of the target dataset, on which the classifier has been trained.
with Riesz SW kernel K between distributions. We com- pare on Table 1 the accuracy of classifier on a test set of MNIST and FMNIST, trained on the synthetic dataset with p E (1, 10, 50) samples by class, either generated with MMD with Riesz SW kernel (MMDSW) or with Distribu- tion Matching (DM), in 4 scenarios: in the ambient space with (v = Ld) and without augmentation (vi" A Id), and with an embedding with (Ar + vi²) and without an augmentation (A* = Td). We solve it with a stochastic gra- dient descent, sampling one augmentation and embedding at each step, for 20K iterations and initializing the samples on true data. The results are averaged over synthetic datasets obtained initializing the flow at different samples, and
We note that in a realistic setting of unsupervised domain adaptation, we would not have access to the labels of the
8

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Table 2: Accuracy of classifier on augmented datasets for k E {1, 10, 10, 100). M refers to MNIST, F to Fashion MNIST, K to KMNIST and U to USPS.
Table 3: Runtime in seconds for the transfer learning experi- ment from MNIST to Fashion MNIST.
Dataset
R-shot
MMDSW
OTDD
(Hua 2023)
Dataset
3:
Train Q
MMDSW
OTDD
(Hua al., 2023)
13.95 + 1.37
294.53
5.21
131.77 +
26.0g5.3
40.5
30.5244
36..4
14.12 + 0.30
1130.89 108
132.98 1.1
5
38.5
61.5246
69.7218
62.7=1.1
10
14.30 +: 0.29
2294.13 + 48
134.35 +: 0.75
t
10
63.9179
65.4g1.5
64.0g14
66.221.
100
47.75 + 0.27
164.19 0.6
100
71.1a15
74.7
73.
18.4ga.
20.9:2.
18.842.1
19.4
M to K
25.9244
37.4g22
31.3g14
39.0.110
data obtained by minimizing OTDD (Alvarez-Melis & Fusi, 2021) or the MMD with product kernel as in (Hua et al., 2023). The results are averaged over training of the networks, and outputs of the flows. Both methods have been reimplemented and we add more details in Appendix D.6. We observe that all three methods improve upon the baseline, with a slight advantage for MMDSW.
10
30.9+46
44.7=1.
34.1
44.1=1.2
100
60.1g1.1
66.8
66.320.
62.4
32.4g79
37.426.
39.5
35.025.
51.4g9.8
73.0=10
73.3:14
69.6
M U
10
60.3g10.1
77.221/2
72.7227
75.6g1.2
100
57.5g0.7
89.7
88.1200
training of the classifier. On a Nvidia v100 GPU, the flow implemented in Jax (Bradbury et al., 2018) runs in around 10 minutes with the embedding and in 30 seconds without it. The baseline "random" refers to the classifier trained data sampled randomly from the original training set, and "full data" to the classifier trained on the full training set. We observe on Table 1 that MMDSW consistently outperforms DM when flowing in the ambient space, and is competitive when adding an embedding. This indicates that adding inter- actions between classes appears to improve the results, possi- bly by distributing the samples more effectively and mitigat- ing the presence of ambiguous samples near class borders.
Complexity. Given P
= P=
- E. Symn and Q
E Egenn with fa and van discrete distributions with samples each, the MMD with a -asse-ssasn based kernel requires to compute C2 Wasserseasen distances, which has a total complexity of O Ln n + d)) using L projections to approximate SW. We report on Table the runtimes for the transfer learning experiment, averaged over 3 outputs of the flows and trained for 5K epochs for each method. MMDSW is much faster than both OTDD and the MMD with product kernel, at least with our imple- mentations in jax detailed in Appendix D.6. Both OTDD and the method of Hua et al. (2023) are implemented us- ing a dimension reduction technique in 2D and a Gaussian approximation to embed the conditional distributions.
Transfer Learning. We now focus on the task of k-shot learning. In this setting, we are interested in training classifier for datasets which have k samples by class, where k is typically small. Following (Alvarez-Melis & Fusi, 2021; Hua et al., 2023), we propose to augment the dataset by generating new synthetic samples for each class. To do this, we will flow a larger source dataset, with possibly different classes, towards the small target dataset, and then concatenate the synthetic and true samples train the classifier on it. More precisely, let Q = EE 8pe.k the target dataset, with yo = + E 6g5 an empirical distribution with k samples, representing the distribution of the class c. Let Po = Σc= 6yc... be a source dataset, n Em day and 200. Then, the goal is to
6. Conclusion
This work provides the first theoretical framework and prac- tical implementation of gradient flows of a suitable MMD objective over the space of random measures, endowed with the Wasserstein over Wasserstein distance. On the thcoret- ical side, we provided a rigorous differential structure on that space and showed that these flows are well-posed. On the numerical side, our results demonstrate that this novel approach provides meaningful dynamics for interpolating between random measures. There are many possible extensions of our study. For instance, it would be interesting to investigate the minimization of alternative functionals over the space of random measures, e.g., MMD with other kernels (Bachoc et al., 2023; Kachaiev & Recanatesi, 2024). integral probability metrics (Müller, 1997; Catalano & Lavenant, 2024) or f-divergences (Csiszár, 1967). Future work could also address the theoretical treatment of non-compact manifolds or derive a continuity equation for Wasserstein over Wasserstein (WoW) gradient flows. Then, another topic of future research would be to provide quan- titative guarantees on the convergence of these schemes.
with y... = 1 flow Po towards Q by minimizing F(P)= = MMDD (P,, Q) with kernel K K(μ,v) v) = -SW2 v). We expect to augment each class of with n samples. Then, we train a classifier with LeNet5 architecture on the dataset obtained as Q 1 EE ogennt with n- where =f = If for i S and of = vi_n for i > n.
We report the results for MNIST as Po and FMNIST, KMNIST and USPS as Q on Table 2 for l E (1,5, 10, 100). compared with the baseline where we train directly on Q. and the baselines where we trained on the synthetic

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Acknowledgements
Bachoe, F., Béthune, L.., Gonzalez-Sanz, A., and Loubes, J.-M. Gaussian Processes on Distributions based Reg- ularized Optimal Transport. In International Conference on Artificial Intelligence and Statistics, pp. 4986-50100 PMLR, 2023. (Cited on p. 9)
We thank the anonymous reviewers for their valuable com- ments. We also thank Quentin Mérigot for fruitful discus- sions, and David Alvarez-Melis for his help on the trans- fer learning experiment. This work was granted access to the HPC resources of IDRIS under the allocation 2024- ADO1 1015891 made by GENCI. CB and AK acknowledge the support of the Agence nationale de la recherche, through and also thank Apple for their academic support through a research funding. CV acknowledges the support of Région île-de- France through the DIM AI4IDF project.
Ballu, M. and Berthet, Q. Mirror Sinkhorm: Fast Online Op- timization on Transport Polytopes. In International Con- ference on Machine Learning, pp. 1595-1613. PMLR, 2023. (Cited on p. 53)
the PEPR PDE-AI AI project (ANR-23-) 1-----------00-
Bellazzi, R., Codegoni, A., Gualandi, S., Nicora, G., and Vercesi, E. The Gene Mover's Distance: Single- cell similarity via Optimal Transport. arXiv preprint arXiv:1102.01218, 2021. (Cited on p. 1)
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Bojchevski, A. and Günnemann, S. Deep Gaussian Em- bedding of Graphs: Unsupervised Inductive Learning via Ranking. In International Conference on Learning Representations, 2018. (Cited on p. 1)
Bonet, C., Malézieux, B., Rakotomamonjy, A., Drumetz, I.., Moreau, T., Kowalski, M., and Courty, N. Sliced- Wasserstein on Symmetric Positive Definite Matrices for M/EEG Signals. In International Conference on Machine Learning, pp. 2777-2805. PMLR, 2023. (Cited on p. 1)
References
Altekrtiger, F., Hertrich, J., and Steidl, G. Neural Wasser- stein Gradient Flows for Discrepancies with Riesz Ker- nels. In International Conference on Machine Learning, pp. 664-666666 690. PMLR, 2023. (Cited on p. 2)
Bonet, C., Uscidda, T., David, A., Aubin-Frankowse Frankowski, P.- C., and Korba, A. Mirror and Preconditioned Gradient Descent in Wasserstein Space. In Thirty-eight Conference on Neural Information Processing Systems, 2024. (Cited on p. 2)
Altschuler, J., Chewi, S., Gerber, P. R., and Stromme, A. Av- eraging on the Bures- Wasserstein manifold: dimension- free convergence of gradient descent. Advances in Neural Information Processing Systems, 34:22132-22145, 2021. (Cited on 52, 55)
Bonet, C., Drumetz, L., and Courty, N. Siteed-Wassesstein Distances and Flows on Cadtan-Hadamard Manifolds. Journal of Machine Learning Research, 26(32):1-76,
Alvarez-Melis, D. and Fusi, N. Geometric Dataset Distances via Optimal Transport. Advances in Neural Information Processing Systems, 15:214228-21499 2020. (Cited on p. 1,51, 54)
2025. (Cited on p. 2, 3, 19, 54)
Bonnabel, S. Stochastic Gradient Descent on Riemannian Manifolds. IEEE Transactions on Automatic Control, 58 11-2217-2229 2013. (Cited on p. 3, 19)
Alvarez-Melis, D. and Fusi, N. Dataset Dynamics via Gra- dient Flows in Probability Space. In International con- ference on machine learning, pp. 219-230. PMLR, 2021. (Cited on p. 1, 2, 7, 8, 9, 54)
Bonneel, N., Rabin, J., Peyré, G., and Pfister, Sliced and Radon Wasserstein Barycenters of Measures. Journal of Mathematical Imaging and Vision, 51:22 45, 2015. (Cited p. 6)
Ambrosio, L.. Gigli, N., and Savaré, G. Gradient Flows: in Metric Spaces and in the Space of Probability Measures. Springer Science & Business Media, 2008. (Cited on p. 2, 3, 4, 5, 19, 20, 24, 25, 26, 27, 29, 33)
Bonnet, B. A Pontryagin Maximum Principle in Wasser- stein Spaces for Constrained Optimal Control Problems. ESAIM: Control, Optimisation and Calculus of Varia- lions, 25:52, 2019. (Cited on p. 17)
Andoni, A., Naor, A., and Neiman, O. Snowflake universal- ity of Wasserstein spaces. Ann. Sei. Éc. Norm. Supér(4), 51(3):557-700 2018 (Cited on p. I)
Arbel, M., Korba, A., Salim, A., and Gretton, A. Maximum Mean Discrepancy Gradient Flow. Advances in Neural Information Processing Systems, 32, 2019. (Cited on p. 2, 52)
Bonnotte, N. Unidimensional and Evolution Methods for Optimal Transportation. PhD thesis, Université Paris Sud- Paris XI; Scuola normale superiore (Pise, Italic), 2013.
(Cited on p. 6, 44)
10

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Boumal, N. An Introduction Optimization Smooth Manifolds. Cambridge University Press, 2023. (Cited p. 35, 36, 37)
Cuturi, M. Sinkhorn Distances:
Lightspeed Computation of Optimal Transport. Advances in neural information processing systems, 26, 2013. (Cited on p. 51)
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang. Q. JAX: composable transformations of Python+NumPy programs, 2018. URL http: /githb/ com/ jaann/aammm///// jaxx. (Cited on p. 9, 52)
Cuturi, M., Meng-Papaxanthos: L., Tian, Y., Bunne, C.,
Davis, G., and Teboul, O. Optimal Transport Tools (OTT); A Jax Toolbox for all things Wasserstein. arXiv preprint arXiv:2201. 12324, 2022. (Cited on p. 51)
Dello Schiavo, L. A Rademacher-type theorem L2- Wasserstein spaces over closed Riemannian manifolds. Journal of Functional Analysis, 278(6):108397, 2020. (Cited on p. 4, 21, 32, 33)
Carriere, M., Cuturi, M., and Oudot, S. Sliced Wasserstein Kernel for Persistence Diagrams. In International con- ference on machine learning, pp. 663-673. PMLR, 2017. (Cited on 6, 44)
Delon, J. and Desolneux, A. A Wasserstein-type distance in the space of Gaussian Mixture Models. SIAM Journal on Imaging Sciences, 13(2):936 2020. (Cited p. 1, 55)
Catalano, M. and Lavenant, H. Hierarchical Integral Probability Metries: A distance random probability measures with low sample complexity. arXiv preprint 2.Xiv:2402.004233, 2024. (Cited on p. 1,
Diao, M. Z., Balasubramanian, K., Chewi, S., and Salim, A. wwward-Backwwrd Gaussian Variational Inference via JKO in the Bures-Wass Warse-Wassrrttein Space. In Interna- tional Conference Machine Learning, pp. 7960-7991. PMLR, 2023. (Cited on p. 57)
Chazal, C., Bach, F., and Korba, A. Statistical and Geometri- cal properties of the Kernel Hellback--eible divergence. In The Tiirtt-eighth Annual Conference on Neural Infor- mation P Processing Systems, 2024. (Cited on 2)
Du, C., Li, T., Pang, T., Yan, S., and Lin, M. Nonpara- metric Generative Modeling with Conditional Sliced- Wasserstein Flows. In International Conference on Ma- chine Learning. pp. 8565 8584. PMLR, 2023. (Cited on p.
Chen, K. and Zhang, Y. Optimal Transport for Mixtures of Radial Functions. arXiv preprint 2XXv:24044088883, 2024. (Cited on p. 1)
Chen, Y., Georgiou, T. T., and Tannenbaum, A. Optimal Transport for Gaussian Mixture Models. JEEE Access, 7: 6269-6278, 2018. (Cited p. 55)
Dukler, Y., Li, W., Lin, A., and Montúfar, G. Wasserstein of Wasserstein Loss for Learning Generative Models. In International conference on machine learning, pp. 1716- 1725. PMLR, 2019. (Cited on p. 1)
Chen, Z., Mustafi, A., Glaser, P., Korba, A., Gretton, A., and Sriperumbudur, B. K. (De)-rcgularized Maxi- mum Mean Discrepancy Gradient Flow. arXiv preprint arXiv:2409. 14980, 2024. (Cited p.
Dusson, G., Ehrlacher, V., and Nouaime, N. A Wasserstein- type metric for generic mixture models, including location-scatter and group invariant measures. arXiv preprint arXiv:2301.07963 2023. (Cited on p. 1)
Choi, J., Choi, J., and Kang, M. Scalable Wasserstein Gra- dient Flow for Generative Modeling through Unbalanced Optimal Transport In Forty-first International Confer- ence on Machine Learning, 2024. (Cited on p. 52)
Dvurechensky, P., Gasnikov, A., and Kroshnin, A. Com- putational Optimal Transport: Complexity by Acceler- ated Gradient Descent is Better than by Sinkhorn's 's Algo- rithm. In International conference on machine learning, pp. 1367-1376. PMLR, 2018. (Cited on p. 54)
Clanuwat, T., Bober-Irizar, M., Kitamoto, A., Lamb, A., Yamamoto, K., and Ha, D. Deep Learning for Classical Japanese Literature, arXiv preprint arXiv: 1XXi::8112.0711: 2018. (Cited on p.
El Hamri, M., Bennani, Y., and Falih, I. Hierarchical Op- timal Transport for Unsupervised Domain Adaptation. Machine Learning, 111(11):4159 4182, 2022. (Cited on p. 1)
Courty, N., Flamary, R., Tuia, D., and Rakotomamonjy, A. Optimal Transport for Domain Adaptation. IEEE transactions on pattern analysis and machine intelligence, 39(9):1853-186, 2016. (Cited on p. 8)
Emami, P. and Pass, B. Optimal transport with optimal trans- port cost: the Kantooohntoo problem on Wasser- stein spaces. Calculus of Variations and Partial Differ- ential Equations, 64(2):43, 2025. (Cited on 4, 21, 26 33)
Csiszár, I. Information- Type Measures of Difference of Probability Distributions and Indirect Observations. Stu- dia Sci. Math. Hungar., 2:299-313 -99,,318, 1967. (Cited on p. 9)

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Erbar, M. The heat equation manifolds a gradient flow in the wasserstein space. In Annales de I'IHP Probabilités et statistiques, volume 46, pp. 1-23, 2010. (Cited on p. 2, 3, 4, 16, 17, 18, 19, 26)
Haviv, D., Remšík, J., Gatie, M., Snopkowski, C., Takizawa, M., Pereira, N., Bashkin, J., Jovanovich, S., Nawy, T. Chaligne, R., et al. The covariance environment defines cellular niches for spatial inference. Nature Biotechnol- ogy, pp. 1-12, 2024b. (Cited on p. 1)
Figalli, A. Existence, Uniqueness, and Regularity of Op- timal Transport Maps. SIAM journal on mathematical analysis, 166):126-1377, 2007. (Cited on p. 2, 16)
He, S., Liu, K., Ji, G., and Zhao, J. Learning to Represent Knowledge Graphs with Gaussian Embedding. In Pro- ceedings of the 24th ACM international on conference on information and knowledge management, pp. 623-632, 2015. (Cited on p. 1)
Flamary, R., Courty, N., Gramfort, A., Alaya, M. Z., Bois- bunon, A., Chambon, S., Chapel, L., Corenflos, A., Fatras, K., Fournier, N., et al. POT: Python Optimal Transport. Journal of Machine Learning Research, 22(78):1-8, 2021. (Cited on p. 48, 51)
Hertrich, J., Gräf, M., Beinert, R., and Steidl, G. Wasser- stein Steepest Descent Flows of Discrepancies with Riesz Kernels. Journal of Mathematical Analysis and Applica- tions, 531 531(1)127829, 2024a. (Cited on 2, 3)
Fornasier, M., Savaré, G., and Sodini, G. Density of sub- algebras of Lipschitz functions in metric Sobolev spaces and applications to Wasserstein Sobolev spaces. Journal of Functional Analysis, (55111):11113,, 2023. (Cited on 4)
Hertrich, J., Wald, C.. Altekrüger, F,, and Hagemann, P. Generative Sliced MMD Flows with Riesz Kernels. In The Twelfth International Conference on Learning Repre- sentations, 2024b. (Cited on p. 2, 6, 7, 47)
Gallouët, T. and Monsaingeon, L. A JKO Splitting Scheme for Kantorovich-Fisher--Be Gradient Flows. SIAM Jour- nal on Mathematical Analysis, 1902):100011133 2017. (Cited on p. 52)
Ho, N., Nguyen, X., Yurochkin, M., Bui, H. H., Huynh, V., and Phung, D. Multilevel Clustering via Wasserstein Means. In International conference machine learning, pp. 1501-1509. PMLR, 2017. (Cited on p. 1)
Geuter, J., Bonet, C., Korba, A., and Divare--Melie D. DDEQs: Distributional Deep Equilibrium Models through Wasserstein Gradient Flows. In The 28th Interna- tional Conference Artificial Intelligence and Statistics, 2025. (Cited p. 1)
Hua, X., Nguyen, T., Le, T., Blanchet, J., and Nguyen, V. A. Dynamic Flows on Curved Space Generated by Labeled Data. In Elkind, E. (ed.), Proceedings of the Thirty-Second International Joint Conference on Artifi- cial Intelligence, IJCAI-23, pp. 3803-3811 International Joint Conferences on Artificial Intelligence Organization, 8 2023. Main Track. (Cited p. 1, 2, 7, 9, 51, 52, 54, 55)
Gigli, N. On the geometry the space of probability mea- sures in IR endowed with the quadratic optimal transport distance. PhD thesis. Scuola Normale Superiore, 2004. (Cited on 3)
Huix, T., Korba, A., Durmus, A. O., and Moulines, E. The- oretical Guarantees for Variational Inference with Fixed- Variance Mixture of Gaussians. In Forty-first Interna- tional Conference on Machine Learning, 2024. (Cited on p. 17)
Gigli, N. On the inverse implication of Brenier-McCann theorems and the structure of (PM) W2). Methods and Applications of Analysis, 18(2):127-158, 2011. (Cited on p. 3. 16, 17, 18, 25)
Glaser, P., Arbel, M., and Gretton, A. KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Support. Advances in Neural Information Processing Systems, 34:8018-8003 2021. (Cited p. 2, 47)
Hull, J. A Database for Handwritten Text Recognition Research. IEEE Transactions on pattern analysis and machine intelligence, 16(5):550-554, 1994. (Cited on p.
7) Jordan, R., Kinderlehrer, D., and Otto, F. The Variational Formulation of the Fokker-Planck Equation. SIAM jour- nal on mathematical analysis, 29(1):1-17, 1998. (Cited on p. 2. 5, 19)
Gretton, A., Borgwardt, K. M., Rasch, M. J., Schölkopf, B., and Smola, A. A kernel two-sample test. The Journal of Machine Learning Research, 13(1):723-773, 2012. (Cited on p. 6)
Haviv, D., Pooladian, A.-A Pe'er, D., and Amos, B. Wasserstein Flow Matching: Generative mod- eling over families of distributions. arXiv preprint Stt::24/1.0069// 2024a. (Cited on p. 1)
Kachaiev. O. and Recanatesi, S. Learning to Embed Dis- tributions via Maximum Kernel Entropy. In The Thirty- eighth Annual Conference on Neural Information Pro- cessing Systems, 2024. (Cited on p. 9)
12

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Kidger, P. and Garcia, C. Equinox: neural networks in JAX via callable PyTres and filtered transformations. Differ- entiable Programming workshop at Neural Information Processing Systems 2021, 2021. (Cited on p. 48)
Lu, Y., Lu, J., and Nolen, J. Accelerating Langevin Sam- pling with Birth-Death. arXiv preprins arXiv: 190vvv990.009866 2019. (Cited on p. 52)
Lu, Y., Slepëev, D., and Wang, L. Birth death dynamics for sampling: global convergence, approximations and their asymptotics. Nonlinearity, 36(11):5731, 2023. (Cited on p. 52)
Kolouri, S., Zou, Y., and Rohde, G. K. Sliced Wasserstein Kernels for Probability Distributions. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pp. 5258-5267, 2016. (Cited on p. 6, 44)
Manupriya, P., Jagarlapudi, S., and Jawanpuria, P. MMD- Regularized Unbalanced Optimal Transport. Transactions on Machine Learning Research, 2024. ISSN 2835-8856. (Cited on p. 53)
Kratsios, A., Debarnot, V., and Dokmanić, I. Small Trans- formers Compute Universal Metric Embeddings. Jour- nal of Machine Learning Research, 24(170): 2023. (Cited on p. 1)
Krizhevsky, A., Hinton, G., et al. Learning Multiple Layers of Features from Tiny Images. 2009. (Cited on p. 7)
McCann, R. J. Polar factorization of maps on Riemannian manifolds. Geometric & Functional Analysis GAFA, 11 (3::5896088, 2001. (Cited on p. 2, 16)
Kusner, M., Sun, Y., Kolkin, N., and Weinberger, K. From Word Embeddings to Document Distances. In Inter- national conference on machine learning, pp. 957-966. PMLR, 2015. (Cited on p. 1)
Meunier, D., Pontil, M., and Ciliberto, C. Distribution Regression with Sliced Wasserstein Kernels. In Inter- national Conference on Machine Learning, pp. 15501- 15523. PMLR, 2022. (Cited on p. 6, 44)
Lambert, M., Chewi, S., Bach, F., Bonnabel, S., and Rigol- let, P. Variational inference via Wasserstein gradient flows. Advances in Neural Information Processing Systems, 35: 14434-14447, 2022. (Cited on p. 55, 57)
Müller, A. Integral Probability Metries and Their Generating Classes of Functions. Advances in applied probability, 29(2):429-444 443, 1997. (Cited on p. 9)
Lanzetti, N., Bolognani, S., and Dörfler, F. First-Order Conditions for Optimization in the Wasserstein Space. SIAM Journal on Mathematics of Data Science, 7(1):
Netzer, Y., Wang, T. Coates, A., Bissacco, A., Wu, B., Ng, A. Y., et al. Reading Digits in Natural Images with Unsupervised Feature Learning. In NIPS workshop on deep learning and unsupervised feature learning, volume 2011, pp. 4. Granada, 2011. (Cited on p. 7)
274-300, 2025. (Cited p. 3, 4, 16, 37)
LeCun, Y. and Cortes, C. MNIST handwritten digit database. 2010. (Cited p.
Neumayer, S., Stein, V., Steidl, G., and Rux, N. Wasserstein Gradient Flows for Moreau Envelopes of f-Divergences in Reproducing Kernel Hilbert Spaces. arXiv preprint wwXhv:2402.04613 2024. (Cited on p. 2)
Lee, J. M. Riemannian Manifolds: An Introduction Cur- vature, volume 176. Springer Science & Business Media, 2006. (Cited p. 36)
Lim, J. N. and Johansen, A. Partiele Semi-Implici Varia- tional Inference. Advances in Neural Information Pro- cessing Systems, 112112395--139999 2024. (Cited p. 57)
Nguyen, K. and Ho, N. Hierarchical Hybrid Sliced Wasser- stein: A Scalable Metric for Heterogeneous Joint Dis- tributions. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. (Cited on p. 54)
Liu, X., Bai, Y., Lu, Y., Soltoggio, A., and Kolouri, S. Wasserstein Task Embedding for Mcasuring Task Simi- larities. Neural Networks, 181:106796, 2025. (Cited on p. 54)
Nguyen, K., Nguyen, H., Pham, T., and Ho, N. Lightspeed Geometric Dataset Distance via Sliced Optimal Transport. arXiv preprint arXiv:2501.1890 2025. (Cited on p. 54)
Liutkus, A., Simsekli, U., Majewski, S., Durmus, A., and Stöter, F.-. Sliced- Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffu- sions. In International Conference on Machine Learning, pp. 4113-4113. PMLR, 2019. (Cited on p. 2)
Nguyen, X. Borrowing strengh in hierarchical Bayes:
Posterior concentration of the Dirichlet base measure. Bernoulli, 22(3):1535 1571, 2016. (Cited on p. 1,
Otto, F. The Geometry of Dissipative Evolution Equa- tions: The Porous Medium Equation. Communications in Partial Differential Equations, 26(1-2):101-77,, 2001. (Cited on p. 2)
Loshchilov, I. and Hutter, F. Decoupled Weight Decay Regularization. In International Conference on Learning Representations, 2019. (Cited on p. 48)
13

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., Vanderplas, J., Passos, A., Cour- napeau, D., Brucher, M., Perrot, M.. and Duchesnay, E. Scikit-learn: Machine learning in Python. Journal of Machine Learning Research, 12:2825-28338 2830, 2011. (Cited on p. 52)
Sodini, G. E., Fornasier, M., and Heid, P. Approximation Theory, Computing, and Deep Learning on the Wasser- stein Space. Mathematical Models and Methods in Ap- plied Sciences, 2025. (Cited on
Sturm, K.-T. Wasserstein Diffusion Multidimensional Spaces. arXiv preprint 2rXiv:2401.12720 2024. (Cited on p. 6)
Persad, S., Choo, Z.-N., Dien, C., Sohail, N., Masilionis, I., Chaligné, R., Nawy, T., Brown, C. C., Sharma, R., Pe'er, I., et al. SEACells infers transcriptional and epigenomic cellular states from single-cell genomics data. Nature Biotechnology, (1(1))1174 1757, 2023. (Cited on 1)
Villani, C. Topics in Optimal Transportation, volume 58. American Mathematical Soc., 2003. (Cited on p. 56)
Villani, C. al. Optimal Transport: Old and New, volume 338. Springer, 2009. (Cited on p. 16, 24, 34)
Petric Maretic, H., El Gheche, M., Chierchia, G., and Frossard, P. GOT: an Optimal Transport framework for Graph comparison. Advances in Neural Information Pro- cessing Systems, 32. 2019. (Cited on p.
Vilnis, L. and McCallum, A. Word Representations via Gaussian Embedding. In Bengio, Y. and LeCun, Y. (eds.), 3rd International Conference tions, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. (Cited on p. 1)
Learning Representa-
Qi, R. Su, H., Mo, K., and Guibas, L. J. PointNet: Deep learning on Point Sets for 3D Classification and Segmentation. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 652-660, 2017. (Cited on p.
Von Renesse, M.-K. and Sturm, K.-T. Entropic Measure and Wasserstein Diffusion. The Annals of Probability, 37(3): 1194-119, 2009. ISSN 0091 1991177 1798. (Cited on p. 4, 6)
Wang. F. 2.,nhan, Z., Sun, L.., Ye, J., and Yan, Y. DiriE: Knowledge Graph Embedding with Dirichlet Distribution. Proceedings of the ACM Web Conference 2022, pp. 3082-3091, 2022. (Cited on p. 1)
Rabin, J., Peyré, G., Delon, J., and Bernot, M. Wasserstein Barycenter and its Application to Texture Mixing. In Scale Space and Variational Methods Computer Vision: Third International Conference, SSVM 2011, Ein-Gedi, Israel, May 29-June 2, 201 I, Revised Selected Papers 3, pp. 435 446. Springer, 2012. (Cited on p. 6, 44)
Wang, T., Zhu, J.-Y,, Torralba, A., and Efros, A. A. Dataset Distillation. arXiv preprint arXiv:1811 10959, 2018. (Cited on p. 1, 1,7,8 7)
Rønning, O., Nalisnick, E., Ley, C., Smyth, P., and Hamel- ryck, T. ELBOing Stein: Variational Bayes with Stein Mixture Inference. In The Thirteenth International Con- ference on Learning Representations, 2025. (Cited on p. 57)
Wibisono, A. Sampling as optimization in the space of measures: The Langevin dynamics as a composite opti- mization problem. In Conference on Learning Theory, pp. 2093-3027. PMLR, 2018. (Cited on p. 2)
Rudin, W. Real and Complex Analysis. McGraw-Hill Sci- ence/Engineering/Mah, 1986. ISBN 0070542341. (Cited on p. 17)
Wilson, M., Needham, T., Park, C., Kundu, S., and Sri- vastava, A. A Wasserstein- Type Distance for Gaussian Mixtures on Vector Bundles with Applications Shape Analysis. SIAM Journal Imaging Sciences, 17(3): 1433-1466, 2024. (Cited on p. 1)
Salim, A., Korba, A., and Luise, G. The Wasserstein Proxi- mal Gradient Algorithm. Advances in Neural Information Processigg Systems, 33:12356: 12366, 2020. (Cited on p. 2)
Yan, Y., Wang, K., and Rigollet, P. Learning Gaussian Mixtures Using the Wasserstein-Fisherr Flow. The Annals of Statistics, 12(4):1774-1795. 2024. (Cited on p. 52)
Gradient
Santambrogio, F. Optimal Transport for Applied Mathe- maticians. Birkäuser, NY, 55(58-63):94, 2015. (Cited on
p. 6, 22)
Yurochkin, M., Claici, S., Chien, E., Mirzazadeh, F,, and Solomon, J. M. Hierarchical Optimal Transport for Doc- ument Representation. Advances in neural information processing systems, 32, 2019. (Cited on p.
Santambrogio, F. (Fuclidean, Metric, and Wasserstein) Gradient Flows: an overview. Bulletin of Mathematical Sciences, 7:8,,154, 2017. (Cited on 2)
Séjoumé,, T., Peyré, G., and Vialard, F-X. Unbalanced Optimal Transport, from Theory to Numerics. Handbook of Numerical Analysis, 24:407-471, 2023. (Cited on p. 53)
Zhao, B. and Bilen, H. Dataset Condensation with Differ- entiable Siamese Augmentation. In International Con- ference on Machine Learning, pp. 12671-1266660 12685. PMLR, 2021. (Cited on p. 8, 50)
14

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Zhao, B. and Bilen, H. Dataset Condensation with Dis- tribution Matching. In Proceedings of the IEEE/CVF Winter Conference Applications of Computer Vision, Pp. 6514 6523, 2023. (Cited on p. 8, 50)
15

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
A. Background on Optimal Transport
A.1. Optimal Transport on (2(M)
Let M be a Riemannian manifold, and denote by d M M + RR4 the associated geodesic distance. We recall that for any μ, v E P2(M), the Wasserstein distance is defined as
gell(a)) / [(c,,)) ().
Wa(m
inf
(24)
Let o : MM- + R. For cost c: M x M R, recall that its c-transform is defined (p(() infzem c(a, y) (zx). is said to be e-concave if there exists Ø .M + R such that y Φe. Here, we focus on c(x, y) 2d(2,, v)². Then, the Wasserstein distance can be written through its dual (see e.g. (Villani et al., 2009, Theorem 5.10)) as
/ | + f
w)(,v)=
sup (EL
sddv,
(25)
with (4) {f: M R, flfldu sc}. The optimal f is called the Kantorovich potential, is noted Spaas and is a c-concave map.
We now recall McCann's theorem, which provides sufficient condition for the existence of an OT map provided that is absolutely continuous with respect to the volume measure. We state the result for connected compact Riemannian manifold. But, note that this result was then extended to other manifolds, see e.g. (Figalli, 2007, Proposition 3.1) for a similar result on complete Riemannian manifolds.
Theorem A.1 (Theorem 9 in (McCann, 2001)). Let M be a connected compact Riemannian manifold. Let E (2) and E P2(M). Then, the optimal coupling 5 E I[(,,v) unique and of the form ÿ = (Id, T) with T(z) expx (x)) for all x E M, where Pu. Kantorovich potential for the pair p.v..
For any x E M, recall that the exponential map exp : TM + M maps tangent vectors v E T back to the manifold at the point reached at time 1 by the geodesic starting at z with initial velocity v. Moreover, when it is well defined, its inverse is the logarithm map log, .M- + M, which satisfies, for any However, the exponential map is not always invertible. For instance, the sphere 5d-1, there are infinite number of geodesics, and thus of directions E M, between M and its antipodal point -rr (see Figure 4). Therefore, the logarithm map logg( (-z) is multivalued.
M, y exp, (v) where
T,M, log, (y) v.
Let μ4, V E (MM. When the exponential map is invertible at μ-almost every x E M, then a (constant-speed geodesic between fa and v can be defined for all E [0, 1] as μr: = =xp,,1 of log, ? # where 7 E IT v), i.e. it satisfies W2(Hs, (t) = It- - s| s/Wa(,,)) v) for all s, E [0, 1]. However, the exponential map might not always be invertible, as described in the last paragraph. One way to circumvent this problem is to consider the space
expp (v) = (7 E P2(M)), 7 exp 7 v, f(l [ell2 dr(x,,) v) = W2 v)).
(26)
This space carries more information than the set of optimal couplings as it precises which geodesic was chosen to move the mass from μ to (Gigli, 2011). Indeed, regarding the previous example the sphere, for z E s m ò₂ and 6_n- (w), while the optimal coupling would simply be given by the map T(x) = -z. Moreover, it allows to define geodesics as t + Pe exp,M (tx*))+ for any Y E exp,, (v) (Gigli, 2011, Theorem 1.11). By Proposition 2.1, if μ E (M)) then there exists a unique E exp, n (V), which is of the form 7 = (Id, PVp ) In this case, the geodesic between μ and v is of the form μ = expia of (( -t
and any u E T.M such that = (v), we have ((2,0) E exp,," "p)),
A.2. Wasserstein Gradient Flows on P2(M)
We provide in this Section some background on Wasserstein gradient flows on P2(M), with M a Riemannian manifold with geodesic distance d. This presentations follows the one from (Lanzetti et al., 2025) on P2(Rd), adapted to (2( using results of (Erbar, 2010).
Differential Structure.
First, we recall sub- and super differentiability on this space.
16

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Figure 4: On the sphere, there an infinite number of geodesics between x and (here 3 are represented). Thus, the logarithm map would be multivalued.
Definition A.2 (Wasserstein sub- and super-differentiablltyy Let P2(M + R lower semi-continuous functional. A m M TM E (M, TM) belongs to the subdifferential 0 FFF( of J at μ for all E P2(M).
/ ((())) dr(x, v) + o (W2(uu))
F(v) 2 F()) +
sup
(27)
yEexP (v)
Similarly, 5 E (H,,TMM) belongs to the superdifferential ot F(u) of at μ - 0- (-F)
Similarly as on (R (Bonnet, 2019; Lanzetti al.. 2025), we say that functional is Wasserstein differentiable if admits sub- and super-differentiaas which coincide.
Definition A.3 (Wasserstein differeetiaability). A functional F: P2(M + Ris Wasserstein differentiable at μ (2
if 0- F(p) n n F(n) * 0. In this case, we say that Ww,F(i) F()) and it satisfies for any E P2(M), exp, (v),
is Wasserstein gradient of F at
F(v) = F())+ / (VwFF(w)(s) v) d7(2,, (W2((,)))
(28)
If J (M), then by Proposition 2.1,7 E exppt((v is unique and of the form 7 (Id, VV.--Veu.ll VVPu.) Thus, in that case, (28) translates as
=(v)=F(H) + /VV,, - (V))( du(x)) + (W22))))),
(29)
which coincides with (Erbar, 2010, Definition 3.1) (for the subdifferential, and up to a sign as they use e-convex maps, and we use Pp,, c-concave map).
If take 14 M ((tm")) 77, x (Gigli, 2011, Theorem 1.11), and thus
7, for y expr'(u) geodesic between v4 then necessarily (s. to )
F(u)= =F))==(()+10 + (WW,F dr(z, + (W2((,,)))
(30)
which implies d di = ((Vw, dr(z,)) e).
priori, the Wasserstein gradient is not unique. Nevertheless, we can always restrict ourselves to a unique gradient belonging to a tangent space whenever an Hilbert space. This is the case for μ absolutely continuous (Gigli, 2011. Corollary 6.6). So, we now focus on (M) C P2(M). In this case, the tangent space can be defined as T,P2(M)
PVP, DE C2 (M)} L (M.TMM This is closed linear subspace of L2, TM) and we can uniquely decompose any & E (M) as
o with Φ E (MM) and & E (MP2 (Rudin, 1986, Theorem 4.11). Since E (M) then by Proposition 2.1, the optimal y is equal to (Id, -Vpp.) with Pa,r a Kantorovich potential between y and In this case, it can be shown that
f (M(6), v) dr(z,v) /W(Ws), (x)) du(x)
(31)
17

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
since VPn. (MM (Erbar, 2010, Lemma 2.6). Thus the only part of the gradient that matters is o, and can show that it is unique.
Proposition A.4. Let F: (2(M) - R. Its gradient at μl E (M), if is exists, the unique element P22 0+ FŌ
Proof. See Appendix C.7.
Another interesting property of gradients belonging to the tangent space that they are actually strong differentials, meaning that they satisfy the Taylor expansion along any coupling, i.e., for any v E P2(M) and y E P2(TTM) such that * 7 = and exp# 7 v, Vw,FFl) E (MM) satisfies
F(v) = F(1)+ (Vw;F(W)))) 2) dr(x,, v) v) to (V/IVIEe | dr(())
(32)
Erbar (2010, Lemma 3.2) showed this property for couplings obtained through maps. We extend it for any coupling in the next Proposition. First, for μ E (2(M) fixed, we define P2(TTM),, == { P2(TMM / = μ}. For every E P2(TMM)p we define I == 11ol222,, v). and we further define its barycentric projection to be the unique vector field B(7) E (1, TM) such that for every E (H,,MM)
/ (5(x)),) o) dr(a, = (5(x),
(())) / = (E.)( (4())
(33)
(see (Gigli, 2011, Chapter 6)). Note that the barycentric projection satisfies (|B()))E>()) 1/
Proposition A.5. Let 0 F(j) (2MM). Then & is (extended) strong subdifferential of J at μu, i.e. for every (TMM),
F(exp 7) 2 F(H) + / (E(x). dz(2z.)) + o(1l/lu))
(34)
By symmetry of the arguments. it also holds for superdifferentials and gradients.
Proof. See Appendix C.8.
We now derive the Wasserstein gradients of well known functionals such potential energies and interaction energies. Proposition A.6. Let V M + be twice difserentiable with Hessian bounded in operator norm by L for all x E M, iec. ||Hsss M V V=))ll max |v ||Hess (x)ulll L. and V p. + f Vdu. Then V is differentiable with gradient V w. (p)) V.VV for any μ E (M(
Proof. See Appendix C.9.
Proposition A.7. Let W M M + R be twice differentiable with Hessian for both arguments bounded in operator norm, and μ + IJ W( ()dx())dp(() Then W is differentiable with gradient Vw.W(u)(r)) (V, (2, + <y, dp(y) for any E P2(M). a M.
Proof. See Appendix C.10.
We also introduce the notion of Hessian on the Wasserstein space, which will be useful to derive smoothness assumptions for the WoW gradients to be well defined.
Definition A.8. Let F (P2M) R. Let E P2(M). The Wasserstein Hessian of F at y exp,;' (v) for some (MM, is map HF,, TM + TM verifying geodesic H4 = exp,, ott>") 47 with 7 exp,"(v)
d² F(r)) (H,F, u) dy(x, v) for constant-speed
18

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Wasserstein Gradient Flows. A Wasserstein gradient flow of F P2(M) + R is defined curve H Me on an interval I, which is weak solution of the continuity equation
Jep = = ((,
(35)
i.e. which satisfies for any & (IxX M),
J.Sm (0)2)(1) - (VMPs(Z), Vw. F(u)(()))) due(z)dt = 0.
(36)
Usually, such equation needs to be approximated by a scheme discretized in time. A common way to do it is through the Otto (JKO) scheme introduced in (Jordan et al., 1998), which is the form
(3( + F(u).
Vk 2 O, p4+1 E argmin #EP2(M)
(37)
2T
Under suitable conditions, this scheme converges towards the Wasserstein gradient flow of F (Ambrosio et al., 2008; Erbar, 2010). However, this scheme is generally costly to compute, as it requires to solve an optimization problem at each iteration. In practice, is more convenient rely on explicit discretization, which can be seen as Riemannian Wasserstein gradient descent (Bonnabel, 2013; Bonet et al., 2025), which is of the form
Vk 2 0, Mk+1
exp pu (- (-TVwF((ue))
(38)
We note that this scheme can be obtained by linearizing the objective in (37). Indeed, if us E P2,ac (37) can be written as
Tk+1 = (MM 27
# [T< (ux(() + +()
(39)
Mk+11 (expoT x++1)##k
TT
Using the coupling = (xd.expoT) #bA M(Ak, (exp oT))) and that F(,x( is a strong differential, then we have that
F((expoT)*nn = F(ux)+ + (Vw,Flun)(t), T(x)) duu(x) +o (1/Iroolle dua(ce)
(40)
Plugging this lincarization in (39), we obtain
1
Tk+1 E argmin (EL((p.TMM 2r
| I'TIIE((M + (Vw,F(uss, L)
(41)
Taking the first order condition, we recover (38) as T = F
Wasserstein Gradient Descent. In practice, we usually work with particles, i.e. we start po 1 update each particle at each iteration h 2 0 as
" Ei= 6e and
Vi E (I,...,n), Ii,t+1 = exp 2. (WTV
(42)
for μi =
1 n El d In particular, for M = RRd, the scheme is obtained as
Vi E (1, (1....,n), n), FiA+1 Titk VW,,F
(43)
Moreover, if the functional F: (2() R has closed-form discrete measures, i.e., there exists F (Rdy" R such
that ( Ei Ōz.) F(x1,.... F. then
backpropagation on F and use that F(, E ) (I..k)
V, N,P(z1 FF t.).
Proposition A.9. Let F (Rd) + Ra Wasserstein differentiable functional, and F : (Rd)" - R such that for any
x (K),,.., In) $ An == (x E (R4ym 3i & j, = 2z} and ur >m 8n F(un) = F F(ri,..., Rw)) Then, for all
e nn..
Vw, FV
(44)
Proof. See Appendix C.11.
19

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
B. Wasserstein over Wasserstein Space
B.I. Function Spaces on P2(M)
In this section we fix P E (2))
Recall that we define the tangent space to Pa(M) at f by for every l E P2(M). We also define the larger tangent space (. where ((M,TMM the space smooth vector fields on M, i.e.
1. (..TMM (M)
(2M2(M = IV EE
P2(M) by P=M) smooth maps from M TM. Our goal is to rigorously define the space (P, (M)) and show that it is indeed a Hilbert space.
= T TM)
Let B C M open, then the map μ E P2(M) + H(B) is Borel, indeed, it is lower semicontinuous by (Ambrosio et al., 2008, Equation (5.1.16)). Thus, the map μ E X + M P2(Y) with X = (2( and Y = .M is a Borel map (in the sense of measure- valued maps), and the formula
P()) Jeaun J. ((,
(45)
defines a probability measure P on P2(M) x M (we follow the same reasoning as in (Ambrosio et al., 2008, Section 5.3)).
We then define P(P,TM to be the quotient of the space of measurable functions f P2(M) x M - T.M, such that 2(,,2) E T.MM for every (H,z) E (2(MM x M, by the equivalence relation corresponding to equality P-almost everywhere, and we equip it with the norm [[I defined by
[ == SPLc (
(46)
(we view f (2( TM) interchangeably as function with signatures P2(M) x M + TMM and (2(M) + (M + TM), hence the notation (u)). It is Hilbert space: indeed, if M is open set U C R" then since TU = U x R"" (P, TU) is a Hilbert space as it is the direct sum of n copies of the Hilbert space (P). and in the general case, we can show that (IP T'M) is complete by showing that Cauchy sequences converge, by using local charts and a partition of unity of M to fall back on the case where M is an open of R".
We can now define L2(P, (M)) as the space of functions f L2(P, such that f(p) T2,P((M for P-ae μ. It is closed in (P, TM) and is therefore Hilbert space. Indeed, if (fr C L2((,, (M)) converges to I E TM,,
lim / [1...Ce) (1. (w))le (P()) =0
(47)
x2 20000
This implies that, up extracting a subsequence, we have (lf ((n) 0 for P-ae 1²² But since the T2 are Hilbert spaces, this implies that (#) E T2,P((M) for IP-ac μa, and indeed belongs to (PP))) We define similarly (IP DDerP and show that it is a Hilbert space. This latter space (2M) is useful in that it allows us to define a notion of differential for Wy-Lipschitz functions on (2(M) as we will see in the next subsection.
B.2. Lipschitz Functions and Rademacher Property
For every smooth vector field w E [[M, TM), let (1SR))))ER be its flow on M, that is, the diffeomorphic flow solution of
((1,2) E x M, a (x)) = wo (o(ssw.
(48)
U/u = ld.
We recall the argument. For every e 0, we have P[][((n) fn (w))l 2 e] - as P([l(()) - fa(u))le230) 2 s
1 ʃ [f(p)) - fa (n)) 2 dP(). So, up to extracting subsequence, we may assume that for every n, P[llf(H) fn(1)) (())) 2
n = .. Then, we can check that the set Nw Un2x (w) - fa(4))lza 2 11 has null P-measure and that for any
1 E A, fn( + +(()) in (p)).
20

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
and denote n the map P2(M) + P2(M) induced by the pushforward by she.
The following definition is taken from (Emami & Pass, 2025, Definition 9).
Definition B.1. (Emami and Pass, 2024) We say that measure Po E (2(PP((M) satisfies the Rademacher property if for every Wz-Lipschitz function U P2(M) + R, there exists Dr. (PPo, (2TTP2(TMM) such that for every I(M,
U ( (.)) - U.) = (DP.U()),
lim 1-0
in (PPo).
(49)
Thus, every time we have reference measure Po E (2((P(MM) satisfying the Rademacher property, we can define for every Wz-Liipschtzz function U measurable section Dr. U of (2) that acts as "differential" of sorts, for perturbations given by smooth vector fields on M.
B.3. Wasserstein Geometry of (2(P2(M))
We recall that the WoW distance between P, Q E (M)() is defined as
/ (W) (T(,,))
Ww, (P,
inf rell(,,)) .))
ar(,,)).
(50)
A natural question is to find the conditions under which this problem admits an OT map. Emami & Pass (2025) showed it is the case for M a compact connected Riemannian manifold, for absolutely continuous measures w.r/ reference measure Po satisfying the following assumption:
Assumption B.2.
Po has no atoms
Po satisfies the following integration by parts formula: for any F,9 E (y)))2 and any smooth vector field w TMM there exists a measurable map μ (()( (2(M) such that
F(n), G(p) dPo(u)= = [ (eam F(u) (1, dPo()))
(51)
Po is quasi-invariant with respect to the action of the flows generated by smooth vector fields, i.e. for any smooth vector field w E T(M, T.M). Po and P = t are mutually absolutely continuous for every E RR and the Radon-Nikodym derivative
dPt s dr dP, @ dr
R.
rER
(52)
satisfies, for Po-a.e. μ1 L essinfre(e.)) R (n) for all 8, E with 8 S
These assumptions were first proposed by Dello Schiavo (2020), and we refer to (Dello Schiavo, 2020) for examples of measures satisfying them. By (Dello Schiavo, 2020, Theorem 2.10), any Po satisfying Assumption B.2 also satisfies the Rademacher property, which Emami& & Pass (2025) leveraged to show the existence of an OT map. In all the following, we fix a reference measure Po E P2 (P2) satisfying Assumption B.2, and when there is no ambiguity, the "differentials" of a Wz-Lipschitz function U will be denoted DU. Moreover, by (Dello Schiavo, 2020, Theorem 2.10 (2)). if E Cy))P2(M)). then its differential coincides with the usual Wasserstein gradient, i.e., DU = V.,
Theorem B.3 (Theorem 13 in (Emami & Pass, 2025)). Let P E P2 (P((Pa (M)) such that P & Po and Q E (Pa((M)) Then, there is unique optimal plan I, which is of the form l' = (Id, P with T Pa(M) + P2(M) satisfying T P = Q. Moreover, T is of the form T(p) = exp ( (H)) 4ts where U is a W3-conavve) Kantorovich potential for P. Q. In fact, for P-a.e. μ. Dr. (p) V.P.T(a) where (P,T(H)) a Gdl F.d--concave Kantorovich potential for p., T(p).
Note that the last two statements on the form of are not part of the statement of the theorem in (Emami & Pass, 2025), but can be found in its proof.
21

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Geodesics on P2 (2)) We recall that constant-speed geodesic between P, P,Q E Pa(P2(M)) is a curve + P. defined on [0, 1J, which satisfies Po P, Pi Q and for all s, [0, Ww, Pt) |t 1i--/Wil w/W,, (Santambrogio, 2015, Box 5.2)).
(,,((e (P, Q) (see e.g.
As we work on manifolds, we introduce similarly on (2(M) generalized inverse of the exponential map, which allows to characterize geodesics even when the optimal coupling not unique. The multivalued inverse of the exponential map between P, Q (P))))) is then
{ r (P(PP((M)), T = P,, P# P Q. /ʃ Ia dr((, = (P, (P. (53)
expp (Q) -
where for any 7 E P2(TM), d - and ((op(T) exp# 7.
Proposition B.4. Let IT E expF (Q). Then the curve + IP. = expom (lD"") + defines geodesic between and Q.
Proof. See Appendix C.12.
B.4. Differentiability on (2(P)(M))
We recall that by Section 3.2, I is Wasserstein differentiable at P. then the WoW gradient Vww, ((P) satisfies for any expp²(Q).
//(Ww., (Ww., (T) o) dy(a, + (PW,,((P,))
F(Q) = F(P)
(54)
Using this formula, we now derive the gradient of potential and interaction energies.
Proposition B.5. Let M be a compact and connected Riemannian manifold, F P2(M) R twice Wasserstein differentiable functional with Hessian bounded in operator norm for all ? P2(TM), i.c. (upppt)))) [|-]]_=1 [H, 1) I L, and F(P) = dP(p) for PP Pa(PP((M)). Then, F is Wasser- stein differentiable, and its gradient is V Vww, F(P) V.
Proof. See Appendix C.13.
Proposition B.6. Let M be a compact and connected Riemannian manifold, W P2(M) x P2(M) - R be Wasserstein differentiable with respect to each of its argument and with bounded Hessian in operator norm for all E (2(TMM as in Proposition B.5. Let IP E (2(P)(M)) and F((P) = [ v) dP(x)dP(v). Then, Vww, IF ((P) = (V.,,, + VwaaW(y, v)) dP(v)..
Proof. See Appendix C.14.
The gradients of the potential and interaction energies derived in the last two propositions are computed by showing that they satisfy the definition of the WoW gradients thanks to coupling arguments. Similarly to the case on (2(M), we expect that they can also be computed as the gradient of the first variation, which is a much simpler way to compute Wasserstein gradient of generic functionals. Let F Pac (P2(M)) - R. Then the first variation oF (IP) (2) + IR at IP is defined as the unique function (up to a constant) satisfying 8F
Relation with first variation.
F(P + +ex -F(P)
/
5F
lim
(P) dx, 3P
(55)
where fd and P EX E Pac (P2(M)) for small. Then, we expect that the WoW gradient of F can be computed as
5F
Vw F(P) Vw,
3P (P).
(56)
We leave for future work to show formally this formula. Nonetheless, we verify that holds for potential and interaction energies. Indeed, for V(P) J F(u) dP dP(u), we have 8% 8p (P) F since
V(P +ex) / F(u)dP(m) + Flu)dx(w),
(57)
22

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
and thus the WoW gradient derived in Proposition B.5 coincides well with Ww 3F W( dP (P())dP(v)
(P). Similarly, for W(F)
W(P + tx) - W(P)
7
+t // w(n dx(e)d(())
+2 /I dx(m)d(l)))) -w()) w(P)
(58)
/ ( WWln dp(N)+ + / W(y, a) dF(w) dx(x).
SP Thus, the first variation is 6 (P) = J W(a, v) dp()) + WWW, n) dP()), and its Wasserstein gradient coincides well
with the WoW gradient derived in Proposition B.6.
Relation with Euclidean gradient. We provide now the analog of Proposition A.9 for WoW gradients, which we use in practice compute them.
We fix here number of classes C' 0 and number of samples > 0, and we consider the class of (fully) discrete measures of )R defined by
C Spzer
Px = C
XE (R4) Nx =
1
Sag, x" E (Rd)².
(59)
n2 i=1
c=1
We also define the space
X {x (Rd)CxI| x² & A, and Vc + d' Px + Hxx
(60)
where A, := (x E (R4)a 3i + j, = 3 is the generalized diagonal of (Rd)", Informally, X is the space of vectors such that the empirical measures Pxe in the support of I'x are all distinct, and are each supported on distinct points of Rd Proposition B.7. Let F P2(P22 - R a functional, and F (R2))xn + R such that for every x E X, F(Px) = F(x). IfU U is Wasserstein differentiable at Px and F differentiable at x ffr some x E X, then for every c E (l,..., C.) and
i e nI.
Vww, ((Px)(Hx))))))) ( CnV,,
(61)
Proof. See Appendix C.15.
B.5. Convexity on P2 (P2(M))
In this section, we focus on (2 and we show the convexity along generalized geodesics of potential energies and interaction energies. Hence, Proposition 4. can be applied to these functionals.
We recall that on P2(R)). generalized geodesic between μ4, v E P22(Rd) is of the form + pt ((1 - f) 1,2 + to 1.3) #* with y E IT(,, μt, v) such that n 1.2 y E II ()) and 1.33 this curve satisfies for all & [0.
# E v). Then a functional F: P22(Rd) -> R X-convex along
(1- +FF() - M( -
W2
(62)
2
When 17 E (Rd)), by Brenier's theorem, there are OT maps 7 between v and and TỤ between and v, and the generalized geodesic translates f = ((1 T) + (T²) n
We define similarly generalized geodesic between Q, O E P2(P (R") as t + P. ((1- (r _ 7 where T E (IP, Q, 0), x 1.2T (P, Q) and # 1.3 T (P, O). We provide sufficient conditions for potential and
interaction energies to be X-convex along generalized geodesics in Pa (R*))
Propusition B.8. Let X > and F: P2(RR) + R be X-convex along generalized geodesics of P2(R'). Then, the potential energy F(P) = F(u)dP(( (μ) is X-convex along generalized geodesics on Pa (2
23

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Proof. See Appendix C.16.
Proposition B.9. Let W P2(RR) x P2(RK) + R be joint convex along generalized geodesics of P2(R).. Then, the interaction energy W(P) = 2/ W(u, v) (P()))P(v) is convex along generalized geodesics P2 P2(P2.uc (R4)).
Proof. See Appendix C.17.
We can also show that F =Q+ 1 Ww, (Q, P) is 1-convex along particular generalized geodesics, which have anchor point P.
Proposition B.10. LerP, Pa
Define the generalized geodesic + Pr = ((1-t))z)+eer)) )
where I' [P,, Q, 0), 12 E (P, Q) and T E [I, 0). Then F Q + 2 (Q, P)2 is 1-convex along this generalized geodesic, i.e, it satisfies for all E [0, 1}.
F(P() & (1 ()) + (0) t(1 -
W W., (Q, 0)2.
(63)
2
Proof. See Appendix C.18.
In particular, Proposition B.10 is the main result allowing to show Proposition 4.1, which can be applied for A-convex potential energies with À 2 0 (or more generally, for any À E R such that + X 2 0, see (Ambrosio et al., 2008, Assumption 4.0.1 and Theorem 4.0.4)), and for convex interaction energies.
C. Proofs
C.1. Proof of Proposition 2.1
In the forthcoming proofs, we will make use of the following selection theorem, whose statement can be found in (Villani al., 2009, Chapter Bibliographical notes):
Theorem C.1. If A + is a Borel surjective map between Polish spaces (L.e. separable complete metric spaces), such that all the fibers (y)) y E B compact, then ff admits a Borel right-inverse.
This allows us to prove
Lemma C.2. There exists a measurable selection 8: ::M² + TM of the map
= (xx,v) E TM, ((x,)xx..(v)) = [|m|l=} f: {
MxM (I, (x))
(64)
(z,
Proof. First, A C TM a Polish space, as a closed subset of TM which is Polish. Second, for every (x, E M the fiber f (2. (x. y) is compact. Indeed, if we let {(x,, 2n) ] S f-'(z,, (z. y). then we have for every n, In = z. (xp) and ||v,||= d(xr, y), so by compactness of the spheres in the tangent space T,M, up to extracting a subsequence there exists v M such that v, + v. and by continuity y (x) so that (a, v) f- w). We can thus apply Theorem C. to to deduce the existence of s.
Now, we can prove Proposition 2.1. If y E exp, (v), we have exp)) E I1( v) by definition of exp,, (v). Furthermore
S.L. d(z, d) exp) (e = f d(I, (().eexp.()))d(xx...))
(65)
TM
[ JTM
|10 v) Wz(H.v)
(66)
so this transport plan is optimal. In particular the inequality is an equality, and we find that d(z, expx(v)) = ||0|| le for y-a.e (x, v). The map thus well defined. To show that it is surjective, take y E To(l, v), and set == s(T1, 42)
24

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
where -M² + TM is the selection map defined in Lemma C.2. Then 7 expx (v). Indeed, by construction, #7 = 11( vi), and also
WH(H JSL d(x, () f d(z, (x) = Sr (1012
(67)
M
This proves the surjectivity of the map.
Now, assume that is absolutely continuous, and that M is compact and connected. Let E exp, (v) and f == (m-M, , exp) #7 E (1, v). By Theorem A.I, 7 is of the form (Id, T) where T is the unique optimal transport map from to v. Furthermore 7' itself is of the form T(x) exp, (-v(()) where yp is a e-concave function M + R (in fact Kantorovich potential for the pair μ. 1). Furthermore, we have, for p-a.e. z E M, T(x) belongs to
0x((x) = {y E M, ppx) + ( = d(x, 3)
(68)
2
This implies, by (Gigli, 2011, Theorem 1.8), that for y.a.e. E M.. (T()) S -0+) where 0 p(x)) is the superdifferential of s. However, since M is compact, &p is Lipschitz and is thus differentiable almost everywhere, so that 0+x(x) = (V)( almost everywhere (in particular this holds fa-ae as μ is absolutely continuous). From this, we conclude that for 7-a.e. (2, v), we have T(x) = expx(v) with exp, (T(x)) = { Vv) so that = -V)(()). Thus, we have proved that
7 (ld, -V
(69)
This finishes the proof.
C.2. Proof of Proposition 3.1
We first state and prove another selection result:
Lemma C.3. There exists a measurable selection s: (P( + P2 of the map
( (P(MM)) E exp M, (exp# + P2(M) x P2(M) (T T exp# 7)-
f:A=
(70)
Proof. We first prove that A A P2(TM) is a Polish space. All we need to do is to prove that it is closed: indeed, since TM is connected Riemannian manifold (with the Sasaki metric), (P2(TM). W2) a Polish space. (See (Ambrosio al., 2008, Proposition 7.1.5). Similarly P2(M) is Polish space.) Note first that if y E A, is supported on the compact K ((x,v) E TM, ||o||= S diam(M)) S TM, || d(æ,, expr, (v)) 7-almost everywhere. Let E () in the W2 metric. Then is also supported on K, and for every n, since
J S A converging My
Th E cxp exp,My, (exp, 7n), we have
(T M / [101Bahha(0,00 (3 na exp# 7m))
(71)
Letting + J0, we find
/ Ielldd((,,)) W exp# 7).
(72)
Indeed, W2(1n, 7) O implies weak convergence of nh to 7, and thus of n MM In and exp# Yn to respectively T and exp# y, and since M is compact, weak convergence P2(M) is the same as W₂ convergence, which in turn implies the convergence of the right-hand side. The left-hand side converges similarly in virtue of the weak convergence of 7n to
(recall that they are supported on the compact set K on which the function (x, v) - >|vi| is bounded). Therefore, (72) implies that 7 E exp xx (exp# 7), so that 7 E A, and
thus closed in P2 (TMM). and a Polish space.
Now, we prove that the fibers of ff are compact. Let p, V E P2(M) and {1 C f v). Since they are supported on K, and P(K) is compact by compactness of K, there exists ? E P(K) C P2(TM) such that, up to extracting a subsequence, Y/2 converges to y, both weakly and in the We metric. We then check as above that y E A, with =MM, = and exp= v. This proves that the fibers of ff are compact. Now, the existence of follows again from Theorem C.I.
25

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
The proof of Proposition 1 is pretty much similar to the previous one. If expp (Q), we have (oM ((ep)) # [I(P, Q)
by definition of expp" -1 (Q). Furthermore,
JALMS Walk, p0x0) ) (p,
J Wa( 7, 7, exP# ()d[())
(73)
P2(TTM)
(74)
JeLams I.x d(e, ex,, (v)) (1))
(75)
J fov I1ll2 v)d Ww.(P,Q)3
(76)
so this transport plan is optimal. In particular all the inequalities are equalities, and we find that for a-a.e. exp,gr, (exp* 7). The map is thus well defined. To show that it surjective, take I E (P, Q), and set i == s(m,, #2) #TI where s: P2(M) +> P2(TM) is the selection map defined in Lemma C.3. Then i E expy (Q) as, by construction, (oM, ) ) = T E II(P, Q), and also
W r, exp# 7) [l hence E exp mgy
(w, / SPACMP Wd(
(77)
JPSTTM Wa( y, exp r)d (7)
(78)
= -Saces L
Jrm
o)dl (7).
(79)
This proves the surjectivity of the map.
Now, assume that IP is absolutely continuous with respect to Po. and that M is compact and connected. Let I E expp (Q) and J = (o-M, (exp) the unique optimal transport map from P to Q. Furthermore, for --a.e. % have
y J (P, Q). By (Emami & Pass, 2025, Theorem 13), I' is of the form (Id, T) where T is
exp,,, (2y), with Hey x My. (Id, + and P is concentrated on absolutely continuous measures (as P < Po). we also have that for Taee 7, My is absolutely continuous and By = T(Hy). Thus, by Proposition 2.1, we have = (Id, - 1wm = (u + (Id, -VFp..(()) (7))) for -aa.c. 7. Therefore, we have proved that
and Vy= == exp# 7. Since
[= (u (Id, - (p.T(())) # MT (nx + (Id, VPuT())))) P.
))
(80)
This finishes the proof.
Remark C.4. As side note, there exists more explicit expression for the unique IT exp; 1 (Q). Let indeed U P2(M) + R be Kantorovich potential for Q (i.e. 1 W3-concave function solving the dual problem). Then, by the same reasoning as in the proof of (Emami & Pass, 2025, Theorem 13), for P-almost every μ E P2(M), we have ViuTTi DU(H). Thus, we have
T = (p + (Id, -DU()) P.
(81)
C.3. Proof of Proposition 3.7
The proof is inspired from (Erbar, 2010, Proposition 2.5) and (Ambrosio et al., 2008, Theorem 8.3.1).
First, fix yp E (M))))(M)) such that (p() := F (ʃ Vidp,..., S Vmdu)) with F (RRM) and Vi,..., Vm C2 (M).
Let us define H: P2(M) (2(M) - IR as
(e())=((e) Walmee)
if
((,
(82)
if u.
We show in Lemma C.5 that H is upper semicontinuous. We want to prove that -> Pr(y)) [yd,, is absolutely
26

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
continuous and bound its metric derivative. For every 8, E I, let ae (P, P.). Then
1 h
1
/
PP (++n(p)
Pa(p)]
(9(N)) - s(v)] dr (,, [ho] P(M)
1
W)( dr ((.
[b]
(83)
Ww, (Ps+hn PP))
SpLnor H2( dl' (t,
[/]
Now, we have I' Tathes (Id, PP when - O, so, since is upper semicontinuous, by (Ambrosio al., 2008, Lemma 5.1.7), we have
[ JpmMp (p, v) dr (U,U)) SPLM) (P dP,(u) = SPLM) [VV,, dPs d/)().).
lim sup h-10
(84)
and thus + (x(()) is absolutely continuous, with metric derivative bounded from above by
(85)
Let now sp E Cyl(I Pa(M)), Q = P2(M), and x J, dP.dt. Then, for any interval J C / such that sp((F)) C J P2(M),
aep di = lim 1/
gtt, μ)) p(t - h. ax(1,
h--0
h Pe((p)) - Prtt((Pe) dt
lim sup h-+01
[Pr(yo,) -P-n((pr)l dt
lim sup
(86)
[h]
ho h-+0
|P' (P'((()) V di
IP1((1) au I.Vissplliees dt.
From this, we infer that the linear form
L(Vw,))==- / da
(87)
(PA..P2(M))
well-defined and Lipschitz continuous. In particular, there exists v E (Vw,y, & E Cyl(T x P2(M))) such that for every to E Cyl(I x P2(M)).
L(V w, p) (v, VwsP)
(wele), Vw,e(())) dP, (p)ddt
(88)
and we have the continuity equation.
Moreover, for a.e. t 04 (Vwai, &o E Cyl(I P2(M))} 12(F,.PP(MM) and for every interval J I, there exists sequence Pn), supported in J P2(M) such that Vw vll. For all n,
(PIC0 de) $ (W dt)
I(Vw,Pn))
(89)
and letting n + 0o, we find
/ | Iulllaaee. dt
at) "(,leelBBeerrr dt)
(90)
J IoulIEare dt (p)) dt.
27

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
We thus conclude that for a.e. LEL I,
( S (P)
(91)
We now show that H is upper semicontinuous.
Lemma C.5. Let up E Cy)(P2 The function H: P2(M) x P2(M) - defined as
[e(a)-p(r]]
fu*v
Hu,,v) = ((WWLe(W)lL>
Wa(ue))
(92)
ff=
is upper semicontinuous.
Proof. We want to show that the function H P2(M) P2(MM - defined by
1p(n)~p(()
i v if μ =
H(,,))
Wa(a.e)
(93)
upper semicontinuous. Let yr, E P2(M), and let (14))et0.1) be the constant speed geodesic from ju to v with velocity field v4 E (..).. Then, we have
J.¹ d di ((e) di
plu) (()))((x)=
(94)
/./ d F (J (Vidue..... / Vindue) de d
(95)
OF Oa:
(f [Vvidee.... /Vndvua)
dt / Vidue dt
(96)
Σ OF
(ʃ (Vvidue.... / Vindua) /
/ (VV., vi) dudd
(97)
Oxs
(Vw2)((e)) ve) (
(98)
$ V V. [oell/a dry S. WWW.eP(HD))
(99)
< Walus Waluresy/ u// / (I
(100)
Hence,
H( s / S. W < 0
(101)
as the right-hand side is finite because the F, V, have compact support. Note that this inequality also true when μ = 1/ as Ne p for every t.
Furthermore, notice that the map uu Vw, is continuous. Indeed, we can check that we have
f(1) fo ( Vidy.... ((vidd
Vindy..)) ds(xx)
(102)
28

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
for some Lipschitz function G with Lipschitz constant L, that if nt" us, we have
/
/ Vindu", ,0) du / G
(/ Vidpi....
/ Vndue du(x)
1f(4*) f(()
G
Vidy", ....,,
(103)
f G
( Vidd"
/ Vindui" ;00 dnm / G (Vidp..... /Vadle du"
(/vides...
(104)
f G
( Vidp.... / V (1)(())= / /a() Vidge.....
(Viidds...)) Vindje,2)
dd(x))
(105)
f Vi d(d -1) - + +15c (ffViee..... Vindenz)) -1)
u))
$ L
12) (1) (2))
+
(106)
Now, if 2² - f and y² - μe, then μ - μ for every (indeed W2(Af, :) S W2(Hf". W2("", = tW2(v", p²") + W2(n", /4) -> 0), and thus by what precedes (W, + [Vw,s(u))lz (2(a) for every L. Therefore, by (101),
we deduce
limsup H("", ,²) S (Waa(())L = H(,,p).
(107)
This proves the upper semicontinuity of H.
C.4. Proof of Proposition 3.8
Let IP (P)P))M)) and define Ps (P(P)(TM)) (T Pa(P2(TMM)) AT P]. Fix (2(P2(TM)) p. we define
// 11013)0(0,
(108)
By the disintegration theorem (see for example (Ambrosio et al., 2008, Theorem 5.3.1)), there exists P-a.e. unique family of probability measures (T (Tu)uEPa(M) such that I is supported on (P)(T and, for every measurable test function f P2())) R+,
[ //
(109)
We can use this family of measures to define the barycentric projection of I: Definition C.6. The barycentric projection of I is the vector field B(T) L2(,, TM) defined by
((T) = B(G)dT (() E (M,
(110)
Note that we work here in the space (P, TM). which is defined in Appendix B. is a larger space than L2(P)TP2(M)), with which should not be confused. The barycentric projection satisfies the following properties: Proposition C.7. For every & E (P,TTM), it holds
// (5(mm) (x), o) u(()) = ((GN) -
((,)(T)) L((P)
(111)
Furthermore [IB(T)]))(()) S ||]|.
Proof. For every & E (M), we have
//
v)T()) /// (() ()) v)du (()))P())
(112)
/// (E(m)(z), da u (p)
(113)
// [f(CEW)) (E()), (P (())))))) (7))P(())
(114)
= ((EW))BBBTTY))))))))P = ((T) (P()))
(115)
29

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Furthermore,
(B())) // / (())(r) 260 du(z)ap(()
(116)
< /// ((T
(117)
$ // (w
(118)
- //] |01|2d wdds μ (7)
(119)
$ ʃʃ [11lBdy(z, vddr (v) ||[||p
(120)
where we used (1(7))1)2(0)) & 1111 = f (r,,) u) to obtain the fourth line.
We now show Proposition 3.8.
Assume by contradiction that & is not a strong subdifferential of F at Then, there exists > 0 and sequence and, for every
(I J (P) such that en |/////= 1-
-F(P)- - /f (5(T v), w)dr, (.()) S -ōen
r(P,)) - -((P)
(121)
with P,, == 2 Now, for every n, fix Fn E expp (PP.). Since & E ((P)) there exists > such that for every n N,
8 F(P,)) -F()) /ʃ u)), (n) -ww, (Pn, P).
d),
(122)
Denoting Vn : B)T,)) and 6,, : B((x), we have, combining (121) and (122), that
8 Ww, (P,, IP)
(c. v (P(P) den 2 (e)
(123)
2
Furthermore, we have Ww, (P,, P) S Ens since
(W, P)2
<
W W3( exp exp4 r)dT (7))
(124)
< //
(125)
< // |1 v)dT (17) = ||_ 2
(126)
Thus, we find for every n >> N
(5. 2, 4 L2(F S --
(127)
Now, since / & a = en and (by optimality of F,,) 11 ( I = = Ww. En for every n, it ensues that, up to extracting a subsequence, there exists V Φ E 1P2 T.)) towards which e,- -, and --1 respectively converge weakly (PP TM). Therefore, dividing (127) by En and passing to the limit, we find
(E. - L2()) < 2
(128)
Now, fix F E ()))P) By applying Lemma C.8 to F, and F,, we find
/ FdP, / FaP (Vw, F, + +(()),
(129)
/ FdP = ( FdP + FWW,F,W))Paa + 0 +
(130)
30

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows Subtracting these two equations, dividing by en and passing the limit, we thus find
(Vw,F, + L) 0.
(131)
and this holds for any F Cyl(P2(M)). However, by assumption, & E T2P2 (P2(M)) and we recall that
L22 (PPIPP(PP((M))
TpPa(P2(MM) = (V..FFFe Cy)(P2(M))
(132)
This implies immediately that (E, 4b 1.(P)= = O, which contradicts (128). Lemma C.8. Let F Cy)(P2(M)), then, for every E P2
P2pP2PP2pp (P(((M)) and T E (P((
o# ( ( FdF // (TW, (
m)) $ C
(133)
for some constant C depending only on F.
Proof. Let F E Ce (RT) and Vi,.. Vn E C (M) be such that
F(R)= =F F VVide......Incce
ME P2(M).
(134)
Since F is compactly supported, there exists C > 0 which only depends on F' such that for every æ, RT
|F(x + h) - F(x) - (VF(2), A)] s CIINI|?
(135)
Fix 7 E MM), and let μ == क y and v == exp # 7. Since the V, are compactly supported, we know by Lemma C.12 that there exists some constant L > O, which depends only the V,, such that for every i,
/ Vidv vidu-Jo J Vidu - (VV(()). o)dr(e, w) 1
(136)
Now, we have
F(v) - F(v)-F(M) F(m) / (Ww,) 0)2
(137)
F(v) F()) Σ OF
m
/ (VV
1) 2)
(138)
i=1 Oz,
Σ OF dzi
m
(fVid= f /vddy)
F(v) F(m)
(139)
m OF
Os, (f Vidu
J |vddu (VVVli//VV(())))))))))))))
(VV.(E), 0) ())()))))))((())))))
(140)
/ V,du /
12 + Σ / Vode Vodu-
/Vodv - Vedu. - (VV((c))
SC
Vidaa
(141)
/ Vidv
/ |Vida 2 + Cl
(142)
S SC 2 / Videe /vilu- - Vidua /Vidu- (VVV(o)) + (
m
2 + +C
(143)
< C Cl
(144)
where we used (135) in the fifth line, with x; = Vidu and h2 = f Vdde- f Vida, we used (136) in the sixth and eighth lines, and we used the Cauchy-Schwarl inequality in the eight line. Throughout the derivation, C denotes constant that
31

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
may change between lines but which only depends on F and the V,. In particular, there exists a constant C which only depends F and the Vi such that for every 7 P2(TM) with μ # and
| S diam(M).
F 7) - F(n)) - / F(u)()), 2) dy v)) <hyll
(145)
Now, if E P2(TMM is such that 1| > diam(M with r == , let v = exp# y and 1y exp, (v). Since
||l| = W2(u. v) I diam(M), this implies
F(v) F(u) - (()(F())()), d7
< F(v) - F(()) - (V ()) d))
(146)
/ (V. F())(o), b)22 d(n 7) (z. w)
(147)
$ <Clwl + 1(V,, B(n) - (()))
(148)
<CImm C + (1(())))))))
(149)
S <0[ + 11vll+ + (|)ll
(150)
s <(1]> (bllp)
(151)
C1
(152)
where we used (145) in the third line, and we obtain the fourth line using the Cauchy-Schwarz inequality and the fact that to0 (since the F and Vi are compactly supported). Again the C's denote constant
(MM) F(,)F(p) depending only on F and the V; (and diam(M)). Thus, we have shown that there exists constant C' depending only on F such that for every 7 E (2TM)
F(exp 7) - F( / (V - (V (1)
(153)
Hence for every I E (P noting P: o and Q = +x T.
/ FdQ -
/ [FaP=
(V.. ) (/ i) widd (7))
(154)
/ Flexp##)) 7) -F(m (V, (z), 8) dy ())
(155)
S C / [11122,, d())) = C CI
(156)
This finishes the proof.
C.5. Proof of Proposition 3.9, and existence of gradients in the tangent space
First, we prove Proposition 3.9. Let &1. E F(P) (P))nP) (T they are also strong gradients of F at IP. Therefore, letting & = {1 - E2, for every v E
Using Proposition 3.8, we know that
we have
/ ((())
= o(| (I(())
(157)
that is, (E, (P)) = = (([]V1((Z(P))) Considering VP = ef, we obtain e||5]12((3) = ole) that is [Ellia()) o(1), and this implies = E2 0, and this finishes the proof.
Now, natural question is to ask whether there is gradient in TP whenever there exists gradient & F-F(P) n (P). While a complete answer to this question is out of the scope of this article, a partial answer can be provided using results laid out in (Dello Schiavo, 2020). First, we consider the following assumption:
Assumption C.9. (Smooth transport property, (Dello Schiavo, 2020, Assumption 2.9)) We say that M satisfies the smooth transport property if, whenever v E (P(M) are absolutely continuous with smooth nowhere vanishing densities, then there exists smooth optimal transport map T M + M from μ to v (in the sense of Theorem A.1).
32

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
This is relatively restrictive assumption M. By (Dello Schiavo, 2020, Theorem 5.9), it holds whenever M satisfies the strong Ma- Ma-Trudinger- Wang condition MTW(K) (K) for some K 0 (we refer to (Dello Schiavo, 2020, Section 5.2) for further details). Under this assumption, we can prove the following result on the existence of gradient in the tangent space:
Proposition C.10. Assume that M satisfies Assumption C.9, and that P satisfies Assumption B.2. Then, if admits WoW gradient at P (i.e. a (P)) n (P) is not empty), then it admits WoW gradient in TpP2((2(M)) (i.c.. 8
is not empty).
Proof. All we need to do is to prove that for every & (PP)) e Pa(2(MM) and U expr
|| (5(T (()) o)), o)d 0.
(158)
Indeed, this ensures that if & E ò F OP)(P) WoW gradient of IF P.. then its orthogonal projection on TpPa Tp (P)( is also WoW gradient.
We thus fix Ç E TpP2 (P2(M) Q (P((MM) and [ expp' Since P satisfies Assumption B.2, by Remark C.4, [ is of the form (u + (Id, - -DpU) o where U' is a Kantorovich potential for the pair P, Q. However, since M satisfies Assumption C.9, and U is Wy-Lipschitz (by (Emami & Pass, 2025, Lemma 12)), (Dello Schiavo, 2020, Theorem 20())) implies that DyU E Pp (P2(M) (as a limit in (,, 2 (M)) of functions of the form Vw, F ())(P) so that
/f ()), d) dd) (7) = -- - ((W))) = O.
(159)
This finishes the proof.
Note that for this proposition to hold, P must not simply be absolutely continuous w.r.t Po. but must itself satisfy Assump- tion B.2. According to (Dello Schiavo, 2020, Proposition 5.2), this is the case whenever, for instance, F = P2PP where xp is a strictly positive Wz-L.ipshhizz function on P2(M).
C.6. Proof of Proposition 4.1
We aim at applying (Ambrosio et al., 2008, Theorem 4.0.4), Since by hypothesis, F is X-convex along the curve IP = ) for I' E II(P, Q, D) and P E P22 (P2.acc (R"). we need to show that G : Q + _ Ww, (Q, P)2 is 1-convex along Pr (see e.g. (Ambrosio et al., 2008, Lemma 9.2.7)). This is well the case by Proposition B.
TT)) ([1--)rıq + + #" #
Now, let P& E (P (R2)) and J(P) = 27 Then, we have
1 W (P, F(P) the functional solved at each step of the JKO scheme.
1
J(P.))
Ww, Px)3 F(P.)
27
C(P) + F(P)
< ((1 - t)G(() + +G(0) - t(1 - t)
W (Q,
(160)
2
At(1 - 1)
+ (1 - (1-t)F((C + (0)
Ww, 0)²
2
Ar +
= -1- (QQ + (J(0)
t(1 - -t)W., (Q, 2)2.
2r
Thus, we conclude that J satisfies well (Ambrosio et al., 2008, Assumption (4)..1)), and then apply (Ambrosio et al., 2008, Theorem 4.0.4).
C.7. Proof of Proposition A.4
Let E.E' E (M) n F+F(() n F(H). By density, for any e > 0, there exist Pe oe E C% such that - (Pell $ 1 and 11g' Ve''I1 S
We rely on the following Lemma, which provides an OT map for any & C (M) for s small enough.
33

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Lemma C.I1. Let μ P2(M) and dp (M). Then, there exists 8 such that expx(+VV((z)) is an OT map between μ and (xpp((VVi)) # for all E] 8.
Proof. Suppose v * 0. Let E 0 and maax, |[ exists as e is supported on a compact. Then, for any E (-5,8), I | S 8|| b e. Then, by (Villani et al., 2009, Theorem 13.5), sib d² /2 convex, and by McCann's theorem, exp V is OT map.
By Lemma C,11, there exists 0 such that = (Id, &V E (x), and y' VVd) (xx,77(( with (expo(>V<pe) 1 and v' (exp (expo(sV ue Thus, have using the definitions of sub-differentials that
F(v) F(a) + { (z)) du(=)+ F(v") 2 F()) + + (z) ( dp(x)) d(s).
(161)
Likewise, by the definition of the super-differentials,, we have
F(V)) F(n) [ dp(x) ((s) F(v) F(u) + + Vp) du(x) ((s).
(162)
Dividing by > O and rearranging the terms, we have
F(v)-F(m) (E,
F(w)-((1) (g', (1)
+1)
(163)
F(w)-((v) (-E, (u +
(())-(()) 2 (-E', (M)) +1())
Summing them, get,
(E - C' VAPETAATTTI + (g' -,, V.dd) + (5 - Vis V + o(1).
(164)
Then, we have
5 / ||5 - 5'IliP - 2(E 5', Vee VVL)) + VV - VPllliA(uTTM))
= 1gg - -5' - -VVe - V
(165)
< ||5 - VePelle((ATTM)) + 1l5 - V
x E.
Taking the limit + 0, we conclude that { g'.
C.8. Proof of Proposition A.5
We assume by contradiction that is not extended strong subdifferential. Then there exists sequence (rn 100 C P2)T))), and d >> 0 such that e.. = || - 0 and for every denoting pa exp# Yr.
F(um)) F(K) ((E)) ()
-ōEn.
(166)
Now, let 17, E <exp7 (.n)) Since & is a subdifferential, there exists N such that for every n >> N.
F(um) F(H)) - (6(x)), (0
u)) (x, 2 We(U,
8
(167)
Since, by optimality of nh have I| 1| W2(H)Hm) s 11 11lp- En. combining these inequalities, we find
/ u)) (2, (E(x))((())))))))))(<)) ()) -
8
(168)
34

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
that is
8 (5, 9n (2) S 2"
(169)
for every n > N, where we have defined Vn ((7)) and 2n := (mn). Since we have || I Imm 1mlll en and likewise 1, (()) S [| S En. up to extracting subsequence we can assume that there exists v. Tb E (1, TM) towards which V and en 22, respectively converge weakly in (1, TM). Thus, dividing (169) by én and passing to the limit, we find
- 2 (E, T 1)w
(170)
Now, fix some & E C. (MM. By Lemma C.12, we have
/ pddu = =f pdd++ + / (VVp(t).))2ddn)k +([mma)1))
O(1lmal)))
(171)
/ pddutt + (V,, (2) + +(()),
(172)
and similarly
I pddp = = pdu++ (Vip, Vn)1a() + ((
(173)
Subtracting these two equations, dividing by en and passing to the limit, we find
(V.. ()
(174)
and this holds for any oo (MM) However, by assumption, & TT,P2(M (Ve, & C (M.TMM implies immediately that (E, 4 (P())) O, which contradicts (170).
This
Lemma C.12. Let sp E C20 (M), then, for every μ E P2(M) and P2(TM)pp
1 / pd pd(oxp 7) / odu - (VVels)))),dnzz,,) v) S [
(175)
where L = (.. | Hessu p(x)((]] <. <0..
Proof. Let (x, v) T.M. Applying (Boumal, 2023, Exercise 5.40) to the geodesic given by c(t) (x), it ensues that there exists E (0, 1) such that
(v))= + (Vp(x),)) v) (Hess (Hess ((t)))))(()(((()))))))))))))))))
(176)
so that, since [|(((t)ll(()) = ||0|a.
le((epp(u)) - p(()) (VV(o),,
u)al S 11u1
(177)
This immediately implies
|f pd(expp 7) - f palu-) (V)(()) = / -(())
-v())
- ( - (Vp(())
(178)
1 11112d((0,0) 1 111111
(179)
35

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
C.9. Proof of Proposition A.6
Let u E P2(M), and y E exp, (v). For any M, v E T.M, let note (2, exp,, the geodesic starting from with direction By (Boumal, 2023, Exercise 5.40), we have that there exists t [0. such that
V - ( + (Hesss ( (c,,,) (t), (t),()),,,,000
(180)
Then.
V(0) (())) = / (t))
- V (x)) d)
(VV(())) 1 (Fr, (10 dr(2, u)
(1)c...
(181)
(VMV(s), ())(,,,)) v)
(C,ssV(c,, (1), (0) dr(
Moreover, using that V has bounded Hessian and that geodesics are constant speed and thus satisfy (9)lec)))))))) (0)ll, = l|llla (Lee, 2006, Lemma 5.5), we have that the last term is bounded by W v) as
/ʃ (HessV (es, (1),,())) dg, (1)) (0)) d7(2, w)
L / IIC.. (t),,(())) dr(z)))
(182)
= L / [1ul2 dr(x,,) = LW3(H,)) v).
Thus, conclude
+ / (VMV(s)),))) dy(,,)) + (W2((1))))
>(v)=(()) =
(183)
Now, let us verify that V (p)).. We denote by PT the parallel transport between T.AM and Ty,M along the geodesic between and (see (Boumal, 2023, Definition 10.35) for the definition). By (Boumal, 2023, Corollary 10.48, 3.), V having its Hessian bounded in operator norm by L is equivalent with having for all x, y E M, for all (z, v) E TM,
(x) - PT (x) +o V (exp, - L||
(184)
Thus, let E P2(M), o some origin, and E ( Then, we have by using sequentially the definition of 7. || + yll < 2||x||2 + 2|ly/l?. (184), that for any (I, v) E supp(?), (v) = o and PTo-w is an isometry (Boumal, 2023, Proposition 10.36), and 7 E P2(TM).
V V.VIII
V VVV(()) dp(z)
/ V d7(r,)))
2 [V - PT (x) V V (v) (exx)))))) dx(x,))
(185)
2 / [PPTepp 122 d>(x,))
< 2
< +00.
Therefore, we can conclude that V(4) = V V by Definition A.3.
36

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
C.10. Proof of Proposition A.7
Let u P2(M), and 7 (v). First, we recall that the product space M M is Riemannian manifold with tangent space TM x T.M. For any (a, v), (z', r') E TM and note (x((t exp,, (tp,,tt), (r exp,, (te') and Cz,z' C (t) (cz,o (t), Cx(,e( (t)) the geodesics starting at (æ,x') with direction v'). Then, by (Boumal, 2023, Exercise 5.40), there exists E)0, 1 such that
((espa(()).xxp_ W(r,='++ (VW(,,))) (WW
(186)
(HessW (Ca, Car (1) (0)
2
Moreover, by the same argument as (182) in Proposition A.6, we have
V/ (Hess W (Cr,v,a', (t))le, (t), (0,)) (t) dr(2, v)dr(e)) 2L, (187)
Then, we have
W())-(()) - // W(,, du(v)du(s)) // W x') du(x))p((²))
W(6) - W(1)
= // (x) - ((xx,1')) dy(a, o)dr(d'.0) v)
1
+ 2 (ess (t) (a.
d,(2,, uddx(ee)
/ T:W(az2/dddee)))))..d... / V2 d7(x', + o (W
(VIW(z,z')+ V dp)(=))) u)( (W22(H))))
(188)
Now, let Ww, = f (VIW(,,z) + V2(2, :) du(x). Using that by Jensen's incquality.
Vw,, WB = //// ( autors du(x)
(189)
2 /ʃ NVW du(=)dp(2)),
and similar reasoning of (185), we find that Vw,W (²))), and we can conclude that V Vw, W is a Wasserstein gradient by Definition A.3.
C.11. Proof of Proposition A.9
Let v E (R and P4n = 1 En 6z,- Since F is Wasserstein differentiable, the Wasserstein gradient WW.F( satisfies for any coupling 7 [I([,n;)) v) (Lanzetti al., 2025, Proposition 2.12).
F(v) = F(un) + (Wws) x) dx(2,)) + (V/ ||0 u1l dz(s,
(190)
37

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Let .i.... , bn R2, Un Ei=1 -z and 7n = (tr,eth))) [[ Pn). Then, since F F(un) and F(x1 , Fin lnn) F(un), get
F(zi + hi In + Mm) - =((n))
(1/ || u112 (2,
=F(pn) + / (VWsF(Hn))))) - dml(x,))+ +o
Σ F(nn)( +o (1)
n = =((un) + n
(1
Σ 1isll
(191)
F(r
(
11.) +o
+o
Σ
11.113
n
Thus, by definition of the gradient of F, we deduce that F21........... .0) n F(un)
C.12. Proof of Proposition B.4
Let I' E expp'(Q) Let s,t E [0, 1], and 0²(7) = (expaM () for E P2(TM). Then, (0", ,') E II(P,, Pt).
Therefore,
(P, S (W3 dī
(192)
Moreover, since for F-a.e. 7. (exp,M o(sT'), exp,M ((Tx))) E ()) o' (7)), we have the following inequality:
/
Ww,
(7)
< /ʃ d( d(exp, (1o)) (
(193)
| s1² // |1 dr(2,,
1 s) 2 w. (P, Q)
where we used that I E expp (Q) and that d( (t)), (su) It s|||||
For the other inequality, we have for any 0 S < 1 1, using the triangle inequality and the previous inequality,
(w,(P Ww, (P, Ww, P.) + Ww, Q) sWW,, (P, Q) Ww, (Ps, P.) (1 (P,
(194)
This is equivalent with
(L W) (P, Q) S Ww, Pc).
(195)
Thus, conclude that Ww, (P,,P.,PPP = |1 s/W,, (I, Q) and thus / ++ IP; is constant-speed geodesic between P and Q.
C.13. Proof of Proposition B.5
We first state a lemma showing relation between 7 E exp and specifically constructed 7. (x) with + (t) geodesic between and v.
Lemma C.13. Let H.V E Pa(M), y exp,"(() and the geodesic benween μ and defined for all t E [0. as H. (t) = (expaM o(t>")) #7 Let 7t = (expam o(tT"), (1 t)PT MM Mexpp oftry) o x²) 47 Then, it E exp exp (v). and, for every s 1), pry (s) (exp,M o(sT")) # i (t + (1 t)s).
38

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Proof. First, we verify the equality (x)(s) pny1 s(1 - ()). Fix 8 [0, 1J and let h M - IR be a bounded measurable map. Then,
f h(v) did,, (s)) = / h( (exx=,, (svr) dr vt)
(196)
/ h( (exPexp,(tt) (tr) (s(1 sts,, dr( v).
Fixing (z,v) E T.,, let c(t) = (tp)( L E [0, 1] be the unique geodesic starting from with a(0) Then, we have PT, (Tzz)o(t)(()) = e(t) by the properties of the parallel transport'. Furthermore, by definition of the exponential map, for every u E [0, l].
exPexp, (tv) (uPT2--oxp,tee x-xp,,,ttt (v)) = cxPa(() (ud(t) = c2(u)
(197)
where 02 is the unique geodesic such that (0) = c() and (0) = e((). By uniqueness of the geodesics, we thus have cu(u)= = ctt + u) = exp,, ((t + u)u) for every 0 S u S 1- - From this, we obtain
ʃ M() a(w,, (s) = 1 (ev) (s(1 - L)PT x- zexp ((n(())) d)(
/ h ((exp,(t + s(1 ())))))d((()) (z)
(198)
/ hly) d(p + s(1 t))) (v),
and thus have proved (s) (1 + (1 s).). In particular, have * 7t n Myy (1) (1) = exp# 7 so v/ has the correct marginals. Moreover, it is optimal
1h, (n(1), and exp# 7t
/ Noulle dro(ae, = [1(1 ((l1-1)PT, (t)) dr(xr, v)
(199)
= (1 - t)² / |1u112 dr(2, = (1 WH, = (t),
where used in the last line that y expp (v) and poy is a geodesic such that ya, μ and μl, (1) v, and in particular, (1) = (3(()()), = (1 - w? (m, (0)
Now, we state a second lemma providing a Taylor remainder theorem on P2(M).
Lemma C.14. Let F P2(M) - R twice Wasserstein differentiable functional, μ. E P2(M) and y expp.'(v) and note pay [0, 1] + MM the geodesic between and v defined as (t,(()) = (exPaM (( 47, % and n E exp (v) given by Lemma 131 Then, there exists EJO, 1[ such that
F(v)= F(u) / (V),FFli)(s), v) d)(2,u) + 2(1 - 1)2 / (1e,, vt)ar dre(se,
-
(200)
Proof. First, let us note that
F(u(11)) - F(u((0) = L'' d
((C di
di
dt
d
F
(m, (1)) 11=0 + L/'' ( (()
T
"t
d
(()))))))))) dt
(201)
d
11=0 + L /'' d2
(1,(t)
F (() dsdt
dt
ds²
1- + f (1 s)
d²
dt F()
F (1y ds.
ds²
'Recall that a vector field X along a smooth curve c is said to be parallel if D,X = 0. where D, is the covariant derivative along and that for every s, f, the parallel transport operator PT PT=())-ec(a) sends every E T Tegy to X (s)) where X is the unique parallel vector field along such that (tt) = v. Then, since the condition for e be geodesic is that D.ê 0, if c geodesic, have PT "Teej--eaa = é(s) for every 8, t.
39

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
For the first term, we get by the chain rule (see (30) F(,, (WW,, o) dr(x,,).
For the second term, using the mean value theorem (since + d.3 (()) is continuous, and 1- O all s [0, 1], there exists t e|0, 1[ such that
d²
J. (u-s) ao2 ds
d²
F (()) / (1 (1-ssss 1 d dr (t).
(202)
dt²
Since, by Lemma C.13, we have Mre (8) = My My(1 (1 + 8(1 - 1)) for every E [0, 1]. we have by Definition A.8
/ (HLFs, (Te, uz), d7e,,, = d²
F( (8))
(203)
ds?
d²
ds² F(un(( (t + +(1 - (1-1)s))]
(204)
d2
(1 1)²
F F( [s)( s=1"
(205)
ds
This finishes the proof.
Let P2Q E (2(P)(M)). Let F Pa(M) + R a Wasserstein differentiable functional, F((P) = F())) dP(p) and
[ E expp (QQ). Let y in the support of U, then we know that (I (I-almost surely), E exp where = MM, and exp# y. In particular, by Lemma C.14, there exists some eJ0, 1[ and y expp)(1) (v) such that
=
=(n)) + / (VW,FFlllE), w) dr(z, v) + 2(1 1)3 / (ILF (Ft, vz), ve)a, dre(de, vz),
(206)
so that, by the assumption on the Hessian of F,
- 2(1 - t)2 / |KHF (Fx, oz), ve)a dx(( 1
F(v) (())) - ((w,) 1)2 d)
(207)
$ 2(1 - t)² f |Ivelle, (rs(
(208)
5 2(1 1)2 LW (1) 1 L
1
-
(209)
/ ||lae dz(
$ 2
(210)
From this, we deduce that
(211)
- =I/ [(exp+ 7) -F 7) / (V.. Mz) (z), u) (()) L // [101]d7(2 v)du [1013d)(z.0)d(())
d)(2,, u) a(())
(212)
WW,, WP, (P, 2)2
(213)
2
Thus, can conclude that
/f (V, (z). u) dr(z, v)dT + (WW,, Q)).
(214)
Moreover, we assumed M compact, F(I) is bounded for any i and thus fIVW, dP(p) < +00.
Therefore, by Definition 3.3, Vww, F(P) = Vw (P).
40

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
C.14. Proof of Proposition B.6
Let Q Pa(P)(M)). Let W P2(M) P2(M) + R a Wasserstein differentiable functional, W(P)
[f v) (P( and I E expp (Q). Let y and 7 be in the support of I, then y E exp,, (v) and y E exp, (1/) exp# 7. v/ exp# 7/ For notation simplicity, we write Vi and V2 instead of
where Jt- T# 7, n¹ = // and Vw2,1 and Vw By the remainder Taylor theorem (Lemma C.14) applied on the product space P2(M) x P2(M), we
get that there exists t EjO, 1[ such that
u) W(v,)) W(), 14)+ (Vi (( (1)) u) o) / (V2W() (x, 1 dp
2 dt W (1)
(215)
The
term is a Hessian term, which can be written as (t) 1y (HW nf (are, vx), (are rf)], (wx, (e.))) ( (1) vf), where we define HW TM x TM +
last
TM x TM the Hessian operator at (7, 72) similarly as in Definition A.8. By the assumption on the Hessian, we thus have
|W(e - w(mn)) / (V WinGu, w)n dz(x, / (()))) v')), d''('',
1
/ITHM vx), (1t rf), (vi, (onn))) dm a ())7%(((1)))) 1f)
2(1 t)2
(f |lvelle, dr( + / luglle drilas, vf)
$ 2(1 t)?
(216)
L
2(1 t)2 (W)(u(()),)) + (B)(,()),
L
(/ |1u112 dx(a, / 10'1 d''(',o))
2
Then, let us bound
/ow W(P) - ((Q)- -a10)///(00 (/ V (()))M))))) dr(a, v)d(())
// (W -(b (n))) ai (>)
-1/l (4
(dd (6/)
\/ʃ ( (M(eea
(>)))(())
( (V o
(11)) d)((z,))d''(a' e')du (P))y (7))
(6'))
(V.W(OM
d7(2, - (ovaaa)))ww/(MMMMeMooooo
do(2', (o) ar (T()))(()))
2/ (fʃ || dr uddr (7) + // 11' dy (((dde)aaa()))) by (216)
= LW2 W. (P, Q) since U E expp=
(217)
This allows conclude by Definition 3.3 that
V Vww. W(P)(u) = / (V,W(H,)) + (2)(s, dP(v).
(218)

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
C.15. Proof of Proposition B.7
We note IP: := Px and je Uxe for every c. Let (Rd)Cxx, for every R, we define Pt == Fx-th, and for every
c, uu = +tc+the, so that P. = Ō Ee Sug. We also consider the transport plan of 1 Ei=1 (lag,n)) (which satisfies Rd # ?i- - μe and #xp# # Tt = 14), and the plan F = ? Een dng (which satisfies d I = IP and o## P.).
It is not difficult to see that for / small enough, for every Tp is actually optimal between μe and fu (that is, of E exp (nt)), and therefore
W3 14) = / |1ul/2 = 12
(219)
Moreover, it is also the case that for L small enough, D. (xpF²(PP)) Indeed, since for every C, Wz (n()
2010'
(2 1e")) which is zero if and only if c = c', it ensues that for t small enough, for every c,
Wa(ne, v&) min Wa(ne,
(220)
Thus, for any I' II(P, Pt), represented by the matrix (Tae)e..'=,.... have
Σ
Wa(ns, ( TTe
(221)
0,4'==
2
Σ
(z 12]) o,
(222)
c,0'=1
1
Wa(m",
(223)
0
n
|M = /f ()) ((7),
(224)
On
so that, by taking the minimum over r, we find If |10l/d u)dr, (7) S (P, Since the reverse inequality always hold, we find that
W w. (P,PP)) lf (2)
|
(225)
On
and we conclude that I is optimal, with Ww, (P,, Pt) = O(). Plugging I, into the definition of the WoW gradient, we
find that
F(P,) = F(P)+ //(Vwin, / (1wr o())
(226)
that is (since th X for t small enough),
C n
F(x + th) F(x) +
( Www, (z)) + o(t).
(227)
nC c=1
From the definition of the gradient, we deduce that for every i,
Vww, ((P)(1)(()) = Cn cx))F.)
(228)
This finishes the proof.
42

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
C.16. Proof of Proposition B.8
Let IP, ([(1-t)Tz (I m #1 where I' 11(1P Q 0), T P E Ll,(P, Q) a.d # 1 [I,(P, 0). Then, we have
V(P.) / F () + (TT) o 7) ( p,
f F(Ne) dr for exp, ((1 ( + +TT)) = ((1 t)TT + +TT) #
(229)
$ (1 - 1) f F(()) dQ(a)+ +1 I ((o))o(v- - At(1-t) t) / W3(4,,) v) aT s (1 () (V(O) 1((1 - 1 (Q,00,))
where we used in the last two lines that F is X-convex along + pc, and Ww (Q,00))) < /W3(u, dr dr
C.17. Proof of Proposition B.9
13 (P, Then, we have Let P, (((1 - (()Ti) TT=') a m²) where I E [I(,, Q, O),, T Io(,, Q) a.d #
1 W(P.) - //2 /ʃ 1(((1- 1)T +IT)) 2 ((1 TyT + + dr (d
{ 1 / // w + t/2// // W(y, dC(v)d0(v))
(230)
(1 (L)+) +tW().
C.18. Proof of Proposition B.10
Let IP, Q, O E (P (B.)). Define the generalized geodesic t + P. = (((1 t)T= + +rr=") x )
+
where I' e
[I,, Q, 0), " _ E (P, Q) and # 1.3T E (P, D). Let show that F Q - Ww, (Q, P)2 is convex along this curve.
To do this, first note that = (TT ((1 t)Tşi + +I= # I' P.). Then, we have
F(P.) =
2 Wws((P, P)2
(231)
S / (3 ((1- ((1-))T% + +TT))#) dl di v, 17).
n= + T is an OT map by Brenier's theorem since it is the gradient of a convex function (as a nonnegative weighted sum of convex functions). Thus, for I'almost every (H,v, W2 (p, ((1 ttTy + +TT) [[ - 1 fT idll2=(0) Then, applying the parallelogram identity on the Hilbert space L²)),,, we get
Note that T (1 - ryT=
1
F(P.))
2 [[11 1)T + (TT Id)l (p,,,,
1
/ [[11 - (T Id) t(T Id) Id) (( u,
2
-t)
/ ||T% -Id (dli dr + t / ||T Idl dF
2 ((1 - 1) | -T dr v,
(232)
2
1)
1(1 (11-1) / |T - TT21 (r(m,
Ww. (P,Q)2 Ww, (P,)) 0)2
2
2
t(1 t) / || T dl
(1 (Q) (0())
43

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Finally, since (Th, #1 (i,,)) we also have Wa(e, n) I -Tllii> Thus, have
ITT Tilliano dr (p, 2 (2( dr 2 (w, 0)? /ITT
(233)
where we applied that a 2 E [I( 0) for the last inequality. Plugging this result in (232), we get
F(P.) - (1- (1) + (0) t(1 2
W. (Q, 0)
(234)
D. Additional Details and Experiments
D.1. Minimization of the MMD
We want to minimize F(P) !MMD? Q) for P,Q (2((P((R)) and kernel K P2(RR) P2( 4 R.
Recall that ((P) V(P) + W(P) + est with V(P) = ((w) dP(u) and V(u) (u,, v) dQ(v), and W(P)
[f K( dP())dP(v). If Ky(u) = k(,v)) is differentiable functional, then the gradient of F is given for all E (R) () E P2(RR)) by
Vww. ((P)(k) Vww, V(P)() + Vwn, ((P)())
= Vw. + (Vw:
(235)
= Vw, Krl) Vw, dQ(v) + / VW, Ko(n) dP(w).
We can choose different kernels, giving different discrepancies. We compare here different kernels based the Sliced- Wasserstein distance (Rabin al., 2012). Let p 1. We recall that the Sliced- Wasserstein distance is defined between nr, E P22(R))a
SWE(GD) = f WHIP
gd-1 HH, P v) (o).
(236)
where Sd (0 E R², |0|| 1} is the sphere, pd (0, x) and a is the uniform measure the sphere. The Sliced- Wasserstein distance allows defining a Gaussian positive definite kernel K(u, v) e JSWWj...))/ (Kolouri et al., 2016; Carriere et al., 2017) and a Laplace positive definite kernel K(H, v) = e SWi(aw)/ (Meunier al., 2022). We also propose in practice use the Riesz SW kernel K ((,, v) SWa( w)" for E (0,2) and inverse multiquadric kemel (IMQ) K( (H,
(cts Note however that the Riesz SW kernel is not positive definite (but conditionally positive definite), and that showing that the IMQ kernel is positive definite is an open question.
The Wasserstein gradient of the Sliced-Wasserseee Wasserstein distance F(u) 1 (S3(H,,) can be computed (Bonnotte, 2013, Proposition 5.1.7)
{ S V6((z, 0;)0 do(O),
V Vw,
(237)
with the the Kantorovich potential between p6 and pa v. and thus (u(2)) - F (Fpga(()) for all R. For the Gaussian kernel, by the chain rule, we have (w.K 1e 4SW3 (VW,
In practice, the integral w.e. is approximated using Monte-Carlo approximation, i.e. we draw oI...., 0r ~ o L independent directions, and approximate the Sliced Wasserstein distance and its gradient as
SW2(n,) I,
Σ t=1
w
ple #
pd v), p1,
Fw,
1 Σ
I, (xx,
(238)
The Wasserstein gradient can also computed using backpropagation as shown in Proposition A.9.
44

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Iter
Iter 10
Iter 25
Iter 100
h=
h=0.05
0=5.5
Figure Gradient flow of MMD with SW Gaussian kernel K(H, = e -SW3(..))((2))
Iter
Iter
Iter 25
100
O O
CO
h=
h=0.5
h=1
Figure 6: Gradient flow of MMD with SW Laplace kernel K Ku,,v) = e SWi(Kw)/
Iter
Iter 10
Iter 25
Iter 100
0.01
0.0.55
0-5.5
Figure 7: Gradient flow of MMD with SW IMQ kernel K(μ, v) = (c + xw)
45

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Iter
Iter
Iter
Iter 100
O
O0
r=
COO
1=5.5
Figure 8: Gradient flow of MMD with SW Riesz kernel K(,,v) = SW2(u,)".
Iter 0
Iter 2
Iter 25
Iter 100
O
Figure 9: Gradient flow of MMD with Riesz kernel k<x,, y) = -1 yl12.
Iter
Iter 10
Iter 25
Iter 100
000 OOO
COO QOO
L=1
O00
L=5
O
QOO
L=
OGR
OOO
L=
5=000
Figure 10: Ablation of the number of projections L for the approximation of the Sliced- Weeedd--ssersteinn distance (with the Sw
Riesz kernel and r = 1).
46

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
3
3
3
8
9
Figure 11: Images along the trajectory of the flow from MNIST to Fashion MNIST. We see that images belonging to the same class in the source dataset are flowed towards images from the same class in the target dataset.
D.2. Ablation of Hyperparameters on Rings
Additionally to Figure 1, we compare in the following figures trajectories of the minimization of the MMD with various kernels and with different hyperparameters. To recall the setting here, the target a mixture of rings (Glaser et al., 2021), and each ring is seen through an empirical distribution fo.
E 5yp- Thus, the target is a mixture of three Dirac:
3 dgs. Each distribution jf contains n = 80 samples. We learn a distribution P f yz + +30 6 + f5po, modeling each μ as n Di 82g- In practice, the distributions and are seen as tensors of size (3, 80, 2).
Q= 63 + 3
55a +
To compute the gradients of the MMD, we use L 500 projections, and 0.1 as learning rate. We plot on Figure results with the Gaussian SW kernel K K(,, v) ( on Figure results with the Laplace SW kernel (u, -sWW( on Figure 7 results with the IMQ kernel K(H,v) v/ = (c + SW3 v)) and on Figure results with the Riesz SW kernel K (,, v) = (p, We also add on Figure 9 a comparison with the flow of the MMD with Riesz kernel k<x,, y)
|| yl1 in (Hertrich et al., 2024b), where the structure of the rings is not taken into account.
On Figure 10, report ablation of the the trajectories with different number of projections for the SW Riesz kernel. More precisely, show the results for L E 51, 10, 100, 500}. This demonstrates that for low dimensional problems such as 2d rings, L = 100 projections already provides good results. However, the scheme is more sensitive to the number of projections in higher dimension as we show on Figure 14.
D.3. Domain Adaptation
We first add on Figure 11 more samples of the flows of the MMD with K(u,v) SW2(1,, v) between MNIST and MashionMNIST. In this experiment, recall that the flow starts from Po D Opan where μ yam is the uniform empirical distribution of samples belonging to the class e E (1,... 10} of MNIST, and targets Q - - De dua,n. We used n 200 samples for each class of the datasets. The Sliced- Wasserstein distance is approximated with L = 500 projections. To speed up the flow, similarly as (Hertrich et al., 2024b), we add a momentum m E [0, 1). ie., at each iteration & 2 the update for each particle (1....,m) in class E {1. fI....,Cy
C) is of the form
fi,++ Vww, F(P&)(M (=f + mWi.k
(239)
x 2f II - TU,,++,
with 0.0 = 0. We choose a step size of T 0.05 and m 0.9.
Complementary to Figure 2, we see on Figure 11 that images from a same class are flowed towards images from a same class in the target dataset. To verify this intuition, we applied domain adaptation experiment which we describe now.
47

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
3
3
3
3
3
3
3
3
3
3
5
5
5
5
5
5
5
5
5
5
5
/
/
/
/
/
/
/
6 6
6
6
6
6
6
6
6
6
6
6
D 0
O
O
O
O
O
O
O
O
O
O
O
O
4
4
4
Y
Y
4
4
4
4
4
4
4
4
4
4
4
9
9
9
9 9 9 9
9
9
9
9
Figure 12: Samples of trajectories starting from Gaussian noise towards MNIST (Left) and CIFARIO (Right).
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
O
O
O
O
O
O
O
O
O
O
O
O
/
/
/
/
/
/
/
/
9
4
9
9
9
9
9
9
9
9
9
6
6
6
6
6
6
6
9
9 9
9
0
O
O
O
O
O
O
O
O
O
O
O
5
S
S
S
S
S
S
S
S
S
4
4
4
Y
Y
4
4
4
Y
4
4
4
4
4
/
/
/ / /
/
/
/
/
9
9
9
9
9
9
9
9
9
9
9
9
}
2
Figure 13: Samples of trajectories starting from Gaussian noise towards MNIST with momentum = 0.9 (Leff) and no momentum (Right). We run the flow for 100K steps, and plot samples every 6667 steps.
We first train a classifier the training set of the MNIST dataset (using 500 samples by class). The classifier is the CNN used in the examples of the equinox library" (Kidger & Garcia, 2021). It is trained for 5000 steps with the A optimizer (Loshchilov & Hutter, 2019) and a batch size of 64. After the training, it has an accuracy of 96% on the test set, and of 100% on the training set.
Then, we flow the dataset FMNIST towards MNIST by minimizing the MMD with kernel (H, v) = SWe(p, v). We run the scheme for 500K steps with a step size of r 0.1 and a momentum of m 0.9. To match the labels of the flowed dataset with the labels of MNIST, we solve an OT problem between P the flowed dataset and the target dataset with the squared 2- Wasserstein distance groundcost, i.e. with P DE Spe... and Q c E D Senin, solve the problem
((W3(w v(r))) 1sk. TEl((P.Q)
min
(240)
using the Python Optimal Transport library (Flamary et al., 2021).
We plot on Figure 3 the accuracy the pretrained classifier along the flow starting from FMNIST. We observe that the accuracy converges to 100% for a sufficient number of iterations. Thus, shows that the classes of the sources datasets are perfectly flowed towards classes of the target dataset, on which the pretrained neural network is trained, and thus has perfect accuracy.
On Figure 3, we also replicate this experiment from SVHN to CIFARIO which are composed of 32 x 32 x 3 dimensional images. The neural network used is the same convolutional network used in Appendix D.5, and is pretrained on CIFARIO 10 during 5000 steps with the Adam optimizer and a batch size of 64. We use here n = 100 samples by class, and run the scheme for 500K steps with step size of 0.1 and m 0.9. We also observe that the accuracy converges to 100%, indicating that it also works in moderately high dimensions.
D.4. Generative Modeling
In this experiment, we generate samples from different datasets starting from Gaussian noise.
We show on Figure 12 trajectories starting from Gaussian noise towards MNIST and CIFARIO. For both datasets, use
hhttps://docs kwwgeiidgiit......ee.eeeeee wwt.//quiiiieee......... nwxp//amp/eeeeeeeeeeeeeeeeeeeepe
48

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
7
(a) - 10
(b) L = 50
0
0
O /
0 /
0
0
0
0 1
|
I
1
/
/
2
2
2
2
2
2
2
2
2
2
2
2
2
2 2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2 3 4 5 2 7
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
4
4
4
4
4
4
4
4
4
4
4
2
4
2
4
4
4
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5 6
6
6
6
6
6
6
6
6
6
6
6
6
6
7
7
7
7
(c) L 100
(d) 200
O
1
1
Z
Z
Z
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
2
3
3
3
3
3
3
3
3
3
3
3
3
3 4
7
3
3
3
3
3
3
3
3
3
3
4
4
4
4
4
4
4
4
4
4
4
4
4
9
4
4
4
4
4
4
4
4
4
4
4
4 5 6
5
5
5
5
5
5
5
5
5
5
5
5 6
5 6
5
5
5
5
5
5
5
5
5
5
5
5
6
6
6
6
7
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
7
2
2
2
2
2
2
7
7
7
7
7
7
(e) T. 300
(
500
0
1
1
7
1
1
1
1
I
2
2
2
2
2
2
2
2
2
3
2
2
2
2
2
2
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
4
4
4
4
4
4
4
4
4
4
4
4
4
4
4
5
0
4
4
4
4
4
4
4
4
4
4
4
4
4
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
5
6
6
6
6
6
6
6
6
6
6
7
7
7
7
7
7
7
(g)
1000
(h) 2000
Figure 14: Ablation over the number of projections for generative modeling MNIST.
49

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Toure Fran Cest Sancel dee Seedeer Ioe
Dre Sande de And
3
5
6
7
8
9
O
2
3
4
5
7
9
0
2
6
7
8
9 9
2
3
4
5
6
7
8
9
0
2
3
5
6
8
O
2
4
5
6
7
8
3
5
6
7
8
9
O
2
3
4
5
6
7
9
9 9
O
2
5
6
7
8
9
O
/
2
3
4
5
7
8
5
6
8
9
O
\
2 2 2 2
4
5
6
7
9 7 9 9
0
2
3
4
5
6
7
9
0
1
3
4
5
P
2
4
5
6
8
9
I
3
4 4
5
6
7
8
5
6
8
9
3
5
6
8
3
4
5
8
9
0
/
à
3
4
S
8
9
(a)
Id
(b)
A 1/
(c) 11 = Id
(d) (d) A + +
Figure 15: Synthetic data for the dataset distillation task on MNIST (Left) and FMNIST (Right) with or without embedding.
momentum of m = 0.9 and a step size of = 1. For MNIST, we run the flow for 18K steps, and plot samples every 1200 step, while for CIFAR 10, we run it for 150K steps and plot samples every 10K step. We used n = 200 samples for each class for MNIST and n 50 for CIFARIO. We also compare trajectories on Figure 13 with using a momentum m 0.9 or no momentum for MNIST, running the flow for 100K steps and showing samples every 6667 step.
In Figure 14, we present an ablation study over the number of projections used to approximate the distance on the MNIST dataset (with the same setting with momentum, i.e. m = 0.9, r = 1 for 18K steps). We observe that to generate sufficiently clear images, we need at least 300 projections. This may be because a higher number of projections provides a better approximation of the gradients.
D.5. Dataset Distillation
In this task, we aim at generating a new dataset allowing to approximate a target distribution Q = à Em=1 dpenn with a distribution IP ?/4 Ee 6ys.r for p < n. in order to be able to train more efficiently neural networks on it. In Table 1, we take Q as the MNIST and Fashion MNIST dataset, with n 5000 samples by class, and C 10 classes, and report the results p {1, 10, 50}. We report the accuracy of ConvNet trained on the synthetic dataset and evaluated on test set, averaged over 5 trainings of the neural network, and 3 synthetic datasets. We use a similar architecture as (Zhao & Bilen, 2023), i.e. the ConvNet includes three repeated convolutional blocks, and each block involves a 128-kernel convolution layer, instance normalization layer, ReLU activation function and average pooling. This forms the backbone part of the network, and the full classifier followed by a linear layer. For the initial distribution PPo = à >e Spc.., each d is chosen as a random subset of the samples of y,"", The results reported in the column "Random" " correspond to the ConvNet trained on the initial data.
Zhao & Bilen (2023) proposed to solve the problem by minimizing
F((n)c) Σ Eew [I/ (
C
d(pe -
Esww
Σ MMD (v) v2 Agva)
(241)
with linear kernel k(x, v) (c, v). where A Rd + Rd is some data augmentation and 1h Rd - Rd with d' < d is a randomly initialized neural network used to embed the data. This loss does not take into account the interaction between the classes and just learn any set of synthetic samples for each class. In this work, we propose to take into account the interaction between the classes, and thus minimize
1 Fow & (P, 04 0- Q)).
((P)
= 2
(242)
with K(n,v) SW),
In practice, for v we use the backbone part of the ConvNet, and for A", we follow the same strategy as (Zhao & Bilen, 2021) (i.e. we sample one augmentation among color jittering, cropping, cutout, scaling and a rotation for MNIST, and also add flipping for Fashion MNIST). We optimize (241) by stochastic gradient descent over the particles, sampling one random network and one random augmentation at each step. We trained it for 20K iterations, a learning rate of = and a momentum of m = 0.9. In practice, we observed numerical instabilities when optimizing in the ambient space with augmentations.
50

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
9
2
8
9
O
8
9
7
8
9
9
O 2
15678
Figure 16: Examples of images output by flows for the transfer learning task with k = 10 for Fashion MNIST (Left), KMNIST (Middle) and USPS (Right).
To optimize (242), we also performed stochastic gradient descent, sampling one random neural network and one random augmentation at each step. We used also 20K iterations, learning rate of r = 1 and a momentum m 0.9. Then, we assign the classes using an OT matching as explained in Appendix D.3. We add on Figure 15 samples learned with this loss, with and without embedding. We observe that images are slightly clearer when using an embedding. To compute the gradient of F in practice, we use autodifferentiation.
D.6. Transfer Learning
We describe the details for the experiment of transfer learning. We recall that the target dataset is of the form Q
C Ee=15pcx Sye. with wotk a uniform empirical distribution of & samples of the class c. In Table 2, the targets datasets are Fashion-MNIST, KMNIST and USPS. Thus, C 10, and we choose k E {1, 5, 10, 100). For the source dataset P T 1 we used the MNIST dataset with n = 200 samples in each class.
I DC
We augment the target dataset by flowing the samples of MNIST on the target. For MMDSW, we minimize ((P) = 2 3MMD Q) with kernel K(μ, v) -W v). by running the forward scheme for 5K steps for k {1,5, 10} and 20K steps for k 100, with step size T 1 and momentum m 0.9. Finally, we align the labels using an OT matching between the flowed samples PP and the target Q, for the dataset distillation experiment.
We compare it with training directly on the small dataset, and with two other methods. The first one, called OTDD (Alvarez- Melis & Fusi, 2020). represents the dataset as probability distribution on Rd x P2(Rd), where the labels are embedded in (Rd) by considering the conditional distribution, i.e., a feature-label pair (Ir, c) is represented as (z, n²). Then, they compare datasets using Optimal Transport with cost d((2; e), (a', .c))² || x' 1/3 + ""(le,""")) The flow of OTDD then minimizes the OT distance with this cost, i.e., the objective is F(u) OOTDDDD (M, v) for u, P2(RR x (2(R))) with
/ (12=-1113 Waln (1 dr((x, (I,)))
OTDD(H, inf 7EIl(H..)
(243)
For big datasets, the conditional distributions H., can be approximated by Gaussian distributions. Alvarez-Melis & Fusi (2021) proposed several schemes to optimize this loss using Wasserstein gradient flows. We did not manage to replicate their results with their code. Thus, we reimplemented it with some differences. First, similarly as (Hua et al., 2023), we used an embedding in dimension 2 of the data to approximate the conditional distributions with Gaussian distributions. Thus, we model the datasets as distributions over Rd x R² S (R), with S + (R) the space of symmetric positive definite matrices. This helps avoiding memory issues and scaling to higher dimensional datasets as it reduces lot the dimension of the samples to flow. For this embedding, we used Principal Component Analysis (but note that we could use other embedding methods such as TSNE (Hua et al., 2023) or Multidimensional Scaling (Liu et al., 2025)). In practice, we approximate OTDD using an entropic regularization, which we compute using the Sinkhorn algorithm (Cuturi, 2013) and ott- jax (Cuturi et al., 2022). We optimize it using A with a learning rate of r = Ie and run it for 5K iterations for & E (1,5, 10, 100). To get the labels, we use an OT matching as in (Hua et al., 2023), which we solve using POT (Flamary et al., 2021). More precisely, for each class e E (1....,C) C} of the target distribution, we can compute mean me and a covariance E... and a
51

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
weight wi = 2s with n the number of samples in the target dataset, and Be the number of samples belonging class e. After flowing samples, we have tuples (2i ma; En )== and want associate each sample a class. To do this, they propose to solve the discrete OT problem between Q = De (.dN( and P = - (m),,
a
min PELI(P Q)
Pie ( Ee),
(244)
and then use as distribution μ²= E ((..)))) with y/4 = Ec-1 cl (P).=max P:}-
The second baseline we use is the one proposed in (Hua et al., 2023). In this work, they first observe that the Gaussian approximation for high dimensional datasets might not scale well in memory. Thus, they propose to use an embedding in a lower dimension space of the conditional distributions, before doing the Gaussian approximation. The datasets are then represented as probability distributions on Rd x Rp x 5p*(R) (R) with p < d. Instead of using an OT cost to compare the datasets, they used the MMD with kernel obtained as product of Gaussian kernel. Then, they applied Wasserstein gradient flow of the MMD (Arbel et al., 2019) to minimize it, with Mures-Wasserstein gradient descent step (Altschuler et al., 2021) for the symmetric positive definite covariance matrix. We note that in contrast with our proposed MMD, it requires many hyperparameters to tune (3 bandwidth of Gaussian kernels and noise to add to make the flow converge). We reimplemented it in jax (Bradbury et al., 2018), used p = and a Principal Component Analysis (using scikit-learn (Pedregosa al., 2011) for the lifting of the conditional distribution (instead of TSNE in (Hua et al., 2023)). The Gaussian of each class is then obtained by computing the mean and variance of each class. We used as bandwidth h = 100 for the feature part, h- = 50 for the mean part and h = 1000 for the covariance part. We ran the flow for 20K steps with a step size of r = 10, and momentum m = 0.9. To get the final labels, we solved (244) as explained in the last paragraph.
In Table 2. we report the accuracy obtained by training a LeNet-5 neural network for 50 epochs with a AdamW optimizer and a learning rate of 10-4. Moreover, we average the results for 5 trainings of the neural network, and 3 outputs of the flows. We add on Figure 16 examples of images returned the end of the flow of the MMD with K (H,,) (,,(
D.7. Handling Different Number of Distributions between the Source and Target
Let P = N Exnn fure and Q = 1 Exe dyرر with M < N. In this situation, the flow might not converge well towards the target distribution since they have a different number of Dirac. This is illustrated on Figure 17, where the target is composed of M rings yk.,, and the source is initialized with N distributions, and we minimize the MMD with a Gaussian SW kernel K(H, v) = e =SW3(uw))/h We see that the flow does not converge to 3 rings, as cannot split the mass because the Wasserstein gradient descent allow only changing the position of particles.
37
This problem could be solved by different solutions. For instance, one could use instead of Wasserstein gradient flow (Gallouët & Monsaingeon, 2017). This flow can be approximated e.g. by using Birth death Langevin algorithms (Lu et al., 2019; 2023) where the Langevin step approximates the Wasserstein gradient flow part, and the Birth death part approximates the Fisher-Rau gradient flow part. The birth death consists at killing and duplicating ran- domly particles at each step. Another solution to approximate the Fisher-Rao flow is to change the weights (Yan et al., 2024).
gradient flow
We propose to perform the Wasserstein gradient flow, but allowing to change the weights of the particles, which is not possible for the Wasserstein gradient descent. Ideally, one would want to solve directly the JKO scheme
(R x/ J' ||x v13 d2(2, v) + (
7 Mk+1 = T M6+1-
(245)
However, if we do not fix the support, it is not possible to directly solve this problem, except we use neural networks. Note that (245) can be seen as a semi-relaxed unbalanced optimal transport problem, where the first marginal is fixed. This has been leveraged to solve the JKO scheme e.g. in (Choi al., 2024).
We propose instead to alternate between Wasserstein gradient descent step, which allows moving the particles without changing the weights, and backward step for which we optimize over the coupling while fixing its support, which allows then to change the weights.
For simplicity, let us describe the procedure more precisely on P(Rd). ). Let E P2(RR)) be a target distribution, and suppose at step k. pk = Σx= of S2, with of 2 0. Σn_ of = Let D be a divergence we want to minimize w.n.f. v, f.e. we want
52

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Iter
Iter 2
Iter 25
Iee
Target
O0
Iter 100
Iter 200
Iter 250
Iter 500
Figure 17: MMD with Gaussian SW kernel and 4 distributions flowed towards 3 rings. The flow does not converge to the rings.
minimize F(p) ((p,)), v). Then, our update is
M [Td VW,, # ul
/ xxsupp(es m 7= 2 ||:: v122 d>(2, v) +
7k-1
Rd)
(M++1 o 7++1-
(246)
The first step is a regular forward step, which moves the position of the particles. The second step learns a coupling E P22(R" x Rd)) which satisfies 4 P and such that < y is supported on the same set of particles. This step can be seen as solving semi-relaxed Unbalanced Optimal Transport problem if the support for both distributions is the same. Suppose that 7 = EE P.)0(r,,zz) and note C' e R=x= the matrix distance. Then, the second step can be rewritten as
Σ DPPT 1n]00,00,
min PEREX" (..
(247)
For D(,,v) KL(p]))), this can be solved using the Sinkhorn algorithm for the semi-relaxed UOT problem, i.e. with (P '(13 (Séjouréé al., 2023). For D MMD
different algorithms solve it such Projected Mirror Descent an Accelerated Gradient Descent (Manupriya et al., 2024). Here, propose to the half step of the Mirror Sinkhorn algorithm (Ballu & Berthet, 2023), which performs first Mirror Descent step with Bregman potential σ( (P) = (P, log P) (for which Pa+1 = VΦ" (V (Pp()))) -TV()) = Ps oe (PPP))) and then perform a (Sinkhorn-like) projection on the constraint, i.e,, noting
f(P) = (C,P) +MMD2 Σ 1n16dr,,
(248)
the objective, the algorithm becomes
PL_ = Ps (
(249)
P diag(a a @ (Pk (P)+11n))
We show on Figure 18 the results on the rings experiment. We observe that the weight the 4th ring set to 0, and thus that the scheme converges to the target.
53

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Iter
Iter 2
Iter 25
Iter
Target
O0
Iter 100
Iter 200
Iter 250
Iter 500
Figure 18: MMD with Gaussian SW kernel and 4 distributions flowed towards rings using the proposed algorithm. The flow converges to the rings by setting the weights of one of the ring 0.
E. Related Works
E.1. Optimal Transport Distance for Datasets
Alvarez-Melis & Fusi (2020) first proposed to compare datasets with dedicated discrepancy, which takes into account features and labels. They proposed to do it by representing datasets as uniform empirical distributions over R x PR)R)) embedding the labels in (R by considering the conditional distributions, i.e., feature-label pair (z, c) is represented as (x., n²) with oe the distribution of samples belonging to the class c. They proposed to compare datasets using an optimal
transport distance with cost d(1r, c), (x', ec))² || 2'112 + Wa(nF, v²"). To summarize, they consider distance between μl, E (2 )2(R")),
/ (1=- W34e ,o))) ((),
OTD inf yEll(u.))
(250)
In practice, Alvarez-Melit & Fusi (2020) approximated the conditional distributions μe by Gaussians to be able to compute the Wasserstein distance in form, which leads to a complexity of (C'nd + & + n²C log(nC)) as requires estimate C means and covariance matrices from samples, to compute C2 Bures-War wessertnsnn distances, and an OT problem between On samples. The final OT problem be approximated using an entropic regularization, which reduces the complexity to O (Crdd Cd + 222 log(nC") (Dvurechensky et al., 2018).
Liu et al. (2025) instead embedded the labels in Rd using Multidimensional Scaling, and further approximated the resulting squared Wasserstein distance with a Wasserstein embedding. Bonet et al. (2025) proposed to embed the labels in hyperbolic space, and used a Sliced-Ve Wasserstein distance to compare distributions on the product space Rd x H.. Nguyen & Ho (2024) used a similar embedding, and a hierarchical hybrid Sliced- Wasserstein distance. More recently, Nguyen al. (2025) introduced sliced optimal transport dataset distance using dedicated projection from Rd x P2((R) to R.
Concerning the task of flowing datasets, Alvarez-Melis & Fusi (2021); Hua et al. (2023) both modeled conditional distributions as Gaussian, and solved flows on Rd R' x - (RR. More precisely, Alvarez-Melis & Fusi (2021) minimized OTDD on Rd x Rd S (R), while Hua et al. (2023) minimized an MMD over Rd IR² S *RR) with a product of Gaussian kernels, and using an embedding on R² for the conditional distributions. In contrast to these works, we encode the labels directly into the discrepancy by using MMD on the space of probability distributions with suitable kernel.
Alvarez-Melis & Fusi (2021) proposed several ways of minimizing F(4) = OTDD((,, 2) for us note 14 i Ei ( (t the dataset at step k and assume <; E {1...., C). For small datasets for which the Wasserstein distance between conditional distributions can be computed efficiently, they just proposed to flow the samples, i.e. computing Fi, = Ig T V... OTDD( 2) and updating the conditional distributions at each step. When using the Gaussian approximation, they proposed to update the mean and covariance at each step (feature driven), or to do gradient satrive---ixed-llaad They also considered the joint-driven variable- -iit-diivve-variable-
Σi (ar, Let
descent step for the C means and covariances
54

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
label, where they decoupled at time the Gaussian, and flowed one mean my and covariance E, by Gaussian, ie., for all i {1,... n and k 0,
Ti,++ = Ji,k - T a.,, (PR, v) Mi,++1 mi,k 7 ~Vn. OTDD (P):
(251)
DEA-1 = VV (TDD((E;))
This however requires to cluster the pairs (m4 Σj) recover labels.
Hua et al. (2023) observed that the embedding of the conditional distribution as Gaussian can be very costly in practice for high-dimensional datasets. Thus, they first proposed to embed the features in R² using TSNE, in order to embed the labels as
Gaussian in R², and therefore represented the datasets as empirical distributions over Rd x RR² x S (R). Then, they proposed minimize the MMD on this space with kernel (zz, m, 2), (a', ,,',,, = - 1//h/ For u. V. E (R x R² x S (B)), (B)), let F(u) = 1/2 MMMD* (H, v) = f Vdu + JJ k(x, y) dp(x)dp(v). with V(z) = - (kka, Its Wasserstein gradient is then for all (r, m, E),
m m-mm ll//mm - --- 1'
Vw, m, x = ((x, / Vik((a, m, M), (a', m', E')) dp((2', m', E')) E Rd X x S2(R). (252)
Using the
geometry for the covariance part, their updates are given by
F1,441 -i,k - - F(u) mike Ex,,)]
mi,+11 = -ii - 7[V,,,((Hx)) (Zi, mi,k Ei,))]2
(253)
21,+1 -xpE,,, ( T| ( T[W,, ((us)) mak )3)))
with exp=(S)= = (La + S)E((a + S) for E + (R), (R) the exponential map on the Bures-Wasserstein space. e.g. (Altschuler et al., 2021, Appendix A.1).
E.2. Variational Inference with Mixture of Gaussians
Lambert al. (2022) considered do Variational Inference with family of Gaussian mixtures. Let's note BW(R") C P2(Rd)) the Observing that there is an identification between BW (R4) and Id S++(() (Chen et al., 2018; Delon & Desolneux, 2020), this amounts at solving the problem, for E (R.
space, ie., the space of Gaussian distributions endowed with the Wasserstein distance.
( (fndwt())l).
min MEP X & (R))
KT.
(254)
where pa (N m, Σ) for 0 = (m, Σ) Rdd x St (R). Equivalently, can be framed as an optimization problem over (R)))) by solving
(f I ((ud(WW)l)
min (BW (R-))
KL
(255)
Note that the KL here is the usual Kullback-Leibler divergence, defined between H, E (R))) as
<T((Ilu>> f log
(a( (x)
du(i))
(256)
where we note P,. and pv the densities of μ and the Lebesgue measure.
They address the problem by solving an ODE on the means and covariances, which characterizes the trajectory of the gradient flow in (P,(BR())) WBW.). Alternatively, they propose to solve the JKO scheme between particles by solving for all & 2 0,
(0 1)
argmin ((1) g(n..
( (00m). IO0
+ 1 S22 2
n Σ
( (()))
Σ 8 Wimm (257)
n
72
We now derive the gradient of this functional using our framework, and make the connections with the formula derived in (Lambert et al., 2022, Appendix F).
55

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
Computation of the gradient. Let E Ps (P(((R)) and g (p We want to do the Taylor expansion of F(cxpp(tE)) (tE)) = [(1 (Id + =E()))
((cxpp(tE)) KT, ((d(((sp((()))))))))
(Idd (p)))) =KL
(258)
= -ʃʃ log
=
[ [ppateston (P) Pra+tg(w)aa (2) dP(u)dr.
(x
By a Taylor expansion, we can write for all x E IRd
(xiv(px()))(((())) o(t), o(t) = pp(x) P(a++)())) (x) (x) + (1) (ot)
(259)
where we used (Villani, 2003, Theorem 5.34) for d.P(aate((a)a = -Idiv (p(((())()))) Plugging this in (258), we get
F(expp((E)) = 1ʃ [Pu(t)dP(v) - div (iv()(()))W)) aP(v) d(t)
P(a+ig())aa dp))d
(x(
= =// (log] (f (culooaa(()) - div (()((())) dl(()) +
(p)(P)))())
--logg loggx(())) P(areew))) (x) dP(w)dr
= -sʃ log [[a(ttaa(e)) (iv( aP(v)
+ +0()) pratesionaaa()) (l(u)dd..
(x(x)
J (c(c))P(())
(260)
Performing the Taylor expansion of the second density, we get
F F(expx(tE)) = // log
[Pu(x)dP(v)
(iv(p)(())(())))) aP(v) + o(t)
(x(2)
(Pv(x)dP(v)
(pu(a) Idiv (d(v(p((())(())())) + o(I)) dP(w)dr
F(P) // //10g (Sm(o))P(())
(ii()) dxdP(n)
(x(
(261)
-t / I div (()))) dP(v) S (P(()))(())
/ (=(=))
=((P++ +1 /1 ( Vlog
[Pu())d((V))
,Ew)(t)) du(x))P(()) + o())
(x())
We used in the last line the integration by part formula, and Vww, F(P) VVP, where Vx(x) log (J (x))P((v)) (.
dr 0. We can conclude that
Computation with 1st variation. as conjectured in Appendix B.4.
We now verify that we would recover the same result by computing the first variation,
Let स E (Rd),, v x e- V with V Rd + R potential. Denote Fx : (2.ac(()) - R. Fa = KL(WUI [T) for all J P2anc(R)) (R.), and for all P E P2 (Pac
F(P) KL (fudF())()) = Fr (f
=Far
(262)
We will now derive the 1tt variation F. First, recall that
6F (1) + log / log + log + V. Thus, we have, ou
56

---

Flowing Datasets with Wasserstein over Wasserstein Gradient Flows
noting p (P)(), pp = J Pa(x) (P()) and X = dx(x).
dF
dFx (P + 1=) 1=0
(f (JudP(n) +1 n mdx(u)) I
df
dt
1=
dF (ā +xx1 ʃ 8F= (i)(x)) dx(=)) by definition of the 1st variation of Fx
di
/ (1 + log (oggp(>) - (x)) dx(x))
(263)
[ (1+ log
p/allo (P(l)))((l)) log pa (x))) SecIee Pu(x) dx(e)dr
Rd
[ Ipp(ee) Joc [++ log
log logp((()) (foale)d(()))
du(x) dx(x).
Therefore, the first variation of IF at P is.
8F 4P E P2()),
(P) = (P)))==
(1+11og) log ((oe))))(()) log log Pm (x) du(xx)
(264)
SP
We note that this coincides with the formula of the Ist variation provided in (Lambert et al., 2022, Appendix F) (in the particular case of mixture of Gaussian).
Now, noting (x) 1 + log (J (xr) dP(v)) log (rr),, the first variation is a potential energy (P) f Ve du.
Thus, the gradient of F P (R))) is obtained by the conjecture in Appendix B.4 as, for all μ E P2(Rd) x E Rd.
Vww, (()) = Vwa OF
(v)
(265)
This is well the same formula obtained by computing the Taylor expansion.
If we want compute the gradient on Pa P2 (B we can take the Bures- Wasse-Wassrrstein gradient of the first variation instead of the Wasserstein gradient. Since it is a potential energy. by (Diao et al., 2023, Lemma 3.1), for any P E P2 (B((BRR)) and y BW(R)), Rd.
ÖF (P))))))) = / VVp du+ (f du) (x mu):
Wew ((P)())()) VBw 8P
(266)
with mn = f zdu(x)) Since the tangent space of the Bures-Waas Wasserstein space is of the form TRBW(( = {a + m + S(x - m2); m,); m E Rd, S {R) (see egg. (Diao et al., 2023, Appendix A)), then we can identify the mean and covariance part of the gradient as (J Vi ] VoV dp), which coincides well with the formula derived in (Lambert et al., 2022, Appendix F).
Lambert et al. (2022) experimented with F in practice by evolving Gaussian particles. Note however that they observed that, even though the KL divergence is (geodesically) convex in P2 for V convex, IF is not convex in Pa (BB(())) as the negative entropy is not.
Also related, Huix al. (2024) considered optimizing the KL over mixtures of Gaussian, but with fixed covariance observing that the objective can be seen as the KL between a mollified distribution and the target. They minimized it using Wasserstein gradient flows over the means of each mixture. Moreover, their scheme in that case can be seen as particular case of the one of (Lambert et al., 2022), described in (Huix al., 2024, Appendix B).
Also to solve Variational Inference problems, Lim & Johansen (2024) considered the family ęe, = k=(()))))(()) with parametric kernels ko satisfying f (o(r))) = for all 2 E Rdd, They solved this problem by minimizing the KI. divergence with a regularizer, using a gradient flow over Rdo x P2 (Rd:). Rønning et al. (2025) considered a mixture family of the form (2((pm) = +/ m k(x 22) with pm # E 82g- and minimized an ELBO using the Stein Variational Gradient Descent.
57


Reviews:
[{"review_id": "ticytQrzna", "reviewer": "Reviewer_asNc", "summary": "This paper introduces a framework for optimizing functionals over probability measures of probability measures by leveraging the Riemannian structure of this space to develop Wasserstein over Wasserstein (WoW) gradient flows. It provides a theoretical foundation for these flows and a practical implementation using Forward Euler discretization. The paper also proposes a functional objective, which is based on Maximum Mean Discrepancy with a Sliced-Wasserstein kernel, enabling computationally efficient gradient flow simulation. The method is applied to dataset modeling, treating datasets as mixtures of probability distributions corresponding to label-conditional distributions.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reading the paper and for your feedback. We answer your comments below. Please do not hesitate if you have other questions.\n\n**The continuity equation relies on the strong assumption that the velocity field vt is Lipschitz.** \n\nIn Proposition 3.7, we show that if $(\\mathbb{P}\\_t)\\_t$ is an absolutely continuous curve, then we have $\\|v_t\\|_{L^2(\\mathbb{P}_t)}\\le |\\mathbb{P}'|(t)$ and the continuity equation. We do not assume that the velocity field is Lipschitz.\n\n**The class alignment via 1-NN and majority vote seems heuristic.**\n\nAnother solution less heuristic to assign the classes is to compute an OT map between $\\mathbb{P}=\\frac{1}{C}\\sum_{c=1}^C \\delta_{\\mu_c}$ and the target $\\mathbb{Q}=\\frac{1}{C}\\sum_{c=1}^C \\delta_{\\nu_c}$ with the 2-Wasserstein distance as ground cost. We replicated the results of Figure 3 with this method, and observe that it gives similar results. We will fix this in the revision.\n\n\n**Recent methods like Dataset Condensation with Gradient Matching (DC) are absent from Table 1.**\n\nWe chose to compare with methods that only optimize a distance to generate new samples, as these are the are the most comparable approaches, and we do not claim to be state of the art for dataset distillation. Thus, we compared in Table 1 with Distribution Matching (DM), introduced in [1]. Notably, [1] already compares DM with DC (see their Table 1), showing that their performances are comparable.\n\n[1] Zhao, B., & Bilen, H. Dataset condensation with distribution matching. WACV 2023.\n\n**Ablation study where the target dataset has non-Gaussian class distributions.**\n\nAs shown in Table 2, our method outperforms in most settings prior works that assume Gaussian class distributions, suggesting that avoiding this assumption leads to better performance on image datasets.\n\n**Related work.**\n\nThank you for pointing to us this work, which seems relevant. \n\n**Can the method scale to larger datasets beyond MNIST variants?**\n\nPlease find the answer in the response of Reviewer LUYC.\n\n**How sensitive is the method to choices of L, tau, momentum , and kernel bandwidth?**\n\nThe method is sensitive to the kernel bandwidth when using the Gaussian SW kernel, as we showed in Figure 5 of Appendix D.2. In practice, we use the Riesz SW kernel, which does not require to tune a bandwidth.\n\nAs we use a gradient descent method, taking the step size too big will make the scheme diverge.\n\nWe use momentum to accelerate the convergence of the scheme, but we notice that it still converges when using no momentum, but slower, see Figure 11 in Appendix D.4.\n\nIn low dimension, the method is not very sensitive to the number of projections $L$ (see e.g. the Figure https://ibb.co/SX2gtntd for the ring experiment). In higher dimension such as for MNIST, it is more sensitive, and a relatively big number of projections improves the convergence as it provides a better approximation of the gradient. We will add the following Figure (https://ibb.co/pr2MwMLH) showing results for different values of $L$ for the generation of MNIST samples.\n\n**How does the computational cost of your method scale with the number of classes and samples per class?**\n\nIn our experiments, we focused on the minimization of the MMD with a kernel based on the Sliced-Wasserstein distance, which can be computed in $O(Ln\\log n)$ between $\\mu_n=\\frac{1}{n}\\sum_{i=1}^n \\delta_{x_i}$ and $\\nu_n=\\frac{1}{n}\\sum_{j=1}^n \\delta_{y_j}$. For $C$ classes with each class having $n$ samples, the MMD has therefore a total runtime complexity of $O(C^2 Ln(\\log n + d))$. Moreover, note that the pairwise Sliced-Wasserstein distances can be computed in parallel. \n\nIn contrast, OTDD requires first to compure $C^2$ pairwise Wasserstein distances which has a complexity of $O(C^2 n^3 \\log n)$ if they are computed between the empirical distributions, and of $O(C^2 (nd^2 + d^3))$ if they are computed using the Gaussian approximation. Then, it requires to compute a second OT problem between the $nC$ samples with cost $d\\big((x,c), (x,c')\\big) = \\|x-x'\\|^2 + W_2^2(\\mu_c,\\nu_{c'})$, which has a complexity in $O(n^3C^3\\log(nC))$ and can be reduced to $O(n^2C^2\\log(nC)/\\varepsilon)$ using the entropic regularized problem. \n\nWe compared the runtime for the transfer learning experiment in https://ibb.co/4R0GtczT. \n\n\n**The paper assumes compactness and connectivity of the manifold for theoretical results.**\n\nWe assume compactness and connectivity of the manifold for theoretical purpose, as we notably rely on the seminal results of [1]. Of course, this is not always the case in practical applications, and thus we acknowledge that there is still a gap here to make the theoretical derivations to hold on non-compact spaces. Nonetheless, it is reasonable to assume that the data lie on a compact space.\n\n[1] Schiavo, L. D. A Rademacher-type theorem on L2-Wasserstein spaces over closed Riemannian manifolds. Journal of Functional Analysis, 2020.", "claims_and_evidence": "The claim Riemannian structure of WoW space enables gradient flow dynamics is proved in Appendix and the effectiveness of SW-based MMD for dataset flows is validated empirically.\n\nThe claim of tractable implementation lacks runtime analysis or complexity comparisons (e.g., runtime, number of projections, number of measures) with alternatives like Sinkhorn, especially for large-scale datasets when authors claim to be tractable.", "methods_and_evaluation_criteria": "The methodology is sound, employing well-established optimal transport techniques and MMD functionals. The evaluation is based on classification accuracy improvements, which are a reasonable measure for transfer learning and dataset distillation tasks.   The use of SW kernels is suitable as SW distances avoid the O(N^2) cost of Wasserstein.  However, additional baselines and comparisons with alternative dataset adaptation methods would strengthen the evaluation.", "theoretical_claims": "Proposition 3.7 assumes the base space M is compact, which is restrictive for unbounded domains. The authors do not address how their results extend to non-compact cases.\n\nThe continuity equation (Eq. 3.6) relies on the strong assumption that the velocity field vt is Lipschitz.\n\nThe class alignment via 1-NN and majority vote seems heuristic.", "experimental_designs_or_analyses": "No ablation study on the number of SW projections L (fixed to 100).\n\nComparisons are limited to OTDD and basic baselines. Recent methods like Dataset Condensation with Gradient Matching are absent from Table 1.\n\nThe paper criticizes prior work for assuming Gaussian class-conditional distributions but not including an ablation study where the target dataset has non-Gaussian class distributions.", "relation_to_broader_scientific_literature": "The paper builds on a rich literature of optimal transport, gradient flows, and dataset adaptation. It extends previous works on dataset dynamics in probability space by introducing a hierarchical structure through the WoW distance. The connection to prior work on Wasserstein gradient flows and MMD functionals is well-discussed.", "essential_references_not_discussed": "A related work on sliced Wasserstein kernel can be useful to discuss \"Unbiased Sliced Wasserstein Kernels for High-Quality Audio Captioning\", Luong et al.", "other_strengths_and_weaknesses": "Strengths: The paper introduces a novel and theoretically grounded approach to dataset adaptation, demonstrating strong empirical results.\n\nWeaknesses: Experiments are only on *NIST datasets. Including more datasets on natural images could make the paper stronger e.g., CIFAR10, Imagenet.", "other_comments_or_suggestions": "No", "questions_for_authors": "Can the method scale to larger datasets beyond MNIST variants?\n\nThe experiments use fixed hyperparameters (e.g., L=500 projections for SW, \\tau=0.1 step size). How sensitive is the method to choices of L, \\tau, momentum , and kernel bandwidth? Are there guidelines for tuning these in practice?\n\nThe WoW gradient flow involves nested Wasserstein computations. How does the computational cost of your method scale with the number of classes  and samples per class? Could you provide a comparison of efficiency (runtime and memory usage) against methods like OTDD on large-scale datasets?  \n\nThe paper assumes compactness and connectivity of the manifold for theoretical results. How do these assumptions impact practical applications where data may lie on non-compact spaces?", "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": "U297jb3WDQ", "reviewer": "Reviewer_LUYC", "summary": "The paper proposes a framework for optimizing functions over probability measure spaces  of probability measures. The approach is based on Wasserstein over Wasserstein gradient flows. The main contribution is a theoretical definition of this flow. The author also introduces objectives that are tractable within this framework. The approach is validated on synthetic datasets and small vision datasets for distribution flows, domain adaptation and dataset distilaltion.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your appraisal and positive comments on our paper. We answer your comments below.\n\n**The qualitative results (Figures) are not entirely convincing and do not clearly suggest a continuous flow from one dataset to another.**\n\nWe observed empirically that the flow goes very fast from the source dataset towards the neighborhood of the target dataset, but takes more time to converge towards clean images of the target dataset. This might explain why the flows do not appear very continuous, e.g. on Figure 2, as the discretization in time is linear. On Figure 9, we reported a finer discretization, which shows slightly better the behaviour.\n\nWe will also add in the zip supplementary materials the following gif showing the evolution of the rings (https://ibb.co/5qgjhgC).\n\n\n**How can the method be applied to larger-scale problems?**\n\n\nAs the computational complexity of computing the MMD with a Sliced-Wasserstein kernel is in $O(C^2 Ln(\\log n + d))$ for datasets with $C$ classes, and $n$ samples in each class, we believe that we can scale this algorithm to larger datasets with more classes and more samples by class.\n\n\n\nThe bottleneck will probably be for higher dimensional datasets. We tried on a (relatively) higher dimensional dataset CIFAR10 (see Figure 10 in Appendix D.4), and observed that the flow scales well. However, it required a lot more optimization steps to converge (150K steps for CIFAR10 against 18K steps for MNIST) with the same number of Monte-Carlo samples.\n\nTo scale to higher dimensional datasets, there could be several solutions. On one hand, one can try to get the number of projections $L$ bigger to get a better approximation of the Sliced-Wasserstein distance, or use better approximations using e.g. control variates as in [1]. We could also use faster optimization algorithms to try to accelerate the speed of convergence or other variants of the Sliced-Wasserstein distance more adapted to images, e.g. using convolution as projections as in [2,3]. We leave these investigations for future works.\n\nWe also replicated the domain adaptation experiment (Figure 3) by flowing from the dataset SVHN to CIFAR10, which are both dataset of shape 32x32x3, see https://ibb.co/C3B3Ffty. It demonstrates that it still works in moderate higher dimension. We will put it in Figure 3 in place of the KMNIST to MNIST example.\n\n\n[1] Leluc, R., Dieuleveut, A., Portier, F., Segers, J., & Zhuman, A. Sliced-Wasserstein estimation with spherical harmonics as control variates. International Conference on Machine Learning (2024).\n\n[2] Nguyen, Khai, and Nhat Ho. \"Revisiting sliced Wasserstein on images: From vectorization to convolution.\" Advances in Neural Information Processing Systems 35 (2022).\n\n[3] Du, C., Li, T., Pang, T., Yan, S., & Lin, M. \"Nonparametric generative modeling with conditional sliced-Wasserstein flows.\" International Conference on Machine Learning (2023).\n\n\n**Typos**\n\nThank you, we corrected the typo.", "claims_and_evidence": "Claim 1: Flowing datasets can be represented as probability distributions over probability distributions.\n\nClaim 2: A differential structure exists in the space of Wasserstein distances over Wasserstein distance.\n\nThese claims are supported by a rigorous mathematical foundation, utilizing optimal transport, Riemannian structures, and geodesics.\n\nClaim 3:\nThis approach can be applied to distribution flows, domain adaptation, and dataset distillation.\n\nThis claim is supported by experimental results (see next section).", "methods_and_evaluation_criteria": "The approach is first illustrated on a synthetic dataset (Three Rings). The qualitative results in this case are convincing.\nThen, the approach is validated on several tasks:\n- Domain Adaptation: Qualitative results of data flowing from MNIST to other datasets are presented. A classification task is then proposed, where a network is pretrained on a dataset, and accuracy is measured as the data flows from the initial dataset to another one (with aligned classes).\n- Dataset Distillation: The WoW gradient flow is used to generate a condensed dataset that maintains high classification accuracy.\n- Transfer Learning: The WoW gradient flow is also applied to transfer learning, generating a condensed dataset that preserves high classification accuracy.\n\nFor all experiments, the quantitative results support the initial claims and are convincing. However, there are some limitations:\n- The qualitative results (Figures) are not entirely convincing and do not clearly suggest a continuous flow from one dataset to another.\n- The results are limited to very small datasets and toy problems.", "theoretical_claims": "I reviewed the proofs, but I am not expert enough to verify their correctness.", "experimental_designs_or_analyses": "The experimental design and analysis are well-suited to the problem, except for the limitation in problem size.", "relation_to_broader_scientific_literature": "The paper is well-situated within the current literature, and the authors demonstrate a strong understanding of the state of the art in the field.", "essential_references_not_discussed": "no.", "other_strengths_and_weaknesses": "Strengths:\n- The paper is very rigorous and mathematically strong.\n- Introducing a tractable gradient flow for probability distributions over probability distributions, even if limited to small problems, is an important contribution.\n\nWeakness:\n- The experiments are restricted to very small datasets.", "other_comments_or_suggestions": "Typos :\n- In this experiment, we generate sample from different datasets starting from Gaussian noise -> In this experiment, we generate **samples** from different datasets starting from Gaussian noise", "questions_for_authors": "How can the method be applied to larger-scale problems?", "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": "0NUFC5Tygp", "reviewer": "Reviewer_yRk5", "summary": "This work presents a novel theoretical framework for handling probability distributions over probability distributions. The authors begin by rigorously establishing the Wasserstein over Wasserstein (WoW) distance metric from a functional analysis perspective, providing a mathematical foundation for measuring distances between distributions of distributions. Building upon this metric structure, they develop a comprehensive theory of gradient flows, incorporating maximum mean discrepancy (MMD) with Sliced-Wasserstein based kernels as a tractable objective functional. The theoretical framework is extensively validated through diverse experimental scenarios. These include transformations between synthetic distributions (demonstrated through a three-ring scenario), domain adaptation tasks, dataset distillation, and transfer learning applications. To ensure accessibility and reproducibility, the authors provide detailed mathematical derivations, comprehensive background material, and thorough theoretical analyses. This includes explicit connections to optimal transport theory, probability theory, functional analysis, and variational inference, making the work both theoretically rigorous and practically applicable.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your appraisal and positive comments on our paper. We address your comments below.\n\n**Domain Specification of the Functional**\n\nFor most of ML applications where we aim at minimizing distances, we agree that the codomain is $\\mathbb{R}_+\\cup \\\\{+\\infty\\\\}$. Nonetheless, the differential structure holds also for $\\mathbb{R}$ as codomain, which may cover e.g. potential energies with $V$ negative.\n\n**Higher-Order Terms**\n\nThe $o(W_{W_2})$ notation can be defined using the Landau notation. The main difficulty when using these terms in the article is whether we can integrate them when using the Taylor expansion. This can be handled using regularity hypothesis as bounded Hessian.\n\n**Continuity Equations in WoW Framework**\n\nWe expect a similar continuity equation $\\partial\\_t \\mathbb{P}\\_t = \\mathrm{div}(\\mathbb{P}\\_t \\nabla_{W\\_{W\\_2}}\\mathbb{F}(\\mathbb{P}\\_t))$. In Proposition 3.7, we took a first step by showing that AC curves satisfy a continuity equation. However, to derive the equation of the WoW gradient flow, we would need to show that the minimizing movement scheme converges towards the right equation (see e.g. (Santambrogio, 2017)), and to define an appropriate notion of divergence operator on this space, e.g. the one proposed in (Schiavo, 2020). We leave these questions for future work.\n\n**Comparisons with relevant baselines.**\n\nThank you for these suggestions. KALE flows and the flows of the KKL divergence are not designed to optimize over $\\mathcal{P}(\\mathcal{P}(\\mathbb{R}^d))$. In Figure 4, we showed an example of the flow of the MMD with Riesz kernel, which is designed on $\\mathcal{P}(\\mathbb{R}^d)$ as KALE flows and the KKL divergence.\n\n**Related works.**\n\nThank you for pointing to us these related works which we were not aware of. \n\n**Limitations and future research directions.**\n\nWe will add a paragraph in the revised version of the paper indicating future research directions (e.g. using other kernels, f-divergences or continuous labels) and limitations (e.g. theory on compacts manifolds, lack of continuity equation for the flow, dimension of xps).\n\n**This work mainly focus on the MMD functional.**\n\nWe focus on the MMD functional as it can be decomposed as a potential and an interaction energy, which we know how to differentiate in the WoW space (see Section 4.2).\n\nUsing this approach for f-divergences could be of great interest to do e.g. sampling. We plan to tackle this in future works, but we note that it is not straightforward, as we would need to identify a base measure w.r.t which the measures need to be AC (see e.g. (Schiavo, 2020)), and then compute the WoW gradient of these functionals.\n\n**Adding a GIF for 3 ring case.**\n\nThank you for this suggestion. We will add https://ibb.co/5qgjhgC.\n\n**1.Labels not categorical?**\n\nHandling continuous labels is an interesting extension of our method. This can be viewed as conditioning with continuous labels, akin to an infinite mixture. More formally, we could define it as $\\mathbb{P}=\\int \\delta_{\\mu_y}\\mathrm{d}\\lambda(y)$ with $\\mu_y$ the conditional distribution given the label $y$ and $\\lambda$ a distribution over the labels. For a discrete distribution $\\lambda$, this would recover the discrete case we use in the paper. In practice,  we would need to discretize $\\lambda$.\n\n**2. Why the authors consider the SW distance rather than wasserstein or sinkhorn?**\n\nWe consider SW for two reasons: \n1. It is much faster to compute than the Wasserstein or Sinkhorn distance (complexity of $O(Ln\\log n)$ for SW versus $O(n^3\\log n)$ for Wasserstein and $O(n^2\\log n/\\varepsilon^2)$ for Sinkhorn, see e.g. [1]).\n2. SW allows to define valid positive definite kernels [2], which is not the case for the Wasserstein distance [1].\n\n[1] Peyré, G., & Cuturi, M. Computational optimal transport: With applications to data science. Foundations and Trends in Machine Learning, 2019.\n\n[2] Kolouri, S., Zou, Y., & Rohde, G. K. Sliced Wasserstein kernels for probability distributions. IEEE Conference on Computer Vision and Pattern Recognition. 2016.\n\n**3. The loss curve in the Rings.ipynb appears to fluctuate.**\n\nThe loss curve may fluctuate as SW is estimated through a Monte-Carlo approximation. We observed empirically that for a sufficient number of iterations and with well chosen hyperparameters, it converges well and it fluctuates around small values, close to the minimum.\n\n**4. Is it possible to extend the proposed approach to constrained support?**\n\nThe proposed flow might be combined with methods to constrain the support such as Mirror Descent, which has been recently extended to Wasserstein gradient flows in [1], or using barrier methods as in [2].\n\n[1] Bonet, C., Uscidda, T., David, A., Aubin-Frankowski, P. C., & Korba, A. Mirror and preconditioned gradient descent in wasserstein space. NeurIPS 2024.\n\n[2] Li, L., Liu, Q., Korba, A., Yurochkin, M., & Solomon, J. Sampling with mollified interaction energy descent. ICLR 2023.", "claims_and_evidence": "The claims made in this submission are thoroughly supported by both theoretical foundations and empirical evidence. Specifically:\n\n1. Theoretical Rigor:\n- The mathematical framework is developed with precise definitions and rigorous derivations\n- All theoretical claims are supported by formal proofs.\n\n2. Empirical Validation:\n- The experimental results comprehensively validate the theoretical framework across multiple scenarios:\n  * Synthetic distributions (three-ring transformation)\n  * Domain adaptation tasks\n  * Dataset distillation\n  * Transfer learning applications\n- Each experiment provides quantitative metrics or qualitative visualizations that convincingly demonstrate the method's effectiveness\n\n3. Technical Soundness:\n- The implementation details are thoroughly documented\n- The experimental protocols are clearly specified and reproducible\n\n\nIn conclusion, the submission maintains a high standard of scientific rigor, with all claims well-supported by both theoretical analysis and experimental results.", "methods_and_evaluation_criteria": "Yes, the proposed methods and evaluation criteria are well-designed and appropriate for demonstrating the effectiveness of the proposed approach.", "theoretical_claims": "As an applied mathematics researcher, I have carefully examined most of the theoretical derivations and proofs in the manuscript. While the mathematical framework is generally sound and well-justified (for me), I would like to raise several technical concerns that merit further discussion:\n\n1. Domain Specification of the Functional:\n- On page 3, the authors define $\\mathcal{F}:\\mathcal{P}_2(\\mathcal{M})\\to\\mathbb{R}$\n- For most machine learning applications, particularly those involving distances or discrepancy measures, shouldn't the codomain be specifically $\\mathbb{R}^+ \\cup$ {0}?\n\n2. Treatment of Higher-Order Terms:\n- The derivations frequently approximate higher-order terms as $o$ terms\n- While this is a common practice in differential geometry, its applicability to $\\text{W}_{\\text{W}_2}$ warrants further justification\n- My understanding is that the finite second moment property of Wasserstein space might justify this approximation, but is this operation still suitable for WoW?\n\n3. Continuity Equations in WoW Framework:\n- An interesting theoretical extension would be the formulation of continuity equations within the $\\text{W}_{\\text{W}_2}$ framework for functional $\\mathcal{F}$. Specifically, for $\\mathcal{F}[\\rho]$ in Wasserstein space, we may define $\\frac{\\partial \\rho}{\\partial \\tau}=-\\nabla\\cdot(\\rho\\nabla\\frac{\\delta \\mathcal{F}[\\rho]}{\\delta \\rho})$. Can we have similar a PDE in WoW ?", "experimental_designs_or_analyses": "I have examined the experimental designs and suggest two specific improvements:\n\n1. Ring Transformation Experiments:\n- Include comparisons with relevant baselines:\n  * KALE flow [1]\n  * Kernel KL divergence [2]\n\n---\n\nReferences:  \n1. KALE Flow: A Relaxed KL Gradient Flow for Probabilities with Disjoint Support  \n2. Statistical and Geometrical properties of regularized Kernel Kullback-Leibler divergence", "relation_to_broader_scientific_literature": "Yes, the key contributions of the paper are foundational and have potential implications for broader applications of the proposed approach.", "essential_references_not_discussed": "The paper does not currently discuss related works on semi-implicit variational inference, which are similar to the proposed approach in appendix E.2. For example: \n1. Variational Bayes with Stein Mixture Inference  \n2. Particle Semi-Implicit Variational Inference", "other_strengths_and_weaknesses": "## Strengths:\n1. The paper is realted to ICML conference.\n2. The experimental results are sufficient.\n\n## Weaknesses:\n1. The limitations and future research directions are not given explictly.\n2. The experiments can be conducted on the recommender system scenario for a wider impact.\n3.  This work mainly focus on the MMD functional, the reviewer doubts that how to implement such approach for $f$-divergence, which is of great importance to design sampling algorithm.", "other_comments_or_suggestions": "1. Presentation: Including a visualization of the gradient for the deep learning backend pipeline in Section 4.1 could improve clarity and make the simulation workflow easier to understand.  \n\n2. Adding an MP4 or a GIF for three ring case could more effectively showcase the superiority of the proposed approach in the supplementary zip file.", "questions_for_authors": "1. **How does the proposed method handle scenarios where the labels are not categorical, such as in transfer learning between regression tasks**?\n2. Why the authors mainly consider the sliced wasserstein distance rather than the vanilla wasserstein distance or sinkhorn distance?\n3. The loss curve in the `Rings.ipynb` appears to fluctuate. How can the convergence of the proposed approach be effectively validated?\n4. Is it possible to extend the proposed approach to constrained support? For example, Dirichlet distribution on the manifold.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: Through motivating applications like domain transfer and distillation, the novel task of dataset flow, where data is itself a measure is considered i.e., flow involving measures of measures. Firstly, the WOW metric is defined, which is nothing but the Wasserstein whose ground metric is another Wasserstein. It geometry is formally studied. Then, JKO flows under this metric are formally studied. simulations on benchmarks show the efficacy of the methodology.

All the reviewers acknowledge that the problem setting is novel and interesting. The technical contribution is enough for a publication.
Decision: Accept (oral)