Optimal Transport

Follow this topic Ask about this topic

Optimal transport studies how probability mass can be moved between distributions at minimum cost, giving a geometric way to compare data even when samples do not correspond one-to-one. Wasserstein distances and transport maps describe this geometry, while Gromov–Wasserstein methods compare relational structure when data live in different spaces. Practical approaches include entropic regularization with Sinkhorn algorithms, sliced and low-rank approximations, and neural transport maps used in generative modeling, domain adaptation, and trajectory inference. A central challenge is balancing computational and statistical scalability with faithful transport, especially when choosing regularization or relaxing exact mass constraints.

Latest papers

289 papers48 papers in the last four weeks, up 17% on the four weeks before. 0.3% of all new papers.

May 7, 2026

Optimal Transport for LLM Reward Modeling from Noisy Preference

Reward models are fundamental to Reinforcement Learning from Human Feedback (RLHF), yet real-world datasets are inevitably corrupted by noisy preference. Conventional training objectives tend to overfit these errors, while existing denoising approaches often rely on homogeneous noise assumptions that fail to capture the complexity of linguistic preferences. To handle these challenges, we propose SelectiveRM, a framework grounded in optimal transport. We first devise a Joint Consistency Discrepancy to align the distribution of model predictions with preference data. Furthermore, to address the limitation of strict mass conservation which compels the model to fit outliers, we incorporate a Mass Relaxation mechanism via partial transport. This enables the autonomous exclusion of samples with noisy preference that contradict semantic consistency. Theoretically, we demonstrate that SelectiveRM optimizes a tighter upper bound on the unobserved clean risk. Extensive experiments validate that our approach significantly outperforms state-of-the-art baselines across diverse benchmarks.

Licheng Pan, Haochen Yang, Haoxuan Li +8

Zhejiang University · Xiaohongshu Inc · National University of Singapore +2

Reward Modeling Noisy-Label Learning Optimal Transport RL from Human Feedback

May 7, 2026

Stability of the Monge Map in Semi-Dual Optimal Transport

This paper shows that the semi-dual formulation of the optimal transport problem has a degenerate saddle-point structure, and that its numerical solution is equivalent to solving a constrained optimization problem. We derive necessary and sufficient conditions for the convergence of Monge maps without requiring optimality of the dual potential. This analysis helps explain why, in practice, numerical algorithms often require more iterations to update the transport map than the potential.

Anton Selitskiy, David Millard

Dept. of Electrical and Computer Engineering Stony Brook University Stony Brook, NY 11790, USA · Dept. of Electrical and Computer Engineering University of Rochester Rochester, NY 14627, USA

Optimal Transport

May 7, 2026

Locality-aware Private Class Identification for Domain Adaptation with Extreme Label Shift

Domain adaptation aims to transfer knowledge from a labeled source domain to an unlabeled target domain with different distributions. In real-world scenarios, the label spaces of the two domains often have an inclusion relationship, where some classes exist only in one domain but not the other. These non-overlapping classes are referred to as private classes. Identifying private class samples and mitigating their adverse effects is critical in the literature. Existing methods rely on the assumption that shifts in private classes are large enough to be considered outliers. However, the variance within a single shared class can be significantly larger than the difference between a private class and another shared class, challenging this assumption. Consequently, private classes substantially increase the difficulty of cross-domain classification. To address these issues, based on local transportation and metric properties of optimal transport (OT), a locality-aware private class identification approach is proposed in the form of a score function on transport mass. The effectiveness of the proposed approach is theoretically proven, highlighting the score function's strong ability to distinguish between shared and private class samples. Building on this, we introduce a reliable OT-based method (ReOT) for domain adaptation under severe label shift. ReOT minimizes classification risk while learning the separated cluster structure between the identified shared classes and private classes, effectively avoiding mismatch between shared-private sample pairs, thus ensuring that important knowledge is reliably transported intra-class to mitigate class-conditional discrepancy. Furthermore, a generalization upper bound of the target risk is provided for extreme label shift scenarios, which can be minimized by ReOT. Extensive experiments on benchmarks validate the effectiveness of ReOT.

Chuan-Xian Ren, Cheng-Jun Guo, Hong Yan

School of Mathematics, Sun Yat-Sen University, Guangzhou 510275, China · Department of Electrical Engineering, City University of Hong Kong, Hong Kong, China

Label Shift Unsupervised Domain Adaptation Domain Adaptation Optimal Transport

May 7, 2026

Optimal Transport Audio Distance with Learned Riemannian Ground Metrics

In audio generation evaluation, Fréchet Audio Distance (FAD) is a 2-Wasserstein distance with structural constraints for both primitives: the cost is a frozen embedding pullback whose invariance set hides severe artifacts, and the coupling is a Gaussian fit that dilutes rank-1 contamination relative to discrete OT. We propose Optimal Transport Audio Distance (OTAD), which corrects each primitive with one dedicated mechanism -- a residual Riemannian ground-metric adapter for the cost and entropic Sinkhorn optimal transport for the coupling. Across eight encoders under a four-axis protocol, coupling-only comparisons at

ε= 0.05

show that Sinkhorn's rank-1 sensitivity exceeds FAD's by a factor of 1.9 to 3.6. Furthermore, OTAD achieves a higher mean Spearman correlation with audio-quality MOS (DCASE 2023 Task 7) than baseline metrics. As an intrinsic benefit of the discrete transport plan, OTAD yields per-sample diagnostics with AUROC

\ge 0.86

, a capability that scalar- or kernel-aggregated metrics structurally lack.

Wonwoo Jeong

Sogang University

Speech Generation Evaluation Metric Learning Riemannian Metric Learning Optimal Transport Generative AI Evaluation

May 6, 2026

Reliable Modeling of Distribution Shifts via Displacement-Reshaped Optimal Transport

Optimal transport (OT) is a central framework for modeling distribution shifts. Because OT compares distributions directly in input space, a well-designed ground metric between observations is essential to ensure that the optimizer does not violate the true geometry of change. We propose Displacement-Reshaped Optimal Transport (ReshapeOT), a method that reshapes the ground metric by integrating observed sample displacements as an additional source of knowledge. Technically, ReshapeOT replaces the Euclidean metric with a Mahalanobis distance estimated from displacement second moments. This effectively carves expressways through the input space, inviting transport solutions that better align with observed displacements. Our method is computationally lightweight, integrates seamlessly into any OT solver that operates on a cost matrix, and can be kernelized for further flexibility. Experiments on synthetic and real-world data show that ReshapeOT achieves substantial gains in transport reliability. We further demonstrate our method's usefulness in two practical use cases.

Philip Naumann, Jacob Kauffmann, Klaus-Robert Müller +1

BIFOLD – Berlin Institute for the Foundations of Learning and Data, 10587 Berlin, Germany · Machine Learning Group, Technische Universität Berlin, 10587 Berlin, Germany · Department of Artificial Intelligence, Korea University, Anam-dong, Seongbuk-gu, Seoul 02841, Korea +2

Distribution Shift Optimal Transport

May 6, 2026

UniVer: A Unified Perspective for Multi-step and Multi-draft Speculative Decoding

Speculative decoding accelerates Large Language Models via draft-then-verify, where verification can be framed as an Optimal Transport (OT) problem. Existing approaches typically handle multi-draft and multi-step aspects in isolation, applying either flat OT to single-step drafts or per-token rejection sampling to tree-structured candidates. This separation leaves the joint regime (where multi-step dependencies meet multi-draft branching) poorly optimized, as local verification rules fail to exploit the coupling between horizontal and vertical dimensions of candidate trees. In this paper, we propose a unified perspective that casts tree-based verification as a conditional OT problem. Our key insight is that vertical dependencies can be abstracted through prefix acceptance probabilities, which act as dynamic scaling factors to actively guide horizontal draft selection. Based on this principle, we introduce UniVer, a verification algorithm that jointly optimizes across tree levels by composing local optimal transport plans under prefix constraints. We prove that UniVer remains lossless and achieves the optimal acceptance rate under the proposed conditional framework. Extensive experiments across different tasks and models demonstrate that UniVer improves acceptance length by 4.2% to 8.5% over standard recursive rejection sampling without replacement, while maintaining exact distributional alignment with the target model.

Yepeng Weng, Qiao Hu, Takehisa Yairi

The University of Tokyo · 2Lenovo AI Technology Center · National Center for Mathematics and Interdisciplinary Sciences (NCMIS), AMSS, CAS

Speculative Decoding LLM Inference Acceleration Optimal Transport

May 5, 2026

Entropic Riemannian Neural Optimal Transport

Many machine learning problems involve data supported on curved spaces such as spheres, rotation groups, hyperbolic spaces, and general Riemannian manifolds, where Euclidean geometry can distort distances, averages, and the resulting optimal transport (OT) problem. Existing manifold OT methods have pursued amortized out-of-sample maps, while entropic regularization has made discrete OT more scalable, but these advantages have remained largely disjoint. We propose Entropic Riemannian Neural Optimal Transport (Entropic RNOT), a unified framework that combines intrinsic entropic OT with amortized out-of-sample evaluation on Riemannian manifolds. Our method learns a single target-side Schrödinger potential through a neural pullback parameterization, recovers the induced Gibbs coupling, and uses the resulting conditional laws to construct intrinsic transport surrogates. These include barycentric projections on Cartan-Hadamard manifolds and heat-smoothed conditional surrogates on stochastically complete manifolds, the latter turning possibly atomic target laws into absolutely continuous ones. For fixed regularization

\varepsilon>0

, we prove that the proposed hypothesis class recovers the entropic optimal coupling in strong probabilistic metrics. As consequences, barycentric surrogates converge in

L^2

, while heat-smoothed surrogates are stable at fixed heat time and asymptotically unbiased as the heat time vanishes. The guarantees hold for compactly supported data on possibly noncompact manifolds. Empirically, our method matches or improves over Euclidean, tangent-space, and log-Euclidean baselines on benchmarks over

\mathbb{S}^2

,

\mathrm{SO}(3)

,

\mathrm{SPD}(3)

,

\mathrm{SE}(3)

, and

\mathbb{H}^2

, scales favorably relative to discrete manifold Sinkhorn, and in a protein-ligand docking application, refines poses on

\mathrm{SE}(3)

without retraining or per-instance optimization.

Alessandro Micheli, Silvia Sapora, Anthea Monod +1

Imperial College London London, UK · University of Oxford Oxford, UK · Statens Serum Institut Copenhagen, Denmark +1

Riemannian Geometry Optimal Transport

May 5, 2026

Globally Solving Unbalanced Optimal Transport and Density Control for Gaussian Distributions

In this article, we study unbalanced optimal transport (UOT) and establish a control-theoretic dynamical extension, which we call the unbalanced density control (UDC), for a class of Gaussian reference measures. In the static setting, we consider UOT with quadratic transport cost and Kullback--Leibler penalties on the marginals relative to prescribed Gaussian measures. We show that the infinite-dimensional variational problem admits an exact Gaussian reduction, yielding a finite-dimensional optimization over masses, means, and covariances, together with a closed-form expression for the optimal transported mass. We then formulate UDC for discrete-time linear systems, where the initial and terminal state measures are imposed softly through KL penalties and the intermediate evolution is governed by controlled linear dynamics with quadratic control cost. For this problem, we prove that any feasible solution can be replaced, without loss of optimality, by a Gaussian initial measure and an affine-Gaussian control policy. This leads to an exact finite-dimensional reformulation and, after a standard covariance-steering lifting, to an SDP-based optimization for fixed mass, again coupled with a closed-form mass update. We further establish existence of optimal solutions and identify a sufficient condition under which the affine-Gaussian UDC policy is deterministic. These results provide globally optimal solution methods for both Gaussian UOT and Gaussian UDC. Finally, we illustrate our results with several numerical examples.

Haruto Nakashima, Siddhartha Ganguly, Kenji Kashima

The Applied Mathematics and Physics Department, Graduate School of Informatics, Kyoto University, Kyoto, Japan · Daniel Guggenheim School of Aerospace, Georgia Institute of Technology, Atlanta, USA

Semidefinite Programming Stochastic Optimal Control Optimal Control Optimal Transport

May 5, 2026

A Provably Convergent and Practical Algorithm for Gromov--Wasserstein Optimal Transport

Gromov--Wasserstein optimal transport (GWOT) aligns metric measure spaces by matching their within-domain relational structures, but large-scale GWOT remains challenging because its objective is nonconvex and projection onto the transport polytope is often solved only approximately in practice. This leads to a gap between practical projected-gradient implementations and convergence theory, which typically assumes exact projections. For squared-loss GWOT, we propose an inexact projected-gradient framework with a verifiable feasibility-residual-based inexact condition for the projection subproblem. This condition is directly computable and avoids unknown quantities such as the exact projection point. Under this implementable condition, we prove subsequential convergence to stationary points and, with a mild tolerance-decay condition, convergence of the whole sequence. The resulting method retains the simplicity and sparsity of projected-gradient schemes while providing rigorous convergence guarantees, turning projected-gradient methods into a principled and scalable approach for GWOT with provable reliability.

Ling Liang, Lei Yang

Department of Mathematics The University of Tennessee Knoxville, TN 37916, USA · School of Computer Science and Engineering Guangdong Province Key Laboratory of Computational Science Sun Yat-sen University Guangzhou, Guangdong 510275, China

Gradient Descent Nonconvex Optimization Constrained Optimization Optimal Transport

May 4, 2026

Sampling-Based Control via Entropy-Regularized Optimal Transport

Sampling-based model predictive control methods like MPPI and CEM are essential for real-time control of nonlinear robotic systems, particularly where discontinuous dynamics preclude gradient-based optimization. However, these methods derive from information-theoretic objectives that are agnostic to the geometry of the control problem, leading to pathological behaviors such as mode-averaging when the cost landscape is complex. We present OT-MPC, a sampling-based algorithm that overcomes these limitations through an entropy-regularized optimal transport formulation. By computing an optimal coupling between candidate control sequences and low-cost proposals, OT-MPC refines candidates toward nearby promising samples while coordinating updates across the ensemble to maintain coverage of the solution space. We derive closed-form, gradient-free updates via the Sinkhorn algorithm, enabling real-time performance. Experiments on navigation, manipulation, and locomotion tasks demonstrate improved success rates over existing methods.

Vincent Pacelli, Akash Ratheesh, Evangelos A. Theodorou

Autonomous Control and Decision Systems Laboratory Georgia Institute of Technology

Robotic Control Nonlinear MPC Model Predictive Control Optimal Control Optimal Transport

May 3, 2026

Analytic Bridge Diffusions for Controlled Path Generation

Most modern bridge-diffusion methods achieve finite-time transport by specifying an interpolation, Schrödinger-bridge, or stochastic-control objective and then learning the associated score or drift field with a neural network. In contrast, we identify a restricted but sufficiently broad and analytically solvable class in which the score, intermediate marginals, and protocol gradients are available in closed form without inner stochastic simulation loops and without neural networks in the optimization loop. We recast the classical linear--quadratic--Gaussian (LQG) stochastic-control structure as a transport problem of the Path Integral Diffusion (PID) type. In classical LQG control, linear dynamics, Gaussian noise, and quadratic costs lead to Riccati equations and closed-form optimal feedback. In LQ-GM-PID, we retain the linear--quadratic stochastic-control backbone, but replace terminal state regulation by a prescribed terminal probability density and allow both the initial and terminal laws to be Gaussian Mixtures (GM). Moreover, LQ-GM-PID turns bridge diffusion from a tool for terminal target matching alone into a tool for path shaping. We demonstrate this on a 2D corridor task, a 2D multi-entrance transport task, and a high-dimensional scaling study with

d=32

and

M=16

Gaussian-mixture terminal modes, all with sub-50,ms analytic precompute on a laptop. We position LQ-GM-PID as an analytically solvable reference model for the state-of-the-art neural bridge-diffusion and generative-transport methods: a controlled setting in which neural approximations, score estimates, path-shaping objectives, and protocol-learning procedures can be tested against exact quantities.

Michael Chertkov

Program in Applied Mathematics & Department of Mathematics University of Arizona, AZ 85721, USA

Stochastic Optimal Control Schrödinger Bridges Linear Quadratic Regulator Optimal Control Optimal Transport

May 2, 2026

Data-Driven, Geometry-Aware Optimal-Transport Calibration of Flavor Tagger

Flavor-tagging calibrations are often provided either as scale factors measured at a finite set of working points or as binned corrections to a chosen one-dimensional discriminant. However, this approach falls short of providing continuous, event-level calibration across the full multicomponent outputs of modern taggers. This limitation leads to information loss in analyses that demand high-performance flavor tagging, restricting analyses to a limited set of predefined variables. In this work, we propose a geometry-aware framework that formulates flavor-tagger calibration as an optimal transport problem on the probability simplex. The transport maps are parameterized and trained in the isometric log-ratio coordinate system. Because the quadratic Euclidean cost of Brenier transport in this coordinate system is equivalent to the Aitchison distance on the simplex, the learned map induces a minimal deformation under the Aitchison geometry. Furthermore, we extract flavor-conditional target distributions directly from control-region data using an expectation-maximization (EM) technique that simultaneously fits multiple control regions, models each flavor component with a normalizing flow, and estimates the regional mixture fractions. The extracted targets are subsequently used to learn flavor-factorized transport maps. Because the joint estimation of mixture fractions and flexible component densities admits weakly constrained directions, we further introduce a linearized feedback-operator analysis that propagates the fitted composition covariance into the extracted component densities, separating data-constrained modes from those dominated by the composition prior. The simulation-based closure study demonstrates improved closure in dedicated control regions and in independent validation mixtures.

Yeonjoon Kim, Un-ki Yang

Department of Physics and Astronomy, Seoul National University, 1 Gwanak-ro, Gwanak-gu, Seoul 08826, Korea

Model Calibration Optimal Transport

Apr 29, 2026

Hyper Input Convex Neural Networks for Shape Constrained Learning and Optimal Transport

We introduce Hyper Input Convex Neural Networks (HyCNNs), a novel neural network architecture designed for learning convex functions. HyCNNs combine the principles of Maxout networks with input convex neural networks (ICNNs) to create a neural network that is always convex in the input, theoretically capable of leveraging depth, and performs reliable when trained at scale compared to ICNNs. Concretely, we prove that HyCNNs require exponentially fewer parameters than ICNNs to approximate quadratic functions up to a given precision. Throughout a series of synthetic experiments, we demonstrate that HyCNNs outperform existing ICNNs and MLPs in terms of predictive performance for convex regression and interpolation tasks. We further apply HyCNNs to learn high-dimensional optimal transport maps for synthetic examples and for single-cell RNA sequencing data, where they oftentimes outperform ICNN-based neural optimal transport methods and other baselines across a wide range of settings.

Shayan Hundrieser, Insung Kong, Johannes Schmidt-Hieber

*Department of Applied Mathematics, University of Twente, Enschede, The Netherlands

Input Convex Neural Networks Neural Network Approximation Theory Optimal Transport

Apr 28, 2026

Block-Wise Differentiable Sinkhorn Attention: Tail-Refinement Gradients with a Gap-Aware Dustbin Bridge

We study long-context balanced entropic optimal transport (OT) attention on TPU hardware through a stopped-base, fixed-depth tail-refinement surrogate. After a stopped

T

-step Sinkhorn solve, we unroll a short refinement tail and differentiate that surrogate exactly. For the reported

R=2

TPU path, the backward pass contains four staircase plan factors. We prove an exact one-reference-tile schedule: the

R=2

score cotangent is a single reference plan tile times an explicit modifier field built from vector cotangents and dual differences. This yields block-wise cost

O((T+R)LW)

,

O(Ld)

input storage, and

O(L)

additional HBM usage for fixed head dimension

d

and band width

W

on the balanced fixed-support path. We also formalize the current dustbin\_block path as the same unit-target surrogate on an augmented support, so the adjoint schedule lifts to the single-active-dustbin path used in our TPU runs; this bridge is algebraic and does not claim a general KL-unbalanced or arbitrary-capacity gap model. We provide a local surrogate-bias bound, an a posteriori bias certificate, and a projective contraction certificate for strictly positive active blocks. On synthetic masked problems, the optimized kernel matches exact autodiff of the same centered surrogate to within

10^{-5}

--

10^{-10}

. On TPU v6e-8, a four-configuration Pfam screen completes end-to-end, and a promoted balanced

R=2

run sustains roughly

8.5

examples per second through a three-hour budget, reaching step

1437

. Held-out Pfam test shards improve reconstruction from

5.57

to

2.05

and sparse CE from

5.53

to

5.30

relative to step

0

, with CE logged diagnostically rather than optimized directly; target-barycenter alignment metrics do not materially improve, and a deterministic diagonal reference remains stronger on those metrics.

Dylan Forde

Independent Researcher

Automatic Differentiation Transformer Attention Differentiable Programming Algorithm Unrolling Attention Mechanisms Optimal Transport

Apr 27, 2026

Sliced-Regularized Optimal Transport

We propose a new regularized optimal transport (OT) formulation, termed sliced-regularized optimal transport (SROT). Unlike entropic OT (EOT), which regularizes the transport plan toward an independent coupling, SROT regularizes it toward a smoothened sliced OT (SOT) plan. To the best of our knowledge, SROT is the first approach to leverage a version of SOT plan as a reference to improve classical OT. We provide a formal definition of SROT, derive its dual formulation, and provide a post-Bayesian interpretation of SROT. We then develop a Sinkhorn-style algorithm for efficient computation, retaining the same scalability advantages as EOT. By incorporating a scalable SOT plan as a prior, SROT yields more accurate approximations of the exact OT plan than EOT under the same level of regularization. Moreover, the resulting transport plan improves upon the reference SOT plan itself. We further introduce the corresponding OT divergence induced by SROT, named SROT divergence, and analyze its topological and computational properties. Finally, we validate our approach through experiments on synthetic datasets and color transfer tasks, demonstrating that SROT is better than both EOT and SOT in approximating exact OT. Additional experiments on gradient flows further highlight the advantages of SROT divergence.

Khai Nguyen

The University of Texas at Austin

Optimal Transport

Apr 26, 2026

Gromov-Wasserstein Methods for Multi-View Relational Embedding and Clustering

Learning low-dimensional representations from multi-view relational data is challenging when underlying geometries differ across views. We propose Bary-GWMDS, a Gromov-Wasserstein-based method that operates directly on distance matrices to learn a consensus embedding preserving shared relational structure. By leveraging intrinsic distances, the approach naturally handles nonlinear distortions across views. We also introduce Mean-GWMDS-C, a clustering-oriented formulation that averages distance matrices and learns reduced-support representations via a consensus Gromov-Wasserstein transport. Experiments on synthetic and real-world datasets show that the proposed framework yields stable and geometrically meaningful embeddings.

Rafael Pereira Eufrazio, Eduardo Fernandes Montesuma, Charles Casimiro Cavalcante

Instituto Federal de Educação, Ciência e Tecnologia do Ceará, Canindé, Ceará, Brazil · Federal University of Ceara, Fortaleza, Ceará, Brazil · Sigma Nova, Paris, France

Representation Learning Multidimensional Scaling Multi-View Learning Multi-View Clustering Dimensionality Reduction Optimal Transport

Apr 22, 2026

Learning to Emulate Chaos: Adversarial Optimal Transport Regularization

Chaos arises in many complex dynamical systems, from weather to power grids, but is difficult to accurately model with data-driven methods such as machine learning emulators. While emulators are promising tools for accelerating simulations and solving inverse problems, they still struggle to learn chaotic dynamics, where sensitivity to initial conditions renders exact long-term forecasts infeasible, especially given noisy data. Recent work instead trains emulators to match the statistical properties of chaotic attractors, but these approaches often rely on handcrafted summary statistics or large, diverse multi-environment datasets. In this work, we propose a family of adversarial optimal transport objectives that can jointly learn high-quality summary statistics and a physically consistent emulator from a single noisy trajectory. We theoretically analyze and experimentally validate a Sinkhorn divergence formulation (2-Wasserstein) and a WGAN-style dual formulation (1-Wasserstein) of our approach. Numerical experiments across a variety of chaotic systems, including ones with high-dimensional spatiotemporal chaos, show that emulators trained using our proposed objectives have significantly improved long-term statistical fidelity.

Gabriel Melo, Leonardo Santiago, Peter Y. Lu

LTCI, Télécom Paris, Institut Polytechnique de Paris, Palaiseau · Department of Mechanical and Aerospace Engineering, North Carolina State University, Raleigh, NC · University of Campinas +1

Neural Surrogate Modeling Chaotic Dynamical Systems Optimal Transport Chaotic Time-Series Prediction

Apr 17, 2026

OT on the Map: Quantifying Domain Shifts in Geographic Space

In computer vision and machine learning for geographic data, out-of-domain generalization is a pervasive challenge, arising from uneven global data coverage and distribution shifts across geographic regions. Though models are frequently trained in one region and deployed in another, there is no principled method for determining when this cross-region adaptation will be successful. A well-defined notion of distance between distributions can effectively quantify how different a new target domain is compared to the domains used for model training, which in turn could support model training and deployment decisions. In this paper, we propose a strategy for computing distances between geospatial domains that leverages geographic information with Optimal Transport methods (GeoSpOT). In our experiments, GeoSpOT distances emerge as effective predictors of cross-domain transfer difficulty. We further demonstrate that embeddings from pretrained location encoders provide information comparable to image/text embeddings, despite relying solely on longitude-latitude pairs as input. This allows users to get an approximation of out-of-domain performance for geospatial models, even when the exact downstream task is unknown, or no task-specific data is available. Building on these findings, we show that GeoSpOT distances can preemptively guide data selection and enable predictive tools to analyze regions where a model is likely to underperform.

Haoran Zhang, Livia Betti, Konstantin Klemmer +2

Harvard University · University of Colorado Boulder · LGND AI, Inc. +3

Distribution Shift Geospatial Representation Learning OOD Generalization Distribution Shift Detection Cross-Domain Transfer Learning Optimal Transport

Apr 16, 2026

One-Shot Generative Flows: Existence and Obstructions

We study dynamic measure transport for generative modeling, focusing on transport maps that connect a source measure

P_0

to a target measure

P_1

by integrating a velocity field of the form

v_t(x) = \mathbb{E}[\dot X_t \mid X_t = x]

, where

X_\bullet = (X_t)_t

is a stochastic process satisfying

(X_0,X_1)\sim{P_0}\otimes{P_1}

and

\dot X_t

is its time derivative. We investigate when

X_\bullet

induces a straight-line flow: a flow whose pointwise acceleration vanishes and is therefore exactly integrable by any first-order method. First, we develop multiple characterizations of straight-line flows in terms of PDEs involving the conditional statistics of the process. Then, we prove that straight-line flows under endpoint independence exhibit a sharp dichotomy. On the one hand, we construct explicit, computable straight-line processes for arbitrary Gaussian endpoints. On the other hand, we show that straight-line processes do not exist for targets with sufficiently well-separated modes. We demonstrate this obstruction through a sequence of increasingly general impossibility theorems that uncover a fundamental relationship between the sample-path behavior of a process with independent endpoints and the space-time geometry of this process' flow map. Taken together, these results provide a structural theory of when straight-line generative flows can, and cannot, exist.

Panos Tsimpos, Daniel Sharp, Youssef Marzouk

Operations Research Center · Massachusetts Institute of Technology · Cambridge, MA 02139 +2

Flow Matching One-Step Generative Modeling Flow-Based Generative Modeling Optimal Transport

Apr 16, 2026

Amortized Optimal Transport from Sliced Potentials

We propose a novel amortized optimization method for predicting optimal transport (OT) plans across multiple pairs of measures by leveraging Kantorovich potentials derived from sliced OT. We introduce two amortization strategies: regression-based amortization (RA-OT) and objective-based amortization (OA-OT). In RA-OT, we formulate a functional regression model that treats Kantorovich potentials from the original OT problem as responses and those obtained from sliced OT as predictors, and estimate these models via least-squares methods. In OA-OT, we estimate the parameters of the functional model by optimizing the Kantorovich dual objective. In both approaches, the predicted OT plan is subsequently recovered from the estimated potentials. As amortized OT methods, both RA-OT and OA-OT enable efficient solutions to repeated OT problems across different measure pairs by reusing information learned from prior instances to rapidly approximate new solutions. Moreover, by exploiting the structure provided by sliced OT, the proposed models are more parsimonious, independent of specific structures of the measures, such as the number of atoms in the discrete case, while achieving high accuracy. We demonstrate the effectiveness of our approaches on tasks including MNIST digit transport, color transfer, supply-demand transportation on spherical data, and mini-batch OT conditional flow matching.

Minh-Phuc Truong, Khai Nguyen

Hanoi University of Science and Technology · The University of Texas at Austin

Optimal Transport

Apr 16, 2026

Enhancing time-frequency resolution with optimal transport and barycentric fusion of multiple spectrogram

Time-frequency representations, such as the short-time Fourier transform (STFT), are fundamental tools for analyzing non-stationary signals. However, their ability to achieve sharp localization in both time and frequency is inherently limited by the Gabor-Heisenberg uncertainty principle. In this paper, we address this limitation by introducing a method to generate super-resolution spectrograms through the fusion of two or more spectrograms with varying resolutions. Specifically, we compute the super-resolution spectrogram as the barycenter of input spectrograms using optimal transport (OT) divergences. Unlike existing fusion approaches, our method does not require the input spectrograms to share the same time-frequency grid. Instead, the input spectrograms can be computed using any STFT parameters, and the resulting super-resolution spectrogram can be defined on an arbitrary user-specified grid. We explore various OT divergences based on different transportation costs. Notably, we introduce a novel transportation cost that preserves time-frequency geometry while significantly reducing computational complexity compared to standard Wasserstein barycenters. We adopt the unbalanced OT framework and derive a new block majorization-minimization algorithm for efficient barycenter computation. We validate the proposed method on controlled synthetic signals and recorded speech using both quantitative and qualitative evaluations. The results show that our approach combines the best localization properties of the input spectrograms and outperforms an unsupervised state-of-the-art fusion method.

David Valdivia, Elsa Cazelles, Cédric Févotte

IRIT, University of Toulouse, CNRS, France

Optimal Transport

Apr 14, 2026

Sliced Wasserstein Steering between Gaussian Measures

Optimal transport with quadratic cost provides a geometric framework for steering an ensemble, modeled by a probability law, with minimal effort. Yet ambient-space formulations become unwieldy in high dimensions, and sensing or actuation in practice often reveals only linear views of the state -- camera silhouettes, LiDAR beams, tomographic slices. We develop a sliced feedback controller for distribution steering: the evolving law is projected onto one-dimensional directions on the sphere, the optimal one-dimensional velocity is synthesized in each projection, and these velocities are averaged to produce a feedback control in the ambient space. The construction reduces to the Benamou--Brenier problem in one dimension. In addition, it is invariant under orthogonal transforms, nonexpansive under projections, and well posed on

\mathcal{P}_2(\mathbb{R}^n)

. Computation proceeds by sampling directions on the sphere and solving independent one-dimensional subproblems, yielding a scalable method aligned with partial observations. In the Gaussian setting, we show that the developed sliced controller steers the law to the prescribed target. Furthermore, we derive an identity relating the energy consumption incurred by the controller to the sliced Wasserstein distance.

Kaito Ito, Anqi Dong

Department of Information Physics and Computing, The University of Tokyo, Tokyo 113-8654, Japan · Department of Decision and Control Systems and Department of Mathematics, KTH Royal Institute of Technology, Stockholm 114 28, Sweden

Optimal Transport Optimal Control Sliced Wasserstein Distance

Apr 9, 2026

Generative optimal transport via forward-backward HJB matching

Controlling the evolution of a many-body stochastic system from a disordered reference state to a structured target ensemble, characterized empirically through samples, arises naturally in non-equilibrium statistical mechanics and stochastic control. The natural relaxation of such a system - driven by diffusion - runs from the structured target toward the disordered reference. The natural question is then: what is the minimum-work stochastic process that reverses this relaxation, given a pathwise cost functional combining spatial penalties and control effort? Computing this optimal process requires knowledge of trajectories that already sample the target ensemble - precisely the object one is trying to construct. We resolve this by establishing a time-reversal duality: the value function governing the hard backward dynamics satisfies an equivalent forward-in-time HJB equation, whose solution can be read off directly from the tractable forward relaxation trajectories. Via the Cole-Hopf transformation and its associated Feynman-Kac representation, this forward potential is computed as a path-space free energy averaged over these forward trajectories - the same relaxation paths that are easy to simulate - without any backward simulation or knowledge of the target beyond samples. The resulting framework provides a physically interpretable description of stochastic transport in terms of path-space free energy, risk-sensitive control, and spatial cost geometry. We illustrate the theory with numerical examples that visualize the learned value function and the induced controlled diffusions, demonstrating how spatial cost fields shape transport geometry analogously to Fermat's Principle in inhomogeneous media. Our results establish a unifying connection between stochastic optimal control, Schrödinger bridge theory, and non-equilibrium statistical mechanics.

Haiqian Yang, Vishaal Krishnan, Sumit Sinha +1

Department of Mechanical Engineering, MIT, Cambridge, MA 02139 · SEAS, Harvard University, Cambridge, MA 02138 · SEAS, Physics and OEB, Harvard University, Cambridge, MA 02138

HJB Equations Stochastic Optimal Control Optimal Transport Schrödinger Bridges

Apr 8, 2026

Region-Graph Optimal Transport Routing for Mixture-of-Experts Whole-Slide Image Classification

Multiple Instance Learning (MIL) is the dominant framework for gigapixel whole-slide image (WSI) classification in computational pathology. However, current MIL aggregators route all instances through a shared pathway, constraining their capacity to specialise across the pathological heterogeneity inherent in each slide. Mixture-of-Experts (MoE) methods offer a natural remedy by partitioning instances across specialised expert subnetworks; yet unconstrained softmax routing may yield highly imbalanced utilisation, where one or a few experts absorb most routing mass, collapsing the mixture back to a near-single-pathway solution. To address these limitations, we propose ROAM (Region-graph OptimAl-transport Mixture-of-experts), a spatially aware MoE-MIL aggregator that routes region tokens to expert poolers via capacity-constrained entropic optimal transport, promoting balanced expert utilisation by construction. ROAM operates on spatial region tokens, obtained by compressing dense patch bags into spatially binned units that align routing with local tissue neighbourhoods and introduces two key mechanisms: (i) region-to-expert assignment formulated as entropic optimal transport (Sinkhorn) with explicit per slide capacity marginals, enforcing balanced expert utilisation without auxiliary load-balancing losses; and (ii) graph-regularised Sinkhorn iterations that diffuse routing assignments over the spatial region graph, encouraging neighbouring regions to coherently route to the same experts. Evaluated on four WSI benchmarks with frozen foundation-model patch embeddings, ROAM achieves performance competitive against strong MIL and MoE baselines, and on NSCLC generalisation (TCGA-CPTAC) reaches external AUC 0.845 +- 0.019.

Xin Tian, Jiuliu Lu, Ephraim Tsalik +3

Centre for Human Genetics, University of Oxford, Oxford, UK · Beckman Coulter Diagnostics, Brea, CA, USA · Danaher Corporation, Washington, DC, USA

Histopathology Image Classification Whole-Slide Image Analysis Mixture-of-Experts Models Multiple Instance Learning Optimal Transport

Apr 8, 2026

SCOT: Multi-Source Cross-City Transfer with Optimal-Transport Soft-Correspondence Objective

Cross-city transfer improves prediction in label-scarce cities by leveraging labeled data from other cities, but it becomes challenging when cities adopt incompatible partitions and no ground-truth region correspondences exist. Existing approaches either rely on heuristic region matching, which is often sensitive to anchor choices, or perform distribution-level alignment that leaves correspondences implicit and can be unstable under strong heterogeneity. We propose SCOT, a cross-city representation learning framework that learns explicit soft correspondences between unequal region sets via Sinkhorn-based entropic optimal transport. SCOT further sharpens transferable structure with an OT-weighted contrastive objective and stabilizes optimization through a cycle-style reconstruction regularizer. For multi-source transfer, SCOT aligns each source and the target to a shared prototype hub using balanced entropic transport guided by a target-induced prototype prior. Across real-world cities and tasks, SCOT consistently improves transfer accuracy and robustness, while the learned transport couplings and hub assignments provide interpretable diagnostics of alignment quality.

Yuyao Wang, Min Yang, Meng Chen +3

Department of Mathematics and Statistics, Boston University, Boston, MA, USA · School of Software, Shandong University, Jinan, China · School of Geography, University of Leeds, Leeds, UK

Transfer Learning Multi-Source Domain Adaptation Geospatial Representation Learning Optimal Transport

Apr 8, 2026

A Generalized Sinkhorn Algorithm for Mean-Field Schrödinger Bridge

The mean-field Schrödinger bridge (MFSB) problem concerns designing a minimum-effort controller that guides a diffusion process with nonlocal interaction to reach a given distribution from another by a fixed deadline. Unlike the standard Schrödinger bridge, the dynamical constraint for MFSB is the mean-field limit of a population of interacting agents with controls. It serves as a natural model for large-scale multi-agent systems. The MFSB is computationally challenging because the nonlocal interaction makes the problem nonconvex. We propose a generalization of the Hopf-Cole transform for MFSB and, building on it, design a Sinkhorn-type recursive algorithm to solve the associated system of integro-PDEs. Under mild assumptions on the interaction potential, we discuss convergence guarantees for the proposed algorithm. We present numerical examples with repulsive and attractive interactions to illustrate the theoretical contributions.

Asmaa Eldesoukey, Yongxin Chen, Abhishek Halder

Department of Aerospace Engineering, Iowa State University, Ames, IA 50011, USA · School of Aerospace Engineering, Georgia Institute of Technology, Atlanta, GA 30332, USA

Schrödinger Bridges Optimal Transport

Apr 7, 2026

Optimal-Transport-Guided Functional Flow Matching for Turbulent Field Generation in Hilbert Space

High-fidelity modeling of turbulent flows requires capturing complex spatiotemporal dynamics and multi-scale intermittency, posing a fundamental challenge for traditional knowledge-based systems. While deep generative models, such as diffusion models and Flow Matching, have shown promising performance, they are fundamentally constrained by their discrete, pixel-based nature. This limitation restricts their applicability in turbulence computing, where data inherently exists in a functional form. To address this gap, we propose Functional Optimal Transport Conditional Flow Matching (FOT-CFM), a generative framework defined directly in infinite-dimensional function space. Unlike conventional approaches defined on fixed grids, FOT-CFM treats physical fields as elements of an infinite-dimensional Hilbert space, and learns resolution-invariant generative dynamics directly at the level of probability measures. By integrating Optimal Transport (OT) theory, we construct deterministic, straight-line probability paths between noise and data measures in Hilbert space. This formulation enables simulation-free training and significantly accelerates the sampling process. We rigorously evaluate the proposed system on a diverse suite of chaotic dynamical systems, including the Navier-Stokes equations, Kolmogorov Flow, and Hasegawa-Wakatani equations, all of which exhibit rich multi-scale turbulent structures. Experimental results demonstrate that FOT-CFM achieves superior fidelity in reproducing high-order turbulent statistics and energy spectra compared to state-of-the-art baselines.

Li Kunpeng, Wan Chenguang, Qu Zhisong +5

Flow-Based Generative Modeling Flow Matching Computational Fluid Dynamics Optimal Transport

Apr 6, 2026

Muon Dynamics as a Spectral Wasserstein Flow

Gradient normalization can stabilize deep-learning optimization, and spectral normalizations are especially natural for matrix-shaped parameter blocks; Muon is the motivating example. We study an idealized deterministic, continuous-time, vanishing-momentum version of this idea in the mean-field regime, where wide models are represented by probability measures on parameter space. Starting from normalized matrix flows, we introduce Spectral Wasserstein distances indexed by norms

γ

on positive semidefinite matrices: the trace norm recovers the classical

W_2

distance, the operator norm gives the Muon geometry, and Schatten norms interpolate between them. We develop a static Kantorovich formulation, a max--min robust-cost representation, and Gaussian reductions extending the Bures formula. For monotone norms, we prove equivalence with a Benamou--Brenier formulation and obtain a gradient-flow interpretation of normalized mean-field training dynamics. Numerical experiments illustrate the resulting flows for MMD objectives, Gaussian reductions, two-layer ReLU models, and shallow attention.

Gabriel Peyré

CNRS and ENS, PSL Université

Neural Network Training Dynamics Wasserstein Distance Muon Optimizer Optimal Transport Wasserstein Gradient Flows

Apr 4, 2026

On the Efficiency of Sinkhorn-Knopp for Entropically Regularized Optimal Transport

The Sinkhorn--Knopp (SK) algorithm is a cornerstone method for matrix scaling and entropically regularized optimal transport (EOT). Despite its empirical efficiency, existing theoretical guarantees to achieve a target marginal accuracy

\varepsilon

deteriorate severely in the presence of outliers, bottlenecked either by the global maximum regularized cost

η\|C\|_\infty

(where

η

is the regularization parameter and

C

the cost matrix) or the matrix's minimum-to-maximum entry ratio

ν

. This creates a fundamental disconnect between theory and practice. In this paper, we resolve this discrepancy. For EOT, we introduce the novel concept of well-boundedness, a local bulk mass property that rigorously isolates the well-behaved portion of the data from extreme outliers. We prove that governed by this fundamental notion, SK recovers the target transport plan for a problem of dimension

n

in

O(\log n - \log \varepsilon)

iterations, completely independent of the regularized cost

η\|C\|_\infty

. Furthermore, we show that a virtually cost-free pre-scaling step eliminates the dimensional dependence entirely, accelerating convergence to a strictly dimension-free

O(\log(1/\varepsilon))

iterations. Beyond EOT, we establish a sharp phase transition for general

(\boldsymbol{u},\boldsymbol{v})

-scaling governed by a critical matrix density threshold. We prove that when a matrix's density exceeds this threshold, the iteration complexity is strictly independent of

ν

. Conversely, when the density falls below this threshold, the dependence on

ν

becomes unavoidable; in this sub-critical regime, we construct instances where SK requires

Ω(n/\varepsilon)

iterations.

Kun He

Renmin University of China

Optimal Transport Entropy Regularization

Apr 4, 2026

Fast Log-Domain Sinkhorn Optimal Transport with Warp-Level GPU Reductions

Entropic regularized optimal transport (OT) via the Sinkhorn algorithm has become a fundamental tool in machine learning, yet existing implementations either suffer from numerical instability for small regularization parameters or incur significant overhead from deep learning frameworks. We present FastSinkhorn, a lightweight, native CUDA implementation of the log-domain Sinkhorn algorithm that combines warp-level shuffle reductions with shared-memory tiling to achieve high GPU utilization without sacrificing numerical stability. Our solver operates entirely in the log-domain, enabling robust computation for regularization parameters as small as epsilon = 10^{-4} where standard-domain methods fail. On dense OT problems with n = m = 8192, our implementation achieves 12x speedup over the widely-used POT library and 5.9x speedup over GPU-accelerated PyTorch baselines, while consuming only 256 MB of GPU memory. We validate our solver on image color transfer, 3D point cloud matching, and convergence analysis, demonstrating that native CUDA kernels with careful numerical treatment provide a practical and efficient foundation for large-scale optimal transport computation.

Hao Xiao

ATLAS AI Lab

Optimal Transport GPU Acceleration GPU Kernel Optimization