Optimal Transport Barycenter via Nonconvex-Concave Minimax Optimization

Thanks for your constructive comments and careful reading. We will fix minor errors in the revision. Please find our response to your main questions below.

The author should emphasize that the problem is not merely concave but strongly concave, as this distinction is crucial for ensuring the validity of the derived results. Additionally, the complexity bound introduced in line 210 does not hold for general nonconvex-concave problems and should be clarified accordingly.

Thanks for pointing out that our algorithmic convergence rate and the existing Euclidean result rely on strong concavity. Indeed, Theorem 1 works for $0 < \alpha < \beta < \infty$, under which Problem (4) is a nonconvex-strongly-concave problem. In the revision, we will clarify in line 106 to "as in the Euclidean nonconvex-strongly-concave optimization problems", revise line 207 to "(Lin et al., 2020) proved that, if $f$ is convex in $x$ and strongly concave in $y$, with suitable choices of step sizes $(\eta, \tau )$, the following bound holds:", and revise line 233 to "Fix $\nu$ and $0 < \alpha < \beta < \infty$, the objective functional $\mathcal{J} (\nu, \varphi_1, \dots, \varphi_n )$ is strongly concave in each $\varphi_i$." to emphasize that the dual program is strongly concave.

The algorithm tested in the numerical section differs from the one analyzed in the theoretical parts. In addition to the discretization of the function, the practical implementation approximates the projection onto the convex function space using the convex envelope. To enhance the rigor of the analysis and the credibility of the numerical results, it would be beneficial for the paper to either provide a more precise theoretical justification for these approximations or conduct numerical experiments on the exact algorithm analyzed.

Results produced by projection and conjugate envelope (i.e., double convex conjugates) are identical if $\alpha=0$ and $\beta=\infty$, and are close if we set $\alpha$ small and $\beta$ large. In Appendix, Subsection C.2, we compared the two algorithms on 1D distributions, where we set $\alpha=0.001$ and $\beta=1000$ in Algorithm 3 (projection), which performs comparably well as Algorithm 4 (convex envelope) for this example, i.e., no significant difference. Specifically, projection achieves $\mathcal{W}_2$-distance $7.610\times 10^{-5}$ (std: $3.367\times 10^{-5}$) and $L^2$-distance 0.608 (std: $0.194$), while convex envelope has $\mathcal{W}_2$-distance $7.771\times 10^{-5}$ (std: $3.32 \times 10^{-5}$) and $L^2$-distance 0.6434 (std: $0.451$).

To compute the projection for 2D distributions, one can use the algorithm proposed by Simonetto. Smooth Strongly Convex Regression. However, this involves solving a large convex quadratically constrained quadratic problem. For 2D distributions discretized on the grid of size $1024^2$, this problem would have $1024^4$ quadratic constraints, potentially taking weeks to solve. Due to the rebuttal time constraint, we are unable to report numeric results of running projection in 2D.

In line 253, I would like to understand why the projection operator onto the convex gradient smooth function class is guaranteed to be unique...

This projection is defined as $\mathcal{P}\_{\mathbb{F}\_{\alpha,\beta}}(\varphi) = \arg\min\_{\psi\in\mathbb{F}\_{\alpha,\beta}} \\| \varphi - \psi \\| \_{\dot{ \mathbb{H}^1}}$ for every $\varphi\in\dot{\mathbb{H}}^1$. First, we note that $\mathbb{F}\_{\alpha,\beta}$ is a convex set. In fact, For every $\varphi_0, \varphi_1\in\mathbb{F}_{\alpha,\beta}$, by definition, we have

This implies that for every $\lambda \in [0, 1]$, with $\varphi_\lambda = \lambda \varphi_1 + (1-\lambda) \varphi_0$, it holds that

So, the convexity of the set $\mathbb{F}\_{\alpha,\beta}$ is proved. Next, we will show that the projection is unique, which follows from the standard argument for projection to convex subset in Hilbert spaces. For $\varphi \in \dot{\mathbb{H}}^1$, if both $\psi\_0$ and $\psi\_1$ in $\mathbb{F}\_{\alpha,\beta}$ minimizes $\\| \cdot - \varphi \\| \_{\dot{\mathbb{H}}^1}$, by taking $\widetilde \psi = (\psi \_0 + \psi\_1)/2 \in\mathbb{F}\_{\alpha,\beta}$, we have

The optimality of $\psi\_0$ and $\psi\_1$ implies that $\psi\_0 =\psi\_1$. Thus, the uniqueness of projection is proved.

Thanks for your constructive comments. In the revision, we will cite the paper [Daaloul 2021], which proposed an interesting notion for approximating and sampling from unregularized Wasserstein barycenter. Please find our response to your main questions below.

the authors also need to clarify which algorithm uses GPUs or all of them used...

All algorithms are executed on Google Colab using the same hardware configuration. We compared the runtimes on L4 GPU and CPU for Example 1:

Note that our WDHA algorithm can naturally be parallelized because the core operations, such as the convex conjugate for the $n$ dual programs (the primary computing bottleneck), can be solved independently on high-performance computing (HPC) clusters (CPU/GPU). Our current implementation is not optimized for HPC-acceleration since the core code is written in C++ with for loops. It is possible that the WDHA algorithm can fully leverage the power of HPC frameworks (e.g., CUDA) and substantially benefit from parallelization. To engineer an HPC-accelerated version of WDHA is a valuable future direction.

For the min-max problem, we can only find for the stationary point but not global opt. Therefore the convergence proof is actually not for the optimal solution for barycenter problems.

If the distributions are regular enough, then a stationary solution is indeed a global optimum because the barycenter functional is convex in the linear structure (Proposition 7.19 in Santambrogio. Optimal Transport for Applied Mathematicians) [even though not geodesically convex]. This implies that any global minimizer $\overline{\nu}$ is a barycenter of $\mu_1, \dots, \mu_n$ iff $0=id - \frac{1}{n} \sum_{i=1}^n \nabla \varphi_{\overline{\nu}}^{\mu_i}$. Moreover, if $\mu_i$'s have densities, the barycenter functional is strictly convex and $\overline{\nu}$ is unique.

Now, suppose that $\varphi\_{\overline{\nu}}^{\mu\_i} \in \mathbb{F}\_{\alpha, \beta}$ for all $i$ (i.e., no approximation error due to assumed regularity), then by Lemma 3.2 and Definition 3.3, any smooth stationary point $\nu$ would satisfy that (Wasserstein gradient of $\mathcal{F}\_{\alpha, \beta}(\nu)$)= $id - \frac{1}{n} \sum\_{i=1}^n \nabla \varphi\_{\nu}^{\mu\_i}$ implies that $\nu$ is a barycenter. This justifies our convergence result to a first-order stationary point solves the global opt problem. This discussion was provided in Remark 3.7.

A simple example with known ground truth is provided in Appendix, Subsection C.1, wherebarycenter of four uniform round disks is computed under the same setting as in Examples 1 and 2. The true barycenter is known as a round disk in the middle. Our method performs uniformly the best to recover the ground truth.

how does the dimension influence the algo, can we apply more numerical studies on 3D cases.

Yes, you are right. A key factor impacting computation speed is the total number of grid points, which grows exponentially in dimension and results in the major computation burden for higher dimensional cases. One solution is to use parallel computing on HPC clusters to accelerate via the decoupled dual programs. Review 9GjF raised the same question to give a 3D example. Please find our reply therein.

How does the grid resolution influence the proposed algo, what's the sample complexity if we hope to extend the method to contionous ones.

A larger number of grid points leads to more accurate approximations but also increases computational cost, creating a tradeoff. Comparing to existing methods in the literature, we demonstrate that a grid size of 500 or 1000 achieves the highest accuracy with the least computation time in Examples 1 and 2. Sample complexity to accommodate discretization error is a challenging question. Given the stability estimate (Lemma 3.1), we conjecture that WDHA sample complexity depends on the smoothness of the continuous dual potentials and the domain dimension $d$ in a nonparametric fashion.

Please clarify the technical contribution for the theoretical results.

We acknowledge that Lemma 3.1 is derived based on the results of [Jacobs & Leger]. However, the other fundational Lemmas 3.2, 3.4, 3.5 are novel and original. Notably, the bound shown in Lemma 3.4 (see eqn (6)) improves upon existing results in literature regarding the quantitative stability of OT maps. This is achieved by restricting $\varphi$ to be $\alpha$-strongly convex and $\beta$-smooth. Lemma 3.2 establishes the first variation of $\max\_{\varphi \in \mathbb{F}\_{\alpha, \beta}} \mathcal{I}\_{\nu}^{\mu}(\varphi)$, which is both original and nontrivial due to the presence of maximization. These lemmas serve as fundamental building blocks for Theorem 3.6, integrating techniques from applied mathematics, OT theory, and optimization in Euclidean space.

Thanks for your constructive comments and careful reading. We will fix minor errors in the revision. Please find our response to your main questions below.

As for experiment 3, here it seems that spatial notions are introduced in the OT problem. Space, understood as the coordinate (x, y) of the four input distributions considered ... the comparison with techniques such as CWB and DSB is unfair.

The experiments 1-3 were arranged on the same $1024 \times 1024$ grid, except that they are uniform distributions on different supports/shapes. We agree that the spatial notion of placing the shapes (e.g., centers of disks/circles) generally implies an embedding of distributions. However, in our experimental setting, the 4 shapes were equilaterally arranged in the corners from the grid center. This means that the coordinate information $(x, y)$ does not impact the barycenter. For instance, we can move all shapes towards the center of the square grid domain (squeezing), and the barycenter does not change. Effectively, all of the 4 input distributions sit in a common domain centered around $(0,0)$ and the problem can be seen as a discretization on a fixed grid. Thus the comparison with CWB and DSB is fair in this particular configuration. To avoid confusion, we will revise line 911 to ``Here, the goal is to compute the barycenter of four uniform distributions, each supported within $[0,1]^2$ and taking the shape of round disks with a radius of 0.15, centered equilaterally at $(0.2, 0.2), (0.2, 0.8), (0.8, 0.2), (0.8, 0.8)$, respectively. Their densities are discretized on a $1024 \times 1024$ grid."

Is not straightforward for me the nonconvexity (w.r.t. the barycenters) and concavity (w.r.t. the potentials) of Problem (4).

We first provide some insights about the nonconvex-concave structure of the barycenter problem. As mentioned in replying to Reviewer 9GjF, the barycenter problem is defined as $\min_\nu \sum_{i=1}^n \alpha_iW_2^2(\nu, \mu_i),$ where the objective function is geodesically nonconvex in $\nu$ for $n \geq 3$ (in fact, it is concave over Gaussians!) However, the barycenter objective is convex in the linear structure (cf. our reply to Reviewer sfZf). This is connected to the concavity in the dual program of potentials, where their $L_2$ gradients are key to defining the Wasserstein gradient of $\mathcal{J}$. Hence, the Wasserstein barycenter functional above can be recast into Problem (4) as a primal-nonconvex and dual-concave program in two different geometries. Below we give more rigorous details about this structure.

Strong concavity in $\varphi$. Lemma 3.1 implies that $\mathcal{I}\_\nu^\mu(\varphi)$ is (strongly) concave on the function space $\mathbb{F}\_{\alpha,\beta}$ w.r.t. $\dot{\mathbb{H}}^1$-norm when $\mu$ admits a density function bounded away from zero and uniformly upper bounded. As the objective functional is $\mathcal{J}(\nu, \varphi) = \frac{1}{n}\sum_{i=1}^n \mathcal{I}\_\nu^{\mu\_i}(\varphi\_i)$ is separable in potentials, we know $\mathcal{J}(\nu,\varphi)$ is strongly concave w.r.t. $\varphi$.

Nonconvexity in $\nu$. Recall that $\mathcal{I}\_\nu^\mu(\varphi) = \int \frac{ \\| x \\| \_2^2}{2} - \varphi(x) d \nu + \int \frac{ \\| x \\|\_2^2}{2} - \varphi^\ast(x) d \mu,$ where $\varphi^\ast$ is the convex conjugate of $\varphi$. Only the first integration in the above definition depends on $\nu$. It can be shown that the functional $\nu\mapsto\int \frac{\rm{id}^2}{2} - \varphi d \nu$ is geodesically convex if and only if the function $\frac{\rm{id}^2}{2} - \varphi$ is convex, which cannot be guaranteed in our WDHA algorithm.

is not clear if parameters alpha and beta, defining the subset of the Sobolev space, are fixed a priori; are related (in some sense) to the input data; or are optimized.

Prameters $\alpha$ and $\beta$ are fixed: they are used for the theoretical purpose for convergence analysis of WDHA algorithm. In practice, we suggested using Algorithm 4 by taking the conjugate envelope (i.e., double convex conjugates) at each iterate, which could help to reduce the computational time than projection. We verified by numeric example that practical differences between small-$\alpha$-large-$\beta$ and conjugate envelope are small (see also our reply to Reviewer rrvS7 on Experimental Designs).

you suggest to replace the dual potentials projections with computing their second convex conjugate. Did you do some experiments to check the difference of the two versions of the algorithm for 2D cases (with visualizations)?

Reviewer rvS7 raised the same question. Please find our reply therein.

Is it possible to test the behaviour of the WDHA solution in a 3D context?

Review 9GjF raised the same question to give a 3D example. Please find our reply therein.

Thanks for your constructive comments. Please find our response to your questions below.

I think $1$ should be discussed that one can resort to using first order methods such as subgradient descent on the dual for exact barycenter calculations.

Thanks for pointing out this relevant first-order method; we will discuss the paper " $1$ Numerical methods for matching for teams and Wasserstein barycenters" in the revised version.

The experiments seem too simplistic... I suggest the authors include more comprehensive experiments...

We extended our WDHA implementation to handle 3D distributions and conducted an example in such setting. The goal is to compute the barycenter of three uniform densities supported on different shapes of ellipsoid contained in $[0,1]^3$. Our simulation setup is: $\frac{(x-0.5)^2}{a^2} + \frac{(y-0.5)^2}{b^2} + \frac{(z-0.5)^2}{c^2} \leq 0.8^2$, where the parameters are

$(a,b,c) \in \{ (0.25, 0.25, 0.5), (0.5, 0.25, 0.25), (0.25, 0.5, 0.25) \}.$

These three densities are discretized on the grid of size $200 \times 200 \times 200$. In addition, we include a weighted barycenter setting, where the weights is set to be $(1/4, 1/4, 1/2)$. (see below for the extension to weighted barycenter). We ran our algorithm on a cluster with 2000 iterations and it took 5 hours to compute the barycenters. The figure showing the computed barycenters by our methods can be accessed using the anonymous link https://drive.google.com/file/d/1k5Yg4YF9FzkEOi7PC03jep5coXKNb2Vz/view?usp=sharing. Moreover, we note that current version of the Python package POT does not support the computation of barycenter for 3D distributions.

Is this algorithm limited to barycenter weights set to $1/n$? How would it handle other barycenter weight settings?

Our algorithm can be generalized to the computation of weighted Wasserstein barycenter with any barycentric coordinates $(\alpha_1, \dots, \alpha_n)$ on the probability simplex

$\overline{\nu} = \arg\min_\nu \sum_{i=1}^n \alpha_iW_2^2(\nu, \mu_i).$

Precisely, we can replace the $\dot{\mathbb{H}}^1$-gradient of $\mathcal{J}$ with \mathbf{\nabla}\_{\varphi\_i} \mathcal{J}(\nu, \mathbf{\varphi}) = \alpha\_i (-\Delta)^{-1} (-\nu + {(\nabla \varphi\_i^{*})}{\\#} \mu\_i) and replace the Wasserstein gradient of $\mathcal{J}(\nu,\varphi)$ with $id - \sum\_{i=1}^n\alpha\_i\varphi\_i$.