Variational Inference for Uncertain Optimal Transport via Sinkhorn Parametrization

Optimal Transport (OT) traditionally relies on a fixed ground cost to produce a single deterministic transport plan—a practice that overlooks the inherent variability and noise in real-world data. While recent sampling based approaches of OT offer a principled way to quantify this uncertainty, these are computationally prohibitive and struggle to scale. In this paper, we introduce Sinkhorn-parameterized Variational Inference, a first scalable variational framework for performing posterior inference over transport plans. Our key insight is that the Sinkhorn map can be treated as a differentiable reparameterization of the set of entropic plans. This enables the use of flexible generative models like normalizing flows to approximate distributions over transport plans while enforcing marginal constraints. We experimentally demonstrate that our method matches the quality of intensive sampling techniques at a fraction of the computational cost, scaling effectively to large-scale problems.

Optimal transport is a mathematical tool for matching one collection of objects to another, such as connecting suppliers to customers, cells at one time point to cells at another, or samples from two different datasets. Usually, this matching is computed using a fixed cost table and produces one best answer. In practice, however, these costs can be noisy or uncertain, so a single matching may hide important ambiguity. We introduce a method that learns not just one transport plan, but a distribution over many plausible transport plans. Our approach uses a procedure called Sinkhorn scaling, which automatically ensures that every sampled matching satisfies the required mass-balance constraints. We combine this with flexible generative models so that the method can represent complex uncertainty while remaining efficient to train. This makes it possible to quantify uncertainty in transport-based predictions without relying on expensive sampling methods. The resulting approach can scale to larger problems and may be useful in applications such as domain adaptation, logistics, and single-cell biology, where understanding several plausible matchings is often more informative than trusting one deterministic solution.