Simplifying Optimal Transport through Schatten-$p$ Regularization

Tyler Maunu

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, barycentric displacements, and cross-covariances in simplified settings. To efficiently solve the proposed program, we develop a mirror descent algorithm with convergence guarantees in the convex setting. Experiments on synthetic and real data demonstrate the method’s efficiency, scalability, and ability to recover low-rank transport structures. In particular, we demonstrate its utility on a machine-learning task in learning transport between high-dimensional cell perturbations for biological applications. All code is publicly available at https://github.com/twmaunu/schatten_ot.

Submission Type: Regular submission (no more than 12 pages of main content)

Changes Since Last Submission:

In addressing our comments to Reviewer PDPV:

We add a revised introduction paragraph in the contributions section to explicitly frame Schatten OT as a potential tool for scaling OT based methods and introducing low-rank structure into OT maps. We also emphasize how it can be used as a drop-in replacement, and add a section on justifying low-rank assumptions in practice for OT (Section 1.3). Finally, we add an introductory paragraph to Section 5.3 that describes how the computation of this map for 4i data plays an important role in ML applied to a biological problem. Finally, our bibliography has been corrected.

In addressing our comments to Reviewer Aqgm:

We have added Section 1.3, justifying the low-rank assumption in practice. We have also added baselines to our experiments as appropriate. We note that it was not appropriate to compare against Subspace Elastic OT in our cell perturbation experiment because it is not designed to recover a low-rank coupling or low-rank barycentric map (it is designed for low-rank displacements). We also emphasize in 4.3 how the nonconvex methods do not have guarantees even in toy models like those discussed in our paper.

In addressing our comments to Reviewer aTux:

We have edited our algorithm section to clarify the steps and computational cost of our method. We then added Section 3.2 discussing alternative algorithms and when they could potentially offer some advantages. We also add a paragraph at the end of 2.3 discussing how one could jointly regularize using entropy and the Schatten penalty we consider.

Code: https://github.com/twmaunu/schatten_ot

Assigned Action Editor: ~Ju_Sun1

Submission Number: 6193