Sparse Optimization on Measures with Over-parameterized Gradient Descent

Introduction

where $\mathcal{M}_{+}(\Theta)$ is the set of nonnegative measures $\nu$ on the parameter space $\Theta$ with finite total mass $\nu(\Theta)<\infty$ and $\lambda>0$ is the regularization strength. This formulation also covers minimization over signed measures with total variation regularization, by replacing $\Theta$ with the disjoint union of two copies of $\Theta$ where $\phi$ takes opposite values, see Appendix A. A large body of research has exhibited the favorable properties of minimizers of such problems with a statistical or variational viewpoint, showing in particular that $\lambda$ favors sparser solutions and increases stability as it gets larger, at the expense of introducing a stronger bias. The present paper deals with the optimization aspect: our goal is to design algorithms that return $\epsilon$-accurate solutions with a guaranteed computational complexity. When the set $\Theta$ is a finite set, this is a finite dimensional convex optimization problem that is well understood . However, convex approaches are generally inefficient when $\Theta$ is a continuous space, such as a $d$-dimensional manifold, where the need to discretize the space leads to a complexity scaling as $\epsilon^{-d}$ in the accuracy $\epsilon$. We consider the following setting:

The algorithm that we analyze in this paper is simple to describe: initialize with a discrete measure and run gradient descent on the positions and weights of the particles. We will see that when the problem (1) admits sparse solutions and is non-degenerate, this over-parameterized non-convex gradient descent has a complexity scaling as $\log(1/\epsilon)$ in the accuracy $\epsilon$. We make the following contributions:

In Section 2, we introduce the conic particle gradient descent algorithm to solve optimization problems in the space of measures and discuss several of its interpretations.

In Section 3, we show under under certain non-degeneracy assumptions that there is a sublevel of $J$ starting from which this algorithm converges exponentially fast to minimizers.

In Section 4, we show that for suitable choices of gradient and initialization, this algorithm converges to global minimizers. The proof combines the result of Section 3 with an analysis of a perturbed mirror descent in the space of measures. The number of iterations required to reach an accuracy $\epsilon$ is polynomial in the characteristics of the problem and logarithmic in $\epsilon$. In contrast, the required number of particles depends exponentially on the dimension $d$, which is unavoidable under our assumptions.

We report results of numerical experiments in Section 5, where the various insights brought by our analysis about local and global behaviors are investigated.

As the problem of finding the simplest linear decomposition over a continuous dictionary is a very natural one, problems of the form (1) appear in a large variety of situations, see for an extensive list. In this paper, our numerical illustrations are focused on two applications, chosen for their practical importance and also because they illustrate the variety of behaviors that can be encountered. We also mention a third example to emphasize on the extreme generality — and thus the intrinsic limits — of our analysis. These three cases are illustrated on Figure 1.

In this application, we want to recover a signal that consists of a mixture of spikes/impulses on $\Theta$ given a noisy and filtered observation $f_{0}$ in the space $\mathcal{F}=L^{2}(\Theta)$ of square-integrable real-valued functions on $\Theta$. When one defines $\phi(\theta):x\mapsto\psi(x-\theta)$ the translations of the filter impulse response $\psi$ and $R$ the squared loss, solving (1) allows to reconstruct the mixture of impulses with some guarantees, see e.g. . In this typically low dimensional application, solving (1) to a high accuracy is crucial. Both the signed and nonnegative case have practical motivations (see Appendix A for how to handle the signed case). Figure 1-(a) illustrates the behavior of particle gradient descent for the signed case on the $1$-torus, where the observed signal is shown in orange. Figure 2 illustrates the unsigned case on the $2$-torus.

2 Related work

Problems with the structure (1) have a long history in optimization when $\Theta$ is discrete, and is typically solved with ISTA , mirror descent or variants of those algorithms. When $\Theta$ is continuous, the one dimensional case can sometimes be dealt with specific algorithms . In higher dimensions, the classical algorithms are conditional gradient algorithms (also known as Franck-Wolfe) , moment methods and adaptive sampling/exchange algorithms . Often, these algorithms are complemented with non-convex updates on the particle positions, which considerably improves their behavior. Given an initial condition that is close to the optimum and with the same structure (i.e. without over-parameterization), the local convergence for non-convex gradient descent is studied in .

The dynamics of two-layer neural networks optimization when the number of hidden units grows unbounded is studied in . This series of work has led to various insights related to stochastic fluctuations and global convergence. The present paper can be seen as a quantitative counterpart to , although we consider a more restrictive settingThe algorithm we study in this paper corresponds to the “$2$-homogeneous case” in . Also, allows non-smooth regularizers and does not require non-degeneracy.. A global rate of convergence is obtained in but for a modified dynamic where particles are re-sampled at each iteration. Instead, we focus on the basic case where particles are only sampled once at the beginning of the algorithm. It should be mentioned that our analysis is different from the line of research on lazy over-parameterized models initiated by , which does not apply to the regularized case and to the unsigned case. Finally, in the parametric case where the unknown measure is assumed to belong to a finite dimensional probability model, Wasserstein natural gradient or accelerated versions have been proposed. Our analysis is however of non-parametric nature because the number of parameters is not fixed a priori in the analysis.

Our framework involves the theory of optimization on manifolds and of Wasserstein gradient flows . Some inspiration and interpretations of the algorithm under consideration come from unbalanced optimal transport theory and in particular, from the lifting construction in . Finally, our local analysis includes a functional and a gradient Łojasiewicz inequality of order $2$ in Wasserstein space. Such inequalities were studied in for displacement convex functions, which does not cover our setting.

3 Notation

Particle gradient descent

Assume now that $h$ has at most quadratic growth, and that the metric is defined on the whole of $\Omega$. One can then see the discrete problem (2) as a discretization of a problem on the space $\mathcal{P}_{2}(\Omega)$ of probability measures on $\Omega$ with finite second moment endowed with the Wasserstein-$2$ metric given by

This point of view leads to insights on the properties of $F_{m}$ that are independent of $m$, which is crucial for our theoretical analysis. For a measure $\mu\in\mathcal{P}_{2}(\Omega)$, we define following the homogeneous projection operator $\mathsf{h}:\mathcal{P}_{2}(\Omega)\to\mathcal{M}_{+}(\Theta)$ where $\mathsf{h}\mu$ is characterized by

There are various ways to optimize (2) with first order methods. Instead of directly focusing on a specific method, we first consider the gradient flow of $F_{m}$, as it is known that (stochastic) gradient descent approximates this dynamics. Let us call $x=(r_{i},\theta_{i})_{i=1}^{m}\in\Omega^{m}$ the variable of $F_{m}$. A gradient flow of $F_{m}$ is an absolutely continuous curve $(x(t))_{t\geq 0}$ in $\Omega^{m}$ that satisfies

for $t\geq 0$, with the gradient given in Eq. (5). Note that if $h^{\prime}(r)\alpha(r)^{-1}$ does not tend to as $r\to 0$, then the non-negativity constraint on $r$ should be explicitly enforced, which requires the notion of subgradient flows, see for details in our setting.

It is also possible to directly study the optimization dynamics in the space $\mathcal{P}_{2}(\Omega)$ for the functional $F$ of Eq. (6). For a measure $\nu\in\mathcal{M}_{+}(\Theta)$, consider the vector field on $\Omega$ with expression

We refer to $g_{\mathsf{h}\mu}$ as the Wasserstein gradient of $F$ at $\mu$ (this notation emphasizes that it only depends on $\mu$ through $\mathsf{h}\mu$). Gradient flows of $F_{m}$ are particular cases of Wasserstein gradient flows of $F$. The latter are defined as the absolutely continuous curves $(\mu_{t})_{t\geq 0}$ in $\mathcal{P}_{2}(\Omega)$ that satisfy

2 The conic case

As seen in Eq. (5), the choice of the homogeneity degree and of the metric on $\Omega$ determine a specific way to combine the vertical and the spatial components of the gradient (along the variable $r$ and $\theta$, respectively). From now on, we focus on what we refer to as the conic case, which corresponds to the following assumption:

The mass parameterization is $h(r)=r^{2}$ and the metric on $\Omega^{*}$ is of the form Eq. (4) with $(\alpha(r),\beta(r))=(\alpha,\beta/r^{2})$ for some $\alpha,\beta>0$.

Plugging the metric into Eq. (5) gives the gradient (extended by continuity to $\{0\}\times\Theta$)

and the Wasserstein gradient is represented by the vector field

Existence of Wasserstein gradient flows under (A1-2), for any initialization in $\mathcal{P}_{2}(\Omega)$ can be proved along the same lines as in , see details in Appendix C.1. Abstracting away its geometric derivation, the important aspects about our choice of gradient (8) are that its leads updates in $r$ which are multiplicative and updates in $\theta$ which are independent of $r$. These two properties are crucial for our local convergence analysis (Section 3). Moreover, multiplicative updates enjoy favorable convergence rates (Section 4). The resulting structure and dynamics admits several interpretations.

First, the projection $\nu_{t}=\mathsf{h}\mu_{t}$ of the gradient flow solves an advection-reaction equation. Importantly, this dynamics depends on $\mu_{t}$ only via the initialization $\mathsf{h}\mu_{0}$, which is a property specific to the conic setting.

Under (A1-2), let $(\mu_{t})_{t\geq 0}$ be a Wasserstein gradient flow for $F$, with $\mu_{0}\in\mathcal{P}_{2}(\Omega)$. Then $\nu_{t}=\mathsf{h}\mu_{t}$ satisfies (in the weak sense)

which is the definition of weak solutions for (9). ∎

When $\beta=0$, we recover the gradient flow of $J$ for the Fisher-Rao (or Hellinger) metric, which also corresponds to continuous time mirror descent on $\mathcal{M}_{+}(\Theta)$ for the entropy mirror map . When $\alpha=0$, this is the gradient flow of $J$ for the Wasserstein metric . When $\alpha,\beta>0$, this is the gradient flow of the functional $J$ for the Wasserstein-Fisher-Rao metric, a.k.a. Hellinger-Kantorovich metric, see e.g. . Under Assumption (A2), the dynamics (7) and (9) are directly related by Proposition 2.1. In the rest of this paper, we present the statements in terms of the projected dynamics $\nu_{t}$, although they also could be stated in terms of $\mu_{t}$. Note that an alternative discretization of the dynamic (9) was proposed in using particle birth-death.

Let us recall the global convergence result of [17, Thm. 3.3], in our setting and notations. We give in Appendix C.1 a simplified proof, enabled by our stronger smoothness assumptions.

Under (A1-2), assume that $R$ is convex, that $\phi$ is $d$-times continuously differentiable, that $\nu_{0}\in\mathcal{M}_{+}(\Theta)$ has full support and that the projected gradient flow $(\nu_{t})_{t\geq 0}$ converges weakly to some $\nu_{\infty}\in\mathcal{M}_{+}(\Theta)$. Then $\nu_{\infty}$ is a global minimizer of $J$.

This theorem can be understood as a consistency result for conic particle gradient descent. It also raises several questions: under which conditions does $\nu_{\infty}$ exist? Can we guarantee a convergence rate ? Can we relax the full support condition on the initialization? In this paper, we answer positively to these questions in the particular case of non-degenerate sparse problems.

3 Conic particle gradient descent algorithm

Just like the continuous-time gradient flow, the discrete time gradient descent has a corresponding projected dynamics in $\mathcal{M}_{+}(\Theta)$. Here the equivalence also relies on the properties of compatible retractions.

The following lemma shows that, for sufficiently small step-sizes, the iterates (11) are well-defined and monotonously decrease the objective. As usual in optimization, this property is useful to convert results on gradient flows into results on gradient descent.

In particular, using this expression with $\psi_{f}(\theta)=\langle\phi(\theta),f\rangle$ where $\|f\|\leq 1$ (which have uniformly bounded norms $\|\psi_{f}\|_{\mathcal{C}^{2}}$ under our assumptions), we get that

By a first order expansion of $R$, we have for $f,f^{\prime}\in\mathcal{F}$, $R(f^{\prime})-R(f)=\langle f^{\prime}-f,\nabla R(f)\rangle+O(\|f^{\prime}-f\|^{2})$. Thus, using the expression of $J^{\prime}_{\nu}$ from Eq. (3), it follows

So there exists $\eta_{\max}$ such that if $\max\{\alpha,\beta\}\leq\eta_{\max}$, we have $J(\nu_{k+1})-J(\nu_{k})\leq-\frac{1}{2}\|g_{\nu_{k}}\|^{2}_{L^{2}(\mu_{k})}$. Finally, since we have assumed that $\lambda>0$ and $\nabla R$ is bounded on sublevel sets, the quantities $\sup_{J(\nu)\leq J(\nu_{k})}\nu(\Theta)$ and $\sup_{J(\nu)\leq J(\nu_{k})}\|J^{\prime}_{\nu}\|_{\mathcal{C}^{2}}$ are finite. By the decrease property we just proved, these quantities decrease after one iteration if $\max\{\alpha,\beta\}\leq\eta_{\max}$. So $\eta_{\max}$, which depends on these quantities, can be chosen independently of $k\geq 0$. ∎

Exponential local convergence

We now proceed to the theoretical analysis of the projected gradient flow (9) and projected gradient descent (12) in the conic setting. In light of Propositions 2.1 and 2.4, these dynamics correspond to the gradient flow and gradient descent of $F$, seen through the projection operator $\mathsf{h}$.

In order to derive global optimality conditions, we assume the following.

Commonly used losses that satisfy the smoothness and convexity conditions are the square loss and the logistic loss. Under this assumption, we have existence of minimizers and a global optimality condition.

Under (A1) and (A3), problem (1) admits minimizers. Moreover, a measure $\nu^{\star}\in\mathcal{M}_{+}(\Theta)$ is a minimizer if and only if it holds $J^{\prime}_{\nu^{\star}}(\theta)\geq 0$ for all $\theta\in\Theta$ and $J^{\prime}_{\nu^{\star}}(\theta)=0$ whenever $\theta$ in the support of $\nu^{\star}$.

Our local analysis requires sparsity of the minimizers of the objective $J$, which can be guaranteed a priori in several settings (e.g. ).

Without loss of generality, we assume $r_{i}>0$ for all $i$ and $\theta_{i}\neq\theta_{i^{\prime}}$ whenever $i\neq i^{\prime}$, so that $(r_{i},\theta_{i})_{i=1}^{m^{\star}}$ is uniquely well-defined, up to re-ordering. Let us fix from now on normal coordinates frames on the neighborhood of each $\theta_{i}$. This allows to identify tensors at $\theta_{i}$ with their expression in coordinates and also induces a set of coordinates on the direct sum of the tangent spaces $T_{\theta_{i}}\Theta$, which is of dimension $m^{\star}\times d$.

where $\bar{\nabla}\phi\coloneqq(2\alpha\phi,\beta\nabla\phi)$ can be interpreted as the gradient of $\mathsf{h}\phi$ at $(1,\theta)$. Remark that $K$ is defined via the quadratic form associated to the Hessian of $R$ at $f^{\star}$. This interaction kernel $K$ appears naturally in the various statistical and optimization analysis of the minimization problem under consideration . We also use the notation for the local kernels for $i\in 1,\dots,m^{\star}$

where here and in the proofs, we use to label the $r$’s coordinate. The local analysis will be carried under the following non-degeneracy assumptions.

The minimizer $\nu^{\star}$ is non-degenerate in the sense that $\nabla^{2}R(f^{\star})$ is positive definite and, calling $\sigma_{\min}(A)$ the smallest singular value of a linear operator $A$, we have global curvature $\sigma_{\min}(K)>0$, local curvature $\sigma_{\min}(H)=\min_{i}\sigma_{\min}(H_{i})>0$, and strict slackness, i.e. the only points where $J^{\prime}_{\nu^{\star}}$ vanishes are $(\theta_{i})_{i=1}^{m^{\star}}$.

The first property is always satisfied if $R$ is strictly convex. The second property is satisfied when the kernel associated to the feature function $\bar{\nabla}\phi$ is positive definite. The last two assumptions unfortunately depend on an a priori unknown object $J^{\prime}_{\nu^{\star}}$, but are often required to perform analysis of Problem (1) . Yet, in some cases, they can be guaranteed to hold, see e.g. . In spite of this drawback, the local analysis leads to interesting qualitative insights on the dynamics in practice, see Section 5.

Θ\mathcal{M}_{+}(\Theta) A first consequence of these assumptions is that convergence in value implies convergence to minimizers. The distance on $\mathcal{M}_{+}(\Theta)$ that naturally appears in the analysis is the Wasserstein-Fisher-Rao, a.k.a. Hellinger-Kantorovich metric $\widehat{W}_{2}$, which is the extension of the Wasserstein $W_{2}$ metric to unnormalized measures. It admits many equivalent definitions , the most suitable to our context being [42, Thm. 7.20]

where the Wasserstein distance on $\Omega$ is defined relative to the cone metric (in this paragraph, with $\alpha=\beta=1$). The proof of the following result involves the construction of a transport map in the lifted space $\mathcal{P}_{2}(\Omega)$ and is postponed to Appendix D.4.

3 Sharpness of the objective

Our first main result is a lower bound on the squared norm of the gradient in terms of the sub-optimality gap, an inequality known as sharpness, or Polyak-Łojasiewicz inequality , which is a special case of Łojasiewicz gradient inequality. It involves the $L^{2}(\nu)$ norm of the gradient, which we denote for $\nu=\mathsf{h}\mu$ by

Under (A1-5), there exists $J_{0}>J^{\star}$ and $\kappa_{0}>0$, such that for all $\nu\in\mathcal{M}_{+}(\Theta)$ satisfying $J(\nu)\leq J_{0}$ and $\alpha,\beta>0$, one has

While the objective is non-convex in the Wasserstein geometry and has typically an infinity of bad stationary points, this inequality guarantees exponential convergence to global minimizers of various gradient-based dynamics as long as their initialization $\nu_{0}$ has a small enough objective value. Crucially, the specific structure of $\nu$ does not matter, beyond the fact that is is close enough to optimality: it applies indifferently to discrete and absolutely continuous measures. Once Theorem 3.3 is established, it is straightforward to prove exponential convergence of gradient flow and gradient descent.

Under (A1-5), let $J_{0}$ and $\kappa_{0}$ be given by Theorem 3.3. Consider $(\nu_{t})_{t\geq 0}$ a projected gradient flow for $J$ as in Eq. (9). If $J(\nu_{0})\leq J_{0}$ then

By Theorem 3.3 and direct computations, one has

and the result follows by Grönwall’s lemma. ∎

By Lemma 2.5, there exists $\eta_{\max}$ such that if $\max\{\alpha,\beta\}\leq\eta_{\max}$, then $J(\nu_{k+1})-J(\nu_{k})\leq-\frac{1}{2}\|g_{\nu_{k}}\|^{2}_{L^{2}(\nu_{k})}$. Combining this inequality with Theorem 3.3, one has $J(\nu_{k+1})-J(\nu_{k})\leq-\kappa_{0}\min\{\alpha,\beta\}(J(\nu_{k})-J^{\star})$. Rearranging the terms, we get $J(\nu_{k+1})-J^{\star}\leq(1-\kappa_{0}\min\{\alpha,\beta\})(J(\nu_{k})-J^{\star})$ and the result follows by recursion. ∎

4 Proof strategy for the sharpness theorem

The proof of Theorem 3.3, in Appendix D, is based on a local expansion of $J(\nu)$ in terms of some local moments of $\nu$. For a radius $\tau>0$ (that shall be fixed at some small enough value in the course of the proof), we define the sets for $i\in\{1,\dots,m^{\star}\}$,

We assume that $\tau$ is smaller than $1$ and small enough so that these sets together with $\Theta_{0}\coloneqq\Theta\setminus\cup_{i=1}^{m^{\star}}\Theta_{i}$ form a partition of $\Theta$ and that the exponential map at $\theta_{i}$ has injectivity radius larger than $\tau$, for $i\in\{1,\dots,m^{\star}\}$. We then say that $\tau$ is an admissible radius.

If $\nu$ has only $1$ atom in each $\Theta_{i}$ then its spatial coordinate is $\bar{\theta}_{i}$ and $\Sigma_{i}=0$. When moreover $\nu(\Theta_{0})=0$, the optimization reduces to a more classical gradient flow in $\Omega^{m^{\star}}$ which local behavior has already been studied , but obtaining measures of this form is typically almost as hard as solving the original problem. This decomposition can be reminiscent of proof techniques used to study log-Sobolev inequalities (another type of sharpness inequality in Wasserstein space ) in the small temperature regime .

It turns out that the local moments of Definition 3.6 are sufficient to characterize the behavior of $J$ near optimality. In particular, we have the following approximations for $J$ and its gradient around optimality. These formulas are obtained as an intermediate step in the proof of Theorem 3.3 and follow by combining the bounds of Proposition D.4 and Proposition D.5 with Lemma D.3.

Assuming (A1-5), for any $\nu\in\mathcal{M}_{+}(\Theta)$ it holds

5 Discussion on the local behavior

Let us now explore what the expansion from Proposition 3.7 teaches us about the local behavior of the dynamics. In order to simplify the discussion, let us fix a small admissible radius $\tau_{0}$ and ignore the error terms in Proposition 3.7.

When there is no over-parameterization ($m=m^{\star}$) and we have a single particle in the neighborhood $\Theta_{i}$ of each optimal particle, then there is no local variance: $\Sigma_{i}=0$ for $i=1,\dots,m^{\star}$. In this case, we recover the Taylor expansion of $F_{m^{\star}}$ around its minimizer

and the local convergence rate is dictated by the conditioning of $(K+H)$. Now, for an arbitrary over-parameterization i.e. $\nu\in\mathcal{M}_{+}(\Theta)$ but with the support of the solution approximately identified, i.e. $\nu(\Theta_{0})=0$, the objective is still entirely characterized locally by the local moments of $\nu$, since

This expression gives a clear picture of the energy landscape, so let us comment on it. If we think of the particles in $\Theta_{i}$ as a cluster, then the first term consists in a global interaction between the clusters, which only depends on the biases of each cluster relatively to their respective ground truth particles. The two other terms are local interactions within each cluster, which are due to the local curvature of $J^{\prime}_{\nu^{\star}}$ at each $\theta_{i}$. Note in particular that the only term in this expansion that penalizes the variance of each cluster $\Sigma_{i}$ consists of local interactions.

In this paper, the assumption that $\lambda$ is non-zero is not crucial as such. Instead the crucial assumption for the local analysis is (A5). Still, this assumption is intimately connected to the regularization: in the signed case (detailed in Appendix A), it is necessary that $\lambda>0$ to have (A5), because with $\lambda=0$, the minimizer $\nu^{\star}$ is a global minimizer in the space of signed measures and thus the global optimality condition $J^{\prime}_{\nu^{\star}}=0$ holds. In fact, a finer analysis of the behavior as $\lambda\to 0$ is possible in the signed case: it can be shown that one has (emphasizing the dependency in $\lambda$ in the notation):

for some $K_{0}$, $H_{0}$ and $J^{\prime}_{0}$ [28, Prop. 1 and Thm. 2] (where the result is proved for $R$ being the square loss but can be directly generalized to $R$ smooth and strongly convex around the minimizer). Under the assumption that $J^{\prime}_{0}$ is non-degenerate in the sense of (A5), as soon as $\nu(\Theta_{0})>0$ or $\Sigma_{i}\neq 0$ for some $i\in\{1,\dots,m^{\star}\}$, the local rate $\kappa_{0}$ is thus of order $\lambda$ and for $\lambda=0$, the exponential convergence rate is lost. This shows that regularization is necessary for fast local convergence in the signed case, and in particular – remembering the previous paragraph – for the variance of each cluster of particles to vanish quickly. Note that it is an open question to even show local convergence when (A5) does not hold.

It can be seen from the proof of Theorem 3.3 that $(J_{0}-J^{\star})^{-1}$ and $\kappa_{0}^{-1}$ depend polynomially on the characteristics of the problem, which are the regularization $\lambda$, the regularity parameters of $\phi$ and $R$, the ratio $\max_{i}r_{i}/\min_{i}r_{i}$, the inverses of the $\sigma_{\min}(\nabla^{2}R(f^{\star}))$, $\sigma_{\min}(H)$, $\sigma_{\min}(K)$ and finally the quantity $v^{\star}$ that quantifies the strict slackness assumption, in the following sense: $v^{*}>0$ is such that for any local minimum $\theta$ of $J^{\prime}_{\nu^{\star}}$, either $\theta=\theta_{i}$ for some $i\in\{1,\dots,m^{\star}\}$ or $J^{\prime}_{\nu^{\star}}(\theta)\geq v^{*}$.

Quantitative global convergence

There are several convex optimization-based algorithms that are known to return approximate minimizers of $J$ which are mixtures of atoms (with typically $m>m^{\star}$) with a guaranteed complexity, see Section 1.2. Starting from any such approximate minimizer, the results of the previous section imply that conic particle gradient descent converges exponentially fast to minimizers of $J$. However, such a “two-algorithms” approach comes with a drawback: one has to decide when to switch from one algorithm to another. In this section, we show that it is possible to reach global optimality by only performing non-convex gradient descent. This is true under two main conditions: (i) the initialization samples $\Theta$ densely enough, and (ii) the ratio $\beta/\alpha$ is small, at least in the early stages of the algorithm.

In order to state the condition on the initialization, we first choose a reference measure $\rho\in\mathcal{M}_{+}(\Theta)$ with a smooth positive density, also denoted by $\rho$, which represents our prior knowledge about the solution $\nu^{\star}$. We introduce the quantity (analogous to a log-likelihood)

It quantifies how good is $\rho$ as a prior for the unknown minimizer $\nu^{\star}$ and we will see that our convergence bounds are better when $\bar{\mathcal{H}}(\nu^{\star},\rho)$ is smaller. If nothing is known about the optimal positions $\theta_{i}$, we should choose $\rho$ as a uniform density $\alpha\!\operatorname{vol}$ over $\Theta$ for some $\alpha>0$. Minimizing $\bar{\mathcal{H}}(\nu^{\star},\alpha\operatorname{vol})$ in $\alpha$ suggests to choose $\alpha=\frac{\nu^{\star}(\Theta)}{\operatorname{vol}(\Theta)}$.

To obtain an implementable algorithm, we then discretize $\rho$ and consider an initialization $\nu_{0}\in\mathcal{M}_{+}(\Theta)$ which is close to $\rho$ in the $W_{\infty}$ distance (our statements do not require $\nu_{0}$ to be discrete but this is necessary to obtain an implementable algorithm). We now state our main theorem.

then the projected gradient flow $(\nu_{t})_{t\geq 0}$ initialized with $\nu_{0}$ converges to the global minimizer $\nu^{\star}$. Denoting $t_{0}=1/\sqrt{\alpha\beta}$ it satisfies, for $t\geq t_{0}$,

We also state a similar result for gradient descent, but without tracking the constants. The proof follows the same lines as that of Theorem 4.1 and is given in Appendix F.

Under (A1-5), let $J_{0}$ and $\kappa_{0}$ be given by Theorem 3.3 and $\rho=\rho\!\operatorname{vol}\in\mathcal{M}_{+}(\Theta)$ an absolutely continuous reference measure with $\log\rho$ Lipschitz. For any $0<\epsilon\leq 1/2$ and $\nu_{0}\in\mathcal{M}_{+}(\Theta)$, there exists $C,C^{\prime}>0$ that depends on the characteristics of the problem and increasingly on $\bar{\mathcal{H}}(\nu^{\star},\nu_{0})$ and $1/\epsilon$, such that if

The non-asymptotic convergence rate does not appear explicitly in Theorem 4.1, because the result is obtained by trading-off various error terms. In an the idealized setting where $\nu_{0}=\rho$ and $\beta=0$, a direct consequence of Lemma 4.3 and Lemma E.1 is that $J(\nu_{t})-J^{\star}$ decreases as $O(\log(t)/t)$ for the gradient flow and in $O(\log(k)/\sqrt{k})$ for the gradient descent in general. For the specific case of the mirror retraction, we show in Appendix G that a faster rate in $O(\log(k)/k)$ holds.

The fact that the sublevel $J_{0}$ from Theorem 3.3 does not depend on the metric parameters $(\alpha,\beta)$ is crucial to prove these theorems. However, the local exponential rate of convergence in Theorem 4.2 may be deceptively bad if $\beta/\alpha$ is extremely small. An natural fix is to start with a small ratio $\beta/\alpha$ as required by Theorem 4.2, and to increase this ratio at each iteration so as to improve the conditioning of $J$ near optimality. The interest of Theorem 4.2 lies mostly in the qualitative insights it brings. In practice, we would advise to choose $W_{\infty}(\nu_{0},\rho)$, $\alpha$ and $\beta$ via heuristics or parameter search rather than trying to derive the constants of Theorem 4.2, which could be deceptively conservative.

2 Proof of global convergence for gradient flows

This is a continuous and decreasing function of $\tau$ that satisfies

which is if and only if $\operatorname{spt}(\nu^{\star})\subset\operatorname{spt}(\nu_{0})$. When $\beta=0$, this function directly controls the rate of convergence of this mirror descent dynamics hence the name mirror rate function.

A direct consequence of this lemma is that $\lim_{t\to\infty}J(\nu_{t})-J^{\star}$ is guaranteed to be small as $\beta$ gets smaller and as $\operatorname{spt}\nu_{0}$ gets closer to $\operatorname{spt}\nu^{\star}$. In Appendix E we give an upper bound on $\mathcal{Qq}$ for the situation of interest here, leading to explicit convergence rates when combined with Lemma 4.3.

For the last integral term, we use the triangular inequality

where the last term is obtained by bounding the integrated flow of the velocity field $(\nabla J^{\prime}_{\nu_{t}})_{t\geq 0}$. Since $\mathcal{H}(\nu_{\epsilon},\nu_{t})\geq 0$ and $J(\nu_{s})$ is decreasing, it follows

Combining this bound with Lemma 4.3, we get that for $t\geq L/(4\alpha B_{\nu_{0}})$,

In particular, for $t=(\alpha\beta)^{-\frac{1}{2}}$, we get

Since this is valid only when $t\geq L/(4\alpha B_{\nu_{0}})$, we require $(\alpha\beta)^{-\frac{1}{2}}\geq L/(4\alpha B_{\nu_{0}})$ which leads to the first condition on $\beta/\alpha$. Now, we want the right-hand side of (16) to be smaller than $\Delta_{0}\coloneqq J_{0}-J^{\star}$ so that we can conclude with Corollary 3.4. To this end, we require, on the one hand $W_{\infty}(\nu_{0},\rho\operatorname{vol})\leq\Delta_{0}/(2B_{\nu_{0}}\nu^{\star}(\Theta))$. On the other hand, we use the bound $\log(u)\leq C_{\epsilon}u^{\epsilon}$ for $\epsilon\in{]0,1/2]}$, require $4B_{\nu_{0}}\sqrt{\alpha/\beta}\geq 1$ and obtain the condition

This leads to the second condition on $\beta/\alpha$ is the theorem. ∎

3 Fully non-convex gradient descent

The results in the previous section require to set $\beta/\alpha$ at a small initial value. This might appear undesirable because the asymptotic convergence result of Theorem 2.2 holds irrespective of the choice of $\beta/\alpha$. Also, in practice, this condition does not seem required, at least in the examples that we have considered (see Section 5). While the proof technique from Section 4.2 fails without controlling $\beta/\alpha$, the question of wether it is possible to obtain convergence rates for any ratio $\beta/\alpha$ is a natural one.

For such a result, the key challenge is to obtain a convergence rate for the gradient flow dynamics (9) when initialized with a positive density, without conditions on $(\alpha,\beta)$. While we were not able to prove such a result, in order to point out at the theoretical difficulty, we show in Appendix H with a proof technique inspired by , that a convergence rate in objective value in $O(1/\sqrt{\eta t})$ holds as long as the density $\nu_{t}$ is lower bounded by some $\eta>0$ (at least on a certain subset of $\Theta$).

Under (A1-3), for any $J_{\max}\geq J^{\star}$, there exists $C>0$ such that for any $\eta,t>0$ and $\nu_{0}\in\mathcal{M}_{+}(\Theta)$ satisfying $J(\nu_{0})\leq J_{\max}$, if the projected gradient flow (9) satisfies for $0\leq s\leq t$,

where $S_{t}=\{\theta\in\Theta\;;\;J^{\prime}_{\nu_{s}}(\theta)\leq 0\text{ for some }s\in{[0,t]}\}$, then $J(\nu_{t})-J^{\star}\leq\frac{C}{\sqrt{\alpha\eta t}}.$

Unfortunately, this result is not sufficient to obtain a convergence rate because the lower bound on the density may decrease too fast. When this happens, the gradient flow may stagnate an a priori unbounded time in neighborhoods of saddle points, although it is guaranteed to eventually escape by Lemma C.1. Note that the result above does not requires $\mathcal{F}$ to be finite dimensional nor $\lambda=0$ while this would be needed for a proof based on the positive definiteness of the tangent kernel .

Numerical experiments

All experiments can be reproduced with the Julia code available onlinehttps://github.com/lchizat/2019-sparse-optim-measures. Our goal here is not to demonstrate the superiority of Algorithm 1 over other algorithms, but rather to illustrate the insights obtained by the analysis. We consider the following problems introduced in Section 1.1 :

We observe on Figure 3 the effect of the regularization parameter $\lambda$ and of the over-parameterization parameter $m$ on the local convergence rates (in $\widehat{W}_{2}$ distance – approximated by mapping each particle to its final position/mass – or in optimality gap). In accordance with the expansion of Proposition 3.7, we observe exponential convergence whenever $\lambda>0$, with a rate that improves as $\lambda$ increases. For sparse deconvolution, we observe fast exponential convergence when $m=m_{0}=3$ which is explained by only the first term in the local expansion (13) being non-zero. By adding just a single particle, the second term comes into play and the behavior is qualitatively similar than with $20$ particles. For Figure 3-(c), the initialization is random and $m_{0}=5$. Here the behavior for $m=m_{0}$ follows that of $m>m_{0}$ which suggests that the first term in the local expansion of Eq. (13) dominates.

We observe on Figure 4 the effect on the success/failure of optimization of the two main parameters that appear in Theorem 4.1: the over-parameterization parameter $m$ (used to decrease the $W_{\infty}$ criterion) and the ratio of the vertical/spatial step-sizes $\beta/\alpha$. In both (a) and (b) we have $m_{0}=5$ and $\lambda>0$, and the final loss is averaged over $5$ random experiments. Without surprise, minimizers cannot be reached when $m$ is too small. It is also observed that increasing $m$ increases the chances of success even when $m\geq m_{0}$. In contrast, these experiments do not reveal a clear role for $\beta/\alpha$, beyond a change in the convergence speed (see Section 4.3).

Finally, we compare on Figure 5 the behavior of mirror descent against that of Euclidean descent (here integrated with ISTA algorithm ). This corresponds respectively to $h(r)=r^{2}$ and $h(r)=r$ in Eq. 2 and $\beta=0$. We consider the problem of recovering a single spike ($m_{0}=1$) for 1D and 2D sparse deconvolution, starting from the uniform measure on $\Theta$ densely sampled on a grid ($m=100$). We report the behavior in early stages of optimization, before the effect of the discretization comes into play. We observe that mirror descent outperforms Euclidean descent and enjoys a convergence rate of order $\sim 1/k$ around iteration number $k=100$. This is in accordance with the result of Appendix G, where we show a convergence rate for mirror descent with continuous densities in $O(\log(k)/k)$, independent of the dimension. The difference in behavior is illustrated on Figure 5-(c) where we plot $\nu_{1000}$ (in the setting of panel (a)).

Conclusion

In this paper, we have studied particle gradient descent for sparse convex optimization on measures and obtained complexity guarantees under non-degeneracy assumptions. One central idea underlying our analysis is to directly study the iterates in Wasserstein space. We believe that this approach, at the crossroads between analysis and optimization, may lead to other insights for over-parameterized and non-convex gradient descent.

An avenue for future research is to study the unregularized case. This may require to exploit finer properties of the problem than mere smoothness and could improve our understanding of the implicit bias of over-parameterized gradient descent. Another important question is to find theoretical explanations for the favorable behavior observed in high dimensions for two layer neural networks optimization.

The author thanks Francis Bach for fruitful discussions related to this work and the anonymous referees for their thorough reading and suggestions.

References

Appendix A Dealing with signed measures

The infima of (17) and (1) are the same and:

Appendix B Generic non-convex minimization

In this section, we show that any smooth optimization problem on a manifold is equivalent to solving a problem of the form (1). This corresponds to the case of a scalar-valued $\phi$.

where $0<\lambda<-2\phi^{\star}$. Then $\emptyset\neq\operatorname{spt}\nu^{\star}\subset\arg\min\phi$ so minimizers of $\phi$ can be built from $\nu^{\star}$. Reciprocally, from a minimizer of $\phi$, one can build a minimizer for (18).

Now suppose that $\nu$ is a global minimizer of $J$. Then the optimality condition in Proposition 3.1 implies that

Solving for $f_{\nu}$ is possible if $\lambda\nu(\Theta)<1$ and leads to $f_{\nu}=\sqrt{1-\lambda\nu(\Theta)}-1$. We also deduce from the fact that $f_{\nu}>-1$ that $\arg\min J^{\prime}_{\nu}=\arg\min\phi$, and so $\operatorname{spt}\nu\subset\arg\min\phi$. It remains to find under which condition $\nu(\Theta)>0$. We use the fact that $f_{\nu}=\phi^{\star}\nu(\Theta)$ in Equation (19), and get

which in particular satisfies $\lambda\nu(\Theta)<1$. Thus, as long as $-2\phi^{\star}>\lambda$, we have $\nu(\Theta)>0$. Finally, we verify that global minimizers exist, so that the above reasoning makes sense. If $-2\phi^{\star}-\lambda\leq 0$, then $\nu=0$ satisfies the global optimality conditions. Otherwise, choose $\theta^{\star}$ a minimizer for $\phi^{\star}$ and define $\nu=\nu(\Theta)\delta_{\theta^{\star}}$ with the value above for $\nu(\Theta)$, which also satisfies the global optimality conditions. ∎

Appendix C Wasserstein gradient flow

In this section, we recall and adapt some results and proofs from , for the sake of completeness.

For this result, we assume (A1-2). For a compactly supported initial condition $\mu_{0}\in\mathcal{P}_{2}(\Omega)$, the proof of existence for Wasserstein gradient flows (Eq. (7)) in goes through, as it is simply based on a compactness arguments which can be directly translated to this Riemannian setting (more precisely, we apply here Arzelà-Ascoli compactness criterion for curves in the Wasserstein space on the cone of $\Theta$, which is a complete metric space ). Note that these arguments do not require convexity of $R$, but in order to guarantee global existence in time, we need to assume that $\nabla R$ is bounded in sub-level sets of $F$.

For the existence of solutions for projected dynamics on $\Theta$ for any $\nu_{0}\in\mathcal{M}_{+}(\Theta)$, consider a measure $\mu_{0}\in\mathcal{M}_{+}(\Omega)$ such that $\mathsf{h}\mu_{0}=\nu_{0}$ (see for such a construction) and the corresponding Wasserstein gradient flow $(\mu_{t})_{t\geq 0}$ for $F$. Then $\mathsf{h}\mu_{t}$ is a solution to (9).

We do not attempt to show uniqueness in the present work. Note that it is proved in for the case where $\Theta$ is a sphere, by applying the theory developed in .

C.2 Asymptotic global convergence

In this section, we give a short proof of Theorem 2.2, adapted from . The next lemma is the crux of the global convergence proof. It gives a criterion to espace from the neighborhood of measures which are not minimizers.

where the first inequality can be seen by using the “characteristic” representation of solutions to (9), see . It follows by Grönwall’s lemma that $\nu_{t}(K_{v})\geq\exp(\alpha v^{\star}t)\nu_{0}(K_{v})$ which implies that $t_{1}$ is finite. Finally, if we had not assumed that is in the range of $J^{\prime}_{\nu}$ in the first place, then we could simply take $K=\Theta$ and conclude by similar arguments. ∎

Appendix D Proof of the gradient inequality

In this whole section, we consider without loss of generality $\alpha=\beta=1$ (we explain in Section D.7 how to adapt the results to arbitrary $\alpha,\beta$). For simplicity, we only track the dependencies in $\nu$ and $\tau$. Any quantity that is independent of $\nu$ and $\tau$ is treated as a constant and represented by $C,C^{\prime},C^{\prime\prime}>0$, and the quantity these symbols refer to can change from line to line.

Given a measure $\nu\in\mathcal{M}_{+}(\Theta)$, we consider the local centered moments introduced in Definition 3.6 and in addition, for $i\in\{1,\dots,m^{\star}\}$,

Finally, we will quantify errors with the following quantity

which also controls the $\widehat{W}_{2}$ distance (introduced in Section 3.1) to the minimizer $\nu^{\star}$ of $J$, as shown in the next proposition.

It holds $\widehat{W}_{2}(\nu,\nu^{\star})\leq W_{\tau}(\nu)(1+O(\tau^{2})+O(W_{\tau}(\nu)^{2}))$.

Note that for $W_{\tau}(\nu)$ small enough, it holds $\nu(\Theta_{i})>0$ for $i\in\{1,\dots,m^{\star}\}$. Let $\mu\in\mathcal{P}_{2}(\Omega)$ be such that $\mathsf{h}\mu=\nu$ and consider the transport map $T:\Omega\to\Omega$ defined as

By construction, it holds $\mathsf{h}(T_{\#}\mu)=\nu^{\star}$. Let us estimate the transport cost associated to this map

The geodesic distance associated to the cone metric is

Let us decompose $T(r,\theta)$ as $(rT^{r}(\theta),T^{\theta}(\theta))$ and estimate the two contributions forming $\mathcal{T}$ separately. On the one hand, we have

As a consequence, we have $\mathcal{T}=W_{\tau}(\nu)(1+O(W_{\tau}(\nu)^{2})+O(\tau^{2}))$. Remark that this estimate does not depend on the chosen lifting $\mu$ satisfying $\mathsf{h}\mu=\nu$. We then conclude by using the characterization in [42, Thm. 7.20] for the distance $\widehat{W}_{2}$:

Thus $\widehat{W}_{2}(\nu,\nu^{\star})^{2}\leq W_{2}(\mu,T_{\#}(\mu))^{2}\leq\mathcal{T}$, and the result follows. ∎

D.2 Local expansion lemma

Let $\psi$ be any (vector or real-valued) smooth function on $\Theta$ and $\nu\in\mathcal{M}_{+}(\Theta)$. If $\tau>0$ is an admissible radius, then the following first and second-order expansions hold

where $M_{k,\psi}(\theta_{i},\theta)$ is the remainder in the $k-1$-th order Taylor expansion of $\psi$ around $\theta_{i}$ in local coordinates (and we recall that $\bar{\nabla}\psi:=(2\psi,\nabla\psi)$).

By a Taylor expansion of $\psi$ around $\theta_{i}$ for $i\in\{1,\dots,m^{\star}\}$, it holds

where we have used a bias-variance decomposition for the quadratic term. The result follows by summing the integrals over each $\Theta_{i}$ and using the expression of $b$. ∎

D.3 Bound on the distance to minimizers

Then there exists $C,C^{\prime}>0$ such that for all $\tau\leq\tau_{0}$ and $\nu\in\mathcal{M}_{+}(\Theta)$ such that $J(\nu)\leq J_{\max}$, it holds

To prove the first claim, we thus have to bound $W_{\tau}(\nu)$ using the terms in the right-hand side of (21).

By a Taylor expansion, one has for $\theta\in\Theta_{i}$ for $i\in\{1,\dots,m^{\star}\}$,

and we deduce a first bound by summing the terms for $i\in\{1,\dots,m^{\star}\}$,

In order to lower bound the integral over $\Theta_{0}$, we first derive a lower bound for $J^{\prime}_{\nu^{\star}}$ on $\Theta_{0}$. This is a continuously differentiable and nonnegative function on a closed domain $\Theta_{0}$ so its minimum is attained either at a local minima in the interior of $\Theta_{0}$ or on its boundary. Using the quadratic lower bound from the previous paragraph, it follows that for $\theta\in\Theta_{0}$,

Thus, if we also assume that $\tau\leq 2\sqrt{v^{\star}/\sigma_{\min}(H)}$ then $J^{\prime}_{\nu^{\star}}(\theta)\geq\tau^{2}\sigma_{\min}(H)/4$ for $\theta\in\Theta_{0}$ and it follows that

Using inequality (21) we have shown so far that

Using the first order expansion of Lemma D.2 then squaring gives

Since we have assumed that $K$ is positive definite, it follows

D.4 Proof of the distance inequality (Proposition 3.2)

Moreover, by Lemma D.3, there exists $\tau_{0}>0$ and $C>0$ such that

Combining these two lemmas, it follows that for some $C^{\prime}>0$, we have

D.5 Local estimate of the objective

We now prove a local expansion formula for $J$.

Using the first order expansion of Lemma D.2 for $\phi$, we get $\|f_{\nu}-f^{\star}\|^{2}_{\star}=b^{\intercal}Kb+O(W_{\tau}(\nu)^{3})$. Also, using the second order expansion of Lemma D.2 for $J^{\prime}_{\nu^{\star}}$ and using the fact that $J^{\prime}_{\nu^{\star}}$ and its gradient vanish for all $\theta_{i}$, we get

D.6 Local estimate of the gradient norm

For $\nu\in\mathcal{P}_{2}(\Omega)$, it holds

where the decomposition follows from Lemma D.2. The expression for the norm of the gradient is as follows:

Here we use the notation $\langle\cdot,\cdot\rangle_{\star}$ to denote the quadratic form associated to $\nabla^{2}R(f^{\star})$. Thanks to the optimality conditions $\bar{\nabla}J^{\prime}_{\nu^{\star}}(\theta_{i})=0$ for $i\in\{1,\dots,m\}$, we get

where $N$ collects the higher order terms and is defined as

where $\|\bar{\nabla}_{j}M_{\phi,3}(\theta_{i},\theta)\|=O(\|\theta-\theta_{i}\|^{2})$ if $j>0$ and $O(\|\theta-\theta_{i}\|^{3})$ if $j=0$. Expanding the square gives the following ten terms:

where the entries of $\bar{K}$ and $\bar{H}$ differ from those of $K$ and $H$ by a factor $\bar{r}_{i}/r_{i}$. More precisely,

and similarly for $\bar{H}-H$. Since $|\bar{r}_{i}/r_{i}-1|=O(|b^{r}_{i}|)$ we have $\sigma_{\max}(\bar{K}-K)=O(W_{\tau}(\nu))$. It follows, by expanding the square, that

Using the expansion $J^{\prime}_{\nu}(\theta)=J^{\prime}_{\nu^{\star}}(\theta)+\langle\phi(\theta),M_{\nabla R,1}(f^{\star},f_{\nu})\rangle$, we get

The result follows by collecting all the estimates above. ∎

D.7 Proof of the sharpness inequality (Theorem 3.3)

By Proposition D.4 we have that for $\tau>0$ small enough

where $C=\sigma_{\max}(K+H)+\|J^{\prime}_{\nu^{\star}}\|_{\infty}$.

Similarly, by Proposition D.5, for $\tau$ small enough, it holds

where $C^{\prime}=\frac{1}{8}\sigma_{\min}(H)^{2}\tau^{4}$. Now fix $\tau>0$ satisfying the hypothesis of Lemma D.3 and the two previous inequalities. By Lemma D.3, $W_{\tau}(\nu)=O((J(\nu)-J^{\star})^{\frac{1}{2}})$. We deduce that there exists $J_{0}>J^{\star}$ and $\kappa_{0}>0$ , such that whenever $\nu\in\mathcal{M}_{+}(\Theta)$ satisfies $J(\nu)<J_{0}$, one has

Finally, notice that if different metric factors $(\alpha,\beta)\neq(1,1)$ are introduced, one can always lower bound the new gradient squared norm as

which proves the statement for any $(\alpha,\beta)$. Note however that if one wants to make a more quantitative bound, then there are values $(\alpha_{0},\beta_{0})$ that would lead to a better conditioning and potentially higher values for $J_{0}$. In this case, the factor appearing in the sharpness inequality should rather be $\min\{\alpha/\alpha_{0},\beta/\beta_{0}\}$.

Appendix E Estimation of the mirror rate function

We provide an upper bound for the mirror rate function $\mathcal{Qq}$ in the situation that is of interest to us, with $\nu^{\star}$ sparse. Note that this approach could be generalized to arbitrary $\nu^{\star}$.

Under (A1), there exists $C_{\Theta}>0$ that only depends on the curvature of $\Theta$, such that for all $\nu^{\star},\nu_{0}\in\mathcal{M}_{+}(\Theta)$ where $\nu^{\star}=\sum_{i=1}^{m^{\star}}r_{i}^{2}\delta_{\theta_{i}}$ and $\nu_{0}=\rho\operatorname{vol}$ where $\log\rho$ is $L$-Lipschitz, then

Moreover, for any other $\hat{\nu}_{0}\in\mathcal{M}_{+}(\Theta)$, it holds $\mathcal{Qq}_{\nu^{\star},\hat{\nu}_{0}}(\tau)\leq\mathcal{Qq}_{\nu^{\star},\nu_{0}}(\tau)+\nu^{\star}(\Theta)\cdot W_{\infty}(\nu_{0},\hat{\nu}_{0}).$

In the context of Lemma E.1, we introduce the quantity,

which measures how much $\rho$ is a good prior for the (a priori unknown) minimizer $\nu^{\star}$. With this quantity, the conclusion of Lemma E.1 reads, for $\tau\geq L$,

Let us build $\nu_{\epsilon}$ in such a way that the quantity defining $\mathcal{Qq}_{\nu^{\star},\nu_{0}}(\tau)$ in Eq. (15) is small. For this, consider a radius $\epsilon>0$ and consider the measure $\nu_{\epsilon}$ defined as the normalized volume measure on each geodesic ball of radius $\tau$ around each $\theta_{i}$, with mass $r_{i}^{2}$ on this ball, and vanishing everywhere else. Using the transport map that maps these balls to their centers $\theta_{i}$, we get if $\Theta$ is flat,

The integral term can be estimated as follows,

Recalling that $-\log V^{(d)}(\epsilon)\leq-d\log(\epsilon)+C$ for some $C$ that only depends on the curvature of $\Theta$, we get that the right-hand side of (15) is bounded by

Let us fix $\epsilon>0$ by minimizing $C\nu^{\star}(\Theta)\epsilon-\nu^{\star}(\Theta)d\log(\epsilon)/\tau$, which gives $\epsilon=d/(C\tau)$. The first claim follows by plugging this value for $\epsilon$ in the expression above.

The claim follows by noticing that, by construction, $W_{\infty}(\nu_{\epsilon},\hat{\nu}_{\epsilon})\leq W_{\infty}(\nu_{0},\hat{\nu}_{0})$ and then by taking the infimum in $\nu_{\epsilon}$. ∎

Appendix F Global convergence for gradient descent

In the following, result, we study the non-convex gradient descent updates $\mu_{k+1}=(T_{k})_{\#}\mu_{k}$ and $\nu_{k}=\mathsf{h}\mu_{k}$ where

As in the proof of Lemma 2.5, we define $(T_{k}^{r}(\theta),T_{k}^{\theta}(\theta))\coloneqq T_{k}(1,\theta)$ and we define recursively $\nu^{\epsilon}_{k+1}=(T^{\theta}_{k})_{\#}\nu^{\epsilon}_{k}$ where $\nu^{\epsilon}_{0}$ is such that $\mathcal{H}(\nu^{\epsilon}_{0},\nu_{0})<\infty$. Using the invariance of the relative entropy under diffeomorphisms (indeed, $T^{\theta}_{k}$ is a diffeomorphism of $\Theta$ for $\beta$ small enough), and doing a first order expansion of $T^{r}_{k}=1-2\alpha J^{\prime}_{\nu_{k}}+O(\alpha^{2})$ it holds for $\beta$ small enough

where the term in $O(\alpha)$ originates from a first order approximation of the retraction. Now, taking $\max\{\alpha,\beta\}$ small enough to ensure decrease of $(J(\nu_{k}))_{k}$ (by Lemma 2.5) so that $C$ above can be chosen independently of $k$, it follows

by bounding each term $\|\nu_{k^{\prime}}^{\epsilon}-\nu_{0}^{\epsilon}\|$ by $B\beta k^{\prime}$. ∎

The proof follows closely that of Theorem 4.1 but we do not track the “constants” (this would be more tedious). By Lemma E.1, there exists $C>0$ (that depends on $\bar{\mathcal{H}}$, the curvature of $\Theta$ and $\nu^{\star}(\Theta)$) such that $\mathcal{Qq}_{\nu^{\star},\hat{\nu}_{0}}(\tau)\leq C(\log\tau)/\tau+\nu^{\star}(\Theta)W_{\infty}(\nu_{0},\hat{\nu}_{0})$. Combining this with Lemma F.1, we get that when $\max\{\alpha,\beta\}\leq\eta_{\max}$,

Our goal is to choose $k_{0},\alpha,\beta$ and $W_{\infty}(\nu_{0},\hat{\nu}_{0})$ so that this is quantity smaller than $\Delta_{0}\coloneqq J_{0}-J^{\star}$. With $\alpha=1/\sqrt{k}$ and $\beta=\beta_{0}/k$ we get

Then, using a bound $\log(u)\leq C_{\epsilon}u^{\epsilon}$, we may choose $k\gtrsim\Delta_{0}^{-2-\epsilon}$, $\beta_{0}\leq\frac{1}{3}\Delta_{0}/B^{2}$ and $W_{\infty}(\nu_{0},\hat{\nu}_{0})\leq\frac{1}{3}\Delta_{0}/(B\nu^{\star}(\Theta))$ in order to have $J(\nu_{k})-J^{\star}\leq\Delta_{0}$. This gives $\alpha\lesssim\Delta_{0}^{1+\epsilon/2}$, $\beta\lesssim\Delta_{0}^{3+\epsilon}$ and the regime of exponential convergence kicks off after $k=\Delta_{0}^{-2-\epsilon}$ iterations. ∎

Appendix G Faster rate for mirror descent

In this section, we show that for a specific choice of retraction, the convergence rate of $O(\log(t)/t)$ for the gradient flow is preserved for the gradient descent.

Assume (A1-4) and consider the infinite dimensional mirror descent update

In particular, combining with Lemma E.1, if $\nu_{0}=\rho\!\operatorname{vol}$ has a smooth positive density, then $J(\nu_{k})-J^{\star}=O(\log(k)/k)$.

Consider $\nu_{\epsilon}\in\mathcal{M}_{+}(\Theta)$ such that $\mathcal{H}(\nu_{\epsilon},\nu_{0})<\infty$. It holds

where the first equality is obtained by rearranging terms in the definition of $\mathcal{H}$, and the second one is specific to the mirror retraction. Let us estimate the two terms in the right-hand side. Using convexity inequalities, we get

Here the term in $O(\alpha\|g_{\nu_{k}}\|^{2}_{L^{2}(\nu_{k})})$ comes from the proof of Lemma 2.5 (note that the iterates remain in a sublevel of $J$ for $\alpha$ small enough). As for the relative entropy term, we have, using the convexity inequality $\exp(u)\geq 1+u$,

We use this inequality in place of the strong convexity of the mirror function used in the usual proof of mirror descent (because there is no Pinsker inequality on $\mathcal{M}_{+}(\Theta)$). Coming back to the first equality we have derived, it holds,

Summing over $K$ iterations and dividing by $K$, we get

Since for $\alpha$ small enough $(J(\nu_{k}))_{k\geq 1}$ is decreasing (by Lemma 2.5), the result follows. ∎

Appendix H Convergence rate for lower bounded densities

In this section, we justify the claim made in Section 4.3 about the convergence without condition on $\beta/\alpha$. Let us recall the result that we want to prove.

Under (A1-3), for any $J_{\max}>J^{\star}$, there exists $C>0$ such that for any $\eta,t>0$ and $\nu_{0}\in\mathcal{M}_{+}(\Theta)$ satisfying $J(\nu_{0})\leq J_{\max}$, if the projected gradient flow (9) satisfies for $0\leq s\leq t$,

where $S_{t}=\{\theta\in\Theta\;;\;J^{\prime}_{\nu_{s}}(\theta)\leq 0\text{ for some }s\in{[0,t]}\}$, then $J(\nu_{t})-J^{\star}\leq\frac{C}{\sqrt{\alpha\eta t}}.$