Scalable Approximate Algorithms for Optimal Transport Linear Models

TL;DR

This work tackles regression problems where the fit is evaluated through an entropy-regularized optimal transport loss, enabling a focus on shape perturbations rather than pure per-sample residuals. It introduces OTLM, a general, scalable framework that extends generalized linear models to include transport-based datafits and convex penalties, solved via Sinkhorn-like scaling with majorization-minimization steps. The authors derive practical MM updates for common penalties and demonstrate the method's scalability on large synthetic problems and real spectroscopic datasets, including muonic X-ray spectra and infrared spectroscopy. The approach offers a flexible, robust, and efficient alternative for linear modeling in contexts where mass transportation provides a more meaningful similarity than conventional Euclidean losses.

Abstract

Recently, linear regression models incorporating an optimal transport (OT) loss have been explored for applications such as supervised unmixing of spectra, music transcription, and mass spectrometry. However, these task-specific approaches often do not generalize readily to a broader class of linear models. In this work, we propose a novel algorithmic framework for solving a general class of non-negative linear regression models with an entropy-regularized OT datafit term, based on Sinkhorn-like scaling iterations. Our framework accommodates convex penalty functions on the weights (e.g. squared-$\ell_2$ and $\ell_1$ norms), and admits additional convex loss terms between the transported marginal and target distribution (e.g. squared error or total variation). We derive simple multiplicative updates for common penalty and datafit terms. This method is suitable for large-scale problems due to its simplicity of implementation and straightforward parallelization.

Figures (5)

• Figure 1: TV regression compared to the balanced OTLM, without weight penalties. Left: target data, three basis vectors, and the best fit using TV regression. Middle: balanced OTLM with the same basis vectors. Right: the OTLM transport plan.
• Figure 3: Scaling of the regularized OTLM algorithm compared to unregularized OTLM solved using linear (LP, left) and quadratic (QP, middle) programming solvers. The right panel shows the difference between the weights obtained with $\epsilon$-regularized and unregularized OTLM.
• Figure 4: Spectrum of the BCR-691-A sample in the region of Cu $K_{\alpha}$ lines. Left: fit with non-convex Bayesian mixture model of skew-Gaussian distributions. Right: fit with convex OTLM. Both fits use 4 components.
• Figure 5: Weights for each X-ray line of the BCR-691-A sample, as calculated by the convex OTLM method (red crosses) and non-convex mixture model of skew-Gaussian distributions (see Section \ref{['sec:muonic_xray_spec']} for details). The mixture model weights were obtained using Baysian sampling. The error-bars show the standard deviation of the weights from the Bayesian posterior. The triangles show the OTLM weights that are outside the y-axis range.
• Figure 6: The 54 blank filter spectra (left panel). These are used to engineer a basis for the interference. Ammonium sulfate spectra profile (middle panel) and 49 filter spectra (right panel). The threshold $t$ selected for the cost matrix is shown as the red dashed line.