Title: Mirror and Preconditioned Gradient Descent in Wasserstein Space
Venue: NeurIPS 2024
Accepted: yes
Keywords: wasserstein gradient flows

mirror descent

preconditioned gradient descent
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

Abstract

As the problem of minimizing functionals on the Wasserstein space encompasses many applications in machine learning, different optimization algorithms on 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 Wassersteingradient-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.

1 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 as it is an infinite-dimensional 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 to 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 a 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) on the position of particles for functionals with a closed-form 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 as the Kullback-Leibler (KL) divergence for sampling, other discretizations can be easily computed such as the Unadjusted Langevin Algorithm (ULA) [106]. This leaves the question open of assessing

Anna Korba

the theoretical and empirical performance of other optimization algorithms relying on alternative geometries and time-discretizations.

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 a cost that is a 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. In 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 algorithms 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) [114], or even Wasserstein gradient descent [113]. Preconditioned Wasserstein gradient descent has been also recently proposed for specific geometries in [31, 44] to minimize the KL in a 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 Sliced-Wasserstein (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 L 2 norms on measures, see Appendix A for more details.

Contributions. We are interested in minimizing a functional F : P 2 ( R d ) → R ∪ { + ∞} over probability distributions, through schemes of the form, for τ > 0 , k ≥ 0 ,

with different costs d : L 2 ( µ k ) × L 2 ( µ k ) → R + , and in providing convergence conditions. While we can recover a map ¯ T = T k ◦ T k -1 · · · ◦ T 1 such that µ k = ¯ T # µ 0 , the scheme (1) proceeds by successive regularized linearizations retaining the Wasserstein structure, since the tangent space to P 2 ( R d ) at µ is a subset of L 2 ( µ ) [96]. This paper is 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 L 2 ( µ ) 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 translation-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 P 2 ( R d ) of probability measures µ on R d with finite second moment and P 2 , ac ( R d ) ⊂ P 2 ( R d ) its subset of absolutely continuous probability measures with respect to the Lebesgue measure. For any µ ∈ P 2 ( R d ) , we denote by L 2 ( µ ) the Hilbert space of functions f : R d → R d such that ∫ ∥ f ∥ 2 d µ < ∞ equipped with the norm ∥·∥ L 2 ( µ ) and inner product ⟨· , ·⟩ L 2 ( µ ) . For a Hilbert space X , the Fenchel transform of f : X → R is f ∗ ( y ) = sup x ∈ X ⟨ x, y ⟩ -f ( x ) . Given a measurable map T : R d → R d and µ ∈ P 2 ( R d ) , T # µ is the pushforward measure of µ by

T ; and T ⋆ µ = ∫ T( · -x )d µ ( x ) . For µ, ν ∈ P 2 ( R d ) , the Wasserstein-2 distance is W 2 2 ( µ, ν ) = inf γ ∈ Π( µ,ν ) ∫ ∥ x -y ∥ 2 d γ ( x, y ) , where Π( µ, ν ) = { γ ∈ P ( R d × R d ) , π 1 # γ = µ, π 2 # γ = ν } with π i ( x 1 , x 2 ) = x i , is the set of couplings between µ and ν , and we denote by Π o ( µ, ν ) the set of optimal couplings. When the optimal coupling is of the form γ = (Id , T ν µ ) # µ with Id : x ↦→ x and T ν µ ∈ L 2 ( µ ) satisfying (T ν µ ) # µ = ν , we call T ν µ the OT map. We refer to the metric space ( P 2 ( R d ) , W 2 ) as the Wasserstein space. We note S ++ d ( R ) the space of symmetric positive definite matrices, and for x ∈ R d , Σ ∈ S ++ d ( R ) , ∥ x ∥ 2 Σ = x T Σ x .

2 Background

In this section, we fix µ ∈ P 2 ( R d ) and introduce first the Bregman divergence on L 2 ( µ ) 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 a functional F : P 2 ( R d ) → R along curves and discuss notions of convexity on P 2 ( R d ) . We refer the reader to Appendix B and Appendix C for more details on L 2 ( µ ) and the Wasserstein space respectively. Finally, we introduce the mirror descent and preconditioned gradient descent on R d .

Bregman divergence on L 2 ( µ ) . Frigyik et al. [54, Definition 2.1] defined the Bregman divergence of Fréchet differentiable functionals. In our case, we only need Gâteaux differentiability. In this paper, ∇ refers to the Gâteaux differential, which coincides with the Fréchet derivative if the latter exists.

Definition 1. Let ϕ µ : L 2 ( µ ) → R be convex and continuously Gâteaux differentiable. The Bregman divergence is defined for all T , S ∈ L 2 ( µ ) as d ϕ µ (T , S) = ϕ µ (T) -ϕ µ (S) -⟨∇ ϕ µ (S) , T -S ⟩ L 2 ( µ ) .

We use the same definition on R d . The map ϕ µ (respectively ∇ ϕ µ ) in the definition of d ϕ µ above is referred to as the Bregman potential (respectively mirror map). If ϕ µ is strictly convex, then d ϕ µ is a valid Bregman divergence, i.e. it is positive and separates maps µ -almost everywhere (a.e.). In particular, for ϕ µ (T) = 1 2 ∥ T ∥ 2 L 2 ( µ ) , we recover the L 2 norm as a divergence d ϕ µ (T , S) = 1 2 ∥ T -S ∥ 2 L 2 ( µ ) . Bregman divergences have received a lot of attention as they allow to define provably convergent schemes for functions which are not smooth in the standard ( e.g. Euclidean) sense [11, 88], and thus for which gradient descent is not appropriate. These guarantees rely on the notion of relative smoothness and relative convexity [88, 89], which we introduce now on L 2 ( µ ) .

Definition 2 (Relative smoothness and convexity) . Let ψ µ , ϕ µ : L 2 ( µ ) → R be convex and continuously Gâteaux differentiable. We say that ψ µ is β -smooth (respectively α -convex) relative to ϕ µ if and only if for all T , S ∈ L 2 ( µ ) , d ψ µ (T , S) ≤ β d ϕ µ (T , S) (respectively d ψ µ (T , S) ≥ α d ϕ µ (T , S) ).

Similarly to the Euclidean case [88], relative smoothness and convexity are equivalent to respectively βϕ µ -ψ µ and ψ µ -αϕ µ 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, we need the L 2 Fenchel transform ϕ ∗ µ of ϕ µ to be differentiable, e.g. to compute its Bregman divergence d ϕ ∗ µ . We show in Lemma 18 that a sufficient condition to satisfy this property is for ϕ µ to be strictly convex, lower semicontinuous and superlinear, i.e. lim ∥ T ∥→∞ ϕ µ (T) / ∥ T ∥ L 2 ( µ ) = + ∞ . Moreover, in this case, ( ∇ ϕ µ ) -1 = ∇ ϕ ∗ µ . When needed, we will suppose that ϕ µ satisfies this assumption.

Differentiability on ( P 2 ( R d ) , W 2 ) . Let F : P 2 ( R d ) → R ∪ { + ∞} , and denote D ( F ) = { µ ∈ P 2 ( R d ) , F ( µ ) < + ∞} the domain of F and D ( ˜ F µ ) = { T ∈ L 2 ( µ ) , T # µ ∈ D ( F ) } the domain of ˜ F µ defined as ˜ F µ (T) := F (T # µ ) for all T ∈ L 2 ( µ ) . In the following, we use the differential structure of ( P 2 ( R d ) , W 2 ) introduced in [17, Definition 2.8], and we say that ∇ W 2 F ( µ ) is a Wasserstein gradient of F at µ ∈ D ( F ) if for any ν ∈ P 2 ( R d ) and any optimal coupling γ ∈ Π o ( µ, ν ) ,

If such a gradient exists, then we say that F is Wasserstein differentiable at µ [17, 74]. Moreover there is a unique gradient belonging to the tangent space of P 2 ( R d ) verifying (2) [74, Proposition

2.5], and we will always restrict ourselves without loss of generality to this particular gradient, see Appendix C.1. The differentiability of F and ˜ F µ are very related, as described in the following proposition.

Proposition 1. Let F : P 2 ( R d ) → R ∪ { + ∞} be a Wasserstein differentiable functional on D ( F ) . Let µ ∈ P 2 ( R d ) and ˜ F µ (T) = F (T # µ ) for all T ∈ D ( ˜ F µ ) . Then, ˜ F µ is Fréchet differentiable, and for all S ∈ D ( ˜ F µ ) , ∇ ˜ F µ (S) = ∇ W 2 F (S # µ ) ◦ S .

The Wasserstein differentiable functionals include c -Wasserstein costs on P 2 , ac ( R d ) [74, Proposition 2.10 and 2.11], potential energies V ( µ ) = ∫ V d µ or interaction energies W ( µ ) = 1 2 ∫∫ W ( x -y ) d µ ( x )d µ ( y ) for V : R d → R and W : R d → R differentiable and with bounded Hessian [74, Section 2.4]. In particular, their Wasserstein gradients read as ∇ W 2 V ( µ ) = ∇ V and ∇ W 2 W ( µ ) = ∇ W ⋆ µ . However, entropy functionals, e.g. the negative entropy defined as H ( µ ) = ∫ log ( ρ ( x ) ) d µ ( x ) for distributions µ admitting a density ρ w.r.t. the Lebesgue measure, are not Wasserstein differentiable. In this case, we can consider subgradients ∇ W 2 F ( µ ) at µ for which (2) becomes an inequality. To guarantee that the Wasserstein subgradient is not empty, we need ρ to satisfy some Sobolev regularity, see e.g. [5, Theorem 10.4.13] or [108]. Then, if ∇ log ρ ∈ L 2 ( µ ) , the only subgradient of H in the tangent space is ∇ W 2 H ( µ ) = ∇ log ρ , 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 KL to a 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 2 ( µ ) and of Wasserstein differentiability enable us to consider alternative Bregman potentials than the L 2 ( µ ) -norm mentioned above. For instance, for V convex, differentiable and L -smooth, we can use potential energies ϕ V µ (T) := V (T # µ ) , for which d ϕ V µ (T , S) = ∫ d V ( T( x ) , S( x ) ) d µ ( x ) where d V is the Bregman divergence of V on R d . Notice that ϕ µ (T) = 1 2 ∥ T ∥ 2 L 2 ( µ ) is a specific example of a potential energy where V = 1 2 ∥ · ∥ 2 . Moreover, we will consider interaction energies ϕ W µ (T) := W (T # µ ) with W convex, differentiable, L -smooth, and satisfying W ( -x ) = W ( x ) ; for which d ϕ µ W (T , S) = 1 2 ∫∫ d W ( T( x ) -T( x ′ ) , S( x ) -S( x ′ ) ) d µ ( x )d µ ( x ′ ) (see Appendix I.3). We will also use ϕ H µ (T) = H (T # µ ) 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 ( P 2 ( R d ) , W 2 ) . 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 ( P 2 ( R d ) , W 2 ) , different such notions have been proposed based on specific choices of curves. The most popular one is to require the functional F to be α -convex along geodesics (see Definition 10), which are of the form µ t = ( (1 -t )Id + t T µ 1 µ 0 ) # µ 0 if µ 0 ∈ P 2 , ac ( R d ) and

µ 1 ∈ P 2 ( R d ) , with T µ 1 µ 0 the OT map between them. In that setting,

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 µ ↦→ 1 2 W 2 2 ( µ, η ) itself, for a given η ∈ P 2 ( R d ) , 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 2 ( R d ) , are of the form µ t = ( (1 -t )T µ η + t T ν η ) # η for T µ η , T ν η OT maps from η to µ and ν . Observing that for ϕ µ (T) = 1 2 ∥ T ∥ 2 L 2 ( µ ) , we can rewrite (3) as α d ϕ µ 0 (T µ 1 µ 0 , Id) ≤ d ˜ F µ 0 (T µ 1 µ 0 , Id) and see that being convex along geodesics boils down to being convex in the L 2 sense for S = Id and T chosen as an OT map. This observation motivates us to consider a more refined notion of convexity along curves.

Definition 3. Let µ ∈ P 2 ( R d ) , T , S ∈ L 2 ( µ ) and for all t ∈ [0 , 1] , µ t = (T t ) # µ with T t = (1 -t )S+ t T . We say that F : P 2 ( R d ) → R is α -convex (resp. β -smooth) relative to G : P 2 ( R d ) → R along t ↦→ µ t if for all s, t ∈ [0 , 1] , d ˜ F µ (T s , T t ) ≥ α d ˜ G µ (T s , T t ) (resp. d ˜ F µ (T s , T t ) ≤ β d ˜ G µ (T s , T t ) ).

Notice that in contrast with Definition 2, Definition 3 is stated for a fixed distribution µ and directions ( S , T ), and involves comparisons between Bregman divergences depending on µ and curves (T s ) s ∈ [0 , 1] depending on S , T . The larger family of S and T for which Definition 3 holds, the more restricted is the notion of convexity of Fα G (resp. of β G-F ) on P 2 ( R d ) . For instance, Wasserstein-2 generalized geodesics with anchor η ∈ P 2 ( R d ) correspond to considering S , T as all the OT maps originating from η , among which geodesics are particular cases when taking η = µ (hence S = Id ). If we furthermore ask for α -convexity to hold for all µ ∈ P 2 ( R d ) and T , S ∈ L 2 ( µ ) ( i.e. , not only OT maps), then we recover the convexity along acceleration free-curves as introduced in [28, 98, 117]. Our motivation behind introducing Definition 3 is that the convergence proofs of MD and preconditioned GD require relative smoothness and convexity properties to hold only along specific curves.

Mirror descent and preconditioned gradient descent on R d . These schemes read respectively as ∇ ϕ ( x k +1 ) -∇ ϕ ( x k ) = -τ ∇ f ( x k ) [12] and y k +1 -y k = -τ ∇ h ∗ ( ∇ g ( y k ) ) [89], where the objectives f, g and the regularizers h, ϕ are convex C 1 functions from R d to R . The algorithms are closely related since, using the Fenchel transform and setting g = ϕ ∗ and h ∗ = f , we see that, for y = ∇ ϕ ( x ) , the two schemes are equivalent when permuting the roles of the objective and of the regularizer. For MD, convergence of f is ensured if f is both 1 / τ -smooth and α -convex relative to ϕ [88, Theorem 3.1]. Concerning preconditioned GD, assuming that h, g are Legendre, ( g ( y k ) ) k converges to the minimum of g if h ∗ is both 1 / τ -smooth and α -convex relative to g ∗ with α > 0 [89, Theorem 3.9].

3 Mirror descent

For every µ ∈ P 2 ( R d ) , let ϕ µ : L 2 ( µ ) → R be strictly convex, proper and differentiable and assume that the (sub)gradient ∇ W 2 F ( µ ) ∈ L 2 ( µ ) exists. In this section, we are interested in analyzing the scheme (1) where the cost d is chosen as a Bregman divergence, i.e. d ϕ µ as defined in Definition 1. This corresponds to a mirror descent scheme in P 2 ( R d ) . For τ > 0 and k ≥ 0 , it writes:

Iterates of mirror descent. In all that follows, we assume that the iterates (4) exist, which is true e.g. for a superlinear ϕ µ k , since the objective is a sum of linear functions and of the continuous ϕ µ k . In the previous section, we have seen that the second term in the proximal scheme (4) can be interpreted as a linearization of the functional F at µ k for Wasserstein (sub)differentiable functionals. Now define for all T ∈ L 2 ( µ k ) , J(T) = d ϕ µ k (T , Id)+ τ ⟨∇ W 2 F ( µ k ) , T -Id ⟩ L 2 ( µ k ) . Then, deriving the first order conditions of (4) as ∇ J(T k +1 ) = 0 , we obtain µ k -a.e.,

Note that for ϕ µ (T) = 1 2 ∥ T ∥ 2 L 2 ( µ ) , the update (5) translates as T k +1 = Id -τ ∇ W 2 F ( µ k ) , 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 in R d [12]. We discuss in Appendix D.2 the continuous formulation of (4), showing it coincides with the gradient flow of the mirror Langevin [3, 130], the limit of the JKO scheme with Bregman groundcosts [104], Information Newton's flows [126], or Sinkhorn's flow [41] for specific choices of ϕ 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 µ, ρ ∈ P 2 , ac ( R d ) and ν ∈ P 2 ( R d ) , setting T µ,ν ϕ µ = argmin T # µ = ν d ϕ µ (T , Id) , U ρ,ν ϕ ρ = argmin U # ρ = ν d ϕ ρ (U , Id) , the functional ϕ µ is such that, for any S ∈ L 2 ( µ ) satisfying S # µ = ρ , we have d ϕ µ (T µ,ν ϕ µ , S) ≥ d ϕ ρ (U ρ,ν ϕ ρ , Id) .

The inequality in Assumption 1 can be interpreted as follows: the 'distance' between ρ and ν is greater when observed from an anchor µ that differs from ρ and ν . We demonstrate that Bregman divergences satisfy this assumption under the following conditions on the Bregman potential ϕ .

Proposition 2. Let µ, ρ ∈ P 2 , ac ( R d ) and ν ∈ P 2 ( R d ) . Let ϕ µ be a pushforward compatible functional, i.e. there exists ϕ : P 2 ( R d ) → R such that for all T ∈ L 2 ( µ ) , ϕ µ (T) = ϕ (T # µ ) . Assume furthermore ∇ W 2 ϕ ( µ ) and ∇ W 2 ϕ ( ρ ) invertible (on R d ). Then, ϕ µ satisfies Assumption 1.

All the maps ϕ V µ , ϕ W µ and ϕ H µ defined in Section 2 satisfy the assumptions of Proposition 2 under mild requirements, see Appendix D.1. The proof of Proposition 2 is given in Appendix H.2. It relies on the definition of an appropriate optimal transport problem

and on the proof of existence of OT maps for absolutely continuous measures (see Proposition 15), which implies W ϕ ( ν, µ ) = d ϕ µ (T µ,ν ϕ µ , Id) with T µ,ν ϕ µ defined as in Assumption 1. From there, we can conclude that ϕ µ satisfies Assumption 1. We notice that the corresponding transport problem recovers previously considered objects such as 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 S = Id and T = T k +1 solutions of (4) for all k ≥ 0 , we derive the following descent lemma, which ensures that ( F ( µ k ) ) k is non-increasing. Its proof can be found in Appendix H.3 and relies on the three-point inequality [29], which we extended to L 2 ( µ ) in Lemma 29.

Proposition 3. Let β > 0 , τ ≤ 1 β . Assume for all k ≥ 0 , F is β -smooth relative to ϕ along t ↦→ ( (1 -t )Id + t T k +1 ) # µ k , which implies β d ϕ µ k (T k +1 , Id) ≥ d ˜ F µ k (T k +1 , Id) . Then, for all k ≥ 0 ,

Assuming additionally the convexity of F along the curves µ t = ( (1 -t )Id + t T µ,ν ϕ µ ) # µ , t ∈ [0 , 1] and that ϕ satisfies Assumption 1, we can obtain global convergence.

Proposition 4. Let ν ∈ P 2 ( R d ) , α ≥ 0 . Suppose Assumption 1 and the conditions of Proposition 3 hold, and that F is α -convex relative to ϕ along the curves t ↦→ ( (1 -t )Id + t T µ k ,ν ϕ µ k ) # µ k . Then, for all k ≥ 1 ,

Moreover, if α > 0 , taking ν = µ ∗ the minimizer of F , we obtain a linear rate: for all k ≥ 0 , W ϕ ( µ ∗ , µ k ) ≤ (1 -τα ) k W ϕ ( µ ∗ , µ 0 ) .

The proof of Proposition 4 can be found in Appendix H.4, and requires Assumption 1 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 R d , because the minimization problem of each iteration (4) happens in a different space L 2 ( µ k ) . We discuss in Section 5 how to verify the relative smoothness and convexity on some examples. In particular, when both F and ϕ are potential energies, it is inherited from the relative smoothness and convexity on R d , and the conditions are similar with those for MD on R d . 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 over 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 ( µ ) = G ( µ ) + H ( µ ) with H non smooth; which includes the KL divergence decomposed as a potential energy plus the negative entropy.

Implementation. We now discuss the practical implementation of MD on ( P 2 ( R d ) , W 2 ) as written in (5). If ϕ µ is pushforward compatible, we have ∇ ϕ µ k (T k +1 ) = ∇ W 2 ϕ ( (T k +1 ) # µ k ) ◦ T k +1 ; but if ∇ ϕ ∗ µ k is unknown, the scheme is implicit in T k +1 . A possible solution is to rely on a root finding algorithm such as Newton's method to find the zero of ∇ J at each step, which we use in Section 5 for ϕ W µ as Bregman potential. However, this procedure may be computationally costly and scale badly w.r.t. the dimension and the number of samples, see Appendix G.1. Nonetheless, in the special case ϕ V µ (T) = ∫ V ◦ T d µ with V differentiable, strongly convex and L -smooth, since ∇ W 2 V ( µ ) = ∇ V and ( ∇ V ) -1 = ∇ V ∗ , the scheme reads as

and can be implemented on particles, i.e. for ˆ µ k = 1 n ∑ n i =1 δ x k i , x k +1 i = ∇ V ∗ ( ∇ V ( x k i ) -τ ∇ W 2 F (ˆ µ k )( x k i ) ) for all k ≥ 0 , i ∈ { 1 , . . . , n } . This scheme is analogous to MD in R d [12] and has been introduced as the mirror Wasserstein gradient descent [113]. Moreover, for V = 1 2 ∥ · ∥ 2 2 , as observed earlier, we recover the usual Wasserstein gradient descent, i.e. T k +1 = Id -τ ∇ W 2 F ( µ k ) . 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, 114, 131] or the Kalman-Wasserstein gradient descent [56]. We refer to Appendix D.4 for more details.

4 Preconditioned gradient descent

As seen in Section 2, preconditioned gradient descent on R d has dual convergence conditions compared to mirror descent. Our goal is to extend these to (1) and P 2 ( R d ) . Let τ > 0 , µ ∈ P 2 ( R d ) and h : R d → R proper and strictly convex on R d . We consider in this section ϕ h µ (T) = ∫ h ◦ T d µ and d(T , Id) = ϕ h µ k ( (Id -T) /τ ) τ = ∫ h ( ( x -T( x )) /τ ) τ d µ k ( x ) . This type of discrepancy is analogous to OT costs with translation-invariant ground cost c ( x, y ) = h ( x -y ) , which have been popular as they induce an OT map [110, Box 1.12]. Such costs have been introduced e.g. in [39, 70] to promote sparse transport maps. More generally, for ϕ µ strictly convex, proper, differentiable and superlinear, we have ( ∇ ϕ µ ) -1 = ∇ ϕ ∗ µ and the following theory is still valid. For simplicity, we leave studying more general ϕ for future works. Here, the scheme (1) results in:

Deriving the first order conditions similarly to (5) in Section 3, we obtain the following update:

Notice that for h = 1 2 ∥· ∥ 2 2 the squared Euclidean norm, ϕ h µ and ϕ h ∗ µ recover the squared L 2 ( µ ) 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 a 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 h = ∥ · ∥ p p with p > 1 [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 a descent lemma on ( ϕ h ∗ µ k ( ∇ W 2 F ( µ k )) ) k under a technical inequality between the Bregman divergences of ϕ h ∗ µ k and ˜ F µ k for all k ≥ 0 . Additionally, we also suppose that F is convex along the curves generated by S = T k +1 and T = Id . This last hypothesis ensures that d ˜ F µ k (T k +1 , Id) ≥ 0 , and thus that ( ϕ h ∗ µ k ( ∇ W 2 F ( µ k )) ) k is non-increasing. Analogously to the Euclidean case, ϕ h ∗ µ 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 e.g.

quantifies the mag[54, Appendix B.7] or Lemma 28) and algebra with the definition of Bregman divergences. Proposition 5. Let β > 0 . Assume τ ≤ 1 β , and for all k ≥ 0 , F convex along t ↦→ ( (1 -t )T k +1 + t Id ) # µ k and d ϕ h ∗ µ k ( ∇ W 2 F ( µ k +1 ) ◦ T k +1 , ∇ W 2 F ( µ k ) ) ≤ β d ˜ F µ k (Id , T k +1 ) . Then, for all k ≥ 0 ,

Under an additional assumption of a reverse inequality between the Bregman divergences of ϕ h ∗ µ k and ˜ F µ k , and assuming that ϕ h ∗ µ attains its minimum in 0, we can show the convergence of the gradient quantified by ϕ h ∗ µ (see Lemma 21), and the convergence of ( F ( µ k ) ) k towards the minimum of F .

Proposition 6. Let α ≥ 0 and µ ∗ ∈ P 2 ( R d ) be the minimizer of F . Assume the conditions of Proposition 5 hold, and that for ¯ T = argmin T , T # µ k = µ ∗ d ˜ F µ k (Id , T) , α d ˜ F µ k (Id , ¯ T) ≤

d ϕ h ∗ µ k ( ∇ W 2 F ( ¯ T # µ k ) ◦ ¯ T , ∇ W 2 F ( µ k ) ) . Then, for all k ≥ 1 , since ∇ W 2 F ( µ ∗ ) = 0 and ϕ h ∗ µ k (0) = h ∗ (0) ,

Moreover, assuming that h ∗ attains its minimum at 0 and α > 0 , F converges towards its minimum at a linear rate, i.e. for all k ≥ 0 , F ( µ k ) -F ( µ ∗ ) ≤ (1 -τα ) k ( F ( µ 0 ) -F ( µ ∗ ) ) .

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 R d for a cost h and an objective function g , that these conditions were equivalent to β -smoothness and α -convexity of the preconditioner h ∗ (analogous to ϕ ∗ µ ) relative to the convex conjugate of the objective g ∗ (analogous to ˜ F ∗ µ ). To write the inequalities we assumed as a relative smoothness/convexity property of ϕ h ∗ µ k w.r.t. ˜ F ∗ µ k , we would need at least to ensure that ˜ F ∗ µ k is differentiable, in order to define its Bregman divergence according to Definition 1. This can be done e.g. by assuming ˜ F µ k 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. Let µ ∈ P 2 ( R d ) and T ∈ L 2 ( µ ) . Assume ˜ F ∗ µ is Gâteaux differentiable and define F ∗ µ on t ↦→ µ t = (T t ) # µ as F ∗ µ ( µ t ) = ˜ F ∗ µ (T t ) for T t = (1 -t )U + t S for all t ∈ [0 , 1] , S , U ∈ L 2 ( µ ) .

If ϕ h ∗ is β -smooth relative to F ∗ µ along t ↦→ ( (1 -t ) ∇ W 2 F ( µ ) + t ∇ W 2 F (T # µ ) ◦ T ) # µ . Then,

Likewise, if ϕ h ∗ is α -convex relative to F ∗ µ along t ↦→ ( (1 -t ) ∇ W 2 F ( µ ) + t ∇ W 2 F (T # µ ) ◦ T ) # µ , then

In particular, for F a potential energy, the conditions coincide with those of [89] in R d . We refer to Appendix E.1 for more details.

5 Applications and Experiments

In this section, we first discuss how to verify the relative convexity and smoothness between functionals 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 1 .

Relative convexity of functionals. To assess relative convexity or smoothness as stated in Definition 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 R d of the underlying potential functions (respectively interaction kernels) to P 2 ( R d ) , 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 are no interactions between the particles. The conditions of convergence then coincide with the ones obtained for MD and preconditioned GD on R d [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, it is less straightforward. Using equivalent notions of convexity (see Proposition 13), we may instead compare their Hessians along the right curves,

Time (             )

τ

iter

Time (             )

τ

iter

·

·

see Appendix E.2 for an example between an interaction and a potential energy. We note also that for the particular case of a functional obtained as a sum F = G + H with ˜ G µ and ˜ H µ convex, since d ˜ F µ = d ˜ G µ + d ˜ H µ , d ˜ F µ ≥ max { d ˜ G µ , d ˜ H µ } , and thus F is 1-convex relative to G and H . This includes e.g. 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 ( µ ) = 1 2 ∫∫ W ( x -y ) d µ ( x )d µ ( y ) with kernel W ( z ) = 1 4 ∥ z ∥ 4 Σ -1 -1 2 ∥ z ∥ 2 Σ -1 , Σ ∈ S ++ d ( 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 to K 4 ( z ) = 1 4 ∥ z ∥ 4 2 + 1 2 ∥ z ∥ 2 2 with β = 4 , and W is β -smooth relative to ϕ µ (T) = 1 2 ∫∫ K 4 ( T( x ) -T( y ) ) d µ ( x )d µ ( y ) . Supposing additionally that the distributions are compactly supported, we can show that W is smooth relative to the interaction energy with K 2 ( z ) = 1 2 ∥ z ∥ 2 2 . For ill-conditioned Σ , 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 K Σ 2 ( z ) = 1 2 ∥ z ∥ 2 Σ -1 and K Σ 4 ( z ) = 1 4 ∥ z ∥ 4 Σ -1 + 1 2 ∥ z ∥ 2 Σ -1 . 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) = ( ∇ K ⋆ T # µ ) ◦ T , the scheme (5) needs to be approximated using Newton's algorithm which can be computationally heavy. Using ϕ V µ (T) = ∫ V ◦ Td µ with V = K Σ 2 , we obtain a more computationally friendly scheme with the same convergence, see Appendix G.2, but for which the smoothness is trickier to show.

MDon KL. We now focus on minimizing F ( µ ) = ∫ V d µ + H ( µ ) for V ( x ) = 1 2 x T Σ -1 x with Σ possibly ill-conditioned, whose minimizer is the Gaussian ν = N (0 , Σ) , and for which Wasserstein gradient descent is slow to converge. We study the MD scheme in (4) with negative entropy H as the Bregman potential (NEM), and compare it on Figure 2 with the Forward-Backward (FB) scheme studied in [43] and the ideally preconditioned Forward-Backward scheme (PFB) with Bregman potential ϕ V µ (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 N (0 , Σ 0 ) , the distributions stay Gaussian over the flow, and their closed-form is 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 R 10 with Σ = UDU T , D diagonal and scaled in log space between 1 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 to a perturbation is a central question in biology. In this context, as the measuring process is destructive, feature descriptions of control and treated cells must be dealt with as (unpaired) source µ and target distributions ν . 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 µ to ν 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)

#iters convergence

F

(

ˆ

µ

)

gradient flow of a functional F ( µ ) = D ( µ, ν ) initialized at µ 0 = µ , where D is a distributional metric, and predict the perturbed population via ˆ µ = min µ F ( µ ) . 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 a latent space learned by the scGen autoencoder [87] for the scRNA data. We use three metrics: the Sliced-Wasserstein distance SW 2 2 [16], the Sinkhorn divergence S 2 ε, 2 [50] and the energy distance ED [59, 60, 105], and we compare the performances when minimizing this functional via preconditioned GD vs. (vanilla) GD. We measure the convergence speed when using a fixed relative tolerance tol = 10 -3 , as 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 (ˆ µ, ν ) . As preconditioner, we use the one induced by h ∗ ( x ) = ( ∥ x ∥ a 2 +1) 1 /a -1 with a > 0 , which is well suited to minimize functionals which grow in ∥ x -x ∗ ∥ a/ ( a -1) near their minimum [119]. We set the step size τ = 1 for all the experiments. Then, we tune the parameter a very simply: for a given metric D and a profiling technology, we pick a random treatment and select a ∈ { 1 . 25 , 1 . 5 , 1 . 75 } by grid search, 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 = S 2 2 ,ε , we converge in 10 times less iterations while providing, on average, a better estimate of the treated population. We also compare our iterative (non parametric) approach with the use of a static (non parametric) map in Appendix G.4.

6 Conclusion

In this work, we extended two non-Euclidean optimization methods on R d to the Wasserstein space, generalizing W 2 -gradient descent to 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 in the Wasserstein geometry, they may be hard to verify in general. Different objectives such as the Sliced-Wasserstein distance or the Sinkhorn divergence, or alternative geometries to the Wasserstein-2 as studied in this work, require to derive specific computations on a case-by-case basis. We leave this investigation for future work.

Acknowledgments and Disclosure of Funding

Clément Bonet acknowledges the support of the center Hi! PARIS and of ANR PEPR PDE-AI. Adam David gratefully acknowledges funding by the BMBF 01|S20053B project SALE. Pierre-Cyril Aubin-Frankowski was funded by the FWF project P 36344-N. Anna Korba acknowledges the support of ANR-22-CE23-0030.

References

Martial Marie-Paul Agueh. Existence of Solutions to Degenerate Parabolic Equations via the Monge-Kantorovich Theory . Georgia Institute of Technology, 2002. (Cited on p. 1, 7)

Byeongkeun Ahn, Chiyoon Kim, Youngjoon Hong, and Hyunwoo J Kim. Invertible Monotone Operators for Normalizing Flows. Advances in Neural Information Processing Systems , 35: 16836-16848, 2022. (Cited on p. 28)

Kwangjun Ahn and Sinho Chewi. Efficient Constrained Sampling via the Mirror-Langevin Algorithm. Advances in Neural Information Processing Systems , 34:28405-28418, 2021. (Cited on p. 2, 5, 21, 29, 44)

David Alvarez-Melis and Nicolò Fusi. Dataset Dynamics via Gradient Flows in Probability Space. In International conference on machine learning , pages 219-230. PMLR, 2021. (Cited on p. 1)

Luigi Ambrosio, Nicola Gigli, and Giuseppe Savaré. Gradient Flows: in Metric Spaces and in the Space of Probability Measures . Springer Science & Business Media, 2005. (Cited on p. 1,

4, 23, 25, 38, 53)

Abdul Fatir Ansari, Ming Liang Ang, and Harold Soh. Refining Deep Generative Models via Discriminator Gradient Flow. In 9th International Conference on Learning Representations, ICLR , 2021. (Cited on p. 2, 20, 44)

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)

Hedy Attouch, Giuseppe Buttazzo, and Gérard Michaille. Variational Analysis in Sobolev and BV Spaces . Society for Industrial and Applied Mathematics, 2014. (Cited on p. 33, 35)

Pierre-Cyril Aubin-Frankowski, Anna Korba, and Flavien Léger. Mirror Descent with Relative Smoothness in Measure Spaces, with Application to Sinkhorn and EM. Advances in Neural Information Processing Systems , 35:17263-17275, 2022. (Cited on p. 2, 6, 21)

Heinz H Bauschke and Patrick L Combettes. Convex Analysis and Monotone Operator Theory in Hilbert Spaces . Springer, 2017. (Cited on p. 33)

Heinz H Bauschke, Jérôme Bolte, and Marc Teboulle. A descent lemma beyond Lipschitz gradient continuity: first-order methods revisited and applications. Mathematics of Operations Research , 42(2):330-348, 2017. (Cited on p. 3, 6, 20, 37)

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)

David M Blei, Alp Kucukelbir, and Jon D McAuliffe. Variational Inference: A Review for Statisticians. Journal of the American statistical Association , 112(518):859-877, 2017. (Cited on p. 1)

Clément Bonet, Nicolas Courty, François Septier, and Lucas Drumetz. Efficient Gradient Flows in Sliced-Wasserstein Space. Transactions on Machine Learning Research , 2022. (Cited on p. 20)

Clément Bonet, Lucas Drumetz, and Nicolas Courty. Sliced-Wasserstein Distances and Flows on Cartan-Hadamard Manifolds. arXiv preprint arXiv:2403.06560 , 2024. (Cited on p. 2)

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)

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 p. 3, 23)

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)

Siwan Boufadène and François-Xavier Vialard. On the global convergence of Wasserstein gradient flow of the Coulomb discrepancy. arXiv preprint arXiv:2312.00800 , 2023. (Cited on p. 2)

Yann Brenier. Polar Factorization and Monotone Rearrangement of Vector-Valued Functions. Communications on pure and applied mathematics , 44(4):375-417, 1991. (Cited on p. 27)

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 p. 9, 10, 43)

Charlotte Bunne, Andreas Krause, and Marco Cuturi. Supervised Training of Conditional Monge Maps. Advances in Neural Information Processing Systems , 35:6859-6872, 2022. (Cited on p. 9)

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-6528. PMLR, 2022. (Cited on p. 1, 9)

Martin Burger, Matthias Erbar, Franca Hoffmann, Daniel Matthes, and André Schlichting. Covariance-modulated optimal transport and gradient flows. arXiv preprint arXiv:2302.07773 , 2023. (Cited on p. 21)

Guillaume Carlier and Chloé Jimenez. On Monge's Problem for Bregman-like Cost Functions. Journal of Convex Analysis , 14(3):647, 2007. (Cited on p. 6, 21)

J. A. Carrillo, M. DiFrancesco, A. Figalli, T. Laurent, and D. Slepˇ cev. Global-in-time weak measure solutions and finite-time aggregation for nonlocal interaction equations. Duke Mathematical Journal , 156(2):229 - 271, 2011. (Cited on p. 1)

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)

Giulia Cavagnari, Giuseppe Savaré, and Giacomo Enrico Sodini. A Lagrangian approach to totally dissipative evolutions in Wasserstein spaces. arXiv preprint arXiv:2305.05211 , 2023. (Cited on p. 5, 26)

Gong Chen and Marc Teboulle. Convergence Analysis of a Proximal-Like Minimization Algorithm Using Bregman Functions. SIAM Journal on Optimization , 3(3):538-543, 1993. (Cited on p. 6)

Zonghao Chen, Aratrika Mustafi, Pierre Glaser, Anna Korba, Arthur Gretton, and Bharath K Sriperumbudur. (De)-regularized Maximum Mean Discrepancy Gradient Flow. arXiv preprint arXiv:2409.14980 , 2024. (Cited on p. 1, 2)

Ziheng Cheng, Shiyue Zhang, Longlin Yu, and Cheng Zhang. Particle-based Variational Inference with Generalized Wasserstein Gradient Flow. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. (Cited on p. 2, 7, 21)

Sinho Chewi, Thibaut Le Gouic, Chen Lu, Tyler Maunu, Philippe Rigollet, and Austin Stromme. Exponential Ergodicity of Mirror-Langevin Diffusions. Advances in Neural Information Processing Systems , 33:19573-19585, 2020. (Cited on p. 2, 21, 44)

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)

Sinho Chewi, Jonathan Niles-Weed, and Philippe Rigollet. Statistical Optimal Transport. arXiv preprint arXiv:2407.18163 , 2024. (Cited on p. 23)

Lenaic Chizat. Sparse Optimization on Measures with Over-parameterized Gradient Descent. Mathematical Programming , 194(1):487-532, 2022. (Cited on p. 5)

Lenaic Chizat and Francis Bach. On the Global Convergence of Gradient Descent for Overparameterized Models using Optimal Transport. Advances in neural information processing systems , 31, 2018. (Cited on p. 1)

Dario Cordero-Erausquin. Transport inequalities for log-concave measures, quantitative forms, and applications. Canadian Journal of Mathematics , 69(3):481-501, 2017. (Cited on p. 21)

Marco Cuturi, Laetitia Meng-Papaxanthos, 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)

Marco Cuturi, Michal Klein, and Pierre Ablin. Monge, Bregman and Occam: Interpretable Optimal Transport in High-Dimensions with Feature-Sparse 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. 7, 9)

Mathieu Dagréou, Pierre Ablin, Samuel Vaiter, and Thomas Moreau. How to compute Hessian-vector products? In ICLR Blogposts 2024 , 2024. (Cited on p. 39)

Nabarun Deb, Young-Heon Kim, Soumik Pal, and Geoffrey Schiebinger. Wasserstein Mirror Gradient Flow as the Limit of the Sinkhorn Algorithm. arXiv preprint arXiv:2307.16421 , 2023. (Cited on p. 5, 21, 29, 56)

Sylvain Delattre and Nicolas Fournier. On the Kozachenko-Leonenko entropy estimator. Journal of Statistical Planning and Inference , 185:69-93, 2017. (Cited on p. 44)

Michael Ziyang Diao, Krishna Balasubramanian, Sinho Chewi, and Adil Salim. Forwardbackward Gaussian variational inference via JKO in the Bures-Wasserstein Space. In International Conference on Machine Learning , pages 7960-7991. PMLR, 2023. (Cited on p. 9, 29, 37, 41, 42, 53)

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 on p. 2, 7, 21)

Chao Du, Tianbo Li, Tianyu Pang, Shuicheng Yan, and Min Lin. Nonparametric Generative Modeling with Conditional Sliced-Wasserstein Flows. In International Conference on Machine Learning (ICML) , 2023. (Cited on p. 2)

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)

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)

Luca Eyring, Dominik Klein, Théo Uscidda, Giovanni Palla, Niki Kilbertus, Zeynep Akata, and Fabian J Theis. Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation. In The Twelfth International Conference on Learning Representations , 2024. (Cited on p. 9)

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 p. 20, 44)

Jean Feydy, Thibault Séjourné, Franç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)

Nic Fishman, Leo Klarner, 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)

Nic Fishman, Leo Klarner, Emile Mathieu, Michael Hutchinson, and Valentin De Bortoli. Metropolis Sampling for Constrained Diffusion Models. Advances in Neural Information Processing Systems , 36, 2024. (Cited on p. 44)

Jean-Yves Franceschi, Mike Gartrell, Ludovic Dos Santos, Thibaut Issenhuth, Emmanuel de Bézenac, Mickaël Chen, and Alain Rakotomamonjy. Unifying GANs and Score-Based Diffusion as Generative Particle Models. Advances in Neural Information Processing Systems , 36, 2023. (Cited on p. 1)

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)

Wilfrid Gangbo and Adrian Tudorascu. On differentiability in the Wasserstein space and wellposedness for Hamilton-Jacobi equations. Journal de Mathématiques Pures et Appliquées , 125:119-174, 2019. (Cited on p. 24)

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 , 19(1):412-441, 2020. (Cited on p. 7, 21, 32)

Xin Guo, Johnny Hong, and Nan Yang. Ambiguity set and learning via Bregman and Wasserstein. arXiv preprint arXiv:1705.08056 , 2017. (Cited on p. 21)

Gabriele Gut, Markus Herrmann, and Lucas Pelkmans. Multiplexed protein maps link subcellular organization to cellular state. Science (New York, N.Y.) , 361, 08 2018. (Cited on p. 10)

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)

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 p. 2, 10, 43)

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. Advances in neural information processing systems , 33:6840-6851, 2020. (Cited on p. 9)

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)

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 ThirtySecond 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. 1)

Qijia Jiang. Mirror Langevin Monte Carlo: the Case Under Isoperimetry. Advances in Neural Information Processing Systems , 34:715-725, 2021. (Cited on p. 2)

Yiheng Jiang, Sinho Chewi, and Aram-Alexandre Pooladian. Algorithms for mean-field variational inference via polyhedral optimization in the Wasserstein space. arXiv preprint arXiv:2312.02849 , 2023. (Cited on p. 37)

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 on p. 1)

Mohammad Reza Karimi, Ya-Ping Hsieh, and Andreas Krause. Sinkhorn Flow: A ContinuousTime Framework for Understanding and Generalizing the Sinkhorn Algorithm. arXiv preprint arXiv:2311.16706 , 2023. (Cited on p. 2, 21)

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)

Dominik Klein, Théo Uscidda, Fabian Theis, and Marco Cuturi. Entropic (Gromov) Wasserstein Flow Matching with GENOT. arXiv preprint arXiv:2310.09254 , 2023. (Cited on p. 9)

Michal Klein, Aram-Alexandre Pooladian, Pierre Ablin, Eugène Ndiaye, Jonathan Niles-Weed, and Marco Cuturi. Learning elastic costs to shape monge displacements, 2023. (Cited on p. 7)

Anna Korba, Adil Salim, Michael Arbel, Giulia Luise, and Arthur Gretton. A Non-Asymptotic Analysis for Stein Variational Gradient Descent. Advances in Neural Information Processing Systems , 33:4672-4682, 2020. (Cited on p. 6, 32)

Anna Korba, Pierre-Cyril Aubin-Frankowski, Szymon Majewski, and Pierre Ablin. Kernel Stein Discrepancy Descent. In International Conference on Machine Learning , pages 57195730. PMLR, 2021. (Cited on p. 2, 24)

Marc Lambert, Sinho Chewi, Francis Bach, Silvère Bonnabel, and Philippe Rigollet. Variational inference via Wasserstein gradient flows. Advances in Neural Information Processing Systems , 35:14434-14447, 2022. (Cited on p. 9, 42)

Nicolas Lanzetti, Saverio Bolognani, and Florian Dörfler. First-Order Conditions for Optimization in the Wasserstein Space. arXiv preprint arXiv:2209.12197 , 2022. (Cited on p. 3, 4, 23, 48)

Emanuel Laude and Panagiotis Patrinos. Anisotropic Proximal Gradient. arXiv preprint arXiv:2210.15531 , 2022. (Cited on p. 7, 20)

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)

Flavien Léger and Pierre-Cyril Aubin-Frankowski. Gradient Descent with a General Cost. arXiv preprint arXiv:2305.04917 , 2023. (Cited on p. 2)

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)

Ruilin Li, Molei Tao, Santosh S Vempala, and Andre Wibisono. The Mirror Langevin Algorithm Converges with Vanishing Bias. In International Conference on Algorithmic Learning Theory , pages 718-742. PMLR, 2022. (Cited on p. 2)

Wuchen Li. Transport Information Bregman Divergences. Information Geometry , 4(2): 435-470, 2021. (Cited on p. 4, 21, 55, 56)

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)

Guan-Horng Liu, Tianrong Chen, Evangelos Theodorou, and Molei Tao. Mirror Diffusion Models for Constrained and Watermarked Generation. Advances in Neural Information Processing Systems , 36, 2024. (Cited on p. 2, 44)

Qiang Liu. Stein Variational Gradient Descent as Gradient Flow. Advances in neural information processing systems , 30, 2017. (Cited on p. 7, 20, 32)

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)

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 on Learning Representations , 2023. (Cited on p. 9)

Antoine Liutkus, Umut Simsekli, Szymon Majewski, Alain Durmus, and Fabian-Robert Stöter. Sliced-Wasserstein Flows: Nonparametric Generative Modeling via Optimal Transport and Diffusions. In International Conference on Machine Learning , pages 4104-4113. PMLR, 2019. (Cited on p. 1, 2)

Mohammad Lotfollahi, F Alexander Wolf, and Fabian J Theis. scGen predicts single-cell perturbation responses. Nature methods , 16(8):715-721, 2019. (Cited on p. 10)

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-354, 2018. (Cited on p. 2, 3, 5, 6, 8, 9, 20, 22, 35, 36, 40, 55)

Chris J Maddison, Daniel Paulin, Yee Whye Teh, and Arnaud Doucet. Dual Space Preconditioning for Gradient Descent. SIAM Journal on Optimization , 31(1):991-1016, 2021. (Cited on p. 2, 3, 5, 7, 8, 20)

Song Mei, Andrea Montanari, and Phan-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 p. 1)

Pierre Monmarché and Julien Reygner. Local convergence rates for Wasserstein gradient flows and McKean-Vlasov equations with multiple stationary solutions. arXiv preprint arXiv:2404.15725 , 2024. (Cited on p. 5)

Arkadij Semenoviˇ c Nemirovskij and David Borisovich Yudin. Problem complexity and method efficiency in optimization . Wiley, 1983. (Cited on p. 2, 20)

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.04613 , 2024. (Cited on p. 2)

Maxence Noble, Valentin De Bortoli, and Alain Durmus. Unbiased constrained sampling with Self-Concordant Barrier Hamiltonian Monte Carlo. In Thirty-seventh Conference on Neural Information Processing Systems , 2023. (Cited on p. 44)

Felix Otto. Double degenerate diffusion equations as steepest descent . Citeseer, 1996. (Cited on p. 1)

Felix Otto. The Geometry of Dissipative Evolution Equations: the Porous Medium Equation. Communications in Partial Differential Equations , 26(1-2):101-174, 2001. (Cited on p. 2)

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):361-400, 2000. (Cited on p. 24)

Guy Parker. Some Convexity Criteria for Differentiable Functions on the 2-Wasserstein Space. arXiv preprint arXiv:2306.09120 , 2023. (Cited on p. 5, 26)

Juan Peypouquet. Convex Optimization in Normed Spaces: Theory, Methods and Examples . Springer, 2015. (Cited on p. 22, 33, 54, 55)

Gabriel Peyré. Entropic approximation of Wasserstein gradient flows. SIAM Journal on Imaging Sciences , 8(4):2323-2351, 2015. (Cited on p. 20)

Aram-Alexandre Pooladian and Jonathan Niles-Weed. Entropic estimation of optimal transport maps. arXiv preprint arXiv:2109.12004 , 2021. (Cited on p. 42, 43)

Julien Rabin, Gabriel Peyré, Julie Delon, and Marc Bernot. Wasserstein Barycenter and its Application to Texture Mixing. In 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. 2, 43)

Cale Rankin and Ting-Kam Leonard Wong. Bregman-Wasserstein Divergence: Geometry and Applications. arXiv preprint arXiv:2302.05833 , 2023. (Cited on p. 6, 21, 29)

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)

Maria L. Rizzo and Gábor J. Székely. Energy Distance. WIREs Computational Statistics , 8(1): 27-38, 2016. (Cited on p. 10)

Gareth O Roberts and Richard L Tweedie. Exponential convergence of Langevin distributions and their discrete approximations. Bernoulli , pages 341-363, 1996. (Cited on p. 1)

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)

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)

Adil Salim, Anna Korba, and Giulia Luise. The Wasserstein Proximal Gradient Algorithm. Advances in Neural Information Processing Systems , 33:12356-12366, 2020. (Cited on p. 37, 38, 42)

Filippo Santambrogio. Optimal Transport for Applied Mathematicians , volume 55. Springer, 2015. (Cited on p. 4, 7, 23)

Filippo Santambrogio. Euclidean, metric, and Wasserstein gradient flows: an overview. Bulletin of Mathematical Sciences , 7:87-154, 2017. (Cited on p. 1)

Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-Transport Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming. Cell , 176(4), 2019. (Cited on p. 9)

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)

Jiaxin Shi, Chang Liu, and Lester Mackey. Sampling with Mirrored Stein Operators. In International Conference on Learning Representations , 2022. (Cited on p. 2, 7, 32, 44)

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations. In International Conference on Learning Representations , 2021. (Cited on p. 9)

Vishwak Srinivasan, Andre Wibisono, and Ashia Wilson. Fast sampling from constrained spaces using the Metropolis-adjusted Mirror Langevin algorithm. In Proceedings of Thirty Seventh Conference on Learning Theory , volume 247, pages 4593-4635, 2024. (Cited on p. 2, 44)

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.05127 , 2023. (Cited on p. 5, 26, 38)

Fuchou Tang, Catalin Barbacioru, Yangzhou Wang, Ellen Nordman, Clarence Lee, Nanlan Xu, Xiaohui Wang, John Bodeau, Brian B Tuch, Asim Siddiqui, et al. mRNA-Seq wholetranscriptome analysis of a single cell. Nature methods , 6(5):377-382, 2009. (Cited on p. 10)

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)

Antonio Terpin, Nicolas Lanzetti, and Florian Dörfler. Learning Diffusion at Lightspeed. Advances in Neural Information Processing Systems , 2024. (Cited on p. 1)

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)

Théo Uscidda and Marco Cuturi. The Monge Gap: A Regularizer to Learn All Transport Maps. In International Conference on Machine Learning , pages 34709-34733. PMLR, 2023. (Cited on p. 9)

Quang Van Nguyen. Forward-backward splitting with Bregman distances. Vietnam Journal of Mathematics , 45(3):519-539, 2017. (Cited on p. 6)

Cédric Villani. Topics in Optimal Transportation , volume 58. American Mathematical Soc., 2003. (Cited on p. 24)

Cédric Villani. Optimal Transport: Old and New , volume 338. Springer, 2009. (Cited on p. 51)

Yifei Wang and Wuchen Li. Information Newton's Flow: Second-Order Optimization Method in Probability Space. arXiv preprint arXiv:2001.04341 , 2020. (Cited on p. 5, 21, 24, 25, 29)

Yifei Wang, Peng Chen, and Wuchen Li. Projected Wasserstein gradient descent for highdimensional Bayesian inference. SIAM/ASA Journal on Uncertainty Quantification , 10(4): 1513-1532, 2022. (Cited on p. 20, 44)

Yifei Wang, Peng Chen, Mert Pilanci, and Wuchen Li. Optimal Neural Network Approximation of Wasserstein Gradient Direction via Convex Optimization. arXiv preprint arXiv:2205.13098 , 2022. (Cited on p. 20, 44)

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, 4, 31, 42)

Andre Wibisono. Proximal Langevin Algorithm: Rapid Convergence under Isoperimetry. arXiv preprint arXiv:1911.01469 , 2019. (Cited on p. 5, 29)

Lantian Xu, Anna Korba, and Dejan Slepcev. Accurate Quantization of Measures via Interacting Particle-based Optimization. In International Conference on Machine Learning , pages 24576-24595. PMLR, 2022. (Cited on p. 7)

Kelvin Shuangjian Zhang, Gabriel Peyré, Jalal Fadili, and Marcelo Pereyra. Wasserstein Control of Mirror Langevin Monte Carlo. In Conference on Learning Theory , pages 38143841. PMLR, 2020. (Cited on p. 2)

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)

Appendix

The appendix is organized as follows. In Appendix A, we discuss related works. In Appendix B and Appendix C, we provide mathematical background respectively on L 2 and on the Wasserstein space. In Appendix D, we provide complementary results for the mirror descent scheme on the Wasserstein space. In Appendix E, we discuss the relative convexity and smoothness between functionals. In Appendix F, we study the Bregman proximal gradient scheme, deriving a 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

A Related works

Mirror descent on R d . Mirror descent has been introduced by Nemirovskij and Yudin [92] to 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 et al. [11], Lu et al. [88] provided convergence guarantees by assuming relative smoothness and convexity.

Preconditioned gradient descent on R d . 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 et al. [68] underlined connections with the mirror descent, and introduced an accelerated version of the preconditioned gradient descent. Laude and Patrinos [75], Laude et al. [76] studied a generalized version of this algorithm by minimizing an anisotropic upper bound and supposing anisotropic smoothness of the objective. In particular, their analysis for the descent lemma is also valid for a 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 are based on time-discretizations of gradient flows with respect to optimal transport metrics, but different than the Wasserstein-2 distance.

To simplify the computation of the backward scheme, Peyré [100] added an entropic regularization into the JKO scheme while Bonet et al. [14] considered using the Sliced-Wasserstein distance instead. More recently, Rankin and Wong [104] suggested using Bregman divergences e.g. when geodesic distances are not known in 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 at each step, which is 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 a reproducing kernel Hilbert space (RKHS), an update in closed-form can be obtained, which is given by the SVGD algorithm [83, 84]. This algorithm

can also be seen as 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 a 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 Kalman-Wasserstein or Covariance Modulated gradient flow [24, 56].

Mirror descent with flat geometry. The space of probability distributions can be endowed with different metrics. When endowed with the Fisher-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. Aubin-Frankowski et al. [9] studied the mirror descent in 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 et al. [41] as a special case.

Bregman divergence on P 2 ( R d ) . Several works introduced Bregman divergences on P 2 ( R d ) . Carlier and Jimenez [25] first studied the existence of Monge maps for the OT problem with Bregman costs c ( x, y ) = d V ( x, y ) and symmetrized Bregman costs c ( x, y ) = d V ( x, y ) + d V ( y, x ) . For Bregman costs, the resulting OT problem was named the Bregman-Wasserstein divergence and its properties were studied in [37, 57, 103]. The Bregman-Wasserstein 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 a geodesically strictly convex F : P 2 ( R d ) → R as

where T µ ν is the OT map between ν and µ w.r.t W 2 . The Bregman divergence used in our work and as defined in Definition 1 is more general as it allows using more general maps and contains as special case (16). Li [80] studied properties of this Bregman divergence for different functionals F and provided closed-forms for one-dimensional distributions or Gaussian, but did not use it to define a mirror scheme.

Mirror descent on P 2 ( R d ) . Deb et al. [41] defined a mirror flow by using the continuous formulation. They focused on KL objectives with Bregman potential ϕ ( µ ) = 1 2 W 2 2 ( µ, ν ) with some reference measure ν ∈ P 2 ( R d ) , and defined the flow as the solution of

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 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 2 ( µ )

B.1 Differential calculus on L 2 ( µ )

We recall some differentiability definitions on the Hilbert space L 2 ( µ ) for µ ∈ P 2 ( R d ) . Let ϕ : L 2 ( µ ) → R . We start by recalling the notions of Gâteaux and Fréchet derivatives.

Definition 4. A function ϕ : L 2 ( µ ) → R is said to be Gâteaux differentiable at T if there exists an operator ϕ ′ (T) : L 2 ( µ ) → R such that for any direction h ∈ L 2 ( µ ) ,

and ϕ ′ (T) is a linear function. The operator ϕ ′ (T) is called the Gâteaux derivative of ϕ at T and if it exists, it is unique.

Definition 5. The Fréchet derivative of ϕ at T ∈ L 2 ( µ ) in the direction h ∈ L 2 ( µ ) , denoted δϕ (T , h ) , is defined implicitly by

If ϕ is Fréchet differentiable, then it is also Gâteaux differentiable, and both derivatives agree, i.e. for all T , h ∈ L 2 ( µ ) , δϕ (T , h ) = ϕ ′ (T)( h ) [99, Proposition 1.26].

As a brief comment on these notions in the context of convexity, if the subdifferential of a convex f at x contains a single element then it is the Gâteaux derivative and we have an inequality f ( y ) ≥ f ( x ) + ⟨∇ f ( x ) , y -x ⟩ . Instead Fréchet différentiability gives an equality (19) corresponding to a series expansion.

Moreover, since L 2 ( µ ) 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 ∇ ϕ ∈ L 2 ( µ ) such that for all h ∈ L 2 ( µ ) , δϕ (T , h ) = ⟨∇ ϕ (T) , h ⟩ L 2 ( µ ) (resp. ϕ ′ (T)( h ) = ⟨∇ ϕ (T) , h ⟩ L 2 ( µ ) ).

B.2 Convexity on L 2 ( µ )

Let ϕ : L 2 ( µ ) → R be Gâteaux differentiable. We recall that ϕ is convex if for all t ∈ [0 , 1] , T , S ∈ L 2 ( µ ) , which is equivalent by [99, Proposition 3.10] with

We now present equivalent definitions of the relative smoothness and relative convexity, which is the equivalent of [88, Proposition 1.1].

Proposition 8. Let ψ, ϕ : L 2 ( µ ) → R be convex and Gâteaux differentiable functions. The following conditions are equivalent:

(a1) ψ is β -smooth relative to ϕ

(a2) βϕ -ψ is a convex function on L 2 ( µ )

(a3) If twice Gâteaux differentiable, ⟨∇ 2 ψ (T)S , S ⟩ L 2 ( µ ) ≤ β ⟨∇ 2 ϕ (T)S , S ⟩ L 2 ( µ ) for all T , S ∈ L 2 ( µ )

(a4) ⟨∇ ψ (T) -∇ ψ (S) , T -S ⟩ L 2 ( µ ) ≤ β ⟨∇ ϕ (T) -∇ ϕ (S) , T -S ⟩ L 2 ( µ ) for all T , S ∈ L 2 ( µ ) .

The following conditions are equivalent:

(b1) ψ is α -convex relative to ϕ

(b2) ψ -αϕ is a convex function on L 2 ( µ )

(b3) If twice differentiable, ⟨∇ 2 ψ (T)S , S ⟩ L 2 ( µ ) ≥ α ⟨∇ 2 ϕ (T)S , S ⟩ L 2 ( µ ) for all T , S ∈ L 2 ( µ )

Proof. We do it only for the smoothness. It holds likewise for the convexity.

(a1) ⇐⇒ (a2) :

For the rest of the equivalences, we apply [99, Proposition 3.10]. Indeed, βϕ -ψ convex is equivalent to

which gives the equivalence between (a2) and (a4) . And if ψ and ϕ are twice differentiables, it is also equivalent to

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 suband super-differentials.

Definition 6 (Wasserstein sub- and super-differential [17, 74]) . Let F : P 2 ( R d ) → ( -∞ , + ∞ ] lower semicontinuous and denote D ( F ) = { µ ∈ P 2 ( R d ) , F ( µ ) < ∞} . Let µ ∈ D ( F ) . Then, a map ξ ∈ L 2 ( µ ) belongs to the subdifferential ∂ -F ( µ ) of F at µ if for all ν ∈ P 2 ( R d ) ,

Similarly, ξ ∈ L 2 ( µ ) belongs to the superdifferential ∂ + F ( µ ) of F at µ if -ξ ∈ ∂ -( -F )( µ ) .

Then, we say that a functional is Wasserstein differentiable if it admits sub and super differentials which coincide.

̸

Definition 7 (Wasserstein differentiability, Definition 2.3 in [74]) . A functional F : P 2 ( R d ) → R is Wasserstein differentiable at µ ∈ P 2 ( R d ) if ∂ -F ( µ ) ∩ ∂ + F ( µ ) = ∅ . In this case, we say that ∇ W 2 F ( µ ) ∈ ∂ -F ( µ ) ∩ ∂ + F ( µ ) is a Wasserstein gradient of F at µ , satisfying for any ν ∈ P 2 ( R d ) , γ ∈ Π o ( µ, ν ) ,

Recall that the tangent space of P 2 ( R d ) at µ ∈ P 2 ( R d ) is defined as where the closure is taken in L 2 ( µ ) , see Ambrosio et al. [5, Definition 8.4.1]. Lanzetti et al. [74, Proposition 2.5] showed that if F is Wasserstein differentiable, then there is always a 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 µ, ν ∈ P 2 ( R d ) , γ ∈ Π( µ, ν ) any coupling and let F : P 2 ( R d ) → R be Wasserstein differentiable at µ with Wasserstein gradient ∇ W 2 F ( µ ) ∈ T µ P 2 ( R d ) . Then,

Under regularity assumptions, the Wasserstein gradient of F can be computed in practice using the first variation δ F δµ [110, Definition 7.12], which is defined, if it exists, as the unique function (up to a constant) such that, for χ satisfying ∫ d χ = 0 ,

Then the Wasserstein gradient can be computed as ∇ W 2 F ( µ ) = ∇ δ F δµ ( µ ) , see e.g. [34, Proposition 5.10].

We now show that we can relate the Fréchet derivative of ˜ F µ (T) := F (T # µ ) with the Wasserstein gradient of F belonging to the tangent space of P 2 ( R d ) at µ .

Proposition 10. Let F : P 2 ( R d ) → R ∪ { + ∞} be a Wasserstein differentiable functional on D ( F ) . Let µ ∈ P 2 ( R d ) and ˜ F µ (T) = F (T # µ ) for all T ∈ D ( ˜ F µ ) . Then, ˜ F µ is Fréchet differentiable, and for all S ∈ D ( ˜ F µ ) , ∇ ˜ F µ (S) = ∇ W 2 F (S # µ ) ◦ 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 L 2 ( µ ) but a lifting L 2 (Ω; R d ) of measures on random variables. They should not be confused.

C.2 Wasserstein Hessians

A natural object of interest in the context of optimization over the Wasserstein space is the Hessian of the objective F , which we define below, according to the original definitions of [97, Section 3] and [124, Chapter 8]. This notion is usually defined along Wasserstein geodesics.

Definition 8 (Chapter 8 in [124]) . The Wasserstein Hessian of F , denoted H F µ is an operator over T µ P 2 ( R d ) verifying ⟨ H F µ v 0 , v 0 ⟩ L 2 ( µ ) = d 2 d t 2 F ( ρ t ) ∣ ∣ t =0 if ( ρ t , v t ) t ∈ [0 , 1] is a Wasserstein geodesic starting at µ .

If µ is absolutely continuous, Wassertein geodesics starting from µ are curves of the form ρ t = (Id + t ∇ ψ ) # µ for ψ ∈ C ∞ c ( R d ) . Using this fact, one can compute the Wasserstein Hessians of Kullback-Leibler 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. µ t = (S + tv ) # µ with S , v ∈ L 2 ( µ ) . Thus, we define analogously the Hessian along such curves.

Definition 9. We define the Hessian operator H F µ,t : L 2 ( µ ) → L 2 ( µ ) as the operator satisfying for all t ∈ [0 , 1] ,

Note that the Hessian H F µ,t is taken at time t but that the vector field v ∈ L 2 ( µ ) is in the tangent space at t = 0 , hence the discrepancy with Definition 8 besides the fact that we can have S = Id .

̸

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 µ t = (S + tv ) # µ with S , v ∈ L 2 ( µ ) . We first provide a lemma computing the derivative of the Wasserstein gradient.

Lemma 11. Let F : P 2 ( R d ) → R be twice continuously differentiable and assume that δ δµ ∇ δ F δµ ( µ ) = ∇ δ 2 F δµ 2 ( µ ) for all µ ∈ P 2 ( R d ) . Let µ ∈ P 2 ( R d ) and for all t ∈ [0 , 1] , µ t = (T t ) # µ where T t is differentiable w.r.t. t with dT t d t ∈ L 2 ( µ ) . Then,

This allows us to define a closed-form for H F µ,t .

Proposition 12. Under the same assumptions as in Lemma 11, let µ t = (T t ) # µ with T t = S + tv , S , v ∈ L 2 ( µ ) , then H F µ,t : L 2 ( µ ) → L 2 ( µ ) (as defined in Definition 9) is defined for all v ∈ L 2 ( µ ) , x ∈ R d as:

While H F µ,t and H F µ t differ in general, in some simple cases their relation boils down to composition with a pushforward. For instance, if S is invertible, we can write µ t as a curve starting from S # µ with a velocity field v ◦ S -1 , i.e. µ t = (Id + tv ◦ S -1 ) # S # µ . Thus, we recover the original definition of the Wasserstein Hessian at t = 0 as H F µ, 0 = H F S # µ . However, in general, this does not need to be the case.

Similarly, if T t is invertible for all t , setting v t = v ◦ T -1 t , we can write

The last line is obtained by leveraging the closed form in Proposition 12 and that µ t = (T t ) # µ , as for all x ∈ supp( µ ) ,

and thus H F µ,t [ v ] ( T -1 t ( x ) ) = H F µ t [ v t ]( x ) .

Here are two examples of F satisfying δ δµ ∇ δ F δµ = ∇ δ 2 F δµ 2 for which Proposition 12 provides an expression of the Wasserstein Hessian.

Example 1 (Potential energy) . Let V ( µ ) = ∫ V d µ with V twice differentiable with bounded Hessian. Then, we have δ V δµ ( µ ) = V and δ 2 V δµ 2 ( µ ) = 0 (using (29) ). Thus, applying Proposition 12, we recover for µ t = (T t ) # µ ,

Example 2 (Interaction energy) . Let W ( µ ) = 1 2 ∫∫ W ( x -y ) d µ ( x )d µ ( y ) with W symmetric, twice differentiable and with bounded Hessian. Then, we have for all x, y ∈ R d , δ W δµ ( µ )( x ) = ( W ⋆ µ )( x ) and δ 2 W δµ 2 ( µ )( x, y ) = W ( x -y ) (see e.g. [126, Example 7]), and thus applying Proposition 12, for µ t = (T t ) # µ , the operator is

and

C.3 Convexity in Wasserstein space

We first recall the definition of α -convex functionals [5, Definition 9.1.1].

Definition 10. F is α -convex along geodesics if for all µ 0 , µ 1 ∈ P 2 ( R d ) ,

where ( µ t ) t ∈ [0 , 1] is a Wasserstein geodesic between µ 0 and µ 1 .

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 β G - F or F α G .

Proposition 13. Let F : P 2 ( R d ) → R be twice continuously differentiable. Let µ ∈ P 2 ( R d ) , T , S ∈ L 2 ( µ ) , µ t = ( T t ) # µ for all t ∈ [0 , 1] where T t = (1 -t )S + t T . Furthermore, denote for t 1 , t 2 ∈ [0 , 1] , ˜ µ t 1 → t 2 t = ( (1 -t )T t 1 + t T t 2 ) # µ . Then, the following statement are equivalent:

(c1) For all t 1 , t 2 , t ∈ [0 , 1] , F (˜ µ t 1 → t 2 t ) ≤ (1 -t ) F ( (T t 1 ) # µ )+ t F ( (T t 2 ) # µ ) , i.e. F is convex along t ↦→ µ t .

(c2) For all t 1 , t 2 ∈ [0 , 1] , we have d ˜ F µ (T t 2 , T t 1 ) ≥ 0 , i.e.

(c3) For all t 1 , t 2 ∈ [0 , 1] ,

(c4) For all s ∈ [0 , 1] , d 2 d t 2 F ( µ t ) ∣ ∣ ∣ t = s ≥ 0 .

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. an 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.

If we require the convexity and the smoothness to hold along any curve of the form µ t = ( (1 -t )S + t T) # µ for µ ∈ P 2 ( R d ) and T , S ∈ L 2 ( µ ) , then it coincides with the transport convexity and smoothness recently introduced by Tanaka [117, Definitions 4.1 and 4.5]. As by Proposition 1, δ ˜ F µ (S , T -S) = ⟨∇ W 2 F (S # µ ) ◦ S , T -S ⟩ L 2 ( µ ) , and thus the convexity of F on such a curve reads as follows

And for ˜ G µ (T) = 1 2 ∥ T ∥ L 2 ( µ ) , the β -smoothness of F relative to G expresses as

This type of convexity is actually a particular case of the notion of convexity along acceleration-free curves introduced by Parker [98] (also introduced by Cavagnari et al. [28] under the name of total convexity). The latter requires convexity to hold along any curve of the form µ t = ( (1 -t ) π 1 + tπ 2 ) # γ with γ ∈ Π( µ, ν ) , µ, ν ∈ P 2 ( R d ) and π 1 ( x, y ) = x , π 2 ( x, y ) = y . The transport convexity of Tanaka [117] is thus a particular case for couplings obtained through maps, i.e. γ = (T , S) # µ . Parker [98] notably showed that this notion of convexity is equivalent to the geodesic convexity for Wasserstein differentiable functionals.

We can also define the strict convexity using strict inequalities in Proposition 13(c1) -(c2) -(c3) , but not in (c4) as there are counter examples already for real functions. For instance, f : R → R , defined as f ( x ) = x 4 for all x ∈ R , is strictly convex but f ′′ (0) = 0 . Thus, for F ( µ ) = ∫ f d µ , choosing µ = δ 0 and T 0 = Id , by Example 1, we have that d 2 d t 2 F ( µ t ) ∣ ∣ t =0 = f ′′ (0) v (0) 2 = 0 . But F ( µ t ) = f ( t T 1 (0) ) < tf ( T 1 (0) ) since f is strictly convex and thus F is well strictly convex.

Finally, as we defined the relative α -convexity and β -smoothness of F relative to G using Bregman divergences in Definition 3, we can show that it is equivalent to F α G and β G - F being convex.

Proposition 14. Let F , G : P 2 ( R d ) → R be two differentiable functionals. Let µ ∈ P 2 ( R d ) , T , S ∈ L 2 ( µ ) and for all t ∈ [0 , 1] , µ t = (T t ) # µ with T t = (1 -t )S + t T . Then, F is α -convex (resp. β -smooth) relative to G along t ↦→ µ t if and only if F α G (resp. β G - F ) is convex along t ↦→ µ t .

Proof. By Definition 3, F is α -convex relative to G along t ↦→ µ t if for all s, t ∈ [0 , 1] , d ˜ F µ (T s , T t ) ≥ α d ˜ G µ (T s , T t ) . This is equivalent to d ˜ F µ -α ˜ G µ (T s , T t ) ≥ 0 , which is equivalent by Proposition 13 (c2) (and using Proposition 1) to F α G convex along t ↦→ µ t . The result for the β -smoothness follows likewise.

D Additional results on mirror descent

D.1 Optimal transport maps for mirror descent

Let ϕ : P 2 ( R d ) → R be a strictly convex functional along all acceleration-free curves µ t = (T t ) # µ , t ∈ [0 , 1] with T t = (1 -t )S + t T for T , S ∈ L 2 ( µ ) . Denote for µ ∈ L 2 ( µ ) , ϕ µ (T) = ϕ (T # µ ) . Since ϕ is strictly convex along all acceleration-free curves, by Proposition 13, for all T = S ∈ L 2 ( µ ) , d ϕ µ (T , S) > 0 and thus ϕ µ is strictly convex. Indeed, recall that

̸

where we used Proposition 1 for the computation of the gradient.

Let us now define for all µ, ν ∈ P 2 ( R d ) ,

This problem encompasses several previously considered objects, as discussed in more detail in Remark 1. Our motivation for introducing Equation (42) is to prove that for ϕ µ 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 , S) # µ ∈ Π(T # µ, S # µ ) , we have d ϕ µ (T , S) ≥ W ϕ (T # µ, S # µ ) . Then, for µ ∈ P 2 , ac ( R d ) , assuming that ∇ W 2 ϕ ( µ ) = ∇ ϕ µ (Id) 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 (T µ,ν ϕ µ , Id) # µ with T µ,ν ϕ µ = argmin T # µ = ν d ϕ µ (T , Id) . Moreover, if ν ∈ P 2 , ac ( R d ) , we also have that T µ,ν ϕ µ is invertible with inverse ¯ T ν,µ ϕ ν = argmin T # ν = µ d ϕ ν (Id , T) .

Proposition 15. Let µ ∈ P 2 , ac ( R d ) , ν ∈ P 2 ( R d ) and assume ∇ W 2 ϕ ( µ ) invertible. Then,

There exists a unique minimizer γ of (42) . Besides, there exists a uniquely determined µ -almost everywhere (a.e.) map T µ,ν ϕ µ : R d → R d such that γ = (T µ,ν ϕ µ , Id) # µ . Finally, there exists a convex function u : R d → R such that T µ,ν ϕ µ = ∇ u ◦ ∇ W 2 ϕ ( µ ) µ -a.e.

Assume further that ν ∈ P 2 , ac ( R d ) . Then there exists a uniquely determined ν -a.e. map ¯ T ν,µ ϕ ν : R d → R d such that γ = (Id , ¯ T ν,µ ϕ ν ) # ν . Moreover, there exists a convex function v : R d → R such that ¯ T ν,µ ϕ ν = ∇ W 2 ϕ ( µ ) -1 ◦ ∇ v ν -a.e., and T µ,ν ϕ µ ◦ ¯ T ν,µ ϕ ν = Id ν -a.e. and ¯ T ν,µ ϕ ν ◦ T µ,ν ϕ µ = Id µ -a.e.

As a corollary, W ϕ ( ν, µ ) = min T # µ = ν d ϕ µ (T , Id) = min T # ν = µ d ϕ ν (Id , T) .

Proof. 1. Observe that problem (42) is equivalent to

Then, since for any γ ∈ Π( ν, µ ) , ( Id , ∇ W 2 ϕ ( µ ) ) # γ ∈ Π ( ν, ∇ W 2 ϕ ( µ ) # µ ) , we have

Let µ ∈ P 2 , ac ( R d ) . Since ∇ W 2 ϕ ( µ ) is invertible, ∇ W 2 ϕ ( µ ) # µ ∈ P 2 , ac ( R d ) . By Brenier's theorem, there exists a convex function u such that ( ∇ u ) # ( ∇ W 2 ϕ ( µ )) # µ = ν and the optimal coupling is of the form ˜ γ ∗ = ( ∇ u, Id) # ∇ W 2 ϕ ( µ ) # µ . Let γ = ( ∇ u ◦ ∇ W 2 ϕ ( µ ) , Id) # µ ∈ Π( ν, µ ) , then

Thus, since γ ∈ Π( ν, µ ) , γ is an optimal coupling for (42).

2. We symmetrize the arguments. Assuming ν ∈ P 2 , ac ( R d ) and ∇ ϕ µ (Id) = ∇ W 2 ϕ ( µ ) invertible, by Brenier's theorem, there exists a convex function v such that ( ∇ v ) # ν = ∇ W 2 ϕ ( µ ) # µ (and such that ∇ u ◦ ∇ v = Id ν -a.e. and ∇ v ◦ ∇ u = Id ∇ W 2 ϕ ( µ ) # µ -a.e.) and the optimal coupling is of the form ˜ γ ∗ = (Id , ∇ v ) # ν . Let γ = (Id , ∇ W 2 ϕ ( µ ) -1 ◦ ∇ v ) # ν ∈ Π( ν, µ ) , then

Thus, since γ ∈ Π( ν, µ ) , γ is an optimal coupling for (42). Moreover, noting T µ,ν ϕ µ = ∇ u ◦∇ W 2 ϕ ( µ ) and ¯ T ν,µ ϕ ν = ∇ W 2 ϕ ( µ ) -1 ◦∇ v , we have µ -a.e., ¯ T ν,µ ϕ ν ◦ T µ,ν ϕ µ = ∇ W 2 ϕ ( µ ) -1 ◦∇ v ◦∇ u ◦∇ W 2 ϕ ( µ ) = Id and ν -a.e., T µ,ν ϕ µ ◦ ¯ T ν,µ ϕ ν = ∇ u ◦ ∇ W 2 ϕ ( µ ) ◦ ∇ W 2 ϕ ( µ ) -1 ◦ ∇ v = Id from the aforementioned consequences of Brenier's theorem.

We continue this section with additional results relative to the invertibility of mirror maps, which are required in Proposition 2. For a potential energy V ( µ ) = ∫ V d µ , since ∇ W 2 V = ∇ V , then ∇ W 2 V is invertible provided ∇ V is. This is the case e.g. for V strictly convex. We now state in the two next lemmas conditions for an interaction energy and for the negative entropy to satisfy the invertibility requirements.

Lemma 16. Let µ ∈ P 2 ( R d ) and let W : R d → R be even, ϵ -strongly convex for ϵ > 0 and differentiable. Then, for W ( µ ) = ∫∫ W ( x -y ) d µ ( x )d µ ( y ) , ∇ W 2 W ( µ ) is invertible.

Proof. On one hand, ∇ W 2 W ( µ ) = ∇ W ⋆ µ . Moreover, W ϵ -strongly convex is equivalent to

̸

which implies for all x, y, z ∈ R d , ⟨∇ W ( x -z ) -∇ W ( y -z ) , x -y ⟩ ≥ ϵ ∥ x -y ∥ 2 2 . By integrating with respect to µ , it implies

Thus, ∇ W ⋆ µ is ϵ -strongly monotone, and in particular invertible [2, Theorem 1].

Lemma 17. Let µ ∈ P 2 , ac ( R d ) such that its density is of the form ρ ∝ e -V with V : R d → R ϵ -strongly convex for ϵ > 0 . Then, for H ( µ ) = ∫ log ( ρ ( x ) ) d µ ( x ) with ρ the density of µ w.r.t the Lebesgue measure, ∇ W 2 H ( µ ) is invertible.

Proof. Let µ such distribution. Then, ∇ W 2 H ( µ ) = ∇ log ρ = -∇ V . Since V is ϵ -strongly convex, then ∇ V is ϵ -strongly monotone and in particular invertible [2, Theorem 1].

We conclude this section with a discussion of (42) with respect to related work.

Remark 1. The OT problem (42) recovers other OT costs for specific choices of ϕ . For instance, for ϕ µ (T) = 1 2 ∥ T ∥ 2 L 2 ( µ ) , it coincides with the squared Wasserstein-2 distance. And more generally, for ϕ V µ (T) = ∫ V ◦ T d µ , since by Lemma 31, for all T , S ∈ L 2 ( µ ) ,

where d V is the Euclidean Bregman divergence, i.e. for all x, y ∈ R d , d V ( x, y ) = V ( x ) -V ( y ) -⟨∇ V ( y ) , x -y ⟩ , W ϕ coincides with the Bregman-Wasserstein divergence [103]

D.2 Continuous formulation

Let ϕ : L 2 ( µ ) → R be pushforward compatible and superlinear. Introducing the (mirror) map φ ( µ ) = ∇ W 2 ϕ ( µ ) , we can write informally the mirror descent scheme (4) and its continuous-time counterpart when τ → 0 as

However, d d t φ ( µ t ) = d d t ∇ W 2 ϕ ( µ t ) = H ϕ µ t ( v t ) where H ϕ µ t : L 2 ( µ t ) → L 2 ( µ t ) is the Hessian operator (defined in Appendix C.2) such that d 2 d t 2 ϕ ( µ t ) = ⟨ H ϕ µ t ( v t ) , v t ⟩ L 2 ( µ t ) and v t ∈ L 2 ( µ t ) is a velocity field satisfying ∂ t µ t +div( µ t v t ) = 0 . Thus, the continuity equation corresponding to the Mirror Flow is given by

For ϕ V µ as Bregman potential, since H ϕ V µ ( v ) = ( ∇ 2 V ) v (see Appendix C.2), the flow is a solution of ∂ t µ t -div ( µ t ( ∇ 2 V ) -1 ∇ W 2 F ( µ t ) ) = 0 . For F ( µ ) = KL( µ || ν ) with ν ∝ e -U , this coincides 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 ϕ = F , this coincides with Information Newton's flows [126]. Note also that Deb et al. [41] defined mirror flows through the scheme τ → 0 of (51), but focused on F ( µ ) = KL( µ || ν ) and ϕ ( µ ) = 1 2 W 2 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 as 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-Wasserstein space, and to use the gradient on the Bures-Wasserstein space [43]. In particular, this space is a submanifold of P 2 , ac ( R d ) 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 ∈ S d ( R ) .

Interaction mirror scheme. Consider as Bregman potential an interaction energy ϕ µ (T) = 1 2 ∫∫ W ( ( T ( x ) -T ( x ′ ) ) d µ ( x )d µ ( x ′ ) . The mirror descent scheme (4) is given by

For the particular case W ( x ) = 1 2 ∥ x ∥ 2 2 , the scheme can be made more explicit as ∇ W ⋆ µ ( x ) = ∫ ∇ W ( x -y ) d µ ( y ) = ∫ ( x -y ) d µ ( y ) = x -m ( µ ) with m ( µ ) = ∫ y d µ ( y ) the expectation, and thus (53) translates as

On one hand, recall from Example 2 that the Hessian of ϕ is given, for µ ∈ P 2 ( R d ) , v ∈ L 2 ( µ ) , by

since ∇ 2 W = I d . On the other hand, the mirror descent scheme (54) can be written as, for all k ≥ 0 ,

Passing to the limit τ → 0 , we get

where d y t d t = d x t d t -d m ( µ t ) d t . Now, by setting v t ( x ) = d x t d t , by integration by part, we have

Combining the latter equation with (55), we obtain as expected that d y t d t = H ϕ µ t [ v t ]( x ) .

Negative entropy mirror scheme. Consider the negative entropy ϕ ( µ ) = ∫ log ( ρ ( x ) ) d µ ( x ) where d µ ( x ) = ρ ( x )d x and ϕ µ (T) = ϕ (T # µ ) . For such Bregman potential, the mirror scheme (4) can be written for all k ≥ 0 as

In general, this scheme is not tractable. Nonetheless, supposing that µ k = N ( m k , Σ k ) for all k ≥ 0 , the scheme translates as

For an objective functional F ( µ ) = H ( µ ) + V ( µ ) with V ( x ) = 1 2 x T Σ -1 x , the scheme is

Assuming m k = 0 for all k , we obtain the following update rule for the covariance matrices:

We illustrate this scheme in Figure 2.

For F ( µ ) = H ( µ ) , we obtain

Assuming m k = 0 for all k , for τ < 1 , we obtain the following update rule for the covariance matrices:

The continuous time analog of this scheme is thus µ t : t ↦→ N (0 , Σ t ) with Σ t = e 2 t Σ 0 and the negative entropy decreases along this curve as

where ( λ i ) i denote the eigenvalues of Σ 0 . This is much faster than the heat flow for which the negative entropy decreases as [129, Appendix E.2]

with the scheme given by [129, Example 6]

With our notations, the heat flow is the continuous time limit of the scheme (4) for the same objective F but for a quadratic Bregman potential ϕ µ (T) = 1 2 ∥ T ∥ 2 L 2 ( µ ) (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 ( µ ) = G ( µ ) + H ( µ ) where G ( µ ) = ∫ U d µ . Then, a natural choice of Bregman potential is also a functional of the form ϕ ( µ ) = Ψ( µ ) + H ( µ ) with Ψ( µ ) = ∫ V d µ , with U α -convex and β -smooth relative to V .

In that case, we obtain the smoothness of F relative to ϕ . Recall we denote ˜ F µ (T) = F (T # µ ) for T ∈ L 2 ( µ ) . Then for all T , S ∈ L 2 ( µ ) , we have α d ˜ Ψ µ (T , S) ≤ d ˜ G µ (T , S) ≤ β d ˜ Ψ µ (T , S) , hence

We now focus on the case where all measures are Gaussian in order to be able to compute a closedform, i.e. U ( x ) = 1 2 ( x -m ) T Σ -1 ( x -m ) , V ( x ) = 1 2 x T Λ -1 x and for all k ≥ 0 , µ k = N ( m k , Σ k ) . In this case, recall that ∇ log µ k ( x ) = -Σ -1 k ( x -m k ) . Then, at each step, the mirror descent scheme (4) writes for x k ∼ µ k , k ≥ 0 as

Thus, we get for the expectation that

We note that the latter update on the means coincides with the forward Euler method in the forwardbackward scheme, see (116) in Appendix F, which uses as Bregman potential ϕ = Ψ . Thus, the entropy does not affect the convergence towards the mean, which can be done simply by (preconditioned) gradient descent.

For the covariance part, we get

Now, supposing that all matrices commute, we get

Denoting the update on covariances is equivalent to

Thus, Σ k +1 = 1 2 ( C ± ( C 2 -4Λ 2 ) 1 2 ) .

D.4 Mirror scheme with non-pushforward compatible Bregman potentials

We study in this Section schemes for which the Bregman potential ϕ µ is not pushforward compatible, and thus for which we cannot apply Proposition 2 and thus Assumption 1 may not hold a priori. An example of such potential is ϕ µ (T) = ⟨ T , P µ T ⟩ L 2 ( µ ) where P µ : L 2 ( µ ) → L 2 ( µ ) is a linear autoadjoint and invertible operator. Since ∇ ϕ µ (T) = P µ T , taking the first order conditions, we obtain the following scheme:

In particular, this includes SVGD [71, 83, 84] if we pose P -1 µ T = ιS µ T with S µ : L 2 ( µ ) → H defined as S µ T = ∫ k ( x, · )T( x )d µ ( x ) which maps functions from L 2 ( µ ) to the reproducing kernel Hilbert space H with kernel k , and with ι : H → L 2 ( µ ) the inclusion operator that is the adjoint of S µ [71]. It also includes the Kalman-Wasserstein gradient descent [56] for which P -1 µ = ∫ ( x -m ( µ ) ) ⊗ ( x -m ( µ ) ) d µ ( x ) is the covariance matrix, where m ( µ ) = ∫ x d µ ( x ) .

More generally, for ϕ µ (T) = ∫ P µ ( V ◦ T)d µ , we can recover their mirrored version, including mirrored SVGD [113, 114], i.e. T k +1 = ∇ V ∗ ◦ ( ∇ V -τP -1 µ k ∇ W 2 F ( µ k ) ) .

Kalman-Wasserstein. We focus now on a particular choice of linear operator P µ . Namely, we take P µ T = C ( µ )T with C ( µ ) = (∫ ( x -m ( µ ) ) ⊗ ( x -m ( µ ) ) d µ ( x ) ) -1 the inverse of the covariance matrix. In this case, (76) corresponds to the discretization of the Kalman-Wasserstein gradient flow [56]. We now show that it satisfies Assumption 1. First, let us compute the Bregman divergence associated to ϕ :

For γ = (T , S) # µ , we can write

Moreover, the problem inf γ ∈ Π( α,β ) ∫ ∥ C ( µ ) 1 2 ( x -y ) ∥ 2 2 d γ ( x, y ) is equivalent to

which is a squared OT problem. Thus, it admits an 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 ∇ W F ( µ ) = P -1 µ ∇ W 2 F ( µ ) [46, 56], i.e. considering

where we assume that ∇ W F ( µ ) ∈ L 2 ( µ ) . In that case, using the first order conditions,

Then, for ϕ µ satisfying Assumption 1, the convergence will hold under relative smoothness and convexity assumptions similarly as for the analysis derived in Section 3.

E Relative convexity and smoothness

E.1 Relative convexity and smoothness between Fenchel transforms

In this Section, we show sufficient conditions to satisfy the inequalities assumed in Proposition 5 and Proposition 6 under the additional assumption that, for all k ≥ 0 , ˜ F µ k is superlinear, lower semicontinuous and strictly convex. In this case, we can show that ˜ F ∗ µ k is Gâteaux differentiable, and thus we can use the Bregman divergence of ˜ F ∗ µ k .

Lemma 18. Let ϕ : L 2 ( µ ) → R be a superlinear, lower semicontinuous and strictly convex function. Then, ϕ ∗ is Gâteaux differentiable.

Proof. Fix g ∈ L 2 ( µ ) . Notice that

So to prove there is a unique element in ∂ϕ ∗ ( g ) , we just need to show that, setting ϕ g ( f ) := -⟨ f, g ⟩ + ϕ ( f ) , the problem inf f ∈ L 2 ( µ ) ϕ g ( f ) has a unique solution. Under our assumptions, ϕ g is lower semicontinuous and strictly convex. Since ϕ is superlinear, ϕ g is coercive, i.e. lim ∥ f ∥→∞ ϕ g ( f ) = + ∞ . There thus exists a solution [8, Theorem 3.3.4], which is unique by strict convexity. Hence ∂ϕ ∗ ( g ) is reduced to a point, which is necessarily the Gâteaux derivative of ϕ ∗ at g .

This allows us to relate the Bregman divergence of ϕ ∗ to the Bregman divergence of ϕ .

Lemma 19. Let ϕ : L 2 ( µ ) → R be a proper superlinear and strictly convex differentiable function, then for all T , S ∈ L 2 ( µ ) , d ϕ ∗ ( ∇ ϕ (T) , ∇ ϕ (S) ) = d ϕ (S , T) .

Proof. By [99, Corollary 3.44], we have ϕ ∗ ( ∇ ϕ (T) ) = ⟨ T , ∇ ϕ (T) ⟩ L 2 ( µ ) -ϕ (T) for all T ∈ L 2 ( µ ) since ϕ is convex and differentiable. By Lemma 18, ϕ ∗ is invertible and by [10, Corollary 16.24], since ϕ is proper, lower semicontinuous and convex, then ( ∇ ϕ ) -1 = ∇ ϕ ∗ .

Thus, for all T , S ∈ L 2 ( µ ) ,

Finally, we can relate the relative convexity of ϕ relative to ψ ∗ by using an inequality between the Bregman divergences of ϕ and ψ . In particular, we recover the assumptions of Propositions 5 and 6 for ϕ h ∗ µ k that is β -smooth and α -convex relative to ˜ F ∗ µ k .

Proposition 20. Let ϕ, ψ : L 2 ( µ ) → R proper, superlinear, strictly convex and differentiable. ϕ is β -smooth (resp. α -convex) relative to ψ ∗ if and only if ∀ T , S ∈ L 2 ( µ ) , d ϕ ( ∇ ψ (T) , ∇ ψ (S) ) ≤ β d ψ (S , T) (resp. d ϕ ( ∇ ψ (T) , ∇ ψ (S) ) ≥ α d ψ (S , T) ).

Proof of Proposition 20. First, suppose that ϕ is β -smooth relative to ψ ∗ . Then, by definition,

In particular,

On the other hand, suppose for all T , S ∈ L 2 ( µ ) , d ϕ ( ∇ ψ (T) , ∇ ψ (S) ) ≤ β d ψ (S , T) . Then, by first using Lemma 19 and then the supposed inequality, we have for all T , S ∈ L 2 ( µ ) ,

Likewise, we can show that ϕ is α -convex relative to ψ if and only if d ϕ ( ∇ ψ (T) , ∇ ψ (S) ) ≥ α d ψ (S , T) for all T , S ∈ L 2 ( µ ) .

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 , ˜ F µ k is strictly convex, differentiable and superlinear. We first note that it implies that F µ k is convex along t ↦→ ( (1 -t )T k +1 + t Id ) # µ k . Moreover, by Lemma 18, ∇ ˜ F ∗ µ k is differentiable.

Note that this assumption is satisfied, e.g. by ϕ µ (T) = ∫ V ◦ T d µ for V η -strongly convex and differentiable. Indeed, in this case, ϕ µ is also η -strongly convex, and satisfies for all T , S ∈ L 2 ( µ ) ,

For S = 0 , and dividing by ∥ T ∥ L 2 ( µ ) the right term diverges to + ∞ when ∥ T ∥ L 2 ( µ ) → + ∞ , and thus lim ∥ T ∥ L 2 ( µ ) →∞ ϕ µ (T) / ∥ T ∥ L 2 ( µ ) = + ∞ , and ϕ µ is superlinear.

This assumption is also satisfied for interaction energies ϕ W µ (T) = ∫∫ W ( T( x ) -T( y ) ) d µ ( x )d µ ( y ) with W η -strongly convex, even and differentiable. Indeed, by strong convexity of W in 0, we have for all x, y ∈ R d ,

Integrating w.r.t. µ ⊗ µ , we get

and dividing by ∥ T ∥ L 2 ( µ ) , we get that ϕ W µ is superlinear.

For a curve t ↦→ µ t , we define F ∗ µ on µ t as F ∗ µ ( µ t ) := ˜ F ∗ µ (T t ) with ˜ F ∗ µ the convex conjugate of ˜ F µ in the L 2 ( µ ) sense. Then, we can apply Proposition 20, and we obtain that the inequality hypothesis of Proposition 5 is implied by the β -smoothness of ϕ h ∗ relative to F ∗ µ k along t ↦→ ( (1 -t ) ∇ W 2 F ( µ k ) + t ∇ W 2 F ( µ k +1 ) ◦ T k +1 ) # µ k since

where we used Proposition 1 to compute the gradient ∇ ˜ F µ k (T k +1 ) = ∇ W 2 F ( µ k +1 ) ◦ T k +1 . Similarly, the condition of Proposition 6

is implied by the α -convexity of ϕ h ∗ relative to F ∗ µ k along t ↦→ ( (1 -t ) ∇ W 2 F ( µ k ) + t ∇ W 2 F (T # µ k ) ◦ T ) # µ k .

These results are summarized in Proposition 7 and shown formally in Appendix H.7.

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 a metrizable topological space, and f : X → R ∪ { + ∞} be strictly convex, τ -lower semicontinuous and with one τ -compact sublevel set. Let x 0 ∈ X be the minimizer of f and take a sequence ( x n ) n ∈ N such that f ( x n ) → f ( x 0 ) . Then, ( x n ) n ∈ N τ -converges to x 0 .

Proof. The existence of the minimum is given by [8, Theorem 3.2.2]. For N large enough, ( x n ) n ≥ N lives in the τ -compact sublevel set of f , since x 0 belongs to it and f ( x 0 ) is minimal. We can then consider a subsequence τ -converging to some x ∗ . By τ -lower semicontinuity, we have f ( x 0 ) ≤ f ( x ∗ ) ≤ lim inf f ( x σ ( n ) ) = f ( x 0 ) , so f ( x 0 ) = f ( x ∗ ) and by strict convexity x 0 = x ∗ . Since all subsequences of ( x n ) n ≥ N converge to x ∗ and the space is metrizable, ( x n ) n ∈ N τ -converges to x 0 .

The typical case is when X is a Hilbert space and τ is the weak topology. One could wish to have strong convergence under a 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 semicontinuity [8, Theorem 3.3.3]. We cannot hope for convergence of the norm of x n to come for free, as the weak convergence would then imply the strong convergence.

E.2 Relative convexity and smoothness between functionals

Let U, V : R d → R be differentiable and convex functions. We recall that V is α -convex relative to U if [88]

Likewise, V is β -smooth relative to U if

Relative convexity and smoothness between potential energies. By Lemma 31, for Bregman potentials of the form ϕ µ (T) = ∫ V ◦ T d µ , the Bregman divergence can be written as

Thus, leveraging this result, we can show that relative convexity and smoothness of ϕ V µ relative to ϕ U µ is inherited by the relative convexity and smoothness of V relative to U .

Proposition 22. Let µ ∈ P 2 ( R d ) , ϕ µ (T) = ∫ V ◦ T d µ and ψ µ (T) = ∫ U ◦ T d µ where V, U : R d → R are C 1 . If V is α -convex (resp. β -smooth) relative to U : R d → R , then ϕ µ is α -convex (resp β -smooth) relative to ψ µ .

Proof. First, observe (Lemma 31) that

Let µ ∈ P 2 ( R d ) , T , S ∈ L 2 ( µ ) . If V is α -convex relatively to U , we have for all x, y ∈ R d ,

and hence by integrating on both sides with respect to µ ,

Likewise, we have the result for the β -smoothness.

Relative convexity and smoothness between interaction energies. Similarly, by Lemma 32, for Bregman potentials obtained through interaction energies, i.e. ϕ µ (T) = 1 2 ∫∫ W ( T( x ) -T( x ′ ) ) d µ ( x )d µ ( x ′ ) , then

It also allows to inherit the relative convexity and smoothness results from R d .

Proposition 23. Let µ ∈ P 2 ( R d ) , W,K : R d → R be symmetric, C 1 and convex. Let ϕ µ (T) = 1 2 ∫∫ W ( T( x ) -T( x ′ ) ) d µ ( x )d µ ( x ′ ) and ψ µ (T) = 1 2 ∫∫ K ( T( x ) -T( x ′ ) ) d µ ( x )d µ ( x ′ ) . If W is α -convex relative to K , then ϕ µ is α -convex relatively to ψ µ . Likewise, if W is β -smooth relatively to K , then ϕ µ is β -smooth relatively to ψ µ .

Proof. We use first Lemma 32 and then that W is α -convex relatively to K :

Likewise, we have the result for the β -smoothness.

Thus, in situations where the objective functional and the Bregman potential are of the same type and either potential energies or interaction energies, we only need to show the convexity and smoothness of the underlying potentials or interaction kernels. For instance, let V : R d → R be a twicedifferentiable convex function, such that ∥∇ 2 V ∥ op ≤ p r ( ∥ x ∥ 2 ) with p r a polynomial function of degree r and ∥ · ∥ op the operator norm. Then, by [88, Proposition 2.1], V is β -smooth relative to h where for all x ∈ R d , h ( x ) = 1 r +2 ∥ x ∥ r +2 2 + 1 2 ∥ x ∥ 2 2 .

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 β G - F or F α G is convex in order to show respectively the β -smoothness and α -convexity of F relative to G . For instance, we can use the characterization through Hessians, and thus we would aim at showing

along the right curve t ↦→ µ t .

For instance, consider an objective functional F ( µ ) = 1 2 ∫∫ W ( x -y ) d µ ( x )d µ ( x ′ ) and another functional G ( µ ) = ∫ V d µ . Then, by Example 1 and Example 2, we have, for µ t = (T t ) # µ and T t = S + tv ,

and

To show the conditions of Proposition 3, we need to take S = Id and v = T k +1 -Id , and to verify for t = s ∈ [0 , 1] the inequality, i.e.

For example, choosing W ( x ) = 1 2 ∥ x ∥ 2 2 , then ∇ 2 W = I d and F is β -smooth relative to G as long as ∇ 2 V ◦ T s ⪰ 1 β I d for any s ∈ [0 , 1] .

F Bregman proximal gradient scheme

In this section, we are interested into minimizing a functional F of the form F ( µ ) = G ( µ ) + H ( µ ) where G is smooth relative to some function ϕ and H is convex on L 2 ( µ ) . Different strategies can 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 et al. [43], Salim et al. [109] use Proximal Gradient algorithms. We focus here on the latter and generalize the Bregman Proximal Gradient algorithm [11], also known as the Forward-Backward scheme. It consists of alternating a forward step on G and then a backward step on H , i.e. for k ≥ 0 ,

The first step of our analysis is to show that this scheme is equivalent to

This is true under the condition that µ k ∈ P 2 , ac ( R d ) implies that ν k +1 ∈ P 2 , ac ( R d ) .

Proposition 24. Let ϕ µ be pushforward compatible, µ 0 ∈ P 2 , ac ( R d ) and assume that if µ k ∈ P 2 , ac ( R d ) then ν k +1 ∈ P 2 , ac ( R d ) . Then the schemes (104) and (105) are equivalent.

Proof.

We are now ready to state the convergence results for the proximal gradient scheme.

Proposition 25. Let µ 0 ∈ P 2 , ac ( R d ) , τ ≤ 1 β and F ( µ ) = G ( µ ) + H ( µ ) . Consider the iterates of the Bregman proximal gradient scheme (104) , equivalently (105) . Let k ≥ 0 . Assume ˜ H µ k is convex on L 2 ( µ k ) and G β -smooth relative to ϕ along t ↦→ ( (1 -t )Id + t ˜ T k +1 ) # µ k . Then, for all T ∈ L 2 ( µ k ) ,

Moreover, for T = Id ,

i.e., the scheme decreases the objective at each iteration. Additionally, let α ≥ 0 , ν ∈ P 2 ( R d ) and suppose that ϕ µ satisfies Assumption 1. If G is α -convex relative to ϕ along t ↦→ ( (1 -t )Id + t T µ k ,ν ϕ µ k ) # µ k , then for all k ≥ 1 ,

Proof.

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 ϕ ( µ ) = ∫ 1 2 ∥ · ∥ 2 2 d µ and G ( µ ) = ∫ U d µ with U (strictly) convex. We extend their result for G ( µ ) = ∫ U d µ and ϕ ( µ ) = ∫ V d µ for V strictly convex and U β -smooth relative to V .

Lemma 26. Let µ ∈ P 2 , ac ( R d ) , G ( µ ) = ∫ U d µ , ϕ µ (T) = ∫ V ◦ T d µ with V strongly convex and U β -smooth relative to V , and T = ∇ V ∗ ◦ ( ∇ V -τ ∇ U ) . Assume τ < 1 β , then T # µ ∈ P 2 , ac ( R d ) .

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 ∇ T | > 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 µ ∈ P 2 , ac ( R d ) , and denote ρ its density w.r.t the Lebesgue measure. For H ( µ ) = ∫ f ( ρ ( x ) ) d x where f : R → R is C 1 and satisfies f (0) = 0 , lim x → 0 xf ′ ( x ) = 0 and x ↦→ f ( x -d ) x d is convex and non-increasing on R + , then by [117, Theorem 4.2], H is convex along curves µ t = ( (1 -t )S + t T) # µ obtained with S and T with positive definite Jacobians. This is the case e.g. for f ( x ) = x log x , for which H corresponds to the negative entropy.

By Remark 3, to be able to apply the three-point inequality (that is necessary to obtain the descent lemma), we actually only need H to be convex along ( (1 -t ) ˜ T k +1 + t Id ) # µ k and along ( (1 -t ) ˜ T k +1 + t T µ k ,ν ϕ µ k ) # µ k for the convergence.

Gaussian target. In what follows, we focus on G ( µ ) = ∫ U d µ with U ( x ) = 1 2 ( x -m ) T Σ -1 ( x -m ) for Σ ∈ S ++ d ( R ) , H the negative entropy and with a Bregman potential of the form ϕ ( µ ) = ∫ V d µ with V ( x ) = 1 2 x T Λ -1 x . Moreover, we suppose µ 0 = N ( m 0 , Σ 0 ) . In this situation, each distribution µ k is also Gaussian, as the forward and backward steps are affine operations.

Assuming the covariances matrices are full rank, ˜ T k +1 is affine and its gradient is invertible. Moreover, by Proposition 15, T µ k ,ν ϕ µ k = ∇ u ◦ ∇ W 2 ϕ ( µ k ) for ∇ u an OT map between ∇ W 2 ϕ ( µ k ) # µ k and ν . Since each distribution is Gaussian, and ∇ W 2 ϕ ( µ k )( x ) = Λ -1 x 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 ( µ ) = ∫ U d µ with U ( x ) = 1 2 ( x -m ) T Σ -1 ( x -m ) , Σ ∈ S ++ d ( R ) , m ∈ R d , and H ( µ ) = ∫ log ( ρ ( x ) ) d µ ( x ) for d µ = ρ ( x )d x . For the Bregman potential, we will choose ϕ ( µ ) = ∫ V d µ for V ( x ) = 1 2 ⟨ x, Λ -1 x ⟩ . Recall that the forward step reads as

Since ∇ V ( x ) = Λ -1 x , and µ k = N ( m k , Σ k ) , we obtain for x k ∼ µ k ,

Thus, the output of the forward step is still a Gaussian of the form ν k +1 = N ( m k + 1 2 , Σ k + 1 2 ) with

Since ∇ V is linear, the output of the backward step stays Gaussian. Moreover, the first order conditions give

Thus, the output is a Gaussian N ( m k +1 , Σ k +1 ) with ( m k +1 , Σ k +1 ) satisfying

Moreover, if Λ and Σ k +1 commute, this is equivalent to

which solution is given by

To sum up, the update is

For Λ = Σ , we call it the ideally preconditioned Forward-Backward scheme (PFB).

G Additional details on experiments

G.1 Implementing the schemes

In this subsection, we sum up how to implement the different schemes in practice, given a finite number of particles. In all cases, we first sample x (0) 1 , . . . , x (0) n ∼ µ 0 , then we apply the scheme to ˆ µ ( k ) n = 1 n ∑ n i =1 δ x ( k ) i .

Mirror descent. In general, for ϕ pushforward compatible, one needs to solve at each iteration k ≥ 0 ,

If ϕ ( µ ) = ∫ V d µ with ∇ V having an analytical inverse, the scheme can be implemented as

Except for this case, one cannot in general invert ∇ W 2 ϕ ( µ k +1 ) ◦ T k +1 directly. A practical workaround is to solve an implicit problem, see e.g. [133]. Here, we use the Newton-Raphson algorithm. Suppose we have µ k = 1 n ∑ n i =1 δ x ( k ) i and we are looking for µ k +1 = 1 n ∑ n i =1 δ x i . Then, the scheme is equivalent to

for

Write G ( x 1 , . . . , x n ) = ( G 1 ( x 1 , . . . , x n ) , . . . , G n ( x 1 , . . . , x n ) ) . Then, at each step k , we perform the following Newton iterations, starting from ( x ( k ) 1 , . . . , x ( k ) n ) :

The Jacobian is of size nd × 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 e.g. [40].

0.0

2.5

5.0

7.5

10.0

Time (

iterations)

0.16

0.14

0.12

0.10

0.08

(

t

)

K

(

x

) =

1

2

x

2

2

K

(

x

) =

1

4

x

4

2

+ 1

2

x

2

2

0

2

4

6

8

10

t

0

2

4

6

8

10

12

14

Time (

Iterations)

10

14

10

10

10

6

10

2

10 2

KL(

t

||

*

)

NEM

PFB

FB

PKLM

KLM

Preconditioned gradient descent. Plugging the empirical measure in (11), the preconditioned scheme can be implemented as

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 ( µ ) = 1 2 ∫∫ W ( x -y ) d µ ( x )d µ ( y ) for W ( z ) = 1 4 ∥ z ∥ 4 2 -1 2 ∥ z ∥ 2 2 . It is well-known that the stationary solution of its gradient flow is a 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 ∇ 2 W ( z ) = 2 zz T + ( ∥ z ∥ 2 2 -1 ) I d , we have for all z ,

with p 2 ( t ) = 3 t 2 +1 . Thus, by [88, Remark 2], W is β -smooth relative to K 4 ( z ) = 1 4 ∥ z ∥ 4 2 + 1 2 ∥ z ∥ 2 2 with β = 4 . Thus, using Proposition 23, ˜ W µ is β -smooth relative to ϕ µ (T) = 1 2 ∫∫ K ( T( x ) -T( x ′ ) ) d µ ( x )d µ ( x ′ ) for all µ , and we can apply Proposition 3.

Under the additional hypothesis that the measures are compactly supported, and thus there exists M > 0 such that ∥ x ∥ 2 2 ≤ M for µ -almost every x , we can also show that W is β -smooth relative to K 2 ( z ) = 1 2 ∥ z ∥ 2 2 . Indeed, on one hand, ∇ 2 K = I d and ∇ 2 W ( z ) = 2 zz T + ( ∥ z ∥ 2 2 -1 ) I d . Thus, for all v, z ∈ R d ,

In Figure 4, we plot the evolution of W along the flows obtained with these two Bregman potential, starting from µ 0 = N (0 , 0 . 25 2 I 2 ) for n = 100 particles, with a step size of τ = 0 . 1 for 120 epochs.

Ill-conditioned interaction energy. We also study the minimization of an interaction energy with an ill-conditioned kernel W ( z ) = 1 4 ( z T Σ -1 z ) 2 -1 2 z T Σ -1 z where Σ ∈ S ++ d ( 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 an ellipsoid instead of a ring. In our experiments, we take Σ = diag(100 , 0 . 1) . For each scheme, we use µ 0 = N (0 , 0 . 25 2 I 2 ) , n = 100 particles and a step size of τ = 0 . 1 .

On Figure 1, we use Bregman potentials which take into account this conditioning, namely we use K Σ 2 ( z ) = 1 2 z T Σ -1 z and K Σ 4 ( z ) = 1 4 ( z T Σ -1 z ) 2 -1 2 ( z T Σ -1 z ) , and we observe that the convergence is much faster compared to the same kernels without preconditioning. For K Σ 2 ( z ) = 1 2 z T Σ -1 z , the scheme becomes

0

5

10

Time (

iterations)

0.05

0.08

0.11

0.14

0.17

(

t

)

K

2

K

4

K

2

K

4

5

0

5

1.0

0.5

0.0

0.5

1.0

K

4

1

0

1

K

4

0

2

4

6

8

10

12

t

0

5

10

Time (

iterations)

0.05

0.08

0.11

0.14

0.17

(

t

)

5

0

5

1.0

0.5

0.0

0.5

1.0

PrecGD

1

0

1

0

2

4

6

8

10

12

t

Thus, we see that Σ -1 has less influence which might explain the faster convergence.

Similarly as in the case without preconditioning, using that ∇ 2 W ( z ) = 2Σ -1 zz T Σ -1 +( z T Σ -1 z -1)Σ -1 , we can show that

For the sake of comparison, we also report on Figure 6 the trajectories of particles for the use of K 4 and K Σ 4 , as well as of the usual Wasserstein gradient descent and the preconditioned Wasserstein gradient descent obtained with h ∗ ( x ) = 1 2 x T Σ x (which is equivalent to the Mirror Descent with ϕ V µ as Bregman potential and V ( x ) = 1 2 x T Σ -1 x ). We observe almost the same trajectories as K 2 , which would indicate that the target is also smooth compared to ϕ V µ .

Runtime. These experiments were run on a personal Laptop with a CPU Intel Core i5-9300H. For the interaction energy as Bregman potential, running the algorithm with Newton's method for n = 100 particles in dimension d = 2 for 120 epochs took about 5mn for K 2 and K Σ 2 , and about 1h for K 4 and K Σ 4 .

G.3 Mirror descent on Gaussians

As the mirror descent scheme cannot be computed in closed-form for Bregman potentials which are not potential energies, and thus are computationally costly, we propose here to restrain ourselves to the Gaussian setting.

We choose as target distribution ν = N (0 , Σ) for Σ a symmetric positive definite matrix in R 10 × 10 , and the functional to be optimized is F ( µ ) = ∫ V d µ + H ( µ ) with V ( x ) = 1 2 x T Σ -1 x . The initial distribution is always chosen as µ 0 = N (0 , I d ) . In all cases, the step size is chosen as τ = 0 . 01 , and we run the scheme for 1500 iterations. For the target distributions, we sample 20 random covariances of the form Σ = UDU T with D evenly spaced in log scale between 1 and 100 , and U ∈ R 10 × 10 chosen as a uniformly random orthogonal matrix, as in [43], and we report the averaged KL divergence over iterations in Figure 2. We add on Figure 5 the same experiments with targets of the form N (0 , D ) where D is a diagonal matrix on R 10 × 10 sampled uniformly over [0 , 50] 10 . We compare here the Forward-Backward (FB) scheme of [43], the ideally preconditioned ForwardBackward scheme (PFB), which uses the closed-form (116) derived in Appendix F with Λ = Σ , and the Mirror Descent with negative entropy Bregman potential (NEM), whose closed-form was derived in Appendix D.3, and which we recall:

We also experiment with the KL divergence as Bregman potential (KLM) and the ideally preconditioned KL 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 as Bregman potential ϕ µ (T) = ∫ ψ ◦ Td µ for ψ ( x ) = 1 2 x T Λ -1 x is equivalent to using a preconditioner with h ∗ ( x ) = 1 2 x T Λ x .

Analysis of the convergence. It is well-known that along the Wasserstein gradient flow of the KL divergence starting from a Gaussian and with a Gaussian target (Ornstein-Uhlenbeck process),

scRNAseq

4i

4i

4i

scRNAseq

scRNAseq

with precond.

without precond.

entropic map

T

ε

the measures stay Gaussian [129]. Thus, the Forward-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, we choose ϕ µ (T) = ∫ ψ ◦ Td µ for ψ ( x ) = 1 2 x T Σ x . In this situation, G ( µ ) = ∫ V d µ is 1-smooth and 1-convex relative to ϕ . Thus, we can apply Proposition 25. We refer to Appendix F for more details on the 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. we project the gradient on the space of affine maps with symmetric linear term, i.e. of the form T( x ) = b + S ( x -m ) with S ∈ S d ( R ) [43]. We refer to [43, 73] for more details on this submanifold. This can be seen as 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 F relative to H along the curve µ t = ( (1 -t )Id + t T k +1 ) # µ k under the hypothesis that the covariances matrices have bounded eigenvalues. Moreover, since d ˜ F µ = d ϕ V µ +d ˜ H µ ≥ d ˜ H µ , F is also 1-convex relative to H .

Proposition 27. Let λ > 0 , F ( µ ) = ∫ V d µ + H ( µ ) with V ( x ) = 1 2 x T Σ -1 x where Σ ∈ S ++ d ( R ) and Σ ⪯ λI d . Suppose that for all k ≥ 0 , (1 -τ )Σ k +1 Σ -1 k + τ Σ k +1 Σ -1 ⪰ 0 . Then, F is smooth relative to H along µ t = ( (1 -t )Id + t T k +1 ) # µ k where µ k = N (0 , Σ k ) with Σ k ∈ S ++ d ( R ) , Σ k ⪯ λI d .

Proof.

G.4 Single-cell experiments

First, we provide more details on the experiment on single cells of Section 5. Then, we detail a second experiment comparing the method with using a static map.

Details on the metrics. We show the benefits of using the polynomial preconditioner over the single-cell datasets for different metrics.

The first one considered is the Sliced-Wasserstein distance [16, 102], defined as

where S d -1 = { θ ∈ R d , ∥ θ ∥ 2 = 1 } , λ denotes the uniform distribution on S d -1 and for all θ ∈ S d -1 , x ∈ R d , P θ ( x ) = ⟨ x, θ ⟩ . For F ( µ ) = 1 2 SW 2 2 ( µ, ν ) , the Wasserstein gradient can be computed as [18]

where, for t ∈ R , ψ ′ θ ( t ) = t -F -1 P θ # ν ( F P θ # µ ( t ) ) with F P θ # µ the cumulative distribution function of P θ # µ . In practice, we compute SW and its gradient using a Monte-Carlo approximation by first drawing L uniform random directions θ 1 , . . . , θ L .

The second one considered is the Sinkhorn divergence [50] defined as

with

the entropic regularized OT. The Wasserstein gradient of S 2 ε, 2 is simply obtained as the potential [50].

Finally, we also consider the energy distance, defined as

To compute its Wasserstein gradient, we use the sliced procedure of [60].

Parameters chosen. For all the metrics, we fixed the step size at τ = 1 . To choose the parameter a of the preconditioner h ∗ ( x ) = ( ∥ x ∥ a 2 +1) 1 /a -1 , we ran a grid search over a ∈ { 1 . 25 , 1 . 5 , 1 . 75 } for a random treatment, and used it for all the others. In particular, we used for the dataset 4i a = 1 . 5 for the Sinkhorn divergence and for SW, and a = 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 d = 50 for scRNAseq. For all 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, we also added 40% of unseen cells following [21]. Note that we reported the results in Figure 3 for 3 different initializations for each treatment, and reported these results with their mean. We report the results using a fixed relative tolerance tol = 10 -3 , i.e. at the first iteration where |F ( µ k ) -F ( µ k -1 ) | / F ( µ k -1 ) ≤ tol , with a maximum value of iterations of 10 4 . For the Sinkhorn divergence, we chose ε as 10 -1 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 L = 1024 projections.

Comparison to an OT static map. We now compare the prediction of the response of cells to a 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 a dynamic procedure, iterating multiple steps to map the unperturbed population µ to the perturbed population ν , 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 T ε . The results are shown in Figure 7.

Runtime. For this experiment, we used a GPU Tesla P100-PCIE-16GB. Depending on the convergence and on the metric considered, each run took in between 30s and 10mn. So in total, it took a few hundred of hours of computation time.

Initial particles

Final particles MD

Final particles MLD

0

50

100

150

200

Iterations

24

25

26

27

28

29

(

t

)

MD

MLD

G.5 Mirror descent on the simplex

We can also leverage the mirror map to perform sampling in 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, 114] or other MCMC algorithms [52, 94].

The goal here is to sample from a Dirichlet distribution, i.e. from a distribution ν ∝ e -V where V ( x ) = -∑ d i =1 a i log( x i ) -a d +1 log ( 1 -∑ d i =1 x i ) . To sample from such a distribution, we minimize the Kullback-Leibler divergence, i.e. F ( µ ) = KL( µ || ν ) = ∫ V d µ + H ( µ ) . To stay on the (open) simplex ∆ d = { x ∈ R d +1 , x i > 0 , ∑ d +1 i =1 x i < 1 } , we use the mirror map ϕ ( µ ) = ∫ ψ d µ with ψ ( x ) = ∑ d i =1 x i log( x i ) + (1 -∑ i x i ) log(1 -∑ i x i ) for which

The scheme here is given by T k +1 = ∇ ψ ∗ ◦ ( ∇ ψ -γ ∇ W 2 F ( µ k )) , where ∇ W 2 F ( µ k ) = ∇ V + ∇ log µ k , with the density of µ k estimated through a Kernel Density Estimator (KDE). We plot on Figure 8a the results obtained for d = 2 , a 1 = a 2 = a 3 = 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 KL over iterations on Figure 8b (where the entropy is estimated using the Kozachenko-Leonenko 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], or using neural networks to learn ratio density estimators [6, 49, 128].

H Proofs

H.1 Proof of Proposition 1

Let µ ∈ P 2 ( R d ) , S , T ∈ D ( ˜ F µ ) , ϵ > 0 . Since F is Wasserstein differentiable at S # µ , applying Proposition 9 at S # µ with ν = ( S + ϵ (T -S) ) # µ and γ = ( S , S + ϵ (T -S) ) # µ ∈ Π(S # µ, ν ) , we

obtain,

Thus, δ ˜ F µ (S , T -S) = ⟨∇ W 2 F (S # µ ) ◦ S , T -S ⟩ L 2 ( µ ) . Note that in the third equality we used that ∇ W 2 F ( µ ) ∈ L 2 ( µ ) .

H.2 Proof of Proposition 2

Let µ, ρ ∈ P 2 , ac ( R d ) and ν ∈ P 2 ( R d ) . Define T µ,ν ϕ µ = argmin T # µ = ν d ϕ µ (T , Id) , U ρ,ν ϕ ρ = argmin U # ρ = ν d ϕ ρ (U , Id) and let S ∈ L 2 ( µ ) such that S # µ = ρ . Then, noticing that γ = (T µ,ν ϕ µ , S) # µ ∈ Π( ν, ρ ) , we have

In the last line, we used Proposition 15, i.e. that the optimal coupling is of the form (U ρ,ν ϕ ρ , Id) # ρ .

H.3 Proof of Proposition 3

Let T k +1 = argmin T ∈ L 2 ( µ k ) τ ⟨∇ W 2 F ( µ k ) , T -Id ⟩ L 2 ( µ k ) +d ϕ µ k (T , Id) . Applying the three-point inequality (Lemma 29) with ψ (T) = τ ⟨∇ W 2 F ( µ k ) , T -Id ⟩ L 2 ( µ k ) which is convex, T 0 = Id and T ∗ = T k +1 , we get for all T ∈ L 2 ( µ k ) ,

which is equivalent to

By the β -smoothness of ˜ F µ k relative to ϕ µ k , we also have

Moreover, since β ≤ 1 τ , this inequality implies (by non-negativity of d ϕ µ k ),

Then, using the inequality (137), we obtain for all T ∈ L 2 ( µ k ) ,

Observing that ˜ F µ k (T k +1 ) = F ( µ k +1 ) and ˜ F µ k (Id) = F ( µ k ) , we get

Finally, setting T = Id , we obtain the result:

H.4 Proof of Proposition 4

Let ν ∈ P 2 ( R d ) , and T = argmin T , T # µ k = ν d ϕ µ k (T , Id) . From the relative convexity hypothesis, we have

Plugging this into (140), we get

Then, by definition of T , note that d ϕ µ k (T , Id) = W ϕ ( ν, µ k ) , and by Assumption 1, we have d ϕ µ k (T , T k +1 ) ≥ W ϕ ( ν, µ k +1 ) , since T # µ k = ν and (T k +1 ) # µ k = µ k +1 . Thus,

Observing that F ( µ k ) ≤ F ( µ ℓ ) for all ℓ ≤ k (by Proposition 3 and non-negativity of d ϕ for ϕ convex) and that W ϕ ( ν, µ ) ≥ 0 , we can apply Lemma 30 with f = F , c = F ( ν ) and g = W ϕ ( ν, · ) , and we obtain

For the second result, from (145), we get for ν = µ ∗ the minimizer of F , since F ( µ k +1 ) -F ( µ ∗ ) ≥ 0 ,

H.5 Proof of Proposition 5

Let k ≥ 0 , by the definition of d ϕ h ∗ µ k and the hypothesis d ϕ h ∗ µ k ( ∇ W 2 F ( µ k +1 ) ◦ T k +1 , ∇ W 2 F ( µ k ) ) ≤ β d ˜ F µ k (Id , T k +1 ) , we have

where we used in the last line that τ ≤ 1 β and the non-negativity of the Bregman divergence since F is convex along t ↦→ ( (1 -t )T k +1 + t Id ) # µ k and thus by Proposition 13, d ˜ F µ k (Id , T k +1 ) ≥ 0 .

Let T ∈ L 2 ( µ k ) . Then, using the three-point identity (Lemma 28) (with S = Id , U = T and T = T k +1 ), and remembering that T k +1 = Id -τ ∇ h ∗ ◦ ∇ W 2 F ( µ k ) , we get

This is equivalent to

Then, using the definition of d ϕ h ∗ µ k ( ∇ W 2 F (T # µ k ) ◦ T , ∇ W 2 F ( µ k ) ) , we obtain

Plugging this into (148), we get

For T = Id , we get

H.6 Proof of Proposition 6

Let µ ∗ ∈ P 2 ( R d ) be the minimizer of F , k ≥ 0 and T = argmin T ∈ L 2 ( µ k ) , T # µ k = µ ∗ d ˜ F µ k (Id , T) . First, observe that since µ ∗ is the minimizer of F , then ∇ W 2 F ( µ ∗ ) = 0 (see e.g. [74, Theorem 3.1]), and thus ϕ h ∗ µ k (0) = h ∗ (0) . Moreover, it induces that d ˜ F µ k (Id , T) = F ( µ k ) - F ( µ ∗ ) and d ˜ F µ k (T k +1 , T) = F ( µ k +1 ) -F ( µ ∗ ) .

Then, applying Lemma 30 with f = ϕ h ∗ · ◦ ∇ W 2 F (which satisfies ϕ h ∗ µ k +1 ( ∇ W 2 F ( µ k +1 ) ) ≤ ϕ h ∗ µ k ( ∇ W 2 F ( µ k ) ) by Proposition 5), c = h ∗ (0) and g = F ( · ) -F ( µ ∗ ) ≥ 0 , we get

Concerning the convergence of F ( µ k ) , if α > 0 and h ∗ attains its minimum in 0 , then necessarily ϕ h ∗ µ (T) ≥ h ∗ (0) for all µ ∈ P 2 ( R d ) and T ∈ L 2 ( µ ) . Thus, using (152), we get

Thus, for all k ≥ 0 ,

H.7 Proof of Proposition 7

Let µ ∈ P 2 ( R d ) . Since ˜ F ∗ µ is Gâteaux differentiable, we can define its Bregman divergence.

For the first point, ϕ h ∗ is β -smooth relative to F ∗ µ along t ↦→ ( (1 -t ) ∇ W 2 F ( µ ) + t ∇ W 2 F (T # µ ) ◦ T ) # µ . Thus, by applying Definition 3 for s = 1 and t = 0 , we have

Using Lemma 19, we finally obtain

which is the desired inequality.

The second point follows similarly.

H.8 Proof of Lemma 11

Let us define ˜ G µ : L 2 ( µ ) × R d → R d as for all T ∈ L 2 ( µ ) , x ∈ R d ,

with for all i , ˜ G i µ : L 2 ( µ ) × R d → R , ˜ G i µ (T , x ) = ∂ ∂x i δ F δµ (T # µ )( x ) . Using the chain rule, for all x ∈ R d ,

On one hand, we have ∇ 2 ˜ G i µ ( T s , T s ( x ) ) = ∇ ∂ ∂x i δ F δµ ( (T s ) # µ )( T s ( x ) ) . On the other hand, let us compute ∇ 1 ˜ G i µ (T , x ) . First, we define the shorthands ˜ g x,i µ (T) = ˜ G i µ (T , x ) = ∂ ∂x i δ F δµ (T # µ )( x ) and g x,i ( ν ) = ∂ ∂x i δ F δµ ( ν )( x ) . Since ˜ g x,i µ (T) = g x,i (T # µ ) , applying Proposition 1, we know that ∇ 1 ˜ G µ (T , x ) = ∇ ˜ g x,i µ (T) = ∇ W 2 g x,i (T # µ ) ◦ T .

Now, let us compute ∇ W 2 g x,i ( ν ) = ∇ δg x,i δµ ( ν ) . Let χ be such that ∫ d χ = 0 , then using the hypothesis that δ δµ ∇ δ F δµ = ∇ δ 2 F δµ 2 and the definition of g x,i ,

Thus, ∇ W 2 g x,i ( ν ) = ∇ y ∂ ∂x i δ 2 F δµ 2 ( ν )( x, y ) .

Putting everything together, we obtain

and thus

H.9 Proof of Proposition 12

First, recall that by using the chain rule and Proposition 1, d d t F ( µ t ) = ⟨∇ W 2 F ( µ t ) ◦ T t , dT t d t ⟩ L 2 ( µ ) . Thus, since d 2 T t d t 2 = 0 ,

H.10 Proof of Proposition 13

(c1) = ⇒ (c2) . Let t > 0 , t 1 , t 2 ∈ [0 , 1] ,

Passing to the limit t → 0 and using Proposition 1, we get ⟨∇ W 2 F ( (T t 1 ) # µ ) ◦ T t 1 , T t 2 -T t 1 ⟩ L 2 ( µ ) ≤ F ( (T t 2 ) # µ ) -F ( (T t 1 ) # µ ) .

(c2) = ⇒ (c3) . Let t 1 , t 2 ∈ [0 , 1] , then by hypothesis,

Summing the two inequalities, we get

(c3) = ⇒ (c4) . Let t 1 , t 2 ∈ [0 , 1] . First, we have,

Let ϵ ∈ (0 , 1) and define t ↦→ ν ϵ t = ˜ µ t 1 → 1 ϵt the interpolation curve between (T t 1 ) # µ and ( T t 1 + ϵ (T -T t 1 ) ) # µ . Then, noting that T t 1 + ϵ (T -T t 1 ) = T t 1 + ϵ (1 -t 1 ) , so

ν ϵ t = ˜ µ t 1 → 1 ϵt = ˜ µ t 1 → t 1 + ϵ (1 -t 1 ) t and we have that

Moreover, by continuity, d 2 d t 2 F ( ν ϵ t ) - - - → ϵ → 0 d 2 d t 2 F ( (T t 1 ) # µ ) = d 2 d t 2 F ( µ t 1 ) . Then, since t ↦→ d 2 d t 2 F ( ν ϵ t ) is continuous on [0 , 1] , it is bounded, and we can apply the dominated convergence theorem. This implies that for all t 1 ∈ [0 , 1] ,

(c4) = ⇒ (c1) . Let t 1 , t 2 ∈ [0 , 1] and φ ( t ) = F (˜ µ t 1 → t 2 t ) for all t ∈ [0 , 1] . From [125, Equation 16.5],

where G is the Green function defined as G ( s, t ) = s (1 -t ) ✶ { s ≤ t } + t (1 -s ) ✶ { t ≤ s } ≥ 0 [125, Equation 16.6]. Then, d 2 d t 2 F ( µ t ) ≥ 0 implies that ∫ 1 0 d 2 d t 2 φ ( s ) G ( s, t ) d s ≥ 0 , and thus φ ( t ) = F (˜ µ t 1 → t 2 t ) ≤ (1 -t ) φ (0)+ tφ (1) = (1 -t ) F ( (T t 1 ) # µ ) + t F ( (T t 2 ) # µ ) . (174)

H.11 Proof of Proposition 24

Let J(T) = d ϕ µ k (T , Id) + τ ( ∇ W 2 G ( µ k ) , T -Id ⟩ L 2 ( µ k ) + H (T # µ k ) ) . Taking the first variation, we get

Thus,

Now, we aim at showing that ˜ T k +1 = T k +1 ◦ S k +1 or

First, by the change of variable formula, since ϕ µ is pushforward compatible, observe that for T ∈ L 2 ( ν k +1 ) , d ϕ ν k +1 (T , Id) + τ H (T # ν k +1 ) = d ϕ µ k (T ◦ S k +1 , S k +1 ) + τ H ( (T ◦ S k +1 ) # µ k ) .

By assumption, ν k +1 ∈ P 2 , ac ( R d ) . Thus, applying Proposition 15, there exists T ν k +1 ,µ k +1 ϕ ν k +1 such that (T ν k +1 ,µ k +1 ϕ ν k +1 ) # ν k +1 = µ k +1 and T ν k +1 ,µ k +1 ϕ ν k +1 = argmin T , T # ν k +1 = µ k +1 d ϕ ν k +1 (T , Id) , and thus d ϕ ν k +1 (T ν k +1 ,µ k +1 ϕ ν k +1 , Id) = W ϕ ( µ k +1 , ν k +1 ) .

By contradiction, we suppose that

On one hand, we have (T ν k +1 ,µ k +1 ϕ ν k +1 ◦ S k +1 ) # µ k = (T ν k +1 ,µ k +1 ϕ ν k +1 ) # ν k +1 = µ k +1 , and therefore H ( (T ν k +1 ,µ k +1 ϕ ν k +1 ◦ S k +1 ) # µ k ) = H ( µ k +1 ) = H ( ( ˜ T k +1 ) # µ k ) . On the other hand, ( ˜ T k +1 , S k +1 ) # µ k ∈ Π( µ k +1 , ν k +1 ) , and thus

Thus,

But T ν k +1 ,µ k +1 ϕ ν k +1 ∈ L 2 ( ν k +1 ) , so this is a contradiction. So, we can conclude that the two schemes are equivalent, and moreover, ˜ T k +1 = T ν k +1 ,µ k +1 ϕ ν k +1 ◦ S k +1 .

H.12 Proof of Proposition 25

Let ψ (T) = τ ( ⟨∇ W 2 G ( µ k ) , T -Id ⟩ L 2 ( µ k ) + H (T # µ k ) ) . Since ˜ H µ k is convex on L 2 ( µ k ) , ψ is convex, and we can apply the three-point inequality (Lemma 29) and for all T ∈ L 2 ( µ k ) ,

Since ˜ G µ k is β -smooth relatively to ϕ µ k along t ↦→ ( (1 -t )Id + t ˜ T k +1 ) # µ k , and τ ≤ 1 β , we also have

Thus, applying first the smoothness of G and then the three-point inequality, we get for all T ∈ L 2 ( µ k ) ,

Now, let ν ∈ P 2 ( R d ) and T µ k ,ν ϕ µ k = argmin T , T # µ k = ν d ϕ µ k (T , Id) , and suppose that ˜ G µ k is α -convex relative to ϕ µ k along t ↦→ ( (1 -t )Id + tT µ k ,ν ϕ µ k ) # µ k . Thus,

d

˜

G

µ k

(T

µ

ϕ

k

,ν

µ k

,

Id)

α

d

ϕ

µ k

(T

µ

ϕ

k

,ν

µ k

⇐⇒ G

(

ν

)

-

α

d

ϕ

µ k

Plugging this into (185), we get

F ( µ k +1 ) ≤ H ( ν )+ G ( ν ) -α d ϕ µ k (T µ k ,ν ϕ µ k , Id)+ 1 τ d ϕ µ k (T µ k ,ν ϕ µ k , Id) -1 τ d ϕ µ k (T µ k ,ν ϕ µ k , ˜ T k +1 ) . (187) Now, note that d ϕ µ k (T µ k ,ν ϕ µ k , Id) = W ϕ ( ν, µ k ) and by Assumption 1, d ϕ µ k (T µ k ,ν ϕ µ k , ˜ T k +1 ) ≥ W ϕ ( ν, µ k +1 ) . Thus,

Using T = Id in (185), we observe that F ( µ k ) ≤ F ( µ ℓ ) for all ℓ ≤ k . Moreover, W ϕ ( ν, µ k ) ≥ 0 . Thus, applying Lemma 30 with f = F , c = F ( ν ) and g = W ϕ ( ν, · ) , we obtain

,

Id)

(T

µ

k

,ν

ϕ

µ k

,

Id)

(

µ

k

) +

⟨∇

W

2

G

(

µ

k

)

,

T

µ

k

,ν

ϕ

µ k

-

Id

⟩

L

2

(

µ

k

)

.

(186)

H.13 Proof of Lemma 26

First, ∇ V ∗ is bijective. Thus, we only need to show that h = ∇ V -τ ∇ U is injective. Take u = V -τU .

Since U is β -smooth relative to V , we have for all x, y ,

Moreover, by definition of d V ,

Summing the two inequalities, we get

This is equivalent to and thus with u being (1 -τβ ) -convex relative to V (for τβ ≤ 1) . For τβ < 1 , it is equivalent to u -(1 -τβ ) V convex, i.e. ⟨∇ u ( x ) -∇ u ( y ) , x -y ⟩ ≥ (1 -τβ ) ⟨∇ V ( x ) -∇ V ( y ) , x -y ⟩ ≥ 0 . Since V is strictly convex, ∇ u is injective.

Moreover, | det ∇ T | = | det ( ∇ 2 V ∗ ◦ ( ∇ V -τ ∇ U ) ) det ∇ 2 u | > 0 because on one hand u is (1 -βτ ) -convex 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], T # µ 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 µ ↦→ ∫ V d µ is smooth relative to H . Using Proposition 13, we need to show that

Recall from (61) that T k +1 ( x ) = ( (1 -τ )Σ k +1 Σ -1 k + τ Σ k +1 Σ -1 ) x + cst , thus ∇ T k +1 is a constant. Using the computations of [43, Appendix B.2],

Assuming (1 -τ )Σ k +1 Σ -1 k + τ Σ k +1 Σ -1 ⪰ 0 , T k +1 is the gradient of a convex function and µ t is a Wasserstein geodesic. Thus, by [43],

Moreover, by [33, Lemma 10], µ ↦→ ∥ Σ µ ∥ op is convex along generalized geodesics, and thus Σ µ t ⪯ λI d and ∥ Σ µ t ∥ op ≤ λ [43]. Hence, noting σ max ( M ) the largest eigenvalue of some matrix M ,

which is equivalent to

From this inequality, we deduce that

So,

I Additional results

I.1 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 ϕ : L 2 ( µ ) → R be Gâteaux differentiable. For all S , T , U ∈ L 2 ( µ ) , we have

Proof. Let S , T , U ∈ L 2 ( µ ) , then using the linearity of the Gâteaux differential,

Lemma 29 (Three-Point Inequality) . Let µ ∈ P 2 ( R d ) , T 0 ∈ L 2 ( µ ) and ϕ µ : L 2 ( µ ) → R convex, and Gâteaux differentiable. Let ψ : L 2 ( µ ) → R be convex, proper and lower semicontinuous. Assume there exists T ∗ = argmin T ∈ L 2 ( µ ) d ϕ µ (T , T 0 ) + ψ (T) . Then, for all T ∈ L 2 ( µ ) ,

Proof. Denote J(T) = d ϕ µ (T , T 0 ) + ψ (T) . Let T ∗ = argmin T ∈ L 2 ( µ ) J(T) , hence 0 ∈ ∂ J(T ∗ ) .

Since ϕ and ψ are proper, convex and lower semicontinuous, and T ↦→ d ϕ µ (T , T 0 ) is continuous (since ϕ µ is continuous), thus by [99, Theorem 3.30], ∂ J(T ∗ ) = ∂ψ (T ∗ ) + ∂ d ϕ µ ( · , T 0 )(T ∗ ) .

Finally, by definition of subgradients and by applying Lemma 28, we get for all T ∈ L 2 ( µ ) ,

Moreover, since ϕ µ is differentiable, ∂ d ϕ µ ( · , T 0 )(T ∗ ) = {∇ T d ϕ µ (T ∗ , T 0 ) } = {∇ ϕ µ (T ∗ ) -∇ ϕ µ (T 0 ) } , and thus ∇ ϕ µ (T 0 ) -∇ ϕ µ (T ∗ ) ∈ ∂ψ (T ∗ )

Remark 3. Actually we can restrict ψ to be convex along ( (1 -t )T ∗ + t T ) # µ . In that case, d ψ (T , T ∗ ) = ψ (T) -ψ (T ∗ ) -⟨ φ, T -T ∗ ⟩ L 2 ( µ ) ≥ 0 for φ ∈ ∂ψ (T ∗ ) (by Proposition 13) and we still have ∂ψ (T ∗ ) + ∂ d ϕ µ ( · , T 0 )(T ∗ ) ⊂ ∂ J(T ∗ ) (see [99, Theorem 3.30]) so that we can conclude.

I.2 Convergence lemma

We first provide a Lemma which follows from [88, Theorem 3.1], and which is useful for the proofs of Propositions 4, 6 and 25.

Lemma 30. Let f : X → R , g : X → R + and ( x k ) k ∈ N a sequence in X such that for all k ≥ 1 , f ( x k ) ≤ f ( x k -1 ) . Assume that there exists β > α ≥ 0 , c ∈ R such that for all k ≥ 0 , f ( x k +1 ) -c ≤ ( β -α ) g ( x k ) -βg ( x k +1 ) , then

Proof. First, observe the f ( x k ) ≤ f ( x ℓ ) for all ℓ ≤ k . Thus, for all k ≥ 1 ,

Now, note that β ∑ k ℓ =1 ( β β -α ) ℓ = α ( β β -α ) k -1 = α ( 1+ α β -α ) k -1 ≤ β -α k since ( 1 + α β -α ) k ≥ 1 + k α β -α (by convexity on R + of x ↦→ (1 + x ) k ). Thus,

I.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 ϕ µ (T) = ∫ V ◦ T d µ . The following Lemma is akin to [80, Proposition 4] which shows it only for OT maps.

Lemma 31. Let V : R d → R convex and ϕ µ (T) = ∫ V ◦ T d µ . Then,

Proof. Let T , S ∈ L 2 ( µ ) , then remembering that ∇ W 2 V ( µ ) = ∇ V , we have

Next, we focus on ϕ µ (T) = 1 2 ∫∫ W ( T( x ) -T( x ′ ) ) d µ ( x )d µ ( x ′ ) , and we generalize the result from [80, Proposition 4].

Lemma 32. Let W : R d → R even ( W ( x ) = W ( -x ) ), convex and differentiable. Let ϕ µ (T) = 1 2 ∫∫ W ( T( x ) -T( x ′ ) ) d µ ( x )d µ ( x ′ ) . Then,

Proof. Let T , S ∈ L 2 ( µ ) , remember that ∇ W 2 W ( µ ) = ∇ W ⋆ µ , and thus ∇ W 2 W (S # µ ) ◦ S = ( ∇ W ⋆ S # µ ) ◦ S . Thus,

Then, note that ∇ W ( -x ) = -∇ W ( x ) and thus the last term can be written as:

Finally, we get

Now, we make the connection with the mirror map used by Deb et al. [41] and derive the related Bregman divergence.

Lemma 33. Let ϕ µ (T) = 1 2 W 2 2 (T # µ, ρ ) for µ, ρ ∈ P 2 , ac ( R d ) . Then, for all T , S ∈ L 2 ( µ ) , such that T # µ, S # µ ∈ P 2 , ac ( R d ) ,

where T ρ T # µ denotes the OT map between T # µ and ρ .

Proof. Let T , S ∈ L 2 ( µ ) such that T # µ, S # µ ∈ P 2 , ac ( R d ) . Remember that ∇ W 2 W 2 2 ( · , ρ ) = Id -T ρ · , then

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 are 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 in 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 it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: On the two theoretical sections (Section 3 and Section 4), we state all the hypotheses to obtain our convergence results. Then, in Section 5, we discuss some couples of target functional F and Bregman potential/preconditioner ϕ for which we can verify these assumptions. However, in general, verifying these assumptions is a hard problem, as stated in the Conclusion, and we leave for future works these investigations.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (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 on 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best

judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The full set of assumptions are 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The information to reproduce the main experimental results are described in Section 5 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 a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make 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 can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to 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 to the data and code, with sufficient instructions 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 https://github.com/clbonet/Mirror_and_Preconditioned_Gradient_ Descent_in_Wasserstein_Space .

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not 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 a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 G. 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 a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: 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.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence 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 a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 on a 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 a 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, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well 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://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The research conducted in this paper is conform with the NeurIPS Code of Ethics.

Guidelines:

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 circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or 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., deployment 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 to 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 a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: 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 a 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's 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 of 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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: [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 we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
comments: 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.

STRENGTHS:
The authors define and apply the Wasserstein mirror descent directions to conduct sampling algorithms. 

The 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.

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? 

I am willing to improve my score after addressing these questions.

LIMITATIONS:
There is no limitation on this paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\n1. **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.\n\nWEAKNESSES:\n1. **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.\n\nQUESTIONS:\n1. **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.\n\nLIMITATIONS:\nI 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nMirror 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nsome minor grammatical errors such as l.158 ($T_\\eta^\\nu$ OT maps from $\\eta$...)\n\nLIMITATIONS:\nThe authors did address the limitations of this work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-09"}

comments: 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.

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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-12

comments: 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.

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. 


1. Notations section should be improved. Many things are not properly defined throughout the paper, e.g., the coupling between distributions, what "Id" is?

2. On line 138, what does it mean that the sets $V$ and $W$ are "differentiable and $L$-smooth with $W$ even?  

3. On line 142, the calligraphic $\mathcal W$ is not defined. 

4. OT map mentioned on line 151 is not defined.

QUESTIONS:
n/a

LIMITATIONS:
n/a
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2024-07-15