Table of Contents
Fetching ...

Simplifying Optimal Transport through Schatten-$p$ Regularization

Tyler Maunu

TL;DR

The paper develops Schatten OT, a convex framework for regularizing optimal transport by penalizing affine images of the transport plan with Schatten-$p$ norms. This unifies and generalizes prior regularizations (e.g., low-rank, elastic, and barycentric penalties) and enables direct theoretical analysis, including recovery guarantees for low-rank couplings and low-rank displacements. A mirror-descent algorithm with KL geometry is proposed to solve the resulting convex programs efficiently, with convergence guarantees for $p,q\ge1$. Empirical results on synthetic and real data demonstrate that Schatten OT yields simpler, interpretable transport maps with modest cost increases and scales to larger problems. The work forges a bridge between OT and compressed sensing, offering principled, scalable tools for interpretable optimal transport.

Abstract

We propose a new general framework for recovering low-rank structure in optimal transport using Schatten-$p$ norm regularization. Our approach extends existing methods that promote sparse and interpretable transport maps or plans, while providing a unified and principled family of convex programs that encourage low-dimensional structure. The convexity of our formulation enables direct theoretical analysis: we derive optimality conditions and prove recovery guarantees for low-rank couplings and barycentric maps in simplified settings. To efficiently solve the proposed program, we develop a mirror descent algorithm with convergence guarantees for $p \geq 1$. Experiments on synthetic and real data demonstrate the method's efficiency, scalability, and ability to recover low-rank transport structures.

Simplifying Optimal Transport through Schatten-$p$ Regularization

TL;DR

The paper develops Schatten OT, a convex framework for regularizing optimal transport by penalizing affine images of the transport plan with Schatten- norms. This unifies and generalizes prior regularizations (e.g., low-rank, elastic, and barycentric penalties) and enables direct theoretical analysis, including recovery guarantees for low-rank couplings and low-rank displacements. A mirror-descent algorithm with KL geometry is proposed to solve the resulting convex programs efficiently, with convergence guarantees for . Empirical results on synthetic and real data demonstrate that Schatten OT yields simpler, interpretable transport maps with modest cost increases and scales to larger problems. The work forges a bridge between OT and compressed sensing, offering principled, scalable tools for interpretable optimal transport.

Abstract

We propose a new general framework for recovering low-rank structure in optimal transport using Schatten- norm regularization. Our approach extends existing methods that promote sparse and interpretable transport maps or plans, while providing a unified and principled family of convex programs that encourage low-dimensional structure. The convexity of our formulation enables direct theoretical analysis: we derive optimality conditions and prove recovery guarantees for low-rank couplings and barycentric maps in simplified settings. To efficiently solve the proposed program, we develop a mirror descent algorithm with convergence guarantees for . Experiments on synthetic and real data demonstrate the method's efficiency, scalability, and ability to recover low-rank transport structures.

Paper Structure

This paper contains 38 sections, 5 theorems, 66 equations, 4 figures.

Key Result

Proposition 1

The coupling $\boldsymbol{P}^\star$ is optimal for eq:schatten_ot if and only if there exists a subgradient $\boldsymbol{G}^\star$ of $\|\mathcal{A}(\boldsymbol{P}^\star)\|_{S_p}^q$ such that

Figures (4)

  • Figure 1: Solution quality of Schatten OT versus regularization parameter for mixture of Gaussian data. Top row: small variance. Bottom row: large variance. On the left, we show the effective rank of the found solution; on the right, we display its transport cost. As we can see, Schatten-1 regularization can greatly simplify the transport plan without substantially increasing transport costs.
  • Figure 2: Solution quality of Schatten OT versus regularization parameter for Gaussian data with a low-rank perturbation. In the left display, we show the effective rank of the found barycentric displacements, and in the right display, we show the transport cost of the found coupling. As we can see, Schatten-1 regularization again simplifies the transport plan, but the transport cost increases substantially across the board.
  • Figure 3: Plot of log excess cost versus iteration for the mirror descent algorithm on the Schatten OT problem. Left: In this experiment, the regularization parameter is small, and the variance of the Gaussian mixture components is large. This shows sublinear convergence of the algorithm, as is expected by the theory. Right: we reduce the variance and increase the regularization parameter, resulting in an optimal low-rank coupling. Here, we see that the geometrically diminishing step size converges linearly.
  • Figure 4: Plots of the performance of Schatten OT on the 4i perturbation data of bunne2023learning. For reference, we compare in all plots with what one gets using a Sinkhorn coupling with a regularization parameter of 1. Top: We examine the performance of Schatten-1 regularization in recovering a low-rank coupling for two different perturbations. As we can see, Schatten OT can reduce the effective rank while not increasing the transport cost too much. Bottom: We now show the performance of Schatten-1 regularization in recovering a low-rank barycentric projection map for two different perturbations. Again, Schatten OT can reduce the effective rank of this map, though the transport cost now increases more.

Theorems & Definitions (7)

  • Proposition 1
  • Theorem 4
  • Theorem 6
  • proof
  • proof
  • Proposition 7: Hard spectral selection
  • Theorem 8