Simplifying Optimal Transport through Schatten-$p$ Regularization
Authors: Tyler Maunu
Venue: Accepted by TMLR ()

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.

Published in Transactions on Machine Learning Research (02/2026)
Simplifying Optimal Transport through Schatten-p Regularization
Tyler Maunu Deportment of Mathematics
maunu@brondeis.edi
Brandeis University
Reviewed on OpenReview: https: openreview. net/ forum? id= DIokkTC5VH
Abstract
We propose 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 our formulation enables direct theoretical analysis: we derive optimality conditions and prove recovery guarantees for low-rank couplings, barycentric displacements, and crvviccovariances in simplified settings. To efficiently solve the proposed program, develop mirror descent algorithm with convergence guarantees in the convex seiling. Experiments on synthetic and real data demonstrate the method's efficiency, scalability, and ability to recover low-rank transport structurcs. In particular, we demonstrate its utility on a machine-learning task in learning transport between high-dimensional cell perturbations for biological applications. All code publicly available at whtttt:://// /ttp::///tthbbb com/ wtw.umu/sccaatteennntt
1 Introduction
Optimal transport (OT) has emerged a fundamental computational tool across many areas, including machine learning, computer vision, statistics, and biology (Arjovsky et al. 2017; Peyré and Cuturi 2019: Schiebinger et al. 2019 Bonneel and Digne 2023). It provides a principled framework for comparing proba- bility distributions, and it has rich mathematical history (Villani 2008). While the combination of practical utility and deep mathematical theory has led to the broad adoption of OT ideas in mathematics, science, and engincering, finding ways to scale OT solutions and make them interpretable remains a fundamental research question (Cuturi et al. 2023 Khamis el al.] 2024). In particular, OT typically suffers from the curse of dimensionality (Chewi et al. 2025), and regularized estimators may lack sparsity (Genevay et al. 2019).
A long line of work has focused on making OT scalable and interpretable through regularization. The most classical of these is entropic regularization, which yields a strictly convex program that can be solved via Sinkhorn scaling (Sinkhorn) 1967 Cuturi 2013). More recent work has sought to increase efficiency and
interpretability through quadratic regularization [Bonnee] et al. 2018 Lorenz et al.) 2021), as well as low-
rank factorizations (Forrow et al.] 2019 Scot et al. 2021). These methods show promise in biological
applications, particularly in single-cell RNA sequencing analysis Klein et al., 2025
Another closely related set of recent works attempts to include sparsity in the OT map using elastic costs (Cuturi et al. 2023 Klein et al.] 2024 Chen et al. 2025). In these works, using different cost modifications can be shown to encourage sparse or low-rank transport displacements. This leads to OT maps with simple, interpretable structures.
Except for entropic regularization, our work simultancously generalizes all of the aforementioned methods in a unified framework. We believe that this unified picture can lead to more principled development of tailored regularization. Furthermore, the theory of OT has not yet fully leveraged the extensive literature

---

Published in Transactions on Machine Learning Research (02/2026)
regularization for scaling and interpretability present in other fields, such compressed sensing. In compressed sensing, the use of & or nuclear-norm penalties proxies for rank minimization has yielded provably efficient algorithms (Eldar and Kutyniok 2012) WWrghtt and Maj 2022p. Our general formulation marrics ideas from OT and compressed sensing, providing a bridge that we expect to fruitful for developing sparse and low-rank optimal-transport models moving forward.
1.1 Contributions
In this work, we present Schatten-p regularized OT, which call Schatten OT, This novel formulation is both general and amenable to direct theoretical analysis. We summarize the main contributions of our work:
We demonstrate how the Schatten OT program simultaneously generalizes a large portion of prior work on low-rank and sparse methods in OT, while also yielding new regularized formulations.
We propose a general mirror descent framework that efficiently solves the Schatten OT problem. Since the problems consider are convex, this algorithm is guaranteod to converge to the global solution.
Since our resulting optimization problems are convex, they allow analysis of the recovered transport plans. In particular, we theoretically demonstrate recovery of low-rank couplings, low-rank transport displacements, and low-rank covariance structures in toy settings.
Experiments on synthetic and real data demonstrate the flexibility and effectiveness of Schatten OT reducing dimensionality while the same time balancing low transport cost.
OT has broad applications in machine learning. For example, in applied workflows, OT can be used to align distributions across domains, time points, or modalities. The fundamental point of our work is to introduce low-dimensionality into OT couplings in a principled and effective way. Low-dimenssonalllyy is important in practice since can reduce memory and computational cost while simultaneously improving the interpretability of the transport structure. This is crucial for downstream tasks like label transfer single-cell data or domain adaptation using fcature embeddings. By making these biases convex, Schatten OT be effectively estimated and analyzed. A final advantage of framework is that it also unifies many prior regularizers (see Section 2.3). From an applications perspective, Schatten OT can be swapped in wherever practitioners currently use regularized OT and wish to encourage low-dimensional structure.
1.2 Related Work
Regularized variants of OT have become increasingly important. in current applied and theoretical re- search. The story of regularized OT begins with entropic regularization (Cuturi, 2003)), which has roots in |cchöddinger (1932). More recent regularizations include quadratic and sparse regularization (Blondel
et al. 2018 Lorenz al. 2021 González-Sanz and Nutz 2024). which seek encourage sparse structures
in the transport plan.
Other work has studied low-rank factorizations in couplings to scale OT. Forrow et al. 2019 define a notion of factored couplings. Scetbon et al. (2021 use this notion of low-rank factorization of the coupling to develop an efficient Sinkhorn algorithm for factored couplings. Later, Lin et al. (2021) use multiple couplings move through anchor points. Halmos et al. (2024) propose a new algorithm to optimize over the LC factorization,
and Halmos et al. 2025) use hierarchical low-rank structures.
Another line of recent work has studied the regularization of displacements. Cuturi et al. 2023 introduce the notion of elastic OT costs and show how to construct maps with sparse or low-rank structure. Later, Klein et al. 2024) introduce learnable parameters into these costs, enabling greater flexibility in selecting the regularizer. Chen et al. (2025) use neural networks to learn maps in these settings.
We note that the incorporation of Jow-dimensional structure in OT displacements dates back earlier subspace-robust notions of OT. Paty and Cuturi 2019) compute Wasserstein distances over worst-case
2

---

Published in Transactions on Machine Learning Research (02/2026)
subspaces in the ambient space. These methods have some practical statistical advantages Niles- Weed and Rigollet 2022).
We can broadly think of regularizing OT as encoding bias in the transport plan. However, there are many other ways the OT problem can be biased. For example, some works seek to encode biases by optimizing the ground cost used within OT. Alvarez-Melis et al. (2019) learn an OT with invariances using an alternating minimization procedure, and focus on optimization over Schatten-p balls. [Sebbouh et al. (2024) learn a
matrix M that defines inner product cost between mcasures on different spaces. Jiin et al. (2021) match
distributions in different spaces using separate linear transformations. We note that these works implicitly regularize transport, as in subspace-robust OT.
In seeking a principled way to regularize and scale OT, we draw connections with compressed sensing. Com- pressed sensing focuses on recovering sparse structures from data. Original foundational works concentrate on recovering sparse vectors using & regularization (Donoho 2006 Candès ct al., 2006). These ideas were later extended to low-rank matrices (Fazel et al. 2008), which used nuclear norm regularization. The use of more general Schatten-p norms followed this (Nie et al.) 2012). The extension of these regularizations to other settings has been fruitful. For example, Scarvelis and Solomon (2024) use it in the context of deep learning.
The ideas of compressed sensing are sceing a resurgence in the age of modern machine learning and AT. Sparse autoencoders have become a primary tool for practitioners studying mechanistic interpretability (Huben et al.] 2024). Sparse coding and rate reduction form recent framework for training deep models to develop "white-box" methods Yuu et al., 2020 2023). Compression as general technique is effective at demonstrating intelligence in simple puzzles (Liao and Gu 2025). These examples show the importance and practicality of developing theoretically principled compression techniques for machine learning and AI problems.
1.3 On the Low-Rank Assumption
Schatten OT is most useful when the transport between two high-dimenniooal distributions is controlled by a small number of latent degrees of freedom. Importantly, in our framework, the Schatten penalty is applied to an affine image A(P) of the coupling P, so "low rank" should be interpreted relative to the chosen transport object A(P) rather than as a blanket property of OT solutions. In our examples, depending on the affine map A, low-dimensional structure may appear as: (i) a coupling that admits a low-rank factorization, (ii) barycentric map (or displacement) that is well-approximated by low-rank operator, or (iii) low-rank sccond-order summaries such a transport covariance. As a concrete motivating example, if the coupling P has rank r, then it can be written as sum of rank-one factors, which can be interpreted as r transport modes FFrrow et al. 2019 Our method can be considered a convex analog of this past work.
Single-cell and high-content perturbation studies. In single-cell perturbation experiments, each con- dition can be viewed as an empirical distribution in a high-dimensional feature space (e.g.. gene expression, metabolic expression, protein markers, or learned image embeddings) (Schiebinger et al., 2019 Bunne et al. 2023 A low-dimensional assumption is most plansible when interventions act through a limited number of latent factors, so that the condition-to-cooniiion shift is largely explained by a few coherent directions. In this setting, promoting low rank in the coupling or in the induced barycentric displacement yields small number of transport archetypes that summarize how subpopulations move under perturbation, providing a compact object that can be inspected and compared across interventions.
Imaging and registration. Optimal transport has many applications in imaging (Haker et al., 2004 Ferradans el al. 2014 Solomon et al.] 2015). In image-based problems like registration, domain a aligmment, or matching related modalities, transport structure is more likely to be low-dimensional when deformations are coherent and can be parameterized by small number of modes. When this holds, encouraging a low- rank coupling or barycentric map promotes alignments that reflect a few dominant deformation directions rather than fully dense plan, while still controlling transport cost.
3

---

Published in Transactions on Machine Learning Research (02/2026)
Generative modeling and low-dimensional data manifolds. Wasserstein-basee machine learning has been effective in a variety of practical settings (Arjovsky et al. 2017 Courty et al. 2016). A low-dimensional bias is most natural when high-dimensional observations concentrate near a dimensionn manifold and the generator itself maps from a low-dimensional latent space to the ambient space. In such cases, the relevant transport between model and data distributions may be well-approximated by a locally low-dimensional mapping. Schatten OT provides a direct way to encode this bias at the level of the transport solution (e.g.. via low-rank barycentric maps or transport covariances), yielding an explicit complexity fidelity tradeoff in OT subroutines used within learning pipelines.
Diagnostic indicators in practice. While there are many applications where a low-rank structure is well-motivated, one should verify that this is indeed good inductive bias when encountering practical problem. A practical diagnostic for low-dimensional structure is spectral decay of the transport object of interest, sunmarized by effective-rank style measures (Roy and Vetterli 2007). Empirically, the low- dimensional regime suggested when principal components or other low-rank summaries capture a large fraction of variance within and across conditions. In practice, tuning the regularization parameter and monitoring how the estimated tool for assessing whether not the low-rank assumption holds.
decreases as the transport cost increases can be an effective
1.4 Notation
Bold lowercase letters are vectors and bold uppercase letters are matrices. We denote the set of integers [n] : [I.....n]. For vectors, II is the standard & (Euclidean) norm. For matrices, I|-Ils, is the Schatten-p norm, ie., the Ep norm of the vector of singular values, and || ||s, = || || is the Frobenius norm. The set of probability measures over 12d with finite 2nd moment is (2 and the subset of absolutely continuous measures is (,
1.5 Organization
First, in Section 2 we give the necessary background and outline our optimization program. Then, in Section we provide our algorithmic framework for solving the Schatten OT problem. After this, Section 4 gives theoretical results about the structure of Schatten OT couplings. Finally, Section presents experiments on synthetic and real data, highlighting the flexibility and advantages of our framework.
2 Background and Method
In this section, we first discuss background ideas in OT and compressed sensing, and then the Schatten OT method. We begin in Section 2.1 by describing discrete OT and its common regularizations. Then, Section 2.2 discusses background on Schatten-p regularization in compressed sensing. After this, we define our Schatten-p norm regularized OT, Schatten OT, in Section 2.3
2.1 OT and Regularization
For simplicity of presentation, we focus on the discrete case; in Appendix A we show how these ideas extend to the continuous sctting. Consider two discrete mcasures μ = "m 0,00, and v = Σj= byoy, where a,,b; by 2 0 and E. b, = We let X E pdxm and Y E pdx= be matrices with the support points as columns. Without loss of generality, assumee 2 m. The transportation polytope U(a, b) is the set of m X n nonnegative matrices whose rows sun to a = (d1,...,,,.. We also refer to these matrices as couplings between /a and v. Tn the transport problem, we are thinking of transporting μ to v. Therefore, we refer to μ as the source distribution and v as the target distribution.
and columns sum to = [br,...,b, bn]
We assume an m n matrix of costs C, where Cij = c( Hj), for some function c Bd Fd oo). The function e is typically called the ground cost. As a concrete example, we can the square of the Euclidean distance,
Cij -|11=-9)) 12
(1)
4

---

Published in Transactions on Machine Learning Research (02/2026)
OT seeks minimum-cost coupling between the measures j and v. It is formulated linear program over the transportation polytope,
OT(H,)) = min (P.C). Pell(a.b)
When the cost corresponds to a power of a metric some underlying space, the resulting OT cost can be used to define metric on P2(R)). For the rest of this paper, we will assume that c(a;,, = || yyll However, we note that our regularization can be applied to OT with any ground cost.
For a review of computational methods related to this linear program, see Peyré and Cuturi 2019 While there are many deep and interesting results related to OT, practice, the direct use of OT in the form presented can suffer. In particular, OT suffers the curse of dimensionality; the worst-case statistical rate of estimation for the p- p-Wasserstein distance is Onn- assuming that a; and y/ are i.id. samples from some population measures. On the computational ional side, for large n x m, the linear program incurs computational cost (n²)), and we must store an (n²²) variable in memory. To combat these issues, various regularizers have been considered, as we mentioned in the introduction. These regularized OT variants solve
min (C, + +R(P),
a.
where R: R - R is the regularization function and À is tunable parameter. An example is the entropy function R(P) = Ei, P.y (log P.y) - -1) (Cuturi 2013).
In this paper, we propose a structured choice of regularizer based on the Schatten norms of affine transforms of P. yielding a convex family of regularizers that can be swapped into any OT-based learning pipeline to encourage a low-dimensional structure.
2.2 Schatten-p Regularization in Compressed Sensing
Schatten-p regularization in compressed sensing served as way to generalize 6 regularization for sparse vector recovery. In particular, using Schatten-p regularization strictly more general than lp regularization because can encode vectors diagonal matrices, in which case ||>||, [|diag(>)]ls,-
Perhaps the most popular regularization is Schatten-1 (nuclear norm) regularization. This is typically used to relax the rank of a matrix, and in a variety of settings, nuclear norm minimization has been shown to recover low-rank matrices (Recht et al.) 2010 Candès and Tao 2010 Candès et al., 2011
These methods have been applied in a variety of settings. Some applications of these methods have included matrix completion and recommender systems (Nie et al. 2012). multitask lcarning (Zhang et al., 2018
high-dimensional covariance estimation (Gavish and Donoho 2017), and image processing (Xie et al.
2016
Optimization with nuclear norms, or more generally Schatten-p norms, involves a variety of algorithms. For example, nuclear norm minimization can involve saddle point or proximal algorithms (Nesterov and Ne- mirovski 2013), the latter of which involves singular value thresholding (Cai et al. 2010). Other algorith paradigms include primal-dual methods (Chambolle and Pock 2011) or ADMM (Yuan and Yang 2013 To solve the more general case of Schatten-p regularized problems for < p S 2, one typically resorts to Lagrangian-style methods (Nie al. 2012) or iteratively reweighted nuclear norm-style methods (Lu al., 2015). Another popular approach to nuclear norm minimization problems involves factor splitting to avoid
SVD computations (Srebro et al. 2004 Fan et al. 2019
2.3 Discrete Schatten-p Regularized OT
We now come to the main innovation of our work. We study new variant of regularized OT problems using Schatlen-p norms. We define the Schatten OT problem as
Sch-OT(j, ((Ai, qi, A.)))
(C,P) + Σ APIA Pel((a.b)
min
(3)
The idea of this program is regularize toward simpler couplings with respect to the maps A, Notice
that the Schatten OT problem relies on three sets of parameters: the regularization strengths X; > 0, the
5

---

Published in Transactions on Machine Learning Research (02/2026)
Schatten powers and exponents P.: 9: > 0, and maps A, : U(a, b) - R Provided that Pa. 2 1 and A, are affine, it is easily seen that Schatten OT is a convex program. With only one regularization term, this simplifies to Sch-OT(u, vi (x,p,q, A)) = (C, + (I)A(P)
We believe that convexity and generality are primary benefits of Schatten OT. It is general enough to cover many existing regularizations in the literature, as we will demonstrate shortly. The convexity of the problem leads to solutions that are easy to characterize, as we will show in our theory section. It also enables efficient solvers using convex optimization techniques.
While many past regularizations for OT fall into this framework, some do not. In particular, classical entropic regularization cannot be recovered from 33, because the negative entropy is not a Schatten norm of an alfine function of P. Nevertheless, entropy can be combined with Schatlen OT in a straightforward way by considering objectives with both the Schatten OT and entropic regularizers. For convex Schatten regularization, this problem remains convex and can be solved by the same KL. mirror descent algorithm: the entropic term contributes an additional term to the gradient, while the KL projection onto the transport polytope is still implemented via Sinkhorn scaling. In this way, can interpolate between standard entropic OT and the purely Schatten-regullrizedd regime studied in this work. Also, in Appendix A to extend Schatten OT to the continuous setting.
we illustrate how
Low-rank and Sparse Couplings: As a first example, consider the affine map A(P) = P. Then, depending on the choice of p. Schatten OT encourages low-rank or sparse couplings. In particular, choosing q p= = promotes low-rank solutions while preserving convexity, while choosing p < encourages low- rank solutions but is non-convex. We do not study this nonconvex case here. This yields a principled, optimization-based analog to the low-rank factorization pursued by works such as Forrow et al. (2019) Son et. al. (2021). ). On the other hand, q= p === 2 corresponds to the case of quadratically regularized OT [Blondel et al.
2018 Lorenz et al. 2021), since the Schatten-2 norm is just the Frobenius norm. Perhaps surprisingly, quadratic regularization in optimal transport tends to encourage sparse solutions (González- Group sparsity (Blondel el. al. 2018 can be achieved through proper choice of the affine maps in and setting P. 2, q
Sanz and Nutz 2024
Elastic costs:
We can also recover some of the elastic cost regularizations of Cuturi et al. 22023 Klein
et al.
2024g: Chen et al. 2025) In particular, can take q = p 1 and let A be the affine map
P diag (Prj(z: u.))
Then, the Schatten OT penalty corresponds to the & elastic cost. Group-sparse elastic costs can be recovered from sums of Schatten regularizations. We can also recover the subspace clastic costs of Klein et al. (2024) by taking q = p 2 and the affine map
P diag diag((Q: (QQ.P.,(a. - 91)) 1.)).......j......)))
where Qr is the projection onto the orthogonal complement of L.. Note that, analogously to Klein et al. (2024), one could include an additional minimization over L in the Schatten OT formulation. This then defines a family of learnable Schatten OT problems. We discuss this possibility further in the appendix.
Barycentric projection maps and displacements: The formulation can be used to penalize map estimators directly. In particular, given transport plan P, one can estimate a transport map using the barycentric projection
1 P. a:
1 (YPT).. 1..- a:
Tp(æ:))
Notice that this map is linear in P. Thus, we can penalize the barycentric projection map in our program by letting A(P) = A where we define A = di) As a further example, we could encourage displacements to be low-rank rather than the map itself. We call Tp(x;) - x, the barycentric displacemend. Then, to encourage these to be simple, we could use A(P) = YP 112 1/2 2 X/2 In both cases, the
additional scaling of 1//2 allows the population limit of the program to be well defined.
6

---

Published in Transactions on Machine Learning Research (02/2026)
Covariance regularization: All of the maps discussed far include zeroth and first moments of the support points (#;; Vj)- However, our formulation is flexible enough to include higher moments of our data distribution. For example, we can take A to be an affine function of the covariance induced by P.
Ep Σ
P
x;
yj
Vj
which is linear in P.
We could penalize the Schatten-1 norm of the cross-covariance Ei, P.ya.y,", which is an affine function of Ep. If the vectors æ; and yj are whitened (i.e., they each have identity covariance), then the singular values of this matrix correspond to the canonical correlations. Minimizing the Schatten-1 norm this case corresponds to minimizing the sum of canonical correlations, which soeks to increase independence between X and Y. Note that if we take the A(P) Eij P.j( (,) 4,) and p = 1, then the regularization is just the quadratic OT cost, Eig (z. - (x) Ils, Eij Pull - yjll?.
These illustrate just a few of the potential choices for adding covariance regularization to OT. In the appendix, we discuss covariance regularization in the context of Schatten OT for Gaussians. An in-depth study of these is left to future work.
3 A Mirror Descent Algorithm
The Schatten OT program in is convex whenever qi 2 and A; are affine, but solving it directly with off-the-sheit convex solvers (e.g.: CVXPY (Diamond and Boyd 2016) or interior point methods) is only feasible for small problems, as the transportation polytope U(a b) involves O(nm) variables and constraints. To address large-scale settings, we turn to first-order optimization methods. A particularly effective choice for optimization over the transport polytope is mirror descent with Kullback-Leiba (KL) geometry (Kemertas) et al. 2025). We use this algorithm for its simplicity and leave the analysis of more general methods, such as primal-dual algorithms or ADMM, to future work.
3.1 Mirror Descent for Schatten OT
Following the approach of Kemertas et al. 2025). we develop mirror descent using the KL geometry on the transport polytope. This choice is natural, since using the KL geometry replaces costly Euclidean projection with efficient Sinkhorn scaling. A few iterations of this method, followed by rounding, effective at projecting to the polytope (Altschuler al. 2017
Assuming that A(P) + 0, we can compute subgradient of the Schatten-p norm term in the Schatten OT problem as
aI|A(P))13) (UE*~IVT) E (1)))
where denotes the subdifferential. This provides a computable subgradient of F(P) at each iteration, provided that we can compute singular value decomposition. For p.q > 1, this is a bona fide gradient.
The mirror descent iteration involves the following steps. Define the cost function we wish to minimize as
F(P) = ( (C, M|A(())1
Then, given an initial coupling pd and sequence step sizes (7 (7 we iterate the following steps:
1. Form the SVD A(P*) U and the subgradient.
G k= -
(U+
Use multiplicative-weightt update
pk+1 Ptz exp(

---

Published in Transactions on Machine Learning Research (02/2026)
3. Project back to the transport polytope
phtl
() (Pi (1)+1).
Here, ()) is the projection onto the transport polytope with respect to the KL divergence, which can be implemented via Sinkhorn scaling. Thus we see that each iteration requires one SVD on A(Pk) (to evaluate the Schatten subgradient) and one Sinkhorn projection.
By standard mirror descent theory (Beck and Teboulle, 2003 Nemirovsky and Yudin 1983 Bubeck 2015 the method achieves an O(1/VT) convergence rate for convex objectives like F when P.q 2 1. The KL geometry ensures that nonnegativity is automatically preserved, and averaging can be used guarantee convergence of the objective values. While more sophisticated primal-dual or ADMM variants can be com- petitive in some regimes, carefully tuning and analyzing them for the full generality of (3) is nontrivial, and we leave a systematic study to future work.
The most expensive parts of the iteration are the SVD computation and the Sinkhorn projection. It is possible to use an adaptive low-rank approximation of A(Ph) throughout the iterations to increase computational efficiency. It would also be interesting to attempt to use sketching methods to approximate low-rank solutions to this problem |uurtsverr et al. 2021).
The choice of step size is essential. In gencral, since the problem is convex but not smooth in general, one could lake n x 1/v/ +1, 1, which yields the O(IvVT) convergence rate in objective value. In our experiments, we can observe faster convergence in specific settings. For example, when p = q= 1 and A has simple structure (A(P)) = P or A(P) = YP A-12), mirror descent can obtain faster convergence with a geometrically diminishing step size, the same schedule used in past work with sharp minima (Davis et al. 2018 Maunu et al., 2019
In practice, must tune the regularization paramcters À and p. For X, we can start from the unregularized OT cost and increase À until the effective rank, [|A(P)) (Il)))(ss/lAA(())) (P)lls drops while the cost increase is modest (see the examples in Figures and 2 On the other hand, the choice of is determined by what type of inductive bias one wants to enforce on A(P). For 1, we see the sharpest rank reduction; = 2 softens shrinkage but can enforce sparsity; p = oo amounts to operator-norn penalization.
3.2 Alternative Algorithms
A core computational bottleneck in large-scale OT solvers is the repeated enforcement of the row row/column constraints, i.e., projecting onto the transportation polytope U(a, b). We adopted mirror descent with Kullback-Lciller (KL) geometry in this work because projection onto ((,, b) coincides with Sinkhorn scaling (Sinkhorn 1967 Cuturi 2013). One could use gradient descent with respect to another geometry, such as the Fuclidean geometry. In this case, the projection corresponds to a constrained least-squares projection, which is substantially more expensive to compute. For the Sinkhorn method, there exist highly optimized implementations and near-linear-tine approximation guarantees (Altschuler et al. 2017). Using KL mirror descent, therefore, lets us reuse the same projection primitive that underlies scalable entropic OT solvers, while optimizing a much more general convex objective.
There are several natural alternatives, including primal-dual splitting methods (Chambolle and Pock 2011 and ADMM (Yuan and Vang. 2013). These approaches can be attractive when the problem admits a convenient splitting (c-g., scparating the transport constraints from a proximal step for the Schatten penalty), but they typically introduce auxiliary variables and dual constraints. In our setting, this can obscure the transport-polytope projection structure and make it harder to directly reuse Sankhorn-based OT primitives.
Regarding SVD computations, the cost is highly formulation-dependeni within the Schatten OT framcwork. In several important cases, the Schatten term admits simple ssabienns/uubgrraienns that avoid a full SVD altogether; for example, when p = = 2 and A(P) = P, the penalty reduces to a Frobenius norm, and similarly for group-sparse Frobenius penaltics on blocks of P. Likewise, for clastic-cos--style constructions in which A(P) is diagonal or block-diagonal, the required spectral computations reduce to clementwise or small-block operations. Even in the nuclear-norm setting (p 1), when solutions iterates are numerically
8

---

Published in Transactions on Machine Learning Research (02/2026)
low-rank, one use thin/truncated SVDs randomized approximations to reduce per-iteration cost. Exploring primal-dual variants that syslematically exploit such structure is interesting direction for future work.
4 Theory
In this section, we present our main theoretical results on the Schatten OT program. First, Section uses convex optimization theory to outline the structure of solutions to the Schatten OT problem. After this, Section 4.2 uses this structure to prove two theorems that demonstrate Schatten OT's ability to recover low-rank couplings and barycentric displacements. Finally, in Section we finish with a discussion of our theoretical results.
4.1
General Structural Theorems
Let P be an optimal solution the optimization problem (3) with p. 2 1. Since this a constrained convex optimization problem in P, we can appeal to standard theory. The solution is characterized by the KKT conditions, which state that there exists a subgradient G" @)],A(P*))13, and dual variables u² R and v* Rn such that
C XG" + +u"1, 1 0,
pd 2020, p11=a, p=
Comparing these conditions the optimality conditions for standard OT, we notice that the only difference is the inclusion of the XG in the first-order stationarity condition. Thercfore, the optimality conditions for Schatten OT are precisely those for an OT problem with the billed cost
S(A, G*) = C' AG =
(4)
where G is some subgradient of [IA()) at P The obstacle to our directly applying this result is that we do not know G*, since that would require knowing P
We can state this characterization as the following proposition.
Proposition 1. The coupling P is optimal for 3 if and only if there exists a subgradient G of ||A('') such that
p E
(S(), G!),
While cannot apply this to the direct computation of P", we use it to characterize solutions to the Schatten OT problem. In the following section, we develop this idea to prove the recovery of low-rank structure in the Schatten OT problem.
4.2 Discrete Recovery Theorems
In this section, prove low-rank recovery theorems for Schatten OT. While these are restrictive toy exam- ples, they represent the first such exact recovery results in the OT literature. We believe this is first step toward applying compressed sensing idcas to rogularized OT problems. is an open qucstion for future work to extend these ideas to the recovery of simple couplings in more complex settings. We begin in Section 4.2.1 with a recovery result for low-rank couplings. Then, Section 4.2.2 presents consequence on the recovery of low-rank set of barycentric displacements. These recovery guarantees play the same role recovery results in compressed and matrix sensing: they certify when the imposed structural bias is correct and will be recovered, enabling principled interpretability.
4.2.1 Low-Rank Coupling Recovery
We now assume that both the source and the target consist of R well-separated clusters, each with the same cardinality. We show that, for nontrivial interval of regularization strengths X, the nuclear-norm penalized

---

Published in Transactions on Machine Learning Research (02/2026)
OT problem recovers rank-R, block-diagonal coupling that matches each source cluster uniformly to its corresponding target cluster. We assume uniform marginals = 1/m, by 1/n for all i,j-
Our first assumption is on the clustered structure of and v.
Assumption 2. For two measures μ -i a,dm, and = ンn by0y₁. Rg, Rg for integers R. 2 1. The source indices [m] and target indices [n] are partitioned into clusters , respectivcly, where |St] = [T = for all t.
SR and andTi..... Ti...., TR.
Let p)) denote the closed Euclidean ball of radius > 0 around 2 E Rd. Our goal is to construct a setting where the cluster S is uniformly matched to Tc. We make the following assumptions about the locations of the source and target points.
Assumption 3. For two measures
Ei= 0;6m, and EE=1 by8y,
The source points lie disjoint balls, as B(e,, p) for E Si, and the target points lie in disjoint balls, Y E B(d..)) for j E Tr.
Jy - for all S₁ and j.j' T fort= Within matched clusters St and Tú, || = ||.
3. The minimam ênter-cluster distance T mingzt C dzll and the mazimumm indra-luuter distance
7 max, l|e, de satisfy
r 7 4p>00 0.
Notice that, under our separation condition 5 the OT coupling actually respects the cluster structure, in the sense that it must match points in St to Tr. Furthermore, any plan that matches z; to Yj within clusters (when i E St, j E T:) optimal. However, these matched clusterings are not low-rank; they are full-rank. On the other hand, as we will show the following theorem, a low-rank matching can be recovered from Schatten OT.
Theorem 4. Let Assumptions and hold, C be the quadratic cost, and let the excess cost for an across- cluster matching be
{lIc: - y3112 - -11e, - 13
min
min S iES.JTT.. a.'ee,,
Then, for any regularization pammeete X satisfying
0 s x g. Ami:
(6)
the minimizer of (C, +\|PPls is rank R coupling supported blockwise (S TT) that is uniform within clusters.
The essential idea of the proof is to ensure that 1) P is the unique coupling that respects the cluster structure and also minimizes the nuclear norm, and 2) there is way to make the nuclear norm term even smaller by using across-cluster matches without incurring more cost.
4.2.2 Low-Rank Displacement Recovery when
We now give a concrete example of how to recover a coupling with low-rank displacements. For our affine map, we consider the weighted barycentric displacement matrix A(P),
A(P) = (Tp(X) 2 =PP A 1-2 1/2
We also again assume that p = == in the Schatten OT formulation. For recovery result, make the following assumptions.
Assumption (Symmetric two-target clusters with separation). Fix orthonormal vectors u, v Rd and integer R 2. Let 0 pr <... < pn be distinct scalars and put m, nhd E IRd for & [R]. Suppose:
10

---

Published in Transactions on Machine Learning Research (02/2026)
1. The source set [m] is partitioned into nonempty clusters Si.... ......,p SR and there crists p such that a; E m + [-P, for all i E St.
2. For some > O, the target support consists of the 2R points
Vt = m, + tE, yt.- m. =v,t E [R]
with target masses b,, = bi, E as.
3. The minimal separation between clusters is lower bounded
A = min ju Mal > 2p.
aft
By symmetry, it is easy to see that all couplings that assign the source points in cluster to yt.+ or yt.- are optimal. In particular, it does not matter how the points assigned within clusters. For Se and + it is also convenient to define the inter-cluster cost gap
Lis .. Us.1 112 - Im: v1,1 (pt Ms)² + 2(1 Ms)Es,
(7)
where m, + E.u and |5 S
We now state the main recovery theorem for this setting. says that, under our assumption, we can exactly recover a coupling with rank-1 barycentric displacement.
Theorem 6. Let Assumption 5 hold and suppose C is the quadratic cost. Defone the explicit threshold
Amax A >
Then for every X E )mux), the unique mánimizer of (C, + NIA(())lss where A(P) is the weighted barycentric displacement, malches æ, to (v.+)} for i E S,. Furthermore, it yields a rank-1 barycentric displacement.
4.3 Discussion
While we demonstrate low-rank recovery only in toy examples here, our methodology highlights the advan- tages of using convex formulations. In particular, it is easier to verify that the recovered solution is low-rank. In fact, these are the first guarantees of low-rank recovery within an OT problem in the literature. This contrasis with nonconvex methods, which currently lack guarantees (Forrow et al. 2019 Klein el al. 2024)
In the future, it would be interesting to demonstrate exact recovery in more general settings using Schatten-p regularization when p 1. It would also be interesting to develop more general recovery conditions to move beyond the toy examples considered here.
5 Experiments
In this section, we give some simulations on real data that demonstrate the advantages of the Schatten OT formulation. First, in Section 5.1 we demonstrate Schatten OT's ability to recover low-rank couplings and barycentric projection maps. Then, Section 5.2 we examine the convergence rate of mirror descent solve the Schatten OT problem. Finally, Section 5.3 gives an experiment on real data that demonstrates the ability of Schatten OT to recover simpler couplings with 4i perturbation data.
In these experiments, also compare against other baseline regularized optimal transport methods. For recovery of low-rank couplings and low-rank barycentric projections, we compare against the Sinkhorn al- gorithm (Cuturi 2013) as well as Low-Rank Sinkhorn Factorization (Scetbon et al. 2021). Both methods are run using the Python Optimal Transport package For the recovery of low-rank displacements, we compare against Sinkhorn and entropically regularized Subspace Elastic OT (Klein al. (024).. Since publicly available code for Subspace Elastic OT was not available, we implement this as alternating mini- mization over couplings and subspaces, where the coupling is found via entropic optimal transport with fixed regularization parameter and the subspace is found via PCA.
https://pyythontt.gteee kit
11

---

Published in Transactions on Machine Learning Research (02/2026)
5.1 Low-rank Recovery
To begin, we examine properties of the Schatten OT problem, and we use CVXPY solve the convex program exactly.
In our first experiment, examine the Schatten OT's ability to recover low-rank couplings. To measure the quality of the recovered P*, we use two metrics: effective rank and transport cost. The latter is defined as (C, P The former is the ratio of the muclcar norm to the operator norm:
| |s |
Effective Rank(B)
In our experiments, the support of consists of two clusters centered at (-2,2) and (-2, -2), and consists of two clusters centered at (2, 2) and (2, 2)). The points within each cluster are Gaussian. We sample 10 points from each cluster, so that n = m = 20. The results are averaged over five randomly generated datasets.
In the first experiment displayed in the top row Figure ! we use a variance of 0.04 for cach Gaussian component. On the left, we show the effective rank 0o. On the right, we display the excess transport cost compared to standard optimal transport with the quadratic cost. As we can see, muclcar norm regularization can significantly reduce the effective rank without substantially increasing the transport cost. Using the Schatten-2 norm can also reduce the effective rank, albeit more gradually. This contrasts with Sinkhorn, which recovers lower-rank couplings since it is biased toward product couplings but incurs a higher transport cost. Low-rank Sinkhorn recovers a good low-rank coupling for well-chosen ranks and a small regularization, but if the regularization parameter too large, it becomes unstable.
regularization strength for p 1,2, and
Schatten
6.6
3.6
2.3
2.0
2.0
Schatten p=1
0.9
0.0
0.1
0.1
0.1
0.1
14
17.5
Schatten p-2
17.9
15.9
12.3
9.7
5.0
Schatten p=2
0.9
0.0
Excess Transport cost:
Bfwcct 000 Rank
10
Schatten pminf-
12.5
Schatten p=inf
B
Sinkboru
3.5
2.6
2.3
Sinkhorn
3.3
0.1
0.1
0.1
0.1
0.8
6
7.5
LR Sinkhorn r-2
2.0
2.0
1.9
1.2
1.0
1.0
LR Sinkhorn r=2
0.1
0.1
15.2
9.3
8.0
8.0
5.0
22
LR Sinkhorn r-4
2.5
LR Sinkhorn r=4
0.1
Sinkboru
174
15.4
11.9
9.1
6.6
3.9
2.6
Sinkhorn
0.2
2.8
14
17.5
I.R Sinkhom r-1
1.0
1.0
LR Sinkhorn r=1
Excess Transport cost
Effectis Rank
LR Sinkhorn r=2
1.8
Sinkhorn r=2
-10
12.5
Sinkhorn r=4
Sinkhorn r-4
8
6
Schatten p-1
7.5
Schatten p=1
0.1
4
Schatten p=2
-5.0
Scballea
0.0
2
Schetten pminf-
20.0 20.0
20.0
20.0
20.0
20.0
20.0
-.5
Schatten p-inf
0.0
0.0
0.0
0.0
Figure 1: Solution quality of Schatten OT versus regularization parameter for Gaussian mixture data. Top row: small variance. Bottom row: large variance. On the left, we show the effective rank of the found coupling for each method; on the right, we display its excess transport cost. As we can see, Schatten-1 regularization can greatly simplify the transport plan without substantially increasing transport costs.
12

---

Published in Transactions on Machine Learning Research (02/2026)
In the bottom row of Figure repeat the experiment from the top with one modification: we use larger within-cluster variance of 2. As can see, it is more challenging to find low-rank transport plan, and when one is found, it increases the transport cost more substantially.
The next experiment simulates the recovery of a low-rank transport displacement. Here, we wish to estimate a transport from measure j to a measure и. We assume that j standard Gaussian, and samples from are generated by y: = 2; + {, where u random unit vector and & standard Gaussian. We now wish to recover coupling with low-rank barycentric map, and repeat the same experiment in Figure but now use A(P) YP A-1 X/1 Figure 2 displays the results of this experiment. As see, Schatten-1 regularization again recovers coupling with lower-rank displacements without substantially increasing transport cost. Here, Sinkhorn does not recover low-rank displacements. Subspace elastic OT performs well, but we note that one needs to specify the rank, while Schatten OT automatically detects low-rank structure.
Sinkhorn
1.7
1.8
Sinkhoru
3.0
1.7
2.5
Excesss Transport coot
Subspace
1.5
Subspace
0.5
1.6
Schatten
144
1.3
1.3
Schatten
0.0
0.0
0.4
0.5
1.5
1.5
Schalten 1.4
1.5
1.5
1.5
1.5
1.5
1.4
Schatten
0.0
0.0
0.0
0.1
022
3.4
0.5
1.0
Schatten p=inf
1.5
1.5
1.6
1.7
1.7
1.8
1.9
1.3
Schatten p=inf
0.0
0.0
0.0
0.1
0.3
0.5
0.6
0.5
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, show the excess transport cost of the found coupling. As we see, Schatten- regularization again simplifies the transport plan, but the transport cost increases substantially across the board. Subspace elastic OT performs well but requires one to know the rank in advance extra parameter.
5.2 Convergence Rates
Next, examine the convergence rate of the mirror descent algorithm with Schatten-1 regularization in two settings. In the first setting, show sublinear convergence; in the second, linear convergence.
In the left display of Figure 3 we use a sctting where we do not expect a low-rank coupling to be casy to find. We set À = 0.1, and the data is a mixture of Gaussians with centers at (-2, =2) and v is a mixture of Gaussians with centers (2, +:2). The variance set to be 1, and n m We observe slow sublinear convergence. Note the log scale on the x-axis.
the same setup before, except now set À = 10 and let the clusters have variance 0.04. Now, since the recovered coupling is low-rank, mirror descent with geometrically diminishing step size converges linearly. This implies that the objective is sufficiently sharp, which can be exploited by this step-size schedule.
In the right display of Figure 3
5.3 4i Perturbation Example
Our final experiments are on transport estimaaiion in a biological example. We consider a standard set ting in single-cell perturbation analysis, where each experimental condition represented by an empirical distri- bution high-dimensional expression space. A central goal is learn how the cellular state distribution changes under intervention, i.e., to model the distributional shift between conditions. In this context, an optimal transport coupling between the two empirical measures can be viewed data-driven alignment
13

---

Published in Transactions on Machine Learning Research (02/2026)
10
4 Hterate)
10
2 INK laveragel
122000 (last
FIP F(PP [1og scale)
1/2 (averige)
FIP) FPP (kog scall)
10-1
10-2
1ick (tarite)
10-3
10
1/2 clest iteratie)
(average)
1/2 (average)
2000
6000
6009
8000
10000
0
200
400
600
800
1000
Iteration
Itecaioo
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 large. This shows sublinear convergence of the algorithm, as expected by the theory. Right: we reduce the variance and increase the regularization parameter, resulting in an optimal low-rank coupling. Here,
that the geometrically diminishing step size converges linearly.
between populations of cells, and the associated barycentric map provides a predictive mapping from source states to target states. From a machine-learning perspective, this map acts as a structured model of the intervention response, and regularization plays the role an inductive bias that controls model complexity.
In the following, we show how Schatten OT can reduce the effective rank of couplings and barycentric projection maps. We use four perturbations from the 4i dataset within the CellOT data of Bunne et al.) (2023). In particular, perturbation data. The processed dataset is publicly available More details on our algorithmic setup for this experiment are given in Appendix D
follow Chen et al. (2025) and consider learning regularized couplings from the 4i
the experiment Chen et al. (0225). the authors fit a displacement-sparse neural OT to 4i perturbation data. In the experiment, we see that dimensionality is reduced, but the error is higher, and the method has high variance because it requires fitting an input convex neural network (ICNN). In contrast, in our work, focus the recovery of low-rank transport structures, which Chen al. (2025 do not target. While our experiments here focus on the compression-fidellly tradeoff within the OT objective itself, they motivate Schatten OT a practical tool for interpretable modeling in high-dimensional biological ML pipelines.
In Figures 4 and 5 we plot the effective rank against À for two different affine maps A(P) = P and A(P) -
We display the result for four different perturbations. For each, we average over five random subsamples of size 500 from the control and perturbation distributions. Since we aim to recover low- rank coupling and low-rank barycentric projection, we compare against Sinkhorn and Low-Rank Sinkhorn Factorization (Cuturi 2013 Scetbon et al., 2021). We do not compare against Subspace Elastic OT (Klein) et al. 2024) since it is not designed to recover low-rank coupling figure, we display both effective rank versus regularization and cost ratio versus regularization. The cost ratio for an estimated P (C, C)/(C, Psink,1), whore Psink, is the solution obtained from Sinkhorn with regularization parameter 1. We compare against this rather than the OT cost for stability.
172
low-rank barycentric map. In each
In Figure f we display the result for low-rank coupling recovery. As we can see, we can drastically reduce the complexity of transport plans without paying much more in transport costs. On the other hand, while Sinkhorn and Low-Rank Sinkhorn reduce the rank, they more dramatically increase the transport cost. Next, in Figure 5 we repeat the previous experiment, focusing on recovering a low-rank barycentric projection map. We see again that Schatten-1 regularization can reduce the effective rank, though the transport cost increases more substantially. However, transport cost increases less substantially for Schatten OT than for Sinkhorn and Low-Rank Sinkhorn.
http:://www.reaarcccoollltttttee
14

---

Published in Transactions on Machine Learning Research (02/2026)
ixazomib
sorafenib
Elsective rank:
120
100
Sckation pa1
Levollask Stekhor (mm2))
Snittmn
3,0
100
Scleciio p=]
Lane-lane Siikton (1=20))
Siabhon
3.5-
ratio
2.5
60
60
2.0
Iffective
-0-
God
2.5-
20
1.5
10
10
10
10
10
Regulariaation
Regularizaiinn
Regulariztion.
Begulacization.
melphalan
cisplatin
100
Siithor
Scaullen p-1
100
Schutioa p=1
Stakhor
Eloective rink
80
Lon-ttank Stekbare (m-02)
00
Lane-aant Strchom (1-10))
rello
60
Iffective
60
Cost
40
40-
20
20
0
Regularization
Regularisation
Regularization
Regularization.
Figure 4: Plots of the performance of regularized OT methods for low-rank coupling recovery across four perturbations from the 4i perturbation data of Bunne et al. (2023). We compare Sinkhorn, Schatten OT with p= 1, and Low-Rank Sinkhorn Factorization. For each perturbation, we plot the effective rank of each method's coupling versus the regularization paramcter, as well as the cost ratio compared to Sinkhorn with parameter 1. As we can see, Schalten OT can reduce the effective rank without increasing the transport cost too much.
ixazomib
sorafenib
2.50
Scturtics
Sckalten
4.0
Eneecive rank
Lote-oone Strchern (m100)
Stekboon
Effecttve rank
2.25
-+-
Stakera
leaeltenk Stakdwm (m-a
2.00
Cost natin
2.00
Cost ratlo
3.0
1.75-
1.75
2.0
1.50
2.0
1.25
1.25
1.5
10
10.
T01
10a
10
10
10a
1.0
10
10
101
Regularizatton
Kegalarizzatinn
Regularizaion
Regulorization
melphalan
cisplatin
Schetter 1-1
2,50
2.25
Stadhom
Etheche runk
2.00
Cost rallo
2.00
Low Sieldom (r-30)
Cost rato
1.75
Eifecttes
3.50
1.50
1.00
1.00
10
10
10a
104
10
10a
10
10
10
10a
10a
10
10
10r
10
10
Regularizatinn A
Regolarizaton X
Hegnsaarssstion A
Hegalarazstion A
Figure 5: Plots of the performance of regularized OT methods in low-rank barycentric map recovery on the 4i perturbation data of Bunne et al. 22023). The setup is the same as Figure | except now we seek to recover a low-rank barycentric map rather than a low-rank coupling. From these figures, we see that Schatten OT can reduce the effective rank of this map, though the transport cost now increases more. On the other hand, while Sinkhorn also reduces the rank of this map, it incurs more drastic increase in the relative transport cost. Low-Rank Sinkhorn is not stable in this experiment.
6 Conclusion
We introduced Schatten-p regularized OT (Schatten OT), a unified convex framework for incorporating low-dinensional structure into OT problems. A key advantage of our formulation lies in its convexity and
15

---

Published in Transactions on Machine Learning Research (02/2026)
generality. Convexity allows us, for the first time, to provide provable recovery results in illustrative yet straightforward examples. Generality allows us to penalize any affine function of the coupling, thereby simultaneously encompassing many existing OT regularizations and enabling new ones.
Theoretically, we established the first recovery guarantees for low-rank couplings and low-rank barycentric displacements, bridging ideas from compressed sensing and OT theory. Algorithmically, we developed an efficient mirror descent method to solve these regularized problems in practice. Empirically, this approach performs well and demonstrates practical utility on 4i cell-perturbation data. Our results show that Schatten OT recovers low-rank structure with only modest increases in transport cost, yielding simpler and more interpretable transport maps.
We believe this work paves the way for more interpretable and scalable OT methods. In particular, the Schatten OT framework may provide a foundation for connecting OT to broader advances in sparse modeling, compressod sensing, and interpretable machine lcarning.
Limitations
While Schatten OT provides a convex mechanism for promoting low-rank transport structure and can be solved efficiently with first-order methods, several limitations remain. First, the recovery guarantecs in Section [ are proved in a stylized setting and should viewed as qualitative guidance rather than a full characterization of when low-rank structure will emerge in practice. Second, although our mirror descent solver avoids explicit semidefinite formulations, each iteration requires computing nuclear-norm subgra- dient, which entails SVD of A(P) and can become computational bottleneck at large scale. In large problems, Iruncated randomized SVD variants are likely necessary. Third, we do not perform exhaus- tive empirical study, and we leave a systematic empirical evaluation across regularizations and formulations to future work. Finally, while OT is a general-purposo tool for machine lcarning practice, we do not conduct an in-depth study of the advantages of this new methodology for specific machine learning lasks. Connec- tions to real machine learning applications, where one can consider learned or task-specific OT models, may further clarify the regimes where Schatten OT offers the greatest practical benefit.
Acknowledgments
T. Maumu was supported by the NSF under Award No. 2305315.
References
Jason Altschuler, Jonathan Niles- Niles-Weed, and Philippe Rigollet. Near-linear time approximation algorithms for optimal transport via Sinkhorn itcration. Advances in Neural Information Processing Systems, 30, 2017.
David Allassz-Melis, Stefanic Jegelka, and Tommi S Jaakkola. Towards optimal transport with global invariances. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 1870- 1879. PMLR, 2019.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International Conference on Machine Learning, pages 214-223. PMLR, 2017.
Amir Bock and Mare Teboulle. Mirror descent and nonlinaa projected subgradient methods for convex optimization. Operations Research Letters, 31(3):167 175, 2003.
Mathicu Blondel, Vivien Seguy, and Antoine Rolet. Smooth and sparse optimal transport. In International Conference on Artificial Intelligence and Statistics, pages 880-889. PMLR, 2018.
Nicolas Bonneel and Julie Digne. A survey of optimal transport for computer graphics and computer vision. In Computer Grophics Forum, volume 42, pages 439-460. Wiley Online Library, 2023.
Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends@ in Machine Learning, 8(3-4):231-38) 357,, 2015.
16

---

Published in Transactions on Machine Learning Research (02/2026)
Charlotte Bunne, Stefan G Stark, Gabriele Gut, Jacobo Sarabia Del Castillo, Mitch Levesque, Kjong- Van Lehmann, Lucas Pelkmans, Andreas Krause, and Gunnar Rätsch. Learning single-cell perturbation responses using neural optimal transport. Nature Methods, 20(11):1759 1768, 2023.
Jian-Feng Cai, Emmanuel Candès, and Zuowei Shen. singular value thresholding algorithm for matrix completion. SIAM Journal on Optimization, 20(4)11966-1982, 2010.
Emmanucl J Candès and Terence Tao. The power of convex relaxation: Ncar-optimal matrix completion. IEEE Transuctions on Information Theory, 56(5):2053 2080, 2010.
Emmanuel J Candès, Justin K Romberg, and Terence Tao. Stable signal recovery from incomplete and inaccurate measurements. Communications Pure and Applied Mathemalics: A Journal Issued by the Courant Institute of Mathematical Sciences, (9(8):2207-122) 2006.
Emmanuel Candès, Xiaodong Li, Yi Ma, and John Wright. Robust principal component analysis? Journal of the ACM (JACM), 58(3):1-37, 2011.
Antonin Chambolle and Thomas Pock. A first-order primal-dual algorithm for convex problems with appli- cations to imaging. Journal of Mathematical Imaging and Vision, 40(1):120 145, 2011.
Peter Chen, Yue Xie, and Qingpeng Zhang. Displacement-sparse neural optimal transport, arXiv preprint arXiv:2502001889, 2025.
Sinho Chewi, Jonathan Niles-Weet and Philippe Rigollet. Statislical oplimal transport. École d'Éié de Probabilités de Saint-Flour. Springer, 2025.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 39(9):1853-1865 1865, 2016.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Information Processing Systems, 26, 2013.
Advances in Neural
Marco Cuturi, Michal Klein, and Pierre Ablin. Monge, Bregman and Occam: interpretable optimal transport bigh-dimenssios with feature-sparse maps. Proceedings of the 40th International Conference on Machine Learning, pages 3671-6689) 2023.
Damek Davis, Dmitriy Drusvyatskiy, Kellie J MacPhee, and Courtney Paquette. Subgradient methods for sharp weakly convex functions. Journal of Optimization Theory and Applications, (79(3))962-8)) 2018.
Steven Diamond and Stephen Boyd. CVXPY: A Python-embedded modeling language for convex optimiza- tion. Journal of Machine Learning Research, 17(83):1 5, 2016.
David L Donoho. Compressed sensing. IEEE Transactions on Information Theory, 52(4))1289 1306, 2006.
Yonina C Eldar and Gitta Kutyniok. Compressed sensing: theory and applications. Cambridge University Press, 2012.
Jicong Fan, Lijun Ding, Yudong Chen, and Madeleine Udell. Factor group-sparse regularization for efficient low-rank matrix recovery. Advances in Neural Information Processing Systems, 32, 2019.
Maryam Fazel, Emmanuel Candès, Ben Recht, and Pablo Parrilo. Compressed sensing and robust recovery of low rank matrices. In 2008 /2nd Asilomar Conference on Signals, Systems and Computers, pages 1043 1047. IEEE, 2008
Sira Ferradans, Nicolas Papadakis, Gabriel Peyré, and Jean-François Aujol. Regularized discrete optimal transport. SIAM Journal on Imaging Sciences, 7(3):1833-1882, 2014.
Aden Forrow, Jan-Christian Hütter, Mor Nitzan, Philippe Rigollet, Geoffrey Schicbinger, and Jonathan Weed. Statistical optimal transport via factored couplings. In The 22nd International Conference on Arlificial Intelligence and Statistics, pages
PMLR, 2019.
17

---

Published in Transactions on Machine Learning Research (02/2026)
Matan Gavish and David L Donoho. Optimal shrinkage of singular values. IEEE Transactions on Information Theory, 63(4):2137-2135 2017.
Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complexity of Sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 1574-1583. PMLR, 2019.
Alberto González-Sanz and Marcel Nutz. Sparsity of quadratically regularized optimal transport: Scalar case. arXiv preprint arXiv:2410. 033553 2024.
Steven Haker, Lei Zhu, Allen Tannenbaum, and Sigurd Angenent. Optimal mass transport for registration and warping. International Journal of Compuler Vision, 60(3):225 240, 2004.
Peter Halmos, Xinhao Liu, Julian Gold, and Bonjamin Raphael. Low-rank optimal transport through factor relaxation with latent coupling. Advances in Neural Information Processing Systems, 37:114374-11444 2024.
Peter Halmos, Julian Gold, Xinhao Liu, and Benjamin J Raphael. Hierarchical refinement: Optimal transport to infinity and beyond. arXiv preprint (53::2503000022 2025.
Robert Huben, Hoagy Cunningham, Logan Riggs Smith, Aidan Ewart, and Lee Sharkey. Sparse autoencoders find highly interpretable features language models. In The Twelfth International Conference on Learning Representations, 2024. URL https: httpss//opeerreven htt.ppparwrriee..ooooooooreeeaa. aet//oorummddeFF/66WWiLee
Kun Jin, Chaoyue Liu, and Cathy Xia. Two-sided Wasserstein Procrustes analysis. In IJCAI, pages 3515- 3521, 2021.
Mete Kemertas, Allan Douglas Jepson, and Amir-massout Farahmand. Efficient and accurate optimal transport with mirror descent and conjugate gradients. Transactions on Machine Learning Research, 2025.
Abdelwahed Khamis, Russell Tsuchida, Mohamed Tarek, Vivien Rolland, and Lars Petersson. Scalable optimal transport methods in machine learning: A contemporary survey. IEEE Transactions Pattern Analysis and Machine Intelligence, 2024.
Dominík Klein, Giovanni Palla, Marius Lange, Michal Klein, Zoe Piran, Manuel Gander, Laetitia Meng- Papaxanthos, Michael Sterr, Lama Saber, Changying Jing, et al. Mapping cells through time and space with moscot. Nature, 638(8052):1065--0777 2025.
Michal Klein, Aram-Alexandro Pooladian, Pierre Ablin, Eugène Ndiaye, Jonathan Niles- Niles-Weed, and Marco Cuturi. Learning elastic costs to shape Monge displacements. Advances in Neural Information Processing Systems, 37:108542 -0006555-1-11655,11111111 2024.
Isaac Liao and Albert Gu. Arc-agi without pretraining, 2025. URL https://11ia023455 1ww.:o2234..gghuhb io/blog. Dosts/arc_age wwwhotttttteeeeeeeeeeee prtptttt_eeeetaain/rrraa aww_aggitttiaaaaae www.thottpeetrriiingg...ooo
Chi-Heng Lin, Mehdi Azabou, and Eva L Dyer. Making transport more robust and interpretable by moving data through small number of anchor points. Proceedings of Machine Learning Research, 139, 2021.
Dirk A Lorenz, Paul Manns, and Christian Meyer. Quadratically regularized optimal transport. Applied Mathematics & Optimization, 83(3):1919 1949, 2021.
Canyi Lu, Jinhui Tang, Shuicheng Yan, and Zhouchen Lin. Nonconvex nonsmooth low rank minimization via iteratively reweighted muclear norm. IEEE Transactions on Image Processing, 25(2):820-88) 2015.
Tyler Maunu, Teng Zhang, and Gilad Lerman. A well-tempered landscape for non-convex robust subspace recovery. Journal of Machine Learning Research, 20(37): 2019.
Arkadij Semenović Nemirovsky and David Borisovich Yudin. Problem complexity and method efficiency in optimization. 1983.
18

---

Published in Transactions on Machine Learning Research (02/2026)
Yurii Nesterov and Arkadi Nemirovski. On first-order algorithms for 11/mullear norm minimization. Acta Numerica, 22:509-577 575, 2013.
Feiping Nie, Heng Huang, and Chris Ding. Low-rank matrix recovery via efficient Schatten p-norm mini- mization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 26, pages 655-656-66 6616, 2012.
Jonathan Niles- Weed and Philippe Rigollet. Estimation of Wasserstein distances in the spiked transport model. Bernoulli, (5(4):66632268)) 2022.
Frangois-Pierre Paty and Marco Cuturi. Subspace robust. Wasserstein distances. In International Conference on Machine Learning, pages 5072 5081. PMLR, 2019.
Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data science. Foundations and Trend&0 in Machine Learning, 11(5-6):355 607, 2019.
Benjamin Recht, Maryam Fazel, and Pablo A Parrilo. Guaranteed minimum-rank solutions of linear matrix equations via muclear norm minimization. SIAM Review, 52(3))471-50)) 2010.
Olivier Roy and Martin Vetterli. The effective rank: A mcasure of effective dimensionality. In 2007 15th European Signal Processing Conference, pages 606-66666000 610. IEEE, 2007.
Christopher Scarvelis and Justin M Solomon. Nuclear norm regularization for deep learning. Advances in Newral Information Processing Systems, 101116222116660 2024.
Meyer Scetbon, Marco Cuturi, and Gabriel Peyré. Low-rank Sinkhorn factorization. In International Con- ference Machine Learning, pages
PMLR, 2021.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Berube, et al. Optimal-transpor analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928 943, 2019.
Erwin Schrödinger. Sur théorie relativiste de l'électron et l'interprétation de mécanique quantique. Tn Annales de l'institut Henri Poincaré, volume 2, pages 269 310, 1932.
Othmane Sebbouh, Marco Cuturi, and Gabriel Peyré. Structured transforms across spaces with cost- regularized optimal transport. In International Conference on Artificial Intelligence and Statistics, pages 586-594. PMLR, 2024.
Richard Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. The American Mathematical Monthly, 74(4):402 405, 1967.
Justin Solomon, Fernando De Goes, Gabriel Peyré, Marco Cuturi, Adrian Butscher, Andy Nguyen, Tao Du, and Leonidas Guibas. Convolutional Wasserstein distances: Efficient optimal transportation on geometric domains. ACM Transactions on Graphics (ToG), 34(4):1-11 11, 2015.
Nathan Srebro, Jason Rennie, and Tommi Jaakkola. Maximum masimuu matrix factorization. Advances in Newral Information Processing Systems, 17, 2004.
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2008.
John Wright and Yi Ma. High-dimennsooal dala analysis with low-diimensional models: Principles, compu- tation, and applications. Cambridge University Press, 2022.
Yuan Xie, Shuhang Gu, Yan Liu, Wangmeng Zuo, Wensheng Zhang, and Lei Zhang. Weighted Schatten p-norm minimization for image denoising and background subtraction. IEEE Transactions on Image Processing, 25(10):4842-48 4857, 2016.
Yaodong Yu, Kwan Ho Ryan Chan, Chong You, Chaobing Song, and Yi Ma. Learning diverse and discrimi- native representations via the principle of maximal coding rate reduction. Advances in Neural Information Processing Systems, 33:9422 9434, 2020.
19

---

Published in Transactions on Machine Learning Research (02/2026)
Yaodong Yu, Sam Buchanan, Druv Pai, Tianzhe Chu, Ziyang Wu, Shengbang Tong, Benjamin Haeffele, and Yi Ma. White-box transformers via sparse rate reduction. Advances in Newral Information Processing Systems, 36:9422 9457, 2023.
Xisoming Yuan and Junfeng Yang. Sparse and low-rank matrix decomposition via alternating direction method. Pacific Journal of Oplimizalion, 9:167, 2013.
Alp Yurtsever, Joel A Tropp, Olivier Fercoq, Madeleine Udell, and Volkan Cevher. Scalable semidefinite programming. SIAM Journal on Mathematics of Data Science, 3(1):171 200, 2021.
Yu Zhang, Ying Wei, and Qiang Yang. Learning to multitask. Advances in Neural Information Processing Systems, 31, 2018.
A
Extension to the Continuous Setting
Up until now, we have focused our attention on formulations in the discrete case. However, there is a direct extension of Schatten OT the continuous setting by taking Schatten-p norms of appropriate linear operators over general Hilbert spaces. In this section, we let p. E P2(RR) be general measures. The set of couplings between these measures is l(,, v).
To define our exlension to the continuous case, we let A P2 x IRa) + B(H) be a map from the space of couplings to the set of bounded linear operators on some Hilbert space H. Let || lls,, now denote the Schatten-p norm over B(H), which is defined as ||T|I, = Ty(()) Then, we define the continuous Schatten OT problem
Sch-OT D: (A.p.q.A)) = min E Ex X Y11 A1I
xell(p.r)
We note that, as in the discrete case, this notion depends on choices of X, p. q, and A. As before, choosing P, and A an affine map makes the problem convex.
Below, we give some examples of affine maps that extend our discrete examples. Throughout, we let 8 be the reference product measure.
Covariance regularization: The most direct connection between the continuous and discrete cases is to penalize moments of the distribution. In the continuous case, this corresponds to regularizing the covariance of т. In this case, all of the regularizations discussed in Section 2.3 are the same except we regularize the
lincar operator over Rdd Σn Ex (x)
Continuous sparse and low-rank regularization: We now discuss the analogs of quadratic and low- rank regularization, the coupling matrix P. Here, we can take A(T) S. (²(v) + I²(()) as the lincar operator
(S+f) = ff f(v) du duly) E+J(() xx]
dp
Note that Sy is affine in æ. Then, regularizing the Schatten-p norm of S, corresponds to using the Schatten-p of P as discussed earlier. The continuous quadratic case has already been studied by Lorenz et al. though not framed Schatten-2 norm of S,.
2021
We can recover the elastic costs by taking the Schatten norm of specifically constructed operator. Choose measurable partition of Rd x Rd given by (Ex)een with p(Ex) 0. Define an orthonormal
Elastic costs:
family in T.2 by
(x(x,)) 1((x, y) E Ex)
Vp(Ex)
20

---

Published in Transactions on Machine Learning Research (02/2026)
where Rd x F {0, 1 is the indicator function of set. Set the diagonal weights to be (T(TT) E. l/y - inI, which are linear in a. Then, letting (ex)keN be basis for f2, can define the linear operator A(T) e + (p) by
A(T)eL= 81 (T)
Notice that this is again linear in and furthermore A(T)eu are orthogonal with [[A(T)eull = Sx(T). There- fore, the singular values of A(T) are (T().. Thus
|A(T) = E [S((TT= = / 11y x idT.
A similar construction yields the subspace elastic costs discussed in Section 2.3 that clastic OT problems can be embodded as Schatten OT regularized problems over appropriate operators.
The general principle here is
Barycentric projection maps and displacements: We also consider Schatten-p regularization of the barycentric projection maps and displacements. In the continuous case, the barycentric projection map Tx : Rd + Rd is (x) Ex = z). Let the displacement map be defined as Dx(x)) (x)) z. Then, we can formulate a Schatten-p penalization of this barycentric displacement map by viewing T. or D, as operators T, Dz : Rd + (p) given by T = (v, T-)) and D₃U = (u, (x)))x())
Supplementary Proofs
B.1 Proof of Theorem 4
Proof. Define, for each I, the cluster indicator mass vectors
i St, otherwise,
(t) E R (a!
Rg 0,
and similarly
1 Rg 0,
j E Tr, otherwise.
g(t) E R", (g(t)))
The rank R coupling we wish to recover is
n p= = Σ a (0 3()) t=1
(10)
Notice that P E U(a. b) is block-diagonal with blocks (S, x Tj) that are uniform (each entry equals 1/(Rg)²), and can explicitly compute P*1|s = F- It will be convenient to define the normalized indicator vectors (+) ("/ (tt) = (+) (000)/l1100 [1] 112, which we stack into matrices U = [u(1). ,..., (R)),
V (x(1) ....., (R)). In this way, the canonical subgradient of ||. lls, P G' UV
Step 1: Lower bound on Ami.. For any E ((c,, p), Vi E B(d,, p), Yout E B(d,, p) with tt,
o youull 2 l|e: d,ll -| - c:ll |out dell > - 2p,
and
|| yin |l $ ||e, dell x cill
11
dell 5 2p.
Thus
x - yout || 1| yum 2 > (T - 2p) - (7 + 2p)2
which is positive when
holds.
Step 2: Across-block exclusion via tilted cost. By convexity of || I|s..
|||s, |P*|s, + (G", P P*), G @|lP* G = UV...
21

---

Published in Transactions on Machine Learning Research (02/2026)
Hence, for any feasible P.
() Plll,s - ((C, Xl (l),, (S)A
For i St and T, with one has Gij O, since the left and right singular vectors block-supported and orthonormal across clusters. On the other hand, for j' Tr,
by 1
ei) Gij = (e) ey)) |o00| ||330|12
g
Therefore, for any i St, 8 + t, j E Ts, and j' E Te:
A g
S1,(A - Siy (2, G*) = (1 - vill2 21 yy 15
9 2
A
If À < gAmin; these gaps are strictly positive, Any minimizer of the original problem must then be block-supported on U,(S, x T.).
S(A, G* -optimall coupling can place mass across clusters.
Step 3: Within-block tie-breaking via the nuclear norm. By the distance equality condition Assumption 3 all within-cluster couplings have equal transport cost. Fix t. We can restrict any feasible to the block (S,, T:), which denote as P.... We note that this can be written
coupling P P b)
PS.T. 1 /9² (M),, where M = 0, (M (M
In other words, we can represent it as rank-1 product coupling plus a perturbation with row wwwcollrmn sums. Choose an orthonormal basis R on the target side with the first vector proportional to 1g. Then M lives entirely in the orthogonal complement of g. The standard inequality I|s 2 || lls, yields
|| P 2 PSs.mlls: / 1 + IM > |1, /
whenever ( O. Summing over shows that among all block-supported couplings, the nuclear norm is uniquely minimized at (t 0, ie., at the uniform block P
Combining these three steps proves the theorem.
B.2 Proof of Theorem 6
Proof. We proceed in four steps. We will show that the coupling recover, P E U(a, b), satisfies the condition (((++) = (i = 20 if i St otherwise P Pi(s.)) = 0 if to S..
Step 1: Feasibility and rank-1 structure. By construction, P El(a, b) and, for cach i E St,
Tp. (x:) 2 ( y1,-)) =4.
m.
Hence Tp- (x:) x; m. x, = - when m. + Eru for leal S p, which is true by assumption. Writing E Rm such that 7 = a,, we have
A(P") = u.
Therefore, rankA(P*) 1, and ||ss = 1yllz
Step 2: A tilted-cost lower bound and across-cluster margin. Let G be a canonical subgradient of the nuclear norm at B* := A(P)::
G¹ 0 8 G uw where Y
|
For any Pe U(aab)) by convexity the nuclear norm,
||A(P)]ss, B'
|B ll, (G', B*).
(11)
22

---

Published in Transactions on Machine Learning Research (02/2026)
Using A(P) B" - P
A-1/2 and cyclicity of the trace,
A 1
(12)
Combining 11 and 12 with the objective FX(P) == (C, + (|))( and the definition of the tilted cost (), G in yields the lower bound
- - 2 G*), P
(13)
We can compute the tilted costs S(A, G*) explicitly: for any and (t,
Yu(to))) (A- TY
_mm mi y1,o) = Tipe
1
Eipe
Vai
1
1h
Therefore, assuming that i E St and for any 8 8 o E )
Ne 112 +1 (-Iri12
&i
) (us me)
(A, G*) - (A, G*) = =||x; ; Vs 2
(14)
1
=Ai,.
By assumption,
D 2 1 - 1x101 - 1all-22 -2p) 2 A(A - 22p) > 0.
(15)
Also, since (l < P, can bound 1/ = E k= a1.52 S p²² Σ
k== 0# p².
This implies that
S 1.
Txlla
Thus, we can extend the lower bound in
(e,,,)) (A, G*) (t,+)) (A, 2 A((A - 2p) 2pp) A)).
(16)
Thus, at G for every X [0, - 2p), across-cluster tilted costs are strictly greater than within-cluster tilted costs. Therefore any tilted cost optimal coupling must match æ; to { for i E S..
Step 3: Within-cluster degeneracy and the muclear-norm tie-break. Fix & E [R]. For E St, any within-cluster move between the symmetric targets (t, +) and (t, -) has zero cost difference,
x Vr 1 = || yt y 113.
Moreover, the tilted cost the same for (t,+) and (t, -), since (u. M1,+)) (u, Hr, pr. Therefore, for any feasible P that sends mass within clusters (i.c., supp PP) S ((i, (), +))
iE S.)),
(S(A, G"), P p') 0.
(17)
For such P, we can write the within-cluster mass split by ps E [0, that
Pa.(t.)) paas, Po(t.-) = (1 p:)a:
A direct computation gives
Tp(x,) = Pillt+ + (1- p.)... = m. + (2p; - 1)80, E 2 0.
Hence, with
B E R", B = :=p= - 1)evai,
we obtain the rank-5 decomposition
A(P) uyy T + 43
(18)
We claim that, for any B +
| 1 ||s, 1
luy ||s, ||
(19)
23

---

Published in Transactions on Machine Learning Research (02/2026)
Indeed, let Q Rddd be orthogonal matrix whose first two columns are and v. Orthogonal invariance of singular values implies
luy + vB 1B Ils: = [r 8 0
0] lls: - - + 02
where 01 2 02 2 O. Tf is not colinear with 7, the matrix has rank 2, so oz > 0, and 01 1|lll2 (since
|| 2 the Euclidean norm of any row). Hence 01 + 02 1rl2. instead B ==oy for some c c 0, then the
matrix has rank with singular value vlhvl22 | 2 v + + > lhll. Thus 19 holds in all cases B + O.
Combining
13 and 17 for any within-chuster feasible P,
F Fx (P*) 2 A(IIA(P)]ss - |ls,, = X(1luy' + v3 Ils, - 11
(20)
which is strictly positive by
19 whenever B + 0, i.c., whenever some p; + 2"
Step 4: Optimality and uniqueness for X E [0, 2p). Let A E -pp).. For any feasible P. decompose P P-P" into an across-cluster part and a within-cluster part. By 16
(S(A,G)) P P') >0
if P sends any mass across clusters, and 13 implics Fx(P) (P*) in that case. Therefore, any minimizer of Fx must. be supported within elusters. For within-cluster couplings, (20) implies (P(P) Fx(P*) unless
Pr 1/2 for all ie., if P P Consequently, P is the unique minimizer of the Schatten OT problem for
À E [0, - 2p)).
C The Gaussian Case
The previous section treated recovery of low-rank structures in discrete OT. We now discuss application
of the continuuus Schatten regularization (9) for Gaussians.
We treat two Gaussian specializations of the Schatten-p programs discussed carlier: (i) a muclear-oorm penally that promotes low-rank cross-covariance, and (ii) a nuclear-norm penalty that promoles low-rank transport. As emphasized in our general framework, the barycentric projection x + E, [Y X x) is an affine map of the coupling n, so the induced Schatten-p penalty is convex in æ for p 2 1. The same holds for Schatten penalties applied to any affine image A(T).
For simplicity, we consider the mean zero case. Let μ = N(0, Eo) and == N
E,) En) on Rd with E. E₁ > 0.
Σn
A Gaussian coupling is a joint Gaussian ा = .
0 1 KT Σ ), parameterized by a cross-covariance
K E gdxd satisfying the feasibility constraint
Σn K
K Σ
E 1/2 KΣ 1/2 < 1.
(21)
We denote the set of all such K as K(Eo, Σ Equivalently, we denote the set of all Gaussian couplings between i and as (H, v).
For the quadratic cost c(x, y) = || - y the transport cost under æ is E Y = tr(Eo) + +tr(Etr +trEE)) - 2tr(K). Moreover the barycentric map induced by π is
(x) A,x, EE E
(22)
C.1 Gaussian Low-rank Croas-Cov
Consider the Gaussian Schatten OT problem
min tr(Ep) + tr(E)) ) 2tr(K) + -|| K E,)
(23)
24

---

Published in Transactions on Machine Learning Research (02/2026)
This is semidefinite program.
We can solve this problem in closed form. Let = 2//2 and write its SVD S Ud.. G.) with 2 ... 2 øa > 0. Feasible K can be written as K = E "ME with [M 12 S 1. By von Neumann's 's trace incquality, tr(K) = tr(SM) I E; 0;8, where 8; are the singular values of M, with equality when M shares singular vectors with S. At optimum we may thus take M = ('d.ag(s and || |s,, = MM Σ, 8;- The objective in (3 reduces (up to constants) to
with 0 I 8, S 1
Σ (1-20))) 0S
min
Σ
min (A -204)s US$51
which is separable and linear in each 8,. Hence, we arrive at the following proposition.
Proposition 7 (Hard spectral selection). The unique minimizer of 23 is obtained by hard thresholding the singular spectrum of S:
8 = 1(0, > a/2),
K, = =// 2 Udiag(s) (s) ///
(24)
In porticular, rank(Kx)) # 0, X/2].
We include examples to illustrate this hard-thresheee thresholding rule.
(i) Isotropic case: E. -2 1 gives S = abl and hence either K₂ abI if < 2ab, or K₁ 0 X 2ob.
(ii) Commuulng covariances: If Σg Udiag(a)UT and Σ, Udiag(b)UT, then o; Vahbi and K₁ ddiag b bllva X/2])UT. Here, is the Hadamard (elementwise) product.
We note that the inclusion of a Schatten-2 penalty in this program results in soft thresholding.
C.2 Gaussian Barycentric Displacements
We now penalize the (weighted) barycentric displacement. In this case, the resulting convex program
min tr(Eo) + tr(E,) 2tr(K) + A| - / (E))(am))
(25)
with A as in (22
We can give closed form when En and Ej commute, which already reveals clear structure. Suppose there exists orthogonal U with Eg = Udiag(a)UT and ΣS Udiag(b)UT Any feasible K aligned with U takes the form K = Udiag(( a c b c m with 0 < my < 1. Then A, = K xg' is diagonal in the same basis with entrics
b Vasmitt mma
04 =
which implics that the diagonal values of (A, - 1)=// are mivb; Vas in this basis as well. Hence [I(A, /)2 ls. Emlmab Varl, and
tr(K) = E Varm
Consequently, 25 decouples into d scalar problems over m, [0. 1]:
min (a. = abm + Ammbb Vol. 0sm<1
(26)
25

---

Published in Transactions on Machine Learning Research (02/2026)
Theorem 8. Fiz (a. b, X) with a, b > 0. The unique minimizer of 26 is
if - a, 1,
m 1,
if > a and 2/a,
Va/b, b > a and z 2/a.
Equivalently, the coupling that solves
25 has barycentric projection map T+(-) Axx where
Vb./ai
if b, S az. if > and X < 2
Ax = "'dia(aa))U,,
a;
=
(27)
1.
if and X 2 2
The proof follows from analyzing the 1-dimensional optimization problem 26
In words, the regularizer does not suppress contracting directions (6: S az), and it prunes expanding direc- tions b, a, back to the identity once À crosses the sharp threshold 2/a,. Thus, rank AA, I) equals the number of contracting eigendirections plus the number of expanding eigendirections with À 2 0z.
As an example, consider the isotropic case. If Eu I Σ = I If 01 S then A = (01/00) for all X. On the other hand, if 01 > do, then A = (a1/an) /0)) I for < 200, and A = for 2 20p-
C.3 Discussion
In the Gaussian case, for 1, the programs remain convex. In the commuting case, we can find separable one-dimensional convex problems similar to those we found in the previous sections. For larger we would observe smooth shrinkage of the spectrum m, rather than hard thresholds at = 1. We leave a detailed analysis of the noncommuting case future work.
Details on 4i Experiment Setup
For reproducibility, we give more details on the experiment here. We subsample source and target points form measures with 500 points cach. We then run mirror descent with the desired regularization, where the initial step size is To 0.1 for low-rank coupling recovery and no 10-4 for low-rank barycentric map is used. We run mirror descent for maximum of 50
recovery. A diminishing step size of 7H no/vk iterations using Sinkhorn projection (with up to 500 Sinkhorn iterations per mirror descent step), unless the marginal error reaches 10-12, This is followed by the rounding procedure of Altschuler et al. (0117) to ensure iterates remain in the transport polytope. We return the averaged iterate as our proposed solution to the Schatten OT problem.
26


Reviews:
[{"review_id": "soLWqyLWht", "reviewer": "Reviewer_PDPV", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The abstract of the paper correctly summarises it. A convex\nformulation for optimal transport is proposed which (like some\nexisting methods) promotes sparse maps. Since the formulation is\nconvex some optimality results are derivable. A mirror-descent\nalgorithm is used since, despite the convex nature of the optimisation\nproblem, it can get too large for off-the-shelf convex optimisation\nalgorithms to solve. Empirical results on both synthetic and real data\nare presented.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Theoretical claims are supported by proofs. Empirical results provide\nevidence for the claimed positive features of the presented method.", "audience_yn": "Yes", "audience_explanation": "OT is used in ML, so those aware of its use in ML may find something\nof interest here. My guess is that this will be a small proportion of\nML researchers, but since ML is so big these days, this small\nproportion could be quite large in absolute terms. Those not already\nfamiliar with OT will not read this paper, I suspect.", "requested_changes": "The paper would be improved if OT could be \"sold\" more to an ML audience. Would it be possible to have an ML example in the Experiments section?\n\nI can't think of any other useful changes without going over the 12 pages. Note that the proofs are in appendices.\n\n\nNote that Wasserstein and Schatten should be capitalised in the bibliography.", "broader_impact_concerns": "None", "additional_comments": "TMLR is a machine learning journal. The submission makes a brief\nargument (by referencing other work) that OT is of use for ML. But\nthis is basically an OT, not an ML paper. I'm not sure why the authors\nhave chosen to submit it to an ML journal.", "extra_scores": {}, "extra_text": {}}, {"review_id": "37zEUkUP9m", "reviewer": "Reviewer_Aqgm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In the manuscript, the authors introduce Schatten-p Regularized Optimal Transport (Schatten OT), a unified convex framework that generalizes many existing regularized OT methods including low-rank, sparse, and quadratic formulations under a single mathematical model. By leveraging ideas from compressed sensing, the authors establish the first theoretical recovery guarantees for low-rank couplings and barycentric displacements in OT. They also develop an efficient mirror-descent algorithm with convergence guarantees for solving the convex formulation at scale. Empirical results on synthetic and biological datasets demonstrate that Schatten OT effectively recovers low-rank, interpretable transport maps with minimal cost increases, thereby bridging OT theory, convex optimization, and compressed sensing in a principled and scalable manner.", "claims_and_evidence_yn": "Yes", "claims_explanation": "Yes. The paper’s claims are well supported by clear theoretical, algorithmic, and empirical evidence. The authors rigorously derive convexity and optimality conditions for the proposed Schatten OT framework and provide formal theorems with proofs establishing low-rank recovery guarantees which is an important novelty in the OT literature. The mirror-descent algorithm is mathematically justified with convergence analysis, and experiments on both synthetic and real-world (4i perturbation) datasets convincingly demonstrate the method’s efficiency, scalability, and ability to recover interpretable low-rank structures. Together, the theoretical results and empirical validations align coherently with the stated contributions and meet TMLR’s standards for clarity, soundness, and reproducibility.", "audience_yn": "Yes", "audience_explanation": "Yes. The paper’s findings would interest TMLR’s audience, particularly researchers in machine learning, optimization, and computational optimal transport, as it introduces a novel convex framework with theoretical guarantees and scalable algorithms for low-rank and interpretable transport.", "requested_changes": "While the paper is very well written and contains many strengths, the following changes could make its findings more robust and broadly impactful.\n\n1. Clarify Practical Impact (Important): I believe the theoretical results from the paper are strong, but the paper would benefit from a clearer discussion of how Schatten OT could be applied in real-world problems like single-cell analysis, imaging, or generative modeling and when the low-rank assumption is most justified.\n2. Expanded Empirical Evaluation (Important): The experiments are convincing but could be strengthened by including comparisons with more recent or diverse OT baselines, such as neural or elastic regularized OT models. This would help demonstrate where Schatten OT provides the most tangible improvements in practice.\n3. Hyperparameter Sensitivity (Minor): Including a short experiment or discussion showing how the regularization strength (lambda) and Schatten parameter (p) affect results would make the method easier to reproduce and tune for new datasets.\n4. Discussion of Limitations and Future Work (Minor): It would be useful for the authors to explicitly mention potential limitations, such as performance in highly nonconvex or noisy settings, and suggest possible future directions.", "broader_impact_concerns": "There are no major ethical or broader concerns associated with this work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "qpSUT7JENo", "reviewer": "Reviewer_aTux", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces \"Schatten OT,\" a novel and general framework for finding low-rank structures in Optimal Transport by using Schatten-p norm regularization. Its primary contribution is providing a single, convex optimization program (for $p, q \\ge 1$) that unifies a wide range of existing methods, such as those for finding low-rank couplings, sparse plans, and elastic costs. A major strength is that this convexity enables, for the first time, provable theoretical guarantees for recovering low-rank solutions in OT, a significant finding that connects OT with compressed sensing principles. The authors also propose a scalable mirror-descent algorithm and provide strong empirical evidence on synthetic and real 4i cell data, showing the method effectively simplifies transport plans with only a modest increase in cost. Key weaknesses include the fact that the novel recovery theorems are currently limited to simplified \"toy\" settings and that the framework, as presented, does not encompass the widely-used entropic regularization.", "claims_and_evidence_yn": "Yes", "claims_explanation": "- **Claim**: The paper proposes a general, convex framework called Schatten OT that unifies many existing regularization methods (like low-rank couplings, quadratic regularization, and elastic costs) and enables new ones.\n- The framework is clearly defined in Equation (3). The authors then explicitly demonstrate how this single equation can recover several existing methods (low-rank couplings, quadratic regularization, group sparsity, elastic costs) simply by changing the affine map $\\mathcal{A}$ and the parameters $p$ and $q$.  They also show how it creates new regularizers for barycentric maps and covariances.\n- **Claim**: The framework's convexity (for $p,q \\ge 1$) allows for direct theoretical analysis and provable recovery guarantees for low-rank solutions.\n- Theorem 4: Proves exact recovery of a low-rank, block-diagonal coupling in a specific clustered setting. Theorem 6: Proves exact recovery of a rank-1 barycentric displacement map in a symmetric two-cluster setting.\n- **Claim**: The proposed mirror descent algorithm efficiently solves the Schatten OT problem and performs well on real data, recovering low-rank structures with only a modest increase in transport cost.\n- Fig. 1 & 2 clearly show that as the regularization strength $\\lambda$ increases, the \"effective rank\" of the solution decreases significantly. Fig. 3 directly supports the algorithmic claims, showing sublinear convergence. Fig. 4 experiment demonstrates practical utility.", "audience_yn": "Yes", "audience_explanation": "The paper directly tackles two of the most significant challenges in computational OT: scalability and interpretability. The paper introduces a unified, convex framework. This convexity allows the authors to provide the first provable recovery guarantees for low-rank solutions in OT, a significant theoretical contribution that would interest the research community. It develops a scalable mirror-descent algorithm to solve the new formulation and demonstrates its effectiveness on real-world 4i cell-perturbation data.", "requested_changes": "- Justification for the choice of mirror descent over other potential convex solvers, such as ADMM or primal-dual methods, would also strengthen the algorithmic contribution by providing clearer context\n- A brief discussion or even a simple experiment on combining the Schatten-p penalty with an entropic regularizer, since (3) cannot recover entropic regularization.", "broader_impact_concerns": "The paper is mostly theoretical, so there are no broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Decision: Accept (poster)