Title: A Riemannian Exponential Augmented Lagrangian Method for Computing the Projection Robust Wasserstein Distance
Venue: NeurIPS 2023
Accepted: yes
Keywords: Barzilai-Borwein method

exponential augmented Lagrangian

inexact gradient

Stiefel manifold

Sinkhorn iteration

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

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 glyph[epsilon1] -stationary point of the original PRW distance problem within O ( glyph[epsilon1] -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.

1 Introduction

The optimal transport (OT) problem has found wide applications in machine learning, representation learning, data sciences, and image sciences; see [21, 5, 42, 1, 14] and the references therein for more details. However, its direct application in 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 a low-dimensional subspace that maximizes the

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 { x 1 , . . . , x n } ⊂ R d and { y 1 , . . . , y n } ⊂ R d , define M ij = ( x i -y j )( x i -y j ) T for each ( i, j ) ∈ [ n ] × [ n ] with [ n ] := { 1 , . . . n } . Let 1 ∈ R n be the all-one vector and δ x be the Dirac delta function at x . Given r = ( r 1 , . . . , r n ) T ∈ ∆ n := { z ∈ R n | 1 T z = 1 , z > 0 } and c = ( c 1 , . . . , c n ) T ∈ ∆ n , define two discrete probability measures µ n = ∑ n i =1 r i δ x i and ν n = ∑ n i =1 c i δ y i . For k ∈ [ d ] , the k -dimensional projection robust Wasserstein (PRW) distance between µ n and ν n is defined as [41]

where 〈· , ·〉 is the standard inner product in R n × n , C ( U ) ∈ R n × n with [ C ( U )] ij = 〈 M ij , UU T 〉 , U = { U ∈ R d × k | U T U = I k } is known as the Stiefel manifold with I k being the k -byk identity matrix, and Π( r, c ) = { π ∈ R n × n | π 1 = r, π T 1 = c, π ≥ 0 } . 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-regularized OT subproblem with dimension n and a full or top k eigendecomposition of a d -byd matrix needs to be solved exactly at each iteration. Very recently, Lin et al. [34] proposed a Riemannian (adaptive) gradient ascent with Sinkhorn (R(A)GAS) algorithm for solving the following entropy-regularized problem with a small regularization parameter η :

where p η ( U ) = min π ∈ Π( r,c ) {〈 π, C ( U ) 〉 -ηH ( π ) } , in which H ( π ) = -∑ ij π ij log π ij is the entropy function. They showed that R(A)GAS can return an glyph[epsilon1] -stationary point of PRW problem (1) within O ( glyph[epsilon1] -4 ) iterations if η = O ( glyph[epsilon1] ) in (2). However, at each iteration, R(A)GAS needs to solve a regularized OT problem in relatively high precision, which results in a 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) as

where L η ( · , · ) is defined in (8) further ahead. By choosing η = O ( glyph[epsilon1] ) in (3), they showed that the whole iteration complexity of R(A)BCD to attain an glyph[epsilon1] -stationary point of PRW problem (1) reduces to O ( glyph[epsilon1] -3 ) , which significantly improves the complexity of R(A)GAS.

However, there are two main issues of R(A)BCD and R(A)GAS. First, to compute a solution of problem (1) with relatively high quality, η 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 a 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 constraints (see problem (6) further ahead), we can resolve the above-mentioned two issues. Our main contributions are summarized as follows. See also Figure 1 for a summary of the related works and the main results of this paper.

(i) We propose a 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.

An

𝝐

-stationary

point of (3)

with

𝜼 = 𝑶 𝝐

is

an

𝝐

-stationary

point of (1)

Compute an

𝝐

-stationary

point of (11) within

𝑶 𝜼

!𝟏

𝝐

!𝟐

𝐢𝐭𝐞𝐫𝐚𝐭𝐢𝐨𝐧𝐬

(Theorems 3.2 & B.5 )

Compute an

𝝐

-stationary

point of (3) within

𝑶 𝜼

!𝟏

𝝐

!𝟐

𝐢𝐭𝐞𝐫𝐚𝐭𝐢𝐨𝐧𝐬

Compute an

𝝐

-stationary

point of (2) within

𝑶 𝜼

!𝟐

𝝐

!𝟐

𝐢𝐭𝐞𝐫𝐚𝐭𝐢𝐨𝐧𝐬

An

𝝐

-stationary point of (11) with

𝜼 = 𝑶 𝝐

is an

𝝐

-stationary point of (1)

(Theorem 2.4)

PRW Problem (1)

Problem (3)

Problem (11)

Problem (2)

An

𝝐

-stationary point of (2) with

𝜼 = 𝑶 𝝐

is an

𝝐

-stationary point of (1)

R(A)GAS [33]

R(A)BCD [27]

iRBBS (this paper)

REALM (this paper)

Converge to

a stationary

point of (1)

(Theorem 2.7)

(ii) To efficiently solve the subproblem in REALM (i.e., problem (11) further ahead), we view it as a 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 a flexible number of Sinkhorn iterations is performed to compute the inexact Riemannian gradient. Compared with R(A)BCD, our proposed iRBBS (applied to problem (11) with fixed π k and η k = O ( glyph[epsilon1] ) ) can not only return a stronger glyph[epsilon1] -stationary point of PRW problem (1) (compared with the definitions in [34, 28]; see Remark 2.2), within O ( glyph[epsilon1] -3 ) iterations (see Theorem 3.2 with glyph[epsilon1] 1 = glyph[epsilon1] 2 = glyph[epsilon1] ), 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 x ∈ R n , Diag( x ) is an n × n diagonal matrix with x being its main diagonal. For a matrix A , denote A min = min ij A ij , A max = max ij A ij , ‖ A ‖ 1 = ∑ ij | A ij | , ‖ A ‖ ∞ = max ij | A ij | , ‖ A ‖ 2 F = ∑ ij A 2 ij , and ‖ A ‖ var = A max -A min . Denote by log( · ) and exp( · ) the element-wise logarithmic and exponential operators, respectively. We use R n × n + and R n × n ++ to denote the nonnegative and positive orthants of R n × n , respectively. Throughout this paper, we define C ∈ R n × n with C ij = ‖ x i -y j ‖ 2 and V π = ∑ ij π ij M ij .

The tangent space at U ∈ U is T U U = { ξ ∈ R d × k | U T ξ + ξ T U = 0 } . Let T U = { ( U, ξ ) | U ∈ U and ξ ∈ T U U} be the tangent bundle of U . A smooth map Retr : T U → U : ( U, ξ ) ↦→ Retr U ( ξ ) is called a retraction if each curve R ( t ) = Retr U ( tξ ) satisfies R (0) = U and R ′ (0) = ξ ; see [9, Definition 3.47] or [2, Definition 4.1.1]. The Riemannian metric 〈· , ·〉 U endowed on the Stiefel manifold U is taken as the standard metric 〈· , ·〉 on R d × k . The Riemannian gradient of a smooth function f : R d × k → R at U ∈ U is defined as grad f ( U ) , which satisfies 〈 grad f ( U ) , ξ 〉 U = 〈∇ f ( U ) , ξ 〉 for all ξ ∈ T U U , where ∇ f ( U ) denotes the Euclidean gradient of f at U . If U T ∇ f ( U ) is symmetric, we have grad f ( U ) = Proj T U U ( ∇ f ( U )) = ( I d -UU T ) ∇ f ( U ) .

The rest of this paper is organized as follows. The proposed REALM is introduced in Section 2. A practical iRBBS for solving the subproblem in REALM is proposed in Section 3. Numerical results are 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 a fixed U ∈ U , consider the OT problem

By adding a redundant constraint ‖ π ‖ 1 = 1 [29, 36], we derive the dual of (4) as

where x = ( α, β, U ) and ϕ ( x ) ∈ R n × n with ϕ ( x ) ij = α i + β j + 〈 M ij , UU T 〉 for each ( i, j ) ∈ [ n ] × [ n ] . Note that the matrix π in (4) can also be understood as the Lagrange multiplier corresponding to the inequalities in (5). Therefore, the value P 2 k ( µ n , ν n ) defined in (1) is equal to the opposite of the optimal objective value of the following optimization problem:

where M = R n × R n ×U . Motivated by the first-order necessary condition of problem (6) (see Appendix A.1 for details), we define the ( glyph[epsilon1] 1 , glyph[epsilon1] 2 ) -stationary point of problem (1) as follows.

Definition 2.1. We call (˜ x , ˜ π ) ∈ M× Π( r, c ) an ( glyph[epsilon1] 1 , glyph[epsilon1] 2 ) -stationary point of PRW problem (1), if ‖ Proj T ˜ U U ( -2 V ˜ π ˜ U ) ‖ F ≤ glyph[epsilon1] 1 and 〈 ˜ π, Z (˜ x ) 〉 ≤ glyph[epsilon1] 2 , where Z (˜ x ) ∈ R n × n with Z (˜ x ) ij = ϕ ( x ) ij -ϕ ( x ) min . If glyph[epsilon1] 1 = glyph[epsilon1] 2 = 0 , we call such (˜ x , ˜ π ) a 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 ( glyph[epsilon1] 1 , glyph[epsilon1] 2 ) -stationary point satisfying the conditions here also satisfies all conditions therein. See Appendix A.2 for more details.

Given π ∈ R n × n ++ and η > 0 , define the function ζ η ( x , π ) ∈ R n × n with

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

For any fixed x , letting ∇ y ˜ L η ( x , y, 11 T ) = 0 , we can obtain the optimal y as y = η log( ‖ ζ η ( x , 11 T ) ‖ 1 ) . By eliminating the variable y in (9), we thus obtain the subproblem (3) of the approach in Huang et al. [28].

It is known that the exponential augmented Lagrangian method (ALM) is usually more stable than the exponential penalty approach; see [18, Tables 3.1-3.3] for a 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] to the manifold case to solve problem (6). Fix the current estimate of the Lagrange multiplier corresponding to the inequality constraints in (6) and the penalty parameter as π k and η k , respectively. Then the subproblem at the k -th iteration is given as

Similar to the way for eliminating y in (9), we obtain an equivalent formulation of (10):

Define the matrix φ η k ( x , π k ) ∈ R n × n with

By the chain rule, we have ∇ α L η k ( x , π k ) = r -φ η k ( x , π k ) 1 , ∇ β L η k ( x , π k ) = c -φ η k ( x , π k ) T 1 , ∇ U L η k ( x , π k ) = -2 V φ η k ( x ,π k ) U , and grad U L η k ( x , π k ) = Proj T U U ( -2 V φ η k ( x ,π k ) U ) . Let e 1 η k ( x , π k ) = ‖ grad U L η k ( x , π k ) ‖ F and e 2 η k ( x , π k ) = ‖∇ α L η k ( x , π k ) ‖ 1 + ‖∇ β L η k ( x , π k ) ‖ 1 . The ( glyph[epsilon1] 1 , glyph[epsilon1] 2 ) -stationary point of (11) and the connections of the approximate stationary points of problems (11) and (1) are given as follows.

Definition 2.3. We say ˜ x ∈ M an ( glyph[epsilon1] 1 , glyph[epsilon1] 2 ) -stationary point of problem (11) (with fixed η k and π k ) if e 1 η k (˜ x , π k ) ≤ glyph[epsilon1] 1 and e 2 η k (˜ x , π k ) ≤ glyph[epsilon1] 2 .

Theorem 2.4. Suppose ˜ x = (˜ α, ˜ β, ˜ U ) ∈ M with ‖ ζ η k (˜ x , π k ) ‖ 1 = 1 is an ( glyph[epsilon1] 1 , glyph[epsilon1] 2 ) -stationary point of problem (11) . Then, we have

where ˆ π := Round ( φ η k (˜ x , π k ) , Π( r, c )) is a feasible matrix returned by running the rounding procedure ' Round ' given in [3, Algorithm 2] with input φ η k (˜ x , π k ) .

define and

Remark 2.5 . By Theorem 2.4, we can see that, given any glyph[epsilon1] > 0 , for fixed π k , by choosing η k = O ( glyph[epsilon1] ) with glyph[epsilon1] 1 = glyph[epsilon1] 2 = glyph[epsilon1] , an glyph[epsilon1] -stationary point with bounded (˜ α, ˜ β ) (which can be found efficiently by iRBBS proposed in Section 3) of problem (11) can recover an glyph[epsilon1] -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 a feasible glyph[epsilon1] -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].

Let x 0 be an initial point satisfying r T α 0 = c T β 0 and ‖ ζ η 1 ( x 0 , π 1 ) ‖ 1 = 1 with π 1 = 11 T . We require that x k = ( α k , β k , U k ) , the ( glyph[epsilon1] k, 1 , glyph[epsilon1] k, 2 ) -stationary point of the subproblem (11) with k ≥ 1 , satisfies the following conditions:

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 k 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 L η k ( x , π k ) = L η k ( x s , π k ) where x s 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 k -1 and x 0 to solve problem (11).

Proposition 2.6. For any x = ( α, β, U ) ∈ M , consider x s = ( α + υ 1 1 , β + υ 2 1 , U ) ∈ M with υ 1 = ( c T β -r T α + η k log( ‖ ζ η k ( x , π k ) ‖ 1 )) / 2 and υ 2 = ( r T α -c T β + η k log( ‖ ζ η k ( x , π k ) ‖ 1 )) / 2 . Then we have that r T ( α + υ 1 1 ) = c T ( β + υ 2 1 ) and ‖ ζ η k ( x s , π k ) ‖ 1 = 1 and also that L η k ( x , π k ) = L η k ( x s , π k ) and ∇ x L η k ( x , π k ) = ∇ x L η k ( x s , π k ) .

With such x k in hand, we compute the candidate of the next estimate π k +1 as

Denote W k ∈ R n × n with W k ij = min { η k ˜ π k +1 ij , ϕ ( x k ) ij } . The penalty parameter η k +1 is updated according to the progress of the complementarity violation [4], denoted by ‖ W k ‖ F . If ‖ W k ‖ F ≤ γ W ‖ W k -1 ‖ F with γ W ∈ (0 , 1) , we keep η k +1 = η k and update π k +1 = ˜ π k +1 ; otherwise we keep π k +1 = π k and reduce η k +1 via

where γ η ∈ (0 , 1) is a constant and glyph[rho1] k → 0 with glyph[rho1] k > 0 .

We summarize the above discussion as the complete algorithm in Algorithm 1, whose convergence is established as follows.

Theorem 2.7. Let { ( x k , ˜ π k ) } be the sequence generated by Algorithm 1 with glyph[epsilon1] 1 = glyph[epsilon1] 2 = glyph[epsilon1] c = 0 and ( x ∞ , ˜ π ∞ ) be a limit point of { ( x k , ˜ π k ) } . Then, ( x ∞ , ˜ π ∞ ) is a stationary point of PRW problem (1) .

Remark 2.8 . Our proposed REALM is 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 (e.g., the one proposed in [20]: (i) 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 a 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.1 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.e., 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 η k in Algorithm 1 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.

Algorithm 1: REALM for solving problem (6).

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 k ‖ F and e k := 〈 ˆ π k , Z ( x k ) 〉 with ˆ π k := Round ( φ η k ( x k , π k ) , Π( r, c )) at the approximate stationary point of the subproblem; (ii) establishing the relationship between η k ˜ π k +1 ij and ϕ ( x k ) ij . Thanks to Theorem 2.4, we can slightly modify Algorithm 1 to establish the iteration complexity. By modifying the 'if' condition in Line 6 of Algorithm 1 as ' ‖ W k ‖ F ≤ γ W ‖ W k -1 ‖ F and e k ≤ γ W e k -1 ', 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 1 will terminate within at most O (max { log glyph[epsilon1] -1 1 , log glyph[epsilon1] -1 2 , T k } ) iterations, where T k := min { k | glyph[rho1] k ≤ glyph[epsilon1] c } .

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., ( α, β, y, 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.

3 A Practical iRBBS for Solving Subproblem (11)

At first glance, (11) 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:

By [34, Lemma 3.1], we know that q ( · ) is differentiable over R d × k . Here, we give a new formulation of grad q ( U ) , which can provide more insights into approximating grad q ( U ) .

Hence we could use the Riemannian gradient descent (RGD) method [2] to solve problem (11). Letting τ t > 0 be some stepsize, the main iterations of RGD are given as

However, RGD (18) needs to calculate ( α ∗ U , β ∗ U ) 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, 7, 52], we propose an inexact RGD framework. Let x t = ( α t , β t , U t ) with ( α t , β t ) ≈ ( α ∗ U t , β ∗ U t ) , wherein the inexactness level is determined by e 2 η k ( x t , π k ) = ‖∇ α L η k ( x t , π k ) ‖ 1 + ‖∇ β L η k ( x t , π k ) ‖ 1 ≤ θ t for given θ t . By Lemma 3.1, we use

to approximate grad q ( U t ) . Then, we perform an inexact RGD with grad q ( U t ) in (18) replaced by ξ t . More specifically, given the inexactness parameter θ t +1 ≥ 0 and the stepsize τ t ≥ 0 , we update x t +1 = ( α t +1 , β t +1 , U t +1 ) with U t +1 and ( α t +1 , β t +1 ) ≈ ( α ∗ U t +1 , β ∗ U t +1 ) satisfying

To make iRGD (19) practical, the first main ingredient is how to compute ( α t +1 , β t +1 ) such that (19b) holds. Given U t +1 ∈ U , α (0) = α t , and β (0) = β t , for glyph[lscript] = 0 , 1 , . . . , we adopt the block coordinate descent method to update

Note that α ( glyph[lscript] +1) and β ( glyph[lscript] +1) admit the closed-form solutions as follows:

α ( glyph[lscript] +1) = α ( glyph[lscript] ) -η k log r + η k log( ζ ( glyph[lscript] ) 1 ) , β ( glyph[lscript] +1) = β ( glyph[lscript] ) -η k log c + η k log(( ζ ( glyph[lscript] + 1 2 ) ) T 1 ) , (21) where ζ ( glyph[lscript] ) := ζ η k ( α ( glyph[lscript] ) , β ( glyph[lscript] ) , U t +1 , π k ) and ζ ( glyph[lscript] + 1 2 ) := ζ η k ( α ( glyph[lscript] +1) , β ( glyph[lscript] ) , U t +1 , π k ) (see (7) for its definition). Note that for fixed U t +1 and π k = 11 T , (21) reduces to the famous Sinkhorn iteration [15]. Therefore, we still name (21) as the Sinkhorn iteration. It is easy to verify that ‖ ζ ( glyph[lscript] + 1 2 ) ‖ 1 = ‖ ζ ( glyph[lscript] +1) ‖ 1 = 1 (see also [28, Remark 3.1]). By (12), we have

From the update of β ( glyph[lscript] +1) , we have ( π ( glyph[lscript] +1) ) T 1 -c = 0 . Therefore, to make condition (19b) hold, we stop the Sinkhorn iteration once

and set α t +1 = α ( glyph[lscript] +1) , β t +1 = β ( glyph[lscript] +1) . Recalling the calculation of the gradient and the definition of e 2 η k ( x t +1 , π k ) after (12), we have

Next, we choose the stepsize τ t in (19a). Since the accurate function and gradient information of q ( U t ) = L η k ( x ∗ U t , π k ) is unavailable, we cannot expect to build the linesearch condition based on q ( U t ) and need to find some appropriate potential function instead. Considering that L η k ( x t , π k ) is an approximation of q ( U t ) and the approximation error is controlled by e 2 η k ( x t , π k ) , it is thus desirable that some combinations of L η k ( x t , π k ) and e 2 η k ( x t , π k ) will be smaller than the corresponding values at the previous iteration. Given ρ ∈ [0 , η k / 2) , we define the potential function as

and require the stepsize τ t to satisfy the following nonmonotone line search condition:

where E r t +1 = ( γQ t E r t + E ρ ( x t +1 )) /Q t +1 and Q t +1 = γQ t +1 with a constant γ ∈ [0 , 1) and E r 0 = E ρ ( x 0 ) , Q 0 = 1 ; see [54]. Such τ t can be found by adopting the simple backtracking line search technique starting from an initial guess of the stepsize τ (0) t . Owing to the excellent performance of the BB method in Riemannian optimization [50, 32, 23, 25, 31, 24], we choose the initial guess τ (0) t for t ≥ 1 as a new Riemannian BB stepsize with safeguards:

where τ max > τ min > 0 are preselected stepsize safeguards and τ BB 1 = τ BB2 1 and for t ≥ 2 , we set τ BB t = min { τ BB2 t -1 , τ BB2 t , max { τ new t , 0 }} if τ BB2 t < κ t τ BB1 t and set τ BB t = τ BB1 t otherwise. Here, τ BB1 t = ‖ U t -U t -1 ‖ 2 F / |〈 U t -U t -1 , ξ t -ξ t -1 〉| , τ BB2 t = |〈 U t -U t -1 , ξ t -ξ t -1 〉| / ‖ ξ t -ξ t -1 ‖ 2 F , and τ new t is chosen according to [30, Eq. (2.15)]. In our numerical tests, we set the initial κ t to be 0.05 and update κ t +1 = κ t / 1 . 02 if τ BB2 t /τ BB1 t < κ t and update κ t +1 = 1 . 02 κ t otherwise.

We are ready to summarize the complete iRBBS in Algorithm 2. The overall complexity of Algorithm 2 to find an ( glyph[epsilon1] 1 , glyph[epsilon1] 2 ) -stationary point of problem (1) is in the same order as that of R(A)BCD.

Theorem 3.2. By choosing glyph[epsilon1] 1 = glyph[epsilon1] ′ 1 / 2 , glyph[epsilon1] 2 = min { glyph[epsilon1] ′ 1 / (4 ‖ C ‖ ∞ ) , glyph[epsilon1] ′ 2 / (4 η k Ψ+6 ‖ C ‖ ∞ ) } with Ψ = ‖ log π k ‖ var +max {‖ log r ‖ var , ‖ log c ‖ var } and η k = glyph[epsilon1] ′ 2 / (4 log n +2 ‖ log π k ‖ var ) , Algorithm 2 can return an ( glyph[epsilon1] ′ 1 , glyph[epsilon1] ′ 2 ) -stationary point of problem (1) in O ( T glyph[epsilon1] ′ 1 ,glyph[epsilon1] ′ 2 ) iterations with

If θ t ≥ 2 R t / ( η k ( glyph[lscript] max -2 + √ 2)) with glyph[lscript] max ≥ 1 and R t = ‖ C ( U t ) ‖ var + η k Ψ , the total number of Sinkhorn iterations is O ( glyph[lscript] max T glyph[epsilon1] ′ 1 ,glyph[epsilon1] ′ 2 ) and the total arithmetic operation complexity is O (( n 2 ( k + glyph[lscript] max ) + ndk + dk 2 ) T glyph[epsilon1] ′ 1 ,glyph[epsilon1] ′ 2 ) .

Algorithm 2: A practical iRBBS for solving problem (11).

Remark 3.3 . The basic idea of proposing Algorithm 2 is sharply different from that of R(A)BCD developed 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 -update in 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 Euclidean 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 be 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 use possible multiple Sinkhorn iterations rather than only one iteration as done in R(A)BCD in updating α and β from the computational point of view. The cost of updating α ( l +1) and β ( glyph[lscript] +1) via one Sinkhorn iteration (21) is O ( n 2 ) . In contrast, the cost of updating U t +1 via performing a RGD step (19a) is O ( ndk + n 2 k + dk 2 ) , wherein the main cost is to compute V φ k ( x t ,π k ) U t , which can be done by observing V π U = X Diag( π 1 ) X T U + Y Diag( π T 1 ) Y T U -XπY T U -Y π T X T U. Considering that the cost of updating α and β is much less than that of updating U , it is reasonable to update α and β multiple times and update U only once.

4 Experimental Results

In this section, we conduct numerical experiments on six Shakespeare operas to evaluate the performance of our proposed approaches; see Dataset C.1 for a 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://github.com/bjiangopt/ReALM .

4.1 Comparison with R(A)BCD on Solving Subproblem (11)

Subproblem (11) with π k = 11 T and a relatively small η k = η is used to compute the PRW distance in [28]. We choose η = 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)BCD. For iRBBS and R(A)BCD, we use the same stopping conditions with glyph[epsilon1] 1 = 2 ‖ C ‖ ∞ glyph[epsilon1] 2 (motivated by (13a)) and glyph[epsilon1] 2 = 10 -6 max {‖ r ‖ ∞ , ‖ c ‖ ∞ } . To make the residual error more comparable, we choose

where the parameter θ is a preselected constant in the proposed iRBBS algorithm. We consider several different values of θ and denote the corresponding version of iRBBS as iRBBSθ . When θ = 0 , we have θ t +1 ≡ glyph[epsilon1] 2 for all t ≥ 0 , which means that we calculate grad q ( U t ) almost exactly; when θ = + ∞ , 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]

where glyph[circledivide] denotes elementwise division, u ( glyph[lscript] ) = exp( -α ( glyph[lscript] ) /η k ) , v ( glyph[lscript] ) = exp( -β ( glyph[lscript] ) /η k ) , and A ∈ R n × n with A ij = π k ij exp( -〈 M ij , UU T 〉 /η k ) . More settings and details on the implementation can be found in Appendix C.

For each instance of Dataset C.1, we randomly generate 20 initial points and report the average performance in Table 1. The term 'nGrad' means the total number of calculating grad U L η k ( x t , π k ) , '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 t returned by some method in hand, we invoke Gurobi 9.5/Mosek 9.3 to compute a more accurate PRW distance, denoted as ̂ P 2 k = min π ∈ Π( r,c ) 〈 π, C ( U t ) 〉 . Note that a larger ̂ P 2 k means a higher solution quality of the corresponding U t . In the last 'A VG' 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 is about 2.8x faster than iRBBSinf and about 1.4x faster than iRBBS-0; (iii) compared with R(A)BCD, iRBBS-0.1 can always take significantly less nGrad and may take a 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 R(A)BCD. More importantly, our methods adopt the adaptive stepsize without needing to tune the best stepsize as done in R(A)BCD.

4.2 Comparison on Computing the PRW Distance (1)

In this subsection, we present numerical results to illustrate the effectiveness and efficiency of our proposed REALM, namely, Algorithm 1. We choose glyph[epsilon1] 1 = 2 ‖ C ‖ ∞ glyph[epsilon1] 2 , glyph[epsilon1] 2 = 10 -6 max {‖ r ‖ ∞ , ‖ c ‖ ∞ } , glyph[epsilon1] c = 10 -3 , and η 1 = 20 . To avoid possible numerical instability, we restrict the maximum number of updating π k as 8. To prevent too small η k , in our implementation, we update η k +1 = min { γ η η k , η min } other than using (16), where η min is a preset positive number. We set γ η = γ glyph[epsilon1] = 0 . 25 . Moreover, we denote by REALM-( η min , γ W ) the REALM with particular parameters η min and γ W . Note that choosing γ W = 0 means that we adopt a continuation technique to solve (11) with π k ≡ π 1 and η k = η min , which is generally better than directly solving a single (11) with

η k = η min . Subproblem (11) in REALM is solved by iRBBS. If max {‖ r ‖ ∞ , ‖ c ‖ ∞ } ≥ 500 η k or ‖ C ( U t ) -η k log π k ‖ var ≥ 900 η k , we set θ = 10 and perform the Sinkhorn iteration (21); otherwise, we set θ = 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 'nSk log ' and 'nSk exp ' mean the total numbers of Sinkhorn iterations (21) and (28), respectively, the pair ' k 1 / k ' in the column 'iter' means that the corresponding algorithm stops at the k -iteration and updates the multiplier matrix k 1 times. To save space, we only report instances where one method can return the value ' ̂ P 2 k ' larger than 1.005 times of the smaller one of the two ̂ P 2 k values returned by the two methods. The better ' ̂ P 2 k ' is marked in bold . Besides, the average performance over all 15 instances is also kept in the 'A VG' line.

From Table 2, we can observe that REALM(0 . 007 , 0 . 9) can not only return better solutions than REALM(0 . 0035 , 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 η k . Second, solving the subproblem with a larger η k is always easier, which enables that REALM(0 . 07 , 0 . 9) computes less grad U L η k ( x t , π k ) and performs less Sinkhorn iterations (21) which involves computing the log-sum-exp function log ∑ i exp( x i /η k ) = x max /η k +log ∑ i exp(( x i -x max ) /η k ) for small η k .

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 a 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 glyph[epsilon1] -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-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 minimax 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, a limitation of our work is that we did not establish the positive lower bound of η k in REALM, despite the fact that REALM performs well in practice and can avoid too small η k in many cases. We shall investigate the conditions under which it is possible to establish a lower bound of η k in REALM. This can be achievable by extending the analysis and conditions in [20] to the Riemannian case.

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 Science Foundation of the Higher Education Institutions of Jiangsu Province under Grant 21KJA110002. The work of Ya-Feng Liu was supported in part by NSFC under Grant 11991020, Grant 11991021, and Grant 12288201.

References

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):A3340A3366, 2020. (Cited on page 1.)

P.-A. Absil, R. Mahony, and R. Sepulchre. Optimization Algorithms on Matrix Manifolds . Princeton University Press, Princeton, 2008. (Cited on pages 3 and 6.)

J. Altschuler, J. Niles-Weed, and P. 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.)

R. Andreani, E. G. Birgin, J. M. Martínez, and M. L. Schuverdt. On augmented Lagrangian methods with general lower-level constraints. SIAM Journal on Optimization , 18(4):1286-1309, 2008. (Cited on page 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.)

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):A1111A1138, 2015. (Cited on page 9.)

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-1518, 2021. (Cited on pages 6 and 8.)

D. P. Bertsekas. Constrained Optimization and Lagrange Multiplier Methods . Academic Press, 1982. (Cited on page 4.)

N. Boumal. An Introduction to Optimization on Smooth Manifolds . Cambridge University Press, 2023. (Cited on page 3.)

N. Boumal, P.-A. Absil, and C. Cartis. Global rates of convergence for nonconvex optimization on manifolds. IMA Journal of Numerical Analysis , 39(1):1-33, 2019. (Cited on page 19.)

R. G. Carter. On the global convergence of trust region algorithms using inexact gradient information. SIAM Journal on Numerical Analysis , 28(1):251-265, 1991. (Cited on pages 6 and 8.)

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 , 4(4):1369-1395, 2022. (Cited on page 14.)

X. Chen, L. Guo, Z. Lu, and J. J. Ye. An augmented Lagrangian method for non-Lipschitz nonconvex programming. SIAM Journal on Numerical Analysis , 55(1):168-193, 2017. (Cited on page 15.)

A. Cherian and S. Aeron. Representation learning via adversarially-contrastive optimal transport. In International Conference on Machine Learning , pages 1820-1830. PMLR, 2020. (Cited on page 1.)

M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in Neural Information Processing Systems , 26:2292-2300, 2013. (Cited on pages 7 and 9.)

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 its use for GANs. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 10648-10656, 2019. (Cited on pages 1 and 2.)

O. Devolder, F. Glineur, and Y. Nesterov. First-order methods of smooth convex optimization with inexact oracle. Mathematical Programming , 146(1):37-75, 2014. (Cited on pages 6 and 8.)

J.-P. Dussault. Augmented non-quadratic penalty algorithms. Mathematical programming , 99 (3):467-486, 2004. (Cited on page 4.)

P. Dvurechensky, A. Gasnikov, and A. Kroshnin. Computational optimal transport: Complexity by accelerated gradient descent is better than by Sinkhorn's algorithm. In International Conference on Machine Learning , pages 1367-1376. PMLR, 2018. (Cited on page 19.)

N. Echebest, M. D. Sánchez, and M. L. Schuverdt. Convergence results of an augmented Lagrangian method using the exponential penalty function. Journal of Optimization Theory and Applications , 168(1):92-108, 2016. (Cited on pages 4, 5, and 10.)

R. Flamary, N. Courty, D. Tuia, and A. Rakotomamonjy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence , 39:1853-1865, 2016. (Cited on page 1.)

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.)

B. Gao, X. Liu, X. Chen, and Y. Yuan. A new first-order algorithmic framework for optimization problems with orthogonality constraints. SIAM Journal on Optimization , 28(1):302-332, 2018. (Cited on page 7.)

B. Gao, X. Liu, and Y. Yuan. Parallelizable algorithms for optimization problems with orthogonality constraints. SIAM Journal on Scientific Computing , 41(3):A1949-A1983, 2019. (Cited on page 7.)

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.)

Y.-Q. Hu and Y.-H. Dai. Inexact Barzilai-Borwein method for saddle point problems. Numerical Linear Algebra with Applications , 14(4):299-317, 2007. (Cited on pages 6 and 8.)

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.)

M. Huang, S. Ma, and L. Lai. A Riemannian block coordinate descent method for computing the projection robust 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.)

M. Huang, S. Ma, and L. Lai. A Riemannian block coordinate descent method for computing the projection robust Wasserstein distance. arXiv:2012.05199, v5 , 2021. (Cited on pages 2, 3, 17, 18, and 19.)

Y.-K. Huang, Y.-H. Dai, and X.-W. Liu. Equipping the Barzilai-Borwein method with the two dimensional quadratic termination property. SIAM Journal on Optimization , 31(4):3068-3096, 2021. (Cited on page 7.)

B. Iannazzo 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-517, 2018. (Cited on page 7.)

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.)

S. Kullback. A lower bound for discrimination information in terms of variation (corresp.). IEEE Transactions on Information Theory , 13(1):126-127, 1967. (Cited on page 18.)

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-9397, 2020. (Cited on pages 2, 3, 4, 6, 8, 14, 17, 21, and 22.)

T. 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 262-270. PMLR, 2021. (Cited on page 2.)

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):1-42, 2022. (Cited on pages 3 and 19.)

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.)

H. Liu, A. M.-C. So, and W. Wu. Quadratic optimization with orthogonality constraint: Explicit Łojasiewicz exponent and linear convergence of retraction-based line-search and stochastic variance-reduced gradient methods. Mathematical Programming , 178(1):215-262, 2019. (Cited on page 19.)

Y.-F. Liu, X. Liu, and S. Ma. On the nonergodic convergence rate of an inexact augmented Lagrangian framework for composite convex programming. Mathematics of Operations Research , 44(2):632-650, 2019. (Cited on pages 6 and 8.)

J. Niles-Weed and P. Rigollet. Estimation of Wasserstein distances in the spiked transport model. Bernoulli , 28(4):2663-2688, 2022. (Cited on pages 1 and 2.)

F.-P. Paty and M. Cuturi. Subspace robust Wasserstein distances. In International Conference on Machine Learning , pages 5072-5081. PMLR, 2019. (Cited on pages 1, 2, and 22.)

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.)

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.)

E. W. Sachs and S. M. Sachs. Nonmonotone line searches for optimization algorithms. Control and Cybernetics , 40(4):1059-1075, 2011. (Cited on page 20.)

S. Samadi, U. Tantipongpipat, J. 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.)

M. Schmidt, N. Roux, and F. Bach. Convergence rates of inexact proximal-gradient methods for convex optimization. Advances in Neural Information Processing Systems , 24:1458-1466, 2011. (Cited on pages 6 and 8.)

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.)

J.-P. Vial. Strong and weak convexity of sets and functions. Mathematics of Operations Research , 8(2):231-259, 1983. (Cited on page 14.)

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.)

Z. Wen and W. Yin. A feasible method for optimization with orthogonality constraints. Mathematical Programming , 142(1):397-434, 2013. (Cited on page 7.)

Y. Xie, X. Wang, R. Wang, and H. 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.)

L. Yang and K.-C. Toh. Bregman proximal point algorithm revisited: A new inexact version and its variant. SIAM Journal on Optimization , 32(3):1523-1554, 2022. (Cited on pages 4, 5, 6, and 8.)

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.)

H. Zhang and W. W. Hager. A nonmonotone line search technique and its application to unconstrained optimization. SIAM Journal on Optimization , 14(4):1043-1056, 2004. (Cited on pages 7 and 20.)

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-61, 2023. (Cited on page 6.)

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 to 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 ( glyph[epsilon1] 1 , glyph[epsilon1] 2 ) -stationary points of the PRW distance problem (1), 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 ¯ x ∈ M and ¯ y ∈ R , suppose that (¯ x , ¯ y ) is a local minimizer of problem (6) , then (¯ x , -ϕ (¯ x ) min ) is a stationary point of problem (6) , namely, there exists ¯ π ∈ Π( r, c ) such that

where Z (¯ x ) ∈ R n × n is defined as

Proof. Since (¯ x , ¯ y ) is a local minimizer of problem (6), there must hold ¯ y = -ϕ (¯ x ) min = max ij {-ϕ (¯ x ) ij } . Moreover, such ¯ x is also a local minimizer of problem min x ∈M h ( x ) with h ( x ) = r T α + c T β -ϕ ( x ) min . For fixed α and β , it is easy to see that -ϕ ( x ) ij + ‖ C ‖ ∞ ‖ U ‖ 2 F is convex with respect to U since ‖ M ij ‖ 2 = C ij ≤ ‖ C ‖ ∞ . We thus know that the function

is convex with respect to U , which means that the function -ϕ (¯ x ) min and thus h ( x ) is ‖ C ‖ ∞ -weakly convex with respect to U [48, Proposition 4.3]. Let A (¯ x ) = { ( i, j ) ∈ [ n ] × [ n ] | ϕ (¯ x ) ij = ϕ (¯ x ) min } . By [48, Proposition 4.6], we have

where e i is the i -th standard unit vector in R n . Moreover, by [53, Theorem 4.1], [53, Theorem 5.1], and T ¯ x M = R n × R n × T ¯ x U , there must hold that 0 ∈ Proj T ¯ x M ∂h (¯ x ) . Putting all the above things together shows that there exists ¯ π ∈ R n × n + with ¯ π ij = 0 for all ( i, j ) glyph[negationslash]∈ A (¯ x ) and ∑ ( i,j ) ∈A (¯ x ) ¯ π i,j = 1 such that ¯ π ∈ Π( r, c ) and (29) hold. The proof is completed.

A.2 Explanations on Remark 2.2

From [28, Section B], we know that the ( glyph[epsilon1] 1 , glyph[epsilon1] 2 ) -stationary point in [28, Definition 3.1] is stronger than that in [34, Definition 2.7]. To show that our Definition 2.1 of the ( glyph[epsilon1] 1 , glyph[epsilon1] 2 ) -stationary point is stronger than both definitions, we only need to verify that 〈 ˜ π, Z (˜ x ) 〉 ≤ glyph[epsilon1] 2 can imply

which is clear since 〈 ˜ π, C ( ˜ U ) 〉 -min π ∈ Π( r,c ) 〈 π, C ( ˜ U ) 〉 ≤ 〈 ˜ π, Z (˜ x ) 〉 . This inequality comes from the fact that problem (5) is a dual formulation of min π ∈ Π( r,c ) 〈 π, 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 ˆ π ∈ Π( r, c ) and

Second, by the triangular inequality and the non-expansive property of the projection operator, we have

where A 1 , A 2 ∈ R d × k . Hence, we have

where the first inequality uses ‖ Proj T ˜ U U ( -2 V φ η k (˜ x ,π k ) ˜ U ) ‖ F ≤ glyph[epsilon1] 1 and the second inequality is due to the fact that ‖ A ˜ U ‖ F ≤ ‖ A ‖ F for any matrix A ∈ R d × d . Moreover, observe that

Combining the above assertions with (31) and (32) yields (13a).

Next, we show that (13b) is also true. By the Cauchy-Schwarz inequality and (31), we have

where the second inequality uses (30), the definition of ϕ ( x ) in (5), and

The remaining is to bound 〈 φ η k (˜ x , π k ) , Z (˜ x ) 〉 . By ‖ ζ η k (˜ x , π k ) ‖ 1 = 1 , (7), and (12), we have ϕ (˜ x ) ij = η k (log π k ij -log φ η k (˜ x , π k ) ij ) and ϕ (˜ x ) min ≥ η k (log π k ) min . Again with (30), we have

which, together with H ( φ η k (˜ x , π k )) ≤ 2 log n and 〈 log π k , φ η k (˜ x , π k ) 〉 ≤ (log π k ) max , further implies 〈 φ η k (˜ x , π k ) , Z (˜ x ) 〉 ≤ (2 log n + ‖ log π k ‖ var ) η k . 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 r T α k ≤ r T α 0 for all k ≥ 0 .

With the first two requirements in (14), we have L η k ( x k , π k ) = 2 r T α k . By the fact ‖ ζ η 1 ( x 0 , π 1 ) ‖ 1 = 1 with π 1 = 11 T , we know from (7) that ϕ ( x 0 ) ij ≥ 0 . Due to the update rules of π k and η k , we know that 0 < π k ij ≤ 1 and 0 < η k ≤ η 1 for all k ≥ 1 . Therefore, we have from (7) that

which further implies ‖ ζ η k ( x 0 , π k ) ‖ 1 ≤ ‖ ζ η 1 ( x 0 , π 1 ) ‖ 1 = 1 . With r T α 0 = c T β 0 , we know L η k ( x 0 , π k ) ≤ 2 r T α 0 . By the third requirement in (14), we further have

Besides, since x k is an ( glyph[epsilon1] k, 1 , glyph[epsilon1] k, 2 ) -stationary point of (11), we have

We next consider two cases, which is motivated by the proofs of Theorem 3.1 in [13].

Case (i). η k is bounded below by some threshold value η > 0 . Due to the update rule of η k +1 , we can see that (16) is invoked only finite times. Besides, due to 0 < π k ij ≤ 1 , without loss of generality,

we assume η k ≡ η for all k ≥ 2 and lim k →∞ π k ij = π ∞ ij . Meanwhile, due to the update rule of π k , we have π k ij = ˜ π k ij for k ≥ 2 and

We now show that { x k } has at least a limit point. By (37), we must have ϕ ( x k ) ij → 0 for each ( i, j ) glyph[negationslash]∈ A ( π ∞ ) := { ( i, j ) ∈ [ n ] × [ n ] | π ∞ ij = 0 } and there must exist K 1 > 0 such that for all k ≥ K 1 , there holds ϕ ( x k ) ij ≥ -1 for each ( i, j ) ∈ A ( π ∞ ) . Recalling the definition of ϕ ( x ) in (5), then we can conclude that there exists K 2 > 0 such that for all k ≥ K 2 , there holds that

Multiplying both sides of the above assertions by c j and then summing the obtained inequality from j = 1 to n , we have

where the second inequality is due to r T α k = c T β k , (35), and (34). Combining (35) and (38), we further have

Similarly, we can establish a similar bound for each β k j . Recalling that U k is in a compact set, it is thus safe to say that the sequence { x k } has at least one limit point, denoted by x ∞ = lim k ∈K ,k →∞ x k . Again with lim k →∞ π k ij = π ∞ , we have from (37) that min { ηπ ∞ ij , ϕ ( x ∞ ) ij } = 0 , which, together with π ∞ ij ≥ 0 and the fact

further implies ϕ ( x ∞ ) min = 0 and 〈 π ∞ , Z ( x ∞ ) 〉 = 0 . This shows that ( x ∞ , 0) is feasible. Moreover, letting k ∈ K go to infinity in (36), we further have π ∞ ∈ Π( r, c ) and ‖ Proj T U ∞ U ( -2 V π ∞ U ∞ ) ‖ F = 0 . From Lemma A.1, we know that ( x ∞ , 0) is a stationary point of problem (6) and ( x ∞ , π ∞ ) , namely, ( x ∞ , ˜ π ∞ ) , is a stationary point of problem (1).

Case (ii). The sequence { η k } is not bounded below by any positive number, namely, lim k →∞ η k = 0 . By the updating rule, we know that η k is updated via (16) infinitely many times. Hence, there must exist k 1 < k 2 < · · · such that η k glyph[lscript] → 0 as glyph[lscript] →∞ and

By (7) and the second assertion in (14), we have π k glyph[lscript] ij exp( -ϕ ( x k glyph[lscript] ) ij /η k glyph[lscript] ) ≤ 1 , which, together with 0 < π k glyph[lscript] ij ≤ 1 and (16), implies

Using the same arguments as in the proof in Case (i), we can show that { x k } is bounded over { k 1 , k 2 , . . . } and thus { x k } has at least a limit point. Without loss of generality, assume lim glyph[lscript] →∞ x k glyph[lscript] = x ∞ . By (39) and glyph[rho1] k → 0 in (16), we know that ϕ ( x ∞ ) ≥ 0 and thus that ( x ∞ , 0) is feasible to problem (6).

glyph[negationslash]

Due to the compactness of ˜ π k in (15), there must exist a subset K ⊆ { k 1 , k 2 , . . . } such that lim k ∈K ,k →∞ ˜ π k ij = ˜ π ∞ . Recalling (15), (7), and the second requirement in (14), we have ˜ π k +1 ij = π k ij exp( -ϕ ( x k ) ij /η k ) for each k ≥ 1 . Let A ( x ∞ ) = { ( i, j ) ∈ [ n ] × [ n ] | ϕ ( x ∞ ) ij = 0 } . We claim A ( x ∞ ) = ∅ . Otherwise, for every ( i, j ) we have ϕ ( x ∞ ) ij > 0 and thus

where the equality uses lim glyph[lscript] →∞ x k glyph[lscript] = x ∞ , K ⊆ { k 1 , k 2 , . . . } , and lim k →∞ η k → 0 . This makes a contradiction with ‖ ˜ π k +1 ‖ 1 = 1 . Moreover, (40) also further implies that 〈 ˜ π ∞ , Z ( x ∞ ) 〉 = 0 . Finally, with (36), we further have ‖ Proj T U ∞ U ( -2 V ˜ π ∞ U ∞ ) ‖ F = 0 and ˜ π ∞ ∈ Π( r, c ) . Putting the above things together, we know that that ( x ∞ , 0) is a stationary point of (6) and ( x ∞ , ˜ π ∞ ) is a stationary point of problem (1). The proof is completed.

B Proofs in Section 3

B.1 Proof of Lemma 3.1

For notational simplicity, we define

It is easy to see that

Therefore, we can write

Since the entropy function is 1-strongly convex with respect to the glyph[lscript] 1 -norm over the probability simplex, the minimization problem in (43) has a unique solution π ∗ U = argmin π ∈ Π( r,c ) h ( π, U ) . Using [34, Lemma B.1], we thus have

By (5), (7), (12), and recalling [ C ( U )] ij = 〈 M ij , UU T 〉 , we have

By the optimality of α ∗ U and β ∗ U , we know that e 1 η k ( x ∗ U , π k ) = e 2 η k ( x ∗ U , π k ) = 0 and thus φ η k ( x ∗ U , π k ) 1 = r and φ η k ( x ∗ U , π k ) T 1 = c . With (45), (41), and the definition of L η k ( · , · ) in (11), by some calculations, we have

By the dual formulation (42), we have -L η k ( x ∗ U , π k ) ≤ h ( π ∗ U , U ) . Hence, we know that h ( φ η k ( x ∗ U , π k ) , U ) ≤ h ( π ∗ U , U ) , which, together with the optimality and uniqueness of π ∗ U , yields φ η k ( x ∗ U , π k ) = π ∗ U . 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 is 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.11] since we can establish the non-increasing property of the feasibility violation of π ( l ) in Lemma B.1.

Lemma B.1. Let { ( α glyph[lscript] , β glyph[lscript] ) } be the sequence generated by (21) and consider { π ( glyph[lscript] ) } and { π ( glyph[lscript] + 1 2 ) } in (22) . Then we have ‖ π (0) 1 -r ‖ 1 + ‖ ( π (0) ) T 1 -c ‖ 1 ≥ ‖ ( π ( 1 2 ) ) T 1 -c ‖ 1 ≥ ‖ π (1) 1 -r ‖ 1 and for each glyph[lscript] ≥ 1 ,

Proof. By the optimality of α ( glyph[lscript] +1) in (20) and the first equation in (22), we have π ( glyph[lscript] + 1 2 ) 1 -r = 0 for glyph[lscript] ≥ 0 . Therefore, we have

where the first inequality uses the Cauchy-Schwarz inequality, the second equality comes from (7), (22), and the definitions of ζ ( glyph[lscript] ) and ζ ( glyph[lscript] + 1 2 ) after (21), and the third equality is due to the second equation in (21) and the first equation in (22). This proves ‖ ( π ( glyph[lscript] + 1 2 ) ) T 1 -c ‖ 1 ≥ ‖ π ( glyph[lscript] +1) 1 -r ‖ 1

for glyph[lscript] ≥ 0 . On the other hand, by the optimality of β ( glyph[lscript] +1) in (20) and the second equation in (22), we have ( π ( glyph[lscript] ) ) T 1 -c = 0 for glyph[lscript] ≥ 1 . Using a similar argument, we can prove ‖ π ( glyph[lscript] ) 1 -r ‖ 1 ≥ ‖ ( π ( glyph[lscript] + 1 2 ) ) T 1 -c ‖ 1 for glyph[lscript] ≥ 1 and ‖ π (0) 1 -r ‖ 1 + ‖ ( π (0) ) T 1 -c ‖ 1 ≥ ‖ ( π ( 1 2 ) ) T 1 -c ‖ 1 . The proof is completed.

Lemma B.2 (Sufficient decrease of L η k ( · , · ) in ( α, β ) ) . Let { ( α glyph[lscript] , β glyph[lscript] ) } be the sequence generated by (21) and consider { π ( glyph[lscript] ) } in (22) . Then, for each glyph[lscript] ≥ 0 , we have

Proof. By [29, Lemma 4.11], for glyph[lscript] ≥ 0 , we have

By Lemma B.1, for glyph[lscript] ≥ 0 , we further have

On the other hand, we have

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, q ∈ ∆ n , the Kullbbback-Leibler divergence between p and q is KL ( p ‖ q ) = ∑ i p i log( p i /q i ) . Combining (48) and (49), we obtain (47). The proof is completed.

We need to use the following elementary result.

Lemma B.3. Let ϑ 1 and ϑ 2 be two given positive constants. Consider two sequences { a glyph[lscript] } , { b glyph[lscript] } ⊆ R + with glyph[lscript] ≥ 0 . If they obey:

for all glyph[lscript] ≥ 1 , then we have b 1 ≤ ϑ 2 /ϑ 1 and

Proof. Applying (50) on glyph[lscript] = 1 , it is easy to see that ϑ 2 b 1 ≥ a 1 ≥ ϑ 1 b 2 1 . Therefore, we have b 1 ≤ ϑ 2 /ϑ 1 and a 1 ≤ ϑ 2 2 /ϑ 1 . For glyph[lscript] ≥ 1 , by (50a), we have

which, together with a 1 ≤ ϑ 2 2 /ϑ 1 , implies

Let glyph[floorleft] glyph[lscript]/ 2 glyph[floorright] be the largest integer smaller than glyph[lscript]/ 2 . For glyph[lscript] ≥ 1 , summing (50a) over i = glyph[floorleft] glyph[lscript]/ 2 glyph[floorright] +1 , . . . , glyph[lscript] , we have

which immediately implies min 1 ≤ i ≤ l +1 b 2 i ≤ ϑ 2 2 /ϑ 2 1 glyph[lscript] ( glyph[lscript] +1) . With the fact that glyph[lscript] ( glyph[lscript] +1) ≥ ( glyph[lscript] + √ 2 -1) 2 for glyph[lscript] ≥ 1 and b 1 ≤ ϑ 2 /ϑ 1 , we thus complete the proof.

With the help of the above lemmas, we can provide a new analysis of the iteration complexity of the classical Sinkhorn iteration based on the decreasing properties developed in Lemma B.1 and B.2. It differs from the proof technique in [19], wherein a switching strategy is adopted to establish the complexity.

Lemma B.4. The total number of Sinkhorn iterations to find a point π ( glyph[lscript] +1) satisfying (23) is at most glyph[ceilingleft] 2 R t +1 / ( η k θ t +1 ) + 2 -√ 2 glyph[ceilingright] , where

where Ψ = ‖ log π k ‖ var +max {‖ log r ‖ var , ‖ log c ‖ var } . Here, glyph[ceilingleft] a glyph[ceilingright] is the smallest nonnegative integer larger than a ∈ R n .

Proof. Let ( α ∗ U t +1 , β ∗ U t +1 ) ∈ argmin α ∈ R n ,β ∈ R n L η k ( α, β, U t +1 , π k ) . By refining the proof in [19] and [36], we can prove

and L η k ( α, β, U t +1 , π k ) is jointly convex with respect to α and β , we have

Given x ∈ R n , let x m = ‖ x ‖ var / 2 + x min , it holds that ‖ x -x m 1 ‖ ∞ = ‖ x ‖ var / 2 . For y ∈ R n with 〈 y, 1 〉 = 0 , we further know that

Applying this assertion with y = π ( glyph[lscript] +1) 1 -r , x = α ∗ U t +1 or x = α ( glyph[lscript] +1) , we obtain from ‖ π ( glyph[lscript] +1) ‖ 1 = ‖ r ‖ 1 = 1 , (52), and (53) that

Applying Lemma B.3 with a glyph[lscript] = L η k ( α ( glyph[lscript] ) , β ( glyph[lscript] ) , U t +1 , π k ) - L η k ( α ∗ U t +1 , β ∗ U t +1 , U t +1 , π k ) and b glyph[lscript] = ‖ π ( glyph[lscript] ) 1 -r ‖ 1 and using (46), (47), and (54), we have ‖ π (1) 1 -r ‖ 1 ≤ 2 R t +1 /η k and ‖ π ( glyph[lscript] +1) 1 -r ‖ 1 ≤ 2 R t +1 / ( η k ( glyph[lscript] + √ 2 -1)) for all glyph[lscript] ≥ 1 . Letting 2 R t +1 / ( η k ( glyph[lscript] + √ 2 -1)) ≤ θ t +1 , 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 t } be the sequence generated by Algorithm 2. If glyph[epsilon1] 1 = glyph[epsilon1] 2 = 0 , we have e 1 η k ( x t , π k ) → 0 and e 2 η k ( x t , π k ) → 0 . If glyph[epsilon1] 1 > 0 and glyph[epsilon1] 2 > 0 , then Algorithm 2 stops within at most glyph[ceilingleft] Υmax { glyph[epsilon1] -2 1 , glyph[epsilon1] -2 2 }glyph[ceilingright] iterations for any θ t +1 ≥ 0 , where Υ is a constant defined in (64) .

Proof. Given α ∈ R n and β ∈ R n , by [29, Lemma 4.8] for any ι ∈ [0 , 1] and U 1 , U 2 ∈ U ,

From [38, 10], we know that the retraction on the Stiefel manifold has the following nice properties, namely, there exist positive constants L 1 and L 2 such that

hold for all ξ ∈ T U U and U ∈ U . With (55) and (56), following the proof of Lemma 3 in [10], for each U ∈ U and ξ ∈ T U U , we have

where

Applying (57) with α = α t , β = β t , U = U t , and ξ = -τ t grad U L η k ( x t , π k ) = -τ t ξ t , with (19a), we have

In addition, by the fact that α t +1 = α ( glyph[lscript] +1) , β t +1 = β ( glyph[lscript] +1) , and e 2 η k ( x t +1 , π k ) = ‖ π ( glyph[lscript] +1) 1 -r ‖ 1 for some glyph[lscript] and α (0) = α t , β (0) = β t , we have from (47) that

Recalling (24), it follows from (59) and (60) that

holds with ρ ∈ [0 , η k / 2) .

Recalling that τ t takes the form τ t = τ (0) t σ s for some nonnegative integer s , with (26), we know from (61) that if s = ⌈ log 2(1 -δ 1 ) τ max L / log σ ⌉ , there must hold that

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 -δ 1 ) τ max L / log σ ⌉ trials and E r t ≥ E ρ ( x t ) . Moreover, since τ (0) t ≥ τ min from (26), we further know that τ t = τ (0) t σ s ≥ τ := 2 σ (1 -δ 1 ) τ min τ max L . Therefore, we know that the nonmonotone linesearch condition (25) is satisfied for some τ t ≥ τ .

Let x ∗ = ( α ∗ , β ∗ , U ∗ ) be the minimizer of problem (11). Using a similar argument as [28, Lemma 4.7], by (34) and -log π k ij ≤ ‖ log π k ‖ ∞ , we have

Let ω t = E r t -E ρ ( x t ) . By (62) and using the argument in [44, Theorem 5], we know that such ω t ≥ 0 satisfies

By (24), the fact that (25) holds with τ t ≥ τ , and ‖ ξ t ‖ F = e 1 η k ( x t , π k ) , we have

Summing (63) over t = 0 , . . . , T and adding the term ( η k / 2 -ρ )( e 2 η k ( x 0 )) 2 on both sides of the obtained equation, and then combining them with L η k ( x T +1 , π k ) ≥ L η k ( x ∗ , π k ) and (62), we have

for any T ≥ 1 . This further implies

Hence, we have e 1 η k ( x t , π k ) → 0 and e 2 η k ( x t , π k ) → 0 as t →∞ .

Suppose e 1 η k ( x t , π k ) ≤ glyph[epsilon1] 1 and e 2 η k ( x t , π k ) ≤ glyph[epsilon1] 2 are fulfilled for t = T but not fulfilled for all t < T . Then there hold that e 1 η k ( x T , π k ) ≤ glyph[epsilon1] 1 , e 2 η k ( x T , π k ) ≤ glyph[epsilon1] 2 , and ( e 1 η k ( x t , π k )) 2 +( e 2 η k ( x t , π k )) 2 > min { glyph[epsilon1] 2 1 , glyph[epsilon1] 2 2 } for all t < T . Setting T in (64) as T yields T min { glyph[epsilon1] 2 1 , glyph[epsilon1] 2 2 } ≤ Υ , which completes the proof.

We are now ready to complete the proof of Theorem 3.2.

Proof of Theorem 3.2. Given glyph[epsilon1] 1 > 0 and glyph[epsilon1] 2 > 0 , by Theorem B.5, we know that Algorithm 2 will return a point x T with T ≤ glyph[ceilingleft] Υmax { glyph[epsilon1] -2 1 , glyph[epsilon1] -2 2 }glyph[ceilingright] satisfying satisfying e 1 η k ( x t , π k ) ≤ glyph[epsilon1] 1 and e 2 η k ( x t , π k ) ≤ glyph[epsilon1] 2 . Let ˆ π T := Round ( φ η k ( x T , π k ) , Π( r, c )) , where ' Round ' is the rounding procedure given in [3, Algorithm 2]. By the second equation in (22) and x T = ( α T , β T , U T ) with α T = α ( glyph[lscript] +1) , β T = β ( glyph[lscript] +1) for some glyph[lscript] , we know that ‖ ζ η k ( x T , π k ) ‖ 1 = 1 . By ‖ C ( U ) ‖ ∞ ≤ ‖ C ‖ ∞ for all U ∈ U , (58), (52), (51), we have ‖ α T ‖ var + ‖ β T ‖ var + ‖ C ‖ ∞ ≤ 3 ‖ C ‖ ∞ +2 η k Ψ . Therefore, we obtain from (13) that

With the choices of η k , glyph[epsilon1] 1 and glyph[epsilon1] 2 in Theorem 3.2, we can immediately know from (13) that such ( x T , -ϕ ( x T ) min ) is an ( glyph[epsilon1] ′ 1 , glyph[epsilon1] ′ 2 ) -stationary point of problem (6). With the expression of Υ 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 η k / 2 ( x 0 ) , Ψ , L 1 , L 2 , log n , and ‖ C ‖ ∞ .

C Experimental Settings

We implemented all methods in MATLAB 2021b and performed all the experiments on a MacBook Pro with a 2.3GHz 8-core Intel Core i9. 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 a matrix over X ∈ R 300 × n X . The values n X of H5, H, JC, MV, O, and RJ are 1303 , 1481 , 910 , 1008 , 1148 , and 1141 , respectively. The weight vector r or c is taken as 1 n X /n X . We set k = 2 .

Initial points and retraction. We choose α -1 = 0 , β -1 = 0 for iRBBS and α 0 = 0 , β 0 = 0 for R(A)BCD. As for U 0 , we choose U 0 ∈ argmax U ∈U 〈 V π 0 , UU T 〉 with π 0 ∈ Π( r, c ) for all methods. Here, π 0 is formulated by firstly generating a matrix ˜ π 0 with each entry randomly drawn from the standard uniform distribution on (0 , 1) and then rounding ˜ π 0 / ‖ ˜ π 0 ‖ 1 to π 0 via [3, Algorithm 2]. The retraction operator in all the above methods is taken as the QR-retraction Retr qr U ( ξ ) = qf( X + ξ ) , where qf( X + ξ ) ∈ U satisfying X + ξ = qf( X + ξ )upp( X + ξ ) with upp( X + ξ ) being an upper triangular k -byk matrix with positive diagonal elements.

Parameters of iRBBS in Section 4.1. We choose τ min = 10 -10 /L , τ max = 10 10 /L , γ = 0 . 85 , τ 0 = 10 -3 , σ = 1 / 2 , δ 1 = 10 -4 , and ρ = 0 . 49 η k .

glyph[negationslash]

Parameters of R(A)BCD in Section 4.1. As stated in Remark 4.1 in [28], to achieve the best performance of R(A)BCD, one has to spend some efforts in 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 τ RBCD = 0.09 if the instance is X/RJ and τ RBCD = 0 . 1 otherwise; τ RABCD = 0.0015 if the instance is X/RJ with X = H, τ RABCD = 0.001 if the instance is H/RJ and τ RABCD = 0 . 0025 otherwise. Moreover, we stop R(A)BCD when e 1 η k ( α t , β t -1 , U t ) ≤ glyph[epsilon1] 1 and e 2 η k ( α t , β t -1 , U t ) ≤ glyph[epsilon1] 2 or the maximum iteration number reaches 5000.

Parameters of REALM in Section 4.2 . Once η k becomes η min , we set the corresponding glyph[epsilon1] k, 1 = glyph[epsilon1] 1 and glyph[epsilon1] k, 2 = glyph[epsilon1] 2 and stop the algorithm if e 1 η k ( x k , π k ) ≤ glyph[epsilon1] 1 and e 2 η k ( x k , π k ) ≤ glyph[epsilon1] 2 . We choose glyph[epsilon1] 1 , 2 = 10 -1 max {‖ r ‖ ∞ , ‖ c ‖ ∞ } and glyph[epsilon1] 1 , 1 = 2 ‖ C ‖ ∞ glyph[epsilon1] 1 , 2 .

100

200

300

400

500

d

10 1

10 2

10 3

10 4

nGrad

100

200

300

400

500

d

10 2

10 3

10 4

nSk

100

200

300

400

500

d

10 0

10 1

10 2

scaled time

(a)

n

= 100

200

400

600

800

1000

n

10 1

10 2

10 3

nGrad

200

400

600

800

1000

n

10 2

10 3

10 4

nSk

200

400

600

800

1000

n

10 0

10 1

10 2

scaled time

(b)

d

= 50

100

200

300

400

500

d

10 1

10 2

10 3

10 4

nGrad

100

200

300

400

500

d

10 2

10 3

10 4

nSk

100

200

300

400

500

d

10 0

10 1

10 2

scaled time

(c)

n

=

d

50

100

150

200

250

d

10 1

10 2

10 3

10 4

nGrad

50

100

150

200

250

d

10 2

10 3

10 4

nSk

50

100

150

200

250

d

10 0

10 1

10 2

scaled time

(d)

n

= 10

d

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 ) glyph[circledot] ∑ k ∗ k =1 e k , where sign( · ) is taken elementwise, k ∗ ∈ { 1 , . . . , d } and { e 1 , . . . , e k ∗ } is the canonical basis of R d . Let µ = U ([ -1 , 1] d ) be the uniform distribution over an hypercube and ν = P # ( µ ) be the pushforward of µ under the map P . We set k = k ∗ = 2 and take both the weight vectors r and c as 1 /n .

Datasets D.2 (Real datasets: digits from MNIST datasets [34]) . For each digit 0 , 1 , . . . , 9 , we extract the 128-dimensional features from a pre-trained convolutional neural network. Each digit corresponds to a matrix X ∈ R 128 × n X . The values n X of 0 , 1 , . . . , 9 are 980 , 1135 , 1032 , 1010 , 982 , 892 , 958 , 1028 , 974 , and 1009 , respectively. The weight vector r or c is taken as 1 n X /n X . We choose k = 2 .

D.1 Comparison with R(A)BCD on Solving Subproblem (11)

We follow the same settings of η in [28]. For Dataset D.1, we choose η = 0 . 2 when d < 250 and η = 0 . 5 otherwise. For Dataset D.2, we set η = 8 .

Parameters of R(A)BCD in this section. For Dataset D.1, we choose τ RBCD = τ RABCD = 0 . 001 as done in [28]. For Dataset D.2, we choose τ RBCD = 0 . 004 , and we do not test RABCD for this dataset since the well-chosen stepsize is not provided in [28].

For Dataset D.1, we randomly generate 10 instances for each ( n, d ) pair, each equipped with 5 randomly generated initial points. The comparison results for Dataset D.1 are plotted as Figure 2. Note that the values ̂ P 2 k returned by different methods are almost the same for this dataset. Therefore, we do not report the values in 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 q ( U t ) in a relatively high precision can help to reduce the whole iteration number of updating U . In particular, iRBBS-0, which computes (almost) the exact grad q ( U t ) , 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 a moderate θ generally achieves the best overall performance. The last rows of Figure 2 show that iRBBS-0.1 is almost the fastest among all iRBBS methods. (ii) Our iRBBSθ is better than R(A)BCD in terms that the former always have smaller nGrad and nSk and are always faster. Particularly, for Dataset D.1, iRBBS-0.1 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.2x faster than RABCD (11.1s) and about 28.6x faster than RBCD (44.4s). It

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) as for the solution quality, iRBBS-0 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-inf takes the most nGrad

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-0, 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 a 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 in R(A)BCD.

D.2 Comparison on Computing the PRW Distance (1)

In this subsection, we present more numerical results on Datasets D.1 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 θ t chosen as in Section 4.1.

The results for Dataset D.1 are presented in Table 4. For each ( n, d ) pair, we randomly generate 10 instances, each equipped with 5 randomly generated initial points. We conside REALM(0 . 055 , 0 . 9) and REALM(0 . 02 , 0) both with η 1 = 1 , γ glyph[epsilon1] = 0 . 25 , and γ η = 0 . 5 . Note that the latter admits a smaller η min and does not update the multiplier matrix. From Table 4, we can observe that REALM(0 . 055 , 0 . 9) can 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(0 . 055 , 0 . 9) updates the multiplier matrix 8 times in 14 total iterations.

The results over 20 runs on the real Dataset D.2 are reported in Table 5. We consider REALM(1 , 0) and REALM(3 , 0 . 9) both with η 1 = 200 and γ η = γ glyph[epsilon1] = 0 . 25 . From Table 5, we can see 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 11.7 total iterations.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nClearly, 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.\n\nWEAKNESSES:\n(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?\n\nQUESTIONS:\n1. 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$?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-04"}

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

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

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

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

LIMITATIONS:
The authors pointed out an important limitation in the theoretical analysis, which is the lack of a lower bound of \eta_k.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-06

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

STRENGTHS:
1.	The authors propose a Riemannian exponential augmented Lagrangian method (REALM) method to efficiently and faithfully compute the PRW distance and establish the global convergence of REALM in the sense that any limit point of the sequence generated by the algorithm is a stationary point of the original problem

2.	To efficiently solve the subproblem in REALM, the authors also propose a novel and practical algorithm, namely, the inexact Riemannian Barzilai-Borwein (BB) method with Sinkhorn iteration (iRBBS).

WEAKNESSES:
1.	The authors should provide a proof sketch for their main theoretical analysis in this paper.

2.	The authors should add more experimental results to verify their theoretical results and their algorithms.

QUESTIONS:
See the above setion

LIMITATIONS:
See Weaknesses
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThis work proposes a new method to compute the PRW distance with solid theoretical guarantees.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nHello 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\nLIMITATIONS:\nIn 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-21"}

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

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

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

2. the writing of this paper needs to be further improved.

QUESTIONS:
see the weakness

LIMITATIONS:
see the weakness.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-31