Variational Inference for Uncertain Optimal Transport via Sinkhorn Parametrization

Ananyapam De, Linus Bleistein, Anton Frederik Thielmann, Benjamin Säfken
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:23250-23268, 2026.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-de26a, title = {Variational Inference for Uncertain Optimal Transport via {S}inkhorn Parametrization}, author = {De, Ananyapam and Bleistein, Linus and Thielmann, Anton Frederik and S\"{a}fken, Benjamin}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {23250--23268}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/de26a/de26a.pdf}, url = {https://proceedings.mlr.press/v306/de26a.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Variational Inference for Uncertain Optimal Transport via Sinkhorn Parametrization %A Ananyapam De %A Linus Bleistein %A Anton Frederik Thielmann %A Benjamin Säfken %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-de26a %I PMLR %P 23250--23268 %U https://proceedings.mlr.press/v306/de26a.html %V 306 %X 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.
APA
De, A., Bleistein, L., Thielmann, A.F. & Säfken, B.. (2026). Variational Inference for Uncertain Optimal Transport via Sinkhorn Parametrization. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:23250-23268 Available from https://proceedings.mlr.press/v306/de26a.html.

Related Material