A Riemannian Exponential Augmented Lagrangian Method for Computing the Projection Robust Wasserstein Distance
Authors: Bo Jiang
Ya-Feng Liu
Venue: NeurIPS 2023 poster (2023)
Keywords: Barzilai-Borwein method
exponential augmented Lagrangian
inexact gradient
Stiefel manifold
Sinkhorn iteration
Wasserstein distance
TL;DR: REALM for Computing the Projection Robust Wasserstein Distance

Abstract: Projection robust Wasserstein (PRW) distance is recently proposed to efficiently mitigate the curse of dimensionality in the classical Wasserstein distance.  In this paper, by equivalently reformulating 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, we propose a Riemannian exponential augmented Lagrangian method (REALM) for solving this problem. Compared with the existing Riemannian exponential penalty-based approaches, REALM can potentially avoid too small penalty parameters and exhibit more stable numerical performance. To solve the subproblems in REALM efficiently, we design an inexact Riemannian Barzilai-Borwein method with Sinkhorn iteration  (iRBBS), which selects the stepsizes adaptively rather than tuning the stepsizes in efforts as done in the existing methods. We show that iRBBS can return an $\epsilon$-stationary point of the original PRW distance problem within  $\mathcal{O}(\epsilon^{-3})$ iterations, which matches the best known iteration complexity result. Extensive numerical results demonstrate that our proposed methods outperform the state-of-the-art solvers for computing the PRW distance.

A Riemannian Exponential Augmented Lagrangian Method for Computing the Projection Robust Wasserstein Distance
Bo Jiang
Ministry of Education Key Laboratory of NSLSCS
School of Mathematical Sciences. Nanjing Normal University Nanjing 210023, China j jiangbo@njmn jnu. edu. cn
Ya-Feng Liu
State Key Laboratory Scientific and Engineering Computing Institute of Computational Mathematics and Scientific/Engieeriin Computing Academy of Mathematics and Systems Science, Chinese Academy of Sciences Beijing 100190, China yafliu@lsec. cc. ac. cn
Abstract
Projection robust Wasserstein (PRW) distance is recently proposed to efficiently mitigate the curse dimensionality in the classical Wasserstein distance. In this paper, by equivalently reformulating the computation of the PRW distance an optimization problem over the Cartesian product of the Stiefel manifold and the Euclidean space with additional nonlinear incquality constraints, we propose Riemannian exponential augmented Lagrangian method (REALM) for solving this problem. Compared with the existing Riemannian exponential penalty- based approaches, REALM can potentially avoid too small penalty parameters and exhibit more stable numerical performance. To solve the subproblems in REALM effi- ciently, we design an inexact Riemannian Barzilai- Borwein method with Sinkhorn iteration (iRBBS), which selects the stepsizes adaptively rather than tuning the step- sizes in efforts as done in the existing methods. We show that iRBBS can return an e-stationary point of the original PRW distance problem within (ce-3) iterations, which matches the best known iteration complexity result. Extensive numerical re- sults demonstrate that our proposed methods outperform the state-of-the -hch--rt solvers for computing the PRW distance.
Introduction
The optimal transport (OT) problem has found wide applications in machine learning, representation learning, data sciences, and image sciences; see [21, 4,2, 1, 14] and the references therein for more details. However, its direct application machine learning may encounter the issue of the curse of dimensionality since the sample complexity of approximating the Wasserstein distance can grow exponentially in dimension [22, 49]. To resolve this issue, by making an important extension to the sliced Wasserstein distance [43], Paty and Cuturi [41]. Deshpande et al. [16]. and Niles- Weed and Rigollet [40] proposed to project the distributions to www-dimeniiona subspace that maximizes the
*Corresponding author.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Wasserstein distance between the projected distribution, which can reduce the sample complexity and overcome the issue of the curse of dimensionality [40, 16, 35].
In this paper, we focus on the discrete probability measure case. For fai..... Fn C Rd and {y1...., 3n C IRd define Mij = (x; (2) - y₁)² for each (i,j) E [n] x [n] with [n] ==
{ [I. Let 1 E R be the all-one vector and dz be the Dirac delta function at x. Given r (ri:...::'') E A" == {z E R 1T 1, > 0} and c (ct,...,
E A", define
two discrete probability measures yrn = >n= rida, and Un = Li=1 Gog.- For k E [d], the k- dimensional projection robust Wasserstein (PRW) distance between jmn and vn is defined as [41]
(M Un)
max min UEul rell(r,e)
(,,)(
where (+-) is the standard inner product in RTxa, C(U) E R=*x with [C(U)] = (My, UUT). U = (U E Bddk U = [x] is known as the Stiefel manifold with I3 being the k-by-k identity matrix, and [I(r, c) = { E R=x= >1= = r, π 1 = c, π 2 03. Problem (1) is a nonconcave-convex
max-min problem over the Stiefel manifold, which makes it very challenging to solve.
Related works and motivations. To compute the PRW distance, Paty and Cuturi [41] proposed two algorithms for solving the subspace robust Wasserstein distance, which is a convex relaxation of problem (1) (without the theoretical guarantee on the relaxation gap). An OT or entropy- regularzze OT subproblem with dimension n and a full or top k cigendecomposition of a d-by-d matrix needs to be solved exactly at each iteration. Very recently, Lin et al. [34] proposed Riemannian (adaptive) gradient ascent with Sinkhorn (R(A) (AS)) algorithm for solving the following entropy-regularized problem with a small regularization parameter TE:
max Pry (U), UEU
where (U) ein (TT, C(U)) - (x)), in which H(T) -by Tay log Rij is the
entropy function. They showed that R(A)GAS can return e-stationary point of PRW problem (1) within O iterations if r = (c) in (2). However, at each iteration, R(A)GAS needs to solve a regularized OT problem in relatively high precision. which results in high computational cost. To reduce the complexity of R(A)GAS, Huang et al. [28, 29] proposed a Riemannian (adaptive) block coordinate descent (R(A)BCD) algorithm for solving an equivalent "min" formulation of (2)
min BER""
(x, 11
(3)
where (..(..) is defined in (8) further ahead. By choosing = Occ) in (3), they showed that the whole iteration complexity of R(A)BCD attain an e-stationary point of PRW problem (1) reduces to which significantly improves the complexity of R(A)GAS.
However, there two main issues of R(A)BCD and R(A)GAS. First, to compute solution of problem (1) with relatively high quality, 77 in problem (2) or (3) has to be chosen small, which makes the corresponding problem ill-conditioned and may cause numerical instability in solving it. Second, the performance of the above algorithms is sensitive to the stepsizes in updating U. Hence, to achieve better performance, one has to spend some efforts tuning the stepsizes carefully. Resolving these two main issues demands some novel approaches from both theoretical and computational points of view, and this is the motivation and focus of our paper.
Contributions. In this paper, by reformulating (1) as an optimization problem defined over the Cartesian product of the Stiefel manifold and the Euclidean space with additional inequality con- straints (see problem (6) further ahead), we can resolve the above-mentioned two issues. Our main contributions are summarized as follows. See also Figure I for a summary of the related works and the main results of this paper.
(i) We propose Riemannian exponential augmented Lagrangian method (REALM) to efficiently and faithfully compute the PRW distance, in which a series of subproblems with dynamically decreasing penalty parameters and adaptively updated multiplier matrices are solved approximately. In theory, we 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; see Theorem 2.7. Numerically, REALM always outperforms the Riemannian exponential penalty approach since it could avoid too small penalty parameters in many cases.
2

---

PRW Problem (3)
intatooorry (1 with -catinnnry point (1)
exctinery of with
16 in eemennnrr peins of
It Cooaan
(oronn 2.4)
Problem (2)
Problum
Problem (11)
Compute pa (2 white "e
Compate
e-padasary
Cornagge asatonerry point ol
whtiee
poize of (1t) olyrte-))
0(a""e-1)
tirrations
licrations
(heerems 8.2%
RJAGAA (38]
RAA/BD (27]
paper)
(this paper)
Figure I: A summary of the related works and the main results of this paper.
(ii) To efficiently solve the subproblem in REALM (i.e., problem (11) further ahead), we view it as one-block optimization problem over the Stiefel manifold and propose a novel and practical algorithm, namely, the inexact Riemannian Barzilai-Borwein (BB) method with Sinkhorn iteration (iRBBS), wherein flexible number of Sinkhorn iterations is performed to compute the inexact Riemannian gradicnt. Compared with R(A)BCD, our proposed iRBBS (applied to problem (I1) with fixed n² and 7} = Oe)) can not only return a stronger e-stationary point of PRW problem (1) (compared with the definitions in [34, 28]; see Remark 2.2), within O(e iterations (see Theorem 3.2 with 61 62 = c), but also has a better numerical performance, which mainly benefits from the adaptive Riemannian BB stepsize (based on the inexact Riemannian gradient information).
Notations. For æ E R", (2iag(z) an n X n diagonal matrix with being its main diagonal. For matrix A, denote Amin min.y Aij, Amax maxij Aij, || Ei, Asyl || = maxij |Ayyl || Ei A and ||A||var Amax Amin Denote by log(-) and exp(-) the element-wise loga- rithmic and exponential operators, respectively. We use 14 and I44 to denote the nonnegative and positive orthants of R==xn respectively. Throughout this paper, we define C E RRax with Cjj 2. y/ and Va Eis Mij.
The tangent space at E 2 is Tyl {e E Bdxk UTE + ET 0}. Let TV ((U.E)
u and & E Tuul be the tangent bundle U. A smooth map Retr TV + u (U.E) Retro(g) called a retraction each curve R())= Retro (IE) satisfies R(0) U and R'(0) = g: see [9, Definition 3.47] or [2. Definition 4.1.1J. The Riemannian metric (::)) endowed on the Stiefel
manifold U is taken as the standard metric (:-) on Rd The Riemannian gradient of a smooth
function Rd + I at U E U is defined as grad f(U), which satisfies (grad
=
(VS(U), for all & E Tull where ((U) denotes the Fuclidean gradient of f at U. IF V(UU) is symmetric, we have grad f(U)= = (V = (La -UUT)
The rest of this paper is organized follows. The proposed REALM is introduced in Section 2. A practical iRBBS for solving the subproblem in REALM is proposed in Section Numerical results presented in Section 4. Finally, we draw some concluding remarks in Section 5.
2 A Riemannian Exponential ALM for Computing the PRW Distance (1)
Given fixed U e Ut consider the OT problem
((U)) s.t. #1 x c, 0. #EReXn
min
(4)
By adding a redundant constraint ||T||| = [29, 36], we derive the dual of (4) as
max oER" BER"
-(r a T + y) s.t. x(x), + + 00 v(i.j) E [m] x [n].
(5)
where X (a, 8, U) and pxx) Rxxn with p(x) aq + + (My, UUT) for each (i.j) e
[n] x [n). Note that the matrix æ in (4) can also be understood as the Lagrange multiplier corresponding
to the inequalities in (5). Therefore, the value PE va) defined in (1) is equal to the opposite of
the optimal objective value of the following optimization problem:
min MERR
T + c. 3 + s.t. p(x) + 2 0, v(..) E [m] x [n].
(6)
3

---

where M = R" X X x U. Motivated by the first-order necessary condition of problem (6) (see
Appendix A.1 for details), we define the (f), E2) -sttttioin point of problem (1) as follows.
Definition 2.1. We call (x, #) E M x [rr c) an ((1, (a)))stttionary point of PRW problem (1),
[Projt 2V2 UF v and (5, (x)) S <2. where ( E IR with 2(x), = (x))
(x() If E1 62 = O, we call such (&, #) stationary point of PRW problem (1).
Remark 2.2. Our Definition 2.1 is stronger than [28, Definition 3.1] and [34, Definition 2.7] in the sense that the
point satisfying the conditions here also satisfies all conditions therein. See Appendix A.2 for more details.
Given E R44* and 1j > O, define the function (7 TT) E RRax" with
(0(x)is
(G](x,)]= = Tisexp
define
[[((x)is + (-
(x(x,,,))= = aa a c. y + 13 Σ Tig exp
"
and
: == (x,
T n log( (x,xx))x)) 1)11))
(8)
One natural approach for solving problem (6) is the Riemannian exponential penalty approach (where the manifold constraints are kept in the subproblem), which aims to solve the penalty subproblem
min xEM.yER
L,,(X, y, 117)
For any fixed H, letting (x, = 0, we can obtain the optimal y as y (x, 111) By eliminating the variable y in (9), we thus obtain the subproblem (3) of the approach Huang et [28].
is known that the exponential augmented I Lagrangian method (ALM) is usually more stable than the exponential penalty approach; see [18, Tables 3. 3-133.3] for detailed example. More specifically, the penalty parameter in the exponential ALM can be chosen as any positive number in the convex case [47, 52] or can be bounded away from zero under some standard assumptions in the general nonlinear case [20], which is in sharp contrast to the exponential penalty approach. Based on the aforementioned knowledge, we thus extend the exponential ALM [8] the manifold case solve problem (6). Fix the current estimate of the Lagrange multiplier corresponding to the inequality constraints (6) and the penalty parameter as T² and 7), respectively. Then the subproblem the k-th iteration is given as
min Eop (x, U, m) xEM
(10)
Similar to the way for eliminating in (9), we obtain an equivalent formulation of (10):
min XEM
f,, 7
(11)
Define the matrix dyp (x, xh) E Raxx with
[pm (x, Tt les Im (x, Mas 15 (x, )11
(12)
By the chain rule, we have VaLo, (x,**) = r- On (x, xt))1 )1, V v (x, Tt) =c-==--- On (x, xx) ) 1 Vy (x, ) -2V -2V (x,+,), U, and grady f.ne (x, x²))
Projtcu (-2V ()). Let
V (x, > + a (x, Th ll.. The (E1, 62) -stattoiyy point of (11) and the connections of the approximate stationary points of problems (11) and (I) are given as follows.
e (x, n²)
grad, Ly (x, 75 ]lF and eñ (x,
Definition 2.3. We say x E M an (E1, (2)-stationary point of problem (11) (with fixed me and x²)
e (x, k) S 61 and e (x, x) S 62.
Theorem 2.4. Suppose x = (ā, B, Ú) E M with 1m (x, * is an (E1, E2) -stationny point of problem (11). Then, we have
Sa 21|C1locE2,
(13a)
(R, (x)) < (2log n + |log [l + (llāllvar + [1,11var + ||C||
(13b)
where â = Round dryy (x, )²,,, ),,,,((r,,)) ()) is feasible matrix returned by running the rounding procedure "Round" given in [3, Algorithm 2] with input on (xx x²))

---

Remark 2.5. By Theorem 2.4, we can see that, given any e >> 0, for fixed n", by choosing 7 with éL E2 e, e-stationary point with bounded (ā, B) (which can be found efficiently by iRBBS proposed in Section 3) of problem (11) can recover an e-stationary point of PRW problem (1). This may also be of independent interest for the case with fixed U. In such case, one can return feasible e-stationary point of the OT problem (4) (with fixed U) evaluated by the primal-dual gap other than the primal gap typically used in the literature, such as [3, Theorem 1].
(e))
Let x" be an initial point satisfying r r = =T and Ilsm (x", TT I1 = 1 with x² = 11". We require that x* = (a*,3*, )² the (Ex,1, Ex,2) -statioooy point of the subproblem (11) with k 2 1, satisfies the following conditions:
a c 3h, 1l5m (x* .*)|| =1, fm fm x²) S m (xx-"1,77 Ly (x", (x2,=*)) ")).
(14)
These conditions are important to establish the convergence of REALM (as shown in Appendix A.4). The following key observation to (11) shows that x satisfying the first two conditions in (14) can be easily obtained. We omit the detailed proof for brevity since they can be verified easily by noticing
that Ln (x, π = n. (x", x²) where x" is defined in Proposition 2.6. The third condition in (14) can be easily satisfied by using any descent-like method starting from the better point of x* to solve problem (11).
and x
Proposition 2.6. For any x= (a, U) E M, consider x = (a + 011,8 12 U) E . with e1 = Bc 3- + Tn log(lIsm (x, (1)) and U2 = (r c. + my log( Sy (x, a )11)) Then we have that rT (a + v11) = (B 12 and 11 (x", T ||2 = and also that L (x, n²)) =
Cy (x', x²) and VxLme (x, = VxLm (x", x*)
With such x* in hand, we compute the candidate of the next estimate x x as
* &n (x** m²")
(15)
*) Denote W E Raxn with W = min (nh k+1 )x)) The penalty parameter 7h+1 is updated
according to the progress of the complementarity violation [4], denoted by ||  If ||W*||F I
TW | |F with "/ E (0, 1), we keep M = 7e and update n # otherwise we keep
π x+1 = x² and reduce 1 via
M6+1 min Th 0 |log x= 1l00
(16)
where 7, E (0, 1) is a constant and @ + with @x > 0.
We summarize the above discussion as the complete algorithm in Algorithm 1, whose convergence is established as follows.
Theorem 2.7. Let {(x*, )} be the sequence generated by Algorithm with 61 = 62 = 6. = 0 and (x00, 70000 be a limit point of (x*,, }²)}. Then, (xx2, ) is a stationary point of PRW problem (1).
Remark 2.8. Our proposed REALM a nontrivial extension of the exponential ALM from the Euclidean case [47, 20, 51, 52] to the Riemannian case. The following two differences distinguish our proposed REALM from the existing exponential ALM (c.g., the one proposed in [20]: ()) Measure of complementarity. The measure of complementarity used in our proposed REALM is motivated by the direct use of the complementarity condition adopted in the classical (quadratic) ALM, while that used in [20] is variant of the measure for the exponential case. (ii) Conditions on global convergence. To guarantee the global convergence of the exponential ALMs, some (strong) constraint qualifications, the boundness of the iterates, and the feasibility of the limit point of the iterates generally need to be assumed; see Proposition 2. and Theorem 2. 1 in [20] for the corresponding results. In contrast, for our considered PRW distance problem (6), we can prove the boundness of the iterates generated by REALM without making the assumption and establish the global convergence of REALM without explicitly dealing with the constraint qualification assumption. This advantage is mainly due to the essential changes in the proposed REALM (compared with the existing exponential ALMs), i.c., specific conditions (14) and (16) (motivated by (13b)) on the solution of subproblems and the adopted measure of complementarity.
Moreover, it would be possible to extend the analysis in [20] to prove that the penalty parameter 7e in A Algorithm I is bounded away from zero if the Riemannian versions of the three conditions hold, including the linear independence constraint qualification, the strict complementarity condition, and the second-order sufficient condition. However, these three conditions might not be easy to check since we do not have prior knowledge of the solution.
5

---

Algorithm 1: REALM for solving problem (6).
Input: Choose <, <1, 2: > O, Yw, "n. Ve E (0, 1), 91 2 f1,$1, 2 62. Choose x = 11 x" E M with r a = c 80, I15m (x, π )ll = 1. Compute W with W = min((n p(x))ss))))))))
for k = 1,2,.... do
3
Compute an (ck,1, <x (x, -sttatonnary point x* of (11) satisfying (14); Compute h according to (15) and compute W*; if [[|] I Ca. e) (x*, x²) S (1, eg (x* x* I 62 then return x² and n=+11 = #ht1; if || S Tw | || then set x m 1 5k-1, 7/++1 ?/s. else set x 7 and
update +1 via (16);
Set 61+1,1 1 and cx+,2 4,.
Remark 2.9. We cannot establish the iteration complexity of Algorithm 1 due to the following two main difficulties: (i) characterizing the connection between the two complementarity measures ||W*|r and e == (a 7 (xx+)) with a := (o (x*, x²),, [[r; e)) at the approximate station- ary point of the subproblem: (ii) establishing the relationship between 7 #++1 and p(xx),; Thanks to Theorem 2.4. we can slightly modify Algorithm 1 to establish the iteration complexity. By modi- fying the "if" condition in Line 6 of Algorithm and leveraging the connection between the approximate stationary points of the subproblem and the original problem as proven in (13), we know that the modified Algorithm will terminate within at most O (max (max(logee E log 62 H).) iterations, where T. =- min(k Ok S Ee).
"I "| |F I Tw |
|| and e S Twe
Remark 2.10. Problem (6) can also be solved by the Riemannian ALMs based on the quadratic penalty function [37, 55]. However, the subproblems therein have four blocks of the variable, i.e., (o, 9, U), and some customized solvers are needed to solve them. Moreover, the connections between the stationary points of problem (1) and the subproblems therein remain unclear.
A Practical iRBBS for Solving Subproblem (11)
At first glance, (1I) is a three-block optimization problem and can be solved by R(A)BCD proposed by [28]. However, as stated therein, tuning the stepsizes for updating U is not easy for R(A)BCD, In sharp contrast, we understand (11) as optimization with only one variable U as follows:
min
{q(U)= == min Ly Ueu
nT") (x,
(17)
aER" "EE"
By [34, Lemma 3. we know that q()) is differentiable over Ra Here, we give a new formulation of grad (U), which can provide more insights into approximating grad (UU). Lemma 3.1. Let (i)) E argmin " f(x, x²) and xi (oi, (i, U). Then we have grad (U) = grady Lyp (Xi, x Projt (-2V U))))
Hence we could use the Riemannian gradient descent (RGD) method [2] to solve problem (11). Letting Ti > 0 be some stepsize, the main iterations of RGD are given as
U++ = Retru - grad qU")).
(18)
However, RGD (18) needs to calculate (at (t)) exactly, which can be challenging (or might be unnecessary) to do. Motivated by the well-established inexact gradient type methods for optimization in the Euclidean space [11, 26, 46, 17,39, 39, 7, 52], we propose an inexact RGD framework. Let x' (aa 3t, ,U) with (a', 8*)
(..)) wherein the inexactness level is determined by m xx) V Cn (xt, m/ > 1 (xf, = S 0, for given Or. By Lemma 3. use
e2
5¹ := g-gddy (x', 7² Projr (-2V 2V (x
approximate grad )U²). Then, we perform inexact RGD with grad q(U*) in (18) replaced by et. More specifically, given the inexactness parameter 0,41 2 and the stepsize T, 2 0, we update
x (a'+1, ge++1 (u+1) with U++1 and (a (a+ gg++) ( satisfying
Ut+ Retry (-T15")
(inexact RGD step)
(19a)
2 (x*+1, x < 0c+1-
(inexactness criterion)
(19b)

---

To make iRGD (19) practical, the first main ingredient is how to compute
such that
(19b) holds. Given U++1 E U. (0)) = a', and B(0) = 8r, for = O, o,l.... we adopt the block
coordinate descent method to update
(f+1) argmin Cye (a,, ((0, utt1 ,T),, g(++l) = argmin Eype (a B, ut ,,)). (20)
oEE"
BER=
Note that (1+11) and g((+1) admit the closed-form solutions follows:
(e++)) = (2) TiK logr + Ti log(<(()1), g(e+1) ((() - The log e + m log (( (21) where ç(0) == Sma
U++1, x²) and ((e+ == San (a 3(00) U++,, x²) (see (7) for
its definition). Note that for fixed U++1 and Th = 111, (21) reduces to the famous Sinkhorn iteration [15]. Therefore, we still name (21) as the Sinkhorn iteration. It is easy to verify that ((c(e+b)]l = ||((tt1111 = 1 (see also [28, Remark 3.1D. By (12), we have x (e+ := -m Pa (tt) ,(0),
,0,,00 u x²) = =((++)), x (8+11) == On
U++1 .Th = (E+1)
(22)
From the update of g((+1), we have (T((+)) c- 0. Therefore, to make condition (19b) hold, we stop the Sinkhorn iteration once
|| 1 r|1, S B,+1,
(23)
and set at1 = all+1), Be+1 = g(++1), Recalling the calculation of the gradient and the definition of e (x!+1,x** x²) after (12), we have
eq (x ++ ,,*) = || ((+1)) 1-rl11
Next, we choose the stepsize 7; in (19a). Since the accurate function and gradient information of q(U*)= = L. (xir,, x²) is unavailable, we cannot expect to build the linesearch condition based on q(U²) and need to find some appropriate potential function instead. Considering that L. (x', x²) is an approximation of (U²) and the approximation error is controlled by e (x*, *²),, ) it is thus desirable that some combinations of Ly (x', =²*) and ei (x", x²) will be smaller than the corresponding values at the previous iteration. Given p E [0, n1/k2)) we define the potential function as
Ex(x') = Lyp (x', ) + ple (x', ))²
(24)
and require the stepsize Tt to satisfy the following nonmonotone line search condition:
Fx((x+1) S F. 11
(7 p) (e (x"=1_ (
Tk
(25)
2
where Fj = (7Q, + (x and Q++1 = 7Q: + with a constant 7 E [0, and E6 Ex(x"), Qo = 1; see [54]. Such TO can be found by adopting the simple backtracking line search technique starting from initial guess of the stepsize 7( Owing to the excellent performance of the BB method in Riemannian optimization [50, 32, 23, 25, 31, 24], we choose the (0) initial guess T3 for / > as new Riemannian BB stepsize with safeguards:
min(max
7minn,, Tmax):
(26)
BB2 and for t 2 2, we set where Tmax Tain > 0 are preselected stepsize safeguards and BB T1
BB2 BB2 max T."h .0}} if T B192 < TT RBI and set T BB = -, otherwise. Here, T.B = min T T,
T. = I - - 1 / - U--1, -- g5----1 ---1) T BB2 = KU - Ut- _e - 11-1)/ - 1--
and 7,ew is chosen according to [30, Eq. (2.15)). In our numerical tests, we set the initial x, to be 0.05 and update x x/1 ///00 if 7,5///777 /-BBI < 2 and update X = 1.02x, otherwise.
We are ready to summarize the complete iRBBS in Algorithm 2. The overall complexity of Algo- rithm 2 to find an (c1, 62) -sttttiion point of problem (1) is in the same order as that of R(A)BCD. Theorem 3.2. By choosing = 2/2 c min[(]/(4|]C]]
/g/ (47/ V + +(]))))>>) with T = / (4/4og n 2ogoog (Ivar)), Algorithm 2 can return an --stttatiinna point of problem (1) (TT).e)) iteratiions with
||log * var + max rllvar: |log v/l and T/s- = <
T Té = mmx {(4) (6)) (())
OfO, 2 2R¹ /// - -2+ /2)) with l 2 1 and Rs |C( + 7h V, the total number of Sinkhorn iterations is (E To ee)) and the total arithmetic operation complexity is Onn2(kk + (max) + nndk- dk2 dk2)Te,,.)) ,T)
7

---

Algorithm 2: A practical iRBBS for solving problem (11).
Input: Choose Tmax > Tmin > 700 > 6,, 62 2 0, o, d1 E (0, 1). p E [0, m2/),, y E [0, 1), and (a _, 1_8 (-) E M. Set a (0) = a a, = 9-1 and perform the Sinkhorn iteration (21)
at U until (23) holds with 0g = for some Set a (l+1), (<+t)),9 ,0 = 3((+1),
for = 0, 1.... do
Compute 5 = grad; Ly (x',*");
it | S Eι and ei (x', Th) S e2 then return x²;
for 8 = 0, 1.......... do
Set U++1 Retry (-TrE") with Ti - 0² and update 0t+1 (e.g., via (27) further
ahead);
Set (1 = a' and ((0) = 3C and perform the Sinkhorn iteration (21) at Ut-1 until (23)
holds for some E; set a²+1 = 11++) and 3i+1 g(++1)
if (25) holds then break:
Remark 3.3. The basic idea of proposing Algorithm is sharply different from that of R(A)BCD de- veloped in [28]. Ours is based on the inexact RGD viewpoint, while the latter is based on the BCD approach. Such an inexact RGD viewpoint enables us to choose the stepsize adaptively via leveraging the efficient BB stepsize. Actually, tuning the best stepsize for the U U-update R(A)BCD is nontrivial. It is remarked in [28, Remark 4.1] that "the adaptive algorithms RABCD and RAGAS are also sensitive to the step size, though they are usually faster than their non-adaptive versions RBCD and RGAS." Numerical results in Section 4.1 show the higher efficiency of our iRBBS over R(A)BCD.
Remark 3.4. Although the inexact gradient type methods have been well explored in the Fuclidean case [11, 26, 46, 17, 39, 7, 52], to our best knowledge, there are little results for the Riemannian case and on how to choose the stepsizes adaptively for general nonlinear objective functions. One exception is R(A)GAS proposed in [34], which can understood as the inexact RGD method. However, it needs to compute the inexact Riemannian gradient with relatively high accuracy and essentially uses the constant-like stepsizes. In contrast, our iRBBS allows to compute the inexact Riemannian gradient with low accuracy and choose the stepsize adaptively.
Remark 3.5. It might be better to possible multiple Sinkhorn iterations rather than only one iteration as done R(A ((D)BCD in updating a and 3 from the computational point view. The cost of updating (1+1) and 3(&+1) via one Sinkhorn iteration (21) (nn)). In contrast, the cost of updating U++1 via performing RGD step (19a) is Onddk + n + dd²),, wherein the main cost is compute V ( U*, which can be done by observing V²U = X Diag( D + Y Diag ( Dia T) XTYTU Yr T.... Considering that the cost of updating and 3 is much less than that of updating U, is reasonable to update a and B multiple times and update U only once.
Experimental Results
In this section, we conduct numerical experiments six Shakespeare operas to evaluate the perfor- mance of our proposed approaches; see Dataset C. for more detailed description of the dataset. All methods are implemented in MATLAB. More numerical results can be found in Appendix D. The codes are available from https: //gtt/////htbbbb// com/bjja wwgapt/tooAllee
4.1 Comparison with R(A)BCD on Solving Subproblem (11)
Subproblem (11) with - = T1 and a relatively small 77 = is used to compute the PRW distance in [28]. We choose 7 = 0.1 as done in [28]. Since [28] has shown the superiority of R(A)BCD over R(A)GAS proposed in [34], we mainly compare our proposed iRBBS, namely, Algorithm 2, with R(A)CCD. For iRBBS and R(A)BCD, we use the same stopping conditions with
bl (motivated by (13a)) and 62 10 max{[|rllae [|]]|>]). To make the residual error more comparable, we choose
e] (xt, at)
8. = 1, 8.11 = max 9
62. t t 0.
(27)
8

---

Table I: Comparison results of R(AJBCD and iRBBS for Dataset C.I. Here, "a", "b" stand for RBCD and RABCD, respectively; "c", "d", and "e" stand for iRBBS- inf, iRBBS-0.1, and iRBBS-0, respectively. For RBCD, RABCD, and iRBBS-inf, nGrad is equal to nSk.
P%
nGrad/nnS
time
data
H5/H
0.049044
0.049074
0,049074
0.049U74
0,049074
871
604
320
64/1296
61/3218
9.6
6.7
3.4
1.2
1.8
HSUC
0.059576
0.059576
0,059576
0.059576
0.059576
904
497
193
63/1571
60/3477
5.3
2.9
1.2
0.6
0.7
HSMV
0.062515
0.062515
0,062515
0.062515
0.062515
800
444
118
63/1373
61/2812
5.4
2.9
0.8
0.6
0.8
H5/O
0,049973
0,049973
0,049973
0.049973
0,049973
397
129
43/1084
45/2410
4.9
3.2
1.0
0.5
H5/RJ
0.180227
0.180226
0.180227
0.180227
0.180227
474
440
319
73/1670
69/341
3.8
3.5
2.4
0,8
1.U
H/IC
0.050081
0,050081
0,050081
0.057156
0,057156
637
523
427
41/1049
42/2020
4.4
3.6
2.9
0.4
H/MV
0.0384166
0,038416
0.038416
0,038415
0,038416
2794
1661
440
117/2045
125/5129
22.9
13.6
3.4
1.2
1.6
INO
0.0140422
0M014022
0,014042
0.014042
0,014042
796
503
212
49/1471
53/3093
7.7
4.7
2.0
0.8
1.2
HRR
0,189475
0,199475
0.189275
0.1894755
0.189475
360
441
265
46933
47/1763
3.4
4.2
2.3
0.6
0.8
JOMV
0.013257
0.013257
0.013256
0.013257
0.0132577
1816
975
216
66/1781
65/3742
7.2
3.8
0.9
0.4
0.6
JCJO
0.009835
0,009K35
0.009835
0,009835
0.009835
1160
45/2454
6.5
3.5
0.5
0.4
JCRJ
0.110336
0.110336
0,110336
0.1103336
0.110336
340
332
235
44/1016
53/2536
1.5
1.5
1.1
0.3
0.5
MVIO
0,010090
0.010090
0,0100990
0.010090
0.010090
2222
1396
154
90/2448
88/5186
13.2
8.2
0.9
0.8
1.1
MV/R
0,168293
0.168293
0.168293
0.168293
0.168293
842
763
632
85/1688
95/4805
5.0
4.5
3.8
0.7
1.1
O/RI
0,125614
0.125614
0.125614
0.124085
0.124085
376
372
210
55/1070
62/2749
2.5
2.4
1.4
0.5
0.8
AVG
0.075387
0,075387
0,075387
0.075757
0,075757
1000
674
265
63/1460
65/3200
6.89
4.62
1,86
0.66
0.90
where the parameter is a preselected constant in the proposed iRBBS algorithm. We consider several different values of f and denote the corresponding version of iRBBS as iRBBS-0. When we have 02+1 = e2 for all 2 O, which means that we calculate grad q(U*) almost exactly; when 0 = +oo, the Sinkhorn iteration always stops in one iteration. For numerical tests in this subsection, we can always perform an equivalent formulation of Sinkhorn iteration (21) as [15, 6]
u (E+
(A)( ((++)) c cA (c+1))
(28)
where @ denotes elementwise division, ( = exp(-a /m.). v(*) = exp(- /nk)), and A E Raxa with A, = nH exp(- UUT) /m.). More settings and details on the implementation can be found in Appendix C.
For each instance of Dataset C.I, we randomly generate 20 initial points and report the average performance in Table 1. The term "nGrad" means the total number of calculating grady Ln2 (x', r²), "nSk" means the total number of Sinkhorn iterations, and "time" represents the running time in seconds evaluated by "tic-toc" commands. Moreover, since we aim to compute the PRW distance, with U returned by some method in hand, we invoke Gurobi 9.5/Mosek 9.3 to compute a more accurate PRW distance, denoted as P7 = minzellrre)) (T, C(U))) Note that a larger P7 means a higher solution quality of the corresponding U*, the last "AVG" line, we summarize the averaged performance.
From Table 1, we have the following observations: (i) as for the solution quality. iRBBS-0.1 performs the best among all methods; (ii) among iRBBS methods, iRBBS-0.1 1 is about 2.8x faster than iRBBS- inf and about 1.4x faster than iRBBS-O: (iii) compared with R(ABCDD, iRBBS-0.1 can always take significantly less nGaa and may take bit more nSk. This makes it about 7x faster than RABCD and about 10.4x faster than RBCD, The results show that iRBBS generally performs much better than BCA)CCD, More importantly, our methods adopt the adaptive stepsize without needing to tune the best stepsize as done in R(A)CDD.
4.2 Comparison on Computing the PRW Distance (1)
In this subsection, we present numerical results illustrate the effectiveness and efficiency of our proposed REALM, namely, Algorithm 1. We choose 6 = =]||C']lso C2 = 10-6 maxf]|r]los. ||cloe }. & 10 and 7/2 = 20. To avoid possible numerical instability, we restrict the maximum number of updating π² as 8. To prevent too small 7%. in our implementation, we update 7 = min [Ti,7k: } other than using (16), where 7/min is a preset positive number. We set 7, = 7c = 0.25. Moreover, we denote by REALM- 'w) the REALM with particular parameters 7mmin and 7w. Note that choosing Tw = means that we adopt a continuation technique to solve (11) with x = nr and T7k- Timin which is generally better than directly solving single (11) with

---

Table 2: Average results of REALM for Dataset C.I, "a" and "b" stand for REALM-(0.0035, 0) and REALM-(0.007, 0.9), respectively.
nGrad
time
iter
data
HSUC
0,09270
0.10985
1081
675
2380M12267
9786/U
84.1
8.6
8.0/15.0
H/MV
0.06378
0.06424
1116
727
3864/13382
30997/0
175.7
14.9
0.0/8.0
8.0/15.0
H/RI
0,21706
0.22607
547
788
309 0091/4062
11747/0
63.3
15.6
0.0/8.0
8.0/15.0
JC/V
0,062700
0.06235
1324
711
12875/9727
17073/0
55.3
5.9
0.08800
8.0/15.0
IC/O
0.04221
0.04277
1650
1100
4997/26612
35031/4500
204.7
46.7
0.0/800
8.0/15.0
MV/O
0.04181
0.03661
822
6101/5289
31651/0
57.0
13.3
0.08.00
8.0/15.0
AVG
0.11338
0.11461
787
7824/7364
20261/300
78.0
0.0/8.0
8.0/15.0
Th Miin- Subproblem (11) in REALM is solved by iRBBS. If max ||c|]>>]] 500m) or ||C(U*) 7/h log 1 2 900ms, we set 0 = 10 and perform the Sinkhorn iteration (21); otherwise, we set f = 0.1 and perform an equivalent formulation of Sinkhorn iteration (28).
The results over 20 runs on Dataset C.1 are reported in Table 2. In this table, the terms "nSkiog" " and "nSkoxp" " mean the total numbers of Sinkhorn iterations (21) and (28), respectively, the pair "ky/k" in the column "iter" means that the corresponding algorithm stops at the k-iteration and updates the multiplier matrix 1 times. To save space, we only report instances where one method can return the value "Ph"" larger than 1.005 times of the smaller one of the two P7 values returned by the two methods. The better "P?" is marked in bold. Besides, the average performance over all 15 instances also kept in the "AVG" line.
From Table 2, we can observe that REALM-(0000%) 0.9) can not only return better solutions than (0.0MM 0) but also is about 5.2x faster. On average, REALM-(0.007, 0.9) updates the multiplier matrix 8 times in 15 total iterations, which shows that updating the multiplier matrix does help. The reasons why REALM with updating the multiplier matrix outperforms REALM without updating the multiplier matrix in terms of solution quality and speed are as follows. First, updating the multiplier matrix in REALM can keep the solution quality even using a larger TJk- Second, solving the subproblem with larger 77 is always easier, which enables that REALM-
0.9) computes
less grady Ly (x', T²) and performs less Sinkhorn iterations (21) which involves computing the log-sum-exp function log exp(ax / = Tmax 1% + log E, exp((2; Fmax))/// /(k3))111) for small Vik-
5 Concluding Remarks
In this paper, we considered the computation of the PRW distance. By reformulating this problem as an optimization problem over the Cartesian product of the Stiefel manifold and the Euclidean space with additional nonlinear inequality constraints, we proposed method called REALM. The convergence of REALM was also established. To solve the subproblem in REALM efficiently, we developed a practical iRBBS method with convergence and iteration complexity guarantees, wherein the Riemannian BB stepsize (based on the inexact Riemannian gradient information) and Sinkhorn iterations are employed. The complexity of iRBBS to attain an e-stationary point of the original PRW distance problem matches the best known iteration complexity result. Numerical results showed that, compared with the state-of-the-ale art methods, our proposed REALM and iRBBS methods have advantages in solution quality and speed.
Moreover, our proposed REALM and iRBBS can also be extended to solve some important mini- max problems over the Riemannian manifolds arising from machine learning, such as the fair PCA problem [45] and the projection robust Wasserstein barycenters [27], etc.
Lastly, limitation of our work that we did not establish the positive lower bound of m in REALM. despite the fact that REALM performs well in practice and can avoid too small in many cases. We shall investigate the conditions under which it is possible to establish lower bound of n in REALM. This can be achivvaale by extending the analysis and conditions in [20] to the Riemannian case.
10

---

Acknowledgments The work of Bo Jiang was supported in part by the National Natural Science Foundation of China (NSFC) under Grant 11971239, Grant 12371314. and in part by the Natural Sci- ence Foundation of the Higher Education Institutions of Jiangsu Province under Grant 21 1IKJA100002 10002. The work of Ya-Feng Liu was supported in part by NSFC under Grant 11991020, Grant 1 11991021, and Grant 12288201.
References
[1] T. Abrishami. N. Guillen, P. Rule, Z. Schutzman, J. Solomon, T. Weighill, and S. Wu. Geometry of graph partitions via optimal transport. SIAM Journal on Scientific Computing, 42(5): A3366, 2020. (Cited page 1.)
[2] P.-A. Absil, R. Mahony, and R. Sepulchre. Optimization Algorithms on Matrix Manifolds. Princeton University Press, Princeton, 2008. (Cited on pages and 6.)
[3] Altschuler, Niles-Wees and Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration. Advances in Neural Information Processing Systems, 30:1961-1971, 2017. (Cited on pages 4, 5, and 21.)
[4] R. Andreani, G. Birgin, J, M. Martínez, and M. Schuverdt, On augmented Lagrangian methods with general lower-level constraints. SIAM Journal on Optimization, 18(4):1286-113 1309, 2008. (Cited on page 5.)
[5] M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, pages 214-223, PMLR, 2017. (Cited on page 1.)
[6] J.-D. Benamou, G. Carlier, M. Cuturi, L. Nenna, and G. Peyré. Iterative Bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2)A11111 A A1138, 2015. (Cited on page 9.)
[7] A. S. Berahas, L. Cao, and K. Scheinberg. Global convergence rate analysis of a generic line search algorithm with noise. SIAM Journal on Optimization. 31(2):1489-155 2021. (Cited on pages and 8.)
[8] D. P. Bertsekas. Constrained Optimization and Lagrange Multiplier Methods. Academic Press, 1982. (Cited on page 4.)
[9] N. Boumal. An Introduction to Optimization on Smooth Manifolds. Cambridge University Press,
2023. (Cited on page 3.)
[10] N. Boumal, P.-A. Absil, and C. Cartis. Global rates of convergence for nonconvex optimization on manifolds. IMA Journal of Numerical Analysis, 39(1): 2019. (Cited page 19.)
[11] R. G. Carter. On the global convergence of trust region algorithms using inexact gradient information. SIAM Journal on Numerical Analysis, 28(1)2251-265, 1991. (Cited on pages
and 8.)
[12] A. Chambolle and J. P. Contreras. Accelerated Bregman primal-dual methods applied to optimal transport and Wasserstein barycenter problems. SIAM Journal on Mathematics of Data Science, 1(4):1369-1395, 2022. (Cited on page 14.)
[13] X. Chen, L. Guo, Z. Lu, and J. J. Ye. An augmented Lagrangian method for non- non-Lipschitz nonconvex programming. SIAM Journal on Numerical Analysis, 55(1):168-193, 2017. (Cited on page 15.)
[14] A. Cherian and S. Aeron. Representation learning via adversariiallycoomttssvve optimal transport. In International Conference on Machine Learning, pages 1820-1830. PMLR, 2020. (Cited on
page 1.)
[15] M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems, 26:2292-223 220000 2013. (Cited on pages 7 and 9.)
[16] I. Deshpande, Y.-T. Hu, R. Sun, A. Pyrros, N. Siddiqui, S. Koyejo, Z. Zhao, D. Forsyth, and A. G. Schwing. Max-sliced Wasserstein distance and use for GANs. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10648-10656, 2019. (Cited on pages and 2.)
[17] O. Devolder, F. Glineur, and Y. Nesterov. First-order methods of smooth convex optimization with inexact oracle. Mathematical PraetticclPoooraaae Programming, 146(1):37-75, 2014. (Cited on pages 6 and 8.)
11

---

[18] J.PP. Dussault. Augmented non con-quadratic penalty algorithms. Mathematical programming, 99 (3):467-486, 2004. (Cited on page 4.)
[19] Dvurechensky, A. Gasnikov, and A. Kroshnin. Computational optimal transport: Complex- ity by accelerated gradient descent is better than by Sinkhorn's 's algorithm. In International Conference on Machine Learning, pages 1367-1376. PMLR, 2018. (Cited on page 19.)
[20] N. Echebest, M. D. Sánchez, and M. Schuverdt. Convergence results of an augmented Lagrangian method using the exponential penalty function. Journal of Optimization Theory and Applications, (68(1)922-008, 2016. (Cited on pages 4, 5, and 10.)
[21] R. Flamary, N. Courty, D. Tuia, and A. Rakotomamoniy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 39:1853-1865,
2016. (Cited on page 1.)
[22] N. Fournier and A. Guillin. On the rate of convergence in Wasserstein distance of the empirical measure. Probability Theory and Related Fields, 162(3):707-738, 2015. (Cited on page 1.)
[23] B. Gao, X. Liu, X. Chen, and Y. Yuan. A new first-order algorithmic framework for optimization problems with orthogonality constraints. SIAM Journal Optimization, 28(1):302 332, 2018. (Cited on page 7.)
[24] B. Gao, X. Liu, and Y. Yuan. Parallelizable algorithms for optimization problems with orthogo- nality constraints. SIAM Journal on Scientific Computing, 41(3):A1949-A1I 2019. (Cited on page 7.)
[25] J. Hu, A. Milzarek, Z. Wen, and Y. Yuan. Adaptive quadratically regularized Newton method for Riemannian optimization. SIAM Journal on Matrix Analysis and Applications, 39(3):
1181-1207, 2018. (Cited on page 7.)
[26] Y.-Q. Hu and Y.-H. Dai. Inexact Barzilai- Borieinn method for saddle point problems. Numerical Linear Algebra with Applications, 14(4):299-311 2007. (Cited pages 6 and 8.)
[27] M. Huang, S. Ma, and L. Lai. Projection robust Wasserstein barycenters. In International Conference on Machine Learning, pages 4456-4465. PMLR, 2021. (Cited on page 10.)
[28] M. Huang, S. Ma, and L. Lai. A Riemannian block coordinate descent method for computing the projection robust W Wasserstein distance. In International Conference on Machine Learning. pages 4446 4455. PMLR, 2021. (Cited on pages 2, 3, 4, 6, 7, 8, 14, 20, 21, and 22.)
[29] M. Huang, S. Ma, and L. Lai. A Riemannian block coordinate descent method for computing the projection robust Wasserstein distance. arXiv:20/2.05199 v.,, 2021. (Cited on pages 2, 3, 17, 18, and 19.)
[30] Y.-K. Huang, Y.-H. Dai, and X.-W. Liu. Equipping the dimensional quadratic termination property. SIAM Journal on Optimization, 31(4):3068330 -0096 2021. (Cited on page 7.)
method with the two
[31] B. lannazzo and M. Porcelli. The Riemannian Barzilai- Borwein method with nonmonotone line search and the matrix geometric mean computation. IMA Journal of Numerical Analysis,
38(1):495-45 517, 2018. (Cited on page 7.)
[32] B. Jiang and Y.H. Dai. A framework of constraint preserving update schemes for optimization on Stiefel manifold. Mathematical Programming, 153(2):535-575, 2015. (Cited on page 7.)
[33] S. Kullback. A lower bound for discrimination information in terms of variation (corresp.). IEEE Transactions on Information Theory, 1361):1261127, 1967. (Cited on page 18.)
[34] T. Lin, C. Fan, N. Ho, M. Cuturi, and M. Jordan. Projection robust Wasserstein distance and Riemannian optimization. Advances in Neural Information Processing Systems, 33:9383-
2020. (Cited on pages 2, 3, 8, 14, 17, 21, and 22.)
[35] Lin, Z. Zheng, E. Chen, M. Cuturi, and M. I. Jordan. On projection robust optimal transport: Sample complexity and model misspecification. In International Conference on Artificial Intelligence and Statistics, pages (62-270. PMLR, 2021. (Cited page 2.)
[36] T. Lin, N. Ho, and M. I. Jordan. On the efficiency of entropic regularized algorithms for optimal transport. Journal of Machine Learning Research, 23(137): 2022. (Cited on pages 3 and 19.)
[37] C. Liu and N. Boumal. Simple algorithms for optimization on Riemannian manifolds with constraints. Applied Mathematics & Optimization, 82:949-981, 2020. (Cited on page 6.)
12

---

[38] H. Liu, A. M.-C. So, and W. Wu. Quadratic optimization with orthogonality constraint: Explicit Lojasiewiez exponent and linear convergence of retraction- based line-search and stochastic variance-reduced gradient methods. Mathematical Programming, 178(1):215-2662 2019. (Cited
on page 19.)
[39] Y.-F. Liu, X. Liu, and S. Ma. On the nonergodic convergence rate of inexact augmented La- grangian framework for composite convex programming. Mathematics of Operations Research, 44(2):632-650, 2019. (Cited on pages and 8.)
[40] J. Niles Weed and P. Rigollet. Estimation of Wasserstein distances in the spiked transport model. Bernoulli, 28(4): 28(4):
2022. (Cited on pages and 2.)
[41] F.-P. Paty and M. Cuturi. Subspace robust Wasserstein distances. In International Conference on Machine Learning, pages 5072-5081. PMLR, 2019. (Cited on pages 2, and 22.)
[42] G. Peyré and M. Cuturi. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 11(5-6):355-607 2019. (Cited on page 1.)
[43] J. Rabin, G. Peyré, J. Delon, and M. Bernot. Wasserstein barycenter and its application to texture mixing. In International Conference on Scale Space and Variational Methods in Computer Vision, pages 435-446. Springer, 2011. (Cited on page 1.)
[44] E. W. Sachs and S. M. Sachs. Nonmonotone line searches for optimization algorithms. Control and Cybernetics, 4054):1059-1075, 2011. (Cited on page 20.)
[45] S. Samadi, U. Tamtipongpipat, H. Morgenstern, M. Singh, and S. Vempala. The price of fair PCA: One extra dimension. Advances in Neural Information Processing Systems, 31:
10999-11010, 2018. (Cited on page 10.)
[46] M. Schmidt, N. Roux, and F. Bach. Convergence rates of inexact proximal-gradienn methods for convex optimization. Advances in Neural Information Processing Systems, 24:1458-1466,
2011. (Cited on pages and 8.)
[47] P. Tseng and D. P. Bertsekas. On the convergence of the exponential multiplier method for convex programming. Mathematical Programming, 60(1):1-19, 1993. (Cited on pages 4 and 5.)
[48] J.-P. Vial. Strong and weak convexity of sets and functions. Mathematics of Operations Research, 8(2):231-22) 13599, 1983. (Cited on page 14.)
[49] J. Weed and F. Bach. Sharp asymptotic and finite-sample rates of convergence of empirical measures in Wasserstein distance. Bernoulli, 25(4A):2620-2648, 2019. (Cited on page 1.)
[50] Z. Wen and W. Yin. A feasible method for optimization with orthogonality constraints. Mathe- matical Programming, 142(1):397-434 2013. (Cited on page 7.)
[51] Y. Xie, X. Wang. R. Wang, and Zha. A fast proximal point method for computing exact Wasserstein distance. In Uncertainty in Artificial Intelligence, pages 433-453. PMLR, 2020.
(Cited on page 5.)
[52] L. Yang and K.-C. Toh. Bregman proximal point algorithm revisited: A new inexact version and its variant. SIAM Journal on Optimization, 32(3):1223-1544 2022. (Cited on pages 4, 5, 6, and 8.)
[53] W. H. Yang. L.-H. Zhang, and R. Song. Optimality conditions for the nonlinear programming problems on Riemannian manifolds. Pacific Journal of Optimization, 10(2):415 434, 2014. (Cited on page 14.)
[54] H. Zhang and W. W. Hager. A nonmonotone line search technique and its application to unconstrained optimization. SIAM Journal on Optimization, 14(4):1043-1066, 2004. (Cited on pages and 20.)
[55] Y. Zhou, C. Bao, C. Ding, and J. Zhu. A semismooth Newton based augmented Lagrangian method for nonsmooth optimization on matrix manifolds. Mathematical Programming. 201: 1 1-61, 2023. (Cited on page 6.)
13

---

The appendices are organized as follows. In Appendix A, we present more details on our proposed REALM in Section 2. Some proofs in Section 3 are given in Appendix B. Some details on the dataset and implementation in Section 4 are shown in Appendix C. Finally, we present additional numerical results to evaluate the performance of our proposed methods compute the PRW distance (1) in Appendix D.
A More Details on Our Proposed REALM in Section 2
In this section, we provide more materials on REALM in Section 2, including the first-order necessary conditions of our reformulation (6), relations between different (E1, -t)-stoninnary points of the PRW distance problem (I), the connections of the approximate stationary points of problems (11) and (1), and the convergence analysis of Algorithm 1.
A.1 First-Order Necessary Conditions of Problem (6)
Lemma A.1 (First-order necessary conditions of problem (6)). Given & E M and y E R, suppose that (x, D) is local minimizer of problem (6), then (x, (X))) is stationary point of problem (6), namely, there exists E [T(,, e) such that
Projtgu (-2V0)) =
(T, Z(x) O,
(29)
where ((x) E R2x is defined as
((x);j = (xx) (xx)
(30)
Proof. Since (x. f) is a local minimizer of problem (6). there must hold y = - ( min
max:)(-p(x)) Moreover, such is also local minimizer of problem minxe h(x) with h(x) = r + 3 mix) For fixed and 88 it is easy
that -x((); + ||C || is
convex with respect to U since || ||2 = Cij I ||C|| We thus know that the function
(x)) + ||C | max { (-x(x); + +|C [1]
convex with respect to U, which means that the function - -x(x) min and thus fxx) is -ee-kally convex with respect to U [48, Proposition 4.3]. Let A(x) = (1.5) E [n] x [n] (x))) yo ( (
By [48, Proposition 4.6), we have
(x( (r, e, 0) + conv ((-e,-E;; -MM; (i,j) E A(x))
where e, is the i-th standard unit vector in RR. Moreover, by [53, Theorem 4. [53, Theorem 5.1]. and T&M = R" x R² Tgll, there must hold that 0 E Oh(i). Putting all the above things together shows that there exists म R Bax" with क. for all (i.j) A(x) and E E0.))A न. = 1 such that न E II(r, c) and (29) hold. The proof is completed.
A.2 Explanations on Remark 2.2
From [28, Section B], we know that the (c1, 62) -ttatioiny point in [28, Definition 3.1] I] is stronger than that in [34, Definition 2.7]. To show that Definition 2.1 of the (21, -t)-stioonna point is stronger than both definitions, we only need to verify that (म, (x)) S e2 can imply
(T, ((O))
min (T, C(Ū)) S E2,
xell(ne))
which is clear since (7, C(Ū) min (m, C(U)) < (ā, Z()). This inequality comes from the fact that problem (5) is dual formulation of min (n,, C(U)) for fixed U and the strong duality theorem.
A.3 Proof of Theorem 2.4
Proof. We first show that (13a) is true. First, by [12, Lemma 3.2], we have Ω E I[(r, c) and
|| om (x, Il1 < 1m (x, 71 r/l1 + 1l0m (x, x² ) c/l1 = (xxx*) n E2. (31)
14

---

Second, by the triangular inequality and the non-expansive property of the projection operator, we have
[|Proitua (Ai)lF (P))iTuu(A)) Projt u (A2)llF+ (Proiteu(Ai))ll
||A11 A2l1F F + P (A1)llF,
where AI, A2 E Bdak Hence, we have
||rojrade (VVVV)) & 2 - Von (8,4t) + t1 52 Ven, (x,**)]FF + 61
(32)
where the first inequality uses [|Projt (-2V 2V (xxx)U))l I 61 and the second inequality is due the fact that ||AUÍ|F I ||A||F for any matrix A Ruxd Moreover, observe that
||V Vong (R, |F
Σ (fts one (8, (x, S ||C|| - -r (X, nT )[1.
Combining the above assertions with (31) and (32) yields (13a).
Next, we show that (13b) is also true. By the Cauchy-Schwarc inequality and (31), we have
(t, (
(At 2n (x, ) ((x)) + (0 (3z, n _(x))
S [Z(())l ( (x, x ))), ((x))
< (l|@]| var + ||:3| /var - ||C )2z + { (0n, (X, x *²),, 2(x))
(33)
where the second inequality uses (30), the definition of (xx) in (5), and
(Msz, ) (C(U)] I Cyj S ||0|| VU u.
(34)
The remaining is bound (0m (X. )² ((x). By 1m (8, *1) 1, (7), and (12), we have cx()),s (log nrJ - log 2n (x, n ),g) and x(xx)min m (log x )min- Again with (30), we have
(dm (8,xx x²) (xx) (o (I, p(()) (3(3)
$ mh (log nh 2n (x, n )) + Hmm. (xx, n x²)) - (log n )
which, together with H (om, (2, T )) - 2 log n and (log Th on (x, x²)) I ()og**)mas
n5h maxx, further
implies (om (x, n²,, ((x) < (2 log + ||log n* [lvar) This, together with (33), yields (13b). The proof is completed.
A.4 Proof of Theorem 2.7
We first use the requirements in (14) to show that T < <T for all k 2 0.
With the first requirements in (14), have L.p (x* ,²*)== Ta
& By the fact ll5m, (x", x >>)l11 1 with n² 11 we know from (7) that (x(x)); 2 0. Due to the update rules of x² and m7s know that 0 ntj I 1 and 0 < mp m71 for all k 2 1. Therefore, have from (7) that
Is (x", m²))ss = mij exp (- x(xx))s)
S exp (- [x(x))s) = =xmm (x",'=")) ))sss
M7
171
which further implies 115ma (x", m I II5m (x", rT )ll = 1. With rT TT, we know L (x", m²) < aa. By the third requirement in (14), we further have
r STa vk 0.
(35)
Besides, since x is an (ex, -tttttaati point of (11), we have
|rojjr (~ U ll < Ek,1, ||r- ---=カ+1 * 1|] ||c - (=*+1)11|11 Ek.2-
(36)
We next consider two cases, which is motivated by the proofs of Theorem 3.1 in [13].
Case (i). 7e is bounded below by some threshold value 17 > 0. Due to the update rule of 7]+1, we that (16) is invoked only finite times. Besides, due to 0 < a/j 1, without loss of generality,
15

---

we assume 7 = 17 for all k 2 2 and lim my m... Meanwhile, due to the update rule of n²,
we have m #j for l > 2 and
+1 *x(x*) (6))
+ 0 A (i,j) [n] x [n].
(37)
We now show that {x*} has at least a limit point. By (37), we must have (xxx)(; - 0 for each (i,j) & A(=*) == ((..)) E [n] x [n] Tiy 0} and there must exist Ki such that for all k 2 K1, there holds ((x*"x ky 2 -1 for each (i,j) E A(x Recalling the definition of (x(x) in (5), then we conclude that there exists K2 > 0 such that for all & 2 K2, there holds that a + 3 (Mi, (U()))
-1 A (i,j) E [n] x [n].
Multiplying both sides of the above assertions by Cj and then summing the obtained inequality from j = 1 to n, we have
a 2 2-11 21-1 ( (UMM,,UUUUU (U())) J -1 || - -Taa
(38)
where the second inequality is due to r rTa c. BB. (35), and (34). Combining (35) and (38), further have
-1 -1- - - rT a r. ((1 + |C + +r + ao)).
Similarly, can establish similar bound for each B Recalling that U is in compact set, it thus safe to say that the sequence [x*} has at least one limit point, denoted by x li x* Again with limg--oo n.j = 2000 we have from (37) that min(g=3 ((xx)x(xxx)))))))))>>>)) );y } = O, which, together with T.j 2 0 and the fact
0 0 (0xx)()00)
=- Tij xx )ij (1xx))mii
15
0. This shows that (x 0) is feasi- ble. Moreover, letting E IC go to infinity in (36), we further have T E [[ c) and [|ProJT (-2V, - problem (6) and (xx2, >)), namely, (xx², #T)), ), is stationary point of problem (1).
further implies =x(x0) min
0 and (TT20 [(x"))))
( || = O. From Lemma A.1, we know that (xx 0) is stationary point of
Case (ii). The sequence {m& is not bounded below by any positive number, namely, lime 7is- 0. By the updating rule, we know that m) is updated via (16) infinitely many times. Hence, there must exist k 122 < ... such that 17,, ->
oc and
The min (iy -1 ek. 1/l log n lloo with ke S s < kets and x = h ky exp(-p((x' )/////////) S 1, which, together with By (7) and the second assertion in (14), we have
0 F. < 1 and (16), implies
p(x*)(;) 2 -Tkk xfpl -T,, |loo th -e- 100 -kc VEZ1. 1.
(39)
Using the same arguments as in the proof in Case (i), we show that fxx) is bounded over (k1, k2.... and thus (x*) has at least limit point. Without loss of generality, assume lim xh* = x By (39) and ex + 0 in (16), we know that (xx*) 2 0 and thus that (xx,0) is feasible problem (6).
Due to the compactness of # in (15), there must exist a subset IC S. {k1, kg....) such that limbex,k->00 #t) = 000, Recalling (15), (7), and the second requirement in (14), we have =k+1 TH exp(-p((*))ss/m) for each & 2 1. Let A(x*)) {ii,j) E [n]] x [r]] sp(x")); = 0). ) * 0. Otherwise, for every (i.j) we have p(x*)(; >> 0 and thus
T.j
We claim A(x*)
exp (- (x(()))) = O,
5 lim kck,k-to
=k+1 < lim 1i
(40)
kER
T
where the equality uses lime-soo xte x K. E (ki, .2,-..), .,, and limg-zoe 7k 0. This makes
contradiction with || = 1. Moreover, (40) also further implies that (7 [xx*)) O. Finally, with (36), we further have || [|Projt (-2V U= = 0 and xx E [I(r, c). Putting the above things together, we know that that (xx0, 0) a stationary point of (6) and (xx0, 700) a stationary point of problem (1). The proof is completed.
16

---

B Proofs in Section
B.1 Proof of Lemma 3.1
For notational simplicity, we define
h(T, U) = (C(U) - 77s log logx*,n) - 17s (T).
(41)
It is easy to see that
max EE BER"
- -c (x, is the dual formulation of
min xell(n.))
h(T, U).
(42)
Therefore, we can write
(0)
min 1( U)) rel((;e))
(43)
Since the entropy function is 1-strongly convex with respect to the normr over the probability simplex, the minimization problem in (43) has a unique solution oे. argmin h(T, U). Using [34, Lemma B.1], we thus have
grad (UU) = ( U).
(44)
By (5), (7), (12), and recalling [C(U)]i (Mej. TUT), we have
(C(U) + tps log On. (xi, x5) vn log m -labli [8tls - T/ log(lISma (x T )l1)) (45) By the optimality of où and fÿ, we know that e (x), xh) = eit (xt,; x²) and thus om (Xy, r²)1 r and on (x7, T
With (45), (41), and the definition of L. (:--) in
(11), by some calculations, we have
h(dm (xy, T5 )))) -r on c. Bg mh6 log(llsm (xG, x Il1) Lor, (xy, m
By the dual formulation (42), have -Ly, (xū, x²) S h(mū U). Hence, know that h(dy, (xy, m²), U) h(zū, U), which, together with the optimality and uniqueness of qe. yields on (xi, xh) M Finally, using (44), we can complete the proof.
B.2 Proof of Theorem 3.2
We first establish the iteration complexity result of iRBBS, which based on the following nice properties of Sinkhorn iteration (21). Note that the sufficient decrease property in Lemma B.2 is quite different from that in [29, Lemmas 4.10 & 4. 4.1 since we can establish the non-increasing property of the feasibility violation of =()) in Lemma B.1.
Lemma B.1. Let {(al, {8)} be the sequence generated by (21) and consider (x(0) and {x( (f((-b) 1)
in (22). Then we have || (0) - r|lı + [1(= - ell1 2 (1( - c|l1 2 (1) - r|11 and
for each 2>
n 1 - r|11 2 [1( (f+ (+)))) cll 2 (++1)] -
(46)
Proof. By the optimality of (f+1) in (20) and the first equation in (22), we have n 1(+4)1-r=0
--r
for & 2 0. Therefore, have
- (e+1)
r|]i
|| -- S (++ ((+11)
(((+12111
g(++)) - 8600
(1+2) Tij
exex
exp
(1+
[ ((tt)++))) (T)++)))
Tij
c/ = -1(( cl1.:
where the first inequality uses the Stuchy-Schward inequality, the second equality comes from (7), (22), and the definitions of (( and ((++$) after (21), and the third equality is due to the second equation in (21) and the first equation in (22). This proves (((=((+)))T cllı 2 [[(()+)] - r|li
17

---

for 8 2 O. On the other hand, by the optimality of g(++1) in (20) and the second equation in (22). we have (x(() ) = 0 for & 2 1. Using a similar argument, we can prove || - ri11 2
[[( " 1 e|lı for t 2 and | - r|] + [[( cll ((( cll1. The proof
is completed.
Lemma B.2 (Sufficient decrease of Len (1.)) in (a, 8)). Let ((a )} be the sequence generated by (21) and consider (x!() in (22). Then, for each l 0, we have
Cyp p((t)), _tt1 ,
$ Lye (al9),(0, U nh)
The (1 rlli ([+21) -rlli).
(47)
2
Proof. By [29, Lemma 4.11], for f 2 0, we have Ly (a (P+1), 1, g(t++) g((+))))))))))) U++1, n²) S f. (a(et)), ,(0, Utt, (( ?k [ 1-cli- cl12
By Lemma B.I, for l 2 O, we further have Lm (a((+)), g(t+1), utt1 x²) fm (a(tt)) ,000 Uutt, ||( - -117
1K
(48)
2
On the other hand, we have
Lm (a(t+1) 300), U x²) - cy ((()) U++1 .²) (r, -a(0)) =- --,,
(og((((00)
mp KL([| (0) S ? (0) -rl.:
(49)
where the first equality uses the first equation in (21), the second equality comes from the second equation in (22), and the last inequality is due to Pinsker's inequality [33]. Here, for p, A", the
Kullbbback-Leiin divergence between p and q is (L(pllg) = E. Ps log(p:/4)). Combining (48) and (49), we obtain (47). The proof is completed.
We need to use the following elementary result.
Lemma B.3, Let 01 and 9-2 be two given positive constants. Consider two sequences {ae], {be} S R+ with 8 2 0. If they obey:
a: - ae+1 2 (1 + 1)+1).
(50a)
at S 8zbE
(50b)
for all l 2 1, then we have bn S 02/01 01 and
02
min 1siset1
b, < 91 - /2
1E 1.
Proof. Applying (50) on l = 1, it is easy to see that 8zb1 2 01 2 0,b%. Therefore, we have
by S 02/01 and aj S 02/ /0. For f 2 1, by (50a), we have
1
at aet1
0, (12 + +1
1)41)
201
z
2
ac+1 at
a
Bbebe+1
03
which, together with a1 S 03/01, implies
03
ac+1 S
42 0.
01
2( 1) -1
Let [e/2] be the largest integer smaller than 6/2. For l. 2 1, summing (50a) over [1/2]+1,....
we have
- 19/2)
min 15i51+1
02 S
(b₂ + 82 <
[L//2]+1
03
11
07
2(12/2) 1) -1'
1/2
which immediately implies minigig141 b3 <
#3/07
With the fact that C(1 +1) 2 (e+ V2 1)²
((+11)
for é 2 and bj 02/01, we thus complete the proof.
18

---

With the help of the above lemmas, we can provide new analysis of the iteration complexity of the classical Sinkhorn iteration based on the decreasing properties developed in Lemma B. and B.2. It differs from the proof technique in [19], wherein switching strategy is adopted to establish the complexity.
Lemma B.4. The total number of Sinkhorn iterations to find a point π (f+11) satisfying (23) is at most [2 /(m00++) + 2 - /21, where
R |[((U++]]| + T7,
(51)
where v = |log nr Ilvar + max log r llaars ||log ellvar)- Here, [a] is the smallest nonnegative integer larger than R"
Proof. Let (att4) 11)+) E argmin DER",BER" Lu (a, B, B,U++, ,T By refining the proof in [19]
and [36], we can prove
max{ (oxx(la(tt))1 lvars | i llvar, || /| var: 11 S R 4E 0.
(52)
Since V Val. (alet1), g((+2), uttt , = - x((+1)1,
mn (et1),,, gletz), Utt1, xt) c c 00
and Lm (a, B, Ut+1 x*) is jointly convex with respect to a and B, we have
Lyp (a(e-1), g(t+1), u++1, nt)) Lm (a U++1 x²) ) (T (e(1))) r, r, - ale+1)).
(53)
Given x E RR, let 2m
||<]| 22 + Tmin. it holds that || - Im 1 11 = || 2. For y E R. with (y. 1) = O, we further know that
||x]|var l/ulla-
(u, x) (, x - xm 1) |Bb|a||0 xm 1o0
Applying this assertion with y = n ((+11) -r, x = ofre or = (l++)), we obtain from [[+111]]]]
||||] = 1, (52), and (53) that
U Cy (a 1f+1) glt+!), u++1, x²)- - fn Se U++1 , S ||*1||=((+1)) r111,
A O.
(54)
Applying Lemma B.3 with as = f. (aa10),00)) U++ , ) - fyp (aitu 12 11, U++1, , m²) and by = 1===11(01-11) and using (46), (47), and (54), we have ||r(1 r|11 S 22R++1 /П4 and [|x(0+1]1-
r//1 < (n + v2- 12-11 1)) for all é 2 1. Letting 2 (E + vτ - 1)) S BeT1, we can complete the proof of Lemma B.4.
We now can establish the convergence result of iRBBS, namely, Algorithm 2.
Theorem B.5. Let {x*) } be the sequence generated by Algorithm ef = 62 = O, we have e (x', ) + 0 and e² (x',x*) + If > and 62 > O, then Algorithm 2 stops within at most [Y max(e] _,22) iterations for any 0.+1 2 0, where Y is constant defined in (64).
Proof. Given R²" and 3 E R", by [29, Lemma 4.8] for any E [0, 1] and U1, U2 U,
V (a, B, U2 T²) Cu (at, B,U0 + (1 - u)U π )|F
(55)
< (1|C1 + 2|(C112 u|U1 - U2llF
From [38, 10], we know that the retraction on the Stiefel manifold has the following nice properties, namely, there exist positive constants L, and L2 such that
(R)ttr - Ul} < |I|lE|F and [|Retre(()) (U + f] < L2|
(56)
hold for all & Tyll and U e u. With (55) and (56), following the proof of Lemma [10], for each U E l and E Tol, we have
S. (a, 3, (e) *t) $ Cm (a, B, U, x²) + (grady Cyp (x, x²), 5) + -/2 11111
L
(57)
2
19

---

where
L 2(1} 12)1C Dc 4 11
(58)
Applying (57) with a a= = B¹, U U², and & regrade Lon, (x', n²h) - with (19a),
we have
Ly (a', ,Utt ) S fi (x( - Tt (1 (1- Ltt
11E* 12.
(59)
In addition, by the fact that a = (t+1), g B((+2), and eit (x++1 x = (1+ r|1
for some and (0) = a², ((00) B², we have from (47) that
Sm (x*+1, n²) s Loy (a', , U++1 ,
Uk ( (x'+1 ,
(60)
2
Recalling (24), follows from (59) and (60) that
(x Ex(x)) Te (1 Lrz
1K'
Iis -p) (e (x(+1 x
(61)
2
holds with p (k))
Recalling that Ta takes the form 7. Tt o² for some nonnegative integer s, with (26), we know from (61) that if = log 2(1-61) log ol. there must hold that
(0)
Tmax
1x (x)( 81 111111= -)) (x'+1 ,
Using similar arguments as [54, Lemma 1.1]. this also shows that the backtracking line search in
Algorithm 2 terminates in at most [log 2(1-6.)
Timea L /log o. trials and E; 2 Ep(x'). Moreover, since
Tt (0) Tmin from (26), further know that 7. Ti
(0) 2 II= 20(1-61 Therefore, we know Toaxh
that the nonmonotone linesearch condition (25) is satisfied for some 7. 2I.
Let x* (a", ,") be the minimizer of problem (11). Using a similar argument as [28, Lemma 4.71. by (34) and - log n S ||log 101ooo, we have
Cye (x", T*) 2 - |0 - mh mmlllog=='ls 1lo
(62)
Let Wi =E
Ex(''). Ex' By (62) and using the argument in [44, Theorem 5], we know that such
we 2 0 satisfies
+oc
We S 7 Σ
1- (Em, + ||C||o + mh Illoo n (loe) +00.
By (24), the fact that (25) holds with ne 2 I. and ||5* = eyy, (xt, mh), we have
5iz(e) (x', ,x",, n ))²² + TJ -p-)) p) (e" (x
(
(63)
S L (x', x² - cy. (x++1, ^T + p((e (x', =T)))² - -e2, (x"+1, ,²))))) wi-
Summing (63) over 0, ... T and adding the term (nx/2 /2 () (x²))² on both sides of the obtained equation, and then combining them with Ly (xT+1 ,Tt) 2 Lyp (x", t and (62), we have
T
(1
1n (xt, T 1)) + (
1k 1) (ern (x", n )²))
<
1
+ |101 + Th melloo ~ l
I-
for any T 1. This further implies
f min f81 mk/2 p)
Σn 2== (e) (x", x²) + (e (x',, x²²))² TT=
(64)
Hence, we have e/y, (x", a/) - 0 and ea (x', x + 0 as
20

---

)²) S are fulfilled for = T but not fulfilled for all t < T. Then there hold that en (x", x²) S eI, e (x" )²) S E2. and (e/, (x', x²))² + (e (x', x²))² > minfei E for all < T. Setting T in (64) as T yields T min ((,,,) e I T, which completes the proof.
Suppose en (x', x²) I and en (x',*")
We now ready complete the proof of Theorem 3.2.
Proof of Theorem 3.2. Given 61 > and 62 > O, by Theorem B.5, we know that Algorithm
will return a point x" with T v [T ma 22 }1 satisfying satistying ey, (x', (x'==") S é1 and e (xxxx* S 62. Let #* == Round (x", ( [[(r,))), where "Round" is the rounding pro- cedure given in [3, Algorithm 2]. By the second equation in (22) and x" = (a "T)) with a" = a (f+1) ,,, ST g(++1) for some E, we know that I15m (x",, x >>) 1. By ||C(U))l S || for all U E U. (58), (52), (51), we have ||a" + 1 ||0 S 3| 2mp V. Therefore, we obtain from (13) that
Proinnu(-2V ]lF S =1 21122//00002 (* 2(x²)) & log + [log TM vark) + (3 + 27p v C2-
With the choices of 7/k- EL and e2 in Theorem 3.2, we can immediately know from (13) that such (x", , -x((") )mix) is an (cí, -taationny point of problem (6). With the expression of Y in (64), Lemma B.4 and Remark 3.5, we can complete the proof. Note that the O-) in this theorem hides the constants related to E,, (xx), V, L1. Lg, log n, and ||C||
C Experimental Settings
We implemented all methods in MATLAB 2021b and performed all the experiments on MacBook Pro with a 2.3GHz 8-core Intel Core 19. We follow the ways in [34] to generate the following real dataset.
Datasets C.1. We consider six Shakespeare operas. Each script is preprocessed according to the way in [34] and corresponds to matrix over X E 1300xxx, The values nx of H5, H, JC, MV, O, and RJ are 1303, 1481, 910, 1008, 1148, and 1141, respectively. The weight vector T or c is taken as 1,x /nx.
We set k 2.
Initial points and retraction. We choose a- 0, 3- = 0 for iRBBS and a = 0, 30 0 for R(A)BCD, As for Uo we choose U E argmax (V=,, UUT) with T" E I[(r e) for all methods. Here, T0 is formulated by firstly generating a matrix 7 with each entry randomly drawn from the standard uniform distribution on (0, 1) and then rounding 50 / to T² via [3, Algorithm 2]. The retraction operator in all the above methods is taken as the QR-retraction Retr qr (E) = qf(X + {), where qf(X + {) l satisfying X E qf(X {Xppp(X {) with upp(X + {) being an upper triangular b-b matrix with positive diagonal elements.
Parameters of iRBBS in Section 4. 1. We choose min 10-10/ /L, 7 max = 1000 /L, 7 0.85.
TO 10-3 σ = 1/2, 61 10-4, and p = 0.491Jp
Parameters of R(A)BCI in Section 4. I. As stated in Remark 4.1 in [28), to achieve the best performance of R(A)BCD, one has to spend some efforts tuning the stepsizes. Here, we adopt the stepsizes used therein (with a slight modification, marked in italic type, to have better performance for some cases). We choose TRaCO 0.09 if the instance is X/RJ and TRECO 0.1 otherwise; TRABCD = 0.0015 if the instance is X/RJ with X * H. TRABCD = 0.001 if the instance is H/RJ and TRABCD 0. 0.0025 otherwise. Moreover, we stop R(A)BCD when e., (a², (U*)) S 6 and
e,, (a', 8 - Ut) & E2 or the maximum iteration number reaches 5000.
Parameters of REALM in Section 4.2. Once Th becomes Mmin. we set the corresponding = = and 6222 = and stop the algorithm if e/y, (x*)**) I and e (x*, x²) S ez. We choose f12 10 max loo. ] and E1,1 2||C'll
21

---

13
20
30
400
200
9.30
1000
100
300
400
500
127
16
10
12
17
8 13
10
11"
10
200
6.0
1000
107
(a)
100
(h) d = 50
(c) m
(d) n 10d
Figure 2: Averaged results for Dataset D.1. The scaled time means the time of each method over the
minimum time of all methods. For (a), d {20, 100, 250, 500); for (b), {100, 250, 500, 1000);
for (c), d E {20, 50, 250, 500]; for (d), d E (10, 20, 100, 250}.
D Additional Numerical Results
In this section, we present additional numerical results to evaluate the performance of our proposed methods to compute the PRW distance (1). We follow the ways in [34, 28, 41] to generate the synthetic and real datasets.
Datasets D.1 (Synthetic dataset: Fragmented hypercube [41]). Define a map P(x) = x 2sign(x))
Ex Ck. where sign(-) is taken elementwise, k* E (1....,d) and (e1,..., Cx= } is the canonical basis of Rd. Let = U[[,,,1] 11 be the uniform distribution over hypercube and Ps (u) be the pushforward of μ under the map P. We set as 1/n.
and take both the weight vectors r and e
Datasets D.2 (Real datasets: digits from MNIST datasets [34]). For each digit 0.1.....9 we extract the 128 Se-ddienssionl corresponds to matrix X E RR 122xxx The values nx of 1,...,9 are 980, 1135, 1032, 1010, 982, 892, 958, 1028, 974, and 1009, respectively. The weight vector r or is taken as Iax /nx. choose k = 2.
from a pre-trained convolutional neural network. Each digir
D.1
Comparison with R(A)BCD on Solving Subproblem (11)
We follow the same settings of y) in [28]. For Dataset D.1, we choose 17 0.2 when d < 250 and 0.5 otherwise. For Dataset D.2, we set = 8.
Parameters of R(A)BCD this section. For Dataset D. choose TaacO -RABCD 0.001 as done in [28]. For Dataset D.2, we choose TRaCD = 0.004, and we do not test RABCD for this dataset since the well-chosen stepsize not provided in [28].
For Dataset D.I, we randomly generate 10 instances for each (n, d) pair, each equipped with randomly generated initial points. The comparison results for Dataset D.1 are plotted as Figure 2. Note that the values P returned by different methods almost the same for this dataset. Therefore, we do not report the values the figure. From the figure, we can draw the following observations. (i) Among iRBBS methods. iRBBS with smaller always has less nGrad but more nSk. This means that computing grad (U*)) in a relatively high precision can help to reduce the whole iteration number of updating U. In particular, iRBBS-O, which computes (almost) the exact grad gU''), takes the least nGrad. Because the complexity of one Sinkhorn iteration is much less than that of updating U (see Remark 3.5), iRBBS with moderate 8 generally achieves the best overall performance. The last rows of Figure 2 show that iRBBS-0. 0.1 is almost the fastest among all iRBBS methods. (ii) Our iRBBS- better than R BCD in terms that the former always have smaller nGrad and nSk and are always faster. Particularly, for Dataset D. BRBBS-0. is always more than 5x faster than RABCD and is about more than 10x faster than RBCD. For the largest instance d 250, n 2500, iRBBS-0.1 (1.6s) is more than 7. 7.2x faster than RABCD (11.1s) and about 28.6x 6x faster than RBCD (44.4s). It
22

---

Table 3: The average results of 20 runs for Dataset D.2. Here, "a", , "c", "d", and "e" stand for RBCD, iRBBS-int iRBBS-0.1, and iRBBS-O, respectively. For RBCD and iRBBS-inf, nGrad is equal to nSk.
10 x P
nGrad/nSk
time
data
d
d
D0/D1
0.974647
0.974647
0.974647
0.974647
519
132
59/506
46/1062
2.7
0.7
0.4
0.4
D0/D2
0.794227
0.794227
0.794227
0.794227
988
187
72/943
74/2655
4.8
1.0
0.5
0.6
D0/33
1.202136
1.202136
1.202136
1.202136
554
177
64/934
44/1386
2.6
0.8
0.4
0.4
D0/44
1.220464
1.220464
1.230833
1.230833
912
243
98/998
81/2947
4.1
1.1
DO/D5
1.026643
1.025569
1.026643
1.026643
1178
252
89/1402
96/4355
4.8
1.1
0.5
0.8
D0/D6
0,802711
0.802711
0,802711
0.802711
919
186
66/1052
54/2387
4.0
0.8
0.4
0.5
DO/D7 0.855776
0,855776
0.855776
0.855776
601
163
63/747
56/1727
2.7
0.8
0.4
0.4
D0/88 1.052850
1.0528500
1.052850
1.052850
709
185
84/1330
72/3302
3.1
0.8
0.5
0.6
DM/D9 1.084633
1.084633
1.084633
1.084633
698
179
68/1108
55/2436
3.3
0.8
0.4
0.5
D1/D2 0.664666
0.664666
0,664666
0.664666
757
114
55/638
47/1511
4.1
0.6
0.4
0.5
D1/D3 0.861031
0.861031
0.861031
0.861031
1649
230
105/982
71/2065
8.8
1.3
0.7
0.8
D1/D4 0.666639
0.666639
0.666639
0.666639
982
153
65/505
$4/1282
5.1
0.8
0.4
0.5
DI/D5 0.839978
0.839978
0.839978
0.839978
3436
540
184/1503
157/3844
15.9
2.4
1.0
1.1
D1/D6 0.795720
0.795720
0.795720
0.795368
2919
331
128/1095
119/3354
15.0
1.7
0.8
1.1
D1/D7 0,572325
0.572325
0.572325
0.572325
961
157
65/525
56/1253
5.2
0.9
0.4
0.5
D1/D8 0,878387
0,878387
0.878387
0.878310
5000
402
200/2001
181/5505
25.8
2.1
1.3
1.6
D1/D9 0.853405
0.853405
0.853405
0.853405
951
168
69/581
58/1321
5.6
1.0
0.6
0.7
D2/D3 0.718821
0.718821
0.718821
0.718821
2070
336
120/1733
100/4897
10.0
1.6
0.8
1.0
D2/D4 1.068421
1.085370
1.085370
1.085370
5000
4602
550/8093
382/16291
24.3
21.4
3.2
3.4
D2/D5 1.076288
1.076728
1.076728
1,076728
1320
272
109/1910
89/3990
5,9
1.2
0.6
0.7
D2/D6 0.899301
0.899301
0.899301
0.899301
1202
245
102/967
87/3205
5.8
1.2
0.6
0.7
D2/D7 0.695023
0.695023
0.695023
0.695023
951
174
67/1156
63/3108
4.7
D2/D8
0.671073
0.671073
0.671073
0.671073
1293
213
72/925
69/3547
6.2
1.0
0.4
0.7
D2/D9
1.0567299
1.056729
1.056729
1.056729
1985
284
120/1907
93/3753
9.4
1.3
0.7
0.8
D3/D4
1.199367
1.200346
1.200346
1.200346
5000
312
76/762
69/2523
22.2
1.4
0.4
0.5
D3/D5
0.581711
0.581711
0.581711
0.581711
1187
181
71/793
58/2430
4.8
0.8
0.4
0.5
D3/D6
1.231675
1.231675
1,231675
1.231675
1028
212
82/772
66/2121
4.5
0.9
0.4
D3/D7 0.724492
0.724492
0.724492
0.724492
1101
202
88/831
776/19184
5.2
0.9
0.5
5.3
D3/D8 0,879795
0,879795
0,879795
0,879795
916
200
84/1113
77/3675
4.0
0.9
0.5
0.7
D3/D9 0.830151
0.830151
0.830151
0.830151
1572
210
72/955
49/2193
7.2
1.0
0.4
0.4
D4/D5 1.006255
1.006255
1.006255
1.006255
936
245
87/863
71/2748
3.9
0.5
D4/D6 0.843616
0.843616
0.843616
0.843616
1312
245
96/1102
82/3463
5.9
1.1
0.5
0.7
D4/D7 0.789906
0.789906
0.789906
0.789906
1351
269
96/912
76/2509
6.4
1.2
0.5
0.6
D4/D8 1.094979
1.098234
1.098234
1.098234
5000
213
52/550
52/1930
21.9
0.9
0.3
0.4
D4/D9 0.489215
0.489215
0.489215
0.489215
1997
281
70/1116
54/2610
9.1
1.3
0.4
0.5
DS/D6 0.717009
0.717009
0.717009
0.717009
1335
220
86/891
69/2536
0.9
0.4
DS/D7 0.912583
0.912583
0,912583
0.912583
899
242
72/652
71/2132
3.7
1.0
0.4
0.5
DS/D8 0.715209
0.715209
0.715209
0.715209
1429
225
71/1060
50/2347
5,9
0.9
0.4
0.4
D5/D9 0.775506
0.775506
0.775506
0.775506
1413
252
86/1216
64/3451
5.9
1.0
0.5
0.6
D6/D7 1.109787
1.109787
1.109787
1.109787
826
195
66/398
62/1318
3.8
0.3
0.4
D6/D8 0,917487
0,917487
0.917487
0.917487
725
186
57/721
57/2305
3.2
0.8
0.3
0.5
D6/D9
1.105550
1.105550
1.105550
1.105550
1117
219
95/1038
79/3046
5.0
1.0
0.5
0.6
D7/D8
1.078740
1.078740
1.078740
1,078740
757
180
68/796
59/2455
3.5
0.8
0.4
0.5
D7/D9 0.607866
0.607866
0.607866
0.607866
1985
266
79/1026
72/3208
9.3
1.3
0.5
0.7
D8/D9
0.867753
0.867753
0.867753
0.867753
745
180
57/713
50/2107
3.2
0.4
AVG
0.884689
0.885180
0.885390
0.885380
1560
326
95/1152
95/3366
7.28
1.52
0.57
0.78
should also be mentioned that for all 800 problem instances, there are 5/21 instances for which RBCD/RABCD meet the maximum iteration numbers and return solutions not satisfying the stopping criteria.
For each instance of Dataset D.2, we randomly generate 20 initial points. The comparison results for Dataset D.2 are reported in Table 3. From the table, we have the following observations: (i) for the solution quality, BRBBS-O and iRBBS-0.1 perform the best among all methods; (ii) among iRBBS methods, iRBBS-0.1 and iRBBS-0 take the least nGrad and iRBBS-in inf takes the most nGrad
23

---

Table 4: The average results of REALM for Dataset D.I. In this table, "a" and "b" stand for REA0M-(0.02, 0) and R00LMM 0.9), respectively.
P
nGrad
nSkexp/nSkiog
time
iter
vlld
b
b
100/20
8.34297
8.34296
182
138
3208/8342
13258/0
1.1
0.1
0.0/7.0
8.0/14.0
100/100
9.17320
9.17340
287
241
3837/3354
11369/0
0.5
0.1
0.0/7.0
8.0/14.0
100/250
10.89561
10.89689
467
348
3200/8858
10550/279
1.3
0.2
0.0/7.0
8.0/14.1
100/500
13.31111
13.31524
680
469
2192/3782
9844/213
0.8
0.2
0.0/7.0
8.0/14.0
100/50
8,59993
8.60005
321
189
3996
11700/0
1.8
0.1
0.0/7.0
8.0/14.0
250/50
8.23339
8,23336
150
156
2480/4168
10828/0
2.2
0.3
0.0/7.0
7.9/13.9
500/50
8.12987
8. 8.12986
122
141
1944/4508
11578/0
7.2
0.8
0.0/7.0
8.0/14.0
1000/50
8.07096
8.07095
111
130
1716/5307
12344/0
30.8
2.3
0.0/7.0
8.0/14.0
20/20
9,30977
9.31036
591
221
4016/23861
11906/0
0.5
0.0
0.0/7.1
8.0/14.0
50/50
9,36270
9.36292
386
248
4656/9416
11819/0
0.5
0.1
0.0/7.0
7.9/13.9
250/250
9.17441
9.17476
282
257
2942/4136
11371/248
2.2
0.5
0.0/7.0
8.0/14.0
500/500
9.11498
9.11540
258
262
2391/4884
9927/873
7.5
2.2
0.0/7.0
7.6/13.6
100/10
8.16197
8.16193
405
173
4635/33570
65772/0
4.6
0.3
0.0/7.0
8.0/14.0
200/20
8.13310
8.13309
180
128
2861/10259
15007/0
3.8
0.3
0.0/7.0
8.0/14.0
1000/100
8.11874
8.11873
124
145
2096/5157
12153/141
30.5
3.1
0.0/7.0
8.0/14.0
2500/250
8.11640
8.11639
117
146
2416/5499
11482/624
373.8
103.3
0.0/7.0
7.9/13.9
AVG
9.01557
9.01602
291
212
3033/9262
15057/149
29.3
7.1
0.0/7.0
8.0/14.0
while iRBBS-0 takes the most nSk and iRBBS-inf takes the least nSk; Compared with iRBBS-inf, IRBBS-0.1 is about 2.7x faster for Dataset D.2. Compared with iRBBS-O, iRBBS-0.1 is about 1.4x faster for Dataset D.2; (iii) compared with R(A)BCD, iRBBS-0.1 can take significantly less nGrad and may take bit more nSk. This makes it about 7x faster than RABCD and about 12.8x faster than RBCD for Dataset D.2. Besides, for instances D1/D8, D2/D4, D3/D4, D4/D8, RBCD meets the maximum iteration number.
From the above results, we can conclude that iRBBS generally performs much better than R(A)BCD for the two datasets. More importantly, our methods adopt the adaptive stepsize without needing to tune the best stepsize as done R(A)BCD.
D.2 Comparison on Computing the PRW Distance (1)
In this subsection, we present more numerical results on Datasets D.I and D.2 to illustrate the effectiveness and efficiency of our proposed REALM, namely, Algorithm 1. The subproblem (11) is
solved by our developed iRBBS algorithm, namely, Algorithm 2, with 0g chosen as in Section 4.
The results for Dataset D.1 are presented in Table 4. For each (nt d) pair. randomly generate 10 instances, each equipped with 5 randomly generated initial points. We conside REALM- (0.055, 0.9) and REALM-(00.0, 0) both with ?/ smaller Mmin and does not update the multiplier matrix. From Table 4, we can observe that REALM- (0.055, 0.9) not only return better solutions than REALM- (0.02,, 0.9) but also is about 4x faster. This shows that updating the multiplier matrix does help. In fact, on average REALM-(00055, 0.9) updates the multiplier matrix 8 times in 14 total iterations.
7 0.25, and 7, 0.5. Note that the latter admits
The results over 20 runs on the real Dataset D.2 are reported Table We consider REALM-(1, 0) and REALM-(3, 0.9) both with 7/ 200 and 7/ % 0.25. From Table 5, can that updating the multiplier matrix also helps. Compared with REALM-(1, 0), REALM-(3, 0.9) can not only return better solutions but is about 2.4x faster. On average, REALM- (3, 0.9) updates the multiplier matrix 6.7 times in 1 11.7 total iterations.
24

---

Table 5: The average results of REALM for Dataset D.2, "a" and "b" stand for REALM-(1, and REALM-(3, 0.9), respectively.
10-3 x PE
nGrad
1SSkepp/nSkpp
time
iter
data
a
b
b
b
b
D0/D1
0.97509
0.97508
238
482
770/1406
4468/0
13.9
5.6
0.0/5.0
8,0/13.0
D0/D2
0.79482
0.79454
340
409
1473/2636
5212/0
15.7
3.3
0.0/5.0
3.0/8.0
D0/33
1.20244
1.20571
359
740
1758/4162
6737/0
25.3
5.7
0.0/5.0
8.0/13.0
D0/44
1.22100
1.23156
428
632
2169/3599
6564/0
20.8
4.8
0.0/5.0
8.0/13.0
D0/D5
1.03638
1.03634
539
693
2385/3840
4666/0
21.0
4.4
0.0/5.2
8.0/13.0
D0/D6
0,80313
0.80683
313
1215
1881/2929
12402/0
17.6
8.7
0.0/5.1
8,0/13.0
D0/D7
0.85624
0.85624
297
562
928/1907
6342/0
11.7
4.5
0.0/5.0
8.0/13.0
DMD8
1.05330
1.05361
292
696
1575/2091
7139/0
11.8
5.0
0.0/5.0
7.0/12.0
DM/D9
1.08510
1.08899
766
1780/2386
7557/0
13.9
5.8
0.0/5.0
8.0/13.0
D1/D2
0.66502
0.66501
193
535
1042/1068
6793/0
10.8
6.7
0.0/5.0
8.0/13.0
D1/D3
0.86173
0.86205
484
1069
1348/2261
10720/0
22.3
12.1
0.0/5.7
7.6/12.6
D1/D4
0.66703
0.66702
267
648
860/1372
5138/0
13.1
7.1
0.0/5.0
8.0/13.0
D1/D5
0.83950
0.83796
639
1330
1800/2966
18003/0
19.2
10.5
0.0/5.0
6.0/11.0
D1/D6
0.79633
0.79614
567
1243
1527/2918
12372/0
26.9
12.8
0.0/5.1
6.0/11.0
D1/D7
0.57276
0.57274
318
862
755/2013
7149/0
19.9
10.1
0.0/5.0
8.0/13.1
D1/D8
0,87915
0.87737
1554
1465
3318/10079
22317/0
92.9
16.1
0.0/S.2
7.0/12.0
D1/D9
0.85394
0.85393
326
712
880/1578
S716/0
15.8
8.3
0.0/S.0
8.0/13.0
D2/D3
0.71916
0.71916
858
1755/1892
16797/0
12.3
7.5
0.0/5.0
8.0/13.0
D2/D4
1.08616
1.08614
348
853
3854/1156
7422/0
8.3
6.5
0.0/5.0
8.0/13.0
D2/D5
1,07706
1.08610
574
1083
707/6040
7035/1642
33.0
15.0
0.0/5.0
8.0/13.0
D2/D6
0.90032
0,90030
342
981
1022/1339
10140/0
8.6
7.5
0.0/5.0
6.0/11.0
D2/D7
0.69549
0.70124
272
538
1890/2284
7972/0
21.0
6.0
0.0/5.0
7.0/12.0
D2/D8
0.67232
0.67159
336
1210
1263/2180
23641/0
13.5
10.2
0.0/5.0
7.0/11.9
D2/D9
1.05704
1.06966
502
721
3027/4541
4172/0
28.0
5.9
0.0/5.0
8.0/13.0
D3/D4
1,20147
1.20145
400
861
1449/1890 896
6847/0
11.8
6.5
0.0/5.0
8.0/13.0
D3/D5
0.58300
0.58299
295
435
1284/1613
5758/0
8.8
3.1
0.0/5.0
4.0/9.0
D3/D6
1.23233
1.23234
750
1166/1290
7664/0
7.9
5.6
0.0/5.0
8.0/13.0
D3/D7
0.72528
0.72529
314
719
836/1263
9008/0
8.4
5.8
0.0/5.0
6.0/1100
D3/D8
0,88074
0,88071
334
781
1433/2175
13688/0
13.0
6.2
0.0/5.2
6.0/11.0
D3/D9
0.83143
0.83145
363
983
1552/2159
20490/0
13.6
8.5
0.0/5.0
8.0/13.0
D4/D5
1.00719
1.00700
302
666
995/1569
8687/0
8.3
4.5
0.0/5.0
6.0/11.5
D4/D6
0.84488
0.84485
650
935/1550
6296/0
9.1
4.4
0.0/5.0
4.5/9.4
D4/D7
0.79120
0.79065
680
856
2701/3931
12068/0
24.3
6.8
0.0/5.2
6.0/1100
D4/D8
1.09898
1.09899
308
566
1307/1616
6729/0
9.5
4.3
0.0/5.0
7.0/12.0
D4/D9
0.49061
0.48902
304
437
1620/1933
10152/0
11.4
3.8
0.0/5.0
4.0/9.0
D5/D6
0.71801
0.71769
628
1275/1716
8743/0
8.9
4.2
0.0/5.0
4.0/9.0
D5/D7
0.91332
0.91333
287
649
744/928
5357/0
5.6
4.4
0.0/5.0
8.0/13.0
DS/D8
0.71623
0.71573
323
823
1633/2654
17832/0
13.8
6.0
0.0/5.0
6.0/1100
D5/D9
0.77656
0.77597
371
811
1952/2633
14998/0
14.4
5.9
0.0/5.0
5.0/10.0
D6/D7
1.11056
1.11051
340
969
843/790
4240/0
5.7
6.8
0.0/5.0
8.0/1.00
D6/D8
0.91820
0,91821
768
1258/1703
14078/0
5.9
0.0/5.0
8.0/13.0
D6/D9
1.10662
1.10607
368
550
1029/1712
4346/0
10.4
3.9
0.0/5.0
3.0/8.0
D7/D8
1.07946
1.07946
298
825
909/1 188
6158/0
7.6
6.2
0.0/5.0
8.0/13.0
D7/D9
0.60736
0.60796
382
606
1794/2195
11288/0
13.8
5.2
0.0/5.0
4.0/9.0
D8/D9
0.86860
0.86855
287
595
1187/1857
8905/0
10.9
4.6
0.0/5.0
5.0/10.0
AVG
0.88606
0.88697
386
783
1504/2378
9551/36
16.1
6.7
0.0/5.0
6.7/11.7
25


Reviews:
[{"review_id": "wGCiu6z5nf", "reviewer": "Reviewer_g8K2", "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.", "questions": "see the weakness", "limitations": "see the weakness.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "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.\n2. 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?\n\n2. the writing of this paper needs to be further improved. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for pointing out the weakness, which gives us an opportunity to clarify it.  \n\nFirst, there are indeed some systematic studies comparing exponential ALM and exponential penalty approaches. We wish to highlight several main results from the existing literature on this matter: \n- Convex case. \n   - Tseng and Bertsekas [TB93], listed as Ref. [45] in our paper, proved that the penalty parameter $\\eta_k$ in the exponential ALM (with subproblems solved exactly) can be chosen as any positive number. This is in sharp contrast to the exponential penalty approach, where the penalty parameter is typically required to be very small or go to zero. Please refer to Proposition 3.1 in [TB93].  Additionally, [TB93]  established the linear convergence rate of the exponential ALM for solving the linear programming; see Proposition 4.1 therein.\n   - More recently, Yang and Toh [YT22], listed as Ref. [50] in our paper, proposed an inexact version of the Bregman proximal point algorithm, which is equivalent to the exponential ALM with subproblems solved inexactly. They proved that the penalty parameter can be chosen as any positive number. Moreover, they claimed that their proposed approach exhibits greater stability than the exponential penalty approach (when applied to solve the standard optimal transport problem). In particular, they mentioned in the paper:  ``*our iBPPA with the entropic proximal term can bypass some numerical instability issues that often plague the popular Sinkhorn's algorithm used in the OT community. This is because in contrast to Sinkhorn's algorithm, our iBPPA does not require the proximal parameter to be very small in order to obtain an accurate approximate solution, as evident from our numerical results*\"\n- General nonlinear case. \n   - Under the regular condition, the strict complementarity condition, and the second-order sufficient condition, Dussault [Dus04] demonstrated that the exponential ALM exhibits one-step superlinear convergence with a rate of $4/3$ (see Theorems 1.1 and 2.1 therein). In contrast, the exponential penalty approach takes two-step superlinear convergence with the same rate (see the second paragraph on Page 475). Additionally, Dussault [Dus04] provided a detailed example illustrating that the exponential ALM can be more stable than the exponential penalty approach. For further information and specific results, please refer to Tables 3.1-3.3 in the referenced paper.\n   - Recently, Echebest et al. [ESS16], listed as Ref. [18] in our paper, showed that the penalty parameter $\\eta_k$ *can be bounded away from zero* under the conditions of the linear independence constraint qualification, the strict complementarity condition, and the second-order sufficient condition. However, no evidence indicating that the exponential penalty approach enjoys a similar property under these conditions.  \n\nBased on the aforementioned studies on the exponential ALM, we stated in our paper ``Based on the knowledge that the exponential ALM is usually more stable than the exponential penalty approach...\". However, we apologize for not providing more detailed explanations and justifications in the current version. We will certainly include the corresponding discussion in the revised version.  \n\nSecond, we acknowledge that the exponential ALM might face a similar numerical instability issue as the exponential penalty approach if the penalty parameter $\\eta_k$ becomes too small. However, we hope the following explanations could convince you that the exponential ALM is potentially more stable than the exponential penalty approach for the PRW problem.\n- To address the potential numerical instability caused by the small penalty parameter $\\eta_k$, we propose to update $\\eta_{k+1}$ as $\\eta_{k+1} = \\max\\\\{\\gamma_{\\eta} \\eta_k, \\eta_{\\min}\\\\}$ as in Line 279.  This update scheme ensures that the penalty parameter does not become too small during the iterations. We observed from our numerical experiments in Table 2 and Tables 4-5 that, compared to the exponential penalty approach, the exponential ALM can accommodate larger values of $\\eta_{\\min}$. This advantage is primarily attributed to the update of the multiplier matrices. By allowing for a larger $\\eta_{\\min}$, the exponential ALM is less likely to suffer from the numerical instability due to the use of a small penalty parameter, as observed in the exponential penalty approach.\n- Theoretically, we can extend the analysis in [ESS16] to prove that the penalty parameter $\\eta_k$ in the proposed Riemannian exponential ALM will also be bounded away from zero if the Riemannian versions of the three conditions hold, including the linear independence constraint qualification, the strict complementarity condition, and the second-order sufficient condition. We will include additional remarks in the revision to provide further clarity on this point.  It should be mentioned that these three conditions might not be easy to check since we do not have prior knowledge of the solution.  However, based on our numerical results, we believe that for certain instances, the three conditions indeed hold.\n\nThird, we shall further enhance the clarity of our writing in the revision to ensure that the results are presented more clearly and the paper becomes easier to follow.   \n\n The corresponding references are listed in order.   \n- [Dus04]  J.-P. Dussault. Augmented non-quadratic penalty algorithms. *Math. Program.*, 99(3):467-486, 2004.\n- [ESS16]  N. Echebest, M. D. Sánchez, and M. L. Schuverdt.  Convergence results of an augmented Lagrangian method using the exponential penalty function. *J. Optim. Theory Appl.*, 168:92-108,2016.\n- [TB93] P. Tseng and D. P. Bertsekas. On the convergence of the exponential multiplier method for convex programming. *Math. Program.*, 60(1-3):1-19, 1993.\n- [YT22]  L. Yang and K.-C. Toh.  Bregman proximal point algorithm revisited: A new inexact version and its inertial variant. *SIAM J. Optim.*, 32(3):1523-1554, 2022."}, {"review_id": "ljUIHG29Kz", "reviewer": "Reviewer_KEsx", "summary": "This work proposes a new method, called REALM, to compute the projection robust Wasserstein (PRW) distance. \n\nThe method REALM is an extension of the exponential augmented Lagrangian method to the Riemannian space. \n\nThe convergence of REALM is established. \n\nTo solve a subproblem during REALM, this work proposes an inexact Riemannian Barzilai-Borwein method with Sinkhorn iteration (iRBBS). \n\nThe 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. \n\nThis work claims the robustness in terms of the parameter-tuning in their proposed algorithms. ", "questions": "Hello authors, I have the following questions, \n\n1. 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. \n\n2. 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? \n\n3. 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. \n", "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. ", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 2, "contribution": 2, "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. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for your positive comments on the basic results achieved in this paper and your helpful comments and suggestions.  \n- Response to Weakness.   \nThanks for your comments.  We want to take this opportunity to elaborate more on the nonmonotone linesearch condition (25) to clarify this issue from the theory perspective.  \nIn particular, the interested problem (11) is a three-block optimization problem and Huang et al. 2021 proposed to solve it by R(A)BCD. The current paper treats it as a one-block optimization problem with $U$ as the only variable, as shown in problem (17). The new perspective allow applying the (inexact) gradient algorithm to solve problem (17); see the proposed iRBBS in Algorithm 2. The key difference here is that the stepsizes in iRBBS can be adaptively chosen with the help of the linesearch condition (25), but tuning the stepsizes for updating $U$ is not easy for R(A)BCD. This is one of the two main contributions of the current paper. The other benefits of treating the interested problem as a one-block optimization problem as in (17) can be found in our responses to your Question 3. \n- Response to Question 1.    \nThanks for your important comment. The intuition in using the proposed exponential ALM lies in its equivalence to a proximal point method for solving the dual problem, incorporating an entropy function. The estimate of the multiplier matrix $\\pi^k$  can be viewed as the proximal point center, and a varying center $\\pi^k$ typically outperforms a fixed center $\\pi^k \\equiv 11^T$ in the proximal point method.  \nTo further improve the numerical performance (in particular the numerical stability) of the proposed REALM, we have carefully exploited the structure of the PRW distance problem and used the problem's special structure into the general exponential ALM framework. More specifically, the first algorithmic innovation is the conditions in (14), which are crucial to guarantee the boundness of the iterates. The second algorithmic innovation is the simple rules of updating the penalty parameters and the multiplier matrices, which take into account of the specific connections between approximate stationary points of the subproblem and the original problem. The latter is of great importance to both the global convergence and the numerical performance/stability of the proposed REALM.\n\n- Response to Question 2.   \nThank you for this important comment, which gives us an opportunity to clarify this point. In addition to the two specific conditions (14) and (16), the following two differences distinguish our proposed REALM from the existing exponential ALM (e.g., the one proposed in Echebest et al. 2016):\n   - **Measure of complementarity.** The measure of complementarity used in our proposed REALM differs from that employed in Echebest et al. 2016. More specifically, the measure in our REALM is motivated by the direct use of the complementarity condition adopted in the classical (quadratic) ALM, while the measure of complementarity used in Echebest et al. 2016 can be regarded as ``an appropriate measure for the exponential case\" as mentioned in the second paragraph on Page 96 of their paper.\n   - **Conditions on global convergence.** To guarantee the global convergence of the exponential ALMs, some (strong) constraint qualifications and the boundness of the iterates generally need to be assumed; see Proposition 2.1 and Theorem 2.1 in  Echebest et al. 2016 for the corresponding results. The current paper extends the exponential ALMs to solve a class of inequality-constrained nonlinear optimization problems with manifold constraints. Based on our best knowledge, this is the first Riemannian version of the exponential ALM. In particular, for our considered PRW distance problem,  we can prove the boundness of the iterates generated by the proposed REALM without making the assumption and establish the global convergence of the proposed REALM without explicitly dealing with the constraint qualification assumption. This advantage is mainly due to the aforementioned *essential* changes in the proposed REALM (compared with the existing methods), i.e., specific conditions (14) and (16) on the solution of subproblems and the adopted measure of complementarity.  \n- Response to Question 3.  \nYour understanding is absolutely correct, and we truly appreciate your valuable suggestion. \nOne motivation of proposing iRBBS is that the condition number of the one-block optimization problem could be smaller than that of the three-block optimization problem. Consequently,  employing an inexact gradient descent method to solve the one-block optimization problem would be more efficient than simply using the block coordinate descent approach for solving the corresponding three-block optimization problem. Moreover, the inexact gradient method also provides important and valuable insights into adaptively choosing the stepsize for updating $U$ via the nonmontone linesearch. Last but not least, the fact that updating $\\alpha$ and $\\beta$ is much easier than updating $U$ also contributes to the superior performance of iRBBS compared to existing block coordinate descent approaches in practice, as you correctly pointed out. Otherwise, there is no necessity to update $\\alpha$ and $\\beta$ multiple times while only updating $U$ once in each iteration of the proposed iRBBS.\n- Response to Limitations.  \nThank you very much for this insightful comment. After a careful and thorough investigation, we can now establish the boundness of the penalty parameter for the proposed REALM by leveraging the analysis in Echebest et al. 2016. To achieve this, we need to consider the Riemannian versions of the following three conditions: the linear independence constraint qualification, the strict complementarity condition, and the second-order sufficient condition. However, it remains unclear under which (easily checkable) conditions the above three conditions hold true, which requires further investigation.\n"}, {"review_id": "tqr6UGcr2z", "reviewer": "Reviewer_X3cc", "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\n", "questions": "See the above setion ", "limitations": "See Weaknesses", "rating": 6, "confidence": 1, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1.\tThe 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\n\n2.\tTo 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).\n", "weaknesses": "1.\tThe authors should provide a proof sketch for their main theoretical analysis in this paper.\n\n2.\tThe authors should add more experimental results to verify their theoretical results and their algorithms.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- Response to Question 1.  \nMany thanks for your suggestion. We agree with you that it is always helpful to provide a proof outline or a preview of the proof before delving into detailed theoretical analysis. We shall take your suggestion into account when we revise the paper.\n\n- Response to Question 2.  \nThanks for this comment. We shall try to provide more simulation results to better verify the obtained theoretical results.\n"}, {"review_id": "rkSzhBSGCi", "reviewer": "Reviewer_iBqS", "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.", "questions": "(1) Can the authors give the iteration complexity for the overall algorithm? What is the main difficulty here?\n(2) What is the dominated computational time in the proposed algorithm and the compared algorithms?\n", "limitations": "The authors pointed out an important limitation in the theoretical analysis, which is the lack of a lower bound of \\eta_k.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "(1) A method (REALM) is proposed and its global convergence is given.\n(2) An inexact Riemannian BB method with Sinkhorn is developed and its iteration complexity is derived. Such a result is interesting by itself. \n(3) From the numerical comparison, the proposed method is most efficient.\n", "weaknesses": "The existing methods have iteration complexities for the corresponding algorithms. This paper only gives the iteration complexity for the subproblem.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "- Response to Question 1.  \nMany thanks for this insightful comment. Regarding the current overall Algorithm 1,  we were not able to establish the iteration complexity due to the following two main difficulties: (i) characterizing the connection between the two complementarity measures $\\\\|W^k\\\\|\\_F$ and $\\langle \\pi^k, Z(\\mathbf{x}^k)\\rangle$ at the approximate stationary point of the subproblem; and (ii) establishing the relationship between $\\eta_k \\tilde{\\pi}^{k+1}\\_{ij}$ and $\\varphi(\\mathbf{x}^k)\\_{ij}$. \n\n   However, thanks to Theorem 2.4, we can now slightly modify Algorithm 1 to establish the iteration complexity. More specifically, let  $e^k = \\langle \\tilde \\pi^{k}, Z(\\mathbf{x}^k)\\rangle$, where $\\tilde \\pi^{k}: = \\mathsf{Round}(\\phi_{\\eta_k}(\\mathbf{x}^k, \\pi^k), \\Pi(r,c))$ is a feasible matrix returned by the rounding procedure mentioned in Theorem 2.4. By modifying the ``if'' condition in Line 6 of Algorithm 1 as \n   \\\\[\n   \\\\|W^{k}\\\\|\\_F \\leq \\gamma\\_W \\\\|W^{k-1}\\\\|\\_F \\\\quad \\\\mathrm{and}\\\\quad  e^k \\leq \\gamma_W e^{k-1},\n   \\\\]\n   we can establish the iteration complexity of the whole exponential ALM without compromising its global convergence. By leveraging the connection between the approximate stationary points of the subproblem and the original problem as proven in Theorem 2.4, we can ensure that the algorithm will terminate within at most \n   \\\\[\n   \\\\mathcal{O}\\\\left(\\\\max\\\\left\\\\{\\\\log \\\\frac{1}{\\epsilon\\_1}, \\\\log \\\\frac{1}{\\epsilon\\_2}, T\\_k\\\\right\\\\}\\\\right)\n   \\\\] iterations, where $T_k := \\min \\\\{k \\mid \\varrho_k \\leq \\epsilon_c\\\\}$.  We will incorporate these improved results in the revised version to address your concerns. Thank you once again for your valuable comment.\n\n- Response to Question 2.   \nThanks very much for this important comment. The main computational cost in the proposed iRBBS and R(A)BCD proposed by (Huang et al. 2021a) lies in the following steps:  \n   - Compute the inexact Riemannian gradient $\\xi^t$, whose cost is $\\mathcal{O}(ndk + n^2k + dk^2)$.\n   - Update  $U^{t+1} = \\mathrm{Retr}_{U^t}(-\\tau_t \\xi^t)$, whose cost is $\\mathcal{O}(dk^2)$. \n   - Compute the matrix $A \\in \\mathbb{R}^{n \\times n}$ with $A_{ij} = \\pi_{ij}^k \\exp(-\\langle M_{ij}, UU^T\\rangle/\\eta_k)$ (shown in Line 255),  whose cost is $\\mathcal{O}(n^2)$.\n   - Perform the Sinkhorn iteration (21) or (27),  whose cost is $\\mathcal{O}(n^2)$.    \n \n   Note that the cost of verifying whether the linesearch condition (25) holds is low since $\\mathcal{L}(\\mathbf{x}^{t+1})$ is always equal to $r^T \\alpha^{t+1} + c^T \\beta^{t+1}$ due to $\\\\|\\zeta^{(\\ell)}\\\\|_1 = 1$, as shown below Eq. (21).\n\n- Response to Limitations.  \nThanks very much for pointing out this important issue.  Now we can overcome this limitation by assuming the Riemannian versions of the following three conditions, including the linear independence constraint qualification, the strict complementarity condition,  and the second-order sufficient condition.  This extension builds upon the results of Echebest et al. (2016) (listed as Ref. [18] in the paper). We will incorporate this result and highlight the significance of considering these Riemannian conditions in the revision.\n\n"}, {"review_id": "9jvww06Dsq", "reviewer": "Reviewer_EBs3", "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.", "questions": "1. Line 65: One ``method\" in this line is redundant.\n\n2. Line 91: There is no definition for $\\nabla f(U)$. The metrics on the left and on the right side are different. \n\n3. Line 145, Proposition 2.6: How to come out with the $x^s$? What is the motivation for the definition of $x^s$?\n\n4. Line 252: Should be $\\theta_t$ in the place of $\\theta$?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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)$?\n\n(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.\n\n(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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Many thanks for your positive comments on the basic results achieved in the paper and your helpful comments and suggestions. The following are our point-to-point responses to your comments. \n- Response to Weakness 1.   \nThank you for your comments. We acknowledge that, due to the page limit, some mathematical details/arguments are missing in the paper. We shall add them either in the revised paper or in the supplementary material. Below let us first provide the detailed derivations for the two examples you mentioned in your above comments.\n   - Due to the optimality of $(\\\\mathbf{x}^*, y^*)$, we have $\\nabla\\_{y} \\widetilde{\\mathcal{L}}\\_{\\eta}(\\mathbf{x}^*, y^*, 11^T) =0.$ Recalling that  $\\widetilde{\\mathcal{L}}\\_{\\eta}(\\mathbf{x}^*, y^*, 11^T) = r^T \\alpha + c^T \\beta + y + \\eta \\sum_{ij} \\exp\\left (- \\frac{\\varphi(\\mathbf{x})\\_{ij} + y}{\\eta} \\right),$ we obtain $\\nabla\\_{y}\\widetilde{\\mathcal{L}}\\_{\\eta}(\\mathbf{x}^*, y^*, 11^T) = 1  - \\sum\\_{ij} \\exp\\left (- \\frac{\\varphi(\\mathbf{x^*})\\_{ij} + y^*}{\\eta} \\right) = 0$,  which gives  $y^* =  \\eta \\log \\sum\\_{ij} \\exp\\left (-\\frac{\\varphi(\\mathbf{x^*})\\_{ij}}{\\eta} \\right)   =  \\eta \\log (\\\\|\\zeta\\_{\\eta}(x^*, 11^T)\\\\|\\_1).$\n    - By the expression \n\\\\[\\mathcal{L}\\_{\\eta}(\\mathbf{x}, \\pi) = r^T \\alpha + c^T \\beta + \\eta \\log \\sum\\_{ij} \\pi\\_{ij} \\exp\\left( -\\frac{\\alpha\\_i + \\beta\\_j + \\langle M\\_{ij}, UU^T\\rangle} {\\eta}\\right)\n\\\\]\nand employing the chain rule,  we can compute the gradient of $\\mathcal{L}\\_{\\eta_k}(\\mathbf{x}, \\pi^k)$ with respect to $U$  as \n\\\\[\\nabla\\_U \\mathcal{L}\\_{\\eta\\_k}(\\mathbf{x}, \\pi^k) = - 2  V\\_{\\phi\\_{\\eta\\_k}(\\mathbf{x}, \\pi^k)} U\\\\] with\n$V\\_{\\phi_{\\eta\\_k}(\\mathbf{x}, \\pi^k)}  = \\sum\\_{ij} [\\phi\\_{\\eta\\_k}(\\mathbf{x}, \\pi^k)]\\_{ij} M\\_{ij}$ (and $V_{\\pi}$ for a given matrix $\\pi$ is defined in Line 85 of the paper). Moreover, the gradient of $\\mathcal{L}_{\\eta_k}(\\mathbf{x}, \\pi^k)$ with respect to $\\alpha$ and $\\beta$ can be found in Lines 124 and 125, respectively. \n\n- Response to Weakness 2.   \nThanks for this important comment. Theoretically, Algorithm 2 allows $\\theta_t$ to be any nonnegative number, which is mainly due to the nice property shown in Lemmas B.1 and B.2 of the Sinkhorn iteration. Please see also Theorem B.5 for the convergence result of Algorithm 2. However, it would be better to clearly specify the choice/update of $\\theta_t$ in Algorithm 2, as you suggested. In the revision, we shall add the update of $\\theta_t$, namely, $$\\theta_{t+1} = \\max\\\\{\\theta/(2\\\\|C\\\\|_\\infty) \\\\|\\xi^t\\\\|_F, \\epsilon_2\\\\},$$  in Line 6 of Algorithm 2 and include the phrase ``Set  $\\theta_0 = 1$ and choose $\\theta > 0$'' in its input part.  The response also relates to your Question 4 below.   \n\n- Response to Weakness 3.   \nThanks for your important comment. There have already been systematic studies in the literature on the advantage of the exponential ALM over the exponential penalty approach in terms of better numerical stability due to its ability of avoiding too small penalty parameters; please check Refs. [18], [45] and [50] in the paper for your reference. However, it is hard to give an explicit estimate on the range ``too small''.  Based on our numerical experience, we have found that when $\\eta_k\\leq 1/500\\max\\\\{||r||\\_\\infty,||c||\\_\\infty\\\\}$ or $\\eta_k \\leq 1/900 \\\\|C(U^t) - \\eta_k \\log \\pi^k\\\\|_{\\mathrm{var}}$ (as shown in Line 283 and Line 284, respectively), the penalty parameter would be considered too small (when REALM is used to solve the PRW distance problem). In such cases, we recommend performing the Sinkhorn iteration with log-domain stabilization, as shown in Eq. (21) and Line 301.\n- Response to Question 1.  \nThanks. We shall correct this typo in the revision. \n- Response to  Question 2.   \nThank you for bringing this to our attention. You are correct, and we apologize for the oversight in not making this point clear. The notation $\\nabla f(U)$ means the Euclidean gradient of $f$ with respect to $U$. Usually, the metrics on the two sides of $\\langle \\mathrm{grad} f(U), \\xi \\rangle = \\langle \\nabla f(U), \\xi \\rangle$ are different: the left-hand side should be $\\langle \\cdot, \\cdot \\rangle_U$ which is a smoothly varied inner product on the tangent space at $U$.  Since the Stiefel manifold is an embedded manifold of a Euclidean space $\\mathcal{E}$, we can simply choose $\\langle \\cdot, \\cdot \\rangle_U$ as the standard inner product in $\\mathcal{E}$, which coincides with the metric in the right-hand side.  We will clarify this issue in the revision. \n- Response to Question 3.   \nThanks for this important comment. It comes out with an observation that  \n\\begin{align}\n&r^T (\\alpha + \\upsilon_1 \\mathbf{1}) + c^T (\\beta + \\upsilon_2 \\mathbf{1}) + \\eta \\log \\sum_{ij} \\pi_{ij} \\exp\\left( -\\frac{(\\alpha_i + \\upsilon_1) + (\\beta_j + \\upsilon_2) + \\langle M_{ij}, UU^T\\rangle} {\\eta}\\right)  \\\\\\\\\n={}& r^T \\alpha   + c^T \\beta  + \\eta \\log \\sum_{ij} \\pi_{ij} \\exp\\left( -\\frac{\\alpha_i  + \\beta_j  + \\langle M_{ij}, UU^T\\rangle} {\\eta}\\right),  \n\\end{align}\nwhere the equality uses the fact that $r^T \\mathbf{1} = c^T \\mathbf{1} = 1$.  Hence, by carefully choosing $\\upsilon_1$ and $\\upsilon_2$, we can derive $\\mathbf{x}^s$ and the corresponding results in Proposition 2.6. \nMoreover, the first two requirements in Eq. (14) are essential to establish the boundness of $\\alpha^k$ and $\\beta^k$, which in turn ensures the convergence of REALM.  This is the main motivation behind defining such $\\mathbf{x}^s$. \n- Response to Question 4.   \nThank you for your comment. The parameter $\\theta$ is a pre-selected constant in the proposed iRBBS algorithm. By choosing different values of $\\theta$, we obtain different versions of iRBBS, denoted as iRBBS-$\\theta$. Usually, a smaller $\\theta$ means that the subproblem is solved more accurately. We will clarify the meaning of $\\theta$ in Line 252 to avoid any potential confusion. This clarification will be included in the revision.\n\n"}]

Author rebuttal:

Meta-review: This paper considers the projection robust optimal transport problem. A new algorithm, i.e., Riemannian counterpart of the exponential multiplier method, is proposed to solve this problem. The authors established the convergence analysis and demonstrated the efficiency of the proposed algorithm via numerical experiments. All reviewers are favorable to the acceptance of this paper.
Decision: Accept (poster)