Title: Non-asymptotic convergence bounds for Wasserstein approximation using point clouds
Authors: Quentin Mérigot
Filippo Santambrogio
Clément SARRAZIN
Venue: NeurIPS 2021
Keywords: Optimal transport
Wasserstein distance
Non-convex optimization
Polyak-Lojasiewicz inequality
Lloyd's algorithm
Point clouds

Abstract: Several issues in machine learning and inverse problems require to generate discrete data, as if sampled from a model probability
distribution. A common way to do so relies on the construction of a uniform probability distribution over a set of $N$ points which
minimizes the Wasserstein distance to the model distribution. This minimization problem, where the unknowns are the positions of the atoms, is non-convex. Yet, in most cases, a suitably adjusted version of Lloyd's algorithm in which Voronoi cells are replaced by Power cells, leads to configurations with small Wasserstein error. This is surprising because, again, of the non-convex nature of the problem, which moreover admits spurious critical points. We provide explicit upper bounds for the convergence speed of this Lloyd-type algorithm, starting from a cloud of points sufficiently far from each other. This already works after one step of the iteration procedure, and similar bounds can be deduced, for the corresponding gradient descent. These bounds naturally lead to a sort of Poliak-Łojasiewicz inequality for the Wasserstein distance cost, with an error term depending on the distances between Dirac masses in the discrete distribution.

TL;DR: We provide explicit error estimates for the output of the Lloyd-type algorithm used in uniform optimal quantization with respect to the sample size.

Review discussion:
Upsides:
+ This paper tackles an important problem of approximation of measures in Wasserstein distance by discrete distributions and gives the first theoretical results for this problem.
+ The PL-like inequality is new and interesting from a nonconvex optimization perspective.
In particular, it gives an approximation error for a nonconvex setting that decreases with the sample size. This type of analysis may have broader applications.
+ The main results and their proofs are clear.
Downsides:
- The experiments are not extensive, and rather serve as a proof of concept.
Furthermore, they do not illustrate the curse of dimensionality encountered in the theorems.
- How does one compute the barycenter of the Voronoi cells in practice? What is the general computational complexity and practicality of implementing this method?
- Theorem 7 and Remark 5 are not sufficiently clear.
In particular, the discrepancy in the rates of Theorem 7 and Corollary 4 is not clearly expounded on. In general, it would be helpful to collect the results for the various methods/theorems together to make them easier to interpret and compare.
- In some sense, for the one-step algorithm, the method seems to be learning the barycentric projection optimal transport map, provided that the initial $Y^0$ is drawn from some density. This would then explain the types of rates that are seen in this setting, due to the rates of estimation for the optimal transport distance. It would be nice if the authors could comment on the connection here.
- The guarantees do not give hints as to what types of measures are easier to approximate than others.
Furthermore, the dimension dependence of the constants is not discussed.
I still find the result interesting and think that it is close to being ready for publication.
However, due to some issues raised by other reviewers, as well as the limited practical experiments, I will leave my score at a 6.
No
Dear Reviewer,
We would like to start by thanking you for your thorough review and critique of our submission. We have tried to answer them in a satisfying fashion. Here are a these which we felt needed a more elaborate response:
First, let us answer the two first similar remarks: "The experiments are not extensive, and rather serve as a proof of concept.
Furthermore, they do not illustrate the curse of dimensionality encountered in the theorems.", "How does one compute the barycenter of the Voronoi cells in practice? What is the general computational complexity and practicality of implementing this method?"
$\rightarrow$ The objective of the experiments is to show that in some (simple) cases, it is possible to get a quantization error of $O((1/N)^{1/d})$ with one step of Lloyd's algorithm instead of the expected optimal one of $O((1/N)^{2/d})$ -- illustrating Appendix E. Since we are in low dimension, we can resort to quadratures for computing the barycenters (we rely on the opensource PySDOT library for this purpose). In high-dimension, this could be done if the source measure is finitely supported (in which case barycenters are simply weighted averages), or using Monte-Carlo methods.
You also mentioned "In some sense, for the one-step algorithm, the method seems to be learning the barycentric projection optimal transport map, provided that the initial is drawn from some density. This would then explain the types of rates that are seen in this setting, due to the rates of estimation for the optimal transport distance. It would be nice if the authors could comment on the connection here."
$\rightarrow$ One step of the Lloyd algorithm is indeed a barycentric projection of the initialization on the measure $\rho$.
However, for these kinds of updates, one would obtain the rates of theorem 3 using Caffarelli stability (we have in mind the estimates of https://arxiv.org/abs/2107.01718). The hypotheses to use such stability are then very demanding whereas our result is still valid for very general measures $\rho$, and in particular measures which are not supported everywhere on the convex domain $\Omega$. We will add a reference to these similar stability rates, highlighting where our hypotheses are weaker.
Another critique was that "The guarantees do not give hints as to what types of measures are easier to approximate than others."
$\rightarrow$ Indeed: the bounds merely depend on the diameter of the support of the measures.
On the other hand, the analysis shows that if the measure is supported on a $k$-dimensional convex set of $R^d$, then the error bounds will behave as in $R^k$. Our belief is that measures supported on smooth and low-dimensional submanifolds of $R^d$ should be easier to approximate than absolutely continuous measures using Lloyd's algorithm, but we currently do not have theorems backing up this impression.
Finally, you wrote that "Furthermore, the dimension dependence of the constants is not discussed."
$\rightarrow$ We note that the constants in Theorem 3 and Corollary 4.(18) are explicit: they depends exponentially on the dimension. We will add a remark to underline this fact.
We will also better highlight the origin of the discrepancies between the two convergence rates for one and several Lloyd steps, as you suggested.
Originality : To my knowledge, this paper is the first to provide accuracy bounds for the Wasserstein distance between a continuous measure and a point cloud after the first step of Lloyd's algorithm, reducing the gap between experimental and theoretical results, the latter being mostly asymptotic (in particular in the optimal quantization community).
In particular, this paper proves the advantage of Lloyd's algorithm, even if the functional $F_N$ to optimise wrt to the point cloud $Y$ is non-convex.
Quality : This paper contains strong theoretically results, for which the assumptions and consequences are clearly stated and explained.
Moreover, some unfamiliar behaviours, such as in Figure 2 (right) are still presented and commented.
Clarity : This paper well written and the problems are clearly stated.
Moreover, natural results following the first part of the paper, such as Corollary 5 (probabilistic case of quantization) are also proven.
However, in my opinion, the introduction of regression problems is not completely justified as it can be more confusing than anything else (it is neither used or leveraged in the applications for example).
Significance : Lloyd's algorithm is largely used in several communities, from optimal transport to optimal quantization problems, even if it only solves provide critical points of a non-convex problem.
Therefore, such a study on the behaviour of Lloyd's algorithm, both in finite time and asymptotically, is of great interest.
No
Dear Reviewer,
We would like to start by thanking you for your thorough review and critique of our submission. We have tried to answer them in a satisfying fashion. Here are a these which we felt needed a more elaborate response:
You mentioned a Limitation : "As the authors discuss in the conclusion, the bound obtained for Lloyd's algorithm does not reach the theoretical rate of convergence, even if the initialisation points are sufficiently far apart."
$\rightarrow$ The error bound obtained after one step of Lloyd's algorithm $(O(N^{-1/d}))$ is below the optimal uniform quantization error $(O(N^{-2/d}))$. It turns out (see Remark 3 and Appendix E) that this is not a limitation of our analysis but rather a limitation of the algorithm itself: it is possible to build examples where the error bound after one step of Lloyd's algorithm is $O(N^{-alpha})$ with alpha arbitrarily close to $1/d$.
You also asked if "Considering the bounds on $W_2^2(\rho,\delta_{B_N(Y)})$, would the gain after iterations of Lloyd's algorithm be significantly better than after a single step?"
$\rightarrow$ This is not possible in general, due to the counter-example of Appendix E (in this example, further iterations of Lloyd's algorithm do not improve the result). One could hope to improve the error bound after several applications of Lloyd's algorithm, and it was the object of our Section 3, but the main result of this section (Theorem 7) leaves a lot of room for improvement (the main difficulty is somehow explained in Remark 5).
We will also correct the typos you highlighted, thank you for bringing them to our attention!
The main finding in this work is that a well-separated discrete measure leads to progress in gradient descent for the optimal quantization problem. This result is achieved by a short argument which proceeds by showing that a well-separted discrete measure must induce well-conditioned power cells, and so their barycenters must have relatively little transport cost. The result is further developed by combination with a previous observation that their objective is semi-concave, which allows some understanding of multiple steps of gradient descent.
While the primary result is interesting and novel, I feel that there is insufficient context and empirical or theoretical development to merit acceptance.
In particular, I wonder whether this result is of questionable significance since in practice, to perform the gradient descent step one must estimate integrals of the density over the relevant power cells. If this is done via Monte-Carlo (which would be standard in a high-dimensional context), then the natural competitor would be the empirical measure on a sample of size $N$, which would achieve rates of the order $N^{-1/d}$ [1]. How does this compare to the estimator proposed here in terms of relative computational/statistical efficiency? If theoretical results are difficult, can an empirical comparison be made?
Another issue is that I feel the context for the obtained rates is not clearly explained. This is somewhat addressed in Remark 3, but I feel it should be more clearly explained that generically we expect for $N$ points and a "nice" density supported a bounded subset of $\mathbb{R}^d$, we cannot do better than $N^{-1/d}$ (this being a strength of the obtained result that is not properly expounded).
Finally, I feel the experimental results section could be improved by a much more thorough substantiation of the claim that Lloyd's algorithm typically finds good local minima. Or, if there is already such a work in the literature, it should be cited (apologies if this has been cited and I missed it).
[1] Fournier and Guillin, "On the rate of convergence in Wasserstein distance of the empirical measure"
No
Dear Reviewer,
We would like to start by thanking you for your thorough review and critiques of our submission. We have tried to answer them in a satisfying fashion. Here are a these which we felt needed a more elaborate response:
Your first critique was that "to perform the gradient descent step one must estimate integrals of the density over the relevant power cells. If this is done via Monte-Carlo (which would be standard in a high-dimensional context), then the natural competitor would be the empirical measure on a sample of size $N$, which would achieve rates of the order $N^{-1/d}$ [1]. How does this compare to the estimator proposed here in terms of relative computational/statistical efficiency? If theoretical results are difficult, can an empirical comparison be made?"
$\rightarrow$ Our experiments are done in a low-dimensional setting, and we can resort to quadratures for computing the barycenters (we rely on the PySDOT library for this purpose). In higher dimensions, this method of sampling a density is indeed less legitimate, as a random sample according to $\rho$ achieves a better convergence rate (although only on average, whereas barycenters are a deterministic way of finding such an optimal sampling). Let us note that in the case of Wasserstein regression or, more generally in machine learning, when the goal is to simulate a sampling of $\rho$, such a random sample may not be available.
You also mentioned: "This is somewhat addressed in Remark 3, but I feel it should be more clearly explained that generically we expect for points and a "nice" density supported a bounded subset of $\mathbb{R}^d$, we cannot do better than (this being a strength of the obtained result that is not properly expounded)."
$\rightarrow$ We agree and will make this remark more visible by incorporating Corollary 5 of the appendix E in the main text.
Finally, you wrote: "I feel the experimental results section could be improved by a much more thorough substantiation of the claim that Lloyd's algorithm typically finds good local minima. Or, if there is already such a work in the literature, it should be cited (apologies if this has been cited and I missed it)."
$\rightarrow$ Notice that we do not claim convergence of the Lloyd algorithm and in general, it is not clear whether local minima can be avoided or even if there exists non-minimizing critical points in which it could fall. To the best of our knowledge, there has been no study of the stability of such degenerate critical points, but experiments seem to show that they all share the unstable, lower-dimensional structure as the one featured in Remark 1.
Thanks for the reply.
Regarding the convergence of Lloyd's algorithm in typical situations, I was referring to the line in the abstract: "Yet, in most cases, a suitably adjusted version of Lloyd’s algorithm — in which Voronoi cells are replaced by Power cells — leads to configurations with small Wasserstein error." It arises in another place, on line 150-151: "Experiments suggest that following the discrete gradient flow of F_N does not bring us to high energy critical points of F_N, such as those described in Remark 1, unless we started from an adversely chosen point cloud". I don't see these experiements or a citation for them.
Regarding the expected rates: here's a more explicit discussion of what I was thinking. The following is somewhere in the matching problem literature no doubt, but I reproduce it here since I don't know a reference off the top of my head. For simplicity let's assume that rho is supported on the unit cube in R^d and there lower bounded by c.
Then in this case, Remark 2 can be turned into a lower bound: the best discrete measure on N points can't get closer than N^{-2/d} in W2^2 distance. Why?
\min_{Y \in (\mathbb{R}^d)^N} W_2^2(\rho, \delta_Y) = \min_{Y \in (\mathbb{R}^d)^N}  \sum_{i = 1}^N \int_{V_i(Y)} \|x - Y_i\|^2 d \rho(x)
\geqslant c \min_{Y \in (\mathbb{R}^d)^N}  \sum_{i = 1}^N \int_{V_i(Y)} \|x - Y_i\|^2 d u(x).
$$ Where $u$ is the uniform measure on $[0,1]^d$. Now put $\Delta_N := \{w \in \mathbb{R}^N_{\geqslant 0} \colon \sum w_i = 1\}$.
Then the above is at least
c \min_{Y \in (\mathbb{R}^d)^N, w \in \Delta_N} W_2^2(u, \sum_i w_i \delta_{Y_i} ) \gtrsim^* \min_{w \in \Delta_N} \sum_i w_i^{1 + 2/d} \gtrsim N^{-2/d}.
$$ The last inequality is Cauchy-Schwarz, and the starred inequality is a sketch but is based off the idea that the most efficient way to pack w_i mass from the uniform measure to a point Y_i is via a disk centered at Y_i with radius w_i^{1/d} up to constants. Anyways this should be somewhere in the matching problem literature and I can help you find it if you want.
But the conclusion is important because it shows that up to constants we can't generally expect better than N^{-2/d} performace when we are restricted to N discrete points.
Thus after one step of gradient descent, Corollary 4 shows that the resulting arrangement is optimal up to a square. This is the sort of theoretical development I think is currently missing from the draft.
This paper considers an interesting problem, i.e., minimizing the Wasserstein distance between a known continuous distribution and a discrete distribution defined over a set of N points. Such a problem cannot be solved in KL divergence, as the KL divergence between any continuous distribution and discrete distribution is always infinite.
However, I am a bit confused by the contribution of the paper; it seems that Lloyd’s algorithm based on power cell has been studied in [11] and [14] under similar setting, see, equation (10) and Proposition 1 in the current paper. The authors should make it clear whether the adjusted Lloyd’s algorithm is a main contribution of the paper or not.
In addition, more details about how to efficiently compute the barycenter for the power cell should be provided.
In general, there are some merits in the paper, but it is not well-written in the current form, and a lot of details need to be clarified.
Minor comments:
1.	In Line 118, the notation $(T_\mu, Id)$# is used without definition
2.	In line 62, should be “In the previous inequality”.
3.	In line 133, the definition of barycenter is not correct, the current definition gives $b_i(Y)=1$, should be $b_i(Y)=N\int_{P_i(Y)}x d\rho(x)$.
4.	The order of Figure 2 and 3 should be swapped.
No
Dear Reviewer,
First off, thank you for your thorough review and critiques of our submission. We have tried to answer them in a satisfying fashion. We wish to adress the following in more details:
You asked that "The authors should make it clear whether the adjusted Lloyd’s algorithm is a main contribution of the paper or not."
$\rightarrow$ This is not a contribution of our article, our contribution is on the analysis of the behaviour of this algorithm. We will mention this more clearly in the future versions.
You also wrote: "In addition, more details about how to efficiently compute the barycenter for the power cell should be provided."
$\rightarrow$ Our experiments is in a low-dimensional setting, where we can resort to quadratures for computing the barycenters (we rely on the PySDOT library for this purpose). In high-dimension, this could be done if the source measure is finitely supported (in which case barycenters are simply weighted averages), or using Monte-Carlo methods.
We will also fix the typos that you listed, thank you for finding them!

Meta-review:
role: PC
data: comment: The reviewers recognize the theoretical interest of the authors' contribution, bounding in particular the error done after one step of gradient descent to minimize Wasserstein distance between a given density and a uniform discrete N-point distribution. They however raise valid concerns, particularly on the complexity needed to compute such gradients, on how it would rely on obtention of samples from the density being approximated. These are important aspects that are not treated in the paper. The reviewers also point to the lack of a discussion contrasting the guarantee they obtain to the one that results from taking as an approximating point cloud an N-sample from the target density. 

It therefore appears that the paper contains a very promising set of results, but is at this stage too preliminary and not ready yet for publication at NeurIPS.

Decision: Accept (Poster)