OPTIMAL TRANSPORT BARYCENTER VIA NONCONVEX CONCAVE MINIMAX OPTIMIZATION

We just want to add the following points.

(iii) doesn't have a ground truth that can convince the readers of the validity of a plot that doesn't look good and has bad Sobolev norm

Ans: we added a new simulation in the appendix subsection C.1 which has gournd truth. The ground truth of this example is the uniform distribution on the disk of radius 0.15 centered at (0.5, 0.5), which is not smooth and has $\dot{\mathbb{H}}^{-1}$-norm (negative Sobolev norm) bounded by 1. Please kindly see our response to point 4.

(i) doesn't look good

Ans: In our newly added simulation example in the appendix of subsection C.1, the Wasserstein distance between the barycenter of our method and the truth is smaller than that of both CWB and DSB. The $L^2$-distance between the density of barycenter of our method and the truth is also the smallest. We think that our barycenter look better visually.

In practice, one would pass the WSB result over a sharpening filter or pass the DSB result over a smoothing filter.

Ans: we indeed conducted an additional thresholding step for WSB and CSB in exmaple 1. Interested readers can check Figure 1 of the paper for the results.

We sincerely appreciate your time, thoughtful consideration, and constructive guidance on our paper. Following your and the other referees' valuable comments, we added three additional simulation studies in the appendix of the revised version to address your questions. If we have successfully addressed some of your concerns, we would kindly ask you to consider reflecting this in your rating. Thank you again for your invaluable feedback.

Below we provide point-by-point response to your comments in italics.

1. How is the differentiation implemented in this work? Terms such as are not continuous and admit discrete proxies. Is it done by finite difference? And, if so, what is the scheme? Or, alternatively, is it done by performing FFT, and is some smoothing filter performed?

Ans: you are right. We used finite difference approach. In particular, let $\Omega = [0,1]^2$ and $\\{(x\_i, y\_j)\\}\_{i,j=0}^{m}$ be the $m^2$ equally spaced grid points, we have $x\_0=0, y\_0=0$ and $x\_i = i/m, y_j = j/m$ for $i,j\neq 0$. Given the evaluations $\{ \varphi\_{i,j} = \varphi(x\_i, y\_j) \}$ of function $\varphi$ on these grid points, we compute the gradient at point $(x\_i, x\_j), i, j \neq 0$ as

$

\nabla \varphi (x_i, y_j) = \left( \begin{array}{c} \frac{\varphi_{i,j} - \varphi_{i-1, j}}{h} \ \frac{\varphi_{i,j} - \varphi_{i, j-1}}{h} \end{array} \right),

$

where $h=1/m$. No smoothing or filtering is performed for computing the gradient. FFT is used to solve the Poisson equation ($\Delta f = u$) with zero Neumann boundary condition.

2. How is the pushforward implemented in this work? Assuming is computed, the term (\nabla \phi)\_{\\#} might exceed the grid specified. How is this going to be resolved in the numerical example? These are questions that need to be addressed within this work. As of now, the reviewer's understanding is that this work only quotes the Jacobs & Leger paper, but that would be insufficient.

Ans: To discuss the implementation of (\nabla \varphi)\_{\\#} \nu, we describe how the mass $\nu(x\_i, y\_j)$ (density value of $\nu$ at a point $(x\_i, y\_j)$) is splitted and mapped \citep{jacobs2020fast} as follows. Since $\varphi$ is convex, we observe that $\nabla\_{x} \varphi (x\_i, y\_j) \leq \nabla\_{x} \varphi (x\_{i+1}, y\_j)$ and $\nabla\_{y} \varphi (x\_i, y\_j) \leq \nabla\_{y} \varphi (x\_{i}, y\_{j+1})$. Let $\mathcal{R}(x\_i,y\_j)$ be the quadrilateral formed by 4 points $\nabla\varphi(x\_i,y\_j),\nabla\varphi(x\_{i+1},y\_j),\nabla\varphi(x\_i,y\_{j+1}),\nabla\varphi(x\_{i+1},y\_{j+1})$ and pick the mesh grids $\\{ (\widetilde{x}\_{i'}, \widetilde{y}\_{j'}) \\}\_{i',j'=1}^k \subset \mathcal{R}(x\_i,y\_j)$, where

with $0 = \alpha_{0}\leq \alpha_1 \leq \dots \leq \alpha_{k} = 1, 0 = \beta_{0}\leq \beta_1 \leq \dots \leq \beta_{k} = 1$. The mass of $\nu(x\_i,y\_j)$ is first uniformly distributed to the meshed grids $\\{(\widetilde{x}\_{i'}, \widetilde{y}\_{j'})\\}\_{i',j'=1}^k$. Then, the mass at each point $(\widetilde{x}\_{i'}, \widetilde{y}\_{j'})$ is distributed to 4 nearest grid points in $\\{ (x\_i, y\_j )\\}$, inversely proportional to their distances. If (\nabla \varphi)\_{\\#} \nu exceeds the grid specified, the mass will be distributed to the boundary points instead. The discussions about computing $\nabla \varphi$, $\varphi^{*}$ and (\nabla \phi)_{\\#} \nu are added to appendix subsection D.2 of the revised version.

3. It is quite confusing as to why the authors use a 1024 $\times$ 1024 grid for Sinkhorn in Figure 1. From the plot, it is clear that the distributions occupy a small region of the grid. Is a smaller size used? This should a priori be mentioned in the numerical experiment section. If somehow the full 1024 $\times$ 1024 space is used, then it would be quite unfair to the Sinkhorn algorithm. Alternatively, does the code for WDHA use the sparsity of the source and target distribution?

Ans: our method is advantageous for high-resolution densities. The $1024 \times 1024$ is the standard pixel size for images and has nothing in particular. Our method wins on $500 \times 500$ images (demonstrated by example 2) as well. All the methods are compared on the full grid and are fed with the same densities, and thus the comparison is fair. In example 1, we put four shapes in the corners solely for visualization purpose, we can enlarge and center the shapes, the results would be the same. Example 2 would correspond to the case that the distributions occupy large regions around the center.

We sincerely appreciate your time, thoughtful consideration, and constructive guidance on our paper. Following your and the other referees' valuable comments, we added three additional simulation studies in the appendix of the revised version to address your questions. If we have successfully addressed some of your concerns, we would kindly ask you to consider reflecting this in your rating. Thank you again for your invaluable feedback.

Below we provide point-by-point response to your comments in italics.

1. Theory: One weakness is that the theoretical analysis is performed for Algorithm 3, while the practical implementation uses Algorithm 4. Also, the theorem assumes lower bounded densities of the distributions, which may not be true in practice (for example, in the case of Figure 1). Even in this ideal setting, there is also an extra assumption on the uniform boundedness of $\nu^t$. Therefore, the convergence guarantee of the algorithm in practice remains unclear.

Ans: Thank you. The same issue also exists in other works that utilize the $\mathbb{H}^1$-gradient if any convergence theories are to be developed; see the assumptions in proposition 1 [Jacobs and L´eger, 2020] and Theorem 2.8 [Jacobs et al., 2021]. We don't think this should be the reason to deny its satisfying empirical performance either. Please check our additional simulation with ground truth. Without additional regularities, generalization of the Theorem can be challenging, the main obstacle is that the map $\nu \rightarrow T\_{\nu}^{\mu}$ is not Lipschitz, see [Berman, 2021] for more details regarding this.

Empirically, the projection can be computed by the approach in [Simonetto, 2021], which is implemented as the Python function reg4opt.regression.convex\_regression in library reg4opt. Unfortunately, this requires solving a large convex quadratically constrained quadratic problem, and the code could take weeks to run. Instead, we implement the algorithm with projection (Algorithm 3) and test the method on 1D distributions. For each repetition $t$ and $i=1,2,3$, we let $\mu\_{i,t}$ be the truncated version of $N(a\_i,{\sigma\_{i}}^2)$ on the domain $[0,1]$, where $a\_i \sim \text{uniform}\left[0.3,0.7\right]$ and $\sigma\_i \sim \text{uniform}\left[0,1\right]$. Let $\overline{\nu}\_t$, $\overline{\nu}\_{1, t}$, $\overline{\nu}\_{2, t}$ be the true barycenter, computed barycernter from Algorithm 3 (with projection onto $\mathbb{F}\_{\alpha, \beta}$), computed barycernter from Algorithm 4 (with double convex conjugates) respectively. We repeat the experiment 300 times and report two types of average distances between the groundtruth and each computed barycenters : average $\mathcal{W}\_2$-distance $\frac{1}{T} \sum\_{t=1}^T \mathcal{W}\_2(\overline{\nu}\_t, \overline{\nu}\_{j, t})$ and average $L^2$-distance between densities $\frac{1}{T} \sum\_{t=1}^T (\int (\overline{\nu}\_t(x)- \overline{\nu}\_{j, t}(x))^2 d x)^{1/2}, j=1,2$. Algorithm 3 performs slightly better with $\mathcal{W}\_2$-distance $7.610\times 10^{-5}$ and $L^2$-distance 0.608, while Algorithm 4 has $\mathcal{W}\_2$-distance $7.771\times 10^{-5}$ and $L^2$-distance 0.6434. We have included this study in subsection C.2 in the appendix of the revised version.

In example 1, we put four shapes in the corners solely for visualization purpose, we can enlarge and center the shapes, the results would be the same. In addition, it's possible to separate the domains of the measures by following precedures in Chartrand et al. [2009] and considering the following formulation

where the support of $\nu$ is in $\Omega$, the suport of $\mu\_i$ is in $\Phi\_i$ and $\varphi\_i : \Omega \rightarrow \mathbb{R}$. The convex conjugate of $\varphi$ is defined to be $\varphi^{*}(y) = \sup\_{x \in \Omega} x^T y - \varphi(x)$ for any $y \in \Phi\_i$. It's then possible to extend the results by using Proposition 2.1 and Lemma 2.2 in Chartrand et al. [2009]. Then, $\mu\_i$ will be bounded only on $\Phi\_i$.

We note that the boundedness assumption of $\nu^t$ is minor, because this can be easily checked in practice. If this assumption is violated, a warning message can be sent to the user.

We sincerely appreciate your time, thoughtful consideration, and constructive guidance on our paper. Following your and the other referees' valuable comments, we added three additional simulation studies in the appendix of the revised version to address your questions. If we have successfully addressed some of your concerns, we would kindly ask you to consider reflecting this in your rating. Thank you again for your invaluable feedback.

Below we provide point-by-point response to your comments in italics.

1. The choice of numerical setting of the main hyperparameters (step-sizes for ascent and descent) is not discussed. What motivated to set the hyperparameters of your algorithm (in particular $\tau\_t$) as you did? Do you have guidelines for general use?

Ans: We revised our discussion of parameters $\tau, \eta$ on line 420 in the old version to ``Theorem 1 suggests that the parameters $\tau$, $\eta$ should be bounded above. Empirically, we find that the above algorithm works better with diminishing step sizes, potentially due to two reasons: (1) diminishing step sizes may reduce the effect of discrete approximations; and (2) diminishing step sizes may be more effective for nonsmooth convex functions. Since the second convex conjugate does not enforce strong convexity and smoothness, Lemma 1 no longer holds, and $\mathcal{I}\_{\nu}^{\mu}$ is only guaranteed to be concave. In addition, diminishing step size may speed up the learning process in the early stages and a inverse time decay for $\tau\_t$, i.e., $\tau\_t = 1/t$, works equally well in the simulation studies." in the revised version.

2. The burden of dimension is not discussed: in the current paper, only 2-dimensional settings are considered while Ot problems are generally of interest in larger scale. It is hard to see if the current method is scalable. The authors should mention in the introduction/related work the other common setting for Wasserstein barycenters related to generative modeling and deep learning, where the input probability distributions are not tractable and are represented by subset of samples. In the image setting, this would correspond to take one image as a sample from a high-dimensional distribution (not a distribution itself). A lot of recent work has been done in this direction recently.

Ans: Thank you for the comment. Using deep neural networks and generative models may be the only way to break the curse of dimensionality, we have added the following discussion in the revised paper on line 58 ``To break the curse of dimensionality, generative models have been investigated for the Wasserstein barycenter problem [Korotin et al., 2022]”". Feel free to let us know if we missed any relevant references.

3. The authors compare their methods to regularized approaches that enable to compute Wasserstein barycenters, but did not consider certain regularized methods that provide the exact same type of numerics (barycenter over images), see [1,2], without proper justification. How do your method compare with methods mentioned in the section ’Weaknesses’?

Ans: We note that [1] is an earlier work of Marco Cuturi who also co-authored CWB and DSB, which are both compared in the paper. DSB is an fixed support method that is more efficient than free support method; see [1] for this conclusion. We tried to add [2] as a comparison method. Unfortunatly, the code in (https://github.com/lingxiaoli94/CWB/tree/master) is not directly applicable to our examples, we thus provide the following justification on line 55 in the revised version ``Also, Li et al. [2020] propose a new dual formulation for the regularized Wasserstein barycenter problem such that discretizing the support is not needed.". Since both methods rely on entropic regularizations, the issue of having blurry output is still expected to exist.

[1] Fast Computation of Wasserstein Barycenters. Cuturi and Doucet. 2014. [2] Continuous Regularized Wasserstein Barycenters. Li et al. 2020.

4. Numerical results do not include uncertainty quantifications, which hurts the statistical significance of the results.

Ans: We added a new exmaple for 1D distributions to compare the algorithm with projection (Algorithm 3) and the algorithm with double convex conjugates (Algorithm 4). In this simulation, we repeat the experiments 300 times and report the average distance to address uncertainty quantification. Please see subsection C.2 in the appendix of the revised version for more details.

Thank you for your detailed answers and new experiments. The newly conducted experiments indeed clarify some of my concerns. However, I agree with the reviewer yaTH that the fact that your approach is tested only in the experiments with dimension $d=2$, raises questions about its applicability to high dimensions.

Still, I appreciate that the authors conduct the experiments which I suggested and increase my score to 6.

We sincerely appreciate your time, thoughtful consideration, and constructive guidance on our paper. Following your and the other referees' valuable comments, we added three additional simulation studies in the appendix of the revised version to address your questions. If we have successfully addressed some of your concerns, we would kindly ask you to consider reflecting this in your rating. Thank you again for your invaluable feedback.

Below we provide point-by-point response to your comments in italics.

1, My major concern is related to some discrepancy between the theoretical investigations of the proposed approach and its use in practical tasks. Namely, in Section 3.4, the authors prove several theoretical results regarding their method, including the important Theorem 1, which shows that the proposed algorithm allows finding stationary points of the objective functional. From the proofs in the Appendix, one can see that all theoretical results are logically connected, i.e. Lemma 2 uses the results of Lemma 1, Lemma 3 uses Lemma 2, and so on. Thus, the fulfillment of Theorem 1 also depends on the fulfillment of Lemma 1. At the same time, in Section 3.5 the authors propose to modify the algorithm for practical use in order to reduce the time complexity. This modification consists in replacing the projecting the potentials on with computing the second convex conjugate. Such a replacement entails non-fulfilment of Lemma 1, as the authors themselves write in lines 422-423. Accordingly, Theorem 1 is also no longer true, i.e. the convergence of the algorithm's solutions to ground-truth solutions is no longer theoretically justified. I will appreciate if the authors could comment on the generalization of Theorem 1 to the case where Lemma 1 is not satisfied or conduct experiments with the original version of the algorithm.

Ans: Thank you for your comments. The same discrepancy issue also exists in other works that utilize the $\mathbb{H}^1$-gradient if any convergence theory is to be developed; see the assumptions in Proposition 1 [Jacobs and Leger, 2020] and Theorem 2.8 [Jacobs et al., 2021. We believe that the satisfying empirical performance of our algorithm should not be discounted merely because of this. Please check our additional simulation result with ground truth. Without additional regularities, the generalization of the Theorem can be challenging. The main obstacle is that the map $\nu \rightarrow T_{\nu}^{\mu}$ is not Lipschitz; see [Berman, 2021] for more details regarding this.

In practice, the projection can be computed following the approach developed by Simonetto [2021], which is implemented as the Python function reg4opt.regression.convex\_regression in library reg4opt. Unfortunately, this projection algorithm requires solving a large convex quadratically constrained quadratic problem so that it takes weeks to run an experiment in 2D.

Instead, we implement the algorithm with projection (Algorithm 3) and test the method on 1D distributions. More precisely, for each repetition $t$ and $i=1,2,3$, we let $\mu\_{i,t}$ be the truncated distribution of $N(a\_i,{\sigma\_{i}}^2)$ on the domain $[0,1]$, where $a\_i \sim \text{uniform}\left[0.3,0.7\right]$ and $\sigma\_i \sim \text{uniform}\left[0,1\right]$. Let $\overline{\nu}\_t$, $\overline{\nu}\_{1, t}$, $\overline{\nu}\_{2, t}$ be the true barycenter, the barycernter computed from Algorithm 3 (with projection onto $\mathbb{F}\_{\alpha, \beta}$), and the barycernter computed from Algorithm 4 (with double convex conjugates) respectively. We repeat the experiment $T=300$ times and report two types of average distances between the groundtruth and each computed barycenters : average $\mathcal{W}\_2$-distance $\frac{1}{T} \sum\_{t=1}^T \mathcal{W}_2(\overline{\nu}\_t, \overline{\nu}\_{j, t})$ and average $L^2$-distance between densities $\frac{1}{T} \sum\_{t=1}^T (\int (\overline{\nu}\_t(x)- \overline{\nu}\_{j, t}(x))^2 \; d x)^{1/2}, j=1,2$. Algorithm 3 performs slightly better with $\mathcal{W}_2$-distance $7.610\times 10^{-5}$ and $L^2$-distance 0.608, while Algorithm 4 has $\mathcal{W}\_2$-distance $7.771\times 10^{-5}$ and $L^2$-distance 0.6434. We have included this study in subsection C.2 in the appendix of the revised version.

We sincerely appreciate your time, thoughtful consideration, and constructive guidance on our paper. Following your and the other referees' valuable comments, we added three additional simulation studies in the appendix of the revised version to address your questions. If we have successfully addressed some of your concerns, we would kindly ask you to consider reflecting this in your rating. Thank you again for your invaluable feedback.

Below we provide point-by-point response to your comments in italics.

1, One point that would deserve more discussion is the property of the nonconvex-concave problem. Unless I am mistaken (I could very well miss something), by reformulating the problem as a nonconvex problem, one loses convergence to the true solution, but only towards a stationary point. This is a bit obfuscated in the current formulation, and would deserve a more honest discussion: efficient computation is obtained by ``sacrificing" some theoretical guarantees, even with good practical performance. (If I missed something and convergence to the true barycenter is guaranteed, then it also deserved to be written!)

Ans: Thank you. In Euclidean space, the goal of minimax problem is to find the minimizer of $\Phi(x)$, where $\Phi(x) = \max_{y \in \mathcal{Y}} f(x, y)$. If $f(x, y)$ is nonconvex in $x$, then $\Phi(x)$ is nonconvex, and thus a feasible goal is to find $x$ such that $\| \nabla \Phi(x) \|_2$ is small.

For the Wasserstein barycenter problem, by definition, $\overline{\nu}$ is a Wasserstein barycenter if $id - \frac{1}{n} \sum_{i=1}^n \nabla \varphi_{\overline{\nu}}^{\mu_i}=0$, where $\varphi_{\overline{\nu}}^{\mu_i}$ is the Kantorovich potential between $\overline{\nu}$ and $\mu_i$. If $\varphi_{\overline{\nu}}^{\mu_i} \in \mathbb{F}\_{\alpha, \beta}$ for all $i$, then $\overline{\nu}$ is a stationary point of $\mathcal{F}\_{\alpha, \beta}$, i.e., Wasserstein gradient of $\mathcal{F}\_{\alpha, \beta}(\overline{\nu})$ is 0. Reversely, if we assume that the Kantorovich potential between the true barycenter and each $\mu_i$ is in $\mathbb{F}\_{\alpha, \beta}$, then Wasserstein gradient of $\mathcal{F}\_{\alpha, \beta}(\overline{\nu}')$ is 0 would mean that $\overline{\nu}'$ is a Wasserstein barycenter. We have included this discussion in Remark 1 of the revised version.