Mirror and Preconditioned Gradient Descent in Wasserstein Space
Authors: Clément Bonet
Théo Uscidda
Adam David
Pierre-Cyril Aubin-Frankowski
Anna Korba
Venue: NeurIPS 2024 spotlight (2024)
Keywords: wasserstein gradient flows
mirror descent
preconditioned gradient descent
Primary area: optimization
TL;DR: We study mirror descent and preconditioned gradient descent on Wasserstein space.

Abstract: As the problem of minimizing functionals on the Wasserstein space encompasses many applications in machine learning, different optimization algorithms on $\mathbb{R}^d$ have received their counterpart analog on the Wasserstein space. We focus here on lifting two explicit algorithms: mirror descent and preconditioned gradient descent. These algorithms have been introduced to better capture the geometry of the function to minimize and are provably convergent under appropriate (namely relative) smoothness and convexity conditions. Adapting these notions to the Wasserstein space, we prove guarantees of convergence of some Wasserstein-gradient-based discrete-time schemes for new pairings of objective functionals and regularizers. The difficulty here is to carefully select along which curves the functionals should be smooth and convex. We illustrate the advantages of adapting the geometry induced by the regularizer on ill conditioned optimization tasks, and showcase the improvement of choosing different discrepancies and geometries in a computational biology task of aligning single-cells.

Mirror and Preconditioned Gradient Descent in Wasserstein Space
Clément Bonet CREST, ENSAE, IP Paris clement. bonet@ensae.fr
Théo Uscidda CREST, ENSAE, IP Paris theo www.uudddeeseaee.......
Adam David Institute of Mathematics Technische Universität Berlin david@math, tu-berlin de
Pierre-Cyril Aubin-Frankowski
Anna Korba CREST, ENSAE, IP Paris anna. kon.kebrbeensae... fr
TU Wien
pierre-cyril aubin@turinn euerre-oyrr..aaiineeui...n ac. at
Abstract
As the problem of minimizing functionals the Wasserstein space encompasses many applications in machine learning, different optimization algorithms on Re have received their counterpart analog the Wasserstein space. We focus here on lifting two explicit algorithms: mirror descent and preconditioned gradient descent. These algorithms have been introduced to better capture the geometry of the function to minimize and are provably convergent under appropriate (namely relative) smoothness and convexity conditions. Adapting these notions the Wasserstein space, we prove guarantees of convergence of some Wasserstein- gradient-based discrete-time schemes for new pairings of objective functionals and regularizers. The difficulty here is to carefully select along which curves the functionals should be smooth and convex. We illustrate the advantages of adapting the geometry induced by the regularizer on ill-conditioned optimization tasks, and showcase the improvement of choosing different discrepancies and geometries in a computational biology task of aligning single-cells.
Introduction
Minimizing functionals on the space of probability distributions has become ubiquitous in machine
learning for e.g. sampling [13, 129], generative modeling [53, 86), learning neural networks [36, 90, 107], dataset transformation [4. 63], or modeling population dynamics [23, 120]. It is a challenging task is dufinite-dimensionee problem. Wasserstein gradient flows [5] provide an elegant way to solve such problems on the Wasserstein space, i.e., the space of probability distributions with bounded second moment, equipped with the Wasserstein-2 distance from optimal transport (OT). These flows provide continuous paths of distributions decreasing the objective functional and can be seen as analog Euclidean gradient flows [111]. Their implicit time discretization, referred to as the JKO scheme [66], has been studied in depth [1, 26, 95, 111]. In contrast, explicit schemes, despite being easier to implement, have been less investigated. Most previous works focus on the optimization of a specific objective functional with time-discretation of its gradient flow with the Wasserstein-2 metrics. For instance, the forward Euler discretization leads to the Wasserstein gradient descent. The latter takes the form of gradient descent (GD) the position of particles for functionals with closed-fomm over discrete measures, e.g. Maximum Mean Discrepancy (MMD), which can be of interest to train neural networks [7, 30]. For objectives involving absolutely continuous measures, such the Kullback-Leibler (KL) divergence for sampling, other discretizations can be easily computed such the Unadjusted Langevin Algorithm (ULA) [106]. This leaves the question open of assessing
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

the theoretical and empirical performance of other optimization algorithms relying on alternative geometries and dime-discretization..
In the optimization community, a recent line of works has focused on extending the methods and convergence theory beyond the Euclidean setting by using more general costs for the gradient descent scheme [77]. For instance, mirror descent (MD), originally introduced by Nemirovskij and Yudin [92] to solve constrained convex problems, uses cost that is divergence defined by a Bregman potential [12]. Mirror descent benefits from convergence guarantees for objective functions that are relatively smooth in the geometry induced by the (Bregman) divergence [88], even if they do not have a Lipschitz gradient, i.e., are not smooth in the Euclidean sense. More recently, a closely related scheme, namely preconditioned gradient descent, was introduced in [89]. It can be seen as a dual version of the mirror descent algorithm, where the role of the objective function and Bregman potential are exchanged.
particular, its convergence guarantees can be obtained under relative smoothness and convexity of the Fenchel transform of the potential, with respect to the objective. This algorithm appears more efficient to minimize the gradient magnitude than mirror descent [68]. The flexible choice of the Bregman divergence used by these two schemes enables to design or discover geometries that are potentially more efficient.
Mirror descent has already attracted attention in the sampling community, and some popular algo- rithms have been extended in this direction. For instance, ULA was adapted into the Mirror Langevin algorithm [3, 32, 62, 64, 79, 121, 132]. Other sampling algorithms have received their counterpart mirror versions such as the Metropolis Adjusted Langevin Algorithm [116], diffusion models [82], Stein Variational Gradient Descent (SVGD) [1 [114), or even Wasserstein gradient descent [113). Pre- conditioned Wasserstein gradient descent has been also recently proposed for specific geometries in [31, 44] to minimize the KL in more efficient way, but without an analysis in discrete time. All the previous references focus on optimizing the KL as an objective, while Wasserstein gradient flows have been studied in machine learning for different functionals such as more general f-divergences [6, 93], interaction energies [19, 78]. MMDs [7, 30, 59, 60, 72] or Seiced-Waseesttein (SW) distances [15, 18, 45, 86). In this work, we propose to bridge this gap by providing a general convergence theory of both mirror and preconditioned gradient descent schemes for general target functionals, and investigate as well empirical benefits of alternative transport geometries for optimizing functionals on the Wasserstein space. We emphasize that the latter is different from [9, 67]. wherein mirror descent is defined in the Radon space of probability distributions, using the flat geometry defined by TV or 2 norms on measures, see Appendix A for more details.
Contributions. We are interested in minimizing functional F P2(RR) + R [+00} over probability distributions, through schemes of the form, for T > 0, k 2 0,
1 -Tddd (2(e)) + a Id), Mk+1 = (Tk+1)
R++1 argmin TEL?(H))
with different costs d : ())())) x (2x + R4, and in providing convergence conditions. While we can recover a map T = Ta o T--1 T1 such that μ = To. the scheme (1) proceeds by successive regularized linearizations retaining the Wasserstein structure, since the tangent space to (2dR at ju is a subset of (p) [96]. This paper organized as follows. In Section 2, we provide some background on Bregman divergences, as well as on differentiability and convexity over the Wasserstein space. In Section 3, we consider Bregman divergences on (,) for the cost in (1), generalizing the mirror descent scheme to the Wasserstein space. We study this new scheme by discussing its implementation, and proving its convergence under relative smoothness and convexity assumptions. In Section 4, we consider alternative costs in (1), that are analogous to OT distances with invariant cost, extending the dual space preconditioning scheme to the latter space. Finally, in Section 5, we apply the two schemes to different objective functionals, including standard free energy functionals such as interaction energies and KL divergence, but also to Sinkhorn divergences [50] or SW [16, 102] with polynomial preconditioners on single-cell datasets.
Notations. Consider the set (2( of probability measures l on Rd with finite second moment and (2,,r((.)) Pa (IR) its subset of absolutely continuous probability measures with respect to the Lebesgue measure. For any μ (2(()),,, we denote by (r²()) the Hilbert space of functions f: .Rd Rd such that f [|
co equipped with the norm ||-11 (2)) and inner product (,, ()))) For Hilbert space X, the Fenchel transform of f X - R is f"
supzex (ar,) - f(x).
Given measurable map T Rd - _ and μ (2( T T u is the pushforward measure of by
2

---

т; and T + = T x)dd(z). For u. P2(R"), the Wasserstein-2 distance is ( (2, 2) = infzeM(a.)) f || -ve-yll yll v), where IT = {7 P(R Rd), TH 7 14, n Y = 4 with (2) vi, is the set of couplings between l and u, and denote by I,, 2) the set of optimal couplings. When the optimal coupling is of the form 7 = (Id, ( with Id x r x and Tи E (2) satisfying #1 = we call Ty the OT map. We refer to the metric space (P2(R"), W2) as the Wasserstein space. We note St (IR) the space of symmetric positive definite matrices, and for I E RR >' 5 || x Ex.
Background
In this section, we fix μ E (R() and introduce first the Bregman divergence (1) along with the notions of relative convexity and smoothness that will be crucial in the analysis of the optimization schemes. Then, we introduce the differential structure and computation rules for differentiating functional F P2 (R)))) + I along curves and discuss notions of convexity on (2 We refer the reader to Appendix B and Appendix C for more details on (n) and the Wasserstein space respectively. Finally, we introduce the mirror descent and preconditioned gradient descent on gd
Bregman divergence on (H)).. Frigyik et al. [54, Definition 2.1] defined the Bregman divergence of Fréchet differentiable functionals. In our case, only need Gâteaux differentiability. In this paper, V refers to the Gâteaux differential, which coincides with the Fréchet derivative if the latter exists.
Let 0₂ (H) - IR be convex and continuously Gâteaux differentiable. The Bregman divergence is defined for all T.S E (1) as do, (T,S) = (T. - o2 -
Definition 1.
LP(N)-
We use the same definition on Rd. The map on (respectively Φ ) in the definition of dg, above is referred as the Bregman potential (respectively mirror map). If o, is strictly convex, then dæ, is valid Bregman divergence, i.e. it is positive and separates maps y--amost everywhere (a.e.). In partic- ular, for (T) = [ we recover the L norm as a divergence de,, (T, = T Bregman divergences have received lot of attention as they allow to define provably convergent schemes for functions which are not smooth in the standard (e.g. Euclidean) sense [1 88), and thus for which gradient descent not appropriate. These guarantees rely on the notion of relative smoothness and relative convexity [88, 89]. which we introduce now on L²()).
Definition 2 (Relative smoothness and convexity). Let Vun oy. (m) + R be conver and continu- ously Gâteaux differentiable. We say that 1b, is B-smooth (respectively a a-convex) relative to on if and only if for all S 1.2 (p),, dy (T, < Bds, (T,S) (respectively dy, (T,S) adl, (T,S)).
Similarly to the Euclidean case [88]. relative smoothness and convexity are equivalent to respectively 36,- v and b a being convex (see Appendix B.2). Yet, proving the convergence of (1) requires only that these properties hold at specific functions (directions), a fact we will soon exploit.
In some situations, compute its Bregman divergence døg- We show Lemma 18 that a sufficient condition to satisfy this property is for Φn to be strictly convex, lower semicontinuous and superlinear, i.e. lim oco 6n (TT//[[]|>>(/))) = +oo. Moreover, in this case, (Vo,)) Vol. When needed, we will suppose that Φμ satisfies this assumption.
need the I Fenchel transform 6% of Φr to be differentiable, e.g. to
Differentiability (P2(R1), W2). Let F. P2(RR) + (+oo), and denote D(F) {u (R( F(r) < +oo) the domain of F and D(F,)== 1²(H)), Ty E D(F)) the domain of Fr, defined as (T) := #(T for all T T.²()).. the following, we use the differential struc- ture of (P2 (R.)), W2) introduced in [17, Definition 2.8], and we say that Fw, is a Wasserstein gradient of F at μ E D(F) if for any & E (2(d)) and any optimal coupling 7 E lI(,,
F(v) F(H)+ + / (VW, x) dy(x, v) + o(W2(H,)))
If such a gradient exists, then we say that F Wasserstein differentiable at ju [17, 74]. Moreover there is unique gradient belonging to the tangent space of (2(d)) verifying (2) [74, Proposition
3

---

2.51, and we will always restrict ourselves without loss of generality to this particular gradient, see Appendix C.I. The differentiabilityy of F and Fr are very related, as described in the following proposition.
be Wasserstein differentiable functional on D(F). Let μ E P2(RR)) and F. (T) #(T for all T E D(Fu). Then, F₁ is Fréchet differentiable, and for all S E p) (S) = Fw S.
Proposition 1. Let F: (2(( + R
The Wasserstein differentiable functionals include e- <-Wasserstein costs P2.a (TE.) [74, Propo- sition 2. 2.10 and 2.11], potential energies V(H) = Vdp or interaction energies W(H) [f (xx - 3) du(x))uuly) for V Rd + IR and W : -R - R differentiable and with bounded Hessian [74, Section 2.4]. In particular. their Wasserstein gradients read as Vw, V(H) VV and V ((w) = VW * μ. However, entropy functionals, e.g. the negative entropy defined as H(u)) = log (p(=)) for distributions f admitting a density w.r.I. the Lebesgue measure, are not Wasserstein differentiable. In this case, we can consider subgradients F()) at μ for which (2) becomes inequality. To guarantee that the Wasserstein subgradient is not empty, we need p to satisfy some Sobolev regularity, see e.g. [5, Theorem 10.4.13] or [108]. Then, if V log p E (H),, the only subgradient of H. in the tangent space is Vw, H(p) = V log p. see [5, Theorem 10.4.17] and [47, Proposition 4.3). Then, free energies are functionals that write as sums of potential, interaction and entropy terms [110, Chapter 7]. It is notably the case for the KI. to fixed target distribution, that is the sum of a potential and entropy term [129], or the MMD as a sum of a potential and interaction term [7].
Examples of functionals. The definitions of Bregman divergences on L L2( and of Wasserstein differentiability enable to consider alternative Bregman potentials than the L (p)) mentioned above. For instance, for V convex, differentiable and L-smooth, we can use potential energies %% (T) (1) for which day (T, Jd ((=))d(()) where dy is the Bregman divergence of V on Rd. Notice that (T) = 11TT122(0) is specific example of a potential energy where V = } || 112. Moreover, we will consider interaction energies @ (T) == W(T # (J) with W convex, differentiable, I-smooth, and satisfying W(-x) = (x)); for which daw (T,S) = [f dw (T(z) T(x'), - ((())dd())) (see Appendix I.3). We will also use 6f (T) #((T#u) with H. the negative entropy. Note that Bregman divergences on the Wasserstein space using these functionals were proposed by Li [80], but only for S = Id and optimal transport maps T.
Convexity and smoothness in (P2(),), W2). In order to study the convergence of gradient flows and their discrete-time counterparts, it is important to have suitable notions of convexity and smoothness. On (P2(R"), W2). different such notions have been proposed based on specific choices of curves. The most popular one is to require the functional F to be a-convex along geodesics (see Definition 10), which are of the form Me = ((1 - t)ld + (TTn # No if po E P2. (RR) and
M1 E P2(dd,, with TH the OT map between them. In that setting,
a 2 W (3((0 111 - [T - TdllE=(wa)) S F(m) - F(vo) - Ww,o), TH: - Td) ((oo)-
(3) For instance, free energies such as potential or interaction energies with convex V or W, or the negative entropy, are convex along geodesics [110, Section 7.3]. However, some popular functionals, such as the Wasserstein-2 distance j + W 77) itself, for a given 7) E P2 (K.)., are not convex along geodesics. Instead Ambrosio et al. [5, Theorem 4.0.4] showed that it was sufficient for the convergence of the gradient flow to be convex along other curves, e.g. along particular generalized geodesics for the Wasserstein-2 distance [5, Lemma 9.2.7], which, for p,, E (R())) are of the form 14: ((1 - T)T + (T²) " for TH Ty OT maps from 1) to μ4 and v. Observing that for
2
(T) = / |ITT we can rewrite (3) as adde, (TA) pag Id) S di Pg Agg Id] and see that being (TH,,
convex along geodesics boils down to being convex in the L sense for Id and T chosen as an OT map. This observation motivates us to consider more refined notion of convexity along curves.
Definition 3. Let μ E (2(Rd), T,S E (n)) and for all E [0, 11, pr = (Te) with T, = (1-1)S-
+ R is a-convex (resp. B-smooth) relative to G P2(RR)) - Th along t + pe for all s, E 00, 1). d (T., TT) 2 adg_ (T., T,) (resp. dz (T., T₁) S Bdg (T.: T.))
IT. We say that F: (R))

---

Notice that in contrast with Definition 2, Definition is stated for a fixed distribution μ and direc- tions (S, T), and involves comparisons between Bregman divergences depending on j and curves (Tx) depending on S, T. The larger family of S and T for which Definition 3 holds, the more re- stricted is the notion of convexity of F- ag (resp. of F) on P2 (R For instance, W generalized geodesics with anchor 7 E (RR))) correspond to considering S, T as all the OT maps originating from '/, among which geodesics are partieular cases when taking 7/ = jl (hence $ = Id). If we furthermore ask for c-convexity to hold for all μ Pa (IR and T, E (1)) (i.e., not only OT maps), then we recover the convexity along acceleration free-curves as introduced in [28, 98, 1 117]. Our motivation behind introducing Definition 3 is that the convergence proof' of MD and precon- ditioned GD require relative smoothness and convexity properties to hold only along specific curves.
Mirror descent and preconditioned gradient descent on Rd. These schemes read respectively as Vo(xi++1) - Vá((a) -T Vf((x) [12] and 3fk++1 - 1/2 = - V (Vy (yx)) [89], where the objec- tives f.g and the regularizers h, & are convex C functions from &d to IR. The algorithms are closely related since, using the Fenchel transform and setting 9 6* and h =
see that, for y = VΦ(#). the two schemes are equivalent when permuting the roles of the objective and of the regularizer. For MD, convergence of f is ensured if is both 1/r-smooth and a-convex relative to & [88, Theorem 3.1]. Concerning preconditioned GD, assuming that h, g Legendre, (9(8))) converges to the minimum of g if h" is both s/r-shooth and a-convex relative g" with > 0 [89, Theorem 3.9].
3 Mirror descent
For every μ E P2 (2 let o (x) + R be strictly convex, proper and differentiable and assume that the (sub)gradient F(H) E (n) exists. In this section, we are interested in analyzing the scheme (I) where the cost dl is chosen as Bregman divergence, i.e. dø, as defined in Definition 1. This corresponds mirror descent scheme in P2 (R.).. For 7 > 0 and k 2 0, it writes:
Tk+1
argmin døpe (T, VV,,
(a)e) Mk++1 (Tx+1)#14-
TETA(M))
Iterates of mirror descent. In all that follows, we assume that the iterates (4) exist, which is true e.g. for a superlinear ópp. since the objective is a sum of linear functions and of the continuous 6nu- In the previous section, interpreted as a linearization of the functional F at μa for Wasserstein (sub)differentiable functionals. Now define for all T E L (px),, J(T)= de (T, (VW, T-)) (((x)) Then, deriving the first order conditions of (4) as VJ(Tx+1) 0, we obtain pa-a.e.
have seen that the second term in the proximal scheme (4) can be
6 (T++1) Vo (Id) - T WW, F(rs) - Ta+1 = Voj @V (Id) - V w, F(w))) (5)
Note that for (T) = 1/2 ||T| (())) the update (5) translates as Tx+1 = Id -V w. F(ps). and our scheme recovers Wasserstein gradient descent [35, 91]. This is analogous to mirror descent recovering GD when the Bregman potential is chosen as the Euclidean squared norm I2 [12]. We discuss in Appendix D.2 the continuous formulation of (4), showing it coincides with the gradient flow of the mirror Langevin [3, 130J, the limit of the JKO scheme with Bregman groundcosts [104], Information Newton's 's flows [126], Sinkhorn's flow [41] for specific choices of 6 and F.
Our proof of convergence of the mirror descent algorithm will require the Bregman divergence to satisfy the following property, which is reminiscent of conditions of optimality for couplings in OT. Assumption 1. For (L, P E P2,ne and u E P2(R), setting T = argmin T de, (T, Id), &. 6 argminl p=w do, (U,Id) Id), the functional 6n is such that, for any S E (²)) satisfying S# = p. we have de, (TH, 5) 2 (U??"" Id).
The inequality in Assumption can be interpreted as follows: the "distance" between p and is greater when observed from an anchor μ that differs from P and v. We demonstrate that Bregman divergences satisfy this assumption under the following conditions on the Bregman potential d.
Let u.p E P2, (IRd) and E P2(IR)). Let oy be a pushforward compatible functional, i.e. there exists & (2 + R. such that for all T E I²² (T) s(T p)). Assume furthermore Vwy) and V W. d(p) invertible (on Rd), Then, op satisfies Assumption 1.
Proposition 2.
5

---

All the maps of of and 6 defined in Section 2 satisfy the assumptions of Proposition 2 under mild requirements, see Appendix D.I. The proof of Proposition 2 is given in Appendix H.2. It relies on the definition of an appropriate optimal transport problem
- (w:o)()))) - -) dy(2, v),
W Wo(
inf @(v) - (())- - 7Ell((.)))
(6)
and on the proof of existence of OT maps for absolutely continuous measures (see Proposition 15), which implies Wa(u, 14) = de. o Id) with To defined as in Assumption 1. From there, we can conclude that Φn satisfies Assumption 1. We notice that the corresponding transport problem recovers previously considered objects such OT problems with Bregman divergence costs [25, 103], but is strictly more general (as our results pertain to the existence of OT maps), as detailed in Appendix D.1.
We now analyze the convergence of the MD scheme. Under a relative smoothness condition along curves generated by = Id and T = Th+1 solutions of (4) for k 2 O, we derive the following de- scent lemma, which ensures that (F(we)), &. is non-inereasing. Its proof can be found in Appendix H.3 and relies on the three-point inequality [29], which we extended to L (p) in Lemma 29.
Proposition 3. Let B > 0, S +. Assume for all k 2 0, F is B-smooth relative & along t + ((1 t)ld + +Tx+1) ) # Mk. which implies Bdon, (Tx+1, Id) 2 dr Id). Then, for all
1 F(m++1) S F(pu) _dem, (Td,Tx+1). Tx+1)-
Assuming additionally the convexity of F along the curves μ4 = ((1- t)ld + (T # #M.. E [0. and that satisfies Assumption I, we can obtain global convergence.
Proposition 4. Let v E (2()))) a 2 O. Suppose Assumption and the conditions of Proposition hold, and that F is a-convex relative to e along the curves + ((1-t)ld + &T 4 4 Then, for all k 2 1,
F(PN) - F(v) (1- -1 Wolv,P) 1- W Wolv. (o)
(8)
Kt
Moreover, if > O. taking = p² the minimizer of F, obtain linear rate: for all k N 0. Wo(n", px) S (1 rat* Wolp", ,po).
The proof of Proposition 4 can be found in Appendix H.4, and requires Assumption to hold so that consecutive distances between iterates and the global minimizer telescope. This is not as direct as in the proofs of [88] over TR., because the minimization problem of each iteration (4) happens in different space L (pxx. We discuss in Section 5 how to verify the relative smoothness and convexity on some examples. In particular, when both F and 6 are potential energies, it is inherited from the relative smoothness and convexity on Rd, and the conditions are similar with those for MD on IR We also note that relative smoothness assumptions along descent directions as stated in Proposition 3 and relative strong convexity along optimal curves between the iterates and a minimizer as stated in Proposition 4 have been used already in the literature of optimization
measures in very specific cases, e.g. for descent results for the KL along SVGD [71] or for Sinkhorn convergence in [9]. We further analyze in Appendix F the convergence of Bregman proximal gradient scheme [11, 123] for objectives of the form F(u)= = (()) + H(1) with H non smooth; which includes the KI. divergence decomposed potential energy plus the negative entropy.
Implementation. We now discuss the practical implementation of MD on (P2(Rd), W2) as written in (5). If o is pushforward compatible, we have Von.. (T&+1) F0nn,(Tx+1] Vot, is unknown, the scheme implicit in Ta+I- possible solution is to rely on a root finding algorithm such as Newton's 's method to find the zero of VJ at each step, which we use in Section 5 for o as Bregman potential. However, this procedure may be computationally costly and scale badly w.E.. the dimension and the number of samples, see Appendix G. Nonetheless, in the special case u (T) = fV T du with V differentiable,, strongly convex and L-smooth, since Vw, V(n) = VV
= Vw, o (T ([T+1)) 44)))) ) o T but
and (VV)-1= VV*, the scheme reads as
Vk 2 0, T = VV* (VV- F(u))(
(9)

---

and can be implemented on particles, i.e. for Pr -Ei=i x = VV (VV(()
r VW, F(i))() for all 2 0, i E [I....,n... This scheme is analogous to MD in R
[12]
and has been introduced as the mirror Wasserstein gradient descent [113]. Moreover, for V V 12 as observed earlier, we recover the usual Wasserstein gradient descent, Le.. TR+1 = Id -i,,>,,,--nn The scheme can also be implemented for Bregman potentials that are not pushforward compatible. For
specific it recovers notably SVGD and its variants [83, 84, 1 114, 131] or the gradient descent [56]. We refer to Appendix D.4 for more details.
Preconditioned gradient descent
As seen in Section 2, preconditioned gradient descent Rd has dual convergence conditions compared to mirror descent. Our goal is to extend these to (1) and P2 Let r > 0, p E (2 and h Fd + R proper and strictly convex on Rd. We consider in this section (T) = f h T du and d(T, Id) = dhe ((dd - T)/T)r f [h((2 - T()))/T)r (x) This type of discrepancy is analogous to OT costs with translation- vraslation-invariant ground cost c(z, y) = h(r - yy). which have been popular as they induce an OT map [110, Box 1.12]. Such costs have been introduced egg. in [39, 70] to promote sparse transport maps. More generally, for or strictly convex, proper, differentiable and superlinear, we have (Vøn)= = Vo and the following theory is still valid. For simplicity, we leave studying more general for future works. Here, the scheme (1) results in:
/ x T(=)
(
T8+11 = argmin TELP(M))
ddu wv,, F(ux). T T-Id) (a(P)) Mk+1 = (Tx+1) ) # Mx-
(10)
Deriving the first order conditions similarly to (5) in Section 3, we obtain the following update:
vk 2 0, T Tx ld -d--rd 1---r(V -Vo (V.,
ld vVh" Vw,F(uu))
(11)
Notice that for h 3 - 1/ the squared Euclidean norm, df and of recover the squared (x) norm, and schemes (4) and (10) coincide. The scheme (10) is analogous to preconditioned gradient descent [68, 75, 76, 89, 119], which provides a dual alternative to mirror descent. For the latter, the goal is to find suitable preconditioner h² allowing to have convergence guarantees, or to speed-up the convergence for ill-conditioned problems. It was recently considered on the Wasserstein space by Cheng et al. [31] and Dong et al. [44] with a focus on the KL divergence as objective F and for - | with p > [31] or h quadratic [44]. Moreover, their theoretical analysis [1] was mostly done using the continuous formulation. Instead we focus on deriving conditions for the convergence of the discrete-time scheme (11) for more general functionals objectives.
Convergence guarantees. Inspired by [89], we now provide descent lemma on (0
under a technical inequality between the Bregman divergences of de and FaR for all ko 2 Additionally, we also suppose that is convex along the curves generated by S = Tx+1 and T = Id. This last hypothesis ensures that de (Tx+1,, Id) 2 0, and thus that (0h, (V..
is non- non-increasing. Analogously to the Fuclidean case, df: quantifies the mag- nitude of the gradient, and provides a second quantifier of convergence leading to possibly different efficient methods compared to mirror descent [68]. The proof relies mainly on the three-point identity (see egg. [54, Appendix B.7] or Lemma 28) and algebra with the definition of Bregman divergences.
Proposition 5. Let B > 0. Assume S B. and for all k 2 0, F convex along ((1 t)T
(Id) Pk and dea- (Vw, TA+1, F(ue)) < Bd Mk
(dd,)))). Then, for all l: 2 0,
1
oy oaars F(Heti)))))) s dh (Vw,F((m:)) - d (Te+1,
(12)
Under an additional assumption of a reverse inequality between the Bregman divergences of of and Foge and assuming that &/"" attains its minimum in O, we can show the convergence of the gradient quantified by o/i' (see Lemma 21), and the convergence of (F(uz)) towards the minimum of F. Proposition 6. Let of Proposition 5 hold, and that for T = Ta d;,, (Id, T), ad (Id, T) S FPx
and n² E P2(Rd)) be the minimizer of F. Assume the conditions
7

---

de-; (W.. T, Vw. m/ (0) (0),
Then, for all k 2 1, since Vw,F(u") = 0 and
- - TK
1 Ta (F(Ho) - F((x)))
ot (Vws)(uss) - (0) - (1 Tα)
(F(uo) - -F(+)))
(13)
Moreover, assuming that h² attains its minimum at and > 0, converges towards its minimum at linear rate, i.e. for all k 2 0, F(ju) F(μ²) S (1 ra)* (F(m))-F(p
The proofs of Proposition 5 and Proposition 6 can be found respectively in Appendix H.5 and Appendix H.6.
We now discuss sufficient conditions to obtain the inequalities between the Bregman divergences required in Proposition 5 and Proposition 6. Maddison et al. [89] showed on Rdd for cost I and an objective function g. that these conditions were equivalent to 8-smoothness and a-convexity of the preconditioner ht (analogous to (%,) relative to the convex conjugate of the objective g* (analogous to D.)) To write the inequalities we assumed as a relative smoothness/convexity property of o% w.E.t Fi Me we would need at least ensure that F is differentiable, in order to define its Bregman H& divergence according to Definition 1. This can be done egg. by assuming Fu,, strictly convex and superlinear (see Lemma 18). The latter is true for several examples of functionals F we already mentioned, such as potential or interaction energies with strongly convex potentials.
Under this assumption, we can show that the inequalities in Proposition 5 and Proposition 6 are implied by relative smoothness and convexity along suitable curves.
Proposition 7. Les μ E (2(R and T E (2) Assume Fi is Gâteaux differentiable and define Fū on t + μe (TT) (P)) FT, for T, (1 t)U + t8 for all t [0, 1], S,U E (1))).
8h is 3-smooth relative to F3 along ((1 F(m) tV,,
T) M. Then,
dop- (Vw,F(TT o T, Vw,F(1)) 1 Ady, (Id, T).
(14)
Likewise, if 6h is a-convex relative to Fj along + ((1 t) Vw, F(u) + +V T) then
#
adz, (Id, T) S døy (Vw, Vw,F(u))
(15)
In particular, for F potential energy, the conditions coincide with those of [89] in Rdd We refer to Appendix E. for more details.
5 Applications and Experiments
In this section, we first discuss how to verify the relative convexity and smoothness between func- tionals in practice. Then, we provide some examples of mirror descent and preconditioned gradient descent on different objectives. We refer to Appendix G for more details on the experiments'
Relative convexity of functionals. To assess relative convexity or smoothness as stated in Defini- tion 3, we need to compare the Bregman divergences along the right curves. When both functionals and F are of the same type, for example potential (respectively interaction) energies, this property is lifted from the convexity and smoothness on Rd of the underlying potential functions (respectively interaction kernels) to P2 (Rd), see Appendix E.2 for more details. When both ø and F are potential energies, the schemes (4) and (10) are equivalent to parallel MD and preconditioned GD since there interactions between the particles. The conditions of convergence then coincide with the ones obtained for MD and preconditioned GD on Rd [88, 89]. In other cases, (4) and (10) provide schemes that are novel to the best of our knowledge.
For functionals which are not of the same type, is less straightforward. Using equivalent notions of convexity (see Proposition 13), we may instead compare their Hessians along the right curves,
'The code is available at https: //gi http://ggtttuub.ttee com/ http:/cloonnttrreeeee and_aanddeeeee Preconditionee i
Descent in_ Wasserstein_Speet
8

---

W(P)
K2
K2
KL(Hilla*)
-208
-111
10-11
NEM (ours)
PFB (ours)
-224
-15
10
12
34
Time ( (T. iter)
Time ( T iter
Figure (Left) Value of W along the flow for two Figure 2: Convergence towards Gaussians difference interaction Bregman potentials, (Middle and Right) Trajectories particles to minimize W.
(0, U DUT) averaged over 20 covariances, with U ~ Unif(O10((R) and D fixed.
see Appendix E.2 for an example between an interaction and potential energy. We note also that for the particular case of functional obtained as a sum F = G + H. with 9r and H. convex, since dj. = dg + d dz, 2 max (dg. d2g and thus JF is 1 1-convex relative to 9 and H. This includes egg. the KL divergence which is convex relative to the potential and the negative entropy.
MD on interaction energies. We first focus on minimizing interaction energies W(n) [f W(x y) du(z))u(y)) with kernel W(z) || =11=1|._.. E Sj (R), whose minimizer is an ellipsoid [27]. Since the Hessian norm of W can be bounded by a polynomial of degree 2, following [88, Section 2]. W is ß-smooth relative Ka(z) 1|l=1l2 + Z ||=11 with = 4, and W 3-smooth relative to 0n (T) = ff Ka - T(y)) du(z))du(y). Supposing additionally that the distributions are compactly supported, we can show that W is smooth relative to the interaction energy with K2(z) ||2112. For coldiiditi E, i.e. for which the ratio between the largest and smallest eigenvalues is large, the convergence can be slow. Thus, we also propose to use (2) = "| and K KF = I|=l| + _|l=ll We illustrate these mirror descent schemes on Figure 1 and observe the convergence we expect for the ones taking into account Σ. In practice, since (T) = (VK= * (1) o T, the scheme (5) needs to be approximated using Newton's algorithm which can be computationally heavy. Using @% (T) = f V T dy with K2. we obtain more computationally friendly scheme with the same convergence, see Appendix G.2, but for which the smoothness is trickier to show.
MD on KL. We now focus on minimizing F(n)= S + H(u) for V(x) V(r)= - (()))=--~>'> with Σ possibly ill-conditioned, whose minimizer is the Gaussian v/ = - (O, M). and for which Wasserstein gradient descent is slow to converge. We study the MD scheme in (4) with negative entropy H the Bregman potential (NEM), and compare it on Figure 2 with the Forward-Backwa
(FB) scheme scheme (PFB) with Bregman
studied in [43] and the ideally preconditioned potential σy (see (116) in Appendix F). For computational purpose, we restrain the minimization in (4) over affine maps, which can be seen as taking the gradient over the submanifold of Gaussians [43, 73]. Starting from (0, Eo), the distributions stay Gaussian over the flow, and their closed-form reported in (62) (Appendix D.3). We note that this might not be the case for the scheme (4), and thus that this scheme does not enter into the framework developed in the previous sections. Nonetheless, it demonstrates the benefits of using different Bregman potentials. We generate 20 Gaussian targets / on IR 10 with Σ U DU', D diagonal and scaled in log space between and 100, and U a uniformly sampled orthogonal matrix, and we report the averaged KL over time. Surprisingly, NEM, which does not require an ideal (and not available in general) preconditioner, is almost as fast to converge as the ideal PFB, and much faster than the FB scheme.
Preconditioned GD for single-cells. Predicting the response of cells perturbation a central question in biology. In this context, the measuring process is destructive, feature descriptions of control and treated cells must be dealt with as (unpaired) source μ and target distributions v. Following [112), OT theory to recover a mapping T between these two populations has been used in [21, 22, 23, 39, 48, 69, 122]. Inspired by the recent success of iterative refinement in generative modeling, through diffusion [61, 115] or flow-based models [81, 85], our scheme (1) follows the idea of transporting l to v via successive and dynamic displacements instead of, directly, with a static map T. We model the transition from unperturbed to perturbed states through the (preconditioned)

---

F(u) SW (1,
F(u)
2 (4,
F(p) ED v)
F(A)
Miters convergence
F(
#iters convergence
F(f))
Hiters convergence
with with precond
13
101
12
10"
with mrccoda
18
10
10
without precond.
without precond.
without precond.
without precond.
without precond.
without precond.
Figure 3: Preconditioned GD vs. (vanilla) GD to predict the responses of cell populations to cancer treatment on 4i (Upper row) and scRNAseq (Lower row) datasets. For each treatment, starting from untreated cells μ4i, we minimize F(u)= D(j,, v1) with is the treated cells. The plot is organized pairs of columns, each corresponding to optimizing specific metric, with two scatter plots displaying points 24 = (I,, 84) where (First column) ys is the attained minima F(ū) = D(ū, v)) with preconditioning and z; that without preconditioning, and (Second column) ys the number of iterations to reach convergence with preconditioning and 2; that without preconditioning. A point below the diagonal column) a better minima or (Second column) faster convergence. We assign a color to each treatment and plot three runs, obtained with three different initializations, along with their mean (brighter point).
x then refers to experiment in which preconditioning provides (First
gradient flow of functional F(u) = D(,, v) initialized at Ai μu, where D is distributional metric, and predict the perturbed population via f = min,, F(p). We focus on the datasets used in [21], consisting of cell lines analyzed using (i) 4i [58], and (ii) scRNA sequencing [118]. For each profiling technology, the response to respectively (i) 34 and (ii) 9 treatments are provided. As in [21], training is performed in data space for the 4i data and in latent space learned by the seGen autoencoder [87] for the scRNA data. We use three metrics: the Sliced- Wasserstein distance SW2 [16], the Sinkhorn divergence $2 [50] and the energy distance ED [59, 60, 105], and compare the performances when minimizing this functional via preconditioned GD vs. (vanilla) GD. We measure the convergence speed when using fixed relative tolerance tol 10-3 well as the attained optimal value F()). Note that we follow [21] and additionally consider 40% of unseen (test) target cells for evaluation, i.e., for computing F()) D(2, v). As preconditioner, we use the one induced by h (z) functionals which grow ||x - x" [=(la/(a-1) near their minimum [119]. We the step size 1 for all the experiments. Then, we tune the parameter a very simply: for a given metric D and profiling technology, we pick random treatment and select and we generalize the selected a for all the other treatments. Results are described in Figure 3: Preconditioned GD significantly outperforms GD over the 43 datasets, in terms of convergence speed and optimal value F()). For instance, for D $2.,, we converge in 10 times less iterations while providing, on average, a better estimate the treated population. We also compare our iterative (non parametric) approach with the use of static (non parametric) map in Appendix G.4.
(||= + +1)
with a O, which is well suited to minimize
(1.25, 1.5, 1.75) by grid search,
6 Conclusion
In this work, we extended two non-Euclidean optimization methods on Ed to the Wasserstein space, generalizing Wy-gradient descent alternative geometries. We investigated the practical benefits of these schemes, and provided rates of convergences for pairs of objectives and Bregman potentials satisfying assumptions of relative smoothness and convexity along specific curves. While these assumptions can be easily checked is some cases (e.g. potential or interaction energies) by comparing the Bregman divergences or Hessian operators the Wasserstein geometry, they may be hard to verify in general. Different objectives such the Sliced-Wasseree Wasserstein distance the Sinkhorn divergence, or alternative geometries to the Wasserstein-2 as studied this work, require to derive specific computations on case-by-card basis. We leave this investigation for future work.
10

---

Acknowledgments and Disclosure of Funding
Clément Bonet acknowledges the support of the center Hi! PARIS and of ANR PEPR PDE-AL. Adam David gratefully acknowledges funding by the BMBF 01IS20053B project SALE. Pierre-Cyri Aubin-Frankowski was funded by the FWF project P 36344-N, Anna Korba acknowledges the support of ANR-22-CE23-0030.
References
[1] Martial Marie- Agueh. Existence of Solutions to Degenerate Parabolic Equations via the Mlonge-Kantorovie Theory. Georgia Institute Technology, 2002. (Cited on 1, 7)
[2] Byeongkeun Ahn, Chiyoon Kim, Youngjoon Hong, and Hyunwoo Kim. Invertible Monotone Operators for Normalizing Flows. Advances in Neural Information Processing Systems, 35: 16836-16848, 2022. (Cited on 28)
[3] Kwangjun Ahn and Sinho Chewi. Efficient Constrained Sampling via the Mirror- Langev-n Algorithm. Advances in Neural Information Processing Systems, 34:28405-288 284188, 2021.
(Cited on p. 2, 5, 21, 29, 44)
[4] David Alvarez-Melis and Nicolò Fusi. Dataset Dynamics via Gradient Flows in Probability Space. In International conference machine learning, pages 219-230. PMLR, 2021. (Cited
on p. 1)
[5] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows: in Metric Spaces and in the Space Probability Measures. Springer Science & Business Media, 2005. (Cited p. 1, 4, 23, 25, 38, 53)
[6] Abdul Fatir Ansari, Ming Liang Ang, and Harold Soh. Refining Deep Generative Models via Discriminator Gradient Flow. In 9th International Conference Learning Representations, ICLR, 2021. (Cited on p. 2, 20, 44)
[7] Michael Arbel, Anna Korba, Adil Salim, and Arthur Gretton. Maximum Mean Discrepancy Gradient Flow. Advances in Neural Information Processing Systems, 32, 2019. (Cited on p. 1,
2, 4, 24)
[8] Hedy Attouch, Giuseppe Buttazzo, and Gérard Michaille. Variational Analysis in Sobolev and
BV Spaces. Society for Industrial and Applied Mathematies, 2014. (Cited p. 33, 35)
[9] Pierre- Smoothness in Measure Spaces, with Application to Sinkhorn and EM. Advances in Neural Information Processing Systems, 35:17263 15:77333177755 2022. (Cited p. 2, 6, 21)
Frankowkii, Anna Korba, and Flavien Léger. Mirror Descent with Relative
[10] Heinz H Bauschke and Patrick L Combettes. Convex Analysis and Monotone Operator Theory in Hilbert Spaces. Springer, 2017. (Cited on p. 33)
[11] Heinz H Bauschke, Jérôme Bolte, and Mare Teboulle. A descent lemma beyond Lipschitz gradient continuity: first-order methods revisited and applications. Mathematics of Operations Research, 42(2):330-345, 2017. (Cited on p. 3, 6, 20, 37)
[12] Amir Beck and Marc Teboulle. Mirror Descent and Nonlinear Projected Subgradient Methods for Convex Optimization. Operations Research Letters, 31(3):167-175, 2003. (Cited on p. 2,
5, 7, 20)
[13] David M Blei, Alp Kucukelbir, and Jon D McAuliffe, Variational Inference: A Review for Statisticians. Journal of the American statistical Association, 11((518):859-877 2017. (Cited
on p. 1)
[14] Clément Bonet, Nicolas Courty, François Septier, and Lucas Drumetz. Efficient Gradient Flows in Sriced-Wasserster Space. Transactions on Machine Learning Research, 2022. (Cited
on p. 20)
[15] Clément Bonet, Lucas Drumetz, and Nicolas Courty. Sliced- iasser-eenn Distances and Flows Cartan-Hadamard Manifolds. arXiv preprint arXiv:2403.000 263600600 2024. (Cited on p. 2)
[16] Nicolas Bonneel, Julien Rabin, Gabriel Peyré, and Hanspeter Pfister. Sliced and Radon Wasserstein Barycenters of Measures. Journal of Mathematical Imaging and Vision, 51:22-45 2015. (Cited on p. 2, 10, 43)

---

[17] Benoît Bonnet. A Pontryagin Maximum Principle in Wasserstein Spaces for Constrained Optimal Control Problems. ESAIM: Control, Optimisation and Calculus of Variations, 25:52, 2019. (Cited on 23)
[18] Nicolas Bonnotte. Unidimensional and Evolution Methods for Optimal Transportation. PhD thesis, Université Paris Sud-Paris XI; Scuola normale superiore (Pise, Italie), 2013. (Cited on
p. 2, 43)
[19] Siwan Boufadène and François-Xavii Xavier Vialard. On the global convergence of Wasserstein gradient flow of the Coulomb discrepancy. arXiv preprint arXi::2312.00800/ 2023. (Cited on
p. 2)
[20] Yann Brenier. Polar Factorization and Monotone Rearrangement of
Functions. Communications on pure and applied mathematics, 44(4):375 1991. (Cited on p. 27)
[21] Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia del Castillo, Kjong-Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Ratsch. Learning Single-Cell Perturbation Responses using Neural Optimal Transport. bioRxiv, 2021. (Cited on 9, 10, 43)
[22] Charlotte Bunne, Andreas Krause, and Marco Cuturi. Supervised Training of Conditional Monge Maps. Advances in Neural Information Processing Systems, 0166859-67722 2022. (Cited on p. 9)
[23] Charlotte Bunne, Laetitia Papaxanthos, Andreas Krause, and Marco Cuturi. Proximal Optimal Transport Modeling of Population Dynamics. In International Conference on Artificial Intelligence and Statistics, pages 6511-6525 PMLR, 2022. (Cited on p. 1,9)
[24] Martin Burger, Matthias Erbar, Franca Hoffmann, Daniel Matthes, and André Schlichting. Covariance-modulatee optimal transport and gradient flows. arXiv preprint arXiv:2302.07773. 2023. (Cited on p. 21)
[25] Guillaume Carlier and Chloé Jimenez. On Monge's Problem for Bregman- Cost Functions. Journal of Convex Analysis, 14(3):647, 2007. (Cited on p. 6, 21)
[26] J. A. Carrillo, M. DiFrancesco, A. Figalli, T. Laurent, and D. Slepčev. Global-in-time weak measure solutions and finite-time aggregation for nonlocal interaction equations. Duke Muthematical Journal, 156()):229 - 271, 201 (Cited on p. I)
[27] José A Carrillo, Katy Craig, Li Wang, and Chaozhen Wei. Primal dual methods for Wasserstein gradient flows. Foundations of Computational Mathematics, pages 1-55, 2022. (Cited on p. 9, 40)
[28] Giulia Cavagnari, Giuseppe Savaré, and Giacomo Enrico Sodini. A Lagrangian approach to totally dissipative evolutions in Wasserstein spaces. arXiv preprint arXiv:2305.052/11 2023. (Cited on p. 5, 26)
[29] Gong Chen and Mare Teboulle. Convergence Analysis of Proximal-Like Minimization Algorithm Using Bregman Functions. SIAM Journal on Optimization, 3(3):538-533, 1993.
(Cited p. 6)
[30] Zonghao Chen, Aratrika Mustafi, Pierre Glaser, Anna Korba, Arthur Gretton, and Bharath K Sriperumbudur: regulaazgl Maximum Mean Discrepancy Gradient Flow. arXiv preprint arXiv:2409. 149800 2024. (Cited on 2)
[31] Ziheng Cheng, Shiyue Zhang, Longlin Yu, and Cheng Zhang. Particle-based Variational Inference with Generalized Wasserstein Gradient Flow. In Thirty-sevnnth Conference Neural Information Processing Systems, 2023. (Cited on p. 2, 2,2
[32] Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, Philippe Rigollet, and Austin Stromme. Exponential Ergodicity of Mirror-LLngevin Diffusions. Advances in Neural Infor-
mation Processing Systems, 33: 13:19933-19585, 2020. (Cited on p. 21, 44)
[33] Sinho Chewi, Tyler Maunu, Philippe Rigollet, and Austin J Stromme. Gradient descent algorithms for Bures- Wasserstein barycenters. In Conference on Learning Theory, pages 1276-1304. PMLR, 2020. (Cited on p. 53)
[34] Sinho Chewi, Jonathan Weed, and Philippe Rigollet. Statistical Optimal Transport. arXiv preprint arXiv:2407. 18163, 2024. (Cited on p. 23)
[35] Lenaic Chizat. Sparse Optimization on Measures with Over-parameterrzed Gradient Descent.
Mathematical Programming, 194(1):487-532, 2022. (Cited on p. 5)
12

---

[36] Lenaic Chizat and Francis Bach, On the Global Convergence of Gradient Descent for Over- parameterized Models using Optimal Transport. Advances in neural information processing systems, 31, 2018. (Cited on p. 1)
[37] Dario Gordero-Eraussqeee Transport incqualities for log-concave measures, quantitative forms, and applications. Canadian Journal of Mathematics, 69(3):481-501, 2017. (Cited on p. 21)
[38] Marco Cuturi, Lactitia Panngapapaa,,soss Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein. arXiv Preprint arXiv:2201. 12324, 2022. (Cited on p. 43)
[39] Marco Cuturi, Michal Klein, and Pierre Ablin. Monge, Bregman and Occam: Interpretable Optimal Transport High-Dimensions with Teaauee-SSprre Maps. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 6671-6682. PMLR, 23-29 Jul 2023. (Cited on p. 9)
[40] Mathicu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. How to compute Hessian-vector products? In ICLR Blogposts 2024, 2024. (Cited on p. 39)
[41] Nabarun Deb, Young-Heor Kim, Soumik Pal, and Geoffrey Schiebinger. Wasserstein Mirror Gradient Flow the Limit of the Sinkhorn Algorithm. arXiv preprint arXiv:2307.1642)
2023, (Cited on p. 5, 21, 29, 56)
[42] Sylvain Delattre and Nicolas Fournier. On the Kozachenko-Leonenkko entropy estimator. Journal of Statistical Planning and Inference, 185:69-93, 2017. (Cited on p. 44)
[43] Michael Ziyang Diao, Krishna Balasubramanian, Sinho Chewi, and Adil Salim. Forward- backward Gaussian variational inference via JKO in the Bures Wasserstein Space. In Interna- tional Conference on Muchine Learning, pages 7960-7991. PMLR, 2023. (Cited on p. 9, 29,
37, 41, 42, 53)
[44] Hanze Dong, Xi Wang, Lin Yong, and Tong Zhang. Particle-based Variational Inference with Preconditioned Functional Gradient Flow. In The Eleventh International Conference on Learning Representations, 2023. (Cited p. 2, 21)
[45] Chao Du, Tianbo Li, Tianyu Pang, Shuicheng Yan, and Min Lin. Nonparametric Generative Modeling with Conditional Siiced-Wasserstein Flows. In International Conference on Machine Learning (ICML), 2023. (Cited on p. 2)
[46] Andrew Duncan, Nikolas Nüsken, and Lukasz Szpruch. On the geometry of Stein variational gradient descent. Journal of Machine Learning Research, 24(56):1-39, 2023. (Cited on p. 21,
32)
[47] Matthias Erbar. The heat equation on manifolds as a gradient flow in the Wasserstein space. Annales de l'Institut Henri Poincaré, Probabilités et Statistiques, 46(1), February 2010. (Cited
on p. 4)
[48] Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian Theis. Unbalancedness Neural Monge Maps Improves Unpaired Domain Translation. In The Twelfth International Conference on Learning Representations, 2024.
(Cited on p. 9)
[49] Xingdong Feng, Yuan Gao, Jian Huang, Yuling Jiao, and Xu Liu. Relative Entropy Gradient Sampler for Unnormalized Distribution. Journal of Computational and Graphical Statistics, pages 1-16, 2024. (Cited on 20, 44)
[50] Jean Feydy, Thibault Séjourné, Trançois-Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between Optimal Transport and MMD using Sinkhorn Divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 2681-2690. PMLR, 2019. (Cited on p. 2, 10, 43)
[51] Nic Fishman, Leo Klamer, Valentin De Bortoli, Emile Mathieu, and Michael John Hutchinson. Diffusion Models for Constrained Domains. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. (Cited on p. 44)
[52] Nic Fishman, Leo Klamer, Emile Mathieu, Michael Hutchinson, and Valentin De Bortoli. Metropolis Sampling for Constrained Diffusion Models. Advances Neural Information Processing Systems, 36, 2024. (Cited p. 44)
13

---

[53] Jean- Yves Franceschi, Mike Gartrell, Ludovic Dos Santos, Thibaut Issenhuth, Emmanuel de Bézenac, Mickaël Chen, and Alain Rukotomamoniy. Unifying GANs and Score-Basr Based Diffusion as Generative Particle Models. Advances in Neural Information Processing Systems, 36, 2023. (Cited on p. 1)
[54] Bela A Frigyik, Santosh Srivastava, and Maya R Gupta. Functional Bregman divergence. In 2008 IEEE International Symposium on Information Theory, pages 1681-1685 IEEE, 2008. (Cited on p. 3,7)
[55] Wilfrid Gangbo and Adrian Tudorascu. On differentiability in the Wasserstein space and well- posedness for Comilton-Jacohfe equations. Journal de Mathématiques Pures et Appliquées, 125:119-1744 2019. (Cited on p. 24)
[56] Alfredo Garbuno-Inigo, Franca Hoffmann, Wuchen Li, and Andrew M Stuart. Interacting Langevin Diffusions: Gradient Structure and Ensemble Kalman Sampler. SIAM Journal on Applied Dynamical Systems, (9(1):122-44) 441, 2020. (Cited on 7, 21, 32)
[57] Xin Guo, Johnny Hong, and Nan Yang. Ambiguity set and learning via Bregman and Wasser- stein. arXiv preprint arXiv: 000.17055.05050 2017. (Cited on p. 21)
[58] Gabriele Gut, Markus Herrmann, and Lucas Pelkmans. Multiplexed protein maps link sub- cellular organization to cellular state. Science (New York, N.Y.). 361, 08 2018. (Cited on p.
10)
[59] Johannes Hertrich, Manuel Gräf, Robert Beinert, and Gabriele Steidl. Wasserstein steepest descent flows of discrepancies with Riesz kernels. Journal of Mathematical Analysis and Applications, 531(1):127829, March 2024. (Cited on p. 2, 10)
[60] Johannes Hertrich, Christian Wald, Fabian Altekrüger, and Paul Hagemann. Generative Sliced MMD Flows with Riesz Kernels. In The Twelfth International Conference on Learning Representations, 2024. (Cited on 2, 10, 43)
[61] Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. Ad- vances in neural information processing systems, 33:6840-6851 2020. (Cited on p. 9)
[62] Ya-Ping Hsieh, Ali Kavis, Paul Rolland, and Volkan Cevher. Mirrored Langevin Dynamics. Advances in Neural Information Processing Systems, 31, 2018. (Cited on p. 2)
[63] Xinru Hua, Truyen Nguyen, Tam Le, Jose Blanchet, and Viet Anh Nguyen. Dynamic Flows on Curved Space Generated by Labeled Data. In Edith Elkind, editor, Proceedings of the Thirty- Second International Joint Conference on Artificial Intelligence, IJCAI-23, pages 3803-3811. International Joint Conferences on Artificial Intelligence Organization, 8 2023. Main Track. (Cited on p.
[64] Qijia Jiang. Mirror Langevin Monte Carlo: the Case Under Isoperimetry. Advances in Neural Information Processing Systems, 34:715-725, 2021. (Cited on p. 2)
[65] Yiheng Jiang, Sinho Chewi, and Aram-Alexandre Pooladian. Algorithms for mean- mean-field variational inference via polyhedral optimization in the Wasserstein space. arXiv preprint #xXiv:2312.02849,
2023. (Cited on 37)
[66] Richard Jordan, David Kinderlehrer, and Felix Otto. The Variational Formulation of the Fokker-Planck Equation. SIAM journal on mathematical analysis, 29(1):1-17, 1998. (Cited
p. 1)
[67] Mohammad Reza Karimi, Ya-Ping Hsich, and Andreas Krause. Sinkhorn Flow: A Continuous- Time Framework for Understanding and Generalizing the Sinkhorn Algorithm. arXiv preprint arXiv:2311.16700 2023. (Cited on p. 2, 21)
[68] Jaeyeon Kim, Chanwoo Park, Asuman Ozdaglar, Jelena Diakonikolas, and Ernest K Ryu, Mirror Duality in Convex Optimization. arXiv preprint arXiv:2311.17296, 2023. (Cited on p.
2, 7, 20)
[69] Dominik Klein, Théo Uscidda, Fabian Theis, and Marco Cuturi. Entropic (Gromov) Wasser- stein Flow Matching with GENOT. arXiv preprint arXiv:2310.09254 2023. (Cited on p.
9)
[70] Michal Klein, Aram-Alexandr Pooladian, Pierre Ablin, Eugène Ndiaye, Jonathan Niles-Weed, and Marco Cuturi. Learning clastic costs to shape monge displacements, 2023, (Cited p. 7)
14

---

[71] Anna Korba, Adil Salim, Michael Arbel, Giulia Luise, and Arthur Gretton. A Non-Asympiotic Analysis for Stein Variational Gradient Descent. Advances in Neural Information Processing Systems, 33:4672-46633 2020. (Cited on 623
[72]
Anna Korba, Pierre-Cyril Aubin-Frankowski, Szymon Majewski, and Pierre Ablin. Kernel Stein Discrepancy Descent. In International Conference Machine Learning, pages 5719-
5730. PMLR, 2021. (Cited on p. 2, 24)
[73] Mare Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Varia- tional inference via Wasserstein gradient flows. Advances in Neural Information Processing Systems, 35:14434-14447, 2022. (Cited on P. 9, 42)
[74] Nicolas Lanzetti, Saverio Bolognani, and Florian Dörfler. First-Order Conditions for Opti- mization in the Wasserstein Space. arXiv preprint 122XX222209.11197, 2022. (Cited on p- 3, 4, 23,48)
[75] Emanuel Laude and Panagiotis Patrinos. Anisotropic Proximal Gradient. arXiv preprint arXiv:2210.15525 2022. (Cited on 7, 20)
[76] Emanuel Laude, Andreas Themelis, and Panagiotis Patrinos. Dualities for Non-Euclidean Smoothness and Strong Convexity under the Light of Generalized Conjugacy. SIAM Journal on Optimization, 33(4):2721-2749, 2023. (Cited on p. 7, 20)
[77] Flavien Léger and Fierre-Cyril Fran--Fra...sski. Gradient Descent with a General Cost. arXiv preprint arXiv:2305.04917 2023. (Cited on p. 2)
[78] Lingxiao Li, Qiang Liu, Anna Korba, Mikhail Yurochkin, and Justin Solomon. Sampling with Mollified Interaction Energy Descent. In The Eleventh International Conference on Learning Representations, 2023. (Cited on p. 2)
[79] Ruilin Li, Molei Tao, Santosh S Vempala, and Andre Wibisono. The Mirror Langevin Algorithm Converges with Vanishing Bias. In International Conference Algorithmic Learning Theory, pages 718-742. PMLR, 2022. (Cited on p. 2)
[80] Wuchen Li. Transport Information Bregman Divergences. Information Geometry, 4(2):
435 470, 2021. (Cited on p. 4, 21, 55, 56)
[81] Yaron Lipman, Ricky T. Q Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow Matching for Generative Modeling. In The Eleventh International Conference on Learning Representations, 2023. (Cited on p. 9)
[82] Guan-Horng Liu, Tianrong Chen, Evangelos Theodorou, and Molei Tao. Mirror Diffusson Models for Constrained and Watermarked Generation. Advances in Neural Information Processing Systems, 36, 2024. (Cited on p. 2, 44)
[83] Qiang Liu. Stein Variational Gradient Descent as Gradient Flow. Advances in neural informa- tion processing systems, 30, 2017. (Cited on p. 7, 20, 32)
[84] Qiang Liu and Dilin Wang. Stein Variational Gradient Descent: A General Purpose Bayesian Inference Algorithm. Advances in neural information processing systems, 29, 2016. (Cited on
p. 7, 20, 32)
[85] Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. In The Eleventh International Conference Learning Representations, 2023. (Cited on p. 9)
[86] Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Rober Stöter. Hliced-Wassersted Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions. In International Conference on Machine Learning, pages 4104-4113, PMLR, 2019. (Cited on 1,
[87] Mohammad Lotfollahi, F Alexander Wolf, and Fabian J Theis. seGen predicts single-cell perturbation responses. Nature methods, 16(8):715-721, 2019. (Cited on p. 10)
[88] Haihao Lu, Robert M Freund, and Yurii Nesterov. Relatively Smooth Convex Optimization by First-Order Methods, and Applications. SIAM Journal on Optimization, 28(1):333-33) 354, 2018. (Cited on p. 3, 5, 6, 8, 9, 20, 22, 35, 36, 40, 55)
[89] Chris J Maddison, Daniel Paulin, Yee Whye Teh, and Arnaud Doucet, Dual Space Precondi- tioning for Gradient Descent. SIAM Journal on Optimization, 31(1):911-1016, 2021. (Cited p. 2, 3, 5, 7, 20)
15

---

[90] Song Mei, Andrea Montanari, and Chan-Minh Nguyen. A Mean Field View of the Landscape of Two-Layer Neural Networks. Proceedings of the National Academy of Sciences, 115(33): E7665-E7671, 2018. (Cited on 1)
[91] Pierre Monmarché and Julien Reygner. Local convergence rates for Wasserstein gradient flows and McKean -eeean-Vllsov equations with multiple stationary solutions. arXiv preprint
1xXiv:2404415723) 2024. (Cited on 5)
[92] Arkadij Semenovič Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization. Wiley, 1983. (Cited on p. 2, 20)
[93] Sebastian Neumayer, Viktor Stein, and Gabriele Steidl. Wasserstein Gradient Flows for Moreau Envelopes of f-Divergences in Reproducing Kernel Hilbert Spaces. arXiv preprint
arXiv:2402.0466 2024. (Cited on p. 2)
[94] Maxence Noble, Valentin De Bortoli, and Alain Durmus. Unbiased constrained sampling with Self-Concordant Barrier Hamiltonian Monte Carlo. In Thirty--eventh Conference on Neural Information Processing Systems, 2023. (Cited on p. 44)
[95] Felix Otto. Double degenerate diffusion equations as steepest descent. Citeseer, 1996. (Cited
on p. 1)
[96] Felix Otto. The Geometry of Dissipative Evolution Equations: the Porous Medium Equation. Communications in Partial Differential Equations, 1611-2)101-177, 2001. (Cited on p. 2)
[97] Felix Otto and Cédric Villani. Generalization of an inequality by Talagrand and links with the logarithmic Sobolev inequality. Journal of Functional Analysis, 173(2)3361-40MM 2000. (Cited on p. 24)
[98] Guy Parker. Some Convexity Criteria for Differentiable Functions on the Wassersstni Space. arXiv preprint arXiv:2306.0911 09120, 2023. (Cited on p. 5, 26)
[99] Juan Peypouquet. Convex Optimization in Normed Spaces: Theory, Methods and Examples.
Springer, 2015. (Cited on p. 22, 33, 54, 55)
[100] Gabriel Peyré. Entropic approximation of Wasserstein gradient flows. SIAM Journal on Imaging Sciences, 8(4):2323 2351, 2015. (Cited on p. 20)
[101]
Aram-Alexandre Pooladian and Jonathan Niles-Weat Entropic estimation of optimal transport maps. arXiv preprint arXiv: 120004409.,2004 2021. (Cited on 42, 43)
[102] Julien Rabin, Gabriel Peyré, Julie Delon, and Mare Bernot. Wasserstein Barycenter and its Application to Texture Mixing. Scale Space and Variational Methods in Computer Vision: Third International Conference, SSVM 2011, Ein-Gedi, Israel, May 29-June 2, 2011, Revised Selected Papers 3, pages 435-446. Springer, 2012. (Cited on p. 23)
Wasserstein [103] Cale Rankin and Ting-Kau Kam Leonard Wong. Bregman-Wassertenn Divergence: Geometry and
Applications. arXiv preprint arXiv:2302. 0583333,.033,.... 2023. (Cited on p. 6, 21, 29)
[104] Cale Rankin and Ting-Kam Leonard Wong. JKO schemes with general transport costs. arXiv
preprint arXiv:2402. 17681, 2024. (Cited on p. 5, 20, 21, 29)
[105] Maria L. Rizzo and Gábor J. Székely. Energy Distance. WIREs Computational Statistics, 8(1):
27-38, 2016. (Cited on p. 10)
[106] Gareth Roberts and Richard Tweedie. Exponential convergence of Langevin distributions
and their discrete approximations, Bernoulli, pages 341-363, 1996. (Cited on p. 1)
[107] Grant M Rotskoff and Eric Vanden-Eijnden. Neural networks as interacting particle systems: Asymptotic convexity of the loss landscape and universal scaling of the approximation error. stat, 1050:22, 2018. (Cited on p. 1)
[108] Adil Salim and Peter Richtárik. Primal dual interpretation of the proximal stochastic gradient Langevin algorithm. Advances in Neural Information Processing Systems, 33:3786-3796, 2020. (Cited on p. 4)
[109] Adil Salim, Anna Korba, and Giulia Luise. The Wasserstein Proximal Gradient Algorithm. Advances in Neural Information Processing Systems, 13356666-1 2020. (Cited on p. 37,
38, 42)
[110] Filippo Santambrogio. Optimal Transport for Applied Mathematicians, volume 55. Springer, 2015. (Cited on p. 4, 7, 23)
16

---

[111] Filippo Santambrogio. Euclidean, metric, and Wasserstein gradient flows: an overview. Bulletin of Mathematical Sciences, 7:87-154, 2017. (Cited on p. 1)
[112] Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-TTanspoon Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell, 176(4), 2019. (Cited on
[113] Louis Sharrock, Lester Mackey, and Christopher Nemeth. Learning Rate Free Bayesian Inference in Constrained Domains. Advances in Neural Information Processing Systems, 36, 2024. (Cited on p. 2, 7, 32, 44)
[114] Jiaxin Shi, Chang Liu, and Lester Mackey. Sampling with Mirrored Stein Operators. In International Conference on Learning Representations, 2022. (Cited on 2, 7, 32, 44)
[115] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score- Generative Modeling through Stochastic Differential Equations. In International Conference on Learning Representattonss 2021. (Cited on 9)
[116] Vishwak Srinivasan, Andre Wibisono, and Ashia Wilson. Fast sampling from constrained spaces using the Metropolis-adjuster Mirror Langevin algorithm. In Proceedings of Thirty Seventh Conference on Learning Theory, volume 247, pages 4593-4635, 2024. (Cited on p. 2, 44)
[117] Ken ichiro Tanaka. Accelerated gradient descent method for functionals of probability measures by new convexity and smoothness based on transport maps. arXiv preprint arXiv:2305.05122 2023. (Cited on 5, 38)
[118] Fuchou Tang, Catalin Barbacioru, Yangzhou Wang, Ellen Nordman, Clarence Lee, Nanlan Xu, Xiaohui Wang, John Bodeau, Brian B Tuch, Asim Siddiqui, al. mRNA-Seq whole- transcriptome analysis of a single cell. Nature methods, 6(5):377-382, 2009. (Cited on
10)
[119] Salma Tarmoun, Stewart Slocum, Benjamin David Haeffele, and Rene Vidal. Gradient Preconditioning for Non-Lipschitz smooth Nonconvex Optimization. 2022. (Cited on p. 7, 10,
20)
[120] Antonio Terpin, Nicolas Lanzetti, and Florian Dörfler. Learning Diffusion at Lightspeed.
Advances in Neural Information Processing Systems, 2024. (Cited on p. 1)
[121] Belinda Tzen, Anant Raj, Maxim Raginsky, and Francis Bach. Variational Principles for Mirror Descent and Mirror Langevin Dynamics. IEEE Control Systems Letters, 7:1542-1547,
2023. (Cited on p. 2)
[122] Théo Uscidda and Marco Cuturi. The Monge Gap: A Regularizer Learn All Transport Maps. In International Conference on Machine Learning, pages 34709-34733. PMLR, 2023. (Cited
p.
[123] Quang Van Nguyen. forward-backwart splitting with Bregman distances. Vietnam Journal of Muthematics, 45(3):519-53,, 2017. (Cited on p. 6)
[124] Cédric Villani. Topics in Optimal Transporaatio,, volume 58. American Mathematical Soc..
2003. (Cited on 24)
[125] Cédric Villani. Optimal Transport: Old and New, volume 338. Springer, 2009. (Cited on
51)
[126] Yifei Wang and Wuchen Li. Information Newton's Flow: Second-Order Optimization Method in Probability Space. arXiv preprint arXiv:2001.0433: 2020. (Cited on p. 5, 21, 24, 25, 29)
[127] Yifei Wang, Peng Chen, and Wuchen Li. Projected Wasserstein gradient descent for high- dimensional Bayesian inference. SIAM/ASA Journal on Uncertainty Quantification, 10(4):
1513-1532, 2022. (Cited on p. 20,44)
[128] Yifei Wang, Peng Chen, Mert Pilanci, and Wuchen Li. Optimal Neural Network Approximation Wasserstein Gradient Direction via Convex Optimization. arXiv preprint arXiv:220s. 2022. (Cited on p. 20, 44)
[129] Andre Wibisono. Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem. In Conference on Learning Theory, pages 2093-3027.
PMLR, 2018. (Cited on p. 1. 1.4. 31, 42)
17

---

[130] Andre Wibisono. Proximal Langevin Algorithm: Rapid Convergence under Isoperimetry. arXiv preprint arXiv:1911.01 2019. (Cited on p. 5, 29)
[131] Lantian Xu, Anna Korba, and Dejan Slepcev. Accurate Quantization of Measures via Interact- ing Particle-based Optimization. In International Conference on Machine Learning, pages 24576-24595. PMLR, 2022. (Cited p. 7)
[132] Kelvin Shuangjann Zhang, Gabriel Peyré, Jalal Fadili, and Marcelo Pereyra. Wasserstein Control of Mirror Langevin Monte Carlo. In Conference on Learning Theory, pages 3814- 3841. PMLR, 2020. (Cited p. 2)
[133] Zico Zolter, David Duvenaud, and Matt Johnson. Deep Implicit Layers- - Neural ODEs, Deep Equilibirum Models, and Beyond. Neurips 2020 Tutorial, 2020. (Cited on p. 39)
18

---

Appendix
The appendix is organized follows. In Appendix A, discuss related works. In Appendix B and Appendix C. we provide mathematical background respectively on L and on the Wasserstein space. In Appendix we provide complementary results for the mirror descent scheme the Wasserstein space. In Appendix E, we discuss the relative convexity and smoothness between functionals. Appendix F, we study the Bregman proximal gradient scheme, deriving convergence result and closed-form updates for the Gaussian case. In Appendix G, we provide more details on the experiments. In Appendix H, we report the proofs of our results. Finally, auxiliary results are given in Appendix I.
Contents
Related works
20
B Background on (p)
21
B.1 Differential calculus on L (p)
21
B.2 Convexity on (²)
22
Background on Wasserstein space
23
C.1 Wasserstein differentials
23
C.2 Wasserstein Hessians
24
C,3 Convexity in Wasserstein space
25
D Additional results mirror descent
27
D.I Optimal transport maps for mirror descent
27
D.2 Continuous formulation
29
D.3 Derivation in specific settings
29
D.4
Mirror scheme with non-pushforward compatible Bregman potentials
32
E Relative convexity and smoothness
33
E. Relative convexity and smoothness between Fenchel transforms
E.2 Relative convexity and smoothness between functionals
35
Bregman proximal gradient scheme
37
Additional details on experiments
39
G.I Implementing the schemes
39
G.2 Mirror descent of interaction energies
40
G.3 Mirror descent on Gaussians
G.4 Single-cell experiments
42
G.5 Mirror descent on the simplex
44
H Proofs
44
H.I Proof of Proposition
44
19

---

H.2 Proof of Proposition 2.
45
H.3 Proof of Proposition 3
45
H.4 Proof of Proposition 4
46
H.5 Proof of Proposition 5
47
H.6 Proof of Proposition 6.
48
H.7 Proof of Proposition
48
H.8 Proof of Lemma 11
49
H.9 Proof of Proposition 12
50
H.10 Proof of Proposition
50
H.11 Proof of Proposition 24
51
H.12 Proof of Proposition 25
52
H.13 Proof of Lemma
H.14 Proof of Proposition 27
53
I Additional results
54
L11 Three-point identity and inequality
54
1.2 Convergence lemma
55
1.3 Some properties of Bregman divergences
J NeurIPS Paper Checklist
58
A Related works
Mirror descent on RR. Mirror descent has been introduced by Nemirovskij and Yudin [92] solve convex optimization problems. Its convergence has been first studied under the assumption that the objective has a Lipschitz gradient, see e.g. [12]. More recently, Bauschke al. [11], Lu et [88] provided convergence guarantees by assuming relative smoothness and convexity.
Preconditioned gradient descent IR The preconditioned gradient descent has first been studied by Maddison et al. [89], providing convergence guarantees under assumptions on the smoothness and convexity of the preconditioner relative to the Legendre transform of the objective. Kim al. [68] underlined connections with the mirror descent, and introduced an accelerated version of the preconditioned gradient descent. Laude and Patrinos [75], Laude al. [76] studied generalized version of this algorithm by minimizing an anisotropic upper bound and supposing anisotropic smoothness of the objective. particular, their analysis for the descent lemma also valid for non-convex smooth objective. Tarmoun et al. [119] also studied preconditioned gradient descent for non-Lipschitz smooth non-convex problems.
Wasserstein Gradient flows with respect to non-Euclidean geometries. Several existing schemes based time-discretizations of gradient flows with respect optimal transport metrics, but different than the Wasserstein-2 distance.
To simplify the computation of the backward scheme, Peyré [100] added entropic regularization into the JKO scheme while Bonet et al. [14] considered using the Sliced-Wasseree Wasserstein distance instead. More recently, Rankin and Wong [104] suggested using Bregman divergences e.g. when geodesic distances are not known closed-forms.
The most popular objective in Wasserstein gradient flows is the KL. However, this can be intricate to compute as it requires the evaluation of the density each step, which not known for particles, and thus requires approximations using kernel density estimators [127] or density ratio estimators [6, 49, 128]. Restricting the velocity field to reproducing kernel Hilbert space (RKHS), an update in closed-form can be obtained, which is given by the SVGD algorithm [83, 84]. This algorithm
20

---

can also be seen using an alternative Wasserstein metric [46]. However, the restriction to RKHS can hinder the flexibility of the method. This motivated the introduction of new schemes based on using the Wasserstein distance with convex translation invariant cost [31, 44]. Particle systems preconditioned by they empirical covariance matrix have also been recently considered, and can be seen as discretization of the Galman-Wasseree or Covariance Modulated gradient flow [24, 56].
Mirror descent with flat geometry. The space of probability distributions can be endowed with different metries. When endowed with the Fisher-Rat Rao metric instead of the Wasserstein distance, the geometry becomes very different. Notably, the shortest path between the two distributions is now a mixture between them. In this situation, the gradient is the first variation. FrabinFFrankowkki et al. [9] studied the mirror descent this space and notably showed connections with the Sinkhorn algorithm when the mirror map and the optimized functionals are KL divergences. Karimi et al. [67] extended the mirror descent algorithm for more general time steps, and notably recovered the "Wasserstein Mirror Flow" proposed by Deb al. [41] as a special case.
Bregman divergence on (2(R)). Several works introduced Bregman divergences on P2 (2R"). Carlier and Jimenez [25] first studied the existence of Monge maps for the OT problem with Bregman costs c(x, y) = dy (zz y) and symmetrized Bregman costs c(I, y) = dy (a, y) + (y. æ). For Bregman costs, the resulting OT problem was named the Bregman- Wasserstein divergence and its properties were studied in [37, 57, 103]. The Bregman- Fregman--asseestein divergence has also been used by Ahn and Chewi [3] to show the convergence of the Mirror Langevin algorithm while Rankin and Wong [104] studied its JKO scheme with KL. objective. Li [80] introduced the notion of Bregman divergence on Wasserstein space for geodesically strictly convex F P2 (2()) + R as
AY, V E PRR),, dy (H, v) = F()) - F(D)) - (V F(u), T Id) 12())
(16)
where Ty is the OT map between v and 1a W.E./ W2. The Bregman divergence used in our work and as defined in Definition 1 more general special case (16). Li [80] studied properties of this Bregman divergence for different functionals F and provided closed-form for one-dimensional distributions or Gaussian, but did not use define a mirror scheme.
allows using more general maps and contains as
Mirror descent on (TRT) Deb et al. [41] defined a mirror flow by using the continuous for- mulation. They focused on KI. objectives with Bregman potential @(μ) (1, v) with some reference measure PR and defined the flow the solution of
(((m)) = Vw,d(u)) d(a)) a ((p)) = - Vw,F(p)))
(17)
We note that @ is pushforward compatible and hence enters our framework. Also related to our work, Wang and Li [126] studied a Wasserstein Newton's flow, which, analogously to the relation between Newton's 's method and mirror descent [32], is another discretization of our scheme for Φ F. We clarify the link with the Mirror Descent algorithm we define in this work with the previous continuous formulation above in Appendix D.2.
B Background on L²(
B.1 Differential calculus (p)(p)
We recall some differentiability definitions the Hilbert space (²) for μ E P2(R'). Let 6: (2() + R. We start by recalling the notions of Gâteaux and Fréchet derivatives.
Definition 4. A function Φ= (p) + R is said to be Gâteaux differentiable at T if there exists an operator (TT (2) + R. such that for any direction h E L²²
Φ(T th) -(TT) t-O
(T())(h)
lim
(18)
and d' (T) is a linear fanction. The operator &' (T) is called the Gâteaux derivative of d at T and if it exists, it is unique.
21

---

Definition 5. The Fréchet derivative of o at T E (p) in the direction h E L2()), denoted 66 h). defined implicitly by
( + th) = ((T) + (T h) + to([))]))-
(19)
Φ is Fréchet differentiable, then it is also Gâteaux differentiable, and both derivatives agree, i.e. for all T, h E (w), S6 h) o' (h) [99, Proposition 1.26].
Moreover, since L (p) is a Hilbert space, and (T and (T) are linear and continuous, if is
Fréchet (resp. Gâteaux) differentiable, by the Riesz representation theorem, there exists Ve E (²² such that for all E (μ), 66 h) = (Vø(T), h) (1)) (resp. (T)) = (V)(T), LP()))
As a brief comment on these notions the context of convexity, if the subdifferential of a convex f contains a single element then it is the Gâteaux derivative and we have an inequality f(v) Z f(z) + (Vf(r), y - Z). Instead Fréchet différentiability gives an equality (19) corresponding to series expansion.
B.2
Convexity on (1)
Let Φ : (p) + E be Gâteaux differentiable. We recall that ø is convex if for all E [0, 1].
T,S E (1)),,
-(11t)TT iS) I (1 - t)o(T) + to()), which is equivalent by [99, Proposition 3.10] with
(20)
VT,S E (p), <(T) 2 <(S) + (Vo(S), T S)) = do 2 0.
(21)
We now present equivalent definitions of the relative smoothness and relative convexity, which is the equivalent of [88, Proposition 1.1].
Proposition 8. Let 1,, Φ (²) + R be convex and Gâteaux differentiable functions. The following
conditions are equivalent:
(al) vb is B-smooth relative to 6
(a2) 3Φ - db is convex function on [1)
(a3) If twice Gâteaux differentiable, (V) & (VV))((). for all T,S e
(2(())
(a4) (Vv(T) - VV()), S
T L all T,S (²²()).
The following conditions are equivalent:
(b1) ab is a-convex relative to
(b2) 1b- aá is a convex function on (²(
(b3) Ifwie differentiable, (T)S, ( for all T,S E 1,3(4)
(b4) (Vv(T) - T- 2 (V Vo(S). T r)) for all T,S E (2))))
Proof. We do it only for the smoothness. It holds likewise for the convexity.
(al) (a2): VT,S E (p),,, (T,S) < (T,S
VT,S (m),,, ((T) v(() -
(22)
$ ( - d(S) (Vo(S), ( S)
VT,S E (4), (80-5) (V(36 ()), T- S) I (BΦ (T).
For the rest of the equivalences, we apply 199, Proposition 3.10]. Indeed, B6 sb convex is equivalent to
VT,S E (p), (V(36 - v]( V)S) S) 2 0 VT,S (2),), B(Vo(T) T 2 (Vv(T) -V SS) (23)
22

---

which gives the equivalence between (a2) and (a4). And if db and e are twice differentiables, it is also equivalent
vT,S E (2),
() 20 - vT,S L2 (VV)) 2 ()) (24)
which gives the equivalence between (a2) and (a3).
C Background on Wasserstein space
C.1 Wasserstein differentials
We recall the notion of Wasserstein differentiability introduced in [5, 17, 74). First, we introduce sub- and super-differentials.
Definition 6 (Wasserstein sub- and super-differential [17, 74)). Let F P2(R2) - (-00, + lower semicontimouss and denote D(F) = {u E P2(R)), F(u) < ooj. Let μ E D(F). Then, map & E I.2 (n)) belongs to the subdifferential å F()) of F at if all v P2(R2))
/ (5(x). x) dr v) + o (W2(u, vEllo(p.u)
F(v) 2 F(N)+ +
sup
(25)
Similarly, & E (2) belongs to the superdifferential Ö+ F(u) of at μ if - O
Then, we say that a functional is Wasserstein differentiable if admits sub and super differentials which coincide.
Definition 7 (Wasserstein differentiability, Definition 2.3 in [74]). A functional F P2(RR) + I
Wasserstein differentiable at μ E P2((2) i F(P) n o F(u) & 0. In this case, we say that V F(p) e F(x)) a Wasserstein gradient of F p. satisfying for any v E (2.
E Vo(,,v).
F(v) = F(u)+ + / (WW,)( 2) dr(x, + (W2(H,,
(26)
Recall that the tangent space (2( at E P2( is defined as
TT (V4, 1 E C (RRd) C (p))
(27)
where the closure is taken in (1)),, see Ambrosio al. [5, Definition 8.4.1]. Lanzetti al. [74, Proposition 2.5] showed that if F is Wasserstein differentiable, then there is always unique gradient living in the tangent space, and we can restrict ourselves without loss of generality to this gradient.
Lanzetti et al. [74] further showed that Wasserstein gradients provide linear approximations even if the perturbations are not induced by OT plans, i.e. differentials are "strong Fréchet differentials". Proposition 9 (Proposition 2.6 in (74]), Let v.vE E P2((R)) Y E 11(,, v) any coupling and let F; (2 + R be Wasserstein differentiable at μ with Wasserstein gradient V F(u) E T (2RR). Then,
/ ||= 1il1-wll6ee 01 (
F(v)= + (WW,) x) dr((x,))
(28)
Under regularity assumptions, the Wasserstein gradient of F can be computed in practice using the first variation AF [110, Definition 7.12), which is defined, if it exists, as the unique function (up to Jy constant) such that, for x satisfying fdxx 0,
d F(u +((+tx) r-
F(p + tx) FF()) lim t-
/ [ (p) dx.
(29)
dt
t
8μ
Then the Wasserstein gradient can be computed as Fw,F(1) = oμ (n)), see egg. [34, Proposition 5.1
65
We now show that we can relate the Fréchet derivative of FT) F(TyN) with the Wasserstein gradient of JF belonging to the tangent space of Pa (IRd) at p.
23

---

Proposition 10. Let F: P2( - [+ be Wasserstein differentiable functional on D(F). Let μ E (2(( and FT) =(T for all E D(Fx). Then, Fy is Fréchet differentiable, and for all S E D(Fw). (S) = Vw, o.S
Proof. See Appendix H.1.
A similar formula can be found in Gangbo and Tudorascu [55, Corollary 3.22], however the space H used there is not (1) but lifting (12; IR4) of measures on random variables. They should not be confused.
C.2 Wasserstein Hessians
natural object of interest in the context of optimization over the Wasserstein space the Hessian of the objective F. which we define below, according to the original definitions of [97, Section 3] and [124, Chapter 8j. This notion is usually defined along Wasserstein geodesics.
Definition 8 (Chapter 8 in [124]. The Wasserstein Hessian of F. denoted HJF, is operator over Tu ( verifying (H,, vo) (2(1) = F(po) if (pr, Ve)te(o.11 is a Wasserstein geodesic starting μ.
If fa is absolutely continuous, Wassertein geodesics starting from μ are curves of the form Pt (Td + + for vú (R)).. Using this fact, one can compute the Wasserstein Hessians of Kullback-Leibe divergence [97], Maximum Mean Discrepancy [7] or Kernel Stein Discrepancy [72] and many other functionals.
However in this work, we are interested in more general curves, which we call acceleration free, i.e.
He = (S + tv) # la with S, v E (p). Thus, we define analogously the Hessian along such curves. Definition 9. We define the Hessian operator H.F (²)))) - (u) as the operator satisfying for all / E 1].
de F(1) (.) dt
(30)
where pe (S (v) y for S, E (p)
Note that the Hessian H.F,,.. is taken at time / but that the vector field (2(()) is in the tangent space at = 0, hence the discrepancy with Definition besides the fact that we can have S * ld.
Wang and Li [126] derived a general closed form of the Wasserstein Hessian on tangent spaces through the first variation of F. Here, we extend their formula along any curve μ4 (S + tu #M with S, U E I (1)) We first provide lemma computing the derivative of the Wasserstein gradient.
Lemma 11. Let F P2(R2) + R be twice continuously differentiable and assume that V 6F on op (n) V 6xz (P) for all μ E P2(RR)). Let μ P2((R)) and for all [0, Mr = (Te) where T, is differentiable w.r.t. L with dT. E (1))).. Then,
d di
(Vw,F(m) T.)(x) = / V V. FF pn ((T)) (T(()).TT())
dT, (v)) du(ب)
dt
+ 22 22
ōu ((T))H)
d_' (z). (31)
di
Proof. See Appendix H.8.
This allows us to define closed-form for HJF p.t-
Proposition 12. Under the same assumptions in Lemma 11, let pt = (T)) with T, S + tw,
S. E L2 (p), then HJF,t L²(()) + L²())) (as defined in Definition 9) is defined for all v E L (H), x E R2d as:
8FF ((T.) (T,()), Tx())) oly) du(s)+ 5y
oF (TT)))))(T((())ell (T+(x))u(()) (32)
HJ ((Fr)((((()>= = V, 542
24

---

Proof. See Appendix H.9.
While HJF n.t and H.F,p differ general, in some simple cases their relation boils down to composition with pushforward. For instance, if S is invertible, we can write μr as curve starting from S $ fl with a velocity field v S-', i.e. l4 = (Id + tv o S $ #H. Thus, we recover the original definition of the Wasserstein Hessian at the case.
HF HF However, in general, this does not need be
Similarly, if T, is invertible for all 1, setting 2 T write
d? dd2 F(u)
| (HLF)))o)((),,))..)))))))))))))) v(x)) du(x))
(33)
= (H.F mulFi (T; (Ix) us (x1)) dpp()r))
(H.F Ut, ve) 12(()))
The last line is obtained by leveraging the closed form in Proposition 12 and that μa (TT) ue as for all x E supp(u),
/
8F
(n) (Tr(T)) T.())) ((y) du(y) + V SF (me (T.(()())))
HJ H.
V y
8n2
ōu
8F
(T) (1) d + V DF (T(())()))
(34)
v V
6p2
óu
HJ m. [v]] (Tr(
and thus HJ, [o] (T, (z)) H.F H.F [v))()).
Here are two examples of F satisfying SF expression of the Wasserstein Hessian.
for which Proposition 12 provides an
of
Example (Potential energy). Let V(H) = f V du with twice differentiable with bounded Hessian.
Then, we have op (n) V and 7 (n) (using (29)). Thus, applying Proposition 12, we recover for μr = (Tt)#14
dr V(HN) - / dp(z).
d²
(35)
Example 2 (Interaction energy). Let ((x) = ff W(x- - dn(z)dply) with W symmetric, twice differentiable and with bounded Hessian. Then, we have for all E Rd, sW (p)(z) (W * (z) by
Fn and 5 (j)((,,)) w) = W(x y) (see e.g. [126, Example 7J), and thus applying Proposition for
He (1,) the operator
(2W(()(x)) - HWuus(u)(s)) -
(((()))u()) du(w) (V2W (T)))))((((()))(()) (36)
and
d²
W(p)) // (V -T(()) (v(x) - v(v)), v(x)) dn(r)d((x). (37)
dta
C.3 Convexity in Wasserstein space
We first recall the definition of a-convex functionals [5, Definition 9.1.1]. Definition 10. F is a-convex along geodesics for all po, f1 (22(
V [0, F(m) < (1 (1 (1-t)F(pp) + tF(u) + t(1 t)
w (Ho. p1),
(38)
2
where (Ho)eejo.1) is Wasserstein geodesic between pa and p1.
25

---

If we want to derive the minimal set of assumptions for the convergence of the gradient descent algorithms on Wasserstein space, we can actually restrict the smoothness and convexity to specific curves. In the next proposition, we characterize the convexity along one curve. The relative smoothness or convexity follows by considering the convexity of respectively 3G or F ag. Proposition 13. Let (2 + R be twice continuously differentiable. Let μ. (2(R)) T,S E (2),), Pt = (Tr) # for all t [0, where T, (1 +JS + t'. Furthermore, denote for t, E [0, 11. i - ((1 -),T + +TTea) # M, Then, the following statement are equivalent:
(cl) For all I 12, E [0, 1). F(1 22 1-
i.e. F is convex +
along t + 14-
(c2) For all 1,,2 12 E [0, 1}. we have dg, (Tep) Tr,) 2 0, i.e.
=((Te)+1)) - F((TA)#M) - (T..F((T))#t)) Thi Tra - -Ta)) 2 0.
(c3) For all t1, t2 E [0.1].
((T) o T2, -Vw,F(TT)))#)) Te,, Tıp Ten)ra 2 0.
(c4) For all 8 E [0,
d² ar F(me) 20.
Proof. See Appendix H.10.
As stated in Section 2, if we require the convexity to hold along all curves with S = Id and T the gradient of some convex function, i.e. OT map, then F is convex along geodesics. Likewise, if the convexity holds for all S, T that are gradients of convex functions, then we obtain the convexity along generalized geodesics.
we require the convexity and the smoothness to hold along any curve of the form μr = ((1
tis + # for μa E (P2(R)) and T,S E Tp), then it coincides with the transport convexity and smoothness recently introduced by Tanaka [117, Definitions 4.1 and 4.5]. As by Proposition 1,
oF,, T S) (V F(S S, ()())) and thus the convexity of J on such curve reads as follows
(
= F(T#R) F() (Vw, oS, - - S) 2 O.
(39)
And for (T( [| the 8-smoothness of F relative to G expresses
8 -1T-SI
F( (o T-S) S
= -dgg (T,S).
(40)
This type of convexity is actually particular case of the notion of convexity along acceleration-fres
curves introduced by Parker [98] (also introduced by Cavagnari et al. [28] under the name of total con- vexity). The latter requires convexity to hold along any curve of the form f4 = ((1- -)) + tt²) #7
with y I[(, v), μ, v (R and x² (z,y) y) x, x² (2, v) = y. The transport convexity of Tanaka [117] is thus particular case for couplings obtained through maps, i.e. = (T. +) Parker [98] notably showed that this notion of convexity is equivalent the geodesic convexity for Wasserstein differentiable functionals.
We can also define the strict convexity using strict inequalities in Proposition (c)()))) but not in (c4) as there are counter examples already for real functions. For instance, f : IR + R, defined as fx) 4 for all æ E R, is strictly convex but f"(() O. Thus, for F(u)
dj, d² F(M)) = (0)))((0)² = 0. But
choosing μ = 6n and To 1d, by Example 1, we have that des F(m)) = ( < 4s (T) since is strictly convex and thus F is well strictly convex.
Finally, as we defined the relative a-convexity and 8-smoothness of J relative to & using Bregman divergences in Definition 3, we can show that it is equivalent to F - ag and BG F being convex.
26

---

Proposition 14. Let F.G. P2(R)) + R be two differentiable functionals Let μ P2(R2),
T,S E (1) and for all [0, 11. br (T,) a with T, (1- + rT. Then, F is a-convex (resp. B-smooth) relative to G along i - μr if and only if - -og (resp. 89 - F) is convex along t t pr-
Proof. By Definition 3, F is a-convex relative to 9 along H Me if for all s,t E [0. 1], dz (T,; Te) 2 adg, (T.: T.). This is equivalent to dg - (T,; ; TT) 2 0, which is equivalent by Proposition 13 (c2) (and using Proposition 1) to F - ag convex along t + pe. The result for the B-smoothness follows likewise.
L
D Additional results mirror descent
D.1 Optimal transport maps for mirror deseent
Let db (22(IR) R be strictly convex functional along all acceleration-free curves 14, (Tz)#1,
E [0, 1] with T. = (1 t)S IT for T,S (m)) Denote for μ. E (²),,, Φu (T) = @T
Since @ is strictly convex along all acceleration-fiee curves, by Proposition 13, for all T + E (w), de (T, S) and thus oy is strictly convex. Indeed, recall that
VT, E (M), do, (T, 0n (T) (.( ((o(s))-VVo_ll) (S), (
(41)
(T # - (( (V (V, o. LPW)
where we used Proposition for the computation of the gradient. Let us now define for all ne v E (2())
W (K) inf d(v) x(w) - (Vw:o())))) (2000ow-/VWWeoWWW0) 7Ell(sp))
y)) dy(x,
(42)
This problem encompasses several previously considered objects, as discussed in more detail in Remark Our motivation for introducing Equation (42) to prove that for on verifying the assumptions of Proposition 2, its associated Bregman divergence dø, satisfies the property given in Assumption 1. First, we can observe that as = (T, # E II S4J4,, we have do, (T, > W # 14). Then, for E P2.ac (IRd), assuming that Vw, ((1)
(Tc) is invertible, we can leverage Brenier's theorem [20], and show in Proposition 15 that the optimal coupling of Equation (42) is of the form (Tv."" Id) with T argmint do, ld). Moreover, if
v E (2aac((Re) we also have that Tuk is invertible with inverse du = a dø (Id, T).
Proposition 15. Let p. E (2, V P2(IE) and assume Vw, invertible. Then,
1. There exists a unique minimizer y of (42). Besides, there exists a uniquely determined p-almost everywhere (a.e.) map TH. .Rd - Rdd such that 9 (TH, Id) Finally, there exists a convex function Rd + R such that Tu." = Vu =VooV p-a.e.
dy
2. Assume further that v E (Rd). Then there exists uniquely determined v-a.e. map Te Rd + Rd such that γ = (Id, T Moreover, there exists a convex function ou Vw, Vu v-a.e. and TH. T ld v-a.e. and bn
v : Rd + I. such that T Tr.p Tu.y Id p-a.e.
3. As a corollary, Wolv, minT da, (T, = minT do (Id, T).
Proof. 1. Observe that problem (42) is equivalent to
/ | dy rEll(e.p)
inf
(43)
Then, since for any y E H(,, (Id, Vw,d(H)) E IIV, Vw, we have
/ |ll--vwa Vw, dr(x.
/ ||. -1==212 d52,) 2). (44)
inf vell(u.u)
inf
sem (ww
27

---

Let E P2.ac (IR... Since (() is invertible, there exists convex function u such that (Vu) # W. 0( # and the optimal coupling is of the form 7" = (Vu, x Vw, (()) Let y Vuo V W; (1)), Id) [T n1), then
Vw, [1 E P2,ae( By Brenier's ' theorem,
/ (VVu(VW; Vw, de(v)
(45)
/ - (w)((u)(l)B d)(
Thus, since [1,,, p). y is an optimal coupling for (42).
2. We symmetrize the arguments. Assuming v E P2.c (R and (od = Vw, invertible, by Brenier's theorem, there exists a convex function such that (Vo)øy= Vw, (()) (and such that Vu o Vu = ld v-a.e. and Vu o Vu- = ld Vw, (()) and the optimal coupling is the form 7"h = (Id, Vu Let = (Id, Vw, (w, V E (T(,,,). then
/ || 2112 d7 (x; = / | vo(a)) du(x)
/ | Vw,
(46)
/ | Vw. dy(r,
Thus, since E [[(,,))) /)) is optimal coupling for (42). Moreover, noting TH. Vuo 6() and T 6. = Vw, (1x) -ioVv, we have u-a.e., TTu On = Vw, Voo VW,VV,WW,WWn x(4) = Id
and v/a.c., TH T Vuo Vw, Vw,o(u) Vo Id from the aforementioned dv consequences of Brenier's theorem.
We continue this section with additional results relative to the invertibility of mirror maps, which required in Proposition 2. For a potential energy V(n) = Vdd,, since Vw, V VV, then Vw, V invertible provided VV is. This is the case e.g. for V strictly convex. We state in the two next lemmas conditions for an interaction energy and for the negative entropy to satisfy the invertibility requirements.
Lemma 16. Let μ E P2(Rd) and let W Rd be even, e-strongly convex for e >> and differentiable. Then, for W(μ) = ff W(x - y) du(()dd(y)) Tw,W(H, is invertible.
Proof. On one hand, Vw, W(n) = VW Moreover, W e-strongly convex is equivalent to
Ya,y E Rdd x + y. (v VW x - v) 2 <||> v112.
(47)
which implies for all x, 2 E Rd (VW(x with respect to μ. implics
VW( 2), v) 2 c| 2112. By integrating
((VW * ()) - (VW +1)(),x=s)) = (VW(x= / -VW
y) du(=) 2 <||00 - y
(48)
Thus, VW * p c-strongly monotone, and in particular invertible [2, Theorem 1].
Lemma 17. Let μ E (R*) such that its density is of the form x eV with V Rd + R e-strongly convex for >> O. Then, for H(1)) = J log (p(z)) dp(xx with the density of p w.r.t the Lebesgue measure, H(p) invertible.
Proof. Let μ such distribution. Then, Vw, H(U) = V log p = VV. Since V is e-strongly convex, then VV is e-strongly monotone and in particular invertible [2, Theorem 1].
We conclude this section with discussion of (42) with respect to related work.
28

---

Remark 1. The OT problem (42) recovers other Of costs for specific choices of o. For instance, for (T) = |Tliiyyy it coincides with the squared Wasserstein-2 distance. And more generally, for f V T du, since by Lemma 31, for all T,S (²),),
(T) (T)
(ox(T. dv du(x)),
(49)
where dy is the Euclidean Bregman divergence, i.e. for all z. R., dv v) ((x) V(y) (VV(y), - v). Wa coincides with the Bregman- Wasserstein divergence [103]
f dv(x,.) dy(a,u). yEll(u.))
By v) inf
v)
(50)
D.2
Continuous formulation
Let Φ (2) + E be pushforward compatible and superlinear. Introducing the (mirror) map ((u) = Vw,o(p), we can write informally the mirror descent scheme (4) and its continuous-time counterpart when - as
(v((u)) Vwz)(uz) (0((++1) Tk+1 = =(((s)= T F(ux)
(p(u))= = Vwgo(u)) (((())===) F(,e(
(51)
However, de ((m))
(.W,o(()) H (vt) where Hop.. : =2((e) - (24()) is the Hessian operator (defined in Appendix C.2) such that d² ar ((se) = (Hop,, (ve), (P)e) and v E (1e()) is velocity field satisfying 0.p4 + div(uevs) = 0. Thus, the continuity equation corresponding to the Mirror Flow is given by
0:44 - div (Ho)))) - -w,, =0
(52)
For oi as Bregman potential, since HoW (v) = VV (see Appendix C.2), the flow is a solution = O. For F(u) (L(u||iv) with x e- this coincides
of 0ep:- dii-div((g (dv((p((V VV, with the gradient flow of the mirror Langevin [3, 130] and with the continuity equation obtained in [104] as the limit of the JKO scheme with Bregman groundcosts. For 6 = F, this coincides with Information Newton's 's flows [126]. Note also that Deb al. [41] defined mirror flows through the scheme T > 0 of (51), but focused on F(a) KL(u|lp) and ((μ) = 2 W (2,
D.3 Derivation in specific settings
In this section, we analyze several novel mirror schemes obtained through the use of different Bregman potential maps in (4), and used in various applications in Section 5. We start by discussing the scheme with an interaction energy Bregman potential. Next, we study mirror descent with negative entropy or KL divergence as Bregman potential. For the last two, we derive closed-forms for the case where every distribution is Gaussian, which is equivalent to working on the Bures- Wasssereenn space, and to use the gradient on the Wasss-ssssnnttiin space [43]. In particular, this space is a submanifold of P2 (IRa) and the tangent space is the space of affine maps with symmetric linear term, i.e. of the form T(x) = b + S(x - m) with S E SaR)
Interaction mirror scheme. Consider as Bregman potential interaction energy (T) [J W - T(x')) du(z)dp(zz). The mirror descent scheme (4) is given by
Ak 2 0, (VW * 11 T = VW * 14 - Vw,
(53)
For the particular case W (x) = |>| the scheme can be made more explicit as VW * u((z) = VW(w y) du(y)) = f( - y) du(()) = x m(μ) with m((a) = ydu(y) the expectation, and thus (53) translates as
Vk 2 0, Øk+1 m( = 2z m(px) + TW,, Ink ~ pk.
(54)
On one hand, recall from Example that the Hessian of di is given, for / E P2(RR)) v E T (p),
by
Vx E Rdd, = - f e(v) du(v) + +(x),
(55)
29

---

since V Id. On the other hand, the mirror descent scheme (54) can be written as, for all l: 2 0,
Uk+1 yk vV, ((4))))) y #s - -((u)); 2k ~ pkk
(56)
Passing to the limit
+ O, we get
dyt:
Vw, F(u)(a))) Ur 2, -m(p)) It pe,
(57)
dt
where dye de
dua
dm(())) de
Now, by setting 24 (x)x)
dحy de
by integration by part, we have
d
m(ue) = =
/ Oept
= a:p:=
f - x div div(p = / ve(v) duz(y).
(58)
di
Combining the latter equation with (55), we obtain as expected that dy (1))))))
dr
Negative entropy mirror scheme. Consider the negative entropy x(1) f log (o(x)) du(x))
where du(x)) = (( and (T) = o(T For such Bregman potential, the mirror scheme (4) can be written for all & 2 0 as
V log Pk+1 Tx+1 V log Phe FVw,F(uk)
(59)
In general, this scheme is not tractable. Nonetheless, supposing that μ4 N( Ex) for all k 2 the scheme translates as
-Epn (TR+1(ZE) - mx+1) = -2xx (Ix - m²) - - -Vw, F(ux) 2y ~ Pk.
(60)
For an objective functional F(u) H(m) + V(u) with V(z) = - x, the scheme
(E m++1) =-x (xx mx) - T( - 2i - (04-mma) + Σ mx))
=-(1 -T) (5x m&) -I-L
(61)
=-((1 ( -11- + - Iz (1 T) m,-
Assuming = 0 for all k; we obtain the following update rule for the covariance matrices:
E ((1-))Σ; + -
En ((1 - +) +
(62)
We illustrate this scheme in Figure 2.
For F(x) H(2)), we obtain
-E (T&+1(T)) m4+1) -11-)) mm);; I& uk-
(63)
Assuming my = 0 for all k, for < 1, we obtain the following update rule for the covariance matrices:
E-1 k+1 (1 - -- i.e.,
(64)
Σ = (1 T)
Σx
Eo 2
(65)
(1 - -)2h 0
The continuous time analog of this scheme is thus l4 : str + N(0, Ex) with Eg = 22 and the negative entropy decreases along this curve
H(())) / log (pr(x)) dx)(
-f log (2T)) /
due(x)
/det E₁
(66)
d log(2T)) log det (e22
1
Ir
12
xx dulls)
2
2 log((TT)) di
log(Ar):
30

---

where (Ae)s denote the eigenvalues of Eg. This is much faster than the heat flow for which the negative entropy decreases as [129, Appendix E.2]
d log(( 2
d
H(p,) = 2
Σ
log(A, + 2(),
(67)
1=1
with the scheme given by [129, Example 6]
m++11 m0
4k 2 0,
(68)
[EE+1 = Ex (Ia + )).
With our notations, the heat flow is the continuous time limit of the scheme (4) for the same objective F but for quadratic Bregman potential (T) = (which recovers the Wasserstein-2 geometry, hence Wasserstein-2 gradient flows).
KL mirror scheme. Suppose we want to optimize the KL divergence, i.e. a functional of the form F(u) - (()) + H(u) where ((u) =f Udu, Then, natural choice of Bregman potential is also functional of the form Φ(μ) = (p) + H(u) with (p) ff Vdd,, with U a-convex and 3-smooth relative to V.
In that case, we obtain the smoothness of F relative to 6. Recall we denote FT = F(T for TE (u)). Then for all S E (u), we have adg, (T, S dg (T,S) - 3dg, (T,S), hence dz, (T.S) = dz (T,S) + dg, (T,S) & dR, (T, + 3dg, (T,S) I m, (T,S). (69) Similarly, dz,, (T, 2 min ml, (T,S)
dd,
We now focus on the case where all measures are Gaussian in order to be able to compute a closed- form, i.e. (xx) = 1 (ez - m)"'= (z m). (zz) = 1 A z and for all k 2 0, μ = N (m,,, Ex). In this case, recall that V log (2) -x (x my). Then, each step, the mirror descent scheme (4) writes for 2& ~ Pk. k 2 0 as
(I(++1)) + V log [c++ (Ik+1)) = (xx) + Vog (Hx(Ex) (VVU(xx)) + V log (#x()x))
A Øk+1 Σn 1
x+1 (Ik+1 mx+11)
A A Σx (0rk mm)) = -0- (IT m) 2x (mx)=))))
(A - E |ak+1 + E mx+1
= -A-1- (1 - 1) - - Iz + (1 - -) m
ma + + m.
(70)
Thus, we get for the expectation that
- -)= (A (A-1 - -1----1- 2 )m++ + 2x 7 -A-1-(1-
--E-1 -TE-))ma( -))=
me ---'m m
> A m (A-1 -) + -m
= M++1 = (Ia (Ia-TAD mm + A
(71)
We note that the latter update on the means coincides with the forward Euler method in the forward- backward scheme, see (116) in Appendix F, which uses as Bregman potential ø = V. Thus, the entropy does not affect the convergence towards the mean, which can be done simply by (precondi- tioned) gradient descent.
For the covariance part, we get
(A - A--1-111 T +++ (A-1 -xr11)
= -A-11 - ---- 11-1-(1-17711 (1- Σx 11-1--7-- (1- -))=))) (72)
Now, supposing that all matrices commute, we get
A - - -2A + 1 k+1 = (A-1 - 7エ - 2(11 - 2- -
+ (1 -)²
(73)
A 2 + Er1 = -A---~55-1 -1 1 (1-)))
Ea+1 + A A3 = =(1a -i--7Ae 10--A>--130 + 2 + ----------111--33--------133333-
31

---

Denoting
C= -= + 2 + (1--))AAE4
(74)
the update on covariances equivalent to
Σ CE + + = -.
(75)
Thus, Σa+1 = -C²-))))))
D.4 Mirror scheme with non-pushforward compatible Bregman potentials
We study in this Section schemes for which the Bregman potential o₂ is not pushforward compatible, and thus for which we cannot apply Proposition 2 and thus Assumption I may not hold a priori. An example of such potential is (x) (T. L) where Pe. [2) + I²( is a linear
autoadjoint and invertible operator. Since V = P. taking the first order conditions, we obtain the following scheme:
4k 2 0, TA+1 = Id --- Vw. F(pn).
(76)
In particular, this includes SVGD [71, 83, 84] if we pose PT iS, with Sy. (z)
H defined as S [(x, (((r,))((=)dp((a)) which maps functions from (2²()) to the reproducing kernel Hilbert space H with kernel k, and with
+ (+) the inclusion operator that is
the adjoint of Sμ [71]. It also includes the Kalman- Wasserstein gradient descent [56] for which
P-- = fz m(4) @ (x m(n) du(z) is the covariance matrix, where m(μ) f du(z).
More generally, for (T) f (V o Tid, we can recover their mirrored version, including mirrored SVGD [113, 114], i.e. Tx41 = VV* o (VV - - ",,
Kalman- Wasserstenn. We focus now particular choice of linear operator Py. Namely, we take P, = C())T with C(p) = (ʃ (x (m))) x (x mm(1)) du(x)) the inverse of the covariance matrix. In this case, (76) corresponds to the discretization of the -assen--asssesttein gradient flow
[56]. We now show that it satisfies Assumption 1. First, let us compute the Bregman divergence associated to d:
VT,S E do, (T,S) = (T, C()))) + 1 (())) (C(N)S, T)
(ITCCW)(T-S)))) + (S ((=))
(77)
(T-))aa
For 7 = #1, write
de (T = ] / (x 1) dy(x,, v).
(78)
Moreover, the problem inf inf (o) J [1C(μ) v dx v) is equivalent to
- C(H)y dv(#,)) v). inf yEll(a.
(79)
which is squared OT problem. Thus, it admits OT map if C(μ) is invertible and or is
absolutely continuous.
Second point of view. Another point of view would be to use the linearization with the gradient corresponding to the associated generalized Wasserstein distance, which is of the form Ww))() P.- (W,F)) [46, 56], i.e. considering
Te41 argmin do (T, Id) + (Vw) T TEL?(n)
(80)
where we assume that Fw)) E (p). In that case, using the first order conditions,
o Tht1 = Vw,0(()) Then, for dh satisfying Assumption the convergence will hold under relative smoothness and convexity assumptions similarly as for the analysis derived in Section 3.
vJ(Tx+1) = + V
(81)
32

---

E Relative convexity and smoothness
E.1 Relative convexity and smoothness between Fenchel transforms
In this Section, show sufficient conditions satisfy the inequalities assumed in Proposition and Proposition 6 under the additional assumption that, for all 2 0, Fuz is superlinear, lower semicontinuous and strictly convex. In this case, we can show that F is Gâteaux differentiable, and pa.
thus we can use the Bregman divergence of psy
Lemma 18. Let ④ (p) - R be superlinear, lower semicontinuous and strictly convex function. Then, of is Gâteaux differentiable.
Proof. Fix g E (p)) Notice that
fE (e000g)) =
-(f) sup (f.g) - o(ff). SEIP(N)
(82)
So to prove there a unique element de (g), we just need to show that, setting (f) == - (f.g)
((f), the problem inf fE1?(()) Og(f) has unique solution. Under our assumptions, Og is lower semicontinuous and strictly convex. Since is superlinear, 0g coercive, i.e. lim Φf(
+oo. There thus exists a solution [8, Theorem 3.3.4], which is unique by strict convexity. Hence 0 (g) is reduced to point, which is necessarily the Gâtcaux derivative of ø at g.
This allows us to relate the Bregman divergence of o² to the Bregman divergence of é.
Lemma 19. Let Φ: (p) + R be a proper superlinear and strictly convex differentiable function, then for all T, S E (p),,, dø- (Vó(T), Vó(S) ds T).
Proof. By [99, Corollary 3,44], we have o" (Vo(T) since é is convex and differentiable. By Lemma 18, Φ² is invertible and by [10, Corollary 16.24), since Φ is proper, lower semicontinuous and convex, then (V6)-1= = VΦ"
Vo(T)) - (T) for all T e L2(
Thus, for all T,S (p))
dø- (Vo(T), Vo(S)) 6z - Vo(())
(V
6" (Vo(T) - - (Vo(s) (S, Vo())
(T((T),))2u - - (V
+ ((S)
(83)
-S, Vo(T) ( (o(S))
=o(S) ((T) (Vo(T), L())
T)
Finally, we can relate the relative convexity of á relative to 30 by using inequality between the Bregman divergences
and v. In particular, we recover the assumptions of Propositions and for dit that 3-smooth and a-convex relative to Fic.
Proposition 20. Let o.b (2) + R proper, superlinear, strictly convex and differentiable. is 3-smooth (resp. (-convex) relative to 1² if and only iV E (2),)), (VV)()))) Vip(S) < 3dg (S, (resp. do ds (T)(T), V*((S) 2 ad,(S,T)) T)).
Proof of Proposition 20. First, suppose that Φ is 3-smooth relative to v Then, by definition,
vT,S E L²),), do S) < 3dg- (T,
(84)
In particular,
de (V)(T), Ve(S)) - dd (VV(T). VD(()) = (S. T),
(85)
using Lemma 19.
33

---

On the other hand, suppose for all T,S E (p), do (V ve(()) & (d, Then, by first using Lemma 19 and then the supposed inequality, we have for all T.S E (1),
Bdg- (VV(T), Vi((S) a T) de (T) Vv)(
(86)
Likewise, we can show that & is a-convex relative to &b if and only (V)(() VV)))) (S, for all T,S E (u).
>
Links with the conditions of Proposition 5 and Proposition 6. Proposition 20 allows to translate the inequality hypothesis of Proposition 5 and Proposition 6. Assume that for all k. Fiee is strictly convex, differentiable and superlinear. We first note that it implies that Frs. convex along + ((1 - +tTTx+1 + fId) # Mk- Moreover, by Lemma 18, V.F. is differentiable.
jlg.
Note that this assumption is satisfied, e.g. by (T) f T differentiable. Indeed, in this case, oy is also 7--strongly convex, and satisfies for all T,S E L2)),
T du for V 7-strongly convex and
døp (T, = oz (S) (S), ( " |T -T
(87)
? " T-
(x) 2 (S) (S), S) +
For S = 0, and dividing by ||]| the right term diverges to too when || + to0, and = too, and 6₂ is superlinear.
thus lim iim/Tilea 6N
This assumption is also satisfied for interaction energies oe (T) fJ (T(xx) -(T(=))T(()) T(y)) du(z)dp(v) with W n-strongly convex, even and differentiable. Indeed, by strong convexity of W in O, we have for all x, y E R..
W(T(x)-Tip (()) - ((v))
" T()) T
2
(88)
" inf [[((x) =1.2-
2
EER
Integrating w.r.t @ nt, we get
W (T) - (())= 2 7 inf
/ /IT(()) - =12 du(x),
(89)
2 sERe
and dividing by [T]]L>(() we get that Th is superlinear.
For curve + Me, we define Fí on pr as Fi(m) = = with Fz the convex conjugate of Fj in the (²()) sense. Then, we can apply Proposition 20, and we obtain that the incquality hypothesis of Proposition 5 is implied by the ß-smoothness of s² relative to F along & + ((1-t) ((,,t)V,,,
Tk+1) # Mx since
day. (( Ta+h, Vw,u(()) S 3dg (Id, Tx+1)
(90)
=d pk
F(me))) TK+1, F(me))
where we used Proposition compute the gradient V (Tk+1) Vw, F(uez)) o Tst1 Similarly, the condition of Proposition
dayg (Vw,F(TTHM) T, Vw, 2 adyy, (Id, T)
adj. (Vw, T, Vw,F(())) (91)
is implied by the a-convexity of sh relative to Fx along + ((1 - () +
WV, o T) Me.
These results are summarized in Proposition 7 and shown formally in Appendix H.7.
34

---

Convergence towards the minimizer in Proposition 6. We add an additional result justifying the convergence towards the minimizer in Proposition 6.
Lemma 21. Let (X, be metrizable topological space, and f X + R R be strictly convex, T-lower semicontinuous and with r-compact sublevel set. Let zo E X be the minimizer of f and take a sequence (In nEN such that (T)) - (xo)). Then, (In) T-converges to Po-
Proof. The existence the minimum is given by [8. Theorem 3.2.2]. For N large enough, (Zn)n2N lives in the T-compact sublevel set of f. since Zo belongs it and f(zo) is minimal. We then consider a subsequence 7-converging to some x². By 7-lower semicontinuily, we have f(xo) & f(=") & lim inf (Fa(()) = f(o), so f(ru) f(x") and by strict convexity 2o x². Since all subsequences of (Zn)o2N converge to x² and the space is metrizable, (In)NEN 7-converges to To.
The typical case when X is a Hilbert space and T the weak topology. One could wish to have strong convergence under coercivity assumption, however "In infinite dimensional spaces, the topologies which are directly related to coercivity are the weak topologies" [8, p86]. Nevertheless Gâteaux differentiability implies continuity, which paired with convexity gives weak lower semiconti- nuity [8, Theorem 3. We cannot hope for convergence of the norm of Ir to come for free, as the weak convergence would then imply the strong convergence.
E.2
Relative convexity and smoothness between fuunctionals
Let U, V Rd be differentiable and convex functions. We recall that a-convex relative to U [88]
D dv(x, 2 adv (x,y).
(92)
Likewise, V is 3-smooth relative to
dy (z S Bdy (2, v))
(93)
Relative convexity and smoothness between potential energies. By Lemma 31, for Bregman potentials of the form (T)) = f V
du. the Bregman divergence can be written as
/ dv (T()), du(x).
VT,S E (²()). do, (T,S) =
(94)
Thus, leveraging this result, we can show that relative convexity and smoothness of 0% relative of is inherited by the relative convexity and smoothness of V relative to U.
P2(RR), (T) SV T du and (T) U T du where V,U Rd + R. are C If is a-convex (resp. B-smooth) relative U Rd - R. then op is a-convex (resp B-smooth) relative to vyp.
Proposition 22. Let
Proof. First, observe (Lemma 31) that
ff d (T(x), du(x).
Vμ P2(RR), T,S LP))), da, (T,S) =
(95)
Let E (2(RR) ",RRRR)TTTEEEEE. L1)) I is a-convex relatively to U, we have for all x, E Rd.
dy (T(x). S(y)) 2 adv (T(x),S(y))
(96)
and hence by integrating both sides with respect to μ.
de (T, 2 adv, (T, S).
(97)
Likewise, we have the result for the 3 3-smoothness.
35

---

Relative convexity and smoothness between interaction energies. Similarly, by Lemma 32, for Bregman potentials obtained through interaction energies, i.e. (T( [J (T(x) T(x') du(x)du(zz)) then
de, (T.S) - /f dw (T(+) -T())) T(x)),(()) S(x')) dn(x)du((²))
& E
(98)
also allows inherit the relative convexity and smoothness results from R2.
Proposition 23. Let E P2(R W,K R + R be symmetric, C² and convex. Let oy (T) (T(x)) - -T(x') dp(x)dd((z") and (T) Jf ((T(x)-T(aa) du(ar)dp(a') IfW is a-convex relative to K, then 0n a-convex relatively 1pp- Likewise, if W is B-smooth relatively to K then Qy is B-smooth relatively to a
Proof. We use first Lemma 32 and then that W is a-convex relatively K:
dø (T,S) - 2 // dw (T(x) - T(x'), S(x')) du(z)du(a²)
a // dx (T(x) S(x)).((x) S(z') dp(x)dd(z))
(99)
2
ad (T.
Likewise, we have the result for the 3-smoothness.
Thus, in situations where the objective functional and the Bregman potential are the same type and either potential energies interaction energies, only need to show the convexity and smoothness of the underlying potentials or interaction kernels. For instance, let V : Rd - be a twice- differentiable convex function, such that V pp([(1=)))) with p,. polynomial function of degree and || llop the operator norm. Then, by [88, Proposition 2.1], V 3-smooth relative to where for all x E RR, h(x)) = - | 1
Relative convexity and smoothness between functionals of different types. When the functionals do not belong to the same type, comparing directly the Bregman divergences is less straightforward in general. In that case, one might instead leverage the equivalence relations given by Proposition 13 and Proposition 14, and show that SG F or F og is convex in order to show respectively the 3-smoothness and a-convexity of relative to G. For instance, we can use the characterization through Hessians, and thus we would aim at showing
d2 de Flue) s (
d²
d² Flue) 2 (D
d²
(100)
along the right curve ( + Pt.
For instance, consider an objective functional F()) Sf W(z y) du(z)du(=z") and another functional ((u) Vdu. Then, by Example and Example 2, we have, for pe (Tr) and T, = S + tv,
d2 ((m)) diz
du(x),
(101)
and
F :l/ ///VAWTT(()) -T(()) v())) v())),,()) du(x)dp(y).
d²
(102)
36

---

To show the conditions of Proposition 3, we need to take S = Id and v T Id, and to verify for E [0, 1] the inqquality i.e.
d² de
3 d2 d
F(m)
t= S
((m))
t=
/| (I
v(v)), v(x)) (pn(t)dum())
$ (VVV(T.()))(()).dd o(()) du)()))
(103)
{ (MOO), / (T - -VW((,,(>) T.(y)) pa ( (1) 2.
- T.(U)) - -VVV(
For example, choosing W(x) = 111111 then vaw = La and F is 8-smooth relative to G as long as V la for any E 1].
F
Bregman proximal gradient scheme
In this section, we are interested into minimizing a functional F of the form form F(P) = 9(m) + (() where G is smooth relative to some function & and H. is convex on (u). Different strategies be used to tackle this problem. For instance, Jiang et al. [65] restrict the space to particular directions along which H. is smooth while Diao al. [43], Salim al. [109] Proximal Gradient algorithms. We focus here on the latter and generalize the Bregman Proximal Gradient algorithm [1 I], also known the
scheme. It consists of alternating forward step on and then backward step on H. i.e. for k 2 O,
Sk+1 argminseL day Id) + VW, LF Tk+1 argmin Aeezas (T, (T(
Mk+1 = (8k+1)#H4
pk+1 (+++1##K444
(104)
The first step of our analysis is to show that this scheme is equivalent to
(T-11 argmin dop,, (T, Id) + +TV T Id) (2())) + (((T
[k+1 = (Tx+1)##k-
(105)
This is true under the condition that μk (2 implies that Pk+1 E (R Proposition 24. Let by be pushforward compatible, p40 E (RR)) and assume that if pm E P2.ac((R) then Vk+1 E (R))) Then the schemes (104) and (105) are equivalent.
Proof. See Appendix H.ll.
We are now ready to state the convergence results for the proximal gradient scheme.
Proposition 25. Let μo E (Rd) T S 1 and F(p) G(u) + H(p)). Consider the iterates of the Bregman proximal gradient scheme (104), equivalently (105). Let 2 Assume Hly, is
convex on (2()) and G B-smooth relative to & along t + ((1 t)ld + x yM Then, for all T (3(()))
1 F(uR+1) ((TTmuR)+ + (Vw. Le) + don, (T, -dam, (T, Tr+1).
(106)
Moreover, for T Id.
1
F(ux+1) S F(pu)
T den, (Id, Tk+1),
(107)
37

---

i.e., the scheme decreases the objective at each iteration. Additionally, let a 2 0, V E P2((R))
and suppose that by satisfies Assumption 1. IfG is a-convex relative to & along + ((1 - -)ld (T Bun * HR, then for all k N 1,
a F(P)) - F(v) 5 (1-Ta) Wolu, po) S 1- aT Wo(v,
1-aT
(108)
kT
Proof. See Appendix H.12.
We verify now that Proposition 24 can be applied for mirror schemes of interest, Salim et al. [109, Lemma 2] showed that it holds for the Wasserstein proximal gradient when using some potential energies, more precisely with Φ(u) = f àll 112 du and G(u) = =f dp with U (strictly) convex. We extend their result for ((u) = U du and ((H) = f V du for V strictly convex and 8-smooth relative to V.
Lemma 26. Let μ E (T. g(v) / U du, (T( fV T du with V strongly convex and U 3-smooth relative to V, and T = VV (VV --VU). Assume T < B. then Tap E a2 (Re).
Sketch of the proof. The proof of the lemma is inspired from [109, Lemma 2]. We apply [5, Lemma 5.5.3), which requires to show that T is injective almost everywhere and that det VT > 0 almost everywhere. See Appendix H.13 for the full proof.
To apply Proposition 25, we need H to be convex along some curve. We discuss here the convexity of the negative entropy along acceleration free curves. Let μ E P2.ac (Rd), and denote p its density w.zsr the Lebesgue measure. For H(p) = f (p(z)) dz where f : IR + R is C' and satisfies f(0) = lim x (x) = 0 and æ + f(z (o )xdd convex and non-increasing on IR,,, then by [117, Theorem 4.21, H convex along curves 14 = ((1 t)S + # obtained with S and T with positive definite log for which H corresponds the negative entropy.
Jacobians. This is the case e.g. for (xx
By Remark 3, to be able to apply the three-point point inequality (that is necessary to obtain the descent lemma), we actually only need H to be convex along ((1 - -t)Tx+1 + tId) # uhs and along ((1
t)Tx+1 + (TTEE") yd for the convergence.
Gaussian target. In what follows, we focus on G( = Udu with Uxx) = (aa--m)
m) for Σ E S++(R), H the negative entropy and with a Bregman potential of the form d(u) = Vdd with V(z) = 1 -
xx. Moreover, we suppose μ4 = N ((m.. Eo). In this situation, each distribution jlg also Gaussian, as the forward and backward steps are affine operations.
Assuming the covariances matrices are full rank, TR+1 is affine and its gradient is invertible. Moreover, by Proposition 15, T Vu Vw,6(uz) for Vu an OT map between Ww, ((1))##* and v.
Since each distribution Gaussian, and Vw, ((x)())x A-- is affine, it has a positive definite Jacobian. Thus, using [ [117, Theorem 4.2], we can conclude that we can apply Proposition 25.
Closed-form for Gaussians. Let G(n)= = Ud with U(x) { m m)) Σ S (B), m E Rd, and H(1)) = f log (p(x)) du(x)) for du = p(r)dæ. For the Bregman potential, we will choose (())= = =J Vdu for V(x) = (x, A 'z). Recall that the forward step reads as
Sk+1 VV (VV -VV,9(ux)) vk+1 (Sk+1)
(109)
Since VV(x) A and p4 = Nmmm, Ex), we obtain for Fy ~ pkk
S (xx) = A(A-1ap by - - -- (xx - m)) = 2x - -AΣ - - m).
(110)
Thus, the output of the forward step is still a Gaussian of the form Uk+1 = N N( Extz) with
mx+3 (Ja - -AE-1 TAE-Im
(111)
E = =(Ia -T--AA--TTTA Sy -TA--1)
38

---

In this subsection, we sum up how to implement the different schemes practice, given a finite number of particles. In all cases, first sample 2(0) ...., (0) po. then apply the scheme to Since VV is linear, the output of the backward step stays Gaussian. Moreover, the first order conditions give
VV o Tk+1 + TV log(ps+1 Tk+1) VV
Vx, A T (T (Tx+ı(z) mx+1)
- Va, x TE+1 - -AE (Tx+1 mx+1).
(+))-
(112)
Thus, the output is Gaussian (m,++1 Ex+1) with (mx+11, Σx-1) satisfying
m++1 = m
(113)
E (Ia -AE Ex+1444 TAI
Moreover, if A and Ex+1 commute, this is equivalent to
E (2T Ex+ )x++ + = =0,
(114)
which solution is given by
Σx+1
(E EKty + + (Ex+ (4TA + Ex+ ()
(115)
To sum up, the update is
TAE (la (4 - Mk+1 ( 2[E++} 2TA (Ex+ (4rA + +k++))
Mk+1 = N (la-
(116)
For A = E, we call it the ideally preconditioned ForwardBBakward scheme (PFB).
G Additional details on experiments
G.1 Implementing the schemes
Em=1 6,000
Mirror descent. In general, for o pushforward compatible, one needs to solve at each iteration 0,
V Vw.0(1441) T Ta+1 = Vw,6(m)) rV w..
(117)
If (()) f V du with VV having analytical inverse, the scheme can be implemented as Vk 2 0, vi (I, x. (k+1) VV (VV((zz T F(gh))
(118)
Except for this case, one cannot in general invert Vw, ((44() Ta+1 directly. practical workaround is to solve an implicit problem, see egg. [133]. Here, we use the Newton-Raphson
algorithm. Suppose we have μa = Σ (.) and we looking for pa+ - Ei 8z Then, the scheme is equivalent to
VjE (1,...,,n)
n), Gy(Z1,........ Fn) =0
(119)
for
GI(ZI,.... Rn.) )- = Vw,Φ
Σ i=1
(xz) - (wwo(l()(zz)) T wwu (120)
Write g(x1,...,zz,) ..) (ZI. Rn).... (r kn)) Then, each step k. we perform
the following Newton iterations, starting from (a ....., (K)
(k)
(2( ..... 1 (ke+1) (zf (ke) ..... x=((t()) (J (kz) x...,((()))) (b)))) g00 (..)), (121)
The Jacobian is of size nd x nd, which does not scale well with the dimension and the number
of samples. We can reduce the complexity of the algorithm by relying on inverse Hessian vector products, see egg. [40].
39

---

KL (u
0..0
KXd - :
-0.10
K(x) 2 + 1
NEM
-0.12
PFB
-0.14
PKLM
10
KLM
2.5
12
Time (r- (terations)
Time (T- lterations)
Figure 4: (Left) Value of W along the flow for two difference interaction Bregman potentials, (Right) Trajectories of particles to minimize W.
Figure 5: Convergence towards Gaussian N (0, D) with D diagonal and uniformly sam- pled on [0, 50
Preconditioned gradient descent. scheme can be implemented as
Plugging the empirical measure in (11), the preconditioned
4 2 0, Vi E n), I Iz (k)
(k+1)
- -Vh (Vw,
(122)
G.2 Mirror descent of interaction energies
Details of Section 5. We detail in this Section the first experiment of Section 5. We aim at
minimizing the interaction energy W(4) = fJ W (x-z)) y) (u(z))u(()) for ((z) 1 2 = well-known that the stationary solution of its gradient flow is Dirac ring [27]. Since the stationary solution is translation invariant, we project the measures to be centered.
We study here two Bregman potentials which are also interaction energies. First, observing that (z) = 2zz" + (11=1 - 1) Ja. we have for all z,
Wllop S 2| 2 + ||=12 1 3 3 +1 p2(11=112),
(123)
with p2(t) = 3t + 1. Thus, by [88, Remark 21, is 3-smooth relative to K2(2) = 1 11 + 11
with 8 = 4. Thus, using Proposition 23, Wμ is 3-smooth relative to (T) 1/2 ʃf (T apply Proposition 3.
T(x') dμ(z)du((z") for all μ., and
Under the additional hypothesis that the measures are compactly supported, and thus there exists M > 0 such that ||>:| < M for y-almost every x, we can also show that W is 3-smooth relative to
K2(2) = =11. Indeed, on one hand, V2K Ia and (²)(zz) = 2zz" + (1 1 1) Thus, for all E Rd.
V W (2)u = u)2,0)² + (11=112 11)1 < 31 < 3M = 3M V K()). (124)
In Figure 4, we plot the evolution of W along the flows obtained with these two Bregman potential, 100 particles, with step size of 7 0.1 for 120 epochs.
starting from μif (N I,) for
Illconditioned interaction energy. an ill-conditioned kernel ((2) ] [2 2)²
We also study the minimization of an interaction energy with
1-''E' where Σ E S **++(R) (R) but is possibly badly conditioned, i.e. the ratio between the largest and smallest eigenvalues is large. In this case, the stationary solution becomes ellipsoid instead of a ring. In our experiments, we take = diag(100 0.1). For each scheme, we use l40 = N (0, 0.252 n = 100 particles and step size of 0.1.
On Figure 1, we use Bregman potentials which take into account this conditioning, namely we use (+(²) -ー and (2)
and we observe that the convergence much faster compared to the same kernels without preconditioning. For (²) scheme becomes
T the
(V pk+1) T8+1 = VK x -Ww,
m -T( = m(ux)) - 7 - 'T 1)Id (125)
- 1a41-m(H441) ld (m(()))()))))))))))))))))))))))- rldd l)ld.
40

---

Maal
Vila
TrasGo
-.00
GB FredtD
-.44
-266
Tondonst
Teritioss
Figure 6: (Left) Value of W over time and trajectory of particles using K2 and KJ as interaction kernels. (Right) Value of W over time and trajectory particles for the Wasserstein gradient descent and preconditioned Wasserstein gradient descent (with the ideal preconditioner h" (x) = _²²² Ex).
Thus, see that E-1 has less influence which might explain the faster convergence.
Similarly as in the case without preconditioning, using that V2 = 2Σ 22 1- 1)Σ show that
V W (2)u = 2(2, -) + (1 1) S 3M
= 3MM'
(126)
For the sake of comparison, and KJ", as well as of the usual Wasserstein gradient descent and the preconditioned Wasserstein gradient descent obtained with h (x:) 1a² Ex (which equivalent to the Mirror Descent with Φy as Bregman potential and Vax) 24 Σー Mー) We observe almost the same trajectories Kg. which would indicate that the target is also smooth compared to oy
also report Figure 6 the trajectories of particles for the use of K4
Runtime. These experiments were run personal Laptop with a CPU Intel Core i5-9300H. For the interaction energy as Bregman potential, running the algorithm with Newton's method for 100 particles in dimension d = 2 for 120 epochs took about 5mn for K2 and KÝ. and about Ih for K₂ and K.F.
G.3 Mirror descent on Gaussians
As the mirror descent scheme cannot be computed in closed-form for Bregman potentials which are potential energies, and thus are computationally costly, we propose here restrain ourselves to the Gaussian setting.
We choose target distribution =NV(0, Σ) for E symmetric positive definite matrix in R 10x and the functional to be optimized is F(u) Vdp + H(1) with V(x) = - distribution is always chosen as 140 (0, Ta). In all cases, the step size is chosen The initial 0.01, and we run the scheme for 1500 iterations. For the target distributions, we sample 20 random covariances of the form Σ UDUT with 1) evenly spaced in log scale between and 100, and U E I. 100x10 chosen as a uniformly random orthogonal matrix, as in [43], and we report the averaged KL divergence over iterations in Figure We add on Figure 5 the same experiments with targets of the form (00 D) where D is a diagonal matrix on IR 10x 10 sampled uniformly over J0, 50) 10 We compare here the Forward Borward-Backward (FB) scheme of [43], the ideally preconditioned Forward- Backward scheme (PFB), which uses the closed-for fosed-foom (116) derived in Appendix F with A = Σ, and the Mirror Descent with negative entropy Bregman potential (NEM), whose dosed-form was derived in Appendix D.3, and which we recall:
Vk O, Er ((1-))Σ;' + (1- +
(127)
We also experiment with the KI. divergence as Bregman potential (KLM) and the ideally precon- ditioned KI. divergence (PKLM). We observe that, even though the objective is convex relative to the Bregman potential, this scheme does not always converge. It might be due to its gradient which might not always be invertible. We leave further investigations for future works.
Remark 2. We note that using Bregman potential th (T) = f vb Tdp vxxx) = 1.T A-1 is equivalent to using preconditioner with he (x) 1 Ax.
Analysis of the convergence. well-known that along the Wasserstein gradient flow of the KL divergence starting from a Gaussian and with Gaussian target COmstein-Uhlenbedd process),
41

---

41
41
scRNAseq
seRNAseg
scRNAseq
F((e) 2))
F(j) (1
F(u)= = ED(,,))
F(J)) 20L,))
F(y)) (0,
F((e) EDIR.
with precond.
10
10-1
18
13
16
10
without procond.
16'
12.
38
107
55:1
3F
16
59.7
entropic map
entropic map
entropic map T2
cntropic map T2
entropic map T2
cntropic map T2
Figure 7: Preconditioned GD and (vanilla) GD vs. the entropic map T. [101] to predict the responses of cell populations cancer treatments 4i and scRNAseq datasets, providing respectively 34 and 9 treatment responses. For each profiling technology and each treatment, we have a pair (J4, v;) of source (untreated) cells and target (treated) cells. For each pair (μi, vi), with both preconditioned GD and vanilla GD, we minimize the functional F(u) (1, w D a metric-to recover the effect of the perturbation. In both cases, the prediction obtained by já: min, F(u). We then fit an entropic map Tg and predict T.tpg, We then compare the objective function values F(A,, and #(:) A point below the diagonal y x then refers an experiment in which (preconditioned) WGD provides better estimate of the perturbed population.
the measures stay Gaussian [129]. Thus, the Herward-Backward scheme has Gaussian iterates at each step [43, 109]. In this work, we also use a linearly preconditioned forward-Backward scheme, whose closed-form is derived in (116) (Appendix F). For the Bregman potential, choose = (T)
go T du for (x)) = 1 T Ex. In this situation, ((p) = Vdu is I-smooth and 1-convex apply Proposition 25. We refer to Appendix F for details the
relative şi. Thus,
convexity of H.
For Bregman potentials whose gradient is not affine, the distributions do not necessarily stay Gaussian along the flows. Thus, we work on the Bures- Wasserstein space and use the Bures- Wasserstein gradient, i.e. project the gradient the space of affine maps with symmetric linear term, i.e. of the form T(x) = + S(x - m) with Sa(R) [43]. We refer to [43, 73] for more details on this submanifold. This can be seen performing Variational Inference. We derive the closed-form of the different schemes in Appendix D.3.
Even though these procedures do not fit exactly the theory developed in this work, we show the relative smoothness of FE relative H along the curve 14 = ((1 t)Id + (T*+1) 4 under the hypothesis that the covariances matrices have bounded eigenvalues. Moreover, since dj. = day + +di 2 dR F is also 1-convex relative to H.
Proposition 27. Let X > F(u) and Σ 3 Xla. Suppose that for all l: 2 0, (1 - T) E-1 + + > O. Then, F is smooth relative to H along pt = ((1 t)ld + (Tx+1) 4 where μa = NO0, Ex) with Ex E S Σx Alg.
Vdp + H(H) with V(x) = - -I-T-111 where
St (R)
Proof. See Appendix H.14.
G.4 Single-cell experiments
First, we provide more details on the experiment on single cells of Section Then, we detail second experiment comparing the method with using static map.
42

---

Details on the metrics. We show the benefits of using the polynomial preconditioner over the single-cell datasets for different metrics.
Sliced-Wasserstele Wasserstein distance [16, 102), defined as The first one considered is the Sliced-
Vp, M E )2(R"), SW2(H -
WH d)(0),
(128)
where gd-1 (0 E Red |01| 1}. A denotes the uniform distribution 5 and for all 0 5d-1, Rd, p0(x) (x) (x, 0)- For F(H) 3S,3(u, the Wasserstein gradient can be computed as [18]
Vw, S & (((()
(129)
where, for E R. (1) = L- Pg. (F
P4u (Fpgg/l (4) with Fpg, the cumulative distribution function of Pe $1.. In practice, compute SW and its gradient using a Monte-Carlo approximation by first drawing L uniform random directions 0,,-..
01. 0,.
The second one considered the Sinkhorn divergence [50] defined as
1 OT - OT,(u,
VP, E P2(R)), (2,2(1, (T=(U,)) 2
(130)
with
/ v1l2 dx <KK( v),,
(T,( inf
(131)
-Ell(a..)
the entropic regularized OT. The Wasserstein gradient of S is simply obtained as the potential [50].
Finally, we also consider the energy distance, defined
Auu E P2(IR), ED(H, v) -===-]] /] ||
vlla d(p v)()) vv)
(132)
To compute its Wasserstein gradient, we use the sliced procedure of [60].
Parameters chosen. For all the metrics, fixed the step size
1. To choose the parameter grid search over {1.25, 1.5, 1.75)
of the preconditioner h (x) (I|==112 + 1) 1/a - 1, we for a random treatment, and used it for all the others. In particular, we used for the dataset 4i a 1.5 1.75 for the energy distance. For the scRNAseq dataset, we used a = 1.25 for the Sinkhorn divergence and SW, and a = 1.5 for the energy distance. We note that for the dataset 4i, the data lie in dimension d 48 and the metrics, we first sampled 4096 particles from the source (untreated) dataset, and used in average between 2000 and 3000 samples from the target dataset. For the test value, also added 40% of unseen cells following [21]. Note that we reported the results in Figure for 3 different initializations for each treatment, and reported these results with their mean. We report the results using fixed relative tolerance tol = 10-3, iee. at the first iteration where {F( F((t-1)l/((uuu) < tol, with maximum value of iterations of 10 For the Sinkhorn divergence, we chose E as 10 time the variance of the target. Finally, for SW and the computation of the gradient of the energy distance, we used a Monte-Carlo approximation with 1024 projections.
for the Sinkhorn divergence and for SW, and
50 for scRNAseq. For all
Comparison to an OT static map. We now compare the prediction of the response of cells to perturbation using Wasserstein gradient descent, with and without preconditioning, to the one provided by a static estimator, the entropic map T, [101]. This experiment motivates the use of dynamic procedure, iterating multiple steps to map the unperturbed population to the perturbed population v, instead of a unique static step. We use the proteomic dataset [21] as the one considered in 3. We use the default OTT-JAX [38] of Tg. The results are shown in Figure 7.
Runtime. For this experiment, we used GPU Tesla P100-PCIE-16GIB Depending on the conver- gence and on the metric considered, each run took between 30s and 10mn. So in total, took few hundred of hours of computation time.
43

---

Initial particles
Final particles
Final particles MLD
Groundtruth
71p4)
MD
20
24
350
Iberations
Figure 8: (Left) Samples from Dirichlet posterior distribution for Mirror Descent (MD) and Mirror
Langevin (MLD). (Right) Evolution of the objective averaged over 20 different initialisations.
G.S Mirror descent on the simplex
We can also leverage the mirror map to perform sampling constrained spaces. This has received a lot of attention recently either through mirror Langevin methods [3, 32, 116], diffusion methods [51, 82], mirror SVGD [113, [14 or other MCMC algorithms [52, 94].
The goal here is to sample from Dirichlet distribution, i.e. from a distribution v x e where V(x) = Σi= a; log(æ;) - aa+1 log 1--i xz) To sample from such a distribution, we minimize the Kullback-Leible divergence, i.e. F(u) (L(u|>>) V du + H(u)). To stay on the (open) simplex Aa fx E Rd+1, 24 > 0, sd+1
Li=l < 1}. we use the mirror map 6(μ) J uddjs
with (xx) Ei= x; log(z;) + (1 Σ; .i) log(1 Σ m:) for which
V+(()) log v; - log (1 - Σ [33)
(1-E,))
Vp
( [1 (1
eth
(133)
The scheme here is given by Tk+1 Vel* (v ? Vw, where Vw, F(u)) VV V log Pk. with the density μ4 estimated through a Kernel Density Estimator (KDE). We plot on Figure 8a the results obtained for d = 2, 0₁ 02 ag 6 and 100 samples. We also report the results for the Mirror Langevin Dynamic (MLD) algorithm, which provide iid samples, which are thus less ordered. We plot the evolution of the KI. over iterations Figure 8b (where the entropy is estimated using the Kozachenko-Leonent
estimator [42]).
The KDE used here will not scale well with the dimension, however, different methods have been recently propose to overcome this issue, such as using projection on lower dimensional subspaces [127], using neural networks to learn ratio density estimators [6, 49, 128).
H Proofs
H.1 Proof of Proposition 1
Let E P2(R'), S,T D(Fx), c 0. Since F is Wasserstein differentiable at Sgp, applying
Proposition at S Syj with = (S + c(T S)) 4 and y (S.S + c(T - -S)) n E IT(S v), we
44

---

obtain.
f ( - = - (1)
+ = + - d7(2,))
+o (1/ 110
F*(S) +e (vw: T(x) S(x)) du(x))
+o (9 [T(()) 5(x) du(x))
FA((S) + (V, -) + coll (l) (134)
Thus, (S, - -) F(u) E (u).
o S) Note that in the third equality we used that
H.2 Proof of Proposition 2
Let p.p E (TR and E (2((R.). Define THE argmint do, (T, Id), U 9,
argminu de, (U,Id) and let S E (²( such that Sym p. Then, noticing that 7 (THi" $) I1(,,p), we have on
/ (WW x- - v) a(T on S)+1((z))) v)
do op
S) (T Φ0 (TT,,
bn ) (111 x(S
=(v) x(p) (V dy(2,9)
2 Wo( p) = (UP) Id).
(135)
In the last line, used Proposition 15, iee that the optimal coupling is of the form (U%," Id) Id)# #p- L
H.3 Proof of Proposition 3
Let Te+1 Te T(V W. F(me), T- Id) (P(N)) + do., (T, (d) Applying the threc-point inequality (Lemma 29) with (T) - (Vu,)((s) Id) (2())) which is convex, To Id and T* Te+1, we get for all T (13(H)).
T(V,,F(m)) T (2) donn (T,
(136)
2 F(ue), Tx41 Id) Id) do (Tx41, Id) + don (T, Tet1),
which is equivalent to
1 (Vw, Tht1 Id)
5 don, IIk+I,
(137)
(Vw,F((mk) T Id)) + don (T, -do.. (T, (T+1).).
By the 3-smoothness of F relative to dhu.. we also have
de up
(T*+1, = Fr (T++1) - Frg (ld) V(,F)(u), Tk-1 - Id) S 3don (Tk+1,
- F (T++1) < F (Td) ((u,F( Ta+1 Id) (z(x) + Bd (Tk+1,
(138)
Moreover, since 3 < J. this inequality implies (by non-negativity of don,)
F (T++)) Frg (Id) (V (Vw, Th+1 Id) (2)))) + 7
dor, (T
(139)
45

---

Then, using the inequality (137), we obtain for all T E (2()),
(Tx+1) S (Td) + (V., T- ca(e) doon (T, Observing that Fu (Tk+1) F(ux+1) and Fr (Id) = F(ue), we get
-dom (T, Tk+1)-
1
7
(140)
F(Ux+1) S F(()) + (wn, - -ddd) L2( + _dou (T,Id) don (T, Tk+1). (141)
Finally, setting Id, we obtain the result:
F(ux+1) S F(ue) 1 dosm (Id, Ta+1).
(142)
H.4 Proof of Proposition 4
Let E P2(RR), and T = argmin T. døne (T,Id) Id). From the relative convexity hypothesis,
we have
d (d) 2 adon, Id.
- Fuz (T) Fiue (Id) - (Vw, T (d))) T-I)
2 ad (T,Id)
(143)
F adon (T, 2 F,,, (Td) + F(wx) Td) (P()))
Flv)= - adoi (T, 2 F(4) + - (d))
Plugging this into (140), we get
F(utt)sF(U)) (d (T, Id) dø (T, Tx+1)) adon, (T, Id).
(144)
Then, by definition of T, note that dæ.,, (T. Wo(v, uu), and by Assumption 1, we have dony (T, TR+1) 2 Woly, 1k+1), since Tw = v and (TR+1) = p441. Thus,
- (--o) Wolu, 1 WolMMuz41
(145)
Observing that F(u)) S F(we) for all < k (by Proposition and non-negativity of dg for á convex) and that W (V, (u) 2 0, we can apply Lemma 30 with f F, F(v) and g W .V,,.)) and we obtain
1 a S Wo(p,
V 2 1, F(px) - F(v) 1
a
F Wolu,po). k
(146)
-1
For the second result, from (145), we get for = 1² the minimizer F. since (114411-FF(( 2
Wolm",u44 I (1 aT) Wolp", s (1 ar) h+1 Wolu",
(147)
46

---

H.S Proof of Proposition 5
Let k 2 O, by the definition of dan- and the hypothesis dø;; ((w,F())))) TK+I, F(ux)) -
Adz (Id, Tk+1). we have FK
ohip (Vw,
o ((px))
+ (Vh" o F(uu)) F(uxt)) o Tat1 - Vwa F(a))
+ day ((T+)((())))))))) o TR+I, Fw,4(1)))
the (Vw,F(uk))
+ (Vh" o
Vw, F(u+t) Tet1 F(ux)) (((N))
+ sd (Id, Tx+1)
pi (Vw,F(uk))
+ (Vh" Vw, F(us), Vw, F(u+t)) T841 Vw,F(ux) LP(H))
d= xk (Id, Tx+1),
(148)
where used in the last line that }/2 and the non-negativity of the Bregman divergence since F is convex along t ++ ((1 (1T+)TT + tId) # j and thus by Proposition 13, df (Id, Tx+1) 2 0.
Let T E (((s)). Then, using the three-point identity (Lemma 28) (with S Id, and T= Tk+1), and remembering that Tx+1 = ld - Vw,F(ux) we get
dz,,, (Id, Tk+1)= = =ds,, (Td, T) - dz,, (T
- -VV,, k+1 ) #1*) Tk+1, Tk+1)
+ (V,, # mk)) T, ld (
dr (Id, T) dj.,, 1/. (T)
+ F(V /b)) o T ((e+t) Te+1, -Te+)
= dr,, (Id, T) - -dz (T
+
T Vw, o Teti, Vh Vw, (2()))) (149)
This is equivalent to
Wo, Vw. T Vw, + d (Id, Tx41)
1 d (Id, T) -din (T)+1,T)
1
7
+ F(T /12)) o TT Vw, Vhi o F(m)))
(150)
Then, using the definition of dan- (Vw2 o T, V Vw F(uu), we obtain t4
(Vh" o Vw, FWF+()))1 Te+1 F(r)( + T
_ds (Id, Te+1)
1 d (Id. T) d (T
- -day- (Vw,F(TyHu) o T, Vw, F(u)) + oie (V o TT) - hh (Vw,F(uk))
(151)
Plugging this into (148), we get
1 dd (Id, T) - -i T)
Diin, (Vw, S of (Vw, o T) +
7
-aay (Vw, T, Vw, F(x)). (152)
47

---

For T = Id. we get
Spins (Vw, F(u++1) S (
(Vw,F(ux)) 1 -ds,, (Tx+1,
(153)
5
H.6 Proof of Proposition 6
Let μ E P2((R)) be the minimizer of of 2 and T (ug d (Id, T).
Pg
First, observe that since μ= is the minimizer of F, then Vw,,F(n") 0 (see e.g. [74, Theorem 3.1J), and thus d (0) h² (0). Moreover, it induces that dg (Id, F(us) - F(u*) and dz,, (TR+1, T) = F(u++1) - F(u*).
Therefore, using (152) and the hypothesis ad, (Id, T) S døs (0, F(u))) we get
1
L
o/ (Vwx)(( - (0) -
(Id,T)
(Tx+1; T) - -da; (0, Vw,
7 Fm
7
1
d Pm
(Id,
(T&+1; T) - -ads,, (Id, T)
-
1
dsn (Id, T)
1 -di. (Tx+1, T)
1
a) (F(m)) -F(x))) (F(uk+2) - F(p)))
(154)
Then, applying Lemma 30 with f Φ.¹ VW,FF (which satisfies @/ (Vw,F((ktt1) s
th ((w,)( by Proposition 5) c = h- (0) and g = F(-) - =(H²) 2 0, we get
1 i = (F(1o)-F(ux)) (155)
oi (Vw)F(uk)) - (0) S
a
(F(uo) - =((())
< k
Concerning the convergence of F F(ux), if a >> 0 and h" attains its minimum in 0, then necessarily A" (T) 2 (0) for all E P2 (R and T E (H). Thus, using (152), we get
0 8/"" peas (Vw - (0)
de (Id, (d,T) I (F(ex) - F((x())
1 T dg. (Tæ41, T) das. (0, ((au(()))
(((m)+1)FF(14
add (Id,
1 (- -a a) (F(H)) - (F(()) -F((") (F(16+1) - F(n)))
(156)
Thus, for all & 2 0,
F(p++1) F(H") (1 -aa (F(ue) (p((us)-aa - F(w"))
(157)
$ (1 Ta) k+1 (F(Ho) - F(H))).
H.7 Proof of Proposition' 7
Let (RR)). Since Fr is Gâteaux differentiable, we can define its Bregman divergence.
For the first point, d² is 8-smooth relative to F% along ((1 w, F(,) IV, F(T T) a n. Thus, by applying Definition 3 for s = and = 0, we have
d day- (Vw,F(THH1) T, F(w)) S Bd (Vw,F(THH) o T, [w,F(p))
(158)
48

---

Using Lemma 19, we finally obtain
døn (Vw:F(TT)) I, F(ue)) - 5,ddg (Vw, T,, Fw,))())
(159)
= =d (Id,T).
which is the desired inequality. The second point follows similarly.
H.8 Proof of Lemma 11
Let us define G2= (p) x Rd Rd
as for all T E (2),,, x RS.
0 of
(T#
(T,
8v du
(T, Fw,
(160)
6
(T (T)
G (T,
8mg 8U
with for all i, G (2) x Rd +R, (T(,, = dz; ou (T)) Using the chain rule, for all Rd.
ds (T.. = (T. Tx())),
dT.
+ (T., T.(T)),
dT's
(1:)
ds
ds
r.(())
(161)
On one hand, have V2G (T., Tx))) V a 6F ((Tx) (x)). On the other hand, let a &J (T) du; oy and (v) = ds; og (x). Since J. (T) g" (T (1), applying Proposition 1, we know that
de; on
us compute V x)). First, we define the shorthands 3 (T) = (T, a 85
VIG,(TT = Vg (T) = Vwag" (T (1)))
Now, let compute Vw,g" hypothesis that bn bu
v (v). Let x be such that f dx 0. then using the by 8 and the definition of g... ofz
v 85
v
ʃ óg (v) dx dx=
/ /
θ 8
ax, 8μ2 (v) v) dx(x).
(162)
8p
ET Thus, Vw.gz V dz. on (z)
Putting everything together, we obtain
dõ: #
i F
dT.
(T,, Tx(x))
(T.))))
ds
Dr, 8n2
ds La
8F
dT.
((T.) (T.(())))
ds (1))
Dz. 8μ
(163)
á F²
dT.
V
((T+) (T. Ts
ds (v) du(z)
Dz, 8p2
0 6F
dT. ds (a)
(TT.))))
Dz; 8y
and thus
(T,, T,(x)) = / V. ((T.) (Tx(2), Tx())) dT.
82F
(v) du(s)
ds
2p2
ds
2 8F
dT
((T.) (T
(z). (164)
ds
49

---

H.9
Proof of Proposition 12
First, recall that by using the chain rule and Proposition 1, F(m)) = (Vw,)(1)) o T,, dt
T de
Thus, since
d²
( (w.F(im) o Te, dt,
F(m))=
dta
dt
dt (2())
(165)
= -(1 dl
dT,
(Vw,F(m)) ( Te), dt (M)
dt
By Lemma 11,
di3 dr2 F(u)) = /1
d²
V V 85F
((T)) Tx())) dT, (v), dT,
dt (m) ap(s)du(()
5n2
dt
+ /
V2
(TT)))))(T(()))
dT, dt (c), di, dt (m) dp(x)
8
((T))# T(())) o(0)) du(z)du(x))
V
(166)
/
8F
+
42
(TT)) du(x)
ōp
V, V,V, 8
((T.) (T.(a) Tx du(u)
5n2
V 8F
(TTT) du(x))
5y
H.10 Proof of Proposition
1. (cl) = (c2). Let t > 0, t2 E [0, 1].
Flpt S - (1F(FFF((TT])
+ +
F( - ((TT))+1)) {(((Te)))) F((Te.)#H) (167)
Passing to the limit t - 0 and using Proposition 1, we get (Vw,F((T,) +1)) o Te, Tes
Te, (2)) S F( #4) - F(TT.,)*))
2. (c2) (c3). Let f1, 12 E 10. then by hypothesis,
(V F (V,,F(( (((T)) T. T2 TT)) ((((Tt))#)) F((T)) (168)
(Vw, ((TT)) Ti, Te, Te)) S F((( )+1))
Summing the two inequalities, we get
(V #1) Te, Vw, #H) Te, Te, TeileA()) 2 0. (169)
3. (c3) = (c4). Let t1, to []. First, we have,
d² Flil di = dt F(n -12
d
dt. F F
de
(VW,, (og,F((TT)) ((TT((
Tr, -w, To, (170)
Th, - Te,)
20
Let c 10 and define t + 1/f = ph the interpolation curve between +H and (Te, + e(T Tr.)) y u. Then, noting that Th, c(T - Tr,) = To)to(1-))) so
50

---

o9 Pet p
56 46 +02--)) and we have that Pt
L de
d2 F(r)) dt 2
(171)
Moreover, by continuity, dr= F(())
d²
c0 dF F ((TT,))
d² F(ue). Then, since
de + ara F(r)) is continuous on [0. 1]. it is bounded, and we can apply the dominated convergence theorem. This implies that for all t E [0, 1].
d²
-im {
da F(u)) dt 0.
Hess F= dt F(m)
(172)
l=t
dt
4. (c4) = (cl). Let Let 1, E [0, 1] and c(t) F(it 22) for all t E [0, 1]. From [125, Equation 16.5),
Vt e ((() (1 ()) tp(() L' d2
(T) ds,
(173)
where is the Green function defined as G(s, t) s(1 -01 s(1-t)1(ss)) + + ()s.)) 2 d² [125, Equation 16.6]. Then, dr2 F(me) implies that fo de (s) ds 0. and thus
(1) = Fiiti $ - (1-1))(()) + (11 = (1)))(((T)))))))))))) + +F ((Tea))))) (174)
H.11 Proof of Proposition 24
Let J(T) douz (T, (VW,,9(m.) T Id) H(( (T)))))) Taking the first variation,
we get
vJ((T++)= Vouu(TTTH Von. (ld) + (V + Vw, ((Txx+ #1x))))))))))))))))))))))))))# Tk+1)
(TH + 7 Vw, H ((Tx+)))mM) o Te+1 - (V (Id) - -VWss
Vonu(TTT)) + vw, H (T Tat1 1
(175)
Thus,
VJ(TR+I) =0 + T2+1 argmin dø (T, + -TH(T (TT)
(176)
TEL"(p)
Now, we aim showing that T8+1 Tk41 S4+1 or
min min don, (T, HH #(x)))))= TEL2(VE+1) TEL((p))
demnn (T, + #+++1) (177)
First, by the change of variable formula, since on is pushforward compatible, observe that for T T (1 dow... (d) + +H(TTu++1)) = darg oTo Sx+1) +H((T Sk+1)#1K))
Since (To T E L3 (vK+1) C L2 (px), we have
min denea (T, TH Pk41) min TEIP(m)
dong (T, Sx+1) + TH(T (178)
(e12(P)11)
By assumption, Vx+1 E (RR)). Thus, applying Proposition 15, there exists Eve-1 such that (T 4/+11 = 1k+1 and TH = argmin a TT domn (T, (T. and thus Pu--1
dovee (T Id) = W (o(44+11, (k41).
By contradiction, suppose that
min derz. (T, + + TH( (b++1) > dop, (Tx+1, S441) + +T(((T+)))))) (179) TEL7(rE41)
On one hand, have (T S4-14 (T 41/411 pk+1. and there- Bezzi fore H((T Sk+1)#H4) = H(mx+1) = H((Tx+1)41s
On the other hand,
(T Sk41)#H E (114641194443)) and thus
dens (Te+1, 2 Wo(peti, P641) = dores 0er41
(T Id).
(180)
51

---

Thus,
min domaas (I, + TH #1++1) > does (Tr+1, Sk-1) + T (Te+1)#uu) ( (E()()))))
(181)
2 d
To Idd TH TH((TT Pk+1).
Beeti
But T E (vK++1) so this is a contradiction. So, we can conclude that the two schemes are 0us41 equivalent, and moreover, T T T44 SK+1-
H.12 Proof of Proposition 25
Let (TT) 7 ((V w, (1x)), T - Id) ((())) + H(T+/14)) Since Hy, is convex on T22())), $ is convex, and can apply the three-point inequality (Lemma 29) and for all T E (P(x).
+(H((THuk) + (V - Id)))
+ done (T,Id)
2 =((H(u0+1 + (Vw,9(us), Th+1 (d) + domz (Tx=I, Id) + dome (T, TRtT), (182)
which is equivalent to
H((x+1) + (VwaG(mu), Tst1 (2))) - dop (Tx+1,
S <H(Twuk) + (Ww,9(ux) T Id) (P(ue) + _d (T,Id) - _o (T, (T,T)+)). (183)
Since 9.m is 8-smoothh relatively to opp along t + ((1 - -JIld (TTx+1) Mk: and r S 1 g. we also have
g(uE+1) S 9((u)) + (V W (wmx) k+1 Td) (P( + Sdon (TK+I,
S g((x) + Ww, ((xx) Txt1 (2)o) + -don (Tk+1,Id).
1
(184)
7
Thus, applying first the smoothness of 9 and then the three-point inequality, we get for all T e I.2 (px), H(pk+1) + (((x+1) & H())++1) + G(m)) + (V, Tx - (2( + don, (Tx+1,
1
s *({Tpkk + +((x)) + ( Ww, 9(wx) T-Id)) + don, (T.
1 donz (T, Tk+1).
(185)
Now, let (R2( and THV argmin Tum don, Id), and suppose that 9p,. is a-convex relative to due along ((1 t)ld + =TTthh Mr.. Thus,
dg (TuE Id) > adoz (TuRi Id) Quk
g(v)= adone (Tras Id) 2 G(ux) + (Vwgg(uu), Tuank [2) (186)
Plugging this into (185), get
1 F((++1) S H(v)) + G(v) - cadens (Tyh." Id] + do (Tiin" Id) doo, oue Tk+1). (187)
1
Now, note that domg (TuAN III, Wolu, and by Assumption 1, daps (Tum Tht1) 2 ONE
Wo(n, (k+1). Thus,
F(ux+1) F(v) & (--a) W Wo(k, WW (V. pk+1).
(188)
Using T= Id in (185), we observe that F(us) S F(ue) for l < k. Moreover, Wolu, ux) 2 Thus, applying Lemma 30 with =F F(v) and g = Wo( we obtain
1 W Wo(u, s "Wo(u.oo)
Wk 1, F(px) F(v) 5
(189)
b
52

---

H.13
Proof of Lemma 26
First, VV- bijective. Thus, we only need to show that VV VVU is injective. Take 2 = V rU.
Since U is 8-smooth relative to V, we have for all 2, y.
U(x)= s U(w) + (VU()) - y) (dv
(190)
which is equivalent to Moreover, by definition of dy.
- -V()) S U(() + (VU()), v) + Bdv (z. v).
(191)
V V(s) - (VV()) - v) - (x.
(192)
Summing the two inequalities, we get
V (v) - -VV x - - (U()
= V(x) - (v) (VV()) -VUUV), 1) - -1- -11-3333 (193)
+ +7.dv(z.z)
This is equivalent to
u(y) S u(s) (u((y), x y) - -(1 -TB)dv
(194)
and thus with u being (1 (1- (1-r9)-coneel
relative to V (for TB < 1). For TB < 1, it is equivalent to (1 - T convex, i.e. (Vu(x) - Vu(y), - 9) 2 (1 TB) ( (z) VV (y), y) 2
Since V is strictly convex, Vu is injective.
Moreover, det VT det (V (VV VU)) det V 0 because on hand is (1 - Br)--covex relative to V which is strictly convex, and on the other hand, V is also strictly convex.
To conclude, applying [5, Lemma 5,5.3), Tyj is absolutely continuous with respect to the Lebesgue measure.
H.14 Proof of Proposition 27
On one hand, H. is 1-smooth relative to H. thus we only need to show that μ + J Vdp is smooth relative to H. Using Proposition 13, need to show that
d2 du
1 ト/2
V(e)
2 (T++1 -di Id)" -ld) duc < B die
H(u)).
(195)
Recall from (61) that T (x) ((1 - -E Sr + + x cst, thus VT is constant. Using the computations of [43, Appendix B.2].
d H(m) = ([VT,- VT - -Ta)-
(196)
dt
Assuming (1 - T)) Ex + +14+1 Σ-1 2 0, Tk+11 is the gradient of a convex function and (4e is
Wasserstein geodesic. Thus, by [43],
d2
dta M(m) 2 1 llop
[ - Id Idli
(197)
Moreover, by [33, Lemma 10], / r ||Eillop convex along generalized geodesics, and thus Ey. 3 Ala and 1op [43]. Hence, noting max (M) the largest eigenvalue of some matrix M,
d²
M((m)) 2 1 |ITE - Idll 2 do (V2V)
(TA+) -
Id)dux
d
2
d²
Aamax (V2V) di2 V(u).
(198)
53

---

From this inequality, we deduce that
max
V2V)
(T)+ max (22)
(H(px+1) H((x)
(Vw, H(Hx), Tx+1 - -Id)
max (V
/(1- -1-t) t) H(m)) dt
d?
(199)
2
2 /
d²
V(Hu)(1 - dt
dt²
dp (T (Tx+I, Id).
So.
d Frz (Tk+1, Id) = drpe (T&+1; 1, Id) + din (Tk+1, Id)
Ao max )
(200)
+
dapp (Tx+1,Id).
O
I Additional results
L11 Three-point identity and inequality
In this Section, we derive results which are useful to show the convergence of mirror descent or preconditioned schemes. Namely, we first derive the three-point identity which we use to show the convergence of the preconditioned scheme in Proposition 5 as well as the three-point inequality, which we use for the convergence of the mirror descent scheme in Proposition 3.
Lemma 28 (Three-Point Identity). Let Φi (1) IR be Gâteaux differentiable. For all S,T,U U (I)),, we have
do U) = (S, T) + dT,T, U) + (V L) - -Vo(U),STT)pau))
(201)
Proof. Let E ( then using the linearity of the Gâteaux differential,
- T) - do(T, = <(S) <(U) (Vo(U), - U)
- -o(S) +SS O(T+ (V(()) T)
T- U)((())
= - + -VooT), - (202)
(VO(U),T-U))uu)
(Vo(T), ST)T)) (V)()), T))
Lemma 29 (Three-Poir) Point Inequality). Let μ E P2(R"), To (2) and on (p) E. convex, and Gâteaux differentiable. Let vp T22(z) T be convex, proper and lower semicontinuaus. Assume there exists T = argminTEL de To) + (T). Then, for all E T.2(M).
(TT) + do, (T, To) 2 dh (TT) + dæy (T*,To) To) + dø (T,T").
(203)
Proof. Denote J(T)= dd. (T, To) + (T). Let T argmin J(T), hence OJ(T*).
Since and s are proper, convex and lower semicontinuous. and T + de, (T, To) is continuous (since ou is continuous), thus by [99, Theorem 3,30), aJ(T*) = (T*) + adg (.. (T*).
Moreover, since oμ differentiable, Odg, (T) = (Vrd,, (T)) = (T*) (T))), and thus (To) (T E @T(TT)
Finally, by definition of subgradients and by applying Lemma 28, get for a E L²(H),
(T) 2 (TT) - = +((T)) de. (T, do. (T,T") + do. (T", ,
T - - (Vo, - -
(204)
54

---

Remark 3. Actually we can restrict ab to be convex along ((1 t)T + +TT) n In that case, dø (T( T*) = (T) (TT)) T2) 2 for 4 E (T*(T) (by Proposition 13) and we still have (Tb(TT) od (s (T)) (T(T)) (see (99, Theorem 3.30]) so that we can conclude.
I.2 Convergence lemma
We first provide Lemma which follows from [88, Theorem 3.1], and which is useful for the proofs of Propositions 4, 6 and 25.
Lemma 30. Let I X + R. g: X + R+ and KENN a sequence in X such that for all k 2 1, [(x)) < [ [x*--1) Assume that there exists 3 > 2 O, c E IR such that for all k 2 0,
[(*++1) - e S (B- (x) ()g(())) - :g( 11) ) then
Vk 2 1, f(xx)) - -
g(00) S k g(wo).
(205)
8
3-a
-1
Proof. First, observe the f(xx) S fZze for all l $ k. Thus, for all k: 2 1,
h 8 ) (f(xx) -c) -
B
Σ
Σ
(f(xz) -c))
B α
f=1
1=1
8
((3 -a)g((ee)) 3g(xz))
3 0
k-1
3
k
B
(206)
3
Σ
g(xx)
-3)
2
g(re)
B - a
t=
8=1
8 k
= 3g(2o)
g(0x))
3 - a
S Bg(zo)
since g 2 0.
Now, note that El (3
<
B-a since (1+
8-a
a
b
2 1 + k B-a
8-
8-0
(1 a-a
(by convexity on R of I
a1+x)
Thus,
B
3 - a
f(xx) -
g(0o) =
g(0o)
< g(zo).
(207)
Σk
===1
3 J-a
-1
k
1.3
Some properties of Bregman divergences
We provide in this Section additional results on the Bregman divergences introduced in Section 3. First, we focus on o (T) = f V T dp. The following Lemma is akin to [80, Proposition 4] which shows it only for OT maps.
Lemma 31. Let V idd + IR convex and (T) f du. Then,
VT, E do, (T,S) / dv (T()))) du(x).
(208)
Proof. Let T,S E (H)),, then remembering that Vw, V(()) = VV, we have de, (T,S) = (T) - (s) (VV S - -)-SSuaa (2(μ)
/ VoT-VoS-
(VV o (V 5) dp
(209)
= d da(x).
55

---

Next, we focus on (x(T)) - // (T(x)) - T(x)) du(z)du(a²), and we generalize the result from [80, Proposition 4].
Lemma 32. Let W : Rd + R even ((x)) w(-x)), convex and differentiable. Let @T( (T(x) - T(x')) du(x)dp(zz). Then,
1 vT,S (2), de (T, _ // (Tz))) S(x') dp(z)du(2²).
(210)
Proof. Let T,S (2)). remember that Ww, VW p, and thus Vw, W(Spp) o (VW * S S. Thus,
dø, (T, (T) 0p(S) - ((VW * SH1) T, 1 12(w))
2 // ((((x)) - T(x')) du(())du(z)) 2 /f
W(S(x)- (Sxx)) dx(x)dp(xx)
/((VW S*))(((x)) T(x) - -(z)) du(x).
(211)
Then, note that (- (=()) and thus the last term can be written as:
/ (VW + S41) T(x) - S(x)) dx(x))
= -// (VW) S(z')) T(x) S(x)) da(z)du(a²)
5 // (VW()(s) - S(x'), T(x) S(x)) dn(xrdp((z))
1
(VW(S(=))) - -())())))((s)),...)))))))))))>)) T(y) - S(w)) du(z)dn((z))
(212)
2 // (VW()(>) - S(z')), S(x) d)(x)dp((x))
2 (VW()()) -((()) T(x') S(z')) dz(z)da(zz)
-/2 // (SW)) - S(x')) du(z)du(zz).
Finally, we get
de (T.S) 1 // (W(T(x)- T(x)) - ((((2) - S(x'))
(213)
((VW - S(x')) T(x) - T(x') (S(x) - S(a'))) dn(r)da((')
1 //2 // dw ( T(E),, S(x) S(x')) dι(z)du(zz).
Now, make the connection with the mirror map used by Deb al. [41] and derive the related Bregman divergence.
Lemma 33, Let ou (T) W2 that Tp, Sqp E P2 (IRd),
#t, p) for ,, P E (R.)). Then, for all T,S T13()), such
1 (T, 2
.T TT
SAn oS - + (IS oS - T TT.n .T TVe
Spp oS) (?())
(214)
where T Tfn" denotes the OT map between T and
56

---

Proof. Let T,S E (p) such that t* Sæp E )R. Remember that Vw, (.(-))
Id T", then
de (T,S) (T( (x(s) (Ww, (( S,T - S) 1264)
W2(T #14, p)
W3(S #1, P) ((Td T -S,T-S) (S)
2
2
1 a
-T
T1
T Szn oS - -Slli2()) + +(Ts S0 oS oS S. -)
2
2
1
2 1 #
TT
T1
1 oS -S11i.(0))
(1))
+ (T * S S,T
T
S) (P(N) + (T'S oS S, T oS (p))
1 T AR
2 (n)
T Sgn
oS - (S-S|l2a(w))
2
2
(T SS
S, T
o$ L3
2 | Ta o oT T 2 (n)
+ T SaH oS SS
2
(T o - S, T
o oS
To T o L2)
(215)
(T oS T /
o T (3)
2 51TT mn TT To
oS - ([--S)[[Zany
(T (T
o S, T
T T
S) (()))
57

---

J
NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: In the abstract and introduction, we claim that we adapt mirror descent and the preconditioned gradient descent to the Wasserstein space, and that we provide guarantees of convergence. These results presented in Section 3 and Section 4 for respectively mirror descent and preconditioned gradient descent. We also claim that we illustrate advantages of such schemes on ill-conditioned optimization tasks and to minimize discrepancies in order to align single-cells, which we do in Section 5.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: On the two theoretical sections (Section and Section 4), we state all the hypotheses to obtain our convergence results. Then, in Section 5, we discuss some couples target functional F and Bregman potentiil//recoonditinnnee e for which we can verify these assumptions. However, in general, verifying these assumptions is a hard problem, stated in the Conclusion, and we leave for future works these investigations.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
58

---

judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: The full set of assumptions provided for each theoretical result, along with a complete proof in Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The information to reproduce the main experimental results are described in Section and Appendix G.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
59

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide a part of the code to reproduce the experiment on Gaussians and on interaction functionals in supplementary materials and in attps://githns cwt.coo/ccooo/o com/ccloonttMMirrrr Descent_ing Www.ewateen
http://rrrrrraaa.eeerdnoottteeed Gradientt.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
casssooonlyyyy for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed reproduce the results. See the NeurIPS code and data submission guidelines (https: //nips.c./////////o/////// wttpsstcccpppccccccccccccccccccc htpppulcccguuieeeeeeeeeepppee
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
All the training and test details for the experiments are provided in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
60

---

The experiment on Gaussians in Figure 2 is run over 20 randomly sampled objective covariances, and the results are plotted with the standard deviation. For the single-cell experiment of Figure 3, we reported the results with 3 different initialization for each treatment, along their mean. For the mirror interaction experiment, the results are deterministic.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to library function, bootstrap, etc.)
The assumptions made should given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative етгог rates).
If error bars reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We report in Appendix G the computer resources and approximate runtime for each experiment. Namely, the Gaussian and mirror interaction experiments were done personal laptop on CPU, and took only few hours to run. The single-cell experiment was performed on GPU as the data are of higher dimension with about 4000 samples, and took few hundred of hours of computational time.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /tttps:/nneurips.. cc hptpalic/tticssuieeeeee....e
Answer: [Yes]
Justification:
The research conducted in this paper is conform with the NeurIPS Code of
Ethics.
Guidelines:
61

---

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work is mostly theoretical and not tied to particular applications.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., dep.oymeet of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
62

---

Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The datasets used are properly cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcom thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information is available online, the authors are encouraged to reach out the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
63

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
64


Reviews:
[{"review_id": "8Yz7HxLmRV", "reviewer": "Reviewer_KGf7", "summary": "The paper provides convergence analysis for Mirror descent (with and without additional preconditioning) in the Wasserstein space. The analysis is performed for relatively smooth and convex functionals. The use of the studied algorithms on computational biology problem is illustrated.", "questions": "n/a", "limitations": "n/a", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "The theory in the paper seems sound and substantially extends the prior work overcoming some technical difficulties related to the discrete nature of the algorithms (explicit discretization) and non-Euclidean geometry in the Wasserstein space. Convergence rates are similar to those in finite dimensional space for MD.", "weaknesses": "The paper is very dense and seems like not well adapted to 9 page format. A lot of explanations and implementation details are referred to the appendix. It is very difficult to read the paper when it is written is such style. For example, although I have spent several hours only trying to understand how the algorithm is implemented, I did not manage to do it even with the help of the appendix. I would advise the authors to revise the paper to make it more accessible. \n\n\n1. Notations section should be improved. Many things are not properly defined throughout the paper, e.g., the coupling between distributions, what \"Id\" is?\n\n2. On line 138, what does it mean that the sets $V$ and $W$ are \"differentiable and $L$-smooth with $W$ even?  \n\n3. On line 142, the calligraphic $\\mathcal W$ is not defined. \n\n4. OT map mentioned on line 151 is not defined.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and time. We have addressed your comments point-by-point below. Please don't hesitate to let us know if you have any further questions.\n\n**The paper is very dense. I would advise the authors to revise the paper to make it more accessible.**\n\nWe understand your concerns about the paper's density and the difficulty in following the implementation details. Taking advantage of the extra page, we will work on revising the paper to make it more accessible and ensure that key explanations are clearer and more integrated into the main text.\n\n**How the algorithm is implemented?**\n\nWe will describe more clearly how the schemes are implemented in practice in a dedicated section of the appendix.\n\nIn practice, we use a particle approximation of the probability measures, i.e. we sample $x_1^{(0)},\\dots,x_n^{(0)}\\sim \\mu_0$ and then we apply the schemes to $\\hat{\\mu}_k^n = \\frac{1}{n} \\sum\\_{i=1}^n \\delta\\_{x\\_i^{(k)}}$ for all $k\\ge 0$.\n\nFor Mirror Descent with a Bregman potential $\\phi$ chosen as a potential energy $\\phi(\\mu) = \\int V\\ \\mathrm{d}\\mu$ with $\\nabla V$ and $\\nabla V^*$ known in closed-forms, we use equation (9) to implement it: for all $k\\ge 0$ and $i\\in \\{1,\\dots,n\\}$, $x\\_i^{(k+1)} = \\nabla V^*\\big(\\nabla V(x\\_i^{(k)}) - \\tau \\nabla\\_{W\\_2}\\mathcal{F}(\\hat{\\mu}\\_k^n)(x\\_i^{(k)})\\big)$. The Wasserstein gradient of $\\mathcal{F}$ in most of the cases considered can be computed in closed-form, e.g. for $\\mathcal{F}(\\mu)=\\int U\\mathrm{d}\\mu$, $\\nabla_{W_2}\\mathcal{F}(\\mu) = \\nabla U$, or for $\\mathcal{F}(\\mu)=\\iint W(x-y)\\ \\mathrm{d}\\mu(x)\\mathrm{d}\\mu(y)$, $\\nabla_{W_2}\\mathcal{F}(\\mu) = \\nabla W\\star \\mu$.\n\nIn the more general case, where the gradient of the Bregman potential cannot be inverted in closed-form, we used Newton's algorithm at each step, which is detailed in Appendix I.1, and we use equation (191) to implement it. \n\nFor the Gaussian experiment, we update directly the means and covariances. The closed-form of NEM is reported in equation (197), and the closed-form of the Forward-Backward scheme are in Appendix H.4 in equation (160). \n\nFinally, for the preconditioned gradient descent scheme, we implemented it using equation (11): for all $k\\ge 0$ and $i\\in\\{1,\\dots,n\\}$, $x_i^{(k+1)} = x_i^{(k)} - \\nabla h^*\\big(\\nabla_{W_2}\\mathcal{F}(\\hat{\\mu}_k^n)(x_i^{(k)})\\big)$. In the experiments of Figure 3, we considered $h^*(x)=(\\|x\\|_2^a + 1)^{\\frac{1}{a}} - 1$ with $a>0$, and its gradient can be computed in closed-form or by backpropagation.\n\n**Notations section should be improved. Many things are not properly defined throughout the paper, e.g., the coupling between distributions, what \"Id\" is?**\n\nWe will complete the notation section. A coupling between $\\mu\\in \\mathcal{P}\\_2(\\mathbb{R}^d)$ and $\\nu\\in\\mathcal{P}\\_2(\\mathbb{R}^d)$ is a probability distribution $\\gamma\\in\\mathcal{P}(\\mathbb{R}^d\\times \\mathbb{R}^d)$ such that $\\pi^1_{\\\\#}\\gamma=\\mu$ and $\\pi^2\\_\\\\#\\gamma=\\nu$ with $\\pi^1:(x,y)\\mapsto x$ and $\\pi^2:(x,y)\\mapsto y$. The $\\mathrm{Id}$ is the identity mapping $x\\mapsto x$.\n\n**On line 138, what does it mean that the sets $V$ and $W$ are \"differentiable and $L$-smooth with $W$ even?**\n\n$V$ and $W$ are functions from $\\mathbb{R}^d$ to $\\mathbb{R}$. $V$ (or $W$) $L$-smooth means that its gradient is $L$-Lipschitz. We will clarify this in the revision of the paper. To ask for $W$ even means that $W(-x)=W(x)$ for all $x\\in \\mathbb{R}^d$.\n\n**On line 142, the calligraphic $\\mathcal W$ is not defined.**\n\n$\\mathcal{W}$ is defined as $\\mathcal{W}(\\mu) = \\iint W(x-y)\\ \\mathrm{d}\\mu(x)\\mathrm{d}\\mu(y)$, see line 127.\n\n**OT map mentioned on line 151 is not defined.**\n\nAn OT map between $\\mu,\\nu\\in\\mathcal{P}\\_2(\\mathbb{R}^d)$ is a map $T:\\mathbb{R}^d\\to\\mathbb{R}^d$ satisfying $T\\_\\\\#\\mu=\\nu$ and $W_2^2(\\mu,\\nu) = \\|T-\\mathrm{Id}\\|_{L^2(\\mu)}^2$. It exists e.g. provided $\\mu\\ll\\mathrm{Leb}$ by Brenier's theorem. We will clarify this in the revised version of the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ctFgyfcs2V", "reviewer": "Reviewer_3j2f", "summary": "The paper generalises mirror descent and preconditioning approaches from optimisation over R^d to optimisation over Wasserstein space, which is applicable in some actual problems, and confirms by the corresponding numerical experiments the efficiency of these approaches.", "questions": "There is an aspect which was not mentioned in this paper, but is interesting from the theoretical point of view: stochastic oracle. The generalisation of stochastic mirror descent to Wasserstein space seems to be the demanding yet possible task, when that of preconditioning is challenging even in R^d case, but it would be interesting to find the proper notions of stochasticity and flexible enough analysis to deal with it in case of Wassestein space. Is your current analysis generalisable to stochastic optimisation and to inexact model of function, more generally?", "limitations": "The limitations are clear from reading.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper completely corresponds to its stated goal, uses the most actual theoretical frameworks of long known approaches from optimisation over R^d, and focuses indeed on the most fundamental methods, which will allow a follower to develop any desired extension based on idea from optimisation over R^d and the detailed examination of technical issues related to the transition to Wasserstein space made in this work. Practical contribution of this paper is also notable, as it demonstrates in numerical experiments, but I guess that thanks to the proposed theoretical basement, many improvements can be made on top of that, which is to be done in future works.", "weaknesses": "The preconditioning is considered in some particular yet practically important case, but the complete generalisation of preconditioning is still to be finished.", "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.\n\n**The preconditioning is considered in some particular yet practically important case, but the complete generalisation of preconditioning is still to be finished.**\n\nWe acknowledge that the complete generalization and study of the preconditioned gradient descent on the Wasserstein space is left for future works. We note that the convergence results would still be valid for $\\phi_\\mu$ strictly convex, and differentiable as described in the first paragraph of Section 4. \n\n\n**Is your current analysis generalisable to stochastic optimisation and to inexact model of function, more generally?**\n\nWe expect our analysis to be generalisable to the stochastic setting. While we leave it for future works, we note that [1, 2] studied stochastic Wasserstein gradient schemes, while [3] studied Mirror Descent on $\\mathbb{R}^d$ in the stochastic setting. Thus, it could be a nice avenue of future works to extend these results to Mirror Descent on the Wasserstein space.\n\n\n[1] Diao, M. Z., Balasubramanian, K., Chewi, S., & Salim, A. (2023, July). Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein space. In International Conference on Machine Learning (pp. 7960-7991). PMLR.\n\n[2] Backhoff-Veraguas, J., Fontbona, J., Rios, G., & Tobar, F. (2022). Stochastic Gradient Descent for Barycenters in Wasserstein Space. arXiv preprint arXiv:2201.04232.\n\n[3] Hanzely, F., & Richtárik, P. (2021). Fastest rates for stochastic mirror descent methods. Computational Optimization and Applications, 79, 717-766.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LEme6umRhq", "reviewer": "Reviewer_mqrt", "summary": "This paper presented a unified view on the functional optimization on the Wasserstein space. Authors proposed a mirror descent algorithm and a preconditioned gradient descent algorithm, both are applied to minimize the objective functional over the Wasserstein space. Many existing algorithms in the Wasserstein space-related literature can be recovered or seen as a special case of the algorithms proposed by the authors.\nConvergence guarantees were provided and backed with numerical results.", "questions": "some minor grammatical errors such as l.158 ($T_\\eta^\\nu$ OT maps from $\\eta$...)", "limitations": "The authors did address the limitations of this work.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "Mirror descent and preconditioned GD are closely related and have been seen as generalizations of GD. Though Wasserstein GD have been proposed, the extension to MD and preconditioned GD is a natural step.\n\nThe formulation in the paper is general and can be applicable to a wide range of studies in both machine learning application and theory.\n\nThe literature review provided by the authors seem to be thorough, and the position of the current manuscript in the literature is well explained. The relationships between the proposed algorithm and existing algorithm were also clear throughout the paper.\n\n\nThe reviewer did not read the full analysis in the appendix, but the technical aspects of this paper seems solid and intuitive.", "weaknesses": "The main concern I have is as presented in the MD optimization literature: what new tool does the MD formulation bring to the table? Previously, the convergence of MD is similar to or worse than GD. The literature on relative strong smoothness and convexity enabled MD to converge even when the objective function is not smooth/convex in the Euclidean sense (which was really mostly KL). The function properties in the Wasserstein space are more complicated and therefore difficult to align, this was briefly discussed by the authors at the start of section 5. But the application offered later is largely limited to existing applications KL or Sinkhorn, whereas these specific cases seem to have been already studied by the literature anyway.\n\nSome sections are not well presented. For instance, the subsection starting at l.136, the statements on the potential energies are confusing and it is very difficult to associate on sentence with either the previous or the latter sentence (such as \"... specific example. In particular, we have xxx. We will also consider...\", is the \"xxx\" with the specific example of what the authors consider?) Similar confusions also exist in this paper, but to a lesser extent.\n\nWhile the contribution is clear from reading the entire paper, the \"Contributions\" paragraph does not help reader and clarify what this paper offers.\n\nThough discussed briefly, the computational complexity and real-life run-time of these type of methods should be better addressed both in theory and experiments, especially on the comparison between the proposed and benchmark methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback and time. We have addressed your comments point-by-point below. Please don't hesitate to let us know if you have any further questions.\n\n**What new tool does the MD formulation bring to the table?**\n\nFirst, the Mirror Descent algorithm on the Wasserstein space allows to perform optimization over probability distributions on constrained spaces (and in particular sampling on constrained spaces for specific objectives). This motivated the development of e.g. Mirror Langevin [1,2,3]. This lens on Mirror Descent on the Wasserstein space was also studied e.g. in [4] to optimize the KL, but without a theoretical analysis of the discrete algorithm.\nHere, we chose to study the convergence of these algorithms, and one of the main contributions of this work is to define appropriate notion of relative smoothness and convexity. This differs a lot from previous works, and is new on the Wasserstein space to the best of our knowledge.\n\nMoreover, we also proposed to use general Bregman potentials, instead of just considering potential energies as done in most previous works focusing on Bregman divergences and Mirror Descent on the Wasserstein space. We believe that this additional flexibility could allow to find new pairs of functionals/Bregman potentials, with theoretical guarantees for the convergence of the Mirror Descent scheme. \nFor instance, a convex functional is 1-smooth and 1-convex relative to itself, and thus could be used as Bregman potential for its optimization, provided we can compute the scheme. We note that this point of view was also recently taken to optimize the MMD in [5], while the MMD is known to be non convex w.r.t the Wasserstein distance [6], which leads MMD flows to converge to spurious local minima in practice [5,6,7].\n\n\n**The function properties in the Wasserstein space are more complicated and therefore difficult to align, this was briefly discussed by the authors at the start of section 5.**\n\nShowing the relative smoothness and convexity of functionals is indeed not straightforward in general. In the first paragraph of Section 5, we described some cases where we can show it leveraging the relative smoothness and convexity on $\\mathbb{R}^d$, and suggested in the general case to compare the Hessians. \n\n\n**But the application offered later is largely limited to existing applications KL or Sinkhorn.**\n\nWe did not focus only on minimizing the KL or Sinkhorn. On Figure 1, we applied the Mirror Descent scheme on an interaction energy with an interaction energy as Bregman potential, in a case where it is well relatively smooth. \nWe focused on the KL for the Gaussian example of Figure 2, but this specific example has not been yet studied in the literature as we use Mirror Descent with negative entropy as Bregman potential.\nFinally, on Figure 3, we applied the preconditioned scheme on different divergences such as the Sliced-Wasserstein distance, the Sinkhorn divergence and the Energy distance. While Wasserstein gradient flows of these divergences have already been proposed and studied, applying the preconditioned gradient descent to them is new to the best of our knowledge.\n\n\n**Some sections are not well presented.**\n\nWe apologize for the confusions and will revise the paper to improve its clarity. \n\n\n**The computational complexity and real-life run-time of these type of methods should be better addressed both in theory and experiments, especially on the comparison between the proposed and benchmark methods.**\n\nCompared to the usual Wasserstein gradient scheme, the preconditioned Wasserstein gradient scheme only requires to evaluate additionaly the gradient of $h^*$, and thus has the same computational complexity as the Wasserstein gradient descent. Nonetheless, we noted in Figure 3 that it required less iterations to converge for a well-chosen preconditioner.\n\nFor Mirror Descent, the complexity depends on the Bregman potential. If we use $\\phi_\\mu(T)=\\int V\\circ T\\ \\mathrm{d}\\mu$ with $V$ a potential which we know in closed form, and from which we can compute $\\nabla V$ and $\\nabla V^*$, then the computational cost is the same as the Wasserstein gradient descent. In the more general case, where we do not know how to invert $\\nabla\\phi_\\mu$, we need to use the Newton method, which is more costly as it requires to invert a Jacobian of size $nd\\times nd$ as stated Section I.1. \n\nThe runtimes are reported in Appendix I. For instance, as stated line 1490, for the experiment of Figure 1, the Mirror Descent used in dimension $d=2$ for $n=100$ particles and $120$ epochs with $K_2$ and $K_2^\\Sigma$ took about 5 minutes while with $K_4$ and $K_4^\\Sigma$, it took about 1h. We leave for future works the development of cheaper optimization methods.\n\n\n**some minor grammatical errors**\n\nThank you, we will correct the grammatical errors.\n\n\n[1] Ahn, K., & Chewi, S. (2021). Efficient constrained sampling via the mirror-Langevin algorithm. Advances in Neural Information Processing Systems, 34, 28405-28418.\n\n[2] Chewi, S., Le Gouic, T., Lu, C., Maunu, T., Rigollet, P., & Stromme, A. (2020). Exponential ergodicity of mirror-Langevin diffusions. Advances in Neural Information Processing Systems, 33, 19573-19585.\n\n[3] Hsieh, Y. P., Kavis, A., Rolland, P., & Cevher, V. (2018). Mirrored langevin dynamics. Advances in Neural Information Processing Systems, 31.\n\n[4] Sharrock, L., Mackey, L., & Nemeth, C. (2023, January). Learning Rate Free Bayesian Inference in Constrained Domains. In NeurIPS.\n\n[5] Gladin, E., Dvurechensky, P., Mielke, A., & Zhu, J. J. (2024). Interaction-Force Transport Gradient Flows. arXiv preprint arXiv:2405.17075.\n\n[6] Arbel, M., Korba, A., Salim, A., & Gretton, A. (2019). Maximum mean discrepancy gradient flow. Advances in Neural Information Processing Systems, 32.\n\n[7] Galashov, A., de Bortoli, V., & Gretton, A. (2024). Deep MMD Gradient Flow without adversarial training. arXiv preprint arXiv:2405.06780.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MqSOj7vqvF", "reviewer": "Reviewer_PBU2", "summary": "In this paper, the authors endeavor to integrate the concepts of mirror descent (MD) and preconditioned gradient descent (PGD) within the framework of Wasserstein distances, from a convergence theory perspective. Initially, they provide a comprehensive background including, but not limited to, Wasserstein distances, Bregman divergence, convexity, and smoothness. With assumptions regarding $\\beta$-smoothness and $\\alpha$-convexity, they successfully demonstrate the iterative procedures for MD and PGD in the Wasserstein space. Subsequently, they conduct experiments focusing on sampling methodologies and single-cell dynamics.", "questions": "1. **About the Proof of Convergence:**\n   From the perspective of gradient flow, assuming we have a velocity field $\\dot{X} = v_t$, we can derive that $v_t = -\\Vert \\nabla\\frac{\\delta \\mathcal{F}}{\\mu} \\Vert_2^2$. In this context, does assuming $\\beta$-smoothness offer additional advantages beyond the conventional benefits?\n\n2. **Alternative Proof Approaches:**\n   Could the convergence location and rate be established from the perspective of a Lyapunov functional? This approach may provide a robust framework for proving stability and convergence.\n\n3. **Riemannian-based Approaches:**\n   Is it feasible to extend the proposed method to Riemannian-based approaches, such as Riemannian SVGD [1]?\n\n4. **About Higher Order Systems and Extensions:**\n   Is it possible to extend this approach to systems characterized by the Fisher-Rao metric [2]? Furthermore, considering that the continuity equation in Wasserstein space represents a first-order system, could this approach be adapted to Hamiltonian Flow [3], which is the second order system to my knowledge?\n---\n### References  \n[1]. Zhang, Ruqi, Qiang Liu, and Xin Tong. \"Sampling in constrained domains with orthogonal-space variational gradient descent.\" Advances in Neural Information Processing Systems 35 (2022): 37108-37120.  \n[2]. Neklyudov, Kirill, et al. \"Wasserstein Quantum Monte Carlo: A Novel Approach for Solving the Quantum Many-Body Schrödinger Equation.\" Advances in Neural Information Processing Systems 36 (2024).  \n[3]. Ambrosio, Luigi, and Wilfrid Gangbo. \"Hamiltonian ODEs in the Wasserstein Space of Probability Measures.\" Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences 61.1 (2008): 18-53.", "limitations": "I appreciate the efforts of the authors. Nevertheless, I recommend that the authors include a dedicated section in the appendix detailing the assumptions such as convexity and smoothness within functionals. This addition would significantly aid in understanding and following the methodologies employed in this work.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. **Solid Mathematical Foundations:** The authors offer a rigorous mathematical derivation concerning the convergence behaviors of MD and PGD, supplemented by a thorough analysis of the theoretical background.\n2. **Relevance to Conference Themes:** The paper presents fundamental theoretical analysis that is highly pertinent to the themes of the conference.\n3. **Comprehensive Background Introduction:** The supplementary materials provide an extensive overview necessary to understand the nuances of the paper. The foundational work done by the authors is commendable.", "weaknesses": "1. **Results Presentation:** The presentation of results could be enhanced by providing more detailed visualizations. For instance, it would be beneficial to include a snapshot illustrating particle evolution for Figure 1, right $K_2$. Additionally, if the authors have access to MATLAB, it is recommended to utilize the `surf` function for the contour of Dirichlet distribution's pdf in Figure 4 to better demonstrate the sampling tasks.\n2. **Clarification of Definitions:** It is recommended that the authors include explicit definitions of convergence and global convergence to preclude any potential misunderstandings in the appendix.\n3. **Visualization of Convergence Data:** I suggest employing seaborn's `sns.lineplot` with a shaded area to depict the error bars, which would support the assertions made in checklist item 7 more robustly.\n4. **Overlooked References:** It appears that significant references relevant to discussions on the convergence of mirror descent may have been overlooked. I recommend incorporating additional reference [1] that provides pertinent insights into these discussions.\n\n\n----\nReferences:\n\n[1]. Tzen, Belinda, et al. \"Variational principles for mirror descent and mirror langevin dynamics.\" IEEE Control Systems Letters 7 (2023): 1542-1547.", "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 presentation of results could be enhanced by providing more detailed visualizations.**\n\nThank you for these suggestions, we will take it into account and add snapshots of particles as well as contours.\n\n**It appears that significant references relevant to discussions on the convergence of mirror descent may have been overlooked. I recommend incorporating additional reference [1] that provides pertinent insights into these discussions.**\n\nThank you for this reference that we did not know. We will add it to the paper. \n\n**From the perspective of gradient flow, assuming we have a velocity field $\\dot{X} = v_t$, we can derive that $v_t = -\\Vert \\nabla\\frac{\\delta \\mathcal{F}}{\\mu} \\Vert_2^2$. In this context, does assuming $\\beta$-smoothness offer additional advantages beyond the conventional benefits?**\n\nFor Wasserstein gradient flows, the velocity field is $v_t=-\\nabla_{W_2}\\mathcal{F}(\\mu_t)$, and we have $\\frac{\\mathrm{d}\\mathcal{F}(\\mu_t)}{\\mathrm{d}t} = \\langle \\nabla_{W_2}\\mathcal{F}(\\mu_t), v_t\\rangle_{L^2(\\mu_t)} = - \\|\\nabla_{W_2}\\mathcal{F}(\\mu_t)\\|_{L^2(\\mu_t)}^2$. Thus, the objective of the continuous formulation is necessarily non-increasing without assuming any smoothness of the objective. \n\nIn discrete time, the Wasserstein gradient descent (obtained as the explicit Euler discretization) $\\mu_{k+1} = (\\mathrm{Id}-\\tau\\nabla_{W_2}\\mathcal{F}(\\mu_k))_\\\\#\\mu_k$ requires smoothness of the objective to obtain descent at each iteration (and eventually global convergence guarantees if we assume furthermore convexity). See for instance analysis of gradient descent [2]. If the step size is too big, it is not guaranteed that the scheme decreases at each iteration, and could actually diverge. Hence in practice, knowing/estimating the smoothness constant will guide the choice of the step size.\n\n**Could the convergence location and rate be established from the perspective of a Lyapunov functional?**\n\nOur rates can indeed be obtained by decreasing the following Lyapunov functional: $L_k = k\\tau(\\mathcal{F}(\\mu_k) - \\mathcal{F}(\\nu)) + \\tau(\\frac{1}{\\tau}-\\alpha) W_\\phi(\\nu,\\mu_k)$. Indeed, $L_{k+1}-L_k \\le 0$ by using the inequalities from equations (85) and (7). Thus $L_k$ is a Lyapunov function [3]. By using a telescopic sum, we recover the inequality $\\mathcal{F}(\\mu_k) - \\mathcal{F}(\\nu) \\le \\frac{1-\\alpha\\tau}{k\\tau} W_\\phi(\\nu,\\mu_0)$ obtained in Proposition 3. We will make this clearer.\n\n**Is it feasible to extend the proposed method to Riemannian-based approaches, such as Riemannian SVGD?**\n\nExtending the proposed methods to Riemannian manifolds is an interesting future direction of works that we did not consider yet. \n\nFor now, Wasserstein gradient flows on manifolds have been sparsely studied. For instance, [4] studied theoretically Wasserstein gradient flows of the negative entropy on manifolds. In a more applied setting, [5] computed Wasserstein gradient flows on Lie group using the JKO scheme and [6] performed Wasserstein gradient descent of the Sliced-Wasserstein distance on Hadamard manifolds. There, the schemes and analyis have been derived on a case-by-case basis.\n\nTo the best of our knowledge, the Mirror Descent algorithm has not been yet extended to Riemannian manifolds. One difficulty might come from the current theoretical analysis which requires to split the inner product, which is not directly possible when using log maps. Moreover, the scheme might also come with computational difficulties, as there might not be closed-forms easy to compute in general.\n\n**Is it possible to extend this approach to systems characterized by the Fisher-Rao metric?**\n\nThe gradient in the space of probability distributions endowed with the Fisher-Rao metric is the 1st variation [7]. Thus, we might define Bregman divergences and the associated Mirror-Descent scheme using this gradient. This was studied e.g. in [8] for Mirror Descent with applications on the minimization of the KL divergence.\n\n\n**Could this approach be adapted to Hamiltonian Flow, which is the second order system to my knowledge?**\n\nThis is an interesting direction for future works which we did not yet consider and which goes beyond our knowledge.\n\n\n**I recommend that the authors include a dedicated section in the appendix detailing the assumptions such as convexity and smoothness within functionals.**\n\nThank you for this suggestion. We will expand our existing sections B.2 and C.3 in the appendix in this direction.\n\n[1] Tzen, B., Raj, A., Raginsky, M., & Bach, F. (2023). Variational principles for mirror descent and mirror langevin dynamics. IEEE Control Systems Letters, 7, 1542-1547.\n\n[2] Garrigos, G., & Gower, R. M. (2023). Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.11235.\n\n[3] Wilson, A. (2018). Lyapunov arguments in optimization. University of California, Berkeley.\n\n[4] Erbar, M. (2010). The heat equation on manifolds as a gradient flow in the Wasserstein space. In Annales de l'IHP Probabilités et statistiques (Vol. 46, No. 1, pp. 1-23).\n\n[5] Bon, D., Pai, G., Bellaard, G., Mula, O., & Duits, R. (2024). Optimal Transport on the Lie Group of Roto-translations. arXiv preprint arXiv:2402.15322.\n\n[6] Bonet, C., Drumetz, L., & Courty, N. (2024). Sliced-Wasserstein Distances and Flows on Cartan-Hadamard Manifolds. arXiv preprint arXiv:2403.06560.\n\n[7] Gallouët, T. O., & Monsaingeon, L. (2017). A JKO splitting scheme for Kantorovich--Fisher--Rao gradient flows. SIAM Journal on Mathematical Analysis, 49(2), 1100-1130.\n\n[8] Aubin-Frankowski, P. C., Korba, A., & Léger, F. (2022). Mirror descent with relative smoothness in measure spaces, with application to sinkhorn and em. Advances in Neural Information Processing Systems, 35, 17263-17275.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1pOHSwZBDy", "reviewer": "Reviewer_P753", "summary": "The authors study the mirror descent method in Wasserstein-2 spaces. This is based on constructing the Bregman divergence functionals in Wasserstein-2 spaces. One can use it to define a mirror descent direction in the sense of pushforward mapping functions, which generalizes the classical Wasserstein gradient descent direction. The authors then studied the convergence analysis of the Wasserstein mirror descent algorithm. Several numerical examples are used to demonstrate the effectiveness of the proposed method.", "questions": "The authors also propose the Wassrsetin-Bregman proximal gradient method. Can authors illustrate the main difficulty of implementing it in practice or some potential computational schemes? How does it compare with the classical sampling schemes if the objective functional is chosen as the KL divergence? \n\nI am willing to improve my score after addressing these questions.", "limitations": "There is no limitation on this paper.", "rating": 9, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "The authors define and apply the Wasserstein mirror descent directions to conduct sampling algorithms. \n\nThe mirror descent method in Wasserstein space is a very natural and novel approach. Congratulations to the authors for using the Wasserstein-Bregman divergences in sampling algorithms. This paper's proof is very clear in the mirror descent method from optimization and optimal transport.", "weaknesses": "On page 8, the authors don't study the mirror descent of KL divergences for non-Gaussian target distributions. This could be a very intriguing example. The authors may think about how to implement this case. This could be compared with the Wasserstein gradient flow of KL divergence.", "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.\n\n**On page 8, the authors don't study the mirror descent of KL divergences for non-Gaussian target distributions. This could be a very intriguing example. The authors may think about how to implement this case. This could be compared with the Wasserstein gradient flow of KL divergence.**\n\nThe difficulty of optimizing the KL divergence with a non-Gaussian target is that we cannot use closed-forms for the mean and covariance anymore (as done in the experiment of Figure 2 for Gaussian, where closed-forms were obtained by projecting the distributions at each step on the space of Gaussian). \n\nA possible solution is to use a particle approximation, which we considered in Appendix I.2, with the goal to sample from a distribution on the simplex. Thus, the Bregman potential was chosen here as a potential energy, with potential a barrier allowing to keep the samples onto the simplex. The density of the distribution at each step was approximated with a Kernel Density Estimator to approximate the Wasserstein gradient of the KL.\n\n\n**The authors also propose the Wasserstein-Bregman proximal gradient method. Can authors illustrate the main difficulty of implementing it in practice or some potential computational schemes? How does it compare with the classical sampling schemes if the objective functional is chosen as the KL divergence?**\n\nThe main difficulty of implementing the Wasserstein-Bregman proximal gradient method is that it is required to solve at each step a backward step (see equation (128)), which does not have a closed-form in general [1]. Thus, at each step, it is required to solve another optimization problem. For the classical JKO scheme, it has been proposed to solve it using approximations with e.g. entropic approximation [2] or neural networks [3,4], but these methods are computationally costly. Nonetheless, note that there is a closed-form when the target and initial distribution are Gaussian [5,6].\n\nIn Figure 2, we considered the Wasserstein-Bregman proximal gradient method with the KL divergence as objective functional, and a Gaussian $\\mathcal{N}(0,\\Sigma)$ as target distribution. For a Bregman potential of the form $\\phi(\\mu) = \\int V\\mathrm{d}\\mu$, with $V(x)=x^T\\Lambda^{-1} x$, we derived the closed-form in Appendix H.4 leveraging the fact that the distributions are always Gaussian. Using $\\Lambda=\\Sigma$, we showed on Figure 2 that it converges much faster than the regular Wasserstein proximal method studied in [1,5] (which corresponds to $\\Lambda=I_d$).\n\n\n[1] Salim, A., Korba, A., & Luise, G. (2020). The Wasserstein proximal gradient algorithm. Advances in Neural Information Processing Systems, 33, 12356-12366.\n\n[2] Peyré, G. (2015). Entropic approximation of Wasserstein gradient flows. SIAM Journal on Imaging Sciences, 8(4), 2323-2351.\n\n[3] Mokrov, P., Korotin, A., Li, L., Genevay, A., Solomon, J. M., & Burnaev, E. (2021). Large-scale wasserstein gradient flows. Advances in Neural Information Processing Systems, 34, 15243-15256.\n\n[4] Alvarez-Melis, D., Schiff, Y., & Mroueh, Y. (2021). Optimizing functionals on the space of probabilities with input convex neural networks. arXiv preprint arXiv:2106.00774.\n\n[5] Diao, M. Z., Balasubramanian, K., Chewi, S., & Salim, A. (2023). Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein space. In International Conference on Machine Learning (pp. 7960-7991). PMLR.\n\n[6] Wibisono, A. (2018). Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem. In Conference on Learning Theory (pp. 2093-3027). PMLR.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their positive comments and common appraisal of the soundness of our approach to lift Euclidean optimization schemes to the Wasserstein space. Following the reviewers' comments, we will improve the clarity of the exposition for the revised version, taking advantage of the allowed extra page.

Meta-review: This paper considers mirror descent and preconditioned gradient descent analogs in the  Wasserstein space of probability distributions. These algorithms better capture the geometry of the optimization has interesting properties explored in optimization literature. Authors explore those in Wasserstein space prove guarantees of convergence of Wasserstein-gradient-based discrete-time schemes.


This paper was reviewed by five reviewers with the following Scores/Confidence: 6/4, 6/3, 9/5, 7/3, 7/3. I think the paper is studying an interesting topic and the results are relevant to NeurIPS community. The following concerns were brought up by the reviewers:

- The paper is very dense and presentation can be friendly to a broader audiance.
  
- The advantages of MD framework should be emphasized.
  
- Definitions should be clarified. See reviewer PBU2's comments.


Authors should carefully go over reviewers' suggestions and address any remaining concerns in their final revision. Based on the reviewers' suggestion, as well as my own assessment of the paper, I recommend including this paper to the NeurIPS 2024 program.
Decision: Accept (spotlight)