A Truncated Newton Method for Optimal Transport

Thank you for appreciating the practical performance of our method and discussion of related work.

Regarding the lack of a general theoretical rate, we acknowledge the lack of a complete global convergence rate theory for MDOT-TruncatedNewton; this was listed as an important future direction at the end of the paper. However, our efforts on local convergence analysis have yielded useful insights and very strong empirical results (more than $10 \times$ improvement in speed at high precision in high dimensions) that suggest the possibility of rates better than the best known theoretical rate $\tilde{O}(n^2 \varepsilon^{-1})$, at least for certain OT problems including some 24 problem sets we have benchmarked on (see comment above directed to all reviewers and Figs. 4-14 in the Appendix of the Rebuttal Revision).

$\epsilon$ vs. $\varepsilon$ notation cleanup. Among other improvements to the notation (listed under our response to 9Fpa), we changed the notation from $(\varepsilon, \epsilon, \epsilon^\prime)$ to $(\varepsilon, \varepsilon_\mathrm{d}, \varepsilon_{\chi})$ to improve clarity. Under the revised notation in the Rebuttal Revision:

• $\varepsilon_\mathrm{d}$ is the tolerance on dual gradient norm,
• $\varepsilon_\chi$ is the threshold on $\chi^2$ divergence and Alg. 4 chooses $\varepsilon_\chi = \varepsilon_{\mathrm{d}}^{2/5}$ (clarified in line 299 of the Rebuttal Revision),
• The variable $\tilde{\epsilon}$ has been removed entirely and Thm. 3.2 is now expressed in terms of $\varepsilon_{\mathrm{d}}$.
• $\varepsilon$ remains an upper bound on the excess cost $\langle P - P^*, C \rangle$ for a feasible plan $P$ returned by an algorithm, in line with notation in prior work (Altschuler et al., 2017; Dvurechensky et al., 2018; Lin et al., 2019).

Response to Questions

We revised Cor. 3.4 to be self-contained ($\lambda_2$ is defined again) in the Rebuttal Revision.

The phrase “helper Sinkhorn routine” referred to Alg. 3 - ChiSinkhorn. We revised the second sentence of Sec. 3: “In Section 3.1, we first develop a technique for obtaining the truncated Newton direction with convergence guarantees (Alg. 2) and introduce a helper Sinkhorn routine (Alg. 3) to improve the convergence rate of this algorithm.” The latter is called a “helper routine”, because it only serves to help control the convergence rate of the former in light of Thm. 3.2.

LogSumExp reductions along rows and columns of a matrix $X$ have been defined explicitly and renamed to $\mathrm{LSE}_r(X)$ and $\mathrm{LSE}_c(X)$ respectively, to improve clarity: $\mathrm{LSE}_r(X) \coloneqq \log \big(\exp(X) \mathbb{1} \big)$ and $\mathrm{LSE}_c(X) \coloneqq \log \big(\exp(X^\top) \mathbb{1} \big)$.

We thank the reviewer for their time and constructive feedback that has helped strengthen the paper in the Rebuttal Revision. We clarify several misunderstandings, and discuss minor fixes and newly added results in the Appendix of the Rebuttal Revision below.

Similarity with prior work in the Background section. We acknowledge the overlap noted by the reviewer, which reflects the shared context of our matching problem setup. In the Rebuttal Revision we have revised the wording and flow to reduce repetition. Note that lines 70–92 are purely expository and introduce no new information.

Convexity of problem (2): We apologize for a missing minus sign in the definition of Shannon entropy in line 66, which has been fixed in the Rebuttal Revision. We suspect the reviewer’s confusion about convexity and $-H(P)$ stems from that typo. Since Shannon entropy is concave in $P$, the negative entropy is convex, so that problem (2) is convex.

Provide the derivation of the sum-of-exponentials form for the dual (eqn’s (3) and (4)) in the Appendix. As suggested, we have provided a full derivation of sum-of-exponentials dual form (and its relationship to the log-sum-exp form) in Appx. 1.2 of the Rebuttal Revision. Note that both forms of the dual widely appear in the literature (see references in line 772 of the Rebuttal Revision; our form sum-of-exponents appears in Altschuler et al. (2017), Dvurechensky et al. (2018), Lin et al. (2019)).

A lot of engineering improvements: We argue that our paper bridges the spectrum between “a collection of engineering improvements” and “an integration of mathematically rigorous ideas.” On the theoretical side, we provide approximately 10 pages of detailed proofs on local convergence rate analysis and the per-iteration cost of the overall algorithm, the computational complexity of key subroutines, and novel connections (to our knowledge) to the Bellman equations in reinforcement learning—connections that at least one reviewer (9Fpa) found compelling, in addition to the rest of the theoretical results. Our theoretical framework using the theory of convergence of stochastic matrices for analyzing a 2nd order method in computational EOT is novel to our knowledge and, as we show, leads to a substantial boost in performance. These contributions go beyond a purely practice-focused approach and aim to advance both understanding and application.

Only two examples in Fig. 2: It is well-established that the cost function plays a critical role in governing convergence behavior. In Fig. 2, we tested on both $L_1$​ and $L_2^2$​ costs across two different tasks: upsampled MNIST and color transfer. This resulted in four distinct problem sets. Still, we added Figs. 5-14 to the Appendix showing the same exact trend as Fig. 2 across 10 different tasks (both with $L_1$ and $L_2^2$ costs), adding 20 more problem sets (see main comment above to all reviewers).

Note also; our study benchmarks against 9 baseline algorithms and transparently lists 3 failed baselines due to slow convergence. This level of comprehensive benchmarking goes beyond what is typically seen in work published recently at similar venues; e.g., Ballu & Berthet (2023), Tang et al. (2024).

Constraint violation not displayed in the y-axis of Fig. 2: For all algorithms (with one exception), Fig. 2 displays the excess cost $\langle P - P^*, C \rangle$ after rounding onto $U(\mathbf{r}, \mathbf{c})$ via Alg. 2 of Altschuler et al. (2017), ensuring there is no constraint violation (within 64-bit machine precision). In other words, constraint violations are already accounted for by the excess cost displayed in Fig. 2 thanks to Lemma 7 of Altschuler et al. (2017) and Hölder’s inequality. The use of rounding is now clarified in the caption of Fig. 2.

The only exception to the above is Alg. 3.5 of Feydy (2020), which performs poorly when run with “no-debiasing” followed by rounding. To give Feydy’s algorithm the best possible representation, we instead display its absolute error of the estimated cost under “debiasing” and without rounding, even though said approach does not return a feasible transport plan. This was described in Appx. D.

Dear Authors,

Thank you for the detailed answers! I am trying to understand the details behind the dual problem, however it is still confusing for me. Next, I list my questions, regarding the dual problem. If we converge to the common ground, I would increase my score.

1. Do you assume that $\sum_{I,j} P_{i,j}=1$ for the original primal problem?

2. I believe that it is not allowed to drop constant terms (line 720) in primal and dual optimization problems as it would break primal-dual properties and mismatch primal function values with dual function values. Otherwise, it should be written using argmin and argmax.

3. What is the benefit of sum-exp formulation over log-sum-exp formulation?

4. Could you kindly advise if I missed it, in what theorem do you prove the primal-dual convergence of the proposed algorithm meaning that the duality gap is getting smaller?

I am still reviewing the experimental part of the rebuttal and will return with it in the next answer.

Q4. Thank you for your question. Our algorithm is not a primal-dual method in the sense of simultaneously minimizing the primal and dual objectives while explicitly monitoring the duality gap. Instead, we focus on minimizing the dual objective until the $L_1$ norm of its gradient is below a prescribed threshold $\varepsilon_{\mathrm{d}}$ (L4 of Alg. 1), then use the rounding procedure of Altschuler et al. (2017) to map the dual solution to a primal solution that is strictly feasible. The rounding procedure ensures that the cost $\langle P, C \rangle$ does not increase by more than the dual gradient norm, thus providing bounded suboptimality. This approach does not rely on directly proving convergence of the duality gap but instead ensures an approximate primal solution with guarantees on feasibility. This is in the same spirit as Alg. 1 of Altschuler et al. (2017) and Alg. 2 of Dvurechensky et al. (2018). See also our high-level overview in lines 150-163, where we discuss the outer loop Alg. 1.

If this misunderstanding arose due to unclear phrasing in the manuscript, we are happy to revise the relevant sections to make this distinction clearer.

For the convergence of unconstrained dual minimization, while we did not prove a global convergence result (which is highly non-trivial for 2nd order methods), we proved a quadratic local convergence result in Thm. 3.5 in terms of the $L_1$ norm of the dual gradient. This is amenable to the aforementioned rounding lemma of Altschuler et al. (2017), which is also expressed in terms of the $L_1$ norm of the dual gradient.

Note that quadratic local convergence here implies that if we initialize sufficiently close to the solution so that step size $\alpha=1$ is admissible (no line search is necessary) and choose $\eta = ||g_k||$, then the gradient norm after a single truncated Newton update becomes $||g_{k+1}||_1 = O(||g_k||_1^2)$ thanks to Eq. (15); e.g., a decrease from order $10^{-3}$ to order $10^{-6}$. This is why the per-iteration computational cost of Alg. 4 in Cor. 3.4 is of importance; because if we can initialize close to the solution, we should expect Alg. 4 to terminate in a handful of iterations at worst and a single iteration at best. Fig. 1 and Table 1 show us that this is indeed possible with some intuitive heuristics (adapting how quickly we decay the temperature $1/\gamma$ during the execution of Alg. 1).

Dear Reviewer,

Thank you for your follow-up and for engaging deeply with the dual problem. We appreciate your openness to revisiting your evaluation and address your questions below.

Q1. In the derivation of the dual in Appx. A.1.2, we assume that $\sum_{ij}P_{ij} = 1$ for the primal problem, as the feasible set $U(\mathbf{r}, \mathbf{c})$ is a subset of the simplex. Any matrix whose row and column sums match probability vectors must have its entries sum to 1. Consequently, adding the term $1 - \sum_{ij}P_{ij}$ to the primal objective (as in line 731) does not affect the objective value for any feasible $P$, nor does it alter the optimal solution.

If we had not switched to the KL divergence with the uniform distribution (but continued with negative Shannon entropy), the Lagrangian in (17) would instead take the form $\gamma \langle P, C \rangle + \langle P, \log P \rangle + \langle \mathbf{u}, \mathbf{r} - \mathbf{r}(P) \rangle + \langle \mathbf{v}, \mathbf{c} - \mathbf{c}(P) \rangle$. Following the same steps as in our derivation, this would again lead to (3) after a reparametrization of the dual variables; see for example, derivation by Lin et al. (2019) in the second half of their Sec. 2.

In Appx. A.1.2, we opted for the KL divergence route in our exposition because of the nice cancellation of constants from line 740 to 741, and the eliminated need to reparametrize $\mathbf{u}, \mathbf{v}$.

Q2. This is correct if we were using the duality gap between the dual and primal objectives. However, in our work, we measure performance using the excess cost $\langle P - P^*, C \rangle$, i.e., the suboptimality of the primal solution returned by the algorithm. For this, we require (and ensure) that the algorithm’s output $P$ is strictly feasible, which we achieve using the rounding procedure of Altschuler et al. (2017). Their Lemma 7 guarantees that the rounding procedure does not increase the cost of $P$ by more than the $L_1$ norm of the gradient of the dual.

In the paper, we do not suggest that the optimal value of the objective in (2) equals the optimal value of the objective in (4). Rather, in line 85 we only say that Problem (4) can be solved instead of Problem (2). That is, if one solves Problem (4) to obtain $\mathbf{u}^*, \mathbf{v}^*$, then plugging these into the closed form in (3) yields the optimal solution $P^*_\gamma$ of Problem (2).

Finally, note that the objective of the unconstrained Problem (4) is in fact scaled by a constant factor of $\gamma$ for convenience (as mentioned in brackets in line 733). Specifically, we have $\gamma \langle P^*_\gamma, C \rangle - H(P^*_\gamma) = -g(\mathbf{u}^*, \mathbf{v}^*)$.

Q3. For our work, the sum-of-exponentials formulation is advantageous as it leads to simpler expressions for the gradient and the Hessian matrix, significantly simplifying the derivations and exposition in Eqs. (7–12), as well as some of the proofs in the Appendix. However, we acknowledge that some authors prefer the log-sum-exp formulation. For example, Lin et al. transitioned from the sum-of-exponentials form in their ICML 2019 work to the log-sum-exp form in their JMLR work. In their Remark 2, Lin et al. (2022) comment on this transition; they note that the log-sum-exp form of the objective is known to be smooth, and leverage this smoothness property in their theoretical analysis.

Note also that we can easily obtain the objective value for the log-sum-exp formulation by simply $L_1$-normalizing $P(\mathbf{u}, \mathbf{v})$ and re-evaluating (4). Let the normalizing constant $N(\mathbf{u}, \mathbf{v}) \coloneqq \sum_{ij} P_{ij}(\mathbf{u}, \mathbf{v})$ for any given $\mathbf{u}, \mathbf{v}$, where $P(\mathbf{u}, \mathbf{v})$ is given by (3). If we update $\mathbf{u}$ by a scalar shift as $\mathbf{u} - \log N(\mathbf{u}, \mathbf{v}) \mathbb{1}_n$ (or, equivalently update $\mathbf{v}$), the resulting $P$ lies on the simplex, and the new value of (4) matches the log-sum-exp objective value prior to updating $\mathbf{u}$.

Dear Reviewer,

Thank you again for your thoughtful questions and engagement with our paper. We have carefully addressed your concerns regarding the sum-of-exponentials form and related aspects in our rebuttal, providing detailed clarifications and adding new results to strengthen the manuscript. We hope these revisions address your concerns and provide the clarity needed to converge on common ground.

As the end of the rebuttal period is approaching, we kindly ask if you might have the opportunity to review our responses and share any further thoughts or updates to your evaluation. If there are any remaining points of confusion or areas needing additional clarification, we would be happy to address them promptly or apply revisions.

Thank you again for your time and valuable feedback.

Best regards,

Authors

We thank the reviewer for their feedback, and their appreciation of both our empirical and theoretical results.

We first clarify a point in Strengths.

$\log(1/\epsilon)$ convergence of Alg. 2 in Thm. 3.2. We would like to clarify that, the dependence on $\epsilon$ in Thm. 3.2, as the reviewer astutely noticed, is indeed possibly worse than $\log(1/\epsilon)$ due to other terms. In particular, we generally have $||\mathbf{r} - \mathbf{r}(P)||_1 \leq \chi(\mathbf{r} | \mathbf{r}(P))$; see Eq. 2.12 of Fill (1991). This is why we run ChiSinkhorn (Alg. 3) each time before Alg. 2 is called: to control the ratio $\chi(\mathbf{r} | \mathbf{r}(P)) / ||\mathbf{r} - \mathbf{r}(P)||_1$. The specific call we make to ChiSinkhorn in line 4 of Alg. 4 results in the $\tilde{O}(\epsilon^{-2/5})$ dependence seen in Corollary 3.4 (proof in Appx. A.1.5). This is now clarified in line 299 of the Rebuttal Revision.

Now, we address the questions, before responding to weaknesses.

Intuition on $P_{\mathrm{rc}}$.

For technical intuition on the matrix $P_{\mathrm{rc}}$, we refer the reviewer to Lemma A.1 in the Appendix and its proof, where we show that $P_{\mathrm{rc}}$ is a reversible, irreducible, row-stochastic matrix whose second largest eigenvalue $\lambda_2$ is strictly less than one, all eigenvalues are real and stationary distribution is given by $\mathbf{r}(P)$.

We have also added a less technical intuitive discussion of $P_{\mathrm{rc}}$ above Lemma A.1 in the Rebuttal Revision (lines 801-813) that the reviewer may find helpful. The reader is now referred to this discussion in the Rebuttal Revision of the main text, in a footnote immediately below Thm. 3.2.

Comparison with other second-order methods. Unfortunately, we were unable to find GPU-parallel implementations of alternative second order methods that can compute feasible, high precision solutions to the OT problem. Available options we could find use scipy’s L-BFGS solver, which is CPU-based and substantially slower.

High level intuition on proofs. The proofs rely heavily on the theory of convergence of stochastic matrices, techniques in the analysis of truncated Newton and linear CG methods in numerical optimization. Unfortunately, we had to defer proofs to the appendix due to space constraints with no room left in the main text even for sketches. For Thm. 3.2, the proof in the Appendix begins with a short high-level intuition; lines 1110-1115 in the Rebuttal Revision.

Now, we discuss Weaknesses:

Lack of control over and precise characterization of the spectral gap $1-\lambda_2$. Possible poor dependence on $n$

Firstly, we agree with the reviewer that a detailed analysis on the spectral gap and its dependence on problem parameters (e.g., properties of the cost matrix and the marginals) is an interesting problem, but we suspect it requires extensive theoretical work beyond the scope of this paper.

To alleviate the reviewers’ concerns around possible poor dependence on problem size $n$, we have revised Fig. 4 in the Appendix, which shows that the dependence on $n$ is no worse than $O(n^2)$, at least for the problems considered in the main text. We also remind the reviewer that Fig. 3 shows the solution of a color transfer problem at a staggering $n \approx 1$ million, which would not have been possible had the method generally suffered from a poor dependence on $n$. For the DOTmark experiments in the Rebuttal Revision (Figs. 5-14), while we have not explicitly ablated $n$, the runtimes at $n=4096$ largely agree with Fig. 2, so we expect similar trends in $n$. While it may be possible to construct edge cases where the spectral gap has poor dependence on $n$, the empirical results strongly suggest that for a wide variety of problems the method remains very robust.

Lastly, we note that our analysis of the algorithm covers the worst-case convergence of diagonally-preconditioned CG based on the condition number of the preconditioned coefficient matrix of the discounted Newton system. However, a more precise characterization of CG convergence accounts for the entire spectrum of the matrix, and shows that it may converge much more quickly than otherwise predicted by the condition number if preconditioned eigenvalues are well-clustered on the real line (Thm. 5.5 of Nocedal & Wright, 2006); see discussion immediately after Thm. 3.2. We suspect the efficiency of our method may be partly owed to this behavior and believe this too is an important area of future research.

Notational clarity and ease of reading. We hope that our efforts to improve readability have alleviated the reviewer’s concerns:

• We included intuition on $P_{\mathrm{rc}}$ in A.1.3 and referred the reader to it in the main text.
• We cleaned up the notation mixing $\varepsilon$ and $\epsilon$ to improve readability and clarity (details in our response to A2Da).
• We note that $\mu$ was defined the first time it was used in the main text. It is the smallest diagonal entry of $P_{\mathrm{rc}}$. This is now repeated the second (and last) time it is mentioned in the main text.
• We revised Cor. 3.4 to be self-contained ($\lambda_2$ is defined again).
• We added derivations for Eqs. (3-4) in the Appendix upon request by reviewer 4av9.
• LogSumExp reductions along rows and columns of a matrix $X$ have been defined explicitly in Sec. 2 (upon request by reviewer A2Da) and renamed to $\mathrm{LSE}_r(X)$ and $\mathrm{LSE}_c(X)$ respectively, to improve clarity.

Dear Reviewer,

In the latest revision, we added a section Appx. F dedicated to the discussion of the spectral gap dependence seen in Thm. 3.2 and Cor. 3.4, as well as a more in-depth empirical analysis of the subject.

First, we provide mathematical rationale (due to Lemma A.6) on why the $(1-\lambda_2)^{-1/2}$ rate can in fact be quite pessimistic, and $(1-\rho)^{-1/2}$ comprises a better bound, where $\rho$ is the discount factor found by Alg. 2 via annealing.

Then, we motivate via an example why it is often the case that, in practice, we have $(1-\rho)^{-1} \ll (1-\lambda_2)^{-1}$ when $(1-\lambda_2)^{-1}$ is very large.

Lastly, we conduct experiments shown in Fig. 16, validating that the convergence behavior of Alg. 2: NewtonSolve is indeed much better captured by $(1-\rho)^{-1/2}$. While $(1-\lambda_2)^{-1/2}$ may indeed explode in some cases, this theoretical dependence is effectively controlled/mitigated by our discounting framework, and is not seen to reflect negatively on the empirical performance of the algorithm (bottom row of Fig. 16).

We thank you for your constructive feedback, which has helped us better present the inner workings of the proposed method.

The pessimism of the spectral gap dependence relative to $(1-\rho)$ is now mentioned immediately after Thm. 3.2 in the main text (with added insight on the proof strategy), and the reader is directed to Appx. F for further discussion and empirical analysis.

In line 290, we have also added mention of the possible dependence on $\varepsilon_d$ for the other terms appearing in (14).

Sincerely,

Authors