Paper: Accelerating Sinkhorn algorithm with sparse Newton iterations (ICLR-2024)

Reviewer comment (Experiments):
It would be great to have comparison with other methods than Sinkhorn, although I agree that most methods can be combined with SNS, it would be beneficial to compare their convergence.

Author response:
The reviewer also comments that we can potentially combine the sparse Newton stage with other warm initializations. The reviewer is correct, and indeed this is a very immediate and potentially fruitful future research direction following our work. The authors chose to use Sinkhorn algorithm as the initialization because Sinkhorn-based warm initialization is the only method with a proven exponential convergence guarantee. As first-order methods will not have exponential convergence, the only warm initialization strategy that might have an exponential convergence guarantee will be GreekHorn (Altschuler et al, 2017), but such a guarantee has not been proven to the best of our knowledge. Thus, we choose the Sinkhorn initialization despite that it may be slower than other methods in practice. As mentioned in the conclusion section, a combination of sparse newton with other warm initialization is easily implementable for future research, and might lead to even better performance on entropic optimal transport.

Reviewer comment (Evaluation):
I believe that to compute the optimality gap (e.g. Fig.1) you need a ground truth (i.e. the minimizer), in this case how is it computed ? Is it the OT plan without entropy regularization ?

Author response:
A1: As Sinkhorn algorithm can be thought of as a matrix scaling task, our termination criterion is when the rows and columns are scaled to machine accuracy. The minimizer is the dual variable which gives rise to the entropic optimal transport solution with the target $\eta$. (Detailed information on how dual variable connects with the entropic transport plan can be found in Section 2 of our paper)

Reviewer comment (Reproducibility):
I might be missing something, but I don't see the target sparsity $\lambda$ in Alg.1.

Author response:
A2: The reviewer's observation is accurate; target sparsity is indeed not specified in Alg.1. In general, one can truncate a matrix $M \in \mathbb{R}^{n \times n}$ by setting the truncation term $\rho$ to be the $k$-th largest entry of the Hessian, so that one would achieve a $k/n^2$ target sparsity. This point is also made in Section 4 under paragraph "Complexity analysis of Algorithm 1".

Reviewer comment (Novelty):
The related work section is great, but maybe you can add other references that speed up Sinkhorn using a factorization of the ground cost [1,3] or Chebyshev polynomials with a sparse graph [2].

Author response:
A4: The authors thank the reviewer for the suggestion. We have incorporated these references into the related works section.

Reviewer comment (Evaluation):
The results in Tab.1 are on the number of iterations until reaching machine accuracy. They don't inform us on the quality of the approximation. It could be interesting to compare the accuracy of the plan in the case where we have a closed form (e.g. entropic OT between Gaussian distributions [4]).

Author response:
A5: The authors thank the reviewer for the suggestion.  In Appendix F, we include results on how entorpic transport solution converges in terms of transport cost for MNIST L1 and L2 examples (see Figure 3 under the revised text). In addition, we thank the reviewer for bringing the work [4] to our attention. We have included this work in the related literature section.

Reviewer comment (Writing):
The authors could add a brief review of Newton’s algorithm since it is an important piece of the paper. In particular, on the importance of solving the Hessian system.

Author response:
A6: We thank the reviewer for the suggestion. Due to space constraints, a detailed review of Newton's algorithm is omitted, but we direct readers to key references for comprehensive understanding, including Convex optimization by Boyd \& Vandenberghe and Numerical Optimization by Wright \& Nocedal, which are excellent references. We have added a comment in the main text that removing sparsification will lead to the Newton's algorithm for the Sinkhorn objective. We also added a comment on the importance of solving the Hessian in achieving super-exponential convergence.

Reviewer comment (Writing):
In eq. 1, you need to define $\eta$ and specify if it should be greater than zero.

Author response:
A7: The authors thank the reviewer for the suggestion. We will edit the manuscript accordingly.

Reviewer comment (Presentation):
Having a preliminary or background section could help. Parts of the introduction could be in this section, and definition such as the optimality gap could be added in that section.

Author response:
A8: We thank the reviewer for the suggestion. We tried to change the organization of the paper but didn't find a viable way to reorganize the text.

Reviewer comment (Writing):
In section 2 "Convergence of Sinkhorn", what is $\alpha$ ? Should it be $\eta$ ?

Author response:
A9: The definition of $\alpha$ is in footnote 2 of the submitted manuscript prior to edit. To improve clarity, we have included the definition of $\alpha$ in the paragraph directly.

Reviewer comment (Experiments):
A potential red flag is that, in the numerical results, only a single entropic regularization coefficient is considered (\eta=1200). I suspect that this is a weak regularization term which results in rather sparse couplings and is therefore favorable to their method. It is OK if the method is not beneficial in the strong-regularization regime. However, this needs to be acknowledged in the paper and not swept under the rug.

Author response:
### Response on the choice of entropy parameter $\eta$

The reviewer comments that the entropy regularization term $\eta$ is too large, therefore the regularization effect is too weak. We appreciate the very detailed feedback and acknowledge this concern. To address it, we added Appendix F with experiments with a wide range of $\eta$. Additionally, we would like to motivate the original choice of the value of the regularization parameter. 

The manuscript maps the (i,j)-th pixel to the point $(i/28, j/28)$, and then use L1 or L2 distance, which respectively corresponds to the Earth mover distance (EMD) and the $W_2$ distance. As a result, for the calculation of EMD (see the L1 distance example in Section 6), our cost matrix C used is 28 times smaller than using pixel distance. To align with studies using pixel distances, our $\eta = 1200$ effectively equates to $\eta \approx 42$ after accounting for the scaling factor. The choice of $\eta \approx 42$ is a reasonable choice in the OT research community. For instance, (Altschuler et al., 2017) uses $\eta = 1,5, 9$ for a 28*28 grid. For the L2 distance, we have not found a choice of $\eta$ among recent work, but we found that taking $\eta \approx 1000$ in the L2 case enables estimation of $W_2$ distance with roughly an 1\% error, which is similar to the desired level of accuracy in estimating EMD in (Altschuler et al., 2017). 

In Appendix F, we systematically address the question with experiments for L1 and L2 MNIST examples with a wider range of practical $\eta$. The SNS algorithm achieves performance improvement similar to the result shown in the main experiment in Section 6. As such, our algorithm shows a significant performance boost in very common instances of OT. If Appendix F and the reasoning behind the original choice address the reviewer's concern, we kindly ask the reviewer to reconsider the review score.

Reviewer comment (Experiments):
The method is not compared to any other current methods aside from the classical methods on which it is based (Sinkhorn and Newton).

Author response:
### Response on Quasi-Newton method

The reviewer comments that we should compare our work more with existing methods beyond Sinkhorn and Newton's method. We include in Appendix E the numerical comparison between our algorithm with quasi-Newton methods, which shows that quasi-Newton method performs much worse than our proposed method. Moreover, we cannot observe significant improvement over even the original Sinkhorn algorithm.

Reviewer comment (Experiments):
For the numerical section, I would like to see numerical results for the same data sets with several regularization strengths (e.g. \eta=10,100,1000). Additionally, I think it is important to compare the proposed method to other current methods for accelerating OT with entropic regularization. In particular to "Massively scalable Sinkhorn distances via the Nystrom method" by Altschuler, Bach, Rudi, and Niles-Weed (2019).

Author response:
### Response on the choice of entropy parameter $\eta$

The reviewer comments that the entropy regularization term $\eta$ is too large, therefore the regularization effect is too weak. We appreciate the very detailed feedback and acknowledge this concern. To address it, we added Appendix F with experiments with a wide range of $\eta$. Additionally, we would like to motivate the original choice of the value of the regularization parameter. 

The manuscript maps the (i,j)-th pixel to the point $(i/28, j/28)$, and then use L1 or L2 distance, which respectively corresponds to the Earth mover distance (EMD) and the $W_2$ distance. As a result, for the calculation of EMD (see the L1 distance example in Section 6), our cost matrix C used is 28 times smaller than using pixel distance. To align with studies using pixel distances, our $\eta = 1200$ effectively equates to $\eta \approx 42$ after accounting for the scaling factor. The choice of $\eta \approx 42$ is a reasonable choice in the OT research community. For instance, (Altschuler et al., 2017) uses $\eta = 1,5, 9$ for a 28*28 grid. For the L2 distance, we have not found a choice of $\eta$ among recent work, but we found that taking $\eta \approx 1000$ in the L2 case enables estimation of $W_2$ distance with roughly an 1\% error, which is similar to the desired level of accuracy in estimating EMD in (Altschuler et al., 2017). 

In Appendix F, we systematically address the question with experiments for L1 and L2 MNIST examples with a wider range of practical $\eta$. The SNS algorithm achieves performance improvement similar to the result shown in the main experiment in Section 6. As such, our algorithm shows a significant performance boost in very common instances of OT. If Appendix F and the reasoning behind the original choice address the reviewer's concern, we kindly ask the reviewer to reconsider the review score.

### Response on comparing performance with (Altschuler et al, 2019)

We greatly agree with the reviewer in that (Altschuler et al, 2019) is a technically very impressive work and constitutes one of the major developments in recent OT algorithm research. However, it is important to our work that kernel approximation of $e^{\eta C}$ of any kind is not used. The work in (Altschuler et al, 2019) achieves speedup in the Sinkhorn task by in part sacrificing the rank of the approximated kernel, and therefore there exists a tradeoff in computational speed and accuracy. In our work, one can conceptualize SNS as introducing an approximated preconditioner, but we do not change the optimization task itself. Thus it will be unfair to compare (Altschuler et al, 2019) to our work, though we can comment that their work and SNS respectively achieve order(s) of magnitude speedup when compared to the Sinkhorn algorithm.

Reviewer comment (Experiments):
Have you considered any quasi-Newton methods? Will these be applicable in the context of your paper?

Author response:
### Response on Quasi-Newton method

The reviewer comments that we should compare our work more with existing methods beyond Sinkhorn and Newton's method. We include in Appendix E the numerical comparison between our algorithm with quasi-Newton methods, which shows that quasi-Newton method performs much worse than our proposed method. Moreover, we cannot observe significant improvement over even the original Sinkhorn algorithm. 

### Aside: Analysis of quasi-Newton's suboptimal performance

The numerical results in Appendix E provide ample evidence that quasi-Newton shows undesirable performance. Furthermore, we reason theoretically why quasi-Newton method is unfavorable in the Sinkhorn setting. The reasons for quasi-Newton working worse than SNS is twofold. 

- Firstly, the use of a quasi-Newton method is only valid if the Hessian of the matrix is hard to obtain, and therefore an approximation of the Hessian is used in place of the true Hessian. In our case, obtaining the Hessian is only of $O(n^2)$ cost, and therefore the cost is dominated by solving the linear system involving the Hessian. Therefore, using a direct approximation of the Hessian has a much better accuracy than keeping track of a Hessian approximation such as in the BFGS algorithm. 

- Secondly, for a quasi-Newton algorithm such as BFGS, one would often be storing a full dense hessian, which leads to a $O(n^3)$ complexity. If a limited memory version of quasi-Newton algorithm is used, then the lack of hessian approximation accuracy hampers fast convergence. Thus, a dense Hessian approximation in a typical quasi-Newton algorithm (e.g. BFGS) is guaranteed to have a worse performance than Sinkhorn-Newton, while a limited memory Hessian approximation (e.g. in LBFGS) also tends to have a worse performance than SNS.

Reviewer comment (Theory):
One notable omission I felt was the lack of guarantees for computational complexity. Although the algorithm performs well numerically, there is no theoretical backing for that.

Author response:
### Clarification regarding computational complexity
The reviewer comments that the algorithm lacks guarantees for computational complexity. 
- We would like to point out that we have a detailed analysis of the per-iteration complexity analysis of Algorithm 1 in Section 4 (paragraph Complexity analysis of Algorithm 1). Our analysis demonstrates that the per-iteration complexity of the Sinkhorn-Newton-Sparse (SNS) algorithm is $O(n^2)$, matching that of the Sinkhorn algorithm.
- As for iteration complexity, we would like to comment that the Newton stage of SNS in practice convergences with significantly fewer (up to 1000-fold less, as can be seen in the numerical experiment) iterations than the Sinkhorn algorithm. As for theoretical analysis on the iteration count, the authors are constrained by the fact that proving strong convergence in convex optimization algorithm carries significant theoretical challenges, which is a fundamental constraint in the tools available to the convex optimization community. In fact, the main theoretical bottleneck for lack of tools to bound the iteration count of the Sinkhorn-Newton algorithm.

Reviewer comment (Writing):
Can the authors explain the notion of sparsity introduced in this paper? It seems to me that $\tau$ is just the proportion of non-zero elements. However, is it usual to approximate $M$ by a sparse $\bar M$?

Author response:
### Clarification on approximate sparsity
**More on approximate sparsity**
As the reviewer pointed out, $\tau$ is just the proportion of non-zero elements for a sparse matrix. For the Hessian matrix M, we use a novel concept of approximate sparsity (see Definition 1). Sparse approximation is a well-discussed topic in applied mathematics, and has been considered in OT community (see Section 2 under paragraph Sparsification in Sinkhorn). An approximately sparse matrix predominantly contains numerically insignificant entries, except for a few significant entries that one needs to keep track of. A $(\lambda, \epsilon)$ approximate sparsity intuitively means that only a $\lambda$ fraction of the entries of the matrix has significant entries, and the combined significance of the rest of the entries is at most $\epsilon$.

**Connection with existing OT work leveraging sparsity**
Existing works mainly concern about a sparse approximation of the entry-wise exponential of the cost matrix $\exp{(-\eta C)}$, and our work concerns the approximate sparsity of the transport plan matrix $P = X\exp{(-\eta C)}Y$ (see first paragraph of Section 1 for Definition). Moreover, our work is the first to combine sparse approximation with second-order method in the field of OT. Through this choice, we are able to achieve quite order(s)-of-manitude speed-up in matrix scaling. Moreover, to the best of our knowledge, our analysis of sparse approximation in OT in Theorem 1 and Theorem 2 is the first systematic study of approximate sparsity in the entropic regularized transport plan.

**We sincerely hope that we have answered your concerns, in which case, we would greatly appreciate it very much if you could increase your scores. Please also reach out to us in the discussion phase if there are additional questions.**

Reviewer comment (Novelty):
The theoretical results in this paper seem to be simple corollaries of existing resluts. For example, Eq. (7) which is crucial to the proof of Theorem 1, follows Corollary 9 of (Weed, 2018).

Author response:
### Clarification on the contribution of Theoerem 1
In this manuscript, the main technical contribution is the acceleration of the existing Sinkhorn algorithm procedure, and the theoretical proof justifies the sparsity argument. For Theorem 1, we synthesize results in (Weed, 2018) (Ghosal \& Nutz, 2022) and (Peyre et al., 2017) into one unified theorem, and we are able to tightly bound the approximate sparsity of the Sinkhorn algorithm after t Sinkhorn steps. Theoerem 1 directly justifies the direction of this paper, which is to use a sparse approximation in the use of the second-order algorithm. This has been corroborated with the numerical experiments, which show a 10-100 fold speedup in compute time. Our main theoretical contribution to Theorem 1 is in the novel synthesis of the presented results. To the best of our knowledge, our work constitutes the very first complete analysis of approximate sparsity of the Hessian matrix of the Lyapunov function (see Section 3, equation (4) for the Hessian, and see equation (2) for the definition of the Lyapunov function). Prior to our analysis, the understanding of this phenomenon was limited, highlighting the novelty of our contribution.

We address the reviewer's comment more directly in terms of novelty. In short, there are two main topics concerning the convergence of entropic optimal transport. First is (a) convergence of the Sinkhorn algorithm to the optimal entropic-regularized solution. Second is (b) convergence of entropic optimal transport to the original linear program solution. The topics (a) and (b) are in their own right fundamental research questions in applied mathematics. The work in (Ghosal \& Nutz, 2022) and (Weed, 2018) are the best analysis in topics (a) and (b) respectively. As shown in (Weed, 2018), its convergence analysis in (b) is shown to have a matching lower bound (Section 3, (Weed, 2018)). Similarly, the analysis in (Ghosal \& Nutz, 2022) is the current strongest result in the exponential convergence of the Sinkhorn algorithm. The authors of this manuscript are successful in providing a tight analysis in approximate sparsity. It is highly unlikely to obtain a stronger bound than the one obtained in this manuscript unless we make substantial improvements in the analysis of either (a) or (b), which would constitute a standalone work on its own.

Reviewer comment (Theory):
Theorem 1 requires that $\eta$ is large enough. However, I believe that a large $\eta$ will bring something negative. Otherwise, why don't people minimize the original problem without the regularization? Thus, I think the author should remark the condition that $\eta$ is large enough of Theorem 1 and clarify how the value of $\eta$ affect the convergence rate.

Author response:
### Response on the entropy regularization $\eta$

The reviewer is also correct in pointing out that a large $\eta$ will be bad for the Sinkhorn algorithm. Indeed, in practice, the Sinkhorn algorithm will suffer from choosing a very large $\eta$, which will slow down the algorithm considerably to a runtime that is unfavorable when compared to using a linear programming solver. However, in the case of large $\eta$, a variant of the Sinkhorn algorithm with annealing-based techniques, such as in (Chen et al., 2023), will lead to much better scaling in terms of $\eta$. 

In this work, we mainly consider the practical choice of $\eta$. For practical cases of $\eta$, it is usually much faster to use the Sinkhorn algorithm to approximate the optimal transport plan than to run the linear programming code. In Appendix F, we give a detailed discussion to show that our choice of $\eta$ in the experiment is similar to the range of entropy regularization chosen by practitioners. Moreover, we show that a bigger $\eta$ than the one we have chosen will not lead to a significant change in the quality of the solution in terms of the transport cost. For reference, in the L2 MNIST example we tested, our choice of parameter $\eta = 1200$ leads to a mere 1\% difference in transport cost from the ground-truth optimal transport cost. 

In addition, one reason for using strong entropy regularization (small $\eta$) is that the transport cost received can be noisy, and entropic regularization can help with dealing with noise, and therefore keeping a small $\eta$ term might help with denoising, as can be seen in the use of entropic regularization in machine learning training, which will follow a similar mechanism as in the paper ``Entropy-SGD: biasing gradient descent into wide valleys".

Reviewer comment (Theory):
The paper claims that ``to achieve possibly super-exponential convergence, we introduce Sinkhorn-Newton-Sparse (SNS)''. However, this paper does not provide any sound convergence analysis of SNS.

Author response:
The reviewer comments that the paper lacks a convergence analysis of SNS. While the point is valid, we would like to point out that the authors are constrained by the fact that proving strong convergence in convex optimization algorithm carries significant theoretical challenges, which is a fundamental constraint in the tools available to the convex optimization community. Conventional analysis in the convergence in convex optimization requires assumptions in the Lipchitz-ness on the utility function f and its Hessian (Boyd \& Vandenberghe, 2004), which is unfortunately unbounded in our case. Moreover, our utility function f, while being concave, does not satisfy self-concordance. Consequently, providing a substantive convergence guarantee for Newton's method in the Sinkhorn objective function is challenging, even without considering sparsity approximation. As such, we do not claim we can prove theoretical super-exponential convergence. Moreover, to address the reviewer's concern, we shall make it more explicit in the manuscript that we cannot prove theoretical super-exponential convergence.

We emphasize that our SNS algorithm has a very strong computational speedup in all of the presented cases. As OT usually deals with geometric transport problems or entropic assignment problems, we have shown strong evidence of speedup across a wide use case of OT. This is also true of Newton's algorithm in general, in the sense that the wide empirical success of Newton's algorithm leads to its wide use by the community even for problems for which it is not clear that the convergence guarantee exists. The same can be said for Sinkhorn's algorithm, for which exponential convergence has only been proven rather recently, but its adoption is due to its good numerical performance and ease of use. In Appendix F, we show that our speedup is robust for different entropy parameters.

Therefore, we emphasize that numerical soundness, aided with a theoretical guarantee for approximate sparsity, lends very strong support for SNS as the future standard of computing OT distance in the machine learning community. 

**We sincerely hope that we have answered your concerns, in which case, we would greatly appreciate it very much if you could increase your scores. Please also reach out to us in the discussion phase if there are additional questions.**