Non-geodesically-convex optimization in the Wasserstein space

We thank the reviewer for the positive comments on our work. We hope the reply below can highlight our contribution.

Strengths

"Computational cost...": as the answer to reviewer 7EW7, we will discuss the computational cost of the JKO in the revised version.

Weaknesses

``the setting considered in this paper is very similar to": our problem setting is much more general than that of ([56], Salim, Korba, Luise - NeurIPS2020). Having some concavity looks simple yet the whole objective structure becomes much more expressive. The expressivity can be acknowledged from the following result: any continuous function can be approximated by a sequence of DC functions over a compact, convex domain [Bačák11].

[Bačák11] Bačák, M., & Borwein, J. M. (2011). On difference convexity of locally Lipschitz functions. Optimization, 60(8-9), 961-978.

By that, we mean that the structure is semantically simple but powerful mathematically.

Saying that the convexity in [56] is replaced by the concavity in our work does not entirely reflect our framework, hence the contribution. We should look at the whole landscape of the objective function (potential energy + internal/interaction energy), and it would be more complete to say that our work replaces the geodesic convexity in [56] by geodesic difference-of-convexity.

Apart from the theoretical contribution mentioned by the reviewer, we think there are two more elegant contributions:

(1) The use of the DC function for the potential energy already allows us to tackle sampling from log-DC density -- which is a large class of densities (instead of log-concave "and" log-smooth density as in ([56] Salim, Korba, Luise, 2020)).

(2) The way we separate $\mathcal{E}_G$ into the JKO part and leave $\mathcal{E}_H$ for the push-forward is non-trivial and is already a contribution (in our opinion). It differs from the common mindset of applying forward to the potential and backward to the internal energy. Our approach makes both steps tractable, i.e., the optimal map of the forward is computable and the optimal map of the backward is characterizable. Moreover, if we started from [56] with the old mindset and with a mere curiosity of what happens if we replace convex $H$ with concave $-H$, we would not do so in the first place and rather abandon the idea because there are not so many ``log-convex" distributions out there.

For these reasons, we also would like to retain the name "semi-FB Euler" to address its non-triviality and distinguish it from [56].

In terms of theory, the use of the concave $-H$ is correctly mentioned by the reviewer. However, we think the existence of $S$ is actually essential because it allows us to work with non-differentiable $H$ (as in Theorem 2). Without this concavity, we have to use more sophisticated subdifferential notions for $H$ like Clarke or Frechet [Clarke90] in the non-differentiability realm. These notions are difficult to work with, and they can be ill-posed (the subdifferential set is empty).

[Clarke90] Clarke, F. H. (1990). Optimization and nonsmooth analysis. Society for Industrial and Applied Mathematics.

``Please make the last step of the proof of Thm 1 more explicit": the equation between lines 734 and 735: $\gamma^{-1}(T_{\mu_*}^{\nu_*} - I) \in \partial (\mathcal{E}_G + \mathscr{H})(\mu^*)$.

By the definition of $\nu^*$ in line 678, $\nu^* := (I+\gamma S)$#$\mu^*$, and since $S$ is the gradient field of a convex function $H$, that push-forward is optimal, i.e., $T_{\mu^*}^{\nu^*} = I + \gamma S$. Plug this in the above formula, we get $S \in \partial(\mathcal{E}_G + \mathscr{H})(\mu^*)$. On the other hand, since $H$ is convex, $S \in \partial \mathcal{E}_H(\mu^*)$ [4, Proposition 4.13]. We will add these explanations in the revised version.

About the computational challenge of the JKO, please refer to our answer to reviewer 7EW7.

Questions

``Could your results be extended ...": We think our results are totally applicable for $L$-smooth potentials with small $\gamma$ as suggested. Note that, any $L$-smooth function is a DC function (see our paragraph Why difference-of-convex structure). Therefore, while the suggestion is nice, it does not extend our work any further. The true extension could be that instead of geodesically-concave potential $-\mathcal{E}_H(\mu)$, we can consider a more general geodesic concavity not necessarily given in the form of potential energy.

What are the "interesting avenues for future work": Those include, for example, the extension mentioned in the previous answer. Moreover, we can also investigate non-geodesically-convex problems in different Wasserstein spaces (either with different base space than $\mathbb{R}^d$, or different cost functions instead of squared Euclidean distance). For example, in [Section 5, Lambert2022], in the case of mixture of Gaussian Variational Inference, the base space for the Wasserstein space is $BW(\mathbb{R}^d)$ (Bures–Wasserstein space) instead of $\mathbb{R}^d$, i.e., we are minimizing over $\mathcal{P}_2(BW(\mathbb{R}^d))$ (Wasserstein space over Bures–Wasserstein space). This adds one layer of complexity and, as a consequence, convexity is lost in that space and, quoting from page 9 [Lambert2022], ``we lose many of the theoretical guarantees". We can extend the discussion when we have one page extra for the final version.

[Lambert2022] Lambert, M., Chewi, S., Bach, F., Bonnabel, S., & Rigollet, P. (2022). Variational inference via Wasserstein gradient flows. Advances in Neural Information Processing Systems, 35, 14434-14447.

Thank you for clarifying what your contributions are. However, this clarification makes the significance of this work less convincing than I initially thought.

The title of the submission "Non-geodesically-convex optimization in the Wasserstein space" indicates that 1. your audience is one that is interested in optimization questions; 2. you consider objective functionals which are non-convex, and the go-to condition in the non-convex optimization literature is smoothness. However your submission is actually about Difference-of-Convex optimization in the Wasserstein space, which is quite different, both in its application cases and in the kind of theory involved for the analysis.

I understand the general appeal of the DC assumption, given the approximation result you cited ("any continuous function $f$ can be approximated by a sequence of DC functions $g-h$ over a compact, convex domain"). But a) this approximation result does not say how to choose the splitting $g-h$ in a usable way for optimization [please correct me if incorrect], and besides b) you do not prove such an approximation result for the Wasserstein space.

Another argument in favor of studying DC functions on the Wasserstein space, instead of smooth ones, might be that your analysis holds under quite few assumptions. But, particularly in a venue such as NeurIPS, a theoretical analysis only has value when it applies (or there are reasons to think it will apply in the future) to (i) an algorithm and (ii) a setting of practical interest. Now

• (i) after reading your answer to Reviewer 7EW7, it is still unclear to me that the algorithm you analyze (using a proximal step in Wasserstein space) has hopes of becoming efficiently implementable. If, as you say, methods proposed 2 years ago allow for faster computation, then it is up to you to implement them.
• (ii) The example setting you promoted most, namely sampling from a log-DC measure, deserves more discussion. Sampling problems of this form typically arise in Bayesian inference, where the target measure is a posterior likelihood which can be quite involved. It is not at all clear how to choose a splitting of it (or rather of its log) in a tractable way.

For these reasons, I downgrade my rating from 4 to 3. Nonetheless I remain open to further discussion. My recommendation to the authors would be to exploit their (technically interesting!) analysis into a full-fledged study of DC optimization in the Wasserstein space. (I realize that this is not easy to incorporate into the current submission unfortunately.)

PS1: I understand that your analysis can be extended to the case where the objective $\mathcal{F}$ is of the form $\tilde{\mathcal{G}} - \tilde{\mathcal{H}}$ where $\tilde{\mathcal{G}}, \tilde{\mathcal{H}}$ are general geodesically convex functionals, instead of $\tilde{\mathcal{G}}$ being linear ($\tilde{\mathcal{G}} = \mathcal{E}_{-H}$ and $\tilde{\mathcal{H}}= \mathcal{H} + \mathcal{E}_G$ in the paper). My comments above implicitly assume this extension.

PS2, regarding the name semi-FB Euler: I do not have a strong opinion on the matter, but I would like to point out that choosing such a name does not contribute at all to convey the non-triviality of the method, and that this name would be quite awkward in the context of the natural extension mentioned in PS1 (with general $\tilde{\mathcal{G}} - \tilde{\mathcal{H}}$).

We thank the reviewer for your positive comments on our work.

Weaknesses

[1.] We will put some material from Part B (discussion on the ICNN approach for the JKO as well as Algorithm 2) into the main text as suggested when we have some extra space.

[2.] We agree with this limitation. Please see also our reply to reviewer 7EW7. As answered to reviewer 7EW7, there are some recent works [A-Melis22, Fan22] proposing new strategies to compute the JKO that scale better.

[A-Melis22] Alvarez-Melis, D., Schiff, Y., & Mroueh, Y. (2022). Optimizing functionals on the space of probabilities with input convex neural networks. Transactions on Machine Learning Research.

[Fan22] Fan, J., Zhang, Q., Taghvaei, A., & Chen, Y. (2022). Variational Wasserstein gradient flow. International Conference on Machine Learning.

[3.] As authors, we also try our best to ensure everything is mathematically correct.

Minor aspect. Thanks for the suggestion. We will replace $\gamma$ with some other symbol and also revise all the sentences carefully.

Question. It is sensible to expect such behaviours of the time-continuous limit of semi-FB-Euler. We think the hypotheses raised by the reviewers are very likely to hold because the semi-FB-Euler is stable and has sufficient regularity. However, we have not really studied that limit and it is not necessarily easy to prove such hypotheses. This could be an interesting future research avenue.

We thank the reviewer for constructive and positive comments on our work.

Weaknesses

the scheme involves to compute the JKO operator, which is computationally costly..., toy examples: Please also refer to our reply to reviewer 7EW7. Since we use [Mokrov21] in our JKO step, our work is equally costly. There are some recent works [A-Melis22, Fan22] that propose new strategies to compute the JKO that scale better.

[Mokrov21] Mokrov, P., Korotin, A., Li, L., Genevay, A., Solomon, J. M., & Burnaev, E. (2021). Large-scale Wasserstein gradient flows. Advances in Neural Information Processing Systems, 34, 15243-15256.

[A-Melis22] Alvarez-Melis, D., Schiff, Y., & Mroueh, Y. (2022). Optimizing functionals on the space of probabilities with input convex neural networks. Transactions on Machine Learning Research.

[Fan22] Fan, J., Zhang, Q., Taghvaei, A., & Chen, Y. (2022). Variational Wasserstein gradient flow. International Conference on Machine Learning.

Question

``Could the analysis be done with more general DC functionals": Yes, our framework can handle interaction energies (e.g., MMD) as discussed in the context paragraph and Appendix A.2. Applying semi-FB Euler to MMD as suggested is a good research direction which we can consider in the near future.

``Assumption 2 supposes that the sublevel sets are compact...": As far as we know, there is no easy check for it when the base space for the Wasserstein space is $\mathbb{R}^n$. If the base space is instead compact ($X \subset \mathbb{R}^n$, $X$ is compact, the Wasserstein space is the space of all finite 2nd-moment probability distributions whose supports are contained in $X$), the compactness of sublevel sets would be easier to verify since the coercivity of $\mathcal{F}$ is a sufficient condition. In turn, to check coercivity, one sufficient condition is that the potential grows fast enough at infinity, e.g., if $\mathscr{H}$ is the entropy, we need $F(x) \gtrsim c \Vert x \Vert^2$. Nevertheless, this is not the case here since we wanted to consider $\mathbb{R}^n$ but not $X$.

Note, however, that the compactness assumption is only used for Theorem 1. In fact, we only need the compactness of the generated sequence {$\mu_n$} in the proof. Therefore, in the final version, we will remove Assumption 2 and instead assume that the sequence {$\mu_n$} is compact inside Theorem 1.

Typos: We will proofread the paper and correct the typos as suggested.

We thank the reviewer for the thoughtful and positive comments on our work.

Weaknesses

$\bullet$ "The Lojasiewicz condition...": We agree that the Lojasiewicz condition is usually used in the case of KL divergence. However, the condition also almost holds for the case of Maximum Mean Discrepancy (MMD) (under mild assumptions). Recall that MMD involves the interaction energy:

$

\mathcal{F}(\mu) = \dfrac{1}{2} D_{\text{MMD}}(\mu^*,\mu)^2:=\dfrac{1}{2} \int \int k(x,y) d\mu(x) d\mu(y) + \int{F(x)} d\mu(x) + C

$

where $D_{\text{MMD}}$ is the MMD distance, $F(x)=-\int{k(x,y)}d\mu^*(y)$, $C$ does not depend on $\mu$ and $k$ is a stationary kernel, i.e., $k(x,y) = W(x-y)$ for some $W$. The Wasserstein gradient of $\mathcal{F}$ is given by

$

\nabla_W \mathcal{F}(\mu) = \nabla F + \nabla W * \mu

$

where * denotes the convolution. Under some regularity conditions of the kernel, it is known that (inequality (62) in [Arbel19] with some notational change)

$

2(\mathcal{F}(\mu)-\mathcal{F}(\mu^*)) \leq \Vert \mu^*-\mu\Vert_{\dot{H}^{-1}(\mu)} \times \int{\Vert \nabla_W \mathcal{F}(\mu) \Vert^2} d\mu(x)

$

where $\Vert \mu^*-\mu\Vert_{\dot{H}^{-1}(\mu)}$ is the weighted negative Sobolev distance. This is almost the Lojasiewicz condition for $\mathcal{F}$. The only requirement to make it real Lojasiewicz is that $\Vert \mu-\mu^*\Vert_{\dot{H}^{-1}(\mu)}$ is bounded uniformly for all $\mu$. However, this is tricky. Nevertheless, the point is that we only use the Lojasiewicz condition for the generated sequence {$\mu_k$}, so all we need is the above quantity to be bounded along that sequence, and this boundedness seems natural and is usually assumed (as in [Prop.7, Arbel19]). We will add this discussion in the revised version.

[Arbel19] Arbel, M., Korba, A., Salim, A., & Gretton, A. (2019). Maximum mean discrepancy gradient flow. Advances in Neural Information Processing Systems, 32.

$\bullet$ "JKO step is very challenging": The bottleneck of the JKO computation using [Mokrov21] is the logdet(Hessian), (scales cubically w.r.t. data dimension) and training an NN at each step. If we trade some accuracy for scalability, we can use stochastic log determinant estimators as in [A-Melis22], leading to quadratic complexity. We are also aware of a recent work that replaces the gradient of ICNN by a residual neural network, leading to remarkable speeding [Fan22]. Therefore, we believe the scalability issue of the JKO and hence our proposed method can be resolved shortly.

[Mokrov21] Mokrov, P., Korotin, A., Li, L., Genevay, A., Solomon, J. M., & Burnaev, E. (2021). Large-scale Wasserstein gradient flows. Advances in Neural Information Processing Systems, 34, 15243-15256.

[A-Melis22] Alvarez-Melis, D., Schiff, Y., & Mroueh, Y. (2022). Optimizing functionals on the space of probabilities with input convex neural networks. Transactions on Machine Learning Research.

[Fan22] Fan, J., Zhang, Q., Taghvaei, A., & Chen, Y. (2022). Variational Wasserstein gradient flow. ICML.

--``only motivated by one application, to maximum mean discrepancy functionals": our work is not only motivated by maximum mean discrepancy but also (and mainly) by sampling from log-DC densities, which existing samplers fail to address. This is a very rich class of (nonconvex) densities and it would be helpful to have another independent survey on DC structures of commonly-used densities.

--"connection to two-layer neural networks": The connection between MMD and the problem of infinite-width one hidden layer neural network is briefly mentioned in [Salim20, Appendix B.2]. We can add a pointer to that and can even briefly present it for completeness.

[Salim20] Salim, A., Korba, A., & Luise, G. (2020). The Wasserstein proximal gradient algorithm. Advances in Neural Information Processing Systems, 33, 12356-12366.

$\bullet$``The proofs don't yield explicit constants": we can discuss the complexity w.r.t. dimensionality (to some extent). However, yielding explicit dimension-constant is very tricky, especially since the ICNN family is only a subfamily of convex functions, and we do not know how close the solution given by ICNN is to the true optimal transport map. So, obtaining the overall complexity that includes the JKO computation is probably impossible in practice.

Question

$\bullet$ ``The thorough introduction...": Thanks for the suggestion. However, we think it would be beneficial for the readers to have Wasserstein's geometry background in the main text, and we also need all those notions to present our results precisely, so moving them to the appendix is a bit difficult. In the final version with one page extra, we will extend the experiment section, e.g., move the practical scheme in the appendix to the main text.

$\bullet$``Please give more space...": as answered above, we can add a pointer to [Salim20, Appendix B.2] or can briefly discuss it for completeness.

$\bullet$ ``Since it is a general fact...": It is known that Log-Sobolev implies Talagrand [Otto00]; consequently, KL divergence controls Wasserstein distance in such a case. We can discuss this connection in our paper. However, Theorem 5 is much more general than that and cannot use the above implication. First, we work with a general regularizer $\mathscr{H}$, so the objective is not KL divergence. Second, even when $\mathscr{H}$ is the negative entropy, resulting in KL divergence, we consider the general Lojasiewicz exponent $\theta \in [0,1)$. At the same time, the above implication applies only for $\theta=1/2$ (the case of Log-Sobolev inequality). Therefore, Theorem 5 is not simply a corollary of Theorem 4.

[Otto00] 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