A Sinkhorn-type Algorithm for Constrained Optimal Transport

The authors thank the reviewer for the reviewer's time. We would like to address the technical and theoretical concerns raised by the reviewer. We anticipate changes to the manuscript, and so all equation number references and line references will be referring to the pre-discussion version submitted to this conference. If the response provides some clarity and is satisfactory to the reviewer, we kindly suggest that the reviewer consider an increase of the score.

Comment 1

The paper is hard to follow. The paper lists technical details and equations without further explanations. For example, which is function L in Line 202? There are no further comments or interpretations of Theorem 2. Sometimes, the author refers to the equation in the appendix, for example, Line 297. The message is not direct. For example, what is Section 3.1 for? Merge some points in Section 1.1 to make the contributions clearer. It is highly recommended that the authors reorganize and polish the paper to make it accessible...

Response:

We would love to have further feedbacks on the parts of the manuscript for which the reviewer thinks deserve enhanced clarification. We give some clarifying explanations here for the specific points raised.

For the function $L$, the function is the primal-dual form, which combines the original function with additional terms taking account of the equality and inequality constraints. The term $L$ is sometimes referred to as the Lagrangian. Following the reviewer's suggestion, we have edited the text by further specifying that $L$ is the primal-dual form.

For the Hessian matrix in line 297, the Hessian blocks would take quite a lot of space, and the main message is to show that the cost for running Algorithm 1 enjoys an $O(n^2)$ scaling. We omit the formula for computing the Hessian so as to not overload the notation. Following the reviewer's suggestion, we substantially changed the wording for the paragraph to improve clarity. We thank the reviewer for this valuable suggestion.

For the contribution section which is Section 1.1, we follow the reviewer's suggestion and moved some points from Section 3.1 to Section 1.1 to further motivate expert readers. To further address the reviewer's question, Section 3.1 involves intuitive though technically intricate acceleration techniques, and it is placed after Algorithm 1 in Section 3 as a clean exposition would not be possible before the technical build-up in Algorithm 1. We appreciate the reviewer's feedback and we thank the reviewer for this suggestion on Section 3.1 that improves the clarity of this manuscript. If the reviewer has more suggestions for reorganizing the paper, we would greatly appreciate any further suggestions.

Comment 2

What is the motivation for focusing on the constrained optimal transport? Although it could be theoretically interesting to extend Sinkhorn to this kind of optimization problem, the author should make the significance of studying the problem clear.

Response: We would love to provide further contexts. As of right now, the field of entropic optimal transport largely focuses on the setting of unconstrained optimal transport, and it has sprung up a lot of applications, such as domain adaptation [1], model fusion [2], unsupervised learning [3], and computer vision [4].

The motivation for constrained optimal transport typically is based on the fact that the cost matrix involved in such OT models may not be the best option. For example, in domain adaptation, the OT map concerns a transport between source data $\{x\}$ and target data $\{y\}$ and the overall goal is to have some alignment between the source and the target data through optimal transport. To model the cost matrix $C$, two natural choices for the distance function would be $c_1(x, y) = \lvert x - y \rvert_{2}$ and $c_2(x, y) = | x - y |_{1}$. One way to consider both distance functions is to linearly combine $c_1$ and $c_2$ to form a new distance function, and the constrained OT perspective would say that the OT coupling would need to optimize over $c_1$ while having some performance lower-bound for $c_2$, or vice-versa. Similarly, all of the aforementioned cases in [1-4] uses some form of distance function but a priori a canonical choice for which one is the best cost.

Therefore, using constrained OT might lead to better coupling matrix design, which could further help with OT applications in machine learning. Importantly, the cost for Algorithm 1 is $O(n^2)$, which means that it is feasible to adapt the proposed Sinkhorn-type algorithm to practical machine learning settings, and it enjoys a good scaling that makes it compatible with large data sets. The scope of this manuscript is that this work is a paper for fast algorithms, and our goal is to put forward an algorithm which enables machine learning practitioners to have more flexibility in using OT.

The authors thank the reviewer for the reviewer's time. We would like to address the technical and theoretical concerns raised by the reviewer. We anticipate changes to the manuscript, and so all equation number references and line references will be referring to the pre-discussion version submitted to this conference. If the response provides some clarity and is satisfactory to the reviewer, we kindly suggest that the reviewer consider an increase of the score.

Comment 1

The numerical experiments are based on (weak) assumptions such as the cost matrix entries being sampled from uniform distribution in case of random assignment problem or the Rademacher distribution in case of Ranking under constraints (appendix A). It is not clear how the proposed algorithm performs when the cost matrix may not conform to simple distributions. For the Random assignment problem, how would the algorithm behave if cost matrix is not generated through uniform sampling and instead has some other distribution? Do you have any observations or thoughts?

Response:

The authors would like to clarify that the choice of random assignment problem as the considered problem class is due to the fact that random assignments are considered to be harder cases of OT in the unconstrained case, and so the experiment is designed to be a challenge to the algorithm. Intuitively, the random assignment case implicitly models optimal transport problems when the data dimension is large. For example, when the distance function is $c(x, y) = \lvert x - y\vert_2^2$ and the source and target distribution sampled from distributions in very high dimensions, then the entries of the cost matrix one forms would behave very much like i.i.d. random variables.

Random assignment are also typically harder instances numerically. Indeed one can see that the ranking under constraint example converges much quicker. The Pareto front profiling problem considers MNIST images under constrained optimal transport, which is a widely considered example in optimal transport, such as in the Cuturi paper [1]. Algorithm 1 converges quicker than random assignment in the MNIST case under the Pareto front profiling example, though we omit the detailed performance in the MNIST case for simplicity. Overall, Algorithm 1 works very well for simple distributions.

For random assignment problems in the unconstrained case, Aldous [2] shows that the choice of random distribution of each entry doesn't matter so much, but rather the tail behavior of the distribution determines the convergence behavior. For the case of this work, the uniform distribution is chosen as the performance is well-understood and the constraints and thresholds can be set to easily ensure feasibility.

Comment 2

Solving large scale problems would help ascertain the claims made by authors about the usability and efficiency of their proposed approach.

Response:

Algorithm 1 and Algorithm 2 do have a $O(n^2)$ scaling. For comparable problems, when one has $K = 1$ and $L = 1$ (one equality and one inequality constraint) the run-time complexity of Algorithm 1 would take 2-3 times the compute time as that of the Sinkhorn iteration, and so the performance of the Algorithm 1 is quite robust.

To address the reviewer's concern, the authors will show further experiments that of larger random assignment problems with constraints. We will provide an update to the paper before the end of the discussion period.

The authors thank the reviewer for the reviewer's time. We anticipate changes to the manuscript, and so all equation number references and line references will be referring to the pre-discussion version submitted to this conference.

Comment 1

Firstly, the problem itself is not motivated well in the paper. The authors should consider working on the introduction to establish the relevance of this work better.

Response:

The authors agree with the reviewer the text lacks motivation that could otherwise make the paper stronger. One limitation faced by this work is that as of right now, the field of entropic optimal transport largely focuses on the setting of unconstrained optimal transport. To the best of our knowledge, the constrained optimal transport is not currently widely considered in the setting of machine learning. As the paper primarily works on fast algorithms, the aim of this work is to provide a satisfactory answer for numerical treatment for constrained optimal transport, and we hope that this work enables machine learning practitioners to have more flexibility in using OT. The authors think a stronger motivation than what was written would not be sufficiently backed by existing research. One notable example of constrained OT used in machine learning is partial optimal transport (POT), and we show that our Algorithm extends nicely into POT and has very good performance.

We propose a possible direction for constrained optimal transport in machine learning, and this is also hinted in the main text in the related work section. A typical use case for constrained optimal transport is when the cost matrix involved in the OT models has more than one good options. For example, in domain adaptation [1], the OT map concerns a transport between source data $\{x\}$ and target data $\{y\}$ and the overall goal is to have some alignment between the source and the target data through optimal transport. To model the cost matrix $C$, two natural choices for the distance function would be $c_1(x, y) = \lvert x - y \rvert_{2}$ and $c_2(x, y) = \lvert x - y \rvert_{1}$. One way to consider both distance functions is to linearly combine $c_1$ and $c_2$ to form a new distance function, and the constrained OT perspective would say that the OT coupling would need to optimize over $c_1$ while having some performance lower-bound for $c_2$, or vice-versa. This type of constrained optimal transport might help with machine learning practitioners obtain better coupling matrices for their down-stream applications. As the scaling of the Algorithm 1 is $O(n^2)$, there is indeed a lot of potential in using this Algorithm for machine learning.

Comment 2

Secondly, the only novelty I see is in Algorithm 1. However, I consider the theoretical contribution itself significant enough to overlook this shortcoming.

Response:

The reviewer thanks the reviewer for the comment. Part of the goal of this work is to provide a general-purpose iterative Bregman projection algorithm which was at the time not fully explored by Benamou et al [2] in the OT community. As the Benamou et al paper has a substantial time gap between now, a large part of the work is to integrate noteworthy ideas from optimal transport to the constrained setting. Incorporating the existing ideas are in itself a rather significant undertaking, and the author agree that the main novelty of this work lies in Algorithm 1 and in the treatment of inequality constraints by an entropic barrier.

We list out two noteworthy novelties of the work. First, while the authors also agree that the acceleration techniques are straightforward extensions of existing work, part of the novelty is in figuring out which techniques are more suitable. One important goal of this work is to consider and compare all noteworthy practical OT acceleration techniques that can be extended to the general setting of unconstrained OT, and we show through numerical experiments that the most reliable techniques for successful implementations are entropic regularization scheduling and sparse Newton iterations. Other techniques such as APDAGD or quasi-Newton methods (BFGS and L-BFGS) are shown to be less suitable.

Moreover, one noteworthy contribution of this work is our numerical treatment for partial optimal transport, which is excluded from the main text as it is a different formulation from the constrained OT considered in the main text. We show that likewise the Siknhorn-type algorithm and the sparse Newton iteration provides great advantage to the APDAGD algorithm considered in [3].

The authors thank the reviewer for the reviewer's time. We would like to address the technical and theoretical concerns raised by the reviewer. We anticipate changes to the manuscript, and so all equation number references and line references will be referring to the pre-discussion version submitted to this conference. If the response provides some clarity and is satisfactory to the reviewer, we kindly suggest that the reviewer consider an increase of the score.

Comment 1

My main concern is that I am not sure if the choice of an entropy regularization makes sense for the additional constraints. For the standard OT problem, this choice is justified as the dual problem can be solved by exact block coordinate descent, but if one needs to employ the Newton's method, then, why should not one use e.g. a logarithmic barrier (instead of entropic regularization), which has the self-concordance property and is potentially better for the doubling strategy.

Response: A rather significant benefit of entropic regularization is that this formulation allows for convergence guarantee by applying the analysis Weed (2018). The authors considered and experimented with a logarithmic barrier, and the performance for inequality constraint is comparable. Additionally, the authors were not able to find a good theoretical framework to analyze an entropic barrier for the coupling matrix while having a log barrier for the inequality constraint. The main work of this text still considers a simplified fixed parameter setting. In contrast, the log barrier framework would necessitate the doubling strategy, as it is typically not possible to achieve exponential convergence in the log barrier case without exponentially tuning the regularization strength. Therefore, the manuscript chooses the entropic barrier as it is much simpler for practitioners for implementation, while having no apparent numerical downside that the authors know of.

Overall, the authors agree with the reviewer that the log barrier is a fruitful direction for further investigation, and the authors' position on barrier function was also highlighted in the Conclusion section of this manuscript.

Comment 2

Most of the analysis seems like a straightforward extensions of previous works Weed (2018) and Altschuler et al. (2017). A general concern about Theorem 1 is that although it is formulated as an exponential decay, it really shows the requirement that the regularization parameter grows proportionally with the inverse of the duality gap. In practice, the gap can be extremely small for large problems, leading to extremely large regularization parameters. On the other hand, the result of Theorem 2 seems to be dimension-free, but it actually depends on the dimension through the regularization parameter. As such, I think that Theorem 2 should clearly reflect the dependency on eta (which is hopefully linear as in Altschuler et al. (2017)).

Response:

The authors agree with the reviewer and wish to comment that the dependence on the optimality gap is an unavoidable consequence of Weed (2018). The entropic LP setting has a natural "slow rate" of $O(1/\eta)$ convergence to the true LP cost, which would be robust to the optimality gap.

As for the algorithmic convergence in Theorem 2, we do not have a concrete answer for the dependence on $\eta$, as the term $c_g$ also has dependence on $\eta$, and we are currently unsure about the scaling, as the variation in the $a$ variable doesn't have a nice mathematical behavior as would be true $x, y$. However, without the dependence on $c_g$, we would have an $O(c_g \varepsilon^{-2}\eta)$ iteration complexity. We hypothesize that same as the case in Altschuler et al (2017), we would have $c_g = O(\eta)$ under reasonable assumptions, but

We can show this scaling by modifying line 1552-1555 of the main text by taking out the constant's dependence on $\eta$. In the proof, the constant $C'$ has a dependence on $\eta$. Explicitly writing down the $\eta$ dependence would show

and one could follow the rest of the proof and show the iteration complexity is $O(c_g \varepsilon^{-2}\eta)$. Assuming $c_g = O(\eta)$, this would be a quadratic dependence on $\eta$, and the numerical performance doesn't show any bad numerical behavior when $\eta$ is large. In fact, the regularization parameter $\eta$ in this work is typically chosen to be large to showcase performance. The analysis in this work does not focus on the $\eta$ dependence, and the authors suspect that a different technique focusing on the $\eta$ dependence would show a linear convergence bound that we observe from the empirical observation.