Research map / Label-efficient and robust learning / Recognition and long-tail
Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a re-scaling scheme that makes any neural Monge map estimator unbalanced via a discrete unbalanced OT coupling, and applies it to OT-ICNN, Monge gap, and OT flow matching (UOT-FM) across single-cell, perturbation, and image translation tasks.
Problem
Classic optimal transport enforces mass conservation, making neural Monge map estimators sensitive to outliers and class imbalance between source and target distributions, which hurts unpaired domain translation.
Approach
Prove that an unbalanced Monge map between two measures equals a balanced Monge map between rescaled measures obtained from the unbalanced OT coupling; this rescaling (estimated batch-wise with entropic UOT) can be plugged into any balanced Monge map estimator, including OT flow matching.
Key results
- On CelebA 256x256 Male→Female, UOT-FM reaches FID 13.9 and T-Cost 56.8 versus OT-FM 14.6 / 59.3 and FM 17.5 / 89.2.
- On EMNIST digits→letters, average FID drops from 36.9 (OT-FM) to 23.0 (UOT-FM); for 8→B it falls from 83.1 to 42.6.
- UOT-ICNN raises correct transitions for Ngn3 EP cells from 7% (OT-ICNN) to 69%, and NEB from 67% to 85%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use unbalanced OT couplings instead of balanced couplings | OT-FM · EMNIST digits → letters (average) | FID: 36.9 -> 23.0 | Helped |
| Use unbalanced OT couplings instead of balanced couplings | OT-FM · EMNIST digits → letters (8 → B) | FID: 83.1 -> 42.6 | Helped |
| Use unbalanced OT couplings instead of balanced couplings | OT-FM · CelebA 256x256 Male → Female | FID: 14.6 -> 13.9 | Helped |
| Use unbalanced OT couplings instead of balanced couplings | OT-FM · CelebA 256x256 Female → Male | FID: 14.8 -> 14.3 | Helped |
| Use unbalanced OT couplings instead of balanced couplings | OT-FM · CelebA 256x256 Remove Glasses | FID: 20.9 -> 18.5 | Helped |
| Use unbalanced OT couplings instead of balanced couplings | OT-FM · CelebA 256x256 Add Glasses | FID: 20.1 -> 18.3 | Helped |
| Use unbalanced OT reweighting | OT-ICNN · Single-cell pancreas trajectory inference (Ngn3 EP) | correct transitions: 7% -> 69% | Helped |
| Use unbalanced OT reweighting | OT-ICNN · Single-cell pancreas trajectory inference (EB) | correct transitions: 54% -> 57% | Helped |
| Use unbalanced OT reweighting | OT-ICNN · Single-cell pancreas trajectory inference (NEB) | correct transitions: 67% -> 85% | Helped |
| Use unbalanced OT instead of balanced OT | OT-FM · CIFAR-10 | FID: 3.59 -> 3.47; 3.42; 3.79 | Mixed |
| Set τ=0.97 (ablation) | UOT-FM · CIFAR-10 | FID: 3.59 -> 3.42 | Helped |
| Lower τ from 0.99 to 0.95 (ablation) | UOT-FM · CIFAR-10 | FID: 3.47 -> 3.79 | Hurt |
Surprising findings
- On sciPlex, the competing unbalanced estimator diverges for lower tau (Sinkhorn divergence up to 27031 at tau=0.7), while UOT-ICNN stays stable.
Limitations
- Choice of the unbalancedness hyperparameter tau requires a grid search and is not evident a priori.
- Unbalanced couplings implicitly remove outliers, which may discard data points that are relevant to training.
Experimental setup
Datasets: EMNIST CelebA CIFAR-10 mouse pancreas scRNA-seq 4i drug perturbation data sciPlex
Benchmarks: EMNIST digits-to-letters CelebA 256x256 CelebA 64x64 CIFAR-10 single-cell trajectory inference (CellRank-based cell transitions) 4i drug response sciPlex
Metrics: FID attribute-wise FID transport cost Sinkhorn divergence L2 drug signature percentage of correct cell transitions
Baselines: OT-ICNN OT-MG OT-FM FM CycleGAN UVCGAN NOT scVelo TrajectoryNet WOT
Builds on: OT-ICNN Monge gap OT flow matching Sinkhorn algorithm unbalanced optimal transport
Closest papers by meaning
- Monge, Bregman and Occam: Interpretable Optimal Transport in High-Dimensions with Feature-Sparse Maps (ICML 2023 · similarity 0.89)
- Unbalanced CO-optimal Transport (AAAI 2023 · similarity 0.89)
- Overcoming Spurious Solutions in Semi-Dual Neural Optimal Transport: A Smoothing Approach for Learning the Optimal Transport Plan (ICML 2025 · similarity 0.86)
- Improving Neural Optimal Transport via Displacement Interpolation (ICLR 2025 · similarity 0.86)
- Hierarchical Refinement: Optimal Transport to Infinity and Beyond (ICML 2025 · similarity 0.85)
- Universal Neural Optimal Transport (ICML 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.