Non-geodesically-convex optimization in the Wasserstein space

Summary

This paper introduces and analyzes an optimization scheme, termed "semi-FB Euler", for a particular minimization problem over the Wasserstein space $P_2(R^d)$: $\min_\mu \mathcal{F}(\mu) = \int (G-H) d\mu + \mathcal{H}(\mu)$, where $G$ and $H$ are convex functions over $R^d$ and $\mathcal{H}$ is a functional that is convex over generalized geodesics. Namely, "semi-FB Euler" is defined by equation (4) of the paper: starting from some $\mu_0$, $$\begin{cases} \nu\_{n+1} = (I + \gamma \nabla H)\_{\sharp} \mu_n, \\\\ \mu\_{n+1} \in \arg\min\_{P\_2(R^d)} \int G d\mu + \mathcal{H}(\mu) + \frac{1}{2\gamma} W\_2^2(\cdot, \nu\_{n+1}). \end{cases}$$ (The second step, called a JKO step, is assumed solvable for the theoretical parts of the paper.) Convergence to critical points of $\mathcal{F}$ is established under mild regularity assumptions, both asymptotically (Theorem 1) and with non-asymptotic rates (Theorems 2, 3). Under a Lojasiewicz inequality assumption, convergence to a minimizer is established with non-asymptotic rates (Theorems 4, 5). Numerical illustrations with synthetic data are provided, in which an input-convex neural network is used to implement the JKO step.

Strengths

This work provides mathematically rigorous grounding for an infinite-dimensional optimization scheme, by working out explicit regularity conditions that ensure well-definedness and asymptotic convergence. The precision and attention to rigor are appreciable. The paper is clearly written and well-structured. It also provides an appropriate level of background information on the Wasserstein geometry. Compared to the closely related work "The Wasserstein proximal gradient algorithm" ([56] Salim, Korba, Luise, 2020), numerical experiments using input-convex neural networks (following [44] Morkov et al., 2021) illustrate that the JKO step can indeed be implemented in general, although computational costs are not discussed.

Weaknesses

The proposed "semi-FB Euler" corresponds to a forward-backward splitting scheme of $\mathcal{F}$ into $\mathcal{F}_1(\mu) = \int (-H) d\mu$ and $\mathcal{F}_2(\mu) = \int G d\mu + \mathcal{H}(\mu)$ in the Wasserstein optimization geometry. (Incidentally, the name "semi-FB Euler" can be misleading, as it is really a forward-backward scheme.) So the setting considered in this paper is very similar to that of "The Wasserstein proximal gradient algorithm" ([56] Salim, Korba, Luise, 2020), which corresponds to the case where $H$ is concave and $G=0$. In fact, it seems that the only way that convexity of $G$ comes into play in the proofs, is to ensure convexity of $\mathcal{F}_2$ over generalized geodesics. (Please do correct me if this is incorrect.) With this in mind, the contributions of this paper are to study what happens when $-H$ is concave in $\mathcal{F}_1(\mu) = \int (-H) d\mu$, as opposed to convex in [56]. It seems that concavity of $-H$ is used for: 1. the existence of a Borel measurable selector $S$ for the subdifferential $\partial H$ (section A.1), but it is not essential to the paper; 2. the descent lemma (Lemma 1); 3. quadratic growth of $\|S\|^2$ under quadratic growth of $H$ (Lemma 4 and proof of Thm 1 line 706); 4. it seems that convexity of $H$ is implicity used in the very last step of the proof of Thm 1 (line 735). (Incidentally, please make the last step of the proof of Thm 1 more explicit.) In summary, compared to the setting of [56], which corresponds to $-H$ being smooth and convex, this work shows that if $-H$ is concave, then the forward-backward algorithm (with a good choice of the selector $S$ when $H$ is non-differentiable) converges for any choice of step-size $\gamma$. I find this theoretical contribution interesting, but I am not sure if it is sufficiently significant. Besides, there is also the weakness, acknowledged by the authors (line 362) that the proposed algorithm relies on a JKO step which is computationally expensive. Hence my overall rating for this submission.

Questions

- Please address the points in "Weaknesses". - The items 1.-4. in "Weaknesses can also be guaranteed if, instead of $-H$ concave, one assumes $H$ is smooth (which is Assumption 5) and $\gamma$ is small enough (see [56, Thm 5] for item 2 and [56, Lemma 2] for item 4). Could your results be extended to the case where $H$ is smooth instead of convex? If yes, then this could be another way to show theoretical guarantees for forward-backward schemes outside of the convex case (in the spirit of your conclusion, line 393). - What are the "interesting avenues for future work" mentioned in the conclusion?

Rating

3

Confidence

4

Soundness

4

Presentation

3

Contribution

1

Limitations

The authors adequately addressed the limitations.

Summary

This work, of a theoretical nature, considers the problem of minimizing a functional $$\mathcal{F}$$ over the space of probability measures of the form $$\mathcal{F}(\mu) = \int (G(x) - H(x)) d\mu(x) + \mathcal{H}(\mu)$$ where $G,H$ are **convex** potentials and $\mathcal{H}$ is (typically) the negative entropy (the work actually extends to functionals $\mathcal{H}$ that are convex along generalized geodesics). This setting substantially differs from standard optimization in the space of probability measure in that $\mathcal{F}$ is not convex along generalized geodesics, the standard assumption to establish convergence of gradient flows in this spaces, following the seminal work of Ambrosio, Gigli and Savaré. Nonetheless, the "difference of convex functions" (DC) structure enables the derivation of a two-step scheme to minimize $\mathcal{F}$. Namely, given a current iterate $\mu_n$, one first computes $\nu_{n+1}$ using a **forward** (explicit) Euler scheme _using only the concave term_ $-H$, then obtain $\mu_{n+1}$ using a **backward** (implicit) iterate on the convex (along generalized geodesics) term $\mu \mapsto \braket{G,\mu} + \mathcal{H}(\mu)$. This approach is referred to as a _semi_ forward-backward (FB) Euler scheme, seen as a variation of the known FB scheme (which would process similarly, but performing an explicit iterate on $G-H$ and an implicit iterate on $\mathcal{H}$). In contrast to the standard FB scheme, this semi-FB scheme that leverages the DC structure enables the derivation of strong theoretical guarantees: existence of accumulation points that are critical points of $\mathcal{F}$, and under additional standard assumptions ($\sim$ PL inequality), convergence rate toward the global minimum of $\mathcal{F}$.

Strengths

The presentation of the work is extremely clear. It is exposed in a pedagogical way, motivations are clear and results are stated formally. The authors obtain strong theoretical guarantees in a scope more general than the restricted _convex along generalized geodesics_ one. The DC framework in the context of Wasserstein-based optimization seems to be a promising playground for further research. The proposed approach is elegant to me, in that it is simple yet seems very powerful, and seems reasonably usable in practice.

Weaknesses

1. While I understand that the paper is mostly theoretical, the "from theory to numeric" side may be enhanced (typically, putting some of the material of Part B, in particular Algorithm 2, to the main body). 2. Still targeting numerical applications, the approach seems somewhat limited in that the numerical experiments, according to the appendix, have been run on a fairly powerful hardware and yet took a couple of hours to run, while they remain at the proof-of-concept level on toy datasets. 3. Though this does not diminishes the intrinsic quality of the paper, the paper relies on technical proofs deferred to the appendix and I could not proofread all of them. I read the proof of Lemma 1 and Theorem 1, which seem correct to me, and quickly check other proofs but cannot guarantee the correctness of the theorems (though they look sounded)---which are nonetheless the main (if not unique) contributions of the work. I understand that this is not the authors fault, but I wanted to stress that point in my review. **Minor aspects:** - It may be better to use $\tau$ instead of $\gamma$ to denote the step-size in JKO schemes, since $\gamma$ is often used to denote a transport plan. - I believe it's better to avoid starting sentences with mathematical symbol. That makes the document a bit easier to read. See for instance line 117 "continuous w.r.t. $\mathcal{L}^d$. $\mu$-a.e. stands..." is a bit confusing.

Questions

1. Genuine question: is there any convergence in the regime $\gamma \to 0$ toward a limit curve $(\mu_t)_t$, and if so does it coincides with the usual Wasserstein gradient flow of $\mathcal{F}$? Said differently, is it true that the (time-continuous) Wasserstein gradient flow of the KL for a log-DC target distribution converges globally? (As a consequence of the fact that the sequence $(\mu_n)_n$ you build do converge globally in that case)

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

NA

Summary

This work focuses on the optimization in the Wasserstein space of a functional which is a sum of an internal energy (convex along a generalized geodesic) and of a potential energy, whose potential is a difference of convex functions. To solve such problem, the authors propose to generalize the semi Forward-Backard Euler scheme to the Wasserstein space, and study it theoretically, showing the convergence in general settings. Finally, they demonstrate the convergence of the algorithm on two toy objectives.

Strengths

- This paper is well written and clear. - To minimize functionals obtained as difference of convex functionals, it introduces the semi-forward backward Euler scheme, and it shows its convergence under different reasonable assumptions. - The theoretical analysis seems rigorous and complete, as it shows the convergence in different settings and under different criteria (e.g. the paper provides a descent lemma, convergence of a subsequence towards a critical point, convergence in gradient, and convergence to the minimum under PL inequality).

Weaknesses

This is mainly a theoretical work. Thus, the following weaknesses on the experiment section and on the method are minor in my opinion. - As underlined in the experiment section, the scheme involves to compute the JKO operator, which is computationally costly as it requires neural networks to solve it - The experiments focus on 2D toy examples

Questions

The focus of the paper is on functionals obtained as a sum of a potential and an internal energy. Could the analysis be done with more general DC functionals (e.g. with interaction energies.). I guess yes since the MMD is presented in the introduction as being such a functional. It would also have been nice to test the semi Forward-Backward scheme on the MMD, as it is not geodesically convex, and does not converge well without tricks such as adding noise (e.g. for Gaussian kernel). Assumption 2 supposes that the sublevel sets are compact w.r.t the Wasserstein topology. Are there easy examples of functionals which satisfy this (e.g. those used in the experiment section)? Typos: - In the abstract, I found the sentence line 3 "When the regularization term is the negative entropy, the optimization problem becomes a sampling problem where it minimizes the Kullback-Leibler divergence between a probability measure (optimization variable) and a target probability measure whose logarithmic probability density is a nonconvex function." weird as it is not specified that the objective function is a potential. - Line 14: "can be considered gradient descent": lack a word? Rating: 7

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have acknowledged the limitations in the main text

Summary

This paper studies minimization algorithms for functionals on the Wasserstein space with the following difference of convex functions (DC) structure: $$ \min_{\mu \in \mathcal{P}_2(X)} \mathcal{F}(\mu) := \int (G(x) - H(x)) d\mu(x) + \mathcal{H}(\mu), $$ where $G$ and $H$ are convex functions and $\mathcal{H}$ is convex along (generalized) geodesics in the Wasserstein space. The main motivation for this problem is that it is the sampling analog to difference of convex functions problems in optimization, giving a problem of intermediate difficulty between convex and true non-convex. The authors also mention an application to maximum mean discrepancy. They study a modified forward Euler sceme which they refer to as "semi FB Euler", where given a current iterate $\mu_n$ they update as $$ \mu_{n + 1/2} := (I + \gamma \nabla H) \mu_n, $$ and then $$ \mu_{n + 1} := JKO_{\gamma(H+ E_G)}(\mu_{n + 1/2}). $$ A similar scheme was considered in [1,2] but with the difference that the function $G$ was included in the first step but not in the second; the motivation for this difference is that it more closely aligns with DC programming. One benefit of this modification is that $\mu_{n + 1/2}$ is pushed forward by an optimal transport map, which makes the problem more tractable. The main focus of the work are various convergence results for the above scheme, including an asymptotic result as well as results on the norm of a gradient analog and the distance between $0$ and the sub-differential. The authors also consider convergence under a gradient domination condition, and conclude with some numerics. [1] Wibisono, Andre. "Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem." Conference on Learning Theory. PMLR, 2018. [2] Salim, Adil, Anna Korba, and Giulia Luise. "The Wasserstein proximal gradient algorithm." Advances in Neural Information Processing Systems 33 (2020): 12356-12366.

Strengths

- This paper is, to my knowledge, the first to consider the difference of convex functions setting. - The above bullet leads the authors to propose a modified algorithm compared to previous literature, and means that their convergence results are generally novel. - The work is rigorous and quite mathematically clear.

Weaknesses

- The Łojasiewicz condition is not verified to hold in any settings other than their the standard case of KL minimization. - The problem is somewhat far from practice, firstly because the JKO step is very challenging to implement in practice, and secondly because the setting the authors study is only motivated by one application, to maximum mean discrepancy functionals. They do mention a connection to two-layer neural networks in the "Context" section on page 2 but this was impossible to understand from the little they wrote. - The proofs don't yield explict constants in their results, which is important because there could be hidden dimension-dependence.

Questions

- The thorough introduction to Wasserstein space is much appreciated, but perhaps it could be largely moved to an appendix so that more space is left in the main text for more discussion of applications and high-level ideas of the proofs. - Please give more space to the connection to two-layer neural networks mentioned in the "Context" section on page 2. - Since it is a general fact (see, e.g. [1]) that a gradient domination condition implies that $d(x, x*)^2 <= C(F(x) - F(x*))$, I feel that Theorem 5 should either be removed, or stated as a corollary in an appendix and just mentioned as a brief remark in the main text. [1] Otto, Felix, and Cédric Villani. "Generalization of an inequality by Talagrand and links with the logarithmic Sobolev inequality." Journal of Functional Analysis 173.2 (2000): 361-400.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have adequately addressed the limitations.

Thanks for your response

With regard to the Lojasiewicz condition, it is interesting that it can almost be satisfied in the MMD case, but I would emphasize that the reason the log-Sobolev inequality is so important is that it holds uniformly. In particular, the gap between a PL inequality at some points versus a PL inequality everywhere can be very large and difficult to bridge. Also, I still don't find the connection to MMD particularly convincing, so this setting is not completely convincing anyways. With regard to the fact that gradient domination implies quadratic growth: sure, the case where $(f(x) - f(x_*))^\theta \leqslant C\|\nabla f(x)\|$ is not explicitly proved in [1], but the proof can be quite easily modified to handle that case. At a minimum, you should add citations to [1] and [2]. Note also that using their results you could likely simplify your proofs and make the constants more explicit. With regard to your remaining points, thanks for your responses. Overall, my opinion of the work has not substantially changed, so I leave my score as-is. [1] Otto, Felix, and Cédric Villani. "Generalization of an inequality by Talagrand and links with the logarithmic Sobolev inequality." Journal of Functional Analysis 173.2 (2000): 361-400. [2] Karimi, Hamed, Julie Nutini, and Mark Schmidt. "Linear convergence of gradient and proximal-gradient methods under the polyak-łojasiewicz condition." Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2016, Riva del Garda, Italy, September 19-23, 2016, Proceedings, Part I 16. Springer International Publishing, 2016.