Title: A Specialized Semismooth Newton Method for Kernel-Based Optimal Transport
Venue: NeurIPS 2023
Accepted: no
Keywords: kernal-based optimal transport estimation

nonsmooth equation model

specialized semi-smooth Newton method
Abstract: Kernel-based optimal transport (OT) estimation is an alternative to the standard plug-in OT estimation. Recent works suggested that kernel-based OT estimators are more statistically efficient than plug-in OT estimators when comparing probability measures in high-dimensions~\citep{Vacher-2021-Dimension}. However, the computation of these estimators relies on the short-step interior-point method for which the required number of iterations is known to be \textit{large} in practice. In this paper, we propose a nonsmooth equation model for kernel-based OT estimation and show that it can be efficiently solved via a specialized semismooth Newton (SSN) method. Indeed, by exploring the special problem structure, the per-iteration cost of performing one SSN step can be significantly reduced in practice. We also prove that our algorithm can achieve a global convergence rate of $O(1/\sqrt{k})$ and a local quadratic convergence rate under some standard regularity conditions. Finally, we demonstrate the effectiveness of our algorithm by conducing the experiments on both synthetic and real datasets.

Abstract

Kernel-based optimal transport (OT) estimation is an alternative to the standard plug-in OT estimation. Recent works suggested that kernel-based OT estimators are more statistically efficient than plug-in OT estimators when comparing probability measures in high-dimensions [59]. However, the computation of these estimators relies on the short-step interior-point method for which the required number of iterations is known to be large in practice. In this paper, we propose a nonsmooth equation model for kernel-based OT estimation and show that it can be efficiently solved via a specialized semismooth Newton (SSN) method. Indeed, by exploring the special problem structure, the per-iteration cost of performing one SSN step can be significantly reduced in practice. We also prove that our algorithm can achieve a global convergence rate of O (1 / √ k ) and a local quadratic convergence rate under some standard regularity conditions. Finally, we demonstrate the effectiveness of our algorithm by conducing the experiments on both synthetic and real datasets.

1 Introduction

Optimal transport (OT) theory [60] has provided a principled framework for comparing probability distributions. It has been extensively adopted in machine learning and related fields, with examples including generative modeling [2, 21, 51, 57], classification and clustering [20, 55, 25], and domain adaptation [9, 10, 49], see also the monograph [43]. It has also had an impact in applied areas such as neuroimaging [27] and cell trajectory prediction [53, 66].

Curse of Dimensionality. In many real application problems, the OT cost is computed for squared Euclidean distance on the sampled distributions with n observations (leading to the 2-Wasserstein distance). It is known that OT estimation suffers from the curse of dimensionality [16, 19, 62]: the standard plug-in estimator, which consists in computing the OT distance between the sampled distributions with n observations, converges to the OT distance between true distributions at a rate of O ( n -1 /d ) , which degrades exponentially in the dimension d . This rate can be improved to O ( n -1 / 2 d ) when true distributions are different [7] but it is still problematic in a high-dimensional regime. This issue can be a barrier to its adoption in machine learning since various application problems arising from image processing and bioengineering are high-dimensional. Practitioners have long been aware of such limitations and proposed efficient computational schemes that not only improve computational complexity but also carry out statistical regularization.

Regularization. In this context, two threads have been investigated to regularize the OT distance: entropic regularization [11, 12, 22, 36] or low-dimensional projection [48, 4, 41, 29, 39, 31, 32, 40].

For the former approach, the sample complexity of entropic OT is bounded by regularization parameter

bounded by

O

(

n

-

1

/k

)

η >

0

O

(

η

-

d/

2

n

-

1

/

2

)

for a

. For the latter approach, the sample complexity of projection OT is for an integer-valued projection dimension

k

d

. Even though these bounds

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

attain the dimension-free dependence on n , they deteriorate when η is small or k is large, either of which is needed to study the sample complexity of OT [7], and which plays a role in real applications.

Leveraging Smoothness. A recent line of works have focused on the wavelet-based OT estimators under a strong smoothness condition [63, 26, 15, 34]. Although these estimators are minimax optimal from a statistical viewpoint, they are algorithmically intractable [59]. In contrast, a specific entropic regularized OT estimator is computationally tractable but still suffers from the curse of dimensionality when the dimension is sufficiently large [44]. Recently, Vacher et al. [59] has closed this statisticalcomputational gap by designing a kernel-based estimator relying on kernel sums-of-squares (SoS) and showed that it can be computed by a short-step interior-point method with polynomial-time complexity guarantee. However, the short-step interior-point method is well known to be ineffective for large number of iterations required as the sample size increases, diminishing their value from both statistical and practical viewpoints 1 . In this context, Muzellec et al. [38] proposed to use the relaxation model and solve it using gradient-based methods. However, the relaxation model may not be a good approximation for kernel-based OT estimator, thereby lacking any statistical guarantee.

Goal: While there is an ongoing debate in the OT literature on the merits of computing the plug-in OT estimators v.s. kernel-based OT estimators, we adopt the perspective that Vacher et al. [59] does introduce a fairly novel approach and we believe that it is worth studying if the kernel-based OT estimation can provide leads for practical use. The goal of this paper is therefore to facilitate the computational aspect by designing new algorithms, and to figure out whether that estimator's theoretical claims is also supported by practical relevance. The statistical analysis of kernel-based OT estimation itself, e.g., the proper choice of penalty parameters, is beyond the scope of this paper.

Contribution: In this paper, we propose a nonsmooth equation model for computing kernel-based OT estimators and show that it has a special problem structure, allowing it to be solved in an efficient manner using semismooth Newton method [37, 47, 46, 58].

We first propose a nonsmooth equation model for computing the kernel-based OT estimator and define an approximate OT value, which allows us to carry out a finite-time analysis of the algorithm. Then, we propose a specialized semismooth Newton method for computing the kernel-based OT estimator and prove a global convergence rate of O (1 / √ k ) (Theorem 3.3) and a local quadratic convergence rate under standard regularity conditions (Theorem 3.4). Notably, we significantly reduce the per-iteration computational cost by exploiting the special problem structure. Finally, we conduct the experiments to evaluate our algorithm on both synthetic and real datasets. Experimental results demonstrate its efficiency for solving the kernel-based OT estimation.

Organization. The remainder of the paper is organized as follows. In Section 2, we present the nonsmooth equation model for computing the kernel-based OT estimators and define the optimality notion based on the residual map. In Section 3, we propose and analyze the specialized semismooth Newton (SSN) algorithm for computing the kernel-based OT estimators and prove that our algorithm achieves the convergence rate guarantee in both global and local sense. In Section 4, we conduct the experiments on both synthetic and real datasets, demonstrating that our algorithm can effectively compute the kernel-based OT estimators and is more efficient than short-step interior-point methods. In Section 5, we conclude this paper. In the supplementary material, we provide further background materials on SSN methods, additional experimental results, and missing proofs for key results.

2 Preliminaries and Technical Background

In this section, we present the basic setup for the kernel-based optimal transport (OT) estimation and propose a nonsmooth equation model for its computation.

2.1 Kernel-based OT estimation

We formally define the OT distance and review the kernel-based OT estimation [59]. Indeed, the OT 81 distance with strong smooth distributions can be estimated at a dimension-free statistical rate with 82 high probability by solving a suitably defined optimization model. 83

Let X and Y be two bounded domains in R d and let P ( X ) and P ( Y ) be the set of Borel probability 84 measures in X and Y . Suppose that µ ∈ P ( X ) , ν ∈ P ( Y ) and Π( µ, ν ) is the set of couplings 85 between µ and ν , the OT distance [60] is given by 86

Its dual formulation is stated as follows, 87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

where C ( R d ) is the space of continuous functions on R d . Note that the supremum can be attained and the corresponding optimal dual functions u glyph[star] and v glyph[star] are referred to as the Kantorovich potentials [52]. This problem is delicate to solve since 1 2 ‖ x -y ‖ 2 ≥ u ( x )+ v ( y ) needs to be satisfied on a continuous set X × Y . A natural approach is to take n points { (˜ x 1 , ˜ y 1 ) , . . . , (˜ x n , ˜ y n ) } ⊆ X × Y and consider the constraints 1 2 ‖ ˜ x i -˜ y i ‖ 2 ≥ u (˜ x i ) + v (˜ y i ) for all 1 ≤ i ≤ n . However, it can not leverage the smoothness of potentials [3], yielding an error of Ω( n -1 /d ) . Vacher et al. [59] has overcome this difficulty by replacing the inequality constraints with equality constraints that are equivalent and considering the equality constraints over n points. Following their works, we impose the following assumption on the support sets X,Y and the densities of µ and ν .

Assumption 2.1 Let d ≥ 1 be the dimension and let m> 2 d +2 be the order of smoothness. Then, we assume that (i) the support sets X,Y are convex, bounded, and open with Lipschitz boundaries; (ii) the densities of µ, ν are finite, bounded away from zero and m -times differentiable.

Assumption 2.1 guarantees that the potentials u glyph[star] and v glyph[star] have a similar order of differentiability [14], leading to an effective way to represent u and v via a reproducing Kernel Hilbert space (RKHS) [42]. In particular, we define H s ( Z ) := { f ∈ L 2 ( Z ) | ‖ f ‖ H s ( Z ) := ∑ | α |≤ s ‖ D α f ‖ L 2 ( Z ) < + ∞} and remark that H s ( Z ) ⊆ C k ( Z ) for any s > d 2 + k , where k ≥ 0 is integer-valued. This implies that H m +1 ( X ) , H m +1 ( Y ) and H m ( X × Y ) are RKHS under Assumption 2.1 and they are associated with three bounded continuous feature maps φ X : X ↦→ H m +1 ( X ) , φ Y : Y ↦→ H m +1 ( Y ) and φ XY : X × Y ↦→ H m ( X × Y ) . For simplicity, we let H X = H m +1 ( X ) , H Y = H m +1 ( Y ) and H XY = H m ( X × Y ) . Vacher et al. [59, Corollary 7] shows that (i) u glyph[star] ∈ H X and v glyph[star] ∈ H Y with

where w µ = ∫ X φ X ( x ) dµ ( x ) and w ν = ∫ Y φ Y ( y ) dν ( y ) are kernel mean embeddings ; (ii) A glyph[star] ∈ S + ( H XY ) 2 exists and satisfies the equality constraint as follows:

Putting these pieces yields a representation theorem for estimating the OT distance. Indeed, under Assumption 2.1, the dual OT problem is equivalent to the RKHS-based problem given by

The above equation offers two advantages: (i) The equality constraint can be well approximated under Assumption 2.1; (ii) RKHSs allow the kernel trick: computing parameters are expressed in terms of kernel functions that correspond to

and

where the kernel functions are explicit and can be computed in O ( d ) given the samples. The final step is to approximate Eq. (2.1) using the data x 1 , . . . , x n sample ∼ µ and y 1 , . . . , y n sample ∼ ν , and the filling points { (˜ x 1 , ˜ y 1 ) , . . . , (˜ x n , ˜ y n ) } ⊆ X × Y . Indeed, we define ˆ µ = 1 n sample ∑ n sample i =1 δ x i and

ˆ ν = 1 n sample ∑ n sample i =1 δ y i , and use 〈 u, w ˆ µ 〉 H X + 〈 v, w ˆ ν 〉 H Y instead of 〈 u, w µ 〉 H X + 〈 v, w ν 〉 H Y where 119 w ˆ µ = 1 n sample ∑ n sample i =1 φ X ( x i ) and w ˆ ν = 1 n sample ∑ n sample i =1 φ Y ( y i ) . We also impose the penalization terms 120 for u , v , and A to alleviate the error induced by sampling the corresponding equality constraints. 121 Then, the resulting problem with regularization parameters λ 1 , λ 2 > 0 is summarized as follows: 122

123

124

125

126

127

128

129

130

131

132

133

134

Focusing on the case n sample = Θ( n ) , we let ˆ u glyph[star] and ˆ v glyph[star] be the unique maximizers of Eq. (2.2). Then, the estimator for OT ( µ, ν ) we consider corresponds to

Remark 2.2 It follows from Vacher et al. [59, Corollary 3] that the norm of empirical potentials can be controlled using λ 1 = ˜ Θ( n -1 / 2 ) and λ 2 = ˜ Θ( n -1 / 2 ) in high probability sense, leading to the sample complexity bound: | ̂ OT n -OT ( µ, ν ) | = ˜ O ( n -1 / 2 ) . In comparison with plug-in estimators, the kernel-based OT estimators are better when the sample size is small and the dimension is high.

Note that Eq. (2.2) is an infinite-dimensional optimization problem and is thus difficult to be solved. Thanks to Vacher et al. [59, Theorem 15], we have that the dual problem of Eq. (2.2) can be presented in a finite-dimensional space and the strong duality holds true. Indeed, we define Q ∈ R n × n with Q ij = k X (˜ x i , ˜ x j ) + k Y (˜ y i , ˜ y j ) , and z ∈ R n with z i = w ˆ µ (˜ x i ) + w ˆ ν (˜ y i ) -λ 2 ‖ ˜ x i -˜ y i ‖ 2 , and q 2 = ‖ w ˆ µ ‖ 2 H X + ‖ w ˆ ν ‖ H Y , where we have

and 135

We define K ∈ R n × n with K ij = k XY ((˜ x i , ˜ y i ) , (˜ x j , ˜ y j )) and R as an upper triangular matrix for 136 the Cholesky decomposition of K . We let Φ i be the i th column of R . Then, the dual problem of 137 Eq. (2.2) reads: 138

Suppose that ˆ γ is one minimizer, we have 139

To our knowledge, the existing method proposed for solving Eq. (2.4) is a short-step interior-point 140 method for which the required number of iterations is known to be large when n is large, which 141 is necessary to guarantee small statistical error. To avoid this issue, Muzellec et al. [38] proposed 142 solving an unconstrained relaxation model which allows for the application of gradient-based methods. 143 However, the estimators obtained from solving such relaxation model lack any statistical guarantee. 144

2.2 Nonsmooth equation model and optimality condition 145

For simplicity, we define the operator Φ : R n × n ↦→ R n and its adjoint Φ glyph[star] : R n ↦→ R n × n by 146

We present the optimality notion for Eq. (2.4) as follows: 147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

Definition 2.1 A point ˆ γ ∈ R n is an optimal solution of Eq. (2.4) if we have Φ glyph[star] (ˆ γ ) + λ 1 I glyph[followsequal] 0 and 1 4 λ 2 ˆ γ glyph[latticetop] Q ˆ γ -1 2 λ 2 ˆ γ glyph[latticetop] z + q 2 4 λ 2 ≤ 1 4 λ 2 γ glyph[latticetop] Qγ -1 2 λ 2 γ glyph[latticetop] z + q 2 4 λ 2 for all γ satisfying that Φ glyph[star] ( γ ) + λ 1 I glyph[followsequal] 0 .

Clearly, Eq. (2.4) can be reformulated as the following optimization problem given by

We denote w = ( γ, X ) as a vector-matrix pair and let R : R n × R n × n → R n × R n × n be given by

where S n + = { X ∈ R n × n : X glyph[followsequal] 0 } . Then, we can measure the optimality of w via appeal to the quantity ‖ R ( w ) ‖ and shows that the notion is the same as used in Definition 2.1.

Proposition 2.3 A point ˆ γ is an optimal solution of Eq. (2.4) if and only if ˆ w = (ˆ γ, ˆ X ) satisfies R ( ˆ w ) = 0 for some ˆ X glyph[followsequal] 0 .

Proposition 2.3 shows that we can compute the kernel-based OT estimators by solving the nonsmooth equation model R ( w ) = 0 . The optimality criterion based on the residual map R ( · ) allows for a global convergence rate analysis for our specialized semismooth Newton method.

3 Algorithm and Convergence Analysis

In this section, we derive our algorithm and provide a convergence rate analysis. The key idea here is to apply the regularized semismooth Newton (SSN) method for solving R ( w ) = 0 and improve the computation of each SSN step by exploring the special structure of generalized Jacobian. We also safeguard the regularized SSN method by min-max method to achieve a global rate.

Generalized Jacobian. We first examine the special structure of the generalized Jacobian of R ( w ) . Indeed, by using the definition of S n + , we have proj S n + ( Z ) = P α Σ α P glyph[latticetop] α where

with Σ = diag ( σ 1 , . . . , σ n ) and the sets of the indices of positive and nonpositive eigenvalues of Z (we denote these sets by α = { i | σ i > 0 } and ¯ α = { 1 , 2 , . . . , n } \ α ). Moreover, we notice that R is Lipschitz continuous. Then, Rademacher's theorem can guarantee that R is almost everywhere differentiable. We introduce the concepts of generalized Jacobian [8].

Definition 3.1 Suppose that R is Lipschitz continuous and D R is the set of differentiable points of R . The B-subdifferential of R at w is given by ∂ B R ( w ) := { lim k → + ∞ ∇ F ( w k ) | w k ∈ D R , w k → w } . The set ∂R ( w ) = conv ( ∂ B R ( w )) is called generalized Jacobian where conv denotes the convex hull.

We define a generalized operator M ( Z ) ∈ ∂ proj S n + ( Z ) using its application to an n × n matrix S :

where the ◦ symbol denotes a Hadamard product and Ω = ( E αα η α ¯ α η glyph[latticetop] α ¯ α 0 ) with E αα being a matrix of ones and η ij = σ i σ i -σ j for all ( i, j ) ∈ α × ¯ α . Note that all entries of Ω lie in the interval (0 , 1] . In general, it is nontrivial to characterize the generalized Jacobian ∂R ( w ) exactly but we can compute an element J ( w ) ∈ ∂R ( w ) using M ( · ) as defined before.

We next introduce the definition of the (strong) semismoothness of an operator.

Definition 3.2 Suppose that R is Lipschitz continuous. Then, R is (strongly) semismooth at w if (i) R is directionally differentiable at w ; and (ii) for any ∆ w and J ∈ ∂R ( w +∆ w ) , we have

Algorithm 1 Solving Eq. (3.2) where r k = ( r 1 k , r 2 k ) ∈ R n × R n × n )

1: Compute a 1 = -r 1 k -1 µ k +1 (Φ( r 2 k + T k [ r 2 k ])) and a 2 = -r 2 k .

2: Use the CG or symmetric QMS method to solve ( 1 2 λ 2 Q + µ k I +Φ T k Φ glyph[star] ) -1 ˜ a 1 = a 1 inexactly and compute ˜ a 2 = 1 µ k +1 ( a 2 + T k [ a 2 ]) , where T k [ · ] is computed using the trick [68].

3: Compute the direction ∆ w k = (∆ w 1 k , ∆ w 2 k ) by ∆ w 1 k = ˜ a 1 and ∆ w 2 k = ˜ a 2 -T k [Φ glyph[star] (˜ a 1 )] .

The following proposition characterizes the residual map given in Eq. (2.6) and its generalized 181 Jacobian matrix. It also guarantees that the SSN method is suitable to solve R ( w ) = 0 . 182

183

Proposition 3.1 The residual map R given in Eq. (2.6) is strongly semismooth.

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

Regularized SSN step. We then discuss how to compute the Newton direction efficiently. In particular, at a given iterate w k , we compute a Newton direction ∆ w k by solving the equation

where J k ∈ ∂R ( w k ) , r k = R ( w k ) and I is an identity operator. The regularization parameter is chosen as µ k = θ k ‖ r k ‖ for stabilizing the semismooth Newton method in practice. From a computational point of view, it is not practical to solve the linear system in Eq. (3.2) exactly. Thus, we seek an approximation step ∆ w k by solving Eq. (3.2) approximately such that

where 0 < τ, κ < 1 are some positive constants and ‖·‖ is defined for a vector-matrix pair w = ( γ, X ) (i.e., ‖ w ‖ = ‖ γ ‖ 2 + ‖ X ‖ F where ‖ · ‖ 2 is Euclidean norm and ‖ · ‖ F is Frobenius norm).

Since J k in Eq. (3.2) is nonsymmetric and its dimension is large, we consider applying the Schur complement trick to transform Eq. (3.2) into a smaller symmetric system. If we vectorize the vector-matrix pair ∆ w 3 , the operators M ( Z ) and Φ can be expressed as matrices:

where ˜ P = P ⊗ P and Γ = diag ( vec (Ω)) .

We next provide a key lemma on the matrix form of J k + µ k I at a given iterate w k = ( γ k , X k ) .

Lemma 3.2 Given an iterate w k = ( γ k , X k ) , we compute Z k = X k -(Φ glyph[star] ( γ k ) + λ 1 I ) and use Eq. (3.1) to obtain P k , Σ k , α k and ¯ α k . We then obtain Ω k , ˜ P k = P k ⊗ P k and Γ k = diag ( vec (Ω k )) . Then, the matrix form of J k + µ k I is given by

where and

with T k = ˜ P k L k ˜ P glyph[latticetop] k where L k is a diagonal matrix with ( L k ) ii = (Γ k ) ii µ k +1 -(Γ k ) ii and (Γ k ) ii ∈ (0 , 1] is then denoted as the i th diagonal entry of Γ k .

As a consequence of Lemma 3.2, the solution of Eq. (3.2) can be obtained by solving one certain symmetric linear system with the matrix 1 2 λ 2 Q + µ k I + AT k A glyph[latticetop] . We remark that this system is well-defined since both Q and AT k A glyph[latticetop] are positive semidefinite and the coefficient µ k is chosen such that 1 2 λ 2 Q + µ k I + AT k A glyph[latticetop] is invertible. This also shows that Eq. (3.2) is well-defined.

Algorithm 2 A specialized SSN method with safeguarding

1: Input: τ, κ , α 2 ≥ α 1 > 0 , β 0 < 1 , β 1 , β 2 > 1 and θ, θ > 0 .

2: Initialization: v 0 = w 0 ∈ R n ×S n + and θ 0 > 0 . Set k = 0 .

3: for k = 0 , 1 , 2 , . . . do

4: Update v k +1 from v k using one-step EG.

5: Select J k ∈ ∂R ( w k ) .

6: Solve the linear system in Eq. (3.2) approximately such that ∆ w k satisfies Eq. (3.3).

7: Compute ˜ w k +1 = w k +∆ w k .

8: Update θ k +1 using Eq. (3.4) accordingly.

9: Set w k +1 = ˜ w k +1 if ‖ R ( ˜ w k +1 ) ‖ ≤ ‖ R ( v k +1 ) ‖ is satisfied. Otherwise, set w k +1 = v k +1 .

We define T k and Q as the operator form of T k = ˜ P k L k ˜ P glyph[latticetop] k and Q and write r k = ( r 1 k , r 2 k ) explicitly 208 where r 1 k ∈ R n and r 2 k ∈ R n × n . Then, we have 209

The next step consists in solving a new symmetric linear system and is given by 210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

227

which leads to

Compared to Eq. (3.2) whose matrix form has size ( n 2 + n ) × ( n 2 + n ) , we remark that the one in the step above is smaller with the size of n × n and can be efficiently solved by conjugate gradient (CG) method or symmetric quasi-minimal residual (QMR) method [28, 50]. The final step is to compute the Newton direction ∆ w k = (∆ w 1 k , ∆ w 2 k ) as follows,

It remains to provide an efficient manner to compute T k [ · ] . Since T k is defined as the operator form of T = ˜ P k L k ˜ P glyph[latticetop] k , we have

where Ψ k is determined by µ k and Ω k . Indeed, we have

where ξ ij = η ij µ k +1 -η ij for all ( i, j ) ∈ α k × ¯ α k . Following Zhao et al. [68], we use the decomposition T k [ S ] = G + G glyph[latticetop] where U = P k (: , α k ) glyph[latticetop] S and

The number of flops required to compute T k [ S ] is 8 | α k | n 2 . For the case of | α k | > ¯ α k , we compute T k [ S ] via T k [ S ] = 1 µ k S -P k (( 1 µ k E -Ψ k ) ◦ ( P glyph[latticetop] k SP k )) P glyph[latticetop] k using 8 | ¯ α k | n 2 flops. This demonstrates that we can obtain an approximate solution of Eq. (3.2) efficiently whenever | α k | or | ¯ α k | is small. We present the scheme for computing an approximate Newton direction in Algorithm 1.

Adaptive strategy. We propose a rule for updating θ k where µ k = θ k ‖ r k ‖ is defined in Eq. (3.2). Indeed, we compute ρ k = -〈 R ( w k ) , ∆ w k 〉 and use it to update θ k +1 . The update rule is summarized as follows:

where β 0 < 1 , β 1 , β 2 > 1 and θ, θ > 0 . 228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

248

249

250

251

252

253

254

255

256

257

258

259

Main scheme. We summarize the complete scheme of our new algorithm in Algorithm 2. Indeed, we generate a sequence of iterates by alternating between extragradient (EG) method [17, 6] and the aforementioned regularized SSN method.

Note that we maintain one auxiliary sequence of iterates { v k } k ≥ 0 . This sequence is directly generated by the EG method for solving the min-max optimization problem in Eq. (2.5) and is used to safeguard the regularized SSN method to achieve a global convergence rate. More specifically, we start with v 0 = w 0 ∈ R n ×S n + and perform the k th iteration as follows,

Update v k +1 from v k using one-step EG.

Update ˜ w k +1 from w k using one-step regularized SSN.

Set w k +1 = ˜ w k +1 if ‖ R ( ˜ w k +1 ) ‖ ≤ ‖ R ( v k +1 ) ‖ and w k +1 = v k +1 otherwise.

In our experiment, we find that the main iterates are mostly generated by regularized SSN steps and the whole algorithm converges at a superlinear rate. This phenomenon is quite intuitive: if the initial point is sufficiently close to one nondegenerate optimal solution, the regularized SSN method can achieve the similar quadratic convergence rate (cf. Theorem 3.4) as shared by other SSN methods in the existing literature [35, 18, 1]. The detailed analysis will be provided in the appendix.

Main results. We establish the convergence guarantee of Algorithm 2 in the following theorems.

Theorem 3.3 Suppose that { w k } k ≥ 0 is a sequence of iterates generated by Algorithm 2. Then, the residuals of { w k } k ≥ 0 converge to 0 at a rate of 1 / √ k , i.e., ‖ R ( w k ) ‖ = O (1 / √ k ) .

Theorem 3.4 Suppose that { w k } k ≥ 0 is a sequence of iterates generated by Algorithm 2. Then, the residuals of { w k } k ≥ 0 converge to 0 at a quadratic rate if the initial point w 0 is sufficiently close to w glyph[star] with R ( w glyph[star] ) = 0 and every element of ∂R ( w glyph[star] ) is invertible.

Remark 3.5 In the context of constrained convex-concave min-max optimization problem, Cai et al. [6] proved the O (1 / √ k ) last-iterate convergence rate of the EG, matching the lower bounds [24, 23]. Since the kernel-based OT estimation can be solved as a min-max problem, the global convergence rate in Theorem 3.3 demonstrates the efficiency of Algorithm 2. It remains unclear whether or not we can improve the convergence result by exploring special structure of Eq. (2.5) .

4 Experiments

We present the results of experiments that evaluate the kernel-based OT estimation with our algorithm. The baseline approach is the short-step interior-point method [59]; we exclude the gradient-based method [38] from our experiment since it only solves the relaxation model. All the experiments were conducted on a MacBook Pro with an Intel Core i9 2.4GHz and 16GB memory.

Following the setup in Vacher et al. [59], we draw n sample samples from µ and n sample samples from ν , 260 where µ is a mixture of 3 d -dimensional Gaussian distributions and ν is a mixture of 5 d -dimensional 261 Gaussian distributions. Then, we sample n filling samples from a 2 d Sobol sequence. We also set the 262 bandwidth σ 2 = 0 . 01 and parameters λ 1 = 1 n and λ 2 = 1 √ n sample . Focusing on the case of d = 1 (i.e., 263

0

100

200

300

400

500

Number of filling samples

0

5

10

15

20

Time (seconds)

The dimension is 2

Short-step IPM

Specialized SSN

0

100

200

300

400

500

Number of filling samples

0

5

10

15

20

Time (seconds)

The dimension is 5

Short-step IPM

Specialized SSN

0

100

200

300

400

500

Number of filling samples

0

5

10

15

20

25

Time (seconds)

The dimension is 10

Short-step IPM

Specialized SSN

By varying the dimension d ∈ { 2 , 5 , 10 } , we also report the computation efficiency results in Figure 3. 266 It indicates that the our new algorithm is more efficient than the IPM as the number of filling points 267 increases, with smaller variance in computation time (seconds). 268

269

270

271

272

273

274

The experiments comparing kernel-based OT estimators with plug-in OT estimators on synthetic datasets have been conducted before [59, 38] and the results demonstrate that the kernel-based OT estimators behave better when the number of samples is small. Here, we repeat such experiment but using the real-world 4i datasets from Bunne et al. [5], which contains single-cell perturbed responses, and which include the unperturbed cells and cells subject to drug perturbations. Our experiments are conducted on 15 datasets with different drug perturbations.

275

276

277

278

279

280

281

282

283

284

Due to space limit, we defer the results to Appendix G (see Figure 4). We can see that the kernel-based OT estimators computed by our algorithm achieve satisfactory performance and behave better in most cases when the number of training samples is small; in particular, they better on 6 datasets, comparable on 5 datasets and worse on 4 datasets. Note that OTT computes the entropic regularized plug-in OT estimators and is heavily optimized to effectively handle noisy data. Therefore, it would be no surprise that OTT outperforms our algorithm when the number of training samples is sufficient. However, the kernel-based OT estimation still provides a fairly effective alternative when the number of training samples is small, which is consistent with the previous observations on synthetic data [59, 38]. Our results also validate the effectiveness of our algorithm for computing kernel-based OT estimators.

5 Concluding Remarks

In this paper, we propose a nonsmooth equation model for computing kernel-based OT estimators 285 and show that it has a special problem structure, allowing it to be solved in an efficient manner using 286 semismooth Newton method. In particular, we propose a specialized semismooth Newton method that 287 achieves low per-iteration computational cost by exploiting the special problem structure, and prove 288 a global sublinear convergence rate and a local quadratic convergence rate under standard regularity 289 conditions. Preliminary experimental results on synthetic datasets show that our algorithm is more 290 efficient than the short-step interior-point method [59], and the results on real data demonstrate the 291 effectiveness of our algorithm. Future work includes the applications of kernel-based OT estimators 292 to deep generative models and other real-world problems. 293

References 294

A. Ali, E. Wong, and J. Z. Kolter. A semismooth Newton method for fast, generic convex 295 programming. In ICML , pages 70-79. PMLR, 2017. (Cited on page 8.) 296

M. Arjovsky, S. Chintala, and L. Bottou. Wasserstein generative adversarial networks. In ICML , 297 pages 214-223, 2017. (Cited on page 1.) 298

P-C. Aubin-Frankowski and Z. Szabó. Hard shape-constrained kernel machines. In NeurIPS , 299 pages 384-395, 2020. (Cited on page 3.) 300

N. Bonneel, J. Rabin, G. Peyré, and H. Pfister. Sliced and radon Wasserstein barycenters of 301 measures. Journal of Mathematical Imaging and Vision , 51(1):22-45, 2015. (Cited on page 1.) 302

C. Bunne, S. G. Stark, G. Gut, J. S. del Castillo, K-V. Lehmann, L. Pelkmans, A. Krause, and 303 G. Ratsch. Learning single-cell perturbation responses using neural optimal transport. BioRxiv , 304 2021. (Cited on pages 9 and 17.) 305

Y. Cai, A. Oikonomou, and W. Zheng. Finite-time last-iterate convergence for learning in 306 multi-player games. In NeurIPS , pages 33904-33919, 2022. (Cited on pages 8 and 16.) 307

L. Chizat, P. Roussillon, F. Léger, F-X. Vialard, and G. Peyré. Faster Wasserstein distance 308 estimation with the Sinkhorn divergence. In NeurIPS , pages 2257-2269, 2020. (Cited on pages 1 309 and 2.) 310

F. H. Clarke. Optimization and Nonsmooth Analysis . SIAM, 1990. (Cited on page 5.) 311

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

N. Courty, R. Flamary, A. Habrard, and A. Rakotomamonjy. Joint distribution optimal trans315 portation for domain adaptation. In NIPS , pages 3733-3742, 2017. (Cited on page 1.) 316

M. Cuturi. Sinkhorn distances: lightspeed computation of optimal transport. In NIPS , pages 317 2292-2300, 2013. (Cited on page 1.) 318

M. Cuturi and A. Doucet. Fast computation of Wasserstein barycenters. In ICML , pages 319 685-693. PMLR, 2014. (Cited on page 1.) 320

M. Cuturi, L. Meng-Papaxanthos, Y. Tian, C. Bunne, G. Davis, and O. Teboul. Optimal transport 321 tools (OTT): A jax toolbox for all things wasserstein. ArXiv Preprint: 2201.12324 , 2022. (Cited 322 on page 17.) 323

G. De Philippis and A. Figalli. The Monge-Ampère equation and its link to optimal trans324 portation. Bulletin of the American Mathematical Society , 51(4):527-580, 2014. (Cited on 325 page 3.) 326

N. Deb, P. Ghosal, and B. Sen. Rates of estimation of optimal transport maps using plug-in 327 estimators via barycentric projections. In NeurIPS , pages 29736-29753, 2021. (Cited on page 2.) 328

R. M. Dudley. The speed of mean glivenko-cantelli convergence. The Annals of Mathematical 329 Statistics , 40(1):40-50, 1969. (Cited on page 1.) 330

331

332

F. Facchinei and J-S. Pang. Finite-Dimensional Variational Inequalities and Complementarity Problems . Springer Science & Business Media, 2007. (Cited on page 8.)

333

334

335

F. Facchinei, A. Fischer, and C. Kanzow. Inexact Newton methods for semismooth equations with applications to variational inequality problems. In Nonlinear Optimization and Applications , pages 125-139. Springer, 1996. (Cited on page 8.)

N. Fournier and A. Guillin. On the rate of convergence in wasserstein distance of the empirical 336 measure. Probability Theory and Related Fields , 162(3):707-738, 2015. (Cited on page 1.) 337

C. Frogner, C. Zhang, H. Mobahi, M. Araya-Polo, and T. Poggio. Learning with a Wasserstein 338 loss. In NIPS , pages 2053-2061, 2015. (Cited on page 1.) 339

A. Genevay, G. Peyré, and M. Cuturi. Learning generative models with Sinkhorn divergences. 340 In AISTATS , pages 1608-1617, 2018. (Cited on page 1.) 341

A. Genevay, L. Chizat, F. Bach, M. Cuturi, and G. Peyré. Sample complexity of Sinkhorn 342 divergences. In AISTATS , pages 1574-1583. PMLR, 2019. (Cited on page 1.) 343

N. Golowich, S. Pattathil, and C. Daskalakis. Tight last-iterate convergence rates for no-regret 344 learning in multi-player games. In NeurIPS , pages 20766-20778, 2020. (Cited on page 8.) 345

N. Golowich, S. Pattathil, C. Daskalakis, and A. Ozdaglar. Last iterate is slower than averaged 346 iterate in smooth convex-concave saddle point problems. In COLT , pages 1758-1784. PMLR, 347 2020. (Cited on page 8.) 348

N. Ho, X. Nguyen, M. Yurochkin, H. H. Bui, V. Huynh, and D. Phung. Multilevel clustering 349 via Wasserstein means. In ICML , pages 1501-1509. PMLR, 2017. (Cited on page 1.) 350

J-C. Hütter and P. Rigollet. Minimax estimation of smooth optimal transport maps. The Annals 351 of Statistics , 49(2):1166-1194, 2021. (Cited on page 2.) 352

353

354

H. Janati, T. Bazeille, B. Thirion, M. Cuturi, and A. Gramfort. Multi-subject MEG/EEG source imaging with sparse multi-task regression. NeuroImage , 220:116847, 2020. (Cited on page 1.)

355

356

357

358

C. T. Kelley. Iterative Methods for Linear and Nonlinear Equations . SIAM, 1995. (Cited on page 7.)

S. Kolouri, K. Nadjahi, U. ¸ Sim¸ sekli, R. Badeau, and G. K. Rohde. Generalized sliced Wasserstein distances. In NIPS , pages 261-272, 2019. (Cited on page 1.)

X. Li, D. Sun, and K-C. Toh. A highly efficient semismooth Newton augmented Lagrangian 359 method for solving Lasso problems. SIAM Journal on Optimization , 28(1):433-458, 2018. 360 (Cited on page 14.) 361

T. Lin, C. Fan, N. Ho, M. Cuturi, and M. I. Jordan. Projection robust Wasserstein distance and 362 Riemannian optimization. In NeurIPS , pages 9383-9397, 2020. (Cited on page 1.) 363

T. Lin, Z. Zheng, E. Chen, M. Cuturi, and M. I. Jordan. On projection robust optimal transport: 364 Sample complexity and model misspecification. In AISTATS , pages 262-270. PMLR, 2021. 365 (Cited on page 1.) 366

Y. Liu, Z. Wen, and W. Yin. A multiscale semismooth Newton method for optimal transport. 367 Journal of Scientific Computing , 91(2):1-29, 2022. (Cited on page 14.) 368

T. Manole, S. Balakrishnan, J. Niles-Weed, and L. Wasserman. Plugin estimation of smooth 369 optimal transport maps. ArXiv Preprint: 2107.12364 , 2021. (Cited on page 2.) 370

J. Martínez and L. Qi. Inexact Newton methods for solving nonsmooth equations. Journal of 371 Computational and Applied Mathematics , 60(1-2):127-145, 1995. (Cited on page 8.) 372

G. Mena and J. Niles-Weed. Statistical bounds for entropic optimal transport: Sample complex373 ity and the central limit theorem. In NIPS , pages 4541-4551, 2019. (Cited on page 1.) 374

R. Mifflin. Semismooth and semiconvex functions in constrained optimization. SIAM Journal 375 on Control and Optimization , 15(6):959-972, 1977. (Cited on pages 2 and 14.) 376

B. Muzellec, A. Vacher, F. Bach, F-X. Vialard, and A. Rudi. Near-optimal estimation of smooth 377 transport maps with kernel sums-of-squares. ArXiv Preprint: 2112.01907 , 2021. (Cited on 378 pages 2, 4, 8, and 9.) 379

K. Nadjahi, A. Durmus, L. Chizat, S. Kolouri, S. Shahrampour, and U. ¸ Sim¸ sekli. Statistical and 380 topological properties of sliced probability divergences. In NeurIPS , pages 20802-20812, 2020. 381 (Cited on page 1.) 382

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

F-P. Paty and M. Cuturi. Subspace robust Wasserstein distances. In ICML , pages 5072-5081. 385 PMLR, 2019. (Cited on page 1.) 386

V. I. Paulsen and M. Raghupathi. An Introduction to The Theory of Reproducing Kernel Hilbert 387 Spaces , volume 152. Cambridge University Press, 2016. (Cited on page 3.) 388

G. Peyré and M. Cuturi. Computational optimal transport: With applications to data science. 389 Foundations and Trends® in Machine Learning , 11(5-6):355-607, 2019. (Cited on page 1.) 390

A-A. Pooladian and J. Niles-Weed. Entropic estimation of optimal transport maps. ArXiv 391 Preprint: 2109.12004 , 2021. (Cited on page 2.) 392

H. Qi and D. Sun. An augmented Lagrangian dual approach for the H-weighted nearest 393 correlation matrix problem. IMA Journal of Numerical Analysis , 31(2):491-511, 2011. (Cited 394 on page 14.) 395

L. Qi and D. Sun. A survey of some nonsmooth equations and smoothing Newton methods. In 396 Progress in Optimization , pages 121-146. Springer, 1999. (Cited on page 2.) 397

L. Qi and J. Sun. A nonsmooth version of Newton's method. Mathematical Programming , 58 398 (1):353-367, 1993. (Cited on pages 2 and 14.) 399

J. Rabin, G. Peyré, J. Delon, and M. Bernot. Wasserstein barycenter and its application to texture 400 mixing. In International Conference on Scale Space and Variational Methods in Computer 401 Vision , pages 435-446. Springer, 2011. (Cited on page 1.) 402

I. Redko, N. Courty, R. Flamary, and D. Tuia. Optimal transport for multi-source domain 403 adaptation under target shift. In AISTATS , pages 849-858. PMLR, 2019. (Cited on page 1.) 404

Y. Saad. Iterative Methods for Sparse Linear Systems . SIAM, 2003. (Cited on page 7.) 405

T. Salimans, H. Zhang, A. Radford, and D. Metaxas. Improving GANs using optimal transport. 406 In ICLR , 2018. URL https://openreview.net/forum?id=rkQkBnJAb . (Cited on page 1.) 407

F. Santambrogio. Optimal Transport for Applied Mathematicians: Calculus of Variations, PDEs, 408 and Modeling , volume 87. Birkhäuser, 2015. (Cited on page 3.) 409

G. Schiebinger, J. Shu, M. Tabaka, B. Cleary, V. Subramanian, A. Solomon, J. Gould, S. Liu, 410 S. Lin, and P. Berube. Optimal-transport analysis of single-cell gene expression identifies 411 developmental trajectories in reprogramming. Cell , 176(4):928-943, 2019. (Cited on page 1.) 412

M. V. Solodov and B. F. Svaiter. A globally convergent inexact Newton method for systems 413 of monotone equations. Reformulation: Nonsmooth, Piecewise Smooth, Semismooth and 414 Smoothing Methods , pages 355-369, 1999. (Cited on page 14.) 415

S. Srivastava, V. Cevher, Q. Dinh, and D. Dunson. WASP: Scalable Bayes via barycenters of 416 subset posteriors. In AISTATS , pages 912-920. PMLR, 2015. (Cited on page 1.) 417

D. Sun and J. Sun. Semismooth matrix-valued functions. Mathematics of Operations Research , 418 27(1):150-169, 2002. (Cited on page 15.) 419

I. Tolstikhin, O. Bousquet, S. Gelly, and B. Schoelkopf. Wasserstein auto-encoders. In ICLR , 420 2018. (Cited on page 1.) 421

M. Ulbrich. Semismooth Newton Methods for Variational Inequalities and Constrained Opti422 mization Problems in Function Spaces . SIAM, 2011. (Cited on pages 2 and 14.) 423

A. Vacher, B. Muzellec, A. Rudi, F. Bach, and F-X. Vialard. A dimension-free computational 424 upper-bound for smooth optimal transport estimation. In COLT , pages 4143-4173. PMLR, 425 2021. (Cited on pages 1, 2, 3, 4, 8, and 9.) 426

C. Villani. Optimal Transport: Old and New , volume 338. Springer, 2009. (Cited on pages 1 427 and 3.) 428

C. Wang, D. Sun, and K-C. Toh. Solving log-determinant optimization problems by a Newton429 CG primal proximal point algorithm. SIAM Journal on Optimization , 20(6):2994-3013, 2010. 430 (Cited on page 14.) 431

J. Weed and F. Bach. Sharp asymptotic and finite-sample rates of convergence of empirical 432 measures in Wasserstein distance. Bernoulli , 25(4A):2620-2648, 2019. (Cited on page 1.) 433

J. Weed and Q. Berthet. Estimation of smooth densities in Wasserstein distance. In COLT , 434 pages 3118-3119. PMLR, 2019. (Cited on page 2.) 435

X. Xiao, Y. Li, Z. Wen, and L. Zhang. A regularized semismooth Newton method with projection 436 steps for composite convex programs. Journal of Scientific Computing , 76(1):364-389, 2018. 437 (Cited on page 14.) 438

J. Yang, D. Sun, and K-C. Toh. A proximal point algorithm for log-determinant optimization 439 with group Lasso regularization. SIAM Journal on Optimization , 23(2):857-893, 2013. (Cited 440 on page 14.) 441

K. D. Yang, K. Damodaran, S. Venkatachalapathy, A. C. Soylemezoglu, G. V. Shivashankar, 442 and C. Uhler. Predicting cell lineages using autoencoders and optimal transport. PLoS Compu443 tational Biology , 16(4):e1007828, 2020. (Cited on page 1.) 444

L. Yang, D. Sun, and K-C. Toh. SDPNAL++: a majorized semismooth Newton-CG augmented 445 Lagrangian method for semidefinite programming with nonnegative constraints. Mathematical 446 Programming Computation , 7(3):331-366, 2015. (Cited on page 14.) 447

X-Y. Zhao, D. Sun, and K-C. Toh. A Newton-CG augmented Lagrangian method for semidefinite 448 programming. SIAM Journal on Optimization , 20(4):1737-1765, 2010. (Cited on pages 6, 7, 449 and 14.) 450

G. Zhou and K-C. Toh. Superlinear convergence of a Newton-type algorithm for monotone 451 equations. Journal of Optimization Theory and Applications , 125(1):205-221, 2005. (Cited on 452 page 14.) 453

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose an implementation of Vacher et. al. (2021) based on a Semi-Smooth Newton (SSN) scheme. They reformulate their optimization problem as a root finding problem (Proposition 3.1) to which they apply the SSN scheme. They provide convergence guarantees (Theorem 3.3) that gives a $O(1/\\sqrt{T})$ convergence rate where $T$ is the number of iterations and provide an efficient way to reduce the cost per iterations (l.184 - l.224). Then they provide numerical experiments to validate that their method is faster than the one proposed in Vacher et. al. (2021).\n\nSTRENGTHS:\nTrying to get a scalable version of kernel based OT is a very legit topic as current implementations are slow and impossible to run on real data sets. Indeed, recall that using an Interior-Point-Method, kernel based OT was solved with a precision $\\epsilon$ in $O(n^{3.5} \\log(n/\\epsilon))$ time where $O(n^3)$ comes from the cost per iteration and $O(\\sqrt{n}\\log(n/\\epsilon))$ is the number of iterations. In this paper, the main contribution is to get rid of the dependency of $n$ in the number of iterations which is indeed a desirable feature. From my understanding, the authors can solve kernel based OT with precision $\\epsilon$ in $O(1/\\epsilon^2)$ iterations.\n\nWEAKNESSES:\nI believe that the authors oversell the work. As can be deduced from my comment above, the proposed method requires $O(1/\\epsilon^2)$ iterations for a precision $\\epsilon$ while previous work requires $O(\\sqrt{n}\\log(n/\\epsilon))$ iterations. When a high precision a sought after $\\epsilon \\to 0$, the proposed algorithm is indeed less efficient. The authors should have explicitly mentioned that.\nFurthermore, nothing precise is said on the cost per-iteration which is a crucial component of the practical efficiency. We can vaguely guess that it is $O(n^3)$ but it is stated nowhere.\n\nThe overall writing is confusing, the whole part on the computational efficiency should be clearly stated in a theorem or a proposition.\n\nQUESTIONS:\nI am actually skeptical on the $O(1/\\sqrt{T})$ convergence rate. What is the dependency in the regularizers $\\lambda_1, \\lambda_2$ and more generally in the condition number? In the case of Vacher et. al. (2021), there is little dependency in the condition number as they use an IPM-like method. Do SSN methods also weakly depend on the conditioning of the problem? Note that this is a crucial aspect as these regularizers implicitly depend on the number of samples $n$.\nOn the experimental side, it is claimed that the proposed method is faster. Yet which stopping criterion was used? Was it the same for both algorithms?\n\nLIMITATIONS:\nThe authors do not compare with enough precision their algorithm with the existing one, both in theory and on practice.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors focus on the problem of approximating OT numerically.\nThey focus on one approximated version of OT which leverages a Sum of Squares approximation to stratify both statistical guarantees and computational amenability.\nWhile the first proposal to solve this SoS approximation relied on interior point methods, the authors focus on a semi-smooth Newton method.\nIt consists in considering KKT optimality as some equation $R(w)=0$ and solve this equation using Newton updates.\nThey derive the algorithm in this specific OT setting, and prove convergence guarantees and rates of their methods.\nThey show experiments on synthetic data to see the approximation impact, and compare with interior point methods.\n\nSTRENGTHS:\nThis recent OT formulation satisfying statistical guarantees and being computationally amenable is an interesting quantity to estimate.\nThe proposal of the authors to propose another algorithm to estimate it and scale it to larger measures would increase the interest of this formulation to practitioners.\n\nWEAKNESSES:\n*The introduction is not precise enough*\n- Line 25, the rate $O(n^{-1/2d})$ is actually worse than the original rate. I think the authors meant a rate $O(n^{-2/d})$.\n- This is a secondary remark, but Line (31,32) another approach which attempts to regularize OT and ease computation is to consider mini batches of input data. I mention the work [FZFGC] and references therein if the authors wish to complement their introduction review.\n- The citation [44] in your paper is irrelevant. It focuses on estimating the OT Monge map when it exists, which is not the problem of estimating the cost, which you consider. Also, saying ‘a specific […] estimator’ is a super vague formulation which should be made precise.\n- Do the authors have references or precise rates to defend the assertion line 45-47 that « interior-point method is well known to be ineffective […] as the sample size increases » ? Similarly, do the authors have references that semi-smooth Newton method have better convergencce/scaling guarantees ? \n- I do not understand the sentence « While there is an ongoing debate in the OT literature on the merits of computing the plug-in OT estimators v.s. kernel-based OT estimators […] ». Which debates it is ? On which aspect does it especially focus ? This sentence is too vague to be insightful.\n- I do not understand the sentence Line 129 « kernel-based OT estimators are better when the sample size is small and the dimension is high ». Does that mean that the fewer samples we have, the better the approximation ?\n\n*The semi-smooth Newton method is not clear to understand* \n- Line 76, I think the authors should have introduced background knowledge on Semi-Smooth Newton methods instead of postponing them in the appendix. Furthermore, what is described by the authors is a review of previous contributions on this method, but no mathematical formulas are detailed. I would have put this part in the main body for related work, especially [33] which is exactly the same method as you, but for unregularized OT, and which you do not mention as related work. Lastly, to provide a self-contained and pedagogical description, I would have ideally wanted a brief description of SSN with a general framework, so that your work is an instantiation of this formulation.\n- I think Definition 2.1 is not extremely useful as it is the definition of optimality in a minimization program, and you can remove it.\n- Something that is not clear for me is whether some matrices are symmetric or not. First the set $S^n_+$ usually represent symmetric, positive matrice, but I see no symmetry in Line 152. The projection over $S^n_+$ of Equation (3.1) is true if Z is symmetric (or X in you context), but I see nowhere that X is assumed to be symmetric (or proved to be symmetric through iterations).  Line 192, you mention a Schur Complement trick to make the Jacobian symmetric, but when the matrix is asymmetric, there is no reason that the Schur complement is symmetric. All in all, the derivation of the method seems unclear and ill-posed. Could you please clarify on this ?\n- Could you please define a quadratic rate of convergence using an equation ?\n\n*Some experimental improvements to suggest*\nYou reproduce the experiments from [59], which is good to establish a comparison. However I think it could be made much clearer with some modifications.\n- In Figure 2, I would be interested in seeing the point wise difference between $c - \\hat{u} - \\hat{v}$ and $c - u_*-v_*$. It would emphasize where the approximation is best performed using this estimator. Reproducing the same experiment using interior point method would be insightful.\n- I don’t understand how time is estimated in Figure 3. Do you report the time to do a given number of iterations ? Is it the time to reach a given level of accuracy ? Without this I cannot make sure the comparison is fair.\n- I think that reproducing Figure 6 from [59] would be insightful. My main question is that you focus on time and approximation error, but I would like to see the statistical error estimation as the number of samples grow. Reproducing this Figure (and comparing with interior point) would illustrate that your computational approach maintains the favorable statistical properties of this OT estimator.\n\n[FZFGC] Fatras, K., Zine, Y., Flamary, R., Gribonval, R., & Courty, N. (2019). Learning with minibatch Wasserstein: asymptotic and gradient properties\n\nQUESTIONS:\nSee my questions above. At the moment I advocate for rejection because I think the paper needs a significant amount of clarification w.r.t. their contributions, background and related work, such that I would not recommend publication in such state. However, I may have misunderstood parts of the paper, and I hope the authors will clarify this by answering my questions.\n\nLIMITATIONS:\nThe authors adressed the societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper focuses on investigating kernel-based optimal transport estimation. The approach involves reformulating the problem as a nonsmooth equation model and utilizing the semismooth Newton method to solve it. The study demonstrates that the associated residual mapping exhibits **strong semismooth** properties, ensuring the applicability of the semismooth Newton method. Additionally, it is verified that the subproblem within the semismooth Newton method is well-defined, as it is equivalent to solving an invertible symmetric linear system. Finally, the proposed algorithm is supported by both theoretical guarantees, including global and local rates, and numerical experiments that highlight its superiority.\n\nSTRENGTHS:\n1. The algorithm is highly practical and can be easily implemented. The paper provides clear instructions on solving the subproblem and updating the parameters, making it accessible for real-world applications.\n2. The theoretical investigation is rigorous and well-founded. The authors define a suitable residual function and present both global and local convergence rates of the proposed semismooth Newton algorithm. \n3. The numerical experiments provide compelling evidence of the algorithm's efficiency compared to existing methods. The results showcase the superior performance and computational advantages of the proposed approach, reinforcing its practical relevance and effectiveness.\n\nWEAKNESSES:\n1. The global convergence rate of the proposed algorithm is dependent on an auxiliary sequence of iterates, which adds extra computational complexity to the algorithm. It would be helpful to provide further clarification in line 238 regarding whether the condition $$w_{k+1}=v_{k+1}$$ always holds. If so, the proposed algorithm will reduced to extragradient method. \n\n2. To show the power of semismooth Newton steps, the proposed algorithm should be compared with the pure extragradient method in numerical experiments.\n\nQUESTIONS:\n1. How to choose the hyperparameters $\\alpha_1,\\alpha_2$, and $\\kappa$ etc. in Algorithm 2?\n2. Is there any intuition to use the adaptive strategy (3.4)?\n3. In the proof of Theorem 3.4, the auxiliary sequence $\\{v_k\\}$ is not considered. It seems that the strategy in line 238 cannot be neglected and the case $w_{k+1}=v_{k+1}$ needs to be precluded under the conditions of Theorem 3.4.\n\nLIMITATIONS:\nSee weakness and questions for further details.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}