Mirror and Preconditioned Gradient Descent in Wasserstein Space

Thank you for your appraisal and positive comments on our paper.

On page 8, the authors don't study the mirror descent of KL divergences for non-Gaussian target distributions. This could be a very intriguing example. The authors may think about how to implement this case. This could be compared with the Wasserstein gradient flow of KL divergence.

The difficulty of optimizing the KL divergence with a non-Gaussian target is that we cannot use closed-forms for the mean and covariance anymore (as done in the experiment of Figure 2 for Gaussian, where closed-forms were obtained by projecting the distributions at each step on the space of Gaussian).

A possible solution is to use a particle approximation, which we considered in Appendix I.2, with the goal to sample from a distribution on the simplex. Thus, the Bregman potential was chosen here as a potential energy, with potential a barrier allowing to keep the samples onto the simplex. The density of the distribution at each step was approximated with a Kernel Density Estimator to approximate the Wasserstein gradient of the KL.

The authors also propose the Wasserstein-Bregman proximal gradient method. Can authors illustrate the main difficulty of implementing it in practice or some potential computational schemes? How does it compare with the classical sampling schemes if the objective functional is chosen as the KL divergence?

The main difficulty of implementing the Wasserstein-Bregman proximal gradient method is that it is required to solve at each step a backward step (see equation (128)), which does not have a closed-form in general [1]. Thus, at each step, it is required to solve another optimization problem. For the classical JKO scheme, it has been proposed to solve it using approximations with e.g. entropic approximation [2] or neural networks [3,4], but these methods are computationally costly. Nonetheless, note that there is a closed-form when the target and initial distribution are Gaussian [5,6].

In Figure 2, we considered the Wasserstein-Bregman proximal gradient method with the KL divergence as objective functional, and a Gaussian $\mathcal{N}(0,\Sigma)$ as target distribution. For a Bregman potential of the form $\phi(\mu) = \int V\mathrm{d}\mu$, with $V(x)=x^T\Lambda^{-1} x$, we derived the closed-form in Appendix H.4 leveraging the fact that the distributions are always Gaussian. Using $\Lambda=\Sigma$, we showed on Figure 2 that it converges much faster than the regular Wasserstein proximal method studied in [1,5] (which corresponds to $\Lambda=I_d$).

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

[2] Peyré, G. (2015). Entropic approximation of Wasserstein gradient flows. SIAM Journal on Imaging Sciences, 8(4), 2323-2351.

[3] 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.

[4] Alvarez-Melis, D., Schiff, Y., & Mroueh, Y. (2021). Optimizing functionals on the space of probabilities with input convex neural networks. arXiv preprint arXiv:2106.00774.

[5] Diao, M. Z., Balasubramanian, K., Chewi, S., & Salim, A. (2023). Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein space. In International Conference on Machine Learning (pp. 7960-7991). PMLR.

[6] Wibisono, A. (2018). Sampling as optimization in the space of measures: The Langevin dynamics as a composite optimization problem. In Conference on Learning Theory (pp. 2093-3027). PMLR.

Thank you for your valuable feedback and time. We have addressed your comments point-by-point below. Please don't hesitate to let us know if you have any further questions.

What new tool does the MD formulation bring to the table?

First, the Mirror Descent algorithm on the Wasserstein space allows to perform optimization over probability distributions on constrained spaces (and in particular sampling on constrained spaces for specific objectives). This motivated the development of e.g. Mirror Langevin [1,2,3]. This lens on Mirror Descent on the Wasserstein space was also studied e.g. in [4] to optimize the KL, but without a theoretical analysis of the discrete algorithm. Here, we chose to study the convergence of these algorithms, and one of the main contributions of this work is to define appropriate notion of relative smoothness and convexity. This differs a lot from previous works, and is new on the Wasserstein space to the best of our knowledge.

Moreover, we also proposed to use general Bregman potentials, instead of just considering potential energies as done in most previous works focusing on Bregman divergences and Mirror Descent on the Wasserstein space. We believe that this additional flexibility could allow to find new pairs of functionals/Bregman potentials, with theoretical guarantees for the convergence of the Mirror Descent scheme. For instance, a convex functional is 1-smooth and 1-convex relative to itself, and thus could be used as Bregman potential for its optimization, provided we can compute the scheme. We note that this point of view was also recently taken to optimize the MMD in [5], while the MMD is known to be non convex w.r.t the Wasserstein distance [6], which leads MMD flows to converge to spurious local minima in practice [5,6,7].

The function properties in the Wasserstein space are more complicated and therefore difficult to align, this was briefly discussed by the authors at the start of section 5.

Showing the relative smoothness and convexity of functionals is indeed not straightforward in general. In the first paragraph of Section 5, we described some cases where we can show it leveraging the relative smoothness and convexity on $\mathbb{R}^d$, and suggested in the general case to compare the Hessians.

But the application offered later is largely limited to existing applications KL or Sinkhorn.

We did not focus only on minimizing the KL or Sinkhorn. On Figure 1, we applied the Mirror Descent scheme on an interaction energy with an interaction energy as Bregman potential, in a case where it is well relatively smooth. We focused on the KL for the Gaussian example of Figure 2, but this specific example has not been yet studied in the literature as we use Mirror Descent with negative entropy as Bregman potential. Finally, on Figure 3, we applied the preconditioned scheme on different divergences such as the Sliced-Wasserstein distance, the Sinkhorn divergence and the Energy distance. While Wasserstein gradient flows of these divergences have already been proposed and studied, applying the preconditioned gradient descent to them is new to the best of our knowledge.

Some sections are not well presented.

We apologize for the confusions and will revise the paper to improve its clarity.

The computational complexity and real-life run-time of these type of methods should be better addressed both in theory and experiments, especially on the comparison between the proposed and benchmark methods.

Compared to the usual Wasserstein gradient scheme, the preconditioned Wasserstein gradient scheme only requires to evaluate additionaly the gradient of $h^*$, and thus has the same computational complexity as the Wasserstein gradient descent. Nonetheless, we noted in Figure 3 that it required less iterations to converge for a well-chosen preconditioner.

For Mirror Descent, the complexity depends on the Bregman potential. If we use $\phi_\mu(T)=\int V\circ T\ \mathrm{d}\mu$ with $V$ a potential which we know in closed form, and from which we can compute $\nabla V$ and $\nabla V^*$, then the computational cost is the same as the Wasserstein gradient descent. In the more general case, where we do not know how to invert $\nabla\phi_\mu$, we need to use the Newton method, which is more costly as it requires to invert a Jacobian of size $nd\times nd$ as stated Section I.1.

The runtimes are reported in Appendix I. For instance, as stated line 1490, for the experiment of Figure 1, the Mirror Descent used in dimension $d=2$ for $n=100$ particles and $120$ epochs with $K_2$ and $K_2^\Sigma$ took about 5 minutes while with $K_4$ and $K_4^\Sigma$, it took about 1h. We leave for future works the development of cheaper optimization methods.

some minor grammatical errors

Thank you, we will correct the grammatical errors.

[1] Ahn, K., & Chewi, S. (2021). Efficient constrained sampling via the mirror-Langevin algorithm. Advances in Neural Information Processing Systems, 34, 28405-28418.

[2] Chewi, S., Le Gouic, T., Lu, C., Maunu, T., Rigollet, P., & Stromme, A. (2020). Exponential ergodicity of mirror-Langevin diffusions. Advances in Neural Information Processing Systems, 33, 19573-19585.

[3] Hsieh, Y. P., Kavis, A., Rolland, P., & Cevher, V. (2018). Mirrored langevin dynamics. Advances in Neural Information Processing Systems, 31.

[4] Sharrock, L., Mackey, L., & Nemeth, C. (2023, January). Learning Rate Free Bayesian Inference in Constrained Domains. In NeurIPS.

[5] Gladin, E., Dvurechensky, P., Mielke, A., & Zhu, J. J. (2024). Interaction-Force Transport Gradient Flows. arXiv preprint arXiv:2405.17075.

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

[7] Galashov, A., de Bortoli, V., & Gretton, A. (2024). Deep MMD Gradient Flow without adversarial training. arXiv preprint arXiv:2405.06780.

Thank you for reading the paper and for your feedback. We answer your comments below. Please do not hesitate if you have other questions.

The presentation of results could be enhanced by providing more detailed visualizations.

Thank you for these suggestions, we will take it into account and add snapshots of particles as well as contours.

It appears that significant references relevant to discussions on the convergence of mirror descent may have been overlooked. I recommend incorporating additional reference [1] that provides pertinent insights into these discussions.

Thank you for this reference that we did not know. We will add it to the paper.

From the perspective of gradient flow, assuming we have a velocity field $\dot{X} = v_t$, we can derive that $v_t = -\Vert \nabla\frac{\delta \mathcal{F}}{\mu} \Vert_2^2$. In this context, does assuming $\beta$-smoothness offer additional advantages beyond the conventional benefits?

For Wasserstein gradient flows, the velocity field is $v_t=-\nabla_{W_2}\mathcal{F}(\mu_t)$, and we have $\frac{\mathrm{d}\mathcal{F}(\mu_t)}{\mathrm{d}t} = \langle \nabla_{W_2}\mathcal{F}(\mu_t), v_t\rangle_{L^2(\mu_t)} = - \|\nabla_{W_2}\mathcal{F}(\mu_t)\|_{L^2(\mu_t)}^2$. Thus, the objective of the continuous formulation is necessarily non-increasing without assuming any smoothness of the objective.

In discrete time, the Wasserstein gradient descent (obtained as the explicit Euler discretization) \mu_{k+1} = (\mathrm{Id}-\tau\nabla_{W_2}\mathcal{F}(\mu_k))_\\#\mu_k requires smoothness of the objective to obtain descent at each iteration (and eventually global convergence guarantees if we assume furthermore convexity). See for instance analysis of gradient descent [2]. If the step size is too big, it is not guaranteed that the scheme decreases at each iteration, and could actually diverge. Hence in practice, knowing/estimating the smoothness constant will guide the choice of the step size.

Could the convergence location and rate be established from the perspective of a Lyapunov functional?

Our rates can indeed be obtained by decreasing the following Lyapunov functional: $L_k = k\tau(\mathcal{F}(\mu_k) - \mathcal{F}(\nu)) + \tau(\frac{1}{\tau}-\alpha) W_\phi(\nu,\mu_k)$. Indeed, $L_{k+1}-L_k \le 0$ by using the inequalities from equations (85) and (7). Thus $L_k$ is a Lyapunov function [3]. By using a telescopic sum, we recover the inequality $\mathcal{F}(\mu_k) - \mathcal{F}(\nu) \le \frac{1-\alpha\tau}{k\tau} W_\phi(\nu,\mu_0)$ obtained in Proposition 3. We will make this clearer.

Is it feasible to extend the proposed method to Riemannian-based approaches, such as Riemannian SVGD?

Extending the proposed methods to Riemannian manifolds is an interesting future direction of works that we did not consider yet.

For now, Wasserstein gradient flows on manifolds have been sparsely studied. For instance, [4] studied theoretically Wasserstein gradient flows of the negative entropy on manifolds. In a more applied setting, [5] computed Wasserstein gradient flows on Lie group using the JKO scheme and [6] performed Wasserstein gradient descent of the Sliced-Wasserstein distance on Hadamard manifolds. There, the schemes and analyis have been derived on a case-by-case basis.

To the best of our knowledge, the Mirror Descent algorithm has not been yet extended to Riemannian manifolds. One difficulty might come from the current theoretical analysis which requires to split the inner product, which is not directly possible when using log maps. Moreover, the scheme might also come with computational difficulties, as there might not be closed-forms easy to compute in general.

Is it possible to extend this approach to systems characterized by the Fisher-Rao metric?

The gradient in the space of probability distributions endowed with the Fisher-Rao metric is the 1st variation [7]. Thus, we might define Bregman divergences and the associated Mirror-Descent scheme using this gradient. This was studied e.g. in [8] for Mirror Descent with applications on the minimization of the KL divergence.

Could this approach be adapted to Hamiltonian Flow, which is the second order system to my knowledge?

This is an interesting direction for future works which we did not yet consider and which goes beyond our knowledge.

I recommend that the authors include a dedicated section in the appendix detailing the assumptions such as convexity and smoothness within functionals.

Thank you for this suggestion. We will expand our existing sections B.2 and C.3 in the appendix in this direction.

[1] Tzen, B., Raj, A., Raginsky, M., & Bach, F. (2023). Variational principles for mirror descent and mirror langevin dynamics. IEEE Control Systems Letters, 7, 1542-1547.

[2] Garrigos, G., & Gower, R. M. (2023). Handbook of convergence theorems for (stochastic) gradient methods. arXiv preprint arXiv:2301.11235.

[3] Wilson, A. (2018). Lyapunov arguments in optimization. University of California, Berkeley.

[4] Erbar, M. (2010). The heat equation on manifolds as a gradient flow in the Wasserstein space. In Annales de l'IHP Probabilités et statistiques (Vol. 46, No. 1, pp. 1-23).

[5] Bon, D., Pai, G., Bellaard, G., Mula, O., & Duits, R. (2024). Optimal Transport on the Lie Group of Roto-translations. arXiv preprint arXiv:2402.15322.

[6] Bonet, C., Drumetz, L., & Courty, N. (2024). Sliced-Wasserstein Distances and Flows on Cartan-Hadamard Manifolds. arXiv preprint arXiv:2403.06560.

[7] Gallouët, T. O., & Monsaingeon, L. (2017). A JKO splitting scheme for Kantorovich--Fisher--Rao gradient flows. SIAM Journal on Mathematical Analysis, 49(2), 1100-1130.

[8] Aubin-Frankowski, P. C., Korba, A., & Léger, F. (2022). Mirror descent with relative smoothness in measure spaces, with application to sinkhorn and em. Advances in Neural Information Processing Systems, 35, 17263-17275.

Thank you for your detailed response, and I appreciate to adjust the score. Additionally, I have a follow-up question regarding lines 1447 and 1448. Why is KDE necessary when conducting mirror descent on the simplex? To my understanding, since the density of the Dirichlet distribution is bounded, ensuring that the velocity field meets the boundary condition $\lim_{\Vert x \Vert\rightarrow\infty} v(x)=0$ should suffice. We could potentially use the equation $\int{\nabla_{x} [p(x)v(x)]} = 0 = \int{\nabla_x p(x) v(x)} + \int{p(x)\nabla v(x)}$ to circumvent the need for explicit density estimation. Could you please clarify this approach?

Thank you for your appraisal and positive comments on our paper.

The preconditioning is considered in some particular yet practically important case, but the complete generalisation of preconditioning is still to be finished.

We acknowledge that the complete generalization and study of the preconditioned gradient descent on the Wasserstein space is left for future works. We note that the convergence results would still be valid for $\phi_\mu$ strictly convex, and differentiable as described in the first paragraph of Section 4.

Is your current analysis generalisable to stochastic optimisation and to inexact model of function, more generally?

We expect our analysis to be generalisable to the stochastic setting. While we leave it for future works, we note that [1, 2] studied stochastic Wasserstein gradient schemes, while [3] studied Mirror Descent on $\mathbb{R}^d$ in the stochastic setting. Thus, it could be a nice avenue of future works to extend these results to Mirror Descent on the Wasserstein space.

[1] Diao, M. Z., Balasubramanian, K., Chewi, S., & Salim, A. (2023, July). Forward-backward Gaussian variational inference via JKO in the Bures-Wasserstein space. In International Conference on Machine Learning (pp. 7960-7991). PMLR.

[2] Backhoff-Veraguas, J., Fontbona, J., Rios, G., & Tobar, F. (2022). Stochastic Gradient Descent for Barycenters in Wasserstein Space. arXiv preprint arXiv:2201.04232.

[3] Hanzely, F., & Richtárik, P. (2021). Fastest rates for stochastic mirror descent methods. Computational Optimization and Applications, 79, 717-766.

Thank you for your valuable feedback and time. We have addressed your comments point-by-point below. Please don't hesitate to let us know if you have any further questions.

The paper is very dense. I would advise the authors to revise the paper to make it more accessible.

We understand your concerns about the paper's density and the difficulty in following the implementation details. Taking advantage of the extra page, we will work on revising the paper to make it more accessible and ensure that key explanations are clearer and more integrated into the main text.

How the algorithm is implemented?

We will describe more clearly how the schemes are implemented in practice in a dedicated section of the appendix.

In practice, we use a particle approximation of the probability measures, i.e. we sample $x_1^{(0)},\dots,x_n^{(0)}\sim \mu_0$ and then we apply the schemes to $\hat{\mu}_k^n = \frac{1}{n} \sum\_{i=1}^n \delta\_{x\_i^{(k)}}$ for all $k\ge 0$.

For Mirror Descent with a Bregman potential $\phi$ chosen as a potential energy $\phi(\mu) = \int V\ \mathrm{d}\mu$ with $\nabla V$ and $\nabla V^*$ known in closed-forms, we use equation (9) to implement it: for all $k\ge 0$ and $i\in \{1,\dots,n\}$, $x\_i^{(k+1)} = \nabla V^*\big(\nabla V(x\_i^{(k)}) - \tau \nabla\_{W\_2}\mathcal{F}(\hat{\mu}\_k^n)(x\_i^{(k)})\big)$. The Wasserstein gradient of $\mathcal{F}$ in most of the cases considered can be computed in closed-form, e.g. for $\mathcal{F}(\mu)=\int U\mathrm{d}\mu$, $\nabla_{W_2}\mathcal{F}(\mu) = \nabla U$, or for $\mathcal{F}(\mu)=\iint W(x-y)\ \mathrm{d}\mu(x)\mathrm{d}\mu(y)$, $\nabla_{W_2}\mathcal{F}(\mu) = \nabla W\star \mu$.

In the more general case, where the gradient of the Bregman potential cannot be inverted in closed-form, we used Newton's algorithm at each step, which is detailed in Appendix I.1, and we use equation (191) to implement it.

For the Gaussian experiment, we update directly the means and covariances. The closed-form of NEM is reported in equation (197), and the closed-form of the Forward-Backward scheme are in Appendix H.4 in equation (160).

Finally, for the preconditioned gradient descent scheme, we implemented it using equation (11): for all $k\ge 0$ and $i\in\{1,\dots,n\}$, $x_i^{(k+1)} = x_i^{(k)} - \nabla h^*\big(\nabla_{W_2}\mathcal{F}(\hat{\mu}_k^n)(x_i^{(k)})\big)$. In the experiments of Figure 3, we considered $h^*(x)=(\|x\|_2^a + 1)^{\frac{1}{a}} - 1$ with $a>0$, and its gradient can be computed in closed-form or by backpropagation.

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

We will complete the notation section. A coupling between $\mu\in \mathcal{P}\_2(\mathbb{R}^d)$ and $\nu\in\mathcal{P}\_2(\mathbb{R}^d)$ is a probability distribution $\gamma\in\mathcal{P}(\mathbb{R}^d\times \mathbb{R}^d)$ such that \pi^1_{\\#}\gamma=\mu and \pi^2\_\\#\gamma=\nu with $\pi^1:(x,y)\mapsto x$ and $\pi^2:(x,y)\mapsto y$. The $\mathrm{Id}$ is the identity mapping $x\mapsto x$.

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

$V$ and $W$ are functions from $\mathbb{R}^d$ to $\mathbb{R}$. $V$ (or $W$) $L$-smooth means that its gradient is $L$-Lipschitz. We will clarify this in the revision of the paper. To ask for $W$ even means that $W(-x)=W(x)$ for all $x\in \mathbb{R}^d$.

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

$\mathcal{W}$ is defined as $\mathcal{W}(\mu) = \iint W(x-y)\ \mathrm{d}\mu(x)\mathrm{d}\mu(y)$, see line 127.

OT map mentioned on line 151 is not defined.

An OT map between $\mu,\nu\in\mathcal{P}\_2(\mathbb{R}^d)$ is a map $T:\mathbb{R}^d\to\mathbb{R}^d$ satisfying T\_\\#\mu=\nu and $W_2^2(\mu,\nu) = \|T-\mathrm{Id}\|_{L^2(\mu)}^2$. It exists e.g. provided $\mu\ll\mathrm{Leb}$ by Brenier's theorem. We will clarify this in the revised version of the paper.