A Riemannian Exponential Augmented Lagrangian Method for Computing the Projection Robust Wasserstein Distance

Summary

This paper considered the problem of computing the projection robust Wasserstein distance between two discrete probability measures. By formulating this problem as an optimization problem over the product space of the Stiefel manifold and the Euclidean space with additional nonlinear inequality constraints, the authors proposed the so-called Riemannian exponential augmented Lagrangian method (REALM) to solve them. Further, the convergence of REALM was given. For solving the subproblems in REALM, the authors designed the so-called inexact Riemannian Barzilai-Borwein method with Sinkhorn iteration (iRBBS), where stepsizes are adaptively chosen. As the authors claimed, the complexity of iRBBS to attain an $\epsilon$-stationary point of the original PRW distance problem matches the best known iteration complexity result.

Strengths

Clearly, this paper generalized some results of the references [26,32]. To the reviewer's best understanding, the core idea is to find feasible points that satisfy the first-order necessary conditions of problem (6) or (11). To solve a three-block optimization problem, the paper transformed it as alternatively minimizing the Stiefel manifold variable $U$ and the Euclidean variables $\alpha$ and $\beta$, where the later can be solved by the well-established inexact gradient methods. Overall, this paper is well-written and mathematically solid.

Weaknesses

(1) Some mathematical details/arguments are missing (please point out if they are added in the supplementary material). For examples, Line 115: why the minimizer $(x^*,y^*)$ of the problem (9) must satisfy the relationship $y^*=\eta\log(\Vert\zeta_\eta(x^*,11^T)\Vert_1)$? Line 124, how to calculate the gradient of $\mathcal{L}_{\eta_k}(x,\pi^k)$? (2) It seems that the update of $\theta_{t}$ in Algorithm 2 is missing because $\theta_{t+1}$ appears in the inexactness criterion (19b). Please correct me if not. (3) The paper may require some ablation studies in the numerical experiments. The authors claimed REALM always outperforms the Riemannian exponential penalty approach since it could avoid too small penalty parameters in many cases (Lines 70-71). What's the range for ``too small" penalty parameters? And which cases?

Questions

1. Line 65: One ``method" in this line is redundant. 2. Line 91: There is no definition for $\nabla f(U)$. The metrics on the left and on the right side are different. 3. Line 145, Proposition 2.6: How to come out with the $x^s$? What is the motivation for the definition of $x^s$? 4. Line 252: Should be $\theta_t$ in the place of $\theta$?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

Summary

This paper reformulates the projection robust Wasserstein distance as an optimization problem over the product of the Stiefel manifolds and a subset of a Euclidean space. A Riemannian exponential augmented Lagrangian method (REALM) is proposed to solve this problem. The proposed method is empirically more stable than the existing Riemannian exponential penalty-based approach. A Riemannian BB method with Sinkhorn iteration (iRBBS) is used for the subproblem. The iteration complexity of iRBBS is given. Numerically, the proposed method outperforms the existing methods.

Strengths

(1) A method (REALM) is proposed and its global convergence is given. (2) An inexact Riemannian BB method with Sinkhorn is developed and its iteration complexity is derived. Such a result is interesting by itself. (3) From the numerical comparison, the proposed method is most efficient.

Weaknesses

The existing methods have iteration complexities for the corresponding algorithms. This paper only gives the iteration complexity for the subproblem.

Questions

(1) Can the authors give the iteration complexity for the overall algorithm? What is the main difficulty here? (2) What is the dominated computational time in the proposed algorithm and the compared algorithms?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors pointed out an important limitation in the theoretical analysis, which is the lack of a lower bound of \eta_k.

Summary

The authors first reformulate the computation of the PRW distance as an optimization problem over the Cartesian product of the Stiefel manifold and the Euclidean space with additional nonlinear inequality constraints. And then they also propose a Riemannian exponential augmented Lagrangian method (REALM) for solving the problem

Strengths

1. The authors propose a Riemannian exponential augmented Lagrangian method (REALM) method to efficiently and faithfully compute the PRW distance and establish the global convergence of REALM in the sense that any limit point of the sequence generated by the algorithm is a stationary point of the original problem 2. To efficiently solve the subproblem in REALM, the authors also propose a novel and practical algorithm, namely, the inexact Riemannian Barzilai-Borwein (BB) method with Sinkhorn iteration (iRBBS).

Weaknesses

1. The authors should provide a proof sketch for their main theoretical analysis in this paper. 2. The authors should add more experimental results to verify their theoretical results and their algorithms.

Questions

See the above setion

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

1: Your assessment is an educated guess. The submission is not in your area or the submission was difficult to understand. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

3 good

Limitations

See Weaknesses

Summary

This work proposes a new method, called REALM, to compute the projection robust Wasserstein (PRW) distance. The method REALM is an extension of the exponential augmented Lagrangian method to the Riemannian space. The convergence of REALM is established. To solve a subproblem during REALM, this work proposes an inexact Riemannian Barzilai-Borwein method with Sinkhorn iteration (iRBBS). The complexity rate of iRBBS to solve the subproblem, whose solution is an $(\epsilon_1,\epsilon_2)$-stationary point of the PRW problem, matches the existing works in the literature. This work claims the robustness in terms of the parameter-tuning in their proposed algorithms.

Strengths

This work proposes a new method to compute the PRW distance with solid theoretical guarantees.

Weaknesses

The complexity rate matches the existing rate. The idea of reformulating the PRW distance computing problem into a Riemannian optimization setting is not new, e.g. Huang et al. 2021. The claimed easier parameter-tuning part seems to need more elaboration, either from the perspective of theory, or from numerical evidences.

Questions

Hello authors, I have the following questions, 1. Generally speaking, the augmented Lagrangian method also suffers from the poor choice of penalty parameters, e.g., see Curtis et al. 2015. In your work, it claims that REALM can "potentially avoid too small penalty parameters". I was wondering if you could provide some intuition here (to explain why) or explicitly point out any efforts on the algorithmic design to have this property. 2. When you extend the exponential ALM, you claim it is a nontrivial extension because the specific problem structure encourages the specific conditions (14) and (16). Besides this, is there anything significantly nontrivial? 3. The iRBBS is discussed in Part 3. In particular, you mention the cost of updating $\alpha$ and $\beta$ is much less than that of updating $U$. This feature appears to be a valuable addition to your framework (17). Could you confirm my understanding? Is it possible to elaborate/support your proposed iRBBS, e.g., with respect to its implementation simplicity, to make your work more distinctive from existing literature.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

1: Your assessment is an educated guess. The submission is not in your area or the submission was difficult to understand. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

In the Euclidean space, some existing works consider the boundedness of the penalty parameters, e.g., see Echebest et al. 2015. While the additional (Riemannian, inequality) constraints might make situation totally different, I was wondering if those existing literature could shed light on further investigating the parameters settings.

Summary

this paper proposed a Riemannian Exponential Augmented Lagrangian Method for solving the projection robust wasserstein distance problem. the authors claimed two contributions compared with the previous works: 1) the proposed algorithm is much more stable as \eta needs to be small in previous works. 2) a Riemannian Barzilai-Borwein method was proposed to adaptively fine tune the step size. numerical experiments compared with previous works were reported.

Strengths

1. the author proposed to do multiple sinkhorn steps and one riemannian gradient step in each iteration because sinkhorn steps are much more cheaper. this idea makes sense and the author provided convergence guarantee for the proposed algorithm. 2. numerical experiments show the advantages of the proposed algorithm compared with RBCD.

Weaknesses

1. the main concern is that the author claimed that the proposed algorithm is numerically more stable than RGAS/RBCD because both of them require \eta to be very small. however, I'm not convinced by simply saying "Based on the knowledge that the exponential ALM is usually more stable than the exponential penalty approach ... ". are there any systematic study to support this claim? the proposed ALM requires the penalty parameter \eta to be exponentially decreasing. such an \eta appears in the denominator of an exponential term (e.g. eq 7). wouldn't this cause the same numerically instability issue? besides, in the proposed algorithm, the author applies the sinkhorn iteration, which will naturally introduces the numerical issues? 2. the writing of this paper needs to be further improved.

Questions

see the weakness

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

see the weakness.

reply to authors' rebuttal

thanks for the clarification. I've increased my score to 5