Mirror and Preconditioned Gradient Descent in Wasserstein Space

Summary

The authors study the mirror descent method in Wasserstein-2 spaces. This is based on constructing the Bregman divergence functionals in Wasserstein-2 spaces. One can use it to define a mirror descent direction in the sense of pushforward mapping functions, which generalizes the classical Wasserstein gradient descent direction. The authors then studied the convergence analysis of the Wasserstein mirror descent algorithm. Several numerical examples are used to demonstrate the effectiveness of the proposed method.

Strengths

The authors define and apply the Wasserstein mirror descent directions to conduct sampling algorithms. The mirror descent method in Wasserstein space is a very natural and novel approach. Congratulations to the authors for using the Wasserstein-Bregman divergences in sampling algorithms. This paper's proof is very clear in the mirror descent method from optimization and optimal transport.

Weaknesses

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.

Questions

The authors also propose the Wassrsetin-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? I am willing to improve my score after addressing these questions.

Rating

9

Confidence

5

Soundness

4

Presentation

4

Contribution

4

Limitations

There is no limitation on this paper.

Summary

In this paper, the authors endeavor to integrate the concepts of mirror descent (MD) and preconditioned gradient descent (PGD) within the framework of Wasserstein distances, from a convergence theory perspective. Initially, they provide a comprehensive background including, but not limited to, Wasserstein distances, Bregman divergence, convexity, and smoothness. With assumptions regarding $\beta$-smoothness and $\alpha$-convexity, they successfully demonstrate the iterative procedures for MD and PGD in the Wasserstein space. Subsequently, they conduct experiments focusing on sampling methodologies and single-cell dynamics.

Strengths

1. **Solid Mathematical Foundations:** The authors offer a rigorous mathematical derivation concerning the convergence behaviors of MD and PGD, supplemented by a thorough analysis of the theoretical background. 2. **Relevance to Conference Themes:** The paper presents fundamental theoretical analysis that is highly pertinent to the themes of the conference. 3. **Comprehensive Background Introduction:** The supplementary materials provide an extensive overview necessary to understand the nuances of the paper. The foundational work done by the authors is commendable.

Weaknesses

1. **Results Presentation:** The presentation of results could be enhanced by providing more detailed visualizations. For instance, it would be beneficial to include a snapshot illustrating particle evolution for Figure 1, right $K_2$. Additionally, if the authors have access to MATLAB, it is recommended to utilize the `surf` function for the contour of Dirichlet distribution's pdf in Figure 4 to better demonstrate the sampling tasks. 2. **Clarification of Definitions:** It is recommended that the authors include explicit definitions of convergence and global convergence to preclude any potential misunderstandings in the appendix. 3. **Visualization of Convergence Data:** I suggest employing seaborn's `sns.lineplot` with a shaded area to depict the error bars, which would support the assertions made in checklist item 7 more robustly. 4. **Overlooked References:** 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. ---- References: [1]. Tzen, Belinda, et al. "Variational principles for mirror descent and mirror langevin dynamics." IEEE Control Systems Letters 7 (2023): 1542-1547.

Questions

1. **About the Proof of Convergence:** 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? 2. **Alternative Proof Approaches:** Could the convergence location and rate be established from the perspective of a Lyapunov functional? This approach may provide a robust framework for proving stability and convergence. 3. **Riemannian-based Approaches:** Is it feasible to extend the proposed method to Riemannian-based approaches, such as Riemannian SVGD [1]? 4. **About Higher Order Systems and Extensions:** Is it possible to extend this approach to systems characterized by the Fisher-Rao metric [2]? Furthermore, considering that the continuity equation in Wasserstein space represents a first-order system, could this approach be adapted to Hamiltonian Flow [3], which is the second order system to my knowledge? --- ### References [1]. Zhang, Ruqi, Qiang Liu, and Xin Tong. "Sampling in constrained domains with orthogonal-space variational gradient descent." Advances in Neural Information Processing Systems 35 (2022): 37108-37120. [2]. Neklyudov, Kirill, et al. "Wasserstein Quantum Monte Carlo: A Novel Approach for Solving the Quantum Many-Body Schrödinger Equation." Advances in Neural Information Processing Systems 36 (2024). [3]. Ambrosio, Luigi, and Wilfrid Gangbo. "Hamiltonian ODEs in the Wasserstein Space of Probability Measures." Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences 61.1 (2008): 18-53.

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

4

Limitations

I appreciate the efforts of the authors. Nevertheless, I recommend that the authors include a dedicated section in the appendix detailing the assumptions such as convexity and smoothness within functionals. This addition would significantly aid in understanding and following the methodologies employed in this work.

Summary

This paper presented a unified view on the functional optimization on the Wasserstein space. Authors proposed a mirror descent algorithm and a preconditioned gradient descent algorithm, both are applied to minimize the objective functional over the Wasserstein space. Many existing algorithms in the Wasserstein space-related literature can be recovered or seen as a special case of the algorithms proposed by the authors. Convergence guarantees were provided and backed with numerical results.

Strengths

Mirror descent and preconditioned GD are closely related and have been seen as generalizations of GD. Though Wasserstein GD have been proposed, the extension to MD and preconditioned GD is a natural step. The formulation in the paper is general and can be applicable to a wide range of studies in both machine learning application and theory. The literature review provided by the authors seem to be thorough, and the position of the current manuscript in the literature is well explained. The relationships between the proposed algorithm and existing algorithm were also clear throughout the paper. The reviewer did not read the full analysis in the appendix, but the technical aspects of this paper seems solid and intuitive.

Weaknesses

The main concern I have is as presented in the MD optimization literature: what new tool does the MD formulation bring to the table? Previously, the convergence of MD is similar to or worse than GD. The literature on relative strong smoothness and convexity enabled MD to converge even when the objective function is not smooth/convex in the Euclidean sense (which was really mostly KL). 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. But the application offered later is largely limited to existing applications KL or Sinkhorn, whereas these specific cases seem to have been already studied by the literature anyway. Some sections are not well presented. For instance, the subsection starting at l.136, the statements on the potential energies are confusing and it is very difficult to associate on sentence with either the previous or the latter sentence (such as "... specific example. In particular, we have xxx. We will also consider...", is the "xxx" with the specific example of what the authors consider?) Similar confusions also exist in this paper, but to a lesser extent. While the contribution is clear from reading the entire paper, the "Contributions" paragraph does not help reader and clarify what this paper offers. Though discussed briefly, 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.

Questions

some minor grammatical errors such as l.158 ($T_\eta^\nu$ OT maps from $\eta$...)

Rating

6

Confidence

4

Soundness

4

Presentation

3

Contribution

3

Limitations

The authors did address the limitations of this work.

Summary

The paper generalises mirror descent and preconditioning approaches from optimisation over R^d to optimisation over Wasserstein space, which is applicable in some actual problems, and confirms by the corresponding numerical experiments the efficiency of these approaches.

Strengths

The paper completely corresponds to its stated goal, uses the most actual theoretical frameworks of long known approaches from optimisation over R^d, and focuses indeed on the most fundamental methods, which will allow a follower to develop any desired extension based on idea from optimisation over R^d and the detailed examination of technical issues related to the transition to Wasserstein space made in this work. Practical contribution of this paper is also notable, as it demonstrates in numerical experiments, but I guess that thanks to the proposed theoretical basement, many improvements can be made on top of that, which is to be done in future works.

Weaknesses

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

Questions

There is an aspect which was not mentioned in this paper, but is interesting from the theoretical point of view: stochastic oracle. The generalisation of stochastic mirror descent to Wasserstein space seems to be the demanding yet possible task, when that of preconditioning is challenging even in R^d case, but it would be interesting to find the proper notions of stochasticity and flexible enough analysis to deal with it in case of Wassestein space. Is your current analysis generalisable to stochastic optimisation and to inexact model of function, more generally?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

The limitations are clear from reading.

Summary

The paper provides convergence analysis for Mirror descent (with and without additional preconditioning) in the Wasserstein space. The analysis is performed for relatively smooth and convex functionals. The use of the studied algorithms on computational biology problem is illustrated.

Strengths

The theory in the paper seems sound and substantially extends the prior work overcoming some technical difficulties related to the discrete nature of the algorithms (explicit discretization) and non-Euclidean geometry in the Wasserstein space. Convergence rates are similar to those in finite dimensional space for MD.

Weaknesses

The paper is very dense and seems like not well adapted to 9 page format. A lot of explanations and implementation details are referred to the appendix. It is very difficult to read the paper when it is written is such style. For example, although I have spent several hours only trying to understand how the algorithm is implemented, I did not manage to do it even with the help of the appendix. I would advise the authors to revise the paper to make it more accessible. 1. Notations section should be improved. Many things are not properly defined throughout the paper, e.g., the coupling between distributions, what "Id" is? 2. On line 138, what does it mean that the sets $V$ and $W$ are "differentiable and $L$-smooth with $W$ even? 3. On line 142, the calligraphic $\mathcal W$ is not defined. 4. OT map mentioned on line 151 is not defined.

Questions

n/a

Rating

6

Confidence

3

Soundness

4

Presentation

2

Contribution

3

Limitations

n/a

Reply to authors

The authors have addressed my questions. Congratulations to the authors for working in a very important direction. More serious computations and analysis can be conducted along this line.

Rebuttal Acknowledgement

I thank the authors for answering my questions and I trust the authors can improve the presentation in the next revision. I have no further concerns.