Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Sinkhorn-type Algorithm for Constrained Optimal Transport

Xun Tang, Holakou Rahmanian, Michael Shavlovsky, Kiran Koshy Thekumparampil, Tesi Xiao, Lexing Ying

TL;DR

This work extends entropic OT by incorporating linear inequality and equality constraints and develops a Sinkhorn-type algorithm with a novel constraint-dual update. By deriving a dual variational form and proving exponential closeness of the entropic solution to the constrained optimum, the approach yields provable convergence behavior; the authors further enhance practical performance via entropy-scheduling and sparse Newton acceleration. Numerical experiments on random assignment with constraints, constrained ranking, and Pareto-front transport illustrate fast convergence and feasible constraint satisfaction, underscoring the method's applicability to complex ML problems. Overall, the framework enables efficient computation of approximately optimal transport plans under rich constraint structures, broadening the use of OT in structured ML tasks.

Abstract

Entropic optimal transport (OT) and the Sinkhorn algorithm have made it practical for machine learning practitioners to perform the fundamental task of calculating transport distance between statistical distributions. In this work, we focus on a general class of OT problems under a combination of equality and inequality constraints. We derive the corresponding entropy regularization formulation and introduce a Sinkhorn-type algorithm for such constrained OT problems supported by theoretical guarantees. We first bound the approximation error when solving the problem through entropic regularization, which reduces exponentially with the increase of the regularization parameter. Furthermore, we prove a sublinear first-order convergence rate of the proposed Sinkhorn-type algorithm in the dual space by characterizing the optimization procedure with a Lyapunov function. To achieve fast and higher-order convergence under weak entropy regularization, we augment the Sinkhorn-type algorithm with dynamic regularization scheduling and second-order acceleration. Overall, this work systematically combines recent theoretical and numerical advances in entropic optimal transport with the constrained case, allowing practitioners to derive approximate transport plans in complex scenarios.

A Sinkhorn-type Algorithm for Constrained Optimal Transport

TL;DR

This work extends entropic OT by incorporating linear inequality and equality constraints and develops a Sinkhorn-type algorithm with a novel constraint-dual update. By deriving a dual variational form and proving exponential closeness of the entropic solution to the constrained optimum, the approach yields provable convergence behavior; the authors further enhance practical performance via entropy-scheduling and sparse Newton acceleration. Numerical experiments on random assignment with constraints, constrained ranking, and Pareto-front transport illustrate fast convergence and feasible constraint satisfaction, underscoring the method's applicability to complex ML problems. Overall, the framework enables efficient computation of approximately optimal transport plans under rich constraint structures, broadening the use of OT in structured ML tasks.

Abstract

Entropic optimal transport (OT) and the Sinkhorn algorithm have made it practical for machine learning practitioners to perform the fundamental task of calculating transport distance between statistical distributions. In this work, we focus on a general class of OT problems under a combination of equality and inequality constraints. We derive the corresponding entropy regularization formulation and introduce a Sinkhorn-type algorithm for such constrained OT problems supported by theoretical guarantees. We first bound the approximation error when solving the problem through entropic regularization, which reduces exponentially with the increase of the regularization parameter. Furthermore, we prove a sublinear first-order convergence rate of the proposed Sinkhorn-type algorithm in the dual space by characterizing the optimization procedure with a Lyapunov function. To achieve fast and higher-order convergence under weak entropy regularization, we augment the Sinkhorn-type algorithm with dynamic regularization scheduling and second-order acceleration. Overall, this work systematically combines recent theoretical and numerical advances in entropic optimal transport with the constrained case, allowing practitioners to derive approximate transport plans in complex scenarios.
Paper Structure (31 sections, 5 theorems, 94 equations, 4 figures, 3 algorithms)

This paper contains 31 sections, 5 theorems, 94 equations, 4 figures, 3 algorithms.

Table of Contents

  1. Introduction
  2. Contributions
  3. Related literature
  4. Constrained optimal transport
  5. OT in machine learning
  6. Acceleration for Sinkhorn
  7. Theoretical foundations of constrained entropic optimal transport
  8. Notations
  9. Background
  10. Entropic optimal transport under constraint
  11. Variational formulation under entropic regularization
  12. Main algorithm
  13. Main idea of the algorithm
  14. Implementation detail
  15. Complexity analysis of Algorithm \ref{['alg:1']}
  16. ...and 16 more sections

Key Result

Theorem 1

For simplicity, assume that $\sum_{i}r_{i} = \sum_{j}c_j = 1$, the LP in equation eqn: constrained OT general form has a unique solution $P^{\star}$, and assume that $\lVert D_{k} \rVert_{\infty} \leq 1$ for $k = 1,\ldots, K$. Denote $P^{\star}_{\eta}$ as the unique solution to equation eqn: entropi

Figures (4)

  • Figure 1: Illustration of 1D optimal transport under different inequality constraints. By incrementally decreasing the upper bound for the Euclidean distance transport cost, the transport plan evolves from minimizing the Euclidean distance transport cost to minimizing the Manhattan distance ($l_1$) transport cost. The transport plan is computed with the Sinkhorn-type procedure in Algorithm \ref{['alg:1']}.
  • Figure 2: Random assignment problem. Plot of the proposed Sinkhorn-type algorithm in terms of assignment cost and constraint violation. Specifically, constraint violation is defined by $\mathrm{Violation}(P) = |\min(0, D_I \cdot \mathrm{Round}(P,\mathcal{U}_{r, c}))| + |D_E \cdot \mathrm{Round}(P,\mathcal{U}_{r, c})|$.
  • Figure 4: Pareto front profiling of Euclidean distance versus Manhattan distance.
  • Figure 5: Random assignment problem. Plot of the proposed Sinkhorn-type algorithm in terms of TV distance to $P_{\eta}^{\star}$.

Theorems & Definitions (12)

  • Theorem 1
  • Theorem 2
  • Theorem 3
  • Definition 1
  • proof
  • proof
  • Remark 1
  • proof
  • Proposition 1
  • proof
  • ...and 2 more