# Supervised Training of Conditional Monge Maps

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

## Key facts

- Cited by 61 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 73 references that are in the map.
- First posted on arXiv 28 Jun 2022.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [Distributed Representations of Words and Phrases and their Compositionality](https://mlsanity.com/paper/1310.4546/) (2013)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [HyperNetworks](https://mlsanity.com/paper/1609.09106/) (2016)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Parameter-efficient Multi-task Fine-tuning for Transformers via Shared Hypernetworks](https://mlsanity.com/paper/2106.04489/) (2021)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [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)
- [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)
- [Entropic estimation of optimal transport maps](https://mlsanity.com/paper/2109.12004/) (2021)
- [Optimal Control Via Neural Networks: A Convex Approach](https://mlsanity.com/paper/1805.11835/) (2018)
- [Large-Scale Wasserstein Gradient Flows](https://mlsanity.com/paper/2106.00736/) (2021)
- [Learning Action Representations for Reinforcement Learning](https://mlsanity.com/paper/1902.00183/) (2019)
- [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)
- [Scalable Unbalanced Optimal Transport using Generative Adversarial Networks](https://mlsanity.com/paper/1810.11447/) (2018)
- [Wasserstein Dictionary Learning: Optimal Transport-based unsupervised non-linear dictionary learning](https://mlsanity.com/paper/1708.01955/) (2017)
- [The Natural Language of Actions](https://mlsanity.com/paper/1902.01119/) (2019)
- [Semi-dual Regularized Optimal Transport](https://mlsanity.com/paper/1811.05527/) (2018)
- [Input Convex Gradient Networks](https://mlsanity.com/paper/2111.12187/) (2021)
- [Multi-subject MEG/EEG source imaging with sparse multi-task regression](https://mlsanity.com/paper/1910.01914/) (2019)
- [Optimal Transport using GANs for Lineage Tracing](https://mlsanity.com/paper/2007.12098/) (2020)

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

- [Multisample Flow Matching: Straightening Flows with Minibatch Couplings](https://mlsanity.com/paper/2304.14772/) (2023)
- [Optimal Flow Matching: Learning Straight Trajectories in Just One Step](https://mlsanity.com/paper/2403.13117/) (2024)
- [How to Build the Virtual Cell with Artificial Intelligence: Priorities and Opportunities](https://mlsanity.com/paper/2409.11654/) (2024)
- [A Computational Framework for Solving Wasserstein Lagrangian Flows](https://mlsanity.com/paper/2310.10649/) (2023)
- [Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold](https://mlsanity.com/paper/2408.14608/) (2024)
- [Aligned Diffusion Schr"odinger Bridges](https://mlsanity.com/paper/2302.11419/) (2023)
- [Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching](https://mlsanity.com/paper/2403.18705/) (2024)
- [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)
- [Scalable Optimal Transport Methods in Machine Learning: A Contemporary Survey](https://mlsanity.com/paper/2305.05080/) (2023)
- [Sampling in Unit Time with Kernel Fisher-Rao Flow](https://mlsanity.com/paper/2401.03892/) (2024)
- [Conditional Optimal Transport on Function Spaces](https://mlsanity.com/paper/2311.05672/) (2023)
- [Dynamic Conditional Optimal Transport through Simulation-Free Flows](https://mlsanity.com/paper/2404.04240/) (2024)
- [Learning diffusion at lightspeed](https://mlsanity.com/paper/2406.12616/) (2024)
- [Optimal transport map estimation in general function spaces](https://mlsanity.com/paper/2212.03722/) (2022)
- [Minimax estimation of discontinuous optimal transport maps: The semi-discrete case](https://mlsanity.com/paper/2301.11302/) (2023)
- [Efficient Neural Network Approaches for Conditional Optimal Transport with Applications in Bayesian Inference](https://mlsanity.com/paper/2310.16975/) (2023)
- [Neural Unbalanced Optimal Transport via Cycle-Consistent Semi-Couplings](https://mlsanity.com/paper/2209.15621/) (2022)
- [Neural Optimal Transport with Lagrangian Costs](https://mlsanity.com/paper/2406.00288/) (2024)
- [Nonlinear Filtering with Brenier Optimal Transport Maps](https://mlsanity.com/paper/2310.13886/) (2023)
- [Learning of Population Dynamics: Inverse Optimization Meets JKO Scheme](https://mlsanity.com/paper/2506.01502/) (2025)
- [DICE: Discrete inverse continuity equation for learning population dynamics](https://mlsanity.com/paper/2507.05107/) (2025)
- [Unbalanced Low-rank Optimal Transport Solvers](https://mlsanity.com/paper/2305.19727/) (2023)
- [A generative flow for conditional sampling via optimal transport](https://mlsanity.com/paper/2307.04102/) (2023)
- [Universal Neural Optimal Transport](https://mlsanity.com/paper/2212.00133/) (2022)
- [Consistent Optimal Transport with Empirical Conditional Measures](https://mlsanity.com/paper/2305.15901/) (2023)
- [Generating Synthetic Datasets by Interpolating along Generalized Geodesics](https://mlsanity.com/paper/2306.06866/) (2023)
- [Improving Hyperparameter Optimization with Checkpointed Model Weights](https://mlsanity.com/paper/2406.18630/) (2024)
- [Sequential Conditional Transport on Probabilistic Graphs for Interpretable Counterfactual Fairness](https://mlsanity.com/paper/2408.03425/) (2024)

## Abstract (arXiv)

Optimal transport (OT) theory describes general principles to define and select, among many possible choices, the most efficient way to map a probability measure onto another. That theory has been mostly used to estimate, given a pair of source and target probability measures (μ, nu), a parameterized map Tθ that can efficiently map μ onto nu. In many applications, such as predicting cell responses to treatments, pairs of input/output data measures (μ, nu) that define optimal transport problems do not arise in isolation but are associated with a context c, as for instance a treatment when comparing populations of untreated and treated cells. To account for that context in OT estimation, we introduce CondOT, a multi-task approach to estimate a family of OT maps conditioned on a context variable, using several pairs of measures left(μi, nuiright) tagged with a context label ci. CondOT learns a global map Tθ conditioned on context that is not only expected to fit all labeled pairs in the dataset leftleft(ci,left(μi, nuiright)right)right, i.e., Tθleft(ciright) sharp μi ≈ nui, but should also generalize to produce meaningful maps Tθleft(c_text new right) when conditioned on unseen contexts c_text new . Our approach harnesses and provides a novel usage for partially input convex neural networks, for which we introduce a robust and efficient initialization strategy inspired by Gaussian approximations. We demonstrate the ability of CondOT to infer the effect of an arbitrary combination of genetic or therapeutic perturbations on single cells, using only observations of the effects of said perturbations separately.