# Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation

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

## Key facts

- Cited by 33 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 63 references that are in the map.
- First posted on arXiv 25 Nov 2023.

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

- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Deep Learning Face Attributes in the Wild](https://mlsanity.com/paper/1411.7766/) (2014)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [MLP-Mixer: An all-MLP Architecture for Vision](https://mlsanity.com/paper/2105.01601/) (2021)
- [StarGAN: Unified Generative Adversarial Networks for Multi-Domain Image-to-Image Translation](https://mlsanity.com/paper/1711.09020/) (2017)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Contrastive Learning for Unpaired Image-to-Image Translation](https://mlsanity.com/paper/2007.15651/) (2020)
- [ILVR: Conditioning Method for Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2108.02938/) (2021)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [On Neural Differential Equations](https://mlsanity.com/paper/2202.02435/) (2022)
- [EMNIST: an extension of MNIST to handwritten letters](https://mlsanity.com/paper/1702.05373/) (2017)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [TrajectoryNet: A Dynamic Optimal Transport Network for Modeling Cellular Dynamics](https://mlsanity.com/paper/2002.04461/) (2020)
- [Optimal Entropy-Transport problems and a new Hellinger-Kantorovich distance between positive measures](https://mlsanity.com/paper/1508.07941/) (2015)
- [UNIT-DDPM: UNpaired Image Translation with Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2104.05358/) (2021)
- [Neural Optimal Transport](https://mlsanity.com/paper/2201.12220/) (2022)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [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)
- [Simulation-free Schr"odinger bridges via score and flow matching](https://mlsanity.com/paper/2307.03672/) (2023)
- [Unbalanced Optimal Transport, from Theory to Numerics](https://mlsanity.com/paper/2211.08775/) (2022)
- [Unpaired Image-to-Image Translation using Adversarial Consistency Loss](https://mlsanity.com/paper/2003.04858/) (2020)
- [Scalable Unbalanced Optimal Transport using Generative Adversarial Networks](https://mlsanity.com/paper/1810.11447/) (2018)

## Built on by (most cited first, 30 of 33)

- [Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models](https://mlsanity.com/paper/2508.09968/) (2025)
- [Wasserstein Flow Matching: Generative modeling over families of distributions](https://mlsanity.com/paper/2411.00698/) (2024)
- [Joint Velocity-Growth Flow Matching for Single-Cell Dynamics Modeling](https://mlsanity.com/paper/2505.13413/) (2025)
- [Multi-Modal and Multi-Attribute Generation of Single Cells with CFGen](https://mlsanity.com/paper/2407.11734/) (2024)
- [Modeling Cell Dynamics and Interactions with Unbalanced Mean Field Schr"odinger Bridge](https://mlsanity.com/paper/2505.11197/) (2025)
- [Variational Regularized Unbalanced Optimal Transport: Single Network, Least Action](https://mlsanity.com/paper/2505.11823/) (2025)
- [Robust Barycenter Estimation using Semi-Unbalanced Neural Optimal Transport](https://mlsanity.com/paper/2410.03974/) (2024)
- [Integrating Dynamical Systems Modeling with Spatiotemporal scRNA-seq Data Analysis](https://mlsanity.com/paper/2503.11347/) (2025)
- [WFR-FM: Simulation-Free Dynamic Unbalanced Optimal Transport](https://mlsanity.com/paper/2601.06810/) (2026)
- [Scalable Simulation-free Entropic Unbalanced Optimal Transport](https://mlsanity.com/paper/2410.02656/) (2024)
- [Inferring stochastic dynamics with growth from cross-sectional data](https://mlsanity.com/paper/2505.13197/) (2025)
- [Disentangled Representation Learning with the Gromov-Monge Gap](https://mlsanity.com/paper/2407.07829/) (2024)
- [Identifiable Shared Component Analysis of Unpaired Multimodal Mixtures](https://mlsanity.com/paper/2409.19422/) (2024)
- [Towards generalizable single-cell perturbation modeling via the Conditional Monge Gap](https://mlsanity.com/paper/2504.08328/) (2025)
- [Diversified Flow Matching with Translation Identifiability](https://mlsanity.com/paper/2511.05558/) (2025)
- [Fast and Large-Scale Unbalanced Optimal Transport via its Semi-Dual and Adaptive Gradient Methods](https://mlsanity.com/paper/2602.10697/) (2026)
- [Beyond Continuity: Simulation-free Reconstruction of Discrete Branching Dynamics from Single-cell Snapshots](https://mlsanity.com/paper/2605.00545/) (2026)
- [Unsupervised Point Cloud Completion through Unbalanced Optimal Transport](https://mlsanity.com/paper/2410.02671/) (2024)
- [GeOT: A spatially explicit framework for evaluating spatio-temporal predictions](https://mlsanity.com/paper/2410.11709/) (2024)
- [Reweighted Flow Matching via Unbalanced OT for Label-free Long-tailed Generation](https://mlsanity.com/paper/2509.25713/) (2025)
- [WFR-MFM: One-Step Inference for Dynamic Unbalanced Optimal Transport](https://mlsanity.com/paper/2601.20606/) (2026)
- [CellBRIDGE: Learning Cellular Trajectories via Interaction-Aware Alignment](https://mlsanity.com/paper/2605.30635/) (2026)
- [TracingFlow: A Simulation-Free Trajectory Inference Framework Based on Second-Order Dynamics](https://mlsanity.com/paper/2608.21070/) (2026)
- [Simulation-free Structure Learning for Stochastic Population Dynamics](https://mlsanity.com/paper/2510.16656/) (2025)
- [Structured Matching via Cost-Regularized Unbalanced Optimal Transport](https://mlsanity.com/paper/2511.19075/) (2025)
- [Multiscale Supervised Unbalanced Optimal Transport Flow Matching](https://mlsanity.com/paper/2605.16529/) (2026)
- [Domain Transfer Becomes Identifiable via a Single Alignment](https://mlsanity.com/paper/2605.17918/) (2026)
- [UOTIP: Unbalanced Optimal Transport Map for Unpaired Inverse Problems](https://mlsanity.com/paper/2605.21094/) (2026)
- [Conditional Neural Optimal Transport for Predicting Cellular Phenotypes from Molecular Structure](https://mlsanity.com/paper/2608.14293/) (2026)
- [Simulation-free Unbalanced Dynamic Optimal Transport with General Growth Penalty](https://mlsanity.com/paper/2609.04710/) (2026)

## Abstract (arXiv)

In optimal transport (OT), a Monge map is known as a mapping that transports a source distribution to a target distribution in the most cost-efficient way. Recently, multiple neural estimators for Monge maps have been developed and applied in diverse unpaired domain translation tasks, e.g. in single-cell biology and computer vision. However, the classic OT framework enforces mass conservation, which makes it prone to outliers and limits its applicability in real-world scenarios. The latter can be particularly harmful in OT domain translation tasks, where the relative position of a sample within a distribution is explicitly taken into account. While unbalanced OT tackles this challenge in the discrete setting, its integration into neural Monge map estimators has received limited attention. We propose a theoretically grounded method to incorporate unbalancedness into any Monge map estimator. We improve existing estimators to model cell trajectories over time and to predict cellular responses to perturbations. Moreover, our approach seamlessly integrates with the OT flow matching (OT-FM) framework. While we show that OT-FM performs competitively in image translation, we further improve performance by incorporating unbalancedness (UOT-FM), which better preserves relevant features. We hence establish UOT-FM as a principled method for unpaired image translation.
