# Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-Couplings

arXiv: https://arxiv.org/abs/2209.15621 · HTML: https://mlsanity.com/paper/2209.15621/

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 65 references that are in the map.
- First posted on arXiv 30 Sep 2022.

## Builds on (most cited first, 30 of 33)

- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [SuperGlue: Learning Feature Matching with Graph Neural Networks](https://mlsanity.com/paper/1911.11763/) (2019)
- [CyCADA: Cycle-Consistent Adversarial Domain Adaptation](https://mlsanity.com/paper/1711.03213/) (2017)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Toward Multimodal Image-to-Image Translation](https://mlsanity.com/paper/1711.11586/) (2017)
- [Style Transfer from Non-Parallel Text by Cross-Alignment](https://mlsanity.com/paper/1705.09655/) (2017)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Learning with a Wasserstein Loss](https://mlsanity.com/paper/1506.05439/) (2015)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Temporal Cycle-Consistency Learning](https://mlsanity.com/paper/1904.07846/) (2019)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](https://mlsanity.com/paper/2103.03606/) (2021)
- [Wasserstein-2 Generative Networks](https://mlsanity.com/paper/1909.13082/) (2019)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Proximal Optimal Transport Modeling of Population Dynamics](https://mlsanity.com/paper/2106.06345/) (2021)
- [Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization](https://mlsanity.com/paper/2012.05942/) (2020)
- [Optimal Control Via Neural Networks: A Convex Approach](https://mlsanity.com/paper/1805.11835/) (2018)
- [On Unbalanced Optimal Transport: An Analysis of Sinkhorn Algorithm](https://mlsanity.com/paper/2002.03293/) (2020)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [Optimizing Functionals on the Space of Probabilities with Input Convex Neural Networks](https://mlsanity.com/paper/2106.00774/) (2021)
- [Scalable Computations of Wasserstein Barycenter via Input Convex Neural Networks](https://mlsanity.com/paper/2007.04462/) (2020)
- [Learning Cross-Domain Correspondence for Control with Dynamics Cycle-Consistency](https://mlsanity.com/paper/2012.09811/) (2020)
- [Scalable Unbalanced Optimal Transport using Generative Adversarial Networks](https://mlsanity.com/paper/1810.11447/) (2018)
- [Unbalanced Optimal Transport through Non-negative Penalized Linear Regression](https://mlsanity.com/paper/2106.04145/) (2021)
- [Fast Optimal Transport Averaging of Neuroimaging Data](https://mlsanity.com/paper/1503.08596/) (2015)
- [Faster Unbalanced Optimal Transport: Translation invariant Sinkhorn and 1-D Frank-Wolfe](https://mlsanity.com/paper/2201.00730/) (2022)
- [Spatio-Temporal Alignments: Optimal transport through space and time](https://mlsanity.com/paper/1910.03860/) (2019)
- [Input Convex Gradient Networks](https://mlsanity.com/paper/2111.12187/) (2021)
- [Reversible Gromov-Monge Sampler for Simulation-Based Inference](https://mlsanity.com/paper/2109.14090/) (2021)

## Built on by (most cited first, 17 of 17)

- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [Action Matching: Learning Stochastic Dynamics from Samples](https://mlsanity.com/paper/2210.06662/) (2022)
- [Generative Modeling through the Semi-dual Formulation of Unbalanced Optimal Transport](https://mlsanity.com/paper/2305.14777/) (2023)
- [Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation](https://mlsanity.com/paper/2311.15100/) (2023)
- [The Monge Gap: A Regularizer to Learn All Transport Maps](https://mlsanity.com/paper/2302.04953/) (2023)
- [On amortizing convex conjugates for optimal transport](https://mlsanity.com/paper/2210.12153/) (2022)
- [GENOT: Entropic (Gromov) Wasserstein Flow Matching with Applications to Single-Cell Genomics](https://mlsanity.com/paper/2310.09254/) (2023)
- [Minimax estimation of discontinuous optimal transport maps: The semi-discrete case](https://mlsanity.com/paper/2301.11302/) (2023)
- [Neural Optimal Transport with Lagrangian Costs](https://mlsanity.com/paper/2406.00288/) (2024)
- [Extremal Domain Translation with Neural Optimal Transport](https://mlsanity.com/paper/2301.12874/) (2023)
- [Unbalanced Diffusion Schr"odinger Bridge](https://mlsanity.com/paper/2306.09099/) (2023)
- [Light Unbalanced Optimal Transport](https://mlsanity.com/paper/2303.07988/) (2023)
- [Denoising Diffusions with Optimal Transport: Localization, Curvature, and Multi-Scale Complexity](https://mlsanity.com/paper/2411.01629/) (2024)
- [Branched Schr"odinger Bridge Matching](https://mlsanity.com/paper/2506.09007/) (2025)
- [Towards generalizable single-cell perturbation modeling via the Conditional Monge Gap](https://mlsanity.com/paper/2504.08328/) (2025)
- [Simulation-free Unbalanced Dynamic Optimal Transport with General Growth Penalty](https://mlsanity.com/paper/2609.04710/) (2026)
- [A Neural JKO Scheme for Hellinger-Kantorovich Gradient Flows via Monge-Growth Pairs](https://mlsanity.com/paper/2610.07602/) (2026)

## Abstract (arXiv)

Comparing unpaired samples of a distribution or population taken at different points in time is a fundamental task in many application domains where measuring populations is destructive and cannot be done repeatedly on the same sample, such as in single-cell biology. Optimal transport (OT) can solve this challenge by learning an optimal coupling of samples across distributions from unpaired data. However, the usual formulation of OT assumes conservation of mass, which is violated in unbalanced scenarios in which the population size changes (e.g., cell proliferation or death) between measurements. In this work, we introduce NubOT, a neural unbalanced OT formulation that relies on the formalism of semi-couplings to account for creation and destruction of mass. To estimate such semi-couplings and generalize out-of-sample, we derive an efficient parameterization based on neural optimal transport maps and propose a novel algorithmic scheme through a cycle-consistent training procedure. We apply our method to the challenging task of forecasting heterogeneous responses of multiple cancer cell lines to various drugs, where we observe that by accurately modeling cell proliferation and death, our method yields notable improvements over previous neural optimal transport methods.