0.0.1 Introduction↩︎

Optimizing over the space of probability measures is an important problem in machine learning, which has received attention to solve problems ranging from variational inference [1]–[3] to generative modeling [4]–[6], reinforcement learning [7], [8], optimization of neural networks [9], [10] or for modeling dynamics of cells [11]–[13]. One prominent way to optimize over the space of probability measures is to equip it with the Wasserstein distance [14], [15], and to discretize the associated Wasserstein gradient flows [16], [17]. This allows to design many optimization algorithms as counterparts of their Euclidean version, such as gradient descent [17], proximal point and gradient algorithms [16], [18], coordinate descent [19] or mirror descent [20], [21].

While these optimization methods have demonstrated good results in minimizing several objectives such as the Kullback-Leibler divergence (KL) [17], \(f\)-divergences [22]–[24], the Energy Distance [25] or the Sliced-Wasserstein distance [26]–[28], most of them are tailored for convex functionals along (generalized) geodesics. However, many popular functionals are known to be non-convex. This is the case for instance of the squared Wasserstein distance itself [29], the Sliced-Wasserstein distance [30], the KL with non-log-concave target [31], the Energy distance [32] or the Maximum Mean Discrepancy (MMD) [33]. An alternative to the convexity assumption is to consider Polyak-Łojaziewicz inequalities [34]–[36], however these are so far only known to hold in restrictive cases such as trivially for the squared Wasserstein distance, the KL w.r.t.a measure satisfying the log-Sobolev inequality [14], Sliced-Wasserstein distance over Gaussians [37], the MMD with Coulomb kernels and smooth initializations [38]. For the MMD specifically, to address its non-convexity in the Wasserstein space, [33] proposed injecting noise to the gradient. However, tuning the amount of noise remains delicate. [39] performed instead the optimization of MMD in the Wasserstein-MMD space, where MMD is convex. Their method improves performance, but it requires to change weights of the distribution, mixing two implicit steps in squared Wasserstein and MMD. More recently , [40] optimized it in the Wasserstein Fisher-Rao space through a fixed-point algorithm, also changing weights. We propose instead a new particle-based algorithm on the Wasserstein space to handle non-convex functionals.

To achieve such a goal, a candidate class of analogous algorithms over \(\R^d\) is the family of methods for objectives that can be written as a difference-of-convex functions (DC) [41], [42]. On \(\R^d\) this includes a large class of functions of interest, in particular twice continuously differentiable functions [43], [44]. One key such algorithm is the Convex-Concave Procedure [43]. These DC algorithms have already been used for many applications in machine learning, such as kernel selection [45], clustering [46], dictionary learning [47], optimal transport [48], or neural network optimization [49], [50] to name a few. However, except for a few works that studied DC algorithms on Riemannian manifolds [51]–[54], and recently on the Wasserstein space [31], [55], optimization algorithms for DC functionals have mostly been studied on Euclidean spaces.

0.0.1.0.0.1 Contributions.

In this work, we focus on developing an optimization scheme on the Wasserstein space tailored for the non-convex objectives that can be written as a difference of convex objectives. To do so, we introduce the Wasserstein Convex-Concave Procedure (WCCCP), and analyze its theoretical convergence, proving almost stationary along its iterates. The closest work to ours is [31] which used a Wasserstein Proximal Gradient algorithm [18] to solve some DC problems, but which restricted their applications to functionals whose concave part is a potential energy, i.e.linear. We argue that it is too restrictive to handle objectives such as Maximum Mean Discrepancies, which can be decomposed as a sum of two non-convex quadratic and linear terms. Thus, we deal with the more general case, where the concave part can be any differentiable functional over the Wasserstein space. Then, for several kernels, we show that with a well-chosen decomposition obtained by splitting the kernel, the introduced scheme can better optimize the MMD than the Wasserstein Gradient Descent. In particular, we provide experiments on the Energy distance and the MMD with Gaussian kernel.

0.0.1.0.0.2 Notation.

We denote by \(\cPr\) the space of probability distributions with second finite moments, and by \(\cPa\) its restriction to absolutely continuous measures with respect to the Lebesgue measure. Given \(\mu,\nu\in\cPr\), we denote by \(\W_2^2(\mu,\nu) = \inf_{\gamma\in\Pi(\mu,\nu)} \;\int \|x-y\|_2^2\;\mathrm{d}\gamma(x,y)\) the squared-Wasserstein distance, where \(\Pi(\mu,\nu)\) is the set of couplings between \(\mu\) and \(\nu\), and \(\Pi_o(\mu,\nu)\) is its subset of optimal couplings. The metric space \((\cPr, \W_2)\) is called the Wasserstein space. For any \(\mu\in\cPr\), we denote by \(L^2(\mu)\) the Hilbert space of functions \(f:\R^d\to\R^d\) such that \(\int \|f(x)\|_2^2\;\mathrm{d}\mu(x) <\infty\) equipped with the norm \(\|f\|_{L^2(\mu)}^2 = \int \|f(x)\|_2^2\;\mathrm{d}\mu(x)\) and with inner product \(\langle \cdot,\cdot\rangle_{L^2(\mu)}\). Given \(\T:\R^d\to\R^d\), \(\T_\#\mu\in\cPr\) is the pushforward measure of \(\mu\).

0.0.2 Background↩︎

We begin by recalling a few facts on optimization on the Wasserstein space. More precisely, we recall the notion of Wasserstein gradient, of total convexity on the Wasserstein space, and some classical optimization schemes on \((\cPr, \W_2)\). For more details about Wasserstein gradient flows, we refer to e.g. [29], [56]. Then, we provide a brief introduction to the convex-concave procedure on Euclidean spaces.

0.0.2.0.0.1 Wasserstein gradient.

Let \(\cF:\cPr\to\R\) be a functional. It admits a Wasserstein gradient \(\gW\cF(\mu)\in L^2(\mu)\) at \(\mu\in\cPr\) if for all \(\nu\in\cPr\), \(\gamma\in\Pi_o(\mu,\nu)\), the following first order Taylor expansion is satisfied [56], [57] \[\cF(\nu) = \cF(\mu) + \int \langle \gW\cF(\mu)(x), y-x\rangle\;\mathrm{d}\gamma(x,y) + o\big(\W_2(\mu,\nu)\big).\] When it exists, the Wasserstein gradient may not be unique in \(L^2(\mu)\) in general. Nonetheless, there is only one gradient living in the tangent space \(T_\mu\cPr\subset L^2(\mu)\) which is a Hilbert space. Hence, by Hilbert’s decomposition theorem, any gradient can be decomposed as a part in \(T_\mu\cPr\) and an orthogonal part \(\xi(\mu)\) which satisfies \(\int \langle\xi(\mu),y-x\rangle\;\mathrm{d}\gamma(x,y)=0\), see [56]. Thus, without loss of generality, we always work with the unique \(\gW\cF(\mu)\in T_\mu\cPr\) for a Wasserstein differentiable functional, using the shorthand W-differentiable for such functionals.

Classical functionals from \(\cPr\) to \(\R\) include potential energies \(\cV(\mu)=\int\V\mathrm{d}\mu\) and interaction energies \(\cW(\mu) = \frac{1}{2} \iint \W(x-y)\;\mathrm{d}\mu(x)\mathrm{d}\mu(y)\) for \(\V,\W:\R^d\to\R\), with \(\W\) symmetric. They are both differentiable provided \(\V\) and \(\W\) are differentiable and smooth enough [56], and their Wasserstein gradients read respectively as \(\gW\cV(\mu) = \nabla \V\) and as the convolution \(\gW\cW(\mu) = \nabla \W * \mu\). In this work, we will mostly focus on functionals obtained as a sum of interaction and potential energies, as they include in particular the Maximum Mean Discrepancy [33].

0.0.2.0.0.2 Convexity in the Wasserstein space.

Let \(\mu\in\cPr\), \(\T,\sS\in L^2(\mu)\). Given \(\phi:L^2(\mu)\to \R\) convex and Gateaux differentiable, the Bregman divergence on \(L^2(\mu)\) between \(\T\), \(\sS\) is defined as [58] \[\label{eq:bregman95div95L2} \D_\phi(\T,\sS) = \phi(\T)-\phi(\sS)-\langle \nabla \phi(\sS), \T-\sS\rangle_{L^2(\mu)}.\tag{1}\] Given \(\cF:\cPr\to\R\) W-differentiable, we can define the Bregman divergence in \(L^2(\mu)\) associated to the lifted functional \(\T\mapsto \cF(\T_\#\mu)\) as \[\D_{\cF}^\mu(\T,\sS) = \cF(\T_\#\mu) - \cF(\sS_\#\mu) - \langle \gW\cF(\sS_\#\mu)\circ \sS, \T-\sS\rangle_{L^2(\mu)},\] using the chain rule for the gradient of \(\tilde{\cF}_\mu:\sS\mapsto \cF(\sS_\#\mu)\) which implies \(\nabla\tilde{\cF}_\mu(\sS)=\gW\cF(\sS_\#\mu)\circ\sS\), see [20]. Note that for \(\cF(\mu)=\int \tfrac12 \|\cdot\|_2^2\;\mathrm{d}\mu\), this reduces to \(\D_\cF^\mu(\T,\sS)=\tfrac12 \|\T-\sS\|_{L^2(\mu)}^2\). Let \(\alpha \ge 0\), we say that \(\cF\) is \(\alpha\)-totally convex [59]–[61] if for all \(\mu\in\cPr\), \(\T,\sS\in L^2(\mu)\), \(\D_{\cF}^\mu(\T,\sS)\ge \frac{\alpha}{2} \|\T-\sS\|_{L^2(\mu)}^2\). Equivalently, it satisfies \[\forall t\in [0,1],\;\cF\big(\big((1-t)\T+t\sS\big)_\#\mu\big)\le (1-t) \cF(\T_\#\mu) + t\cF(\sS_\#\mu) - \alpha \frac{t(1-t)}{2} \|\T-\sS\|_{L^2(\mu)}^2.\] If this result holds only for \(\sS=\id\) and for \(\T\) the gradient of a convex function, this corresponds to the less restrictive notion of strong convexity along geodesics [29].

0.0.2.0.0.3 Wasserstein Gradient Descent.

Optimization on \(L^2(\mu)\) and on \(\cPr\) are very much intertwined in practice, see e.g. [20], [62]. For instance, the Wasserstein Gradient Descent (WGD) over a W-differentiable functional \(\cF:\cPr\to\R\) which is defined for all \(k\ge 0\), \(\tau>0\), as \(\mu_{k+1}=\big(\id-\tau\gW\cF(\mu_k)\big)_\#\mu_k\), can be written at each iteration in two steps: first solving an optimization problem on \(L^2(\mu_k)\) to get a map \(\T_{k+1}\in L^2(\mu_k)\), then pushing forward \(\mu_k\) by \(\T_{k+1}\), i.e. \[\label{eq:wgd} \begin{cases} \T_{k+1} = \mathop{\mathrm{argmin}}_{\T\in L^2(\mu_k)}\;\tfrac{1}{2\tau} \|\T-\id\|_{L^2(\mu_k)}^2 + \langle \gW\cF(\mu_k), \T-\id\rangle_{L^2(\mu_k)} \\ \mu_{k+1} = (\T_{k+1})_\#\mu_k. \end{cases}\tag{2}\] It can be shown to converge if \(\cF\) is smooth along \(t\mapsto \big((1-t)\id + t\T_{k+1})_\#\mu_k\) and convex along geodesics. Other first-order algorithms have been lifted from \(\R^d\) to \(\cPr\). For instance, replacing the squared \(L^2\) distance in 2 by a Bregman divergence 1 allows to lift the Mirror descent algorithm [63], [64] to \(\cPr\) [20].

0.0.2.0.0.4 The Convex-Concave Procedure.

A function \(f:\R^d\to\R\) is DC if it can be written as the difference of two convex functions, i.e. if there exists \(\fp, \fm\) two convex functions such that \(\f=\fp-\fm\). Every \(C^{1}\) function with Lipschitz gradient is DC [44]. To minimize such functions, a popular algorithm is the Convex-Concave Procedure (CCCP) [43]. This algorithm amounts to linearizing the concave part around the current iterate \(x_k\in\R^d\), which by convexity gives the lower bound, \[\forall x\in \R^d,\;\fm(x) \ge \fm(x_k) + \langle \nabla \fm(x_k), x-x_k\rangle,\] which entails an upper bound on \(f=\fp-\fm\), and the majorization-minimization \[x_{k+1} = \mathop{\mathrm{argmin}}_x\;\fp(x) - \fm(x_k) - \langle \nabla \fm(x_k), x-x_k\rangle, \, \forall k\ge 0.\] When both \(\fp\) and \(\fm\) are differentiable, the iterates satisfy \(\nabla \fp(x_{k+1}) = \nabla \fm(x_k)\) by the first order conditions. This algorithm belongs to the more general family of Difference-of-Convex algorithms (DCA) [41], [42], and is related to different optimization algorithms including Frank-Wolfe [65], the Mirror and Bregman proximal descent [66] or the Proximal gradient algorithm [67]. The first convergence analysis of CCCP focused on obtaining asymptotic convergence, showing that it converges towards a stationary point under some assumptions [68], [69]. Then, several works such as [65], [66], [70] derived non-asymptotic convergence rates. In particular, the algorithm was shown to converge in \(O(1/k)\) in terms of the squared norm of the gradient. More recently, [71] provided an analysis of CCCP under a generalized convexity perspective, though in finite dimensions. Linear rates were also derived under Polyak-Łojaziewicz inequalities adapted to DC functions [66], [70]–[72].

0.0.3 Wasserstein CCCP↩︎

In this section, we first introduce the Wasserstein Convex-Concave Procedure (WCCCP) to minimize difference-of-convex functions on the Wasserstein space as well as our assumptions. Then we provide a theoretical analysis in several settings, including the convex and non-convex ones. Finally, we discuss how we can implement these schemes in practice. All the proofs are deferred to Appendix 0.0.12.

0.0.3.1 Convex-Concave Procedure in the Wasserstein Space

We focus on the problem of minimizing \(\bF:\cPr\to\R\) where \(\bF\) can be decomposed as \[\bF(\mu)=\cFp(\mu)-\cFm(\mu),\quad \forall \mu\in\cPr,\] with \(\cFp,\cFm:\cPr\to\R\) both totally convex. Moreover, we assume that \(\cFm\) is W-differentiable, and will assume \(\cFp\) W-differentiable on a case-by-case basis. In contrast to [31], we do not restrict \(\cFm\) to be a potential energy.

Let \(\mu_{}\in\cPr\). Since \(\cFm\) is totally convex, for any \(\T\in L^2(\mu_{})\), \[\D_{\cFm}^{\mu_{}}(\T,\id) \ge 0 \iff \cFm(\T_\#\mu_{}) \ge \cFm(\mu_{}) + \langle \gW\cFm(\mu_{}), \T-\id\rangle_{L^2(\mu_{})}.\] Hence, we have the following upper bound on \(\bF\): \[\label{eq:upper95bound95dc} \bF(\T_\#\mu_{}) = \cFp(\T_\#\mu_{}) - \cFm(\T_\#\mu_{}) \le \cFp(\T_\#\mu_{}) - \cFm(\mu_{}) - \langle \gW\cFm(\mu_{}), \T-\id\rangle_{L^2(\mu_{})}.\tag{3}\]

This extends CCCP [43] to the Wasserstein space, CCCP being recovered when \(\cFp\) and \(\cFm\) are potential energies. Restricting the optimization to measures of the form \(\T_\#\mu_k\) for \(\T \in L^2(\mu_k)\) is without loss of generality in two important cases: i) as done in our experiments, for empirical target and initial distributions with the same number of particles; ii) if \(\mu_k\in\cPa\), as by Brenier’s theorem [73], there always exists an OT map between \(\mu_k\) and any \(\nu\in\cPr\). While for greater generality we could instead optimize over couplings with first marginal \(\mu_k\), the subproblems on \(L^2(\mu_k)\) are more tractable and reflect practical implementations.

We assume existence and uniqueness in [eq:argmin95WDCA95maps] for simplicity of exposition. Sufficient conditions are \(\T\mapsto\cFp(\T_\#\mu_k)\) strictly convex, coercive and lower semicontinuous over \(L^2(\mu_k)\), for all \(k\). We now study the theoretical convergence of the WCCCP in several settings.

0.0.3.2 Theoretical Analysis in the Non-Convex Case

We first observe that in general, similarly to the analysis of [66] in the Euclidean case, [eq:argmin95WDCA95maps] is equivalent to both a Mirror Descent and a Bregman Proximal Descent in the Wasserstein space on the objective \(\bF\), with Bregman potential respectively \(\cFp\) and \(\cFm\), and with step size \(\tau=1\).

These equivalences allow us to provide a convergence rate in the case where \(\bF\) satisfies convexity assumptions relative to \(\cFp\) or \(\cFm\), leveraging the convergence analysis of Mirror descent [20] and of Bregman Proximal Descent. We refer to Appendix 0.0.8 for these results and we focus only on the non-convex case here.

If both \(\cFp,\cFm\) are W-differentiable, then we can take the first order conditions in [eq:argmin95WDCA95maps] which yield the equivalent update \[\label{eq:foc95wdca} \gW\cFp(\mu_{k+1})\circ \T_{k+1} = \gW\cFm(\mu_k).\tag{4}\] Leveraging this result, we obtain the following relation between \(\cF(\mu_{k+1})\) and \(\cF(\mu_k)\) involving Bregman divergences with Bregman potential \(\cFp\) and \(\cFm\).

The term \(\Ck\), used also in the non-smooth and Euclidean case in [70] or [65], is merely a proxy for \(\D_{\cFp}^{\mu_k}(\id, \T_{k+1})\) when \(\cFp\) is possibly not W-differentiable. Note that \(\Ck\ge 0\), as taking \(\T=\id\) to upper bound the min, the difference vanishes. Hence, ?? implies that \(\cF\) is non-increasing along the WCCCP scheme. The condition \(\Ck=0\) can thus be used as a termination criterion as it implies that \(\mu_k\) is a critical point of \(\bF\).

0.0.3.2.0.1 Sublinear rates.

We now derive sublinear rates for the non-convex case. All the next results build on the next immediate result (obtained by telescoping), which provides a bound for any arbitrary sequence and corresponds actually to the more involved [66], \[\label{eq:base95nonconvex95WDCA} \min_{k\in \{0,\dots,K-1\}}\;\bF(\mu'_k) - \bF(\mu'_{k+1}) \le \frac{\bF(\mu'_0)-\bF(\mu'_K)}{K}, \quad \forall (\mu'_k)_{k \ge 0}, \, K\ge 1.\tag{5}\]

In the next proposition, we show a sublinear convergence rate for the termination criterion \(\Ck\), in accordance with results over \(\R^d\) from [65], [66], [70].

Assuming that either \(\cFp\) or \(\cFm\) is strongly convex along iterates, we can also obtain a sublinear convergence result for the squared distance between iterates and the norm of the W-gradient, which guarantees that at least one iterate of WCCCP is almost a stationary point.

0.0.3.3 Computing WCCCP

To solve the WCCCP scheme [eq:argmin95WDCA95maps], we need to minimize at each iteration \(\J\) on \(L^2(\mu_k)\). First, we observe that if \(\cFp\) is convex along curves of the form \(t\mapsto \big((1-t)\T+t\sS\big)_\#\mu_k\) for any \(\sS,\T\in L^2(\mu_k)\), then \(\J\) is convex on \(L^2(\mu_k)\) as \(\D_{\cFp}^{\mu_k}(\T,\sS)\ge 0\), see 0.0.2. If \(\cFp\) is W-differentiable, then \(\T\mapsto \cFp(\T_\#\mu_k)\) is Fréchet differentiable and taking the first order conditions in [eq:argmin95WDCA95maps], we get the equivalent update 4 . If \(\cFp(\mu)=\int \V\mathrm{d}\mu\) with \(\V\) strictly convex, then \(\T_{k+1}=\nabla \V^* \circ \gW\cFm(\mu_k)\) with \(\V^*(x)=\sup_y \langle x,y\rangle - \V(y)\) the Legendre transform of \(\V\). For more general functionals, the problem is implicit and, to the best of our knowledge, cannot be found in closed-form.

Nonetheless, since \(\J\) is convex on \(L^2(\mu_k)\), under the Lipschitz condition ?? (corresponding to \(L\)-smoothness of \(\T\mapsto \cFp(\T_\#\mu_k)\) for all \(k\ge 0\)), we can perform a gradient descent on \(L^2(\mu_k)\), which is of the form, for \(0<\tau\le 1/L\) and \(\tilde \T_k^0 = \id\), \[\forall \ell\ge 0,\;\tilde \T_k^{\ell+1} = \tilde \T_k^\ell - \tau \big(\gW\cFp\big((\tilde \T_k^\ell)_\#\mu_k\big)\circ \tilde \T_k^\ell - \gW\cFm(\mu_k)\big). \label{eq:inner}\tag{6}\] Then, we approximate the solution of [eq:argmin95WDCA95maps] by \(\tilde \T_{k+1}=\tilde \T^L_k\) and \(\mu_{k+1} = (\tilde \T_{k+1})_\#\mu_k\). If \(\mu_k^\ell = (\tilde \T_k^\ell)_\#\mu_k=\frac{1}{n}\sum_{i=1}^n \delta_{x_i^{k,\ell}}\), then the update translates to the particles as \[\label{eq:inner95gd} \forall k\ge 0, \ell\ge 0,\;x_i^{k,\ell+1} = x_i^{k,\ell} - \tau \big(\gW\cFp(\mu_k^\ell)(x_i^{k,\ell}) - \gW\cFm(\mu_k)(x_i^{k,0})\big).\tag{7}\] In practice, one can follow this procedure for any \(\tau\) small enough, and we run the algorithm for \(0< k \le K\) outer steps [eq:argmin95WDCA95maps], and, at each \(k\), \(0<\ell \le M\) inner steps of 7 to minimize \(\J\). Alternatively one could rely on root-finding algorithms instead, such as Newton’s method [74].

0.0.3.4 Connection with Wasserstein Proximal Gradient Algorithm

[31] proposed to solve the DC problem on \(\cPr\) by using a Wasserstein Proximal Gradient Descent scheme [18], alternating between a gradient descent step on \(-\cFm\) and a JKO step on \(\cFp\), i.e. using \[\label{eq:wasserstein95proximal95dc} \begin{cases} \nu_{k+1} = \big(\id + \tau \gW\cFm(\mu_k)\big)_\#\mu_k \\ \mu_{k+1} = \mathop{\mathrm{argmin}}_{\mu\in\cPr}\;\frac{1}{2\tau}\W_2^2(\mu, \nu_{k+1}) + \cFp(\mu). \end{cases}\tag{8}\] In 6, we show that this is equivalent to minimizing an upper bound similar to that of WCCCP 3 , with an additional quadratic cost.

Hence, this scheme can be seen as a lifting of CCCP algorithms with additional quadratic terms [75] which have been in particular used to deal with the stochastic setting [76]–[78]. Note that it would also be possible to use instead a Bregman divergence on \(L^2(\mu_k)\), and that it would be equivalent with a Bregman proximal gradient scheme, see e.g. [20]. Furthermore the WCCCP algorithm can be seen as minimizing in ?? a regularization \(\cF(\T_\#\mu_k)+\D_{\cFm}^{\mu_k}(\T,\id)\) of \(\cF(\T_\#\mu_k)\), and ?? hence corresponds to replacing \(\cF(\T_\#\mu_k)\) by \(\cF(\T_\#\mu_k)+\frac{1}{2\tau}\|\T-\id\|_{L^2(\mu_k)}^2\) in our analysis. [31] focused on functionals whose concave part is a potential energy, i.e. \(\cFm(\mu) = \int \Vm\mathrm{d}\mu\) for \(\Vm:\R^d\to\R\) a convex function. This includes a large class of functionals, and in particular the KL divergence with a possibly non log-concave target if the potential of the target admits a DC decomposition \(\V=\Vp-\Vm\). Our theory in 0.0.3 would also cover these objectives, up to replacing the gradient by the unique subgradient of the negative entropy in the tangent space, as discussed in [29], and verifying that the measure stays regular enough at each iteration, which can be enforced through a regularization [79] to avoid failures [80]. In the next section, we focus instead on the Maximum Mean Discrepancy, which can be decomposed as a DC function, whose concave part is a sum of an interaction and of a potential energies.

0.0.4 DC Decomposition for the Maximum Mean Discrepancy↩︎

We now focus on finding a DC decomposition to the Maximum Mean Discrepancy (MMD) [81]. Given \(k:\R^d\times\R^d\to \R\) a kernel, the squared MMD is defined as \[\forall \mu,\nu\in\cPr,\;\mmd_k^2(\mu,\nu) = \iint k(x,y)\;\mathrm{d}(\mu-\nu)(x)\mathrm{d}(\mu-\nu)(y).\] It is well known that the squared MMD distance can be decomposed as a sum of an interaction energy and a potential energy [33], i.e. \[\label{eq:mmd} \bF(\mu) = \frac{1}{2} \mmd_k^2(\mu,\nu) = \frac{1}{2}\iint k(x,y)\;\mathrm{d}\mu(x)\mathrm{d}\mu(y) + \int \V\;\mathrm{d}\mu + c(\nu),\tag{9}\] with \(\V(x) = - \int k(x,y)\;\mathrm{d}\nu(y)\) and \(c(\nu)=\frac{1}{2}\iint k(x,y)\;\mathrm{d}\nu(x)\mathrm{d}\nu(y)\). The first term is an interaction term and the second a potential. This objective is in general not (geodesically) convex, but only semi-convex [33], i.e. \(\lambda\)-totally convex with \(\lambda\in\R\). Moreover, the performance of WGD to minimize it depends heavily on the kernel as observed in e.g. [25], [33], [82]. For an \(L\)-smooth kernel \(k\), i.e. satisfying \(\|\nabla k(x,y) - \nabla k(x',y')\|_2^2 \le L\big(\|x-x'\|_2^2 + \|y-y'\|_2^2\big)\) for all \(x,x',y,y'\in \R^d\), [31] proposed to use, for \(\alpha \ge L\), \[\label{eq:decomposition95luu} \cFp(\mu) = \alpha \int \|x\|_2^2\;\mathrm{d}\mu(x) + \frac{1}{2}\iint k(x,y)\;\mathrm{d}\mu(x)\mathrm{d}\mu(y) + c(\nu), \, \cFm(\mu) = \int \big(\alpha\|x\|_2^2 - \V(x)\big)\;\mathrm{d}\mu(x)\tag{10}\] as a DC decomposition of 9 . Instead, we propose to obtain a Difference-of-Convex function for this objective by decomposing the kernel itself. For this, we focus on translation-invariant kernels.

0.0.4.0.0.1 MMD with translation-invariant kernel.

A large class of useful kernels are the translation-invariant one, i.e. those of the form \(k(x,y)=\psi(x-y)\) for some symmetric function \(\psi:\R^d\to \R\) [83], [84]. This class includes, among others, the Riesz (a.k.a negative distance) kernel with \(\psi(z)=-\|z\|_2\), the Gaussian kernel with \(\psi(z) = e^{-\|z\|_2^2/(2h)}\) or the inverse multiquadric kernel with \(\psi(z)=(c^2 + \|z\|_2^2)^{-\alpha}\) and \(\alpha >1\) [84]. Assuming that \(\psi\) admits a DC decomposition \(\psi=\psip-\psim\), we obtain the following DC decomposition of the MMD.

We now focus on the subclass of radial kernels, i.e. those for which there exists \(q:\R_+\to\R\) such that \(\psi(z)=q(\|z\|_2^2)\). For such kernels, DC decompositions can be obtained via a Jordan decomposition\[\label{eq:dc95decomposition95jordan} \begin{cases} q_+(x) = q(0) + (q'(0)+A)x + \int_{0}^{x} (x-t) \max\big(0, q''(t)\big) \, \mathrm{d}t, \, \\ q_-(x) = Ax- \int_{0}^{x} (x-t) \min\big(0, q''(t)\big) \, \mathrm{d}t, \end{cases}\tag{11}\] whenever these integrals of the second derivative are computable, and taking \(A\ge \max(0,-q'(0))\) to ensure that \(q_+\) and \(q_-\) are nondecreasing (hence \(\psi_+\) and \(\psi_-\) are convex). Alternatively if \(q\) is analytic, i.e.\(q(x)=\sum_{i\in \mathbb{N}} a_i x^i\), then we can set \(I^+\mathrel{\vcenter{:}}= \{i \, | \, a_i\ge 0\}\) and \(I^-\mathrel{\vcenter{:}}= \{i \, | \, a_i\le 0\}\), \(\qp(x)=\sum_{i\in I^+} a_i x^i\), \(\qm(x)=\sum_{i\in I^-} a_i x^i\). These two choices are explored for the Gaussian kernel (Gauss Jordan vs cosh/sinh) in 0.0.5, and their forms are detailed in 1.

Let \(\Omega=\R^d\) or a compact convex subset and \(S_*=\sup_{x,y\in\Omega}\;\|x-y\|_2\). Define for \(\psi(z)=q(\|z\|_2^2)\), \[\underline\lambda[q]\mathrel{\vcenter{:}}=\inf_{0\le s\le S_*}\min\bigl\{2q'(s),\,2q'(s)+4sq''(s)\bigr\}, \quad \overline{\Lambda}[q] \mathrel{\vcenter{:}}= \sup_{0\le s\le S_*}\;\max \bigl\{2q'(s), 2q'(s)+4s q''(s)\bigr\}.\]

Having \(\underline\lambda[q_{\pm}]\ge 0\) and finite \(\overline{\Lambda}[q_{\pm}]\), which are related to bounds on the Hessian of \(\cF^\pm\), allows for a sufficient condition on a DC decomposition of \(q\) and therefore of the MMD. Owing to this decomposition we can apply 7 and all the convergence results of 0.0.3.

We refer to 1 for decompositions of the Gaussian and (smoothed) Riesz kernels satisfying 8. For more discussion about kernels satisfying 8, we refer to Appendix 0.0.10. In Appendix 0.0.10.3, we also provide sufficient conditions under which the WCCCP scheme converges locally towards a critical point of the MMD in the compact (Prop 14) and non-compact cases (Prop 15).

0.0.5 Numerical Experiments↩︎

We now apply the WCCCP algorithm on the Energy Distance and on the MMD with Gaussian kernel, and compare its performance with the Wasserstein Gradient Descent (WGD) and the Wasserstein Proximal Gradient 8 proposed in [31] denoted Forward-Backward (FB). For more numerical experiments and implementation details, we refer to Appendix 0.0.11¹.

Energy distance. The Energy distance (ED) [85] corresponds to 9 induced by \(\psi(z)=-\|z\|_2\), i.e. \(k(x,y)=-\|x-y\|_2\) and, for \(c(\nu)=-\frac{1}{2} \iint \|x-y\|_2\;\mathrm{d}\nu(x)\mathrm{d}\nu(y)\), \[\frac{1}{2}\ed(\mu,\nu) = -\frac{1}{2}\iint \|x-y\|_2\;\mathrm{d}\mu(x)\mathrm{d}\mu(y) + \int \V\;\mathrm{d}\mu + c(\nu),\;\V(\cdot) = \int \|\cdot -y\|_2\;\mathrm{d}\nu(y).\] While not convex, its Wasserstein gradient flow has a good behavior [86] and has demonstrated good results for different machine learning applications [25], [32], [87]–[89]. Nonetheless, it can be naturally decomposed as a DC functional using 7 with \(\psip=0\) and \(\psim(z) = \|z\|_2\). Thus, we propose to apply the WCCCP algorithm to minimize it.

On [fig:cv_mmd_riesz_shapes_spiral_cat,fig:cv_mmd_riesz_cifar], we minimize the Energy distance with respect to \(\nu_n\) an empirical distribution of \(n=500\) samples drawn uniformly over spiral or cat shapes, and from CIFAR10 [90]. In both cases, we start from an empirical distribution of \(n\) samples of \(\mu_0=\cN(0, I_d)\), and use the stepsize \(\tau=1\) for WGD and FB, as WCCCP can be seen as Mirror Descent algorithm with \(\tau=1\). It allows comparing the three schemes in the same setting, assuming access to an oracle for WCCCP and FB. We add in 6 a comparison with different step sizes for WGD and the same computational budget for WCCCP, i.e. with the same number of gradient evaluations. Results were averaged over 100 different source and target samples on 1 and 5 on [fig:cv95mmd95riesz95cifar]. For both FB and WCCCP, the inner optimization of \(\J\) is performed via gradient descent, possibly with momentum. We display the number of outer iterations, focusing on the behaviour in objective values on 1 (FB and WCCCP having \(M=50\) extra inner iterations). On [fig:cv95mmd95riesz95cifar] instead, to be fair on the computational time (1h30 for WGD, 1h15 for WCCCP on a Nvidia V100 GPU) for \(x\in \R^d\) and \(d\approx \text{3K}\), in total 200K iterations were performed for WGD against 40K outer iterations for WCCCP, and we show snapshots every 20K iterations for WGD, and every 4K iterations for WCCCP. Even with this rescaling, the convergence of WCCCP remains much faster.

MMD with Gaussian kernel. Another natural translation-invariant kernel is the Gaussian kernel \(k(x,y) = e^{-\|x-y\|_2^2/(2h)}\) for which \(q(s)=e^{-s/(2h)}\). However, minimizing the MMD with this kernel has been proven challenging, as its convergence depends a lot on the value of its bandwidth \(h\) [33]. On the one hand, we can use the decomposition 10 proposed by [31] with optimal choice \(\alpha=\tfrac{1}{h}\) (see 1). On the other hand, we can use the DC decomposition from 8 with \(\qp,\qm\) chosen either with the Jordan decomposition 11 or the algebraic one for which \(\qp(2hs)=\cosh(s), \qm(2hs)=\sinh(s)\). Their properties are detailed in 1 and derived in Appendix 0.0.10.

We compare on 2 the minimization of the squared MMD with Gaussian kernel and bandwidth \(h=10\) between WGD and WCCCP with these three decompositions. We use \(n=500\) particles initially sampled from \(\mu_0 = \cN(5\mathbb{1}_2, I_2)\) and set the target as an empirical distribution \(\nu_n = \frac{1}{n}\sum_{i=1}^n \delta_{y_i}\) with \(y_i\) sampled independently from \(\nu=\cN(0, \Sigma)\) with \(\Sigma = \bigl(\begin{smallmatrix} 1 & 0.5 \\ 0.5 & 1 \end{smallmatrix}\bigr)\). The results are averaged over 25 runs with different source and target samples. The inner problems of the WCCCP scheme with \(\cosh/\sinh\) are solved using a gradient descent with momentum, with \(m=0.9\) and step size \(\tau=5\cdot 10^{-4}\) for \(M=250\) iterations, while for the WCCCP scheme with Jordan decomposition, we use a gradient descent with \(\tau=0.1\) and \(M=500\) iterations. In particular, this decomposition is smoother and less prone to numerical instabilities than the \(\cosh/\sinh\) one. We observe that WGD and WCCCP with the decomposition 10 get stuck in a local minimum where some samples drift away of the target distribution, while WCCCP with \(\cosh/\sinh\) and Jordan decompositions converge much better. We note that the Jordan decomposition can still sometimes be stuck in local minima, which are nonetheless better than WGD. We hypothesize that this is due to the inexact solver used at each iteration of [eq:argmin95WDCA95maps]. On the other hand, all the samples of WCCCP with \(\cosh/\sinh\) appear to eventually converge, instead of being sent away, but at a slower rate. We also compare the results with the FB algorithm in Appendix 0.0.11, which also showcases results that depend heavily on the choice of DC decomposition.

0.0.6 Conclusion↩︎

In this work, we lifted the convex-concave procedure to the Wasserstein space, and analyzed its convergence in the convex and non-convex settings. Then, we used it to minimize the Maximum Mean Discrepancy, and showed improved performance over the Wasserstein Gradient Descent for the negative distance and Gaussian kernel. Nonetheless, the improved convergence of WCCCP strongly depends on the choice of the DC decomposition of the kernel. Future work will therefore focus on better understanding the impact of different DC decompositions on the performance, with the goal of designing more effective, automatic, and adaptive decomposition strategies [91]. On the theoretical side, taking into account the inexact solvers used for the inner optimization problems would also be a natural avenue for future work.

Appendix↩︎

The appendix is organized as follows. In Appendix 0.0.7, we discuss some limitations of the paper. In Appendix 0.0.8, we detail the theoretical analysis of the WCCCP algorithm in the convex case, leveraging the mirror and Bregman proximal descent formulations. In Appendix 0.0.9, we discuss the derivations of Polyak-Łojaziewicz inequalities for DC functionals. In Appendix 0.0.10, we provide the full theoretical analysis of DC decompositions of the MMD. In Appendix 0.0.11, we include more details about the numerical experiments. Finally, in Appendix 0.0.12, we state all the proofs.

0.0.7 Limitations↩︎

Our empirical findings show that the improved convergence of WCCCP on MMD strongly depends on the choice of the DC decomposition of the kernel \(k\). Our current strategy is restricted to a few simple decomposition of the kernels, without taking into account the geometry of the optimization problem. Hence, future work could focus on developing automatic and adaptative ways to find DC decompositions of the kernels better suited to the problem of minimizing the MMD, e.g. based on algebraic decompositions and polynomial DC decompositions minimizing a suitable objective [91]–[93].

This work also focused on DC decompositions of MMD, which can be written as a sum of simple functionals, which are potential and interaction energies. For MMD with the ubiquitous radial kernels, we were thus able to derive DC decompositions based on decompositions of functions on \(\R\to\R\). Future work could focus on the problem of deriving DC decomposition of more complex functionals.

On the theoretical part, as highlighted in 0.0.1, there are versions of Polyak-Łojaziewicz (PL) inequalities for both the Wasserstein space [34], [35] and DC functions [66], [70]–[72]. So far our framework does not include them despite preliminary research. We discuss reasons for this in 0.0.9.

We also did not take into account that we actually solve each inner scheme approximately, which leads to a gap between theory and practice.

0.0.8 Theoretical Analysis of WCCCP in the Convex Case↩︎

We focus in this section on the case where \(\cF=\cFp-\cFm\) is also convex along a curve of interest which we will detail. For this, we will leverage 1, where we showed that [eq:argmin95WDCA95maps] is equivalent to both a mirror descent and a Bregman proximal descent in the Wasserstein space. Consequently, on the one hand, we will inherit the results from [20] for mirror descent, and on the other hand we will derive novel convergence results for Bregman proximal descent.

0.0.8.0.0.1 Relative convexity and smoothness.

First, we need to introduce the notions of relative convexity and smoothness. Let \(\alpha,\beta \ge 0\). Following [20], we say that \(\cF\) is \(\alpha\)-convex relative to \(\cG:\cPr\to\R\) along \(t\mapsto \big((1-t)\T+t\sS\big)_\#\mu\) if \(\D_\cF^\mu(\T,\sS)\ge \alpha\D_\cG^\mu(\T,\sS)\). Equivalently, we have that \(\cF-\alpha\cG\) is convex along this curve, and \[\forall t\in [0,1],\;\cF(\mu_t)\le (1-t) \cF(\T_\#\mu) + t\cF(\sS_\#\mu) - \alpha t(1-t) \D_\cG^\mu(\T,\sS).\] Likewise, \(\cF\) is \(\beta\)-smooth relative to \(\cG\) along this curve if \(\D_\cF^\mu(\T,\sS)\le \beta \D_\cG^\mu(\T,\sS)\). These notions enable lifting the notion of relative convexity and smoothness [64], [94] to \(\cPr\) for differentiable functionals.

If the convexity holds for all \(\T,\sS\in L^2(\mu)\), \(\mu\in\cPr\) and \(\cG(\mu) = \int \tfrac12 \|\cdot\|_2^2\mathrm{d}\mu\), then we say that \(\cF\) is \(\alpha\)-totally convex [59]–[61]. Following [29], if convexity only holds for \(\T=\id\) and \(\sS\) gradient of convex functions, it coincides with strong convexity along geodesics [29]. Convexity along generalized geodesics corresponds to both \(\T\) and \(\sS\) being the gradients of convex functions. Examples of totally convex functionals include potential and interaction energies, provided \(\V,\W\) are convex, lower semi-continuous and have a negative part with quadratic growth [29]. Note also that all three notions of convexity are equivalent for continuous functionals (for \(d\ge 2\)) [59], [61].

0.0.8.0.0.2 Bregman Wasserstein distance.

Similarly to [20], let us introduce the Bregman Optimal Transport problem \(\W_\phi\) associated with a Wasserstein differentiable functional \(\phi:\cPr\to\R\), defined for \(\mu,\nu\in\cPr\) as \[\W_\phi(\nu,\mu) \mathrel{\vcenter{:}}= \inf_{\gamma\in\Pi(\nu,\mu)}\;\phi(\nu) - \phi(\mu) - \int \langle \gW\phi(\mu)(y), x-y\rangle\;\mathrm{d}\gamma(x,y).\] By [20], if \(\mu\in\cPa\), then there exists \(\T\in L^2(\mu)\) such that \(\W_\phi(\nu,\mu) = \D_{\phi}^\mu(\T,\id)\), i.e. this problem admits an optimal transport map.

0.0.8.0.0.3 Convergence results.

We now provide a first convergence result for \(\bF\) smooth and strongly convex relative to \(\cFp\), relying on the mirror descent formulation ?? and the results of [20]. More precisely, \(\bF\) needs to be smooth along iterates and convex along what would be the analog of geodesics on the space \((\cPr, \W_{\cFp})\).

Assuming analog conditions as for the convergence of the JKO scheme adapted to the Bregman Proximal Gradient Descent ?? , we obtain the following linear rate convergence result.

The convergence results are given in the geometry induced by the Bregman divergence with potential given by \(\cFm\).

0.0.9 Enquiry on Polyak-Łojaziewicz Inequality for DC Functionals↩︎

An analog of the DC PL inequality of [71], writing it only using the Bregman divergences of \(\cFp\) and \(\cFm\), reads as follow

Under this ideal condition, we show that we have the following linear convergence rate, analogously to [66] and [71].

If \(\eta_1 \ge 1\), then we can set \(\eta_1=1\) and we have convergence in one step.

Although [71] in the Euclidean case shows that [66] implies their [71], the latter is difficult to check directly for a given functional and is more of a target inequality to establish. We thus want to argue that [66] is more promising to stand as DC PL inequality and has received a more developed discussion of when it holds, e.g. under strong convexity. However this definition rests upon Fenchel duality on \(\R^d\), for which there is no publicly available counterpart on the Wasserstein space at the time of this submission. Consequently it is unclear for now whether 1 holds for functionals with strongly convex DC decompositions.

While [31] do rest upon a Łojaziewicz inequality, the latter is unrelated to Bregman divergences. [31] achieve instead their linear rates using the extra Wasserstein regularization appearing in ?? . While we discussed in the similarity of the two approaches 0.0.3.4, one cannot use their theory without the Wasserstein regularization as otherwise their constant become vacuous.

0.0.10 DC Theory for MMD↩︎

Let \(k\) be a translation-invariant kernel of the form \(k(x,y)=\psi(x-y)\) for all \(x,y\in\R^d\), with \(\psi\) admitting a DC decomposition \(\psi=\psip-\psim\). Then by 7, the squared MMD 9 admits the DC decomposition \(\bF=\cFp-\cFm\) where for all \(\mu\in\cPr\), \[\label{slukarqn} \left\{ \begin{array}{ll} \cFp(\mu) = \tfrac12 \iint \psip(x-y)\;\mathrm{d}\mu(x)\mathrm{d}\mu(y) + \int \Vm\mathrm{d}\mu + c, & \Vm(\cdot) = \int \psim(\cdot-y)\;\mathrm{d}\nu(y), \\ \cFm(\mu) = \tfrac12 \iint \psim(x-y)\;\mathrm{d}\mu(x)\mathrm{d}\mu(y) + \int \Vp\mathrm{d}\mu, & \Vp(\cdot) = \int \psip(\cdot-y)\;\mathrm{d}\nu(y), \end{array} \right.\tag{12}\] and \(\cFp,\cFm\) are both totally convex. Moreover, is \(\psip\) and \(\psim\) are respectively \(\ap\ge 0\) and \(\am\ge 0\) strongly convex, then \(\cFp\) is \(\am\)-totally convex and \(\cFm\) is \(\ap\)-totally convex as they inherit the strong convexity of the potential, while the interaction terms are only convex, see [29].

We will now focus on radial kernels. Let \(\Omega \in \mathbb{R}^d\) be a nonempty compact convex set. Consider the geodesically convex set \(\mathcal{H}=\mathcal{P}_2(\Omega)\). Let \[S_* \mathrel{\vcenter{:}}= \sup_{x,y \in \Omega} \|x-y\|_2.\]

We consider in what follows: \[\psi^{\pm}(z) := q_{\pm}(\|z\|_2^2),\quad z \in \Omega -\Omega\mathrel{\vcenter{:}}= \{x-y \, | \, x,y\in\Omega\},\] where \(q_{\pm}\in C^2([0,S_*])\). The MMD functional \(\mathcal{F}\) is therefore defined as a function of \(q^{\pm}.\)

0.0.10.1 Strong Convexity of \(\cFm\) and \(\cFp\) for Radial Kernels

We first study the Hessian of radial functions, in order to derive conditions for \(\psi_{\pm}\) to be (strongly) convex and being able to apply 7.

Let \(q\in C^2([0,S_*])\) and define: \[\underline\lambda[q] \mathrel{\vcenter{:}}= \inf_{0\le s\le S_*}\min\bigl\{2q'(s),\,2q'(s)+4sq''(s)\bigr\}, \quad \overline{\Lambda}[q] \mathrel{\vcenter{:}}=\sup_{0\le s\le S_*}\max\bigl\{2q'(s),\,2q'(s)+4sq''(s)\bigr\}.\] We use the shorthands \(\lambda_\pm:=\underline\lambda[q_\pm]\) and \(\Lambda_\pm:=\overline{\Lambda}[q_\pm]\). First, we compute the Hessian of \(\psi(z)=q(\|z\|_2^2)\) and bound its eigenvalues using \(\underline\lambda[q]\) and \(\overline{\Lambda}[q]\).

Now, building on the previous Lemma, we deduce sufficient condition under which \(\psi\) is convex.

This lemma provides sufficient conditions to get a DC decomposition of the form ?? by 7.

Building on these two lemmas, we show that several usual kernels admit such decompositions. We begin by discussing several DC decompositions for the Gaussian kernel that we can use for MMD. The first one is based on a remark from [31], and the other ones are based on observing that it is a radial kernel with \(q(t)=e^{-t/(2h)}\). Hence, we simply need to find a DC decomposition of \(q=\qp-\qm\) which satisfies the assumptions in 2.

0.0.10.1.0.1 Gaussian kernel.

Recall that the Gaussian kernel is a radial kernel with \(q(t)=e^{-\alpha t}\) with \(\alpha\ge 0\) (often taken as \(\alpha=1/(2h)\), \(h\) being the bandwidth).

[31] observed in their Appendix A.2 that for \(k\) differentiable and \(L\)-smooth, i.e. satisfying \(\|\nabla k(x,y) - \nabla k(x',y')\|_2^2 \le L\big(\|x-x'\|_2^2 + \|y-y'\|_2^2)\) for all \(x,x',y,y'\in \R^d\), a DC decomposition of 9 is given by \[\label{eq:decomposition95luu95v2} \cFp(\mu) = \alpha \int \|\cdot\|_2^2\;\mathrm{d}\mu + \frac{1}{2}\iint k(x,y)\;\mathrm{d}\mu(x)\mathrm{d}\mu(y) + c(\nu), \quad \cFm(\mu) = \int \big(\alpha \|x\|_2^2 - \V(x)\big)\;\mathrm{d}\mu(x)\tag{13}\] for any \(\alpha \ge L\). Actually we can specify in the next example the best choice of \(\alpha\) for the Gaussian kernel.

Another natural DC decomposition is based on the observation that \(e^{-t}=\cosh(t) - \sinh(t)\). We show in the next lemma that this gives a valid decomposition for 7.

Applying 7, we can deduce that \(\cFp\) is \(2\alpha>0\)-totally convex while \(\cFm\) is only totally convex. Hence, we can apply 5 and obtain convergence of WCCCP towards a stationary point in a sublinear rate.

Next, we also consider looking at \(q(t)=e^{-\alpha t}\) directly. Its derivative is \(q'(t)=-\alpha e^{-\alpha t}\) and second derivative \(q''(t)=\alpha^2 e^{-\alpha t}\). We note that it is convex but it does not satisfy 2 as \(q'<0\). However, adding a linear term large enough, we can obtain a DC decomposition of \(q\) satisfying the assumptions of 2. For this, we apply 11 with \(A=\alpha\).

Note that for this decomposition the definition of the bounds \(\lambda\) do not need to be restricted to a compact set.

0.0.10.1.0.2 Smoothed Riesz Kernel.

We now deal with the negative distance (Riesz) kernel \(k(x,y)=-\|x-y\|_2\). Since it is not differentiable in \(x=y\), we instead focus on a smoothed version of it, defined as \(k_\varepsilon(x,y) = -\sqrt{\varepsilon + \|x-y\|_2^2}\). In this case, we have \(q(t) = -\sqrt{t + \varepsilon}\) for all \(t\in\R\).

0.0.10.1.0.3 Rational Quadratic Kernel.

We now turn to the rational quadratic kernel \(k(x,y)= \psi(x-y),\) where \(\psi(z)=\frac{1}{(c^2 + \|z\|_2^2)^{\alpha}}, \alpha \geq 1\).

0.0.10.2 Lipschitz Continuity of \(\gW\mathcal{F}^+\) and Stationarity

We now show that for a DC decomposition based on radial kernels as in 2, the Wasserstein gradient of \(\gW\cFp\) satisfies a Lipschitz condition, which corresponds to the assumption in 5.

Using this Proposition, we can apply 5 to obtain a sublinear rate over the minimum of the Wasserstein gradient over the scheme.

0.0.10.3 Critical points and Local Convergence

In the following we analyze the conditions under which we obtain local convergence of WCCCP for the MMD functional with base space \(\Omega\). All the statements are in the weak topology and assume a continuous kernel \(k\).

First,we consider a condition on the initialization w.r.t. the critical value gap as follows:

Under Assumption 1, we have the following local convergence result, interpreted as: if one starts with a near optimal initialization, by descent, there must be convergence. The condition \(\cF(\mu_0)< c_{*}\) entails that \(\nu\) is an isolated critical point.

The following proposition does not require compactness but assumes local quadratic growth of the MMD and uniqueness of critical points in a local neighborhood of \(\nu\), and it guarantees local convergence:

0.0.11 Numerical Applications on MMD↩︎

We now detail the experiments of 0.0.4, as well as extra experiments. We first detail the experiments on the Energy Distance, and then focus on MMD with the Gaussian kernel. All the numerical applications are done on a Nvidia V100 GPU.

0.0.11.1 Energy Distance

We recall that the Energy distance [85] is of the form, for \(\mu,\nu\in\cPr\), \[\label{eq:def95ED} \ed(\mu,\nu) = - \iint \|x-y\|_2\;\mathrm{d}(\mu-\nu)(x)\mathrm{d}(\mu-\nu)(y),\tag{14}\] which by 7, for a fixed target \(\nu\), can be decomposed as \(\tfrac12\ed(\mu,\nu) = \cFp(\mu) - \cFm(\mu)\) where \[\cFp(\mu) = \iint \|x-y\|_2\;\mathrm{d}\nu(y)\mathrm{d}\mu(x) + c(\nu), \quad \cFm(\mu) = \frac{1}{2} \iint \|x-y\|_2 \;\mathrm{d}\mu(x)\mathrm{d}\mu(y),\] with \(c(\nu)=-\frac{1}{2} \iint \|x-y\|_2 \;\mathrm{d}\nu(x)\mathrm{d}\nu(y)\). The functional \(\cFp\) is a potential energy with \(V(x)=\int \|x-y\|_2\;\mathrm{d}\nu(y)\) and \(\cFm\) an interaction energy, and both are totally convex. If we use the smoothed Riesz kernel \(k(x,y)=-\sqrt{\varepsilon + \|x-y\|_2^2}\) and assume that \(\mu,\nu\in\cP_2(\Omega)\) for \(\Omega\) a compact convex set, then by 5 and 1, \(\cFp\) is also \(\alpha\)-totally convex with \(\alpha=\frac{\varepsilon}{(\varepsilon + S_*)^{\frac{3}{2}}}\) with \(S_*=\sup_{x,y\in\Omega}\;\|x-y\|_2\). In practice, we use mostly the non-smooth Riesz kernel \(k(x,y)=-\|x-y\|_2\) as it works well in practice [25]. Nonetheless, some smoothed versions based on convolutions have been also shown to have more favorable theoretical properties [95].

0.0.11.1.0.1 Experiments on shapes.

For the experiment of 1, we compare the Wasserstein gradient descent 2 , the Forward-Backward scheme 8 from [31], and our scheme WCCCP [eq:argmin95WDCA95maps]. We use a spatial discretization as described in 0.0.3.3, i.e. we first sample \(n=500\) particles from \(\mu_0=\cN(0,I_2)\), and \(n\) uniform independent particles from the target shape to obtain the target distribution \(\nu_n\). Then, at each iteration, we compute the map \(\T_{k+1}\) and apply it to move each of the particles. For WGD and FB, we use as step size \(\tau=1\). For WCCCP and FB, we solve each inner optimization problem with a gradient descent with momentum \(m=0.9\) for \(M=50\) iterations and step size \(0.1\). We report on 1 the values of the objective \(\ed^2(\mu_k, \nu_n)\) for the cat and spiral shapes depending on the number of outer iterations (hence WGD is much faster but still converges to a local minimum and does not improve further). We performed the experiment 100 times, and report the average values with standard deviation to get confidence intervals. We add on 4 the results for the heart and disk shapes, as well as particles for the WCCCP at iterations 0, 10 and 500. On 5, we performed the same experiment with target samples from a Gaussian \(\nu=\cN(0,\Sigma)\) for \(\Sigma = \bigl(\begin{smallmatrix} 1 & 0.5 \\ 0.5 & 1 \end{smallmatrix}\bigr)\) (top), and with target samples from a mixture of 3 Gaussian with equal weights, means \(m_1=(0,0)\), \(m_2=(3,-1)\) and \(m_3=(1,4)\), and covariances \(\Sigma_1=\bigl(\begin{smallmatrix} 1 & 0.5 \\ 0.5 & 2\end{smallmatrix}\bigr)\), \(\Sigma_2=I_2\) and \(\Sigma_3=\bigl(\begin{smallmatrix} 3 & 0.5 \\ 0.5 & 1\end{smallmatrix}\bigr)\). For both, we use the same initial distribution \(\mu_0=\cN(5,I_2)\) and the same hyperparameters as the shape experiment.

We also report on 6 the convergence for the cat target uniform distribution, with the exact same number of iterations between WCCCP and WGD with different step sizes. We observe that for \(\tau=0.1\), WGD converges better but slower than for \(\tau=1\). With the same computational budget for WCCCP with \(M=50\), \(K=400\), we observe that the algorithm converges in similar compuational time.

0.0.11.1.0.2 Experiments on images.

On [fig:cv95mmd95riesz95cifar], we performed the same experiment with target samples from the CIFAR10 dataset, whose images are of size \(3\times 32 \times 32\). More precisely, we sampled 50 points by class, and hence also worked with \(n=500\) particles. We started the flows from \(\mu_0=\cN(0,I_d)\). We compared WGD with stepsize \(\tau=1\) and WCCCP with a gradient descent to solve the inner optimization scheme with \(\tau=1\) and \(20\) iterations. We ran WGD for 200K iterations and WCCCP for 40K iterations, with 20 iterations to solve each of the subproblems [eq:argmin95WDCA95maps]. We choose to use a different number of iterations to be fair in comparing the two methods, as WCCCP solves each inner problem in closed-form, and is thus less computationally expensive. With a Nvidia v100 GPU, WCCCP took about 1h15 while WGD took 1h30. On [fig:cv95mmd95riesz95cifar], we plot samples along the scheme of WGD every 20K iterations, and samples along the scheme of WCCCP every 4K iterations. Below, we plot the value of the loss across iterations, where we rescaled the abscissa for WCCCP to match the abscisse of WGD as done on the image samples. We observe overall a much faster convergence of WCCCP compared to WGD on this high dimensional challenging dataset. On 7, we add more samples along the schemes of WGD and WCCCP, every 2K iterations for WCCCP and 10K iterations for WGD.

We also performed the experiment on MNIST and report the results on 7. Here, the samples are reported every 1K iterations for WCCCP and 5K iterations for WGD. The schemes took about 18 minutes to run for both WGD and WCCCP.

We show on 8 the evolution of the Energy distance along the scheme, averaged over 5 runs with different samples of the source and of the target. We observe that on the MNIST experiment, both schemes converge, but WCCCP converges faster (even though the iterations are rescaled). On the CIFAR10 experiment, even after 200K iterations, WGD is very far from converging, whereas WCCCP has already converged. Hence, WCCCP can be promising to accelerate convergence in high dimensions (here \(d=3\times 32 \times 32=3072\)).

0.0.11.2 MMD with Gaussian Kernel

We now focus on MMD with Gaussian kernel \(k(x,y)=e^{-\|x-y\|_2^2/(2h)}\) for a bandwidth \(h>0\).

0.0.11.2.0.1 Decomposition based on the radial kernel.

By 1, we can have a DC decomposition of the squared MMD if we find \(\qp,\qm\) satisfying \(\qp',\qm',\qp'',\qm''\ge 0\) and \(q=\qp-\qm\). We discuss here two natural decompositions of \(q:t\mapsto e^{-t/(2h)}\).

The first one is based on the remark that \(e^{-z} = \cosh(z) - \sinh(z)\) since \(\cosh(z) = \big(e^z+e^{-z})/2\) and \(\sinh(z)=\big(e^z - e^{-z})/2\), and hence \(\qp(t)=\cosh\big(t/(2h)\big)\), \(\qm(t)=\sinh\big(t/(2h)\big)\). This is based on the algebraic decomposition described in 0.0.4 as \(\cosh(t)=\sum_{k \;\mathrm{even}} \frac{t^{k}}{k!}\) and \(\sinh(t)=\sum_{k\;\mathrm{odd}} \frac{t^k}{k!}\). This decomposition is valid to apply 7 as showed in 3. We show this decomposition on 3 (a). We note that both \(\qp\) and \(\qm\) tend to take very large values, which might be prone to numerical instabilities, in particular for small bandwidths.

The second one is based on decomposing the second derivative, and taking the convex part as the function having as second derivative the non-negative part, and the convex part as the function having as second derivative the opposite of the non-positive part. This can be done using the Jordan decomposition 11 . The exact formula gives \(\qp(s)=e^{-\alpha s} + \alpha s\) and \(\qm(s)=\alpha s\). This also satisfies 7 as showed in 12. We plot on 3 (b) the resulting DC decomposition. We observe here that this decomposition does not blow up, and should be numerically more stable.

0.0.11.2.0.2 Experiments.

The Wasserstein gradient descent on \(\cF(\mu)=\frac{1}{2} \mmd_k^2(\mu,\nu)\) with \(k(x,y)=e^{-\|x-y\|_2^2/(2h)}\) is known to heavily depend on the choice of the bandwidth [25], [33], and might not converge well in practice. In particular, on simple examples such as a single Gaussian distribution, [33], [39] observed that some particles get stuck further from the mode, and never converge. In our experiments, we compare several DC decomposition in the WCCCP algorithm and test for which decompositions we obtain a better convergence or not.

For this, we choose the same setting as [39]. We first take \(n=500\) samples from \(\mu_0=\cN(5, I_2)\) to get the initial distribution. Then, we take \(n=500\) samples from the target distribution \(\nu\). We experiment with 2 choices for \(\nu\). The first one is a Gaussian \(\nu=\cN(0,\Sigma)\) with \[\Sigma = \begin{pmatrix} 1 & 0.5 \\ 0.5 & 1 \end{pmatrix}.\] The second is a mixture of 3 Gaussian with uniform weights, means \(m_1=(0, 0)\), \(m_2=(3, -1)\), \(m_3=(1, 4)\) and covariances \[\Sigma_1 = \begin{pmatrix} 1 & 0.5 \\ 0.5 & 1\end{pmatrix}, \quad \Sigma_2=I_2,\quad \Sigma_3=\begin{pmatrix} 3 & 0.5 \\ 0.5 & 1 \end{pmatrix}.\] Moreover, the bandwidth is fixed to \(h=10\).

In each experiment, we use for WGD and FB a stepsize \(\tau=1\). For FB and WCCCP, we optimize the inner problem with a gradient descent with momentum \(m=0.9\), step size \(\tau=5\cdot 10^{-4}\) and for 250 iterations.

On 2, we show the evolution of the loss over 200K iterations of the algorithms, comparing WGD with WCCCP for the DC decomposition 13 , and the DC decomposition of the radial kernel based on \(\cosh\) and \(\sinh\) as presented in 3 as well as the decomposition based on the Jordan decomposition presented in 12. We show the evolution of the objective over the iterations and average the results over 100 different set of initial and target samples.

We observe that the two radial decompositions perform better than WGD and the DC decomposition 13 . The \(\cosh/\sinh\) decomposition seems to converges more consistently than the one based on the Jordan decomposition, even though it takes longer to converge. Indeed, the Jordan decomposition converges much faster, but also shows a higher variance, and sometimes, particles get stuck away from the target mode, which does not seem to happen with the \(\cosh/\sinh\) decomposition.

We also show the result on 10 of the same experiment on the Gaussian mixture target. On this target, the Jordan decomposition outperforms all the other. On 9, we add the loss for one run and \(n=50\). We observe that WCCCP with Jordan decomposition can plateau during the training, and seems to have several regime of convergences for the Gaussian target. Moreover, it is much faster in its first iterations than the \(\cosh/\sinh\) decomposition.

On 11, we plot the evolution of the loss for the Gaussian and Gaussian mixture targets, comparing WGD with WCCCP and the Forward-Backward (FB) algorithm studied in [31]. We compare WCCCP and FB with two DC decompositions: the one based on the DC decomposition of the radial kernel with \(\cosh/\sinh\) as presented in 3 and the one based on the DC decomposition 13 . We observe that both WCCCP and FB performance heavily depend on the choice of the DC decomposition. In particular, the \(\cosh/\sinh\) decomposition seems to work well for both algorithms, while the iterates of the decomposition 13 gets stuck in a local minimum as WGD. We add on 12 and 13 the evolutions of particles over the flows fo WGD, and WCCCP and FB for both decompositions.

0.0.12 Proofs↩︎

0.0.12.1 Proof of 1

Let \(k\ge 0\), \(\mu_k\in\cPr\). On one hand, we have for any \(\T\in L^2(\mu_k)\), setting \(\J\) as the r.h.s.of 3 , \[\begin{align} \J(\T) &\mathrel{\vcenter{:}}= \cFp(\T_\#\mu_k) - \cFm(\mu_k) - \langle \gW\cFm(\mu_k), \T-\id\rangle_{L^2(\mu_k)} \\ &= \cFp(\T_\#\mu_k) - \cFp(\mu_k) - \langle \gW\cFp(\mu_k), \T-\id\rangle_{L^2(\mu_k)} \\ &\qquad\qquad\qquad +\cFp(\mu_k) + \langle \gW\cFp(\mu_k), \T-\id\rangle_{L^2(\mu_k)} \\ &\quad - \cFm(\mu_k) - \langle \gW\cFm(\mu_k), \T-\id\rangle_{L^2(\mu_k)} \\ &= \D_{\cFp}^{\mu_k}(\T,\id) + \cFp(\mu_k)-\cFm(\mu_k) + \langle \gW\cFp(\mu_k)-\gW\cFm(\mu_k), \T-\id\rangle_{L^2(\mu_k)} \\ &= \D_{\cFp}^{\mu_k}(\T,\id) + \bF(\mu_k) + \langle \gW\bF(\mu_k), \T-\id\rangle_{L^2(\mu_k)}. \end{align}\] Hence, [eq:argmin95WDCA95maps] is equivalent with a mirror descent on \(\bF\) with geometry induced by the Bregman divergence with Bregman potential \(\cFp\).

On the other hand, we also have, for any \(\T\in L^2(\mu_k)\), \[\begin{align} \J(\T) &= \cFp(\T_\#\mu_k) - \cFm(\mu_k) - \langle \gW\cFm(\mu_k), \T-\id\rangle_{L^2(\mu_k)} \\ &= \cFp(\T_\#\mu_k) - \cFm(\T_\#\mu_k) + \cFm(\T_\#\mu_k) - \cFm(\mu_k) - \langle \gW\cFm(\mu_k), \T-\id\rangle_{L^2(\mu_k)} \\ &= \bF(\T_\#\mu_k) + \D_{\cFm}^{\mu_k}(\T,\id). \end{align}\] Hence, [eq:argmin95WDCA95maps] is also equivalent with a Bregman proximal descent on \(\bF\) with geometry induced by the Bregman divergence with Bregman potential \(\cFm\).

0.0.12.2 Proof of 2

\[\begin{align} \bF(\mu_{k+1})&=\bF(\mu_k)+\cFp(\mu_{k+1})-\cFp(\mu_{k})-\cFm(\mu_{k+1})+\cFm(\mu_{k}) \\ &=\bF(\mu_k)-\D_{\cFm}^{\mu_k}(\T_{k+1}, \id)+\cFp(\mu_{k+1})-\cFp(\mu_{k}) - \langle \gW\cFm(\mu_k), \T_{k+1}-\id\rangle_{L^2(\mu_k)} \\ &= \bF(\mu_k) - \D_{\cFm}^{\mu_k}(\T_{k+1},\id) - \Ck. \end{align}\] If \(\cFp\) is W-differentiable, observe that by the first order conditions in [eq:argmin95WDCA95maps], \(\gW\cFp(\mu_{k+1})\circ\T_{k+1} = \gW\cFm(\mu_k)\), and thus \(\Ck=\D_{\cFp}^{\mu_k}(\id,\T_{k+1})\).

0.0.12.3 Proof of 3

If \(\Ck=0\), then \(\id\in \mathop{\mathrm{argmin}}_{\T\in L^2(\mu_k)}\;\cFp(\T_\#\mu_k) - \cFm(\mu_k) - \langle \gW\cFm(\mu_k), \T-\id\rangle_{L^2(\mu_k)}\). Hence the first order condition gives \(0 = \gW\cFp(\id_\#\mu_{k})\circ \id - \gW\cFm(\mu_k)=\gW\cF(\mu_k)\).

0.0.12.4 Proof of 4

We just use a telescopic sum on ?? , discarding \(\D_{\cFm}^{\mu_k}(\T_{k+1}, \id)\) (nonnegative by convexity of \(\cFm\)), and then use the observation 5 .

0.0.12.5 Proof of 5

Since \(\cFp\) and \(\cFm\) are respectively \(\ap\) and \(\am\)-convex along iterates with respect to \(\mu\mapsto \int \tfrac12 \|\cdot\|_2^2\;\mathrm{d}\mu\), then \(\D_{\cFp}^{\mu_k}(\id,\T_{k+1}) \ge \frac{\ap}{2}\|\T_{k+1} - \id\|_{L^2(\mu_k)}^2\) and \(\D_{\cFm}^{\mu_k}(\T_{k+1},\id) \ge \frac{\am}{2}\|\T_{k+1}-\id\|_{L^2(\mu_k)}^2\). Thus, for all \(k\in\{0,\dots,K-1\}\), \[\begin{align} \frac{\ap + \am}{2}\|\T_{k+1}-\id\|_{L^2(\mu_k)}^2 &\le \D_{\cFp}^{\mu_k}(\id, \T_{k+1}) + \D_{\cFm}^{\mu_k}(\T_{k+1}, \id) \\ &\stackrel{ \eqref{eq:iterates95difference95gap} }{=} \cF(\mu_k) - \cF(\mu_{k+1}). \end{align}\] Hence, by 5 , \[\begin{align} \min_{k\in\{0,\dots,K-1\}}\;\|\T_{k+1}-\id\|_{L^2(\mu_k)}^2 &\le \frac{2}{\ap + \am} \min_{k\in\{0,\dots,K-1\}}\;\bF(\mu_k)-\bF(\mu_{k+1}) \\ &\le \frac{2}{\ap+\am} \cdot \frac{\bF(\mu_0)-\bF(\mu_K)}{K}. \end{align}\]

Under the additional assumption that \(\|\gW\cFp(\mu_k)\circ\T_{k+1}-\gW\cFp(\mu_k)\|_{L^2(\mu_k)} \le L \|\T_{k+1}-\id\|_{L^2(\mu_k)}\), then by the first order condition 4 , \[\begin{align} \|\gW\bF(\mu_k)\|_{L^2(\mu_k)} &= \|\gW\cFp(\mu_k) - \gW\cFm(\mu_{k})\|_{L^2(\mu_k)} \\ &= \|\gW\cFp(\mu_k) - \gW\cFm(\mu_{k+1})\circ\T_{k+1}\|_{L^2(\mu_k)} \\ &\le L\|\T_{k+1}-\id\|_{L^2(\mu_k)}. \end{align}\] Hence, combining this with the previous result, \[\min_{k\in \{0,\dots,K-1\}} \;\|\gW\cF(\mu_k)\|_{L^2(\mu_k)}^2 \le \frac{2L^2}{\ap+\am} \cdot \frac{\bF(\mu_0)-\bF(\mu_K)}{K}.\]

0.0.12.6 Proof of 6

Let \(\sS_{k+1}\mathrel{\vcenter{:}}=\id + \tau \gW\cFm(\mu_k)\). First, notice that since \(\nu_{k+1}\in\cPa\) by assumption, then 8 is equivalent to \[\left\{ \begin{array}{ll} \sS_{k+1} = \mathop{\mathrm{argmin}}_{\sS\in L^2(\mu_k)}\;\frac{1}{2\tau} \|\sS-\id\|_{L^2(\mu_k)}^2 - \langle \gW\cFm(\mu_k), \sS-\id\rangle_{L^2(\mu_k)}, & \nu_{k+1} = (\sS_{k+1})_\#\mu_k, \\ \T_{k+1} = \mathop{\mathrm{argmin}}_{\T\in L^2(\nu_{k+1})}\;\frac{1}{2\tau}\|\T-\id\|_{L^2(\nu_{k+1})}^2 + \cFp(\T_\#\nu_{k+1}), & \mu_{k+1} = (\T_{k+1})_\#\mu_k. \end{array} \right.\] The first equation is obtained by the first order condition, and the second by observing that \(\T_{k+1}\) is necessarily an OT map between \(\nu_{k+1}\) and \(\mu_{k+1}\). Indeed, as \(\nu_{k+1}\in\cPa\), by Brenier’s theorem, there exists a unique OT map between \(\nu_{k+1}\) and \(\mu_{k+1}\), and if \(\T_{k+1}\) is not such that \(\|\T_{k+1}-\id\|_{L^2(\nu_{k+1})}^2 = \W_2^2(\mu_{k+1}, \nu_{k+1})\), then we can replace \(\T_{k+1}\) by the OT map \(\T_{\nu_{k+1}}^{\mu_{k+1}}\), which won’t change \(\cFp\big((\T_{\nu_{k+1}}^{\mu_{k+1}})_\#\mu\big)=\cFp(\mu_{k+1})\) but will give a better transport cost.

Let \(J(\T)=\cFp(\T_\#\mu_k) - \langle \gW\cFm(\mu_k), \T-\id\rangle_{L^2(\mu_k)} + \frac{1}{2\tau}\|\T-\id\|_{L^2(\mu_k)}^2\). The gradient of \(J\) gives \[\begin{align} \nabla J(\T) &= \gW\cFp(\T_\#\mu_k)\circ \T - \gW\cFm(\mu_k) + \frac{1}{\tau}(\T-\id) \\ &= \gW\cFp(\T_\#\mu_k) \circ \T + \frac{1}{\tau} \big(\T - (\id + \tau \gW\cFm(\mu_k))\big) \\ &= \gW\cFp(\T_\#\mu_k)\circ \T + \frac{1}{\tau} \big(\T - \sS_{k+1}\big). \end{align}\] Taking the first order conditions, \[\label{eq:foc95tilde} \nabla J(\tilde{\T}_{k+1}) = 0 \iff \tilde{\T}_{k+1} = \mathop{\mathrm{argmin}}_{\tilde\T\in L^2(\mu_k)}\;\cFp(\tilde\T_\#\mu_k) + \frac{1}{2\tau}\|\tilde\T-\sS_{k+1}\|_{L^2(\mu_k)}^2.\tag{15}\]

Now let us show that \(\tilde{\T}_{k+1} = \T_{k+1}\circ\sS_{k+1}\), i.e. \[\label{eq:ineq95min} \min_{\tilde\T\in L^2(\mu_k)}\;\cFp(\tilde\T_\#\mu_k) + \frac{1}{2\tau}\|\tilde\T-\sS_{k+1}\|_{L^2(\mu_k)}^2 = \min_{\T\in L^2(\nu_{k+1})}\;\cFp(\T_\#\nu_{k+1}) + \frac{1}{2\tau}\|\T-\id\|_{L^2(\nu_{k+1})}^2.\tag{16}\]

On one hand, notice by a change of variables, since \(\nu_{k+1}=(\sS_{k+1})_\#\mu_k\), that for all \(\T\in L^2(\nu_{k+1})\), \[\frac{1}{2\tau}\|\T-\id\|_{L^2(\nu_{k+1})} + \cFp(\T_\#\nu_{k+1}) = \frac{1}{2\tau} \|\T\circ \sS_{k+1} - \sS_{k+1}\|_{L^2(\mu_k)} + \cFp\big((\T\circ \sS_{k+1})_\#\mu_k\big).\] Since \(\{\T\circ \sS_{k+1},\;\T\in L^2(\nu_{k+1})\}\subset L^2(\mu_k)\), \[\begin{align} &\min_{\T\in L^2(\nu_{k+1})} \;\frac{1}{2\tau}\|\T-\id\|_{L^2(\nu_{k+1})} + \cFp(\T_\#\nu_{k+1})\\ &= \min_{\T\in L^2(\nu_{k+1})} \; \frac{1}{2\tau} \|\T\circ \sS_{k+1} - \sS_{k+1}\|_{L^2(\mu_k)} + \cFp\big((\T\circ \sS_{k+1})_\#\mu_k\big) \\ &\ge \min_{\T\in L^2(\mu_k)}\;\frac{1}{2\tau}\|\T-\sS_{k+1}\|_{L^2(\mu_k)}^2 + \cFp(\T_\#\mu_k) \\ &\stackrel{\eqref{eq:foc95tilde}}{=} \frac{1}{2\tau}\|\tilde{\T}_{k+1} - \sS_{k+1}\|_{L^2(\mu_k)}^2 + \cFp\big((\tilde{\T}_{k+1})_\#\mu_k\big). \end{align}\]

Now, let us suppose by contradiction that it is a strict inequality, i.e. \[\min_{\T\in L^2(\nu_{k+1})} \;\frac{1}{2\tau}\|\T-\id\|_{L^2(\nu_{k+1})} + \cFp(\T_\#\nu_{k+1}) > \frac{1}{2\tau}\|\tilde{\T}_{k+1} - \sS_{k+1}\|_{L^2(\mu_k)}^2 + \cFp\big((\tilde{\T}_{k+1})_\#\mu_k\big).\] Define \(\eta = (\tilde{\T}_{k+1})_\#\mu_k\). Since \(\nu_{k+1}\in\cPa\), there exists an OT map \(\T^\eta\in L^2(\nu_{k+1})\) such that \(\W_2^2(\eta, \nu_{k+1}) = \|\T^\eta - \id\|_{L^2(\nu_{k+1})}^2\). Notice that \((\tilde{\T}_{k+1}, \sS_{k+1})_\#\mu_k \in \Pi(\eta, \nu_{k+1})\), hence \(\|\tilde{\T}_{k+1} - \sS_{k+1}\|_{L^2(\mu_k)}^2 \ge \W_2^2(\eta, \nu_{k+1})\) and \[\begin{align} \min_{\T\in L^2(\nu_{k+1})} \;\frac{1}{2\tau}\|\T-\id\|_{L^2(\nu_{k+1})} + \cFp(\T_\#\nu_{k+1}) &> \frac{1}{2\tau}\|\tilde{\T}_{k+1} - \sS_{k+1}\|_{L^2(\mu_k)}^2 + \cFp\big((\tilde{\T}_{k+1})_\#\mu_k\big) \\ &\ge \frac{1}{2\tau}\|\T^\eta-\id\|_{L^2(\nu_{k+1}} + \cFp((\T^\eta)_\#\nu_{k+1}). \end{align}\] However, \(\T^\eta \in L^2(\nu_{k+1})\) and thus this is a contradiction. Therefore, we necessarily have an equality in 16 and \(\tilde{\T}_{k+1} = \T_{k+1}\circ \sS_{k+1}\).

0.0.12.7 Proof of 7

We start with two preliminary lemmas.

We now move to proving 7. First, let us focus on the potential term of MMD, i.e. \[\cV(\mu) = \int \V\;\mathrm{d}\mu, \quad \V(\cdot) = -\int k(\cdot, y)\;\mathrm{d}\nu(y).\] Since for all \(x,y\in \R^d\), \(k(x,y) = \psi(x-y) = \psip(x-y) - \psim(x-y)\), we can rewrite \(\V\) as, for all \(x\in \R^d\), \[\begin{align} \V(x) &= -\int \psi(x-y)\;\mathrm{d}\nu(y) \\ &= - \int \big(\psip(x-y) - \psim(x-y)\big)\;\mathrm{d}\nu(y) \\ &= \int \psim(x-y)\;\mathrm{d}\nu(y) - \int \psip(x-y)\;\mathrm{d}\nu(y) \\ &= \Vm(x) - \Vp(x). \end{align}\] Moreover, as \(\psim\) and \(\psip\) are respectively \(\am\) and \(\ap\)-strongly convex and locally Lipschitz, \(\Vm\) and \(\Vp\) are respective \(\am\) and \(\ap\)-strongly convex by 7 and continuous.

We now show \(\psim\) and \(\psip\) have more than a negative quadratic growth. Since they have full domain, there have convex subdifferentials everywhere, in particular in \(0\). Take for instance \(p\in\partial \psim(0)\), then \(\psim(x) \ge \psim(0)+\langle p, x \rangle\). If \(p=0\), then \(\psim\) is lower bounded. Assume \(p\neq 0\), then \[\psim(x) \ge \psim(0)+\langle p, x \rangle \ge \psim(0)-\|p\|_2 \|x\|_2 \ge \psim(0) -\|p\|_2(1+\|x\|_2^2).\] The same reasoning applies to \(\psip\). Consequently there exists \(a^+, a^-, b^+, b^-\in \R\) such that \(\psip(\cdot) \ge -a^+-b^+\|\cdot\|_2^2\) and \(\psim(\cdot) \ge -a^--b^-\|\cdot\|_2^2\). Consequently, \(\Vm\) and \(\Vp\) have a negative part with a quadratic growth using 8.

Hence, by [29], \(\cV\) can be decomposed as a difference of two strongly totally convex potential energies \(\cV = \cVm - \cVp\) with \(\cVm=\int\Vm\mathrm{d}\mu\) \(\am\)-totally convex, \(\cVp=\int\Vp\mathrm{d}\mu\) \(\ap\)-totally convex.

Similarly, the interaction energy term \(\cW(\mu) = \frac{1}{2} \iint k(x,y)\;\mathrm{d}\mu(x)\mathrm{d}\mu(y)\) can be decomposed as a difference of totally convex interaction energies \(\cW=\cWp-\cWm\) (by [29]) with \[\cWp(\mu) = \frac{1}{2} \iint \psip(x-y)\;\mathrm{d}\mu(x)\mathrm{d}\mu(y),\quad \cWm(\mu) = \frac{1}{2} \iint \psim(x-y)\;\mathrm{d}\mu(x)\mathrm{d}\mu(y).\]

Defining \(\cFp(\mu) = \cWp + \cVm + c\) and \(\cFm(\mu) = \cWm + \cVp\), we have for all \(\mu\in\cPr\), \[\bF(\mu) = \cFp(\mu) - \cFm(\mu),\] where \(\cFp\) is \(\am\)-totally convex and \(\cFm\) is \(\ap\)-totally convex, as the sum of a convex term (the interactions) and of a strongly-convex term (the potentials).

0.0.12.8 Proof of 8

The assumption \(\underline\lambda[\qp],\underline\lambda[\qm]\ge 0\) allows to deduce that \(\psip:z\mapsto \qp(\|z\|_2^2)\) and \(\psim:z\mapsto \qm(\|z\|_2^2)\) are \(\underline\lambda[\qp]\ge0\) and \(\underline\lambda[\qm]\ge 0\) strongly convex. Hence, we can apply 7 and obtain that \(\cFp\) is \(\underline\lambda[\qm]\)-totally convex, and \(\cFm\) is \(\underline\lambda[\qp]\)-totally convex.

Then applying 13, we obtain the second result.

0.0.12.9 Proof of 10

We recall that for all \(k\ge 0\), \[\T_{k+1} = \mathop{\mathrm{argmin}}_{\T\in L^2(\mu_k)}\;\bF(\T_\#\mu_k) + \D_{\cFm}^{\mu_k}(\T,\id).\] Taking the first order conditions, we obtain \[\begin{gather} \label{eq:first-order95cnd} \gW\bF(\mu_{k+1})\circ \T_{k+1} + \gW\cFm(\mu_{k+1})\circ\T_{k+1} - \gW\cFm(\mu_k) = 0\\ \iff \gW\bF(\mu_{k+1})\circ \T_{k+1} = -\big(\gW\cFm(\mu_{k+1})\circ\T_{k+1} - \gW\cFm(\mu_k)\big). \end{gather}\tag{17}\]

Notice that \(\T_{k+1}=\mathop{\mathrm{argmin}}_{\T_\#\mu_k=\mu_{k+1}}\;\D_{\cFm}^{\mu_k}(\T,\id)\), and define \(\T_k^* = \mathop{\mathrm{argmin}}_{\T_\#\mu_k=\mu^*}\;\D_{\cFm}^{\mu_k}(\T,\id)\). Both exist as by assumption, \(\mu_k\in\cPa\) for all \(k\ge 0\).

By hypothesis, \(\bF\) is \(\alpha\)-convex relative to \(\cFm\) along \(t\mapsto \big((1-t)\T_k^*+t\T_{k+1}\big)_\#\mu_k\), thus we have \[\begin{gather} \D_{\bF}^{\mu_k}(\T_k^*, \T_{k+1}) \ge \alpha \D_{\cFm}^{\mu_k}(\T_k^*, \T_{k+1}) \\ \iff \bF(\mu^*) - \bF(\mu_{k+1}) - \langle \gW\bF(\mu_{k+1})\circ \T_{k+1}, \T_k^* - \T_{k+1}\rangle_{L^2(\mu_k)} \ge \alpha \D_{\cFm}^{\mu_k}(\T_k^*, \T_{k+1}). \end{gather}\] By definition of \(\mu^*\), \(\bF(\mu^*)-\bF(\mu_{k+1}) \le 0\). Using the first order conditions, we get the inequality \[\label{ineq:612} \langle \gW\cFm(\mu_{k+1})\circ \T_{k+1}-\gW\cFm(\mu_k), \T_k^* - \T_{k+1}\rangle_{L^2(\mu_k)} \ge \alpha \D_{\cFm}^{\mu_k}(\T_k^*, \T_{k+1}).\tag{18}\] By the 3-point equality (see [20]) applied with \(\T:=\T_{k+1}\), \(\sS:=\T_k^*\) and \(\mathrm{U}:=\id\), \[\begin{gather} \label{eq:app95three95pt95ineq} \langle \gW\cFm(\mu_{k+1})\circ \T_{k+1}-\gW\cFm(\mu_k), \T_k^* - \T_{k+1}\rangle_{L^2(\mu_k)} \\ = \D_{\cFm}^{\mu_k}(\T_k^*,\id) - \D_{\cFm}^{\mu_k}(\T_k^*,\T_{k+1}) - \D_{\cFm}^{\mu_k}(\T_{k+1}, \id). \end{gather}\tag{19}\] Plugging 19 into 18 , we get \[\begin{align} \D_{\cFm}^{\mu_k}(\T_k^*,\id) - \D_{\cFm}^{\mu_k}(\T_{k+1},\id) \ge (\alpha + 1) \D_{\cFm}^{\mu_k}(\T_k^*, \T_{k+1}). \end{align}\] Using that \(\cFm\) is convex along \(t\mapsto \big((1-t)\id + t \T_{k+1})_\#\mu_k\), we get that \(\D_{\cFm}^{\mu_k}(\T_{k+1},\id) \ge 0\). Using the definition of \(\T_k^*\), we have \(\D_{\cFm}^{\mu_k}(\T_k^*,\id)=\W_{\cFm}(\mu^*,\mu_k)\) and since \(\gamma=(\T_k^*, \T_{k+1})_\#\mu_k \in\Pi(\mu^*, \mu_{k+1})\), we also have \(\W_{\cFm}(\mu^*,\mu_{k+1}) \le \D_{\cFm}^{\mu_k}(\T_k^*,\T_{k+1})\). Thus, we obtain by induction \[\W_{\cFm}(\mu^*,\mu_{k+1}) \le \D_{\cFm}^{\mu_k}(\T_k^*, \T_{k+1}) \le \frac{1}{1+\alpha}\W_{\cFm}(\mu^*,\mu_k) \le \left(\frac{1}{1+\alpha}\right)^{k+1} \W_{\cFm}(\mu^*,\mu_0).\]

Moreover, by definition of \(\T_{k+1}\), \[\bF(\mu_{k+1}) + \D_{\cFm}^{\mu_k}(\T_{k+1}, \id) \le \bF(\mu^*) + \D_{\cFm}^{\mu_k}(\T_k^*, \id).\] Hence, \[\bF(\mu_{k+1}) - \bF(\mu^*) \le \W_{\cFm}(\mu^*,\mu_k) \le \left(\frac{1}{1+\alpha}\right)^k \W_{\cFm}(\mu^*,\mu_0).\]

0.0.12.10 Proof of 1

Let \(s=\|z\|_2^2\) and recall that \(\psi(z)=q(\|z\|_2^2)\). Differentiating, we have \(\nabla \psi(z)= 2q'(s)z\), hence differentiating a second time we obtain \[\nabla^2 \psi(z)= 2 q'(s)I_d + 4q''(s) zz^{\top}.\] Fix \(z\) and let \(e\perp z\) and \(\|e\|_2=1\) then \(zz^{\top}e=0\), and \[\nabla^2\psi(z)e= 2q'(s)e,\] then \(e\) is a tangential eigenvector, with eigenvalue \(2q'(s)\). We can find \(d-1\) vectors that are linearly independent and \(\perp z\), hence \(2q'(s)\) is an eigenvalue with multiplicity \(d-1\). Now consider in the radial direction \(v = \frac{z}{\|z\|_2}\) (when \(z\neq 0\)), we have: \[\begin{align} \nabla^2\psi(z) v &= \left(2 q'(s)I_d + 4q''(s) zz^{\top}\right) v\\ &= 2 q'(s)v + 4q''(s) zz^{\top} \frac{z}{\|z\|_2} \\ & = 2 q'(s)v + 4 q''(s) \|z\|_2 z\\ & =2 q'(s)v + 4q''(s) \|z\|_2^2 v\\ & = \big(2q'(s) + 4q''(s)s\big)v. \end{align}\] Hence \(v\) is an eigenvector of \(\nabla^2\psi(z)\) with eigenvalue \(\left(2q'(s) + 4q''(s)s\right)\) of multiplicity 1. Hence we have identified the full spectrum of \(\nabla^2\psi(z)\), and the smallest eigenvalue is \[\min\bigl\{2q'(s),\,2q'(s)+4sq''(s)\bigr\},\] and the largest eigenvalue is \[\max\bigl\{2q'(s),\,2q'(s)+4sq''(s)\bigr\}.\]

0.0.12.11 Proof of 3

Recall that \(k(x,y) = \psi(x-y)\), where \(\psi(x-y) = \exp(-\alpha \|x-y\|_2^2)\). Let \(q(s) = \exp(-\alpha s)\), we have \(q(s)=\qp(s)-\qm(s)\) where \[\qp(s) = \cosh(\alpha s), \quad \qm(s)=\sinh(\alpha s).\] Taking the derivatives, we have \[\qp'(s)= \alpha \sinh(\alpha s), \quad \qp''(s)= \alpha^2 \cosh(\alpha s),\] and \[\qm'(s)= \alpha \cosh(\alpha s), \quad \qm''(s)= \alpha^2 \sinh(\alpha s).\] It is easy to see that \(q'_{\pm}(s)\geq 0\) and \(q''_{\pm}(s)\geq 0\) for \(s\geq 0\). Therefore \(z\mapsto \psi_{\pm}(z)= q_{\pm}(\|z\|_2^2)\) are convex on \(\Omega - \Omega\) by 2, and we have a DC decomposition of the Gaussian kernel.

On \([0,S^*]\) we have: \[\lambda_+ = \inf_{s \in [0,S^*]} \min \big(2\alpha \sinh(\alpha s), 2\alpha \sinh(\alpha s)+ 4\alpha^2 s \cosh(\alpha s) \big)= \inf_{s\in [0,S^*]} 2\alpha \sinh(\alpha s) = 0,\] and \[\begin{align} \Lambda_+ &= \sup_{s\in [0,S^*]} \max\left( 2\alpha \sinh(\alpha s), 2\alpha \sinh(\alpha s)+ 4\alpha^2 s \cosh(\alpha s) \right)\\ &= \sup_{s\in [0,S^*]} 2\alpha \sinh(\alpha s)+ 4\alpha^2 s \cosh(\alpha s) \\ &= 2\alpha \sinh(\alpha S^*) + 4\alpha^2\cosh(\alpha S^*). \end{align}\] Similarly we have: \[\lambda_{-}= 2\alpha, \quad \Lambda_{-} = 2\alpha \cosh(\alpha S^*) + 4\alpha^2\sinh(\alpha S^*).\]

0.0.12.12 Proof of 12

Recall that \(k(x,y)=\psi(x-y)\), where \(\psi(x-y)=e^{-\alpha \|x-y\|_2^2}\). Let \(q(s) = e^{-\alpha s}\), then \(q(s)=\qp(s)-\qm(s)\) where for all \(s\ge 0\), \[\qp(s) = e^{-\alpha s} + \alpha s,\quad \qm(s)=\alpha s.\] Taking the derivatives, we get \[\qp'(s) = \alpha(1- e^{-\alpha s}),\quad \qp''(s) = \alpha^2 e^{-\alpha s},\] and \[\qm'(s) = \alpha,\quad \qm''(s) = 0.\] Hence, for \(s\ge 0\), \(q_{\pm}(s) \ge 0\) and thus \(z\mapsto \psi_{\pm}(\|z\|_2^2)\) are convex by 2.

Moreover, we get the following minimal and maximal eigenvalues, \[\lambda_+ = \inf_{s\in [0,+\infty)}\;\min \big\{ 2\alpha (1-e^{-\alpha s}), 2\alpha (1-e^{-\alpha s}) + 4s\alpha^2 e^{-\alpha s}\big\} = \inf_{s\in [0,+\infty)}\;2\alpha (1-e^{-\alpha s}) = 0,\] and \[\begin{align} \Lambda_+ &= \sup_{s\in [0,+\infty)} \max \big\{ 2\alpha (1-e^{-\alpha s}), 2\alpha (1-e^{-\alpha s}) + 4s\alpha^2 e^{-\alpha s}\big\} \\ &= \sup_{s\in [0,+\infty)}\;2\alpha (1-e^{-\alpha s}) + 4s\alpha^2 e^{-\alpha s} = 2\alpha + 2\alpha(2s\alpha - 1) e^{-\alpha s}. \end{align}\] Let \(f(s)=2\alpha + 2\alpha(2s\alpha - 1) e^{-\alpha s}\), its derivative give \(f'(s)=4\alpha^2 e^{-\alpha s} - 2\alpha^2(2s\alpha -1) e^{-\alpha s} = 2\alpha^2 e^{-\alpha s} (3-2s\alpha) = 0 \iff s= \frac{3}{2\alpha}\). Moreover, \(f''(s)=-2\alpha^3 e^{-\alpha s}(3-2s\alpha) -4\alpha^3 e^{-s\alpha} = -2\alpha^3 e^{-s\alpha} (5-2s\alpha) \le 0 \iff s\le \frac{5}{2\alpha}\). Hence \(s=\frac{3}{2\alpha}\) is the maximizer, and \[\begin{align} \Lambda_+ &= 2\alpha (1+2 e^{-\frac{3}{2}}). \end{align}\] Similarly, \(\lambda_-=\Lambda_-=2\alpha\).

0.0.12.13 Proof of 5

The derivatives of \(\qm\) give, for all \(s\in \R\), \[\qm'(s) = \frac{1}{2 \sqrt{\varepsilon + s}}, \quad \qm''(s)= - \frac{1}{4(\varepsilon+s)^{3/2}}.\] For the minimum eigenvalue we have \[\begin{align} \min\bigl\{2\qm'(s),\,2\qm'(s)+4s\qm''(s)\bigr\} &= 2\qm'(s)+4s\qm''(s) \\ &= \frac{1}{\sqrt{\varepsilon+ s}} - \frac{s}{(s+\varepsilon)^{\frac{3}{2}}} \\ & = \frac{\varepsilon}{(s+ \varepsilon)^{3/2}}. \end{align}\] Therefore for \(s\in [0,S_*]\), we have \[\lambda_- =\frac{\varepsilon}{(S_*+ \varepsilon)^{3/2}}.\] For the maximum eigenvalue we have \[\max \bigl\{2\qm'(s),\,2\qm'(s)+4s\qm''(s)\bigr\} = 2\qm'(s) = \frac{1}{\sqrt{\varepsilon+s}} .\] Therefore for \(s\in [0,S_*]\), we have \[\Lambda_{-} = \frac{1}{\sqrt{\varepsilon}}.\]

0.0.12.14 Proof of 6

\(\psi(z)=\frac{1}{(c^2 + \|z\|_2^2)^{\alpha}},\) this corresponds to \(q(s)=\frac{1}{(c^2 + s)^{\alpha}}, \alpha \geq 1.\) We have for \(s\geq 0\), \[q'(s)= -\alpha \left(c^2 + s \right)^{-\alpha -1}, \quad q''(s)= \alpha(\alpha+1)\left(c^2 +s\right)^{-\alpha -2}.\] Hence we have \(q'(s)<0\) and \(q''(s)>0\), and thus \(q(s)\) is convex on \(s \geq 0\), but \(\psi\) is not convex since \(q'(s)<0\). Set \(A= \max\big(0, -q'(0)\big)=\alpha c^{-2(\alpha+1)},\) hence \[\qm(s)= A s - \int_{0}^s(s-t) \min\big(0,q''(t)\big) \;\mathrm{d}t = \alpha c^{-2(\alpha+1)} s,\] and \[\qp(s) = q(s) + \qm(s) = \frac{1}{(c^2 + s)^{\alpha}} +\alpha c^{-2(\alpha+1)} s.\] Now turning to minimum of the Hessian we have \[\underline\lambda[\qp] \mathrel{\vcenter{:}}= \inf_{s\geq 0}\min\bigl\{2\qp'(s),\,2\qp'(s)+4s\qp''(s)\bigr\}, \quad \underline\lambda[\qm] \mathrel{\vcenter{:}}= \inf_{s\geq 0}\min\bigl\{2\qm'(s),\,2\qm'(s)+4s\qm''(s)\bigr\}.\] By construction \(\qp'(s)\geq 0\) and \(\qp''(s)\geq 0\), and hence \(\qp'\) is non decreasing and \(\underline\lambda[\qp] = \inf_{s\geq 0} 2\qp'(s) = 2\qp'(0)=-A+A=0.\) On the other hand \(\qm'(s)= \alpha c^{-2(\alpha+1)}\) and \(\qm''(s)=0\) and hence \(\underline\lambda[\qm] =2\alpha c^{-2(\alpha+1)}.\) For the maximum of the Hessian we have as well, \[\overline{\Lambda}[\qp] =\sup_{s \geq 0}\max\bigl\{2\qp'(s),\,2\qp'(s)+4s\qp''(s)\bigr\} = \sup_{s\geq 0 }\;2\qp'(s)+4s\qp''(s).\] Let us note for \(s\geq 0\), \[f(s)= 2 \qp'(s) + 4 s \qp''(s), \quad f'(s)= 6 \qp''(s) + 4s \qp'''(s).\] We have \[\begin{align} f'(s) &= 6\alpha(\alpha+1)(c^2+s)^{-\alpha-2} - 4 \alpha(\alpha+1)(\alpha+2) s(c^2+s)^{-\alpha-3} \\ &=\alpha(\alpha+1)(c^2+s)^{-\alpha-2} \left( 6 - 4 (\alpha+2)\frac{s}{c^2+s} \right). \end{align}\] Let \(s^*\) such that \[6- 4 (\alpha+2)\frac{s^*}{c^2+s^*}=0,\] equivalently \[s^* = \frac{6c^2}{4(\alpha+2)-6} \geq 0.\] We have \(f'(s)\geq 0\) for \(s\in [0,s^*]\) and \(f'(s)\leq 0\) for \([s^*,+\infty[\), hence \(f(s^*)\) is the global sup. The result for \(\overline{\Lambda} [\qm]\) is immediate.

0.0.12.15 Proof of 13

Recall that based on the DC decomposition of 7, \(\cFp = \cWp + \cVm + c\) where \[\cWp(\mu) = \frac{1}{2} \iint \psip(x-y)\;\mathrm{d}\mu(x)\mathrm{d}\mu(y),\quad \cVm(\mu)=\int\Vm(x)\mathrm{d}\mu(x),\] and \[\Vm(x)= \int \psim(x-y)\;\mathrm{d}\nu(y).\] As \(\cWp\) is an interaction energy, and \(\cVm\) a potential energy, their Wasserstein gradients at \(\mu\in\cP_2(\Omega)\) read as, for all \(x\in\R^d\), \[\begin{align} &\gW \cWp(\mu)(x) = \nabla(\psip * \mu)(x) \\ &\gW \cVm(\mu)(x) = \nabla \Vm(x) = \nabla (\psim * \nu)(x). \end{align}\] Hence, the Wasserstein gradient of \(\cFp\) at \(\mu\) is, for all \(x\in\R^d\), \[\begin{align} \gW \cFp(\mu)(x) &=\gW \cWp(\mu)(x)+ \gW \cVm(\mu)(x)\\ &= \nabla(\psip * \mu)(x) + \nabla (\psim * \nu)(x). \end{align}\]

Let \(\T\in L^2(\mu)\) such that \(\T_\#\mu=\sigma\). Define \(a, b:\R^d\to\R^d\) as \[a(x) \mathrel{\vcenter{:}}= \nabla(\psip * \mu)(x) - \nabla(\psip * \sigma)\big(\T(x)\big), \quad b(x) \mathrel{\vcenter{:}}= \nabla (\psim * \nu)(x) - \nabla (\psim * \nu)\big(\T(x)\big).\] Then, we have the foloowing relation between the Wasserstein gradient of \(\cFp\) at \(\mu\) and \(\sigma\): \[\gW\cFp(\mu) - \gW\cFp(\sigma)\circ \T = a(x )+ b(x).\]

Bounding the term \(\|a\|_{L^2(\mu)}\). Looking at the term \(a\) we have: \[\begin{align} a(x) &= \int \nabla \psip(x-y) \;\mathrm{d}\mu(y)- \int \nabla \psip\big(\T(x)-z\big)\;\mathrm{d}\sigma(z)\\ &= \int \left(\nabla \psip(x-y) - \nabla \psip\big(\T(x)-\T(y)\big) \right)\;\mathrm{d}\mu(y), \end{align}\] since \(\T_{\#}\mu = \sigma\). \(\nabla \psi_{+}\) is \(\Lambda_+\)-Lipschitz by assumption on \(\Omega -\Omega\) and hence we have: \[\|a(x)\|_2 \leq \Lambda_{+} \int \|(x- y) -( \T(x)-\T(y))\|_2\;\mathrm{d}\mu(y).\] Now, Jensen inequality and rearranging terms we obtain: \[\|a(x)\|_2^2 \leq \Lambda^2_{+} \int \|(x- \T(x)) -( y-\T(y))\|_2^2\;\mathrm{d}\mu(y)\] Integrating on \(x\) w.r.t \(\mu\), \[\label{eq:terma} \int \|a(x)\|_2^2\;\mathrm{d}\mu(x) \leq \Lambda^2_{+} \int \|(x- \T(x)) -( y-\T(y))\|_2^2\;\mathrm{d}\mu(x) \mathrm{d}\mu(y) .\tag{20}\]

Let us now provide an intermediary lemma before bounding \(a\).

Using Lemma 9 in Equation 20 for \(H(x)= x-\T(x)\) we obtain: \[\begin{align} \int \|a(x)\|_2^2\;\mathrm{d}\mu(x) &\leq 2\Lambda^2_{+} \int \|x-\T(x) \|_2^2\;\mathrm{d}\mu(x) - 2 \Lambda^2_{+} \left| \int (x-\T(x))\;\mathrm{d}\mu(x)\right|^2\\ &\leq 2\Lambda_+^2 \|\T-\id\|^2_{L^2(\mu)}. \end{align}\] Therefore we have: \[\label{eq:Boundnorma} \|a\|_{L^2(\mu)} \leq \sqrt{2} \Lambda_{+} \|\T-\id\|_{L^2(\mu)}.\tag{21}\]

Bounding the term \(\|b\|_{L^2(\mu)}\). Note that \(\nabla \Vm\) is \(\Lambda_-\)-Lipschitz since \(\nabla^2 \psim \preceq \Lambda_{-}I_{d}\) and hence for all \(x\in \Omega\), \[\nabla^2 \Vm(x) = \int \nabla^2 \psim(x-y)\;\mathrm{d}\nu(y) \preceq \Lambda_{-} I_{d}.\] Hence \[\|b(x)\|_2 = \| \nabla \Vm(x) -\nabla \Vm(\T(x) )\|_2\leq \Lambda_{-} \|x-\T(x)\|_2.\] Integrating on \(x\) and by Jensen inequality, we have finally \[\label{eq:termb} \|b\|_{L^2(\mu)} \leq \Lambda_- \|\T-\id\|_{L^2(\mu)}.\tag{22}\]

Final Bound. Combining Equations 20 and 22 we obtain: \[\begin{align} \| \gW\cFp(\mu) - \gW\cFp(\sigma)\circ \T \|_{L^2(\mu)} &= \|a(x )+ b(x)\|_{L^2(\mu)} \\ &\leq \|a\|_{L^2(\mu)} + \|b\|_{L^2(\mu)} \\ &\leq \left( \sqrt{2}\Lambda_+ + \Lambda_- \right)\|\T-\id\|_{L^2(\mu)}. \end{align}\]

0.0.12.16 Proof of 1

In Wasserstein CCCP we have \(\gW\cFp(\mu_{k+1})\circ\T_{k+1} = \gW\cFm(\mu_k)\) and hence \[\begin{align} \gW\cF(\mu_k) &= \gW\cFp(\mu_k)-\gW\cFm(\mu_k)\\ &=\gW\cFp(\mu_k) - \gW\cFp(\mu_{k+1})\circ\T_{k+1}. \end{align}\] Hence we have \[\begin{align} \| \gW\cF(\mu_k)\|_{L^2(\mu_{k})}& = \|\gW\cFp(\mu_k) - \gW\cFp(\mu_{k+1})\circ\T_{k+1} \|_{L^2(\mu_{k})} \\ &\leq L \|\T_{k+1}-\id\|_{L^2(\mu_k)}, \label{eq:lipch} \end{align}\tag{23}\] where we used 13 and \(L=\sqrt{2}\Lambda_+ + \Lambda_-\). Since \(\lambda_{\pm}\ge 0\) and \(\lambda_+ + \lambda_- > 0\), we can apply 5 with \(\alpha_{\pm}=\lambda_{\pm}\). Taking minimum on both sides in inequality 23 , and applying 5 we have finally \[\begin{align} \min_{0\le k\le K-1}\;\| \gW\cF(\mu_k)\|^2_{L^2(\mu_{k})} &\leq L^2 \min_{0\le k\le K-1}\;\|\T_{k+1}-\id\|^2_{L^2(\mu_k)} \\ & \le \frac{2 L^2}{\ap+\am} \frac{\big(\bF(\mu_0)-\bF(\mu_K)\big)}{K}. \end{align}\] As \(K\to \infty\) we have a stationary point on \(\bF\).

References↩︎

---

1. Code available at https://github.com/clbonet/Wasserstein_Convex_Concave_Procedure↩︎