# Distribution-Conditioned Transport

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 69 references that are in the map.
- First posted on arXiv 5 Mar 2026.

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

- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [StarGAN: Unified Generative Adversarial Networks for Multi-Domain Image-to-Image Translation](https://mlsanity.com/paper/1711.09020/) (2017)
- [Stochastic Interpolants: A Unifying Framework for Flows and Diffusions](https://mlsanity.com/paper/2303.08797/) (2023)
- [Improving and generalizing flow-based generative models with minibatch optimal transport](https://mlsanity.com/paper/2302.00482/) (2023)
- [Multimodal Unsupervised Image-to-Image Translation](https://mlsanity.com/paper/1804.04732/) (2018)
- [Wasserstein GAN](https://mlsanity.com/paper/1701.07875/) (2017)
- [Kernel Mean Embedding of Distributions: A Review and Beyond](https://mlsanity.com/paper/1605.09522/) (2016)
- [Discrete Flow Matching](https://mlsanity.com/paper/2407.15595/) (2024)
- [Learning Generative Models with Sinkhorn Divergences](https://mlsanity.com/paper/1706.00292/) (2017)
- [Unsupervised Machine Translation Using Monolingual Corpora Only](https://mlsanity.com/paper/1711.00043/) (2017)
- [Generative Moment Matching Networks](https://mlsanity.com/paper/1502.02761/) (2015)
- [Generalized Sliced Wasserstein Distances](https://mlsanity.com/paper/1902.00434/) (2019)
- [ProGen2: Exploring the Boundaries of Protein Language Models](https://mlsanity.com/paper/2206.13517/) (2022)
- [Unsupervised Neural Machine Translation](https://mlsanity.com/paper/1710.11041/) (2017)
- [Phrase-Based & Neural Unsupervised Machine Translation](https://mlsanity.com/paper/1804.07755/) (2018)
- [Generative Models and Model Criticism via Optimized Maximum Mean Discrepancy](https://mlsanity.com/paper/1611.04488/) (2016)
- [Simulation-free Schr"odinger bridges via score and flow matching](https://mlsanity.com/paper/2307.03672/) (2023)
- [Fisher Flow Matching for Generative Modeling over Discrete Data](https://mlsanity.com/paper/2405.14664/) (2024)
- [Universal Approximation of Functions on Sets](https://mlsanity.com/paper/2107.01959/) (2021)
- [Categorical Flow Matching on Statistical Manifolds](https://mlsanity.com/paper/2405.16441/) (2024)
- [Meta Flow Matching: Integrating Vector Fields on the Wasserstein Manifold](https://mlsanity.com/paper/2408.14608/) (2024)
- [Towards a mathematical theory of trajectory inference](https://mlsanity.com/paper/2102.09204/) (2021)
- [Set Norm and Equivariant Skip Connections: Putting the Deep in Deep Sets](https://mlsanity.com/paper/2206.11925/) (2022)
- [Multimarginal generative modeling with stochastic interpolants](https://mlsanity.com/paper/2310.03695/) (2023)
- [Wasserstein Wormhole: Scalable Optimal Transport Distance with Transformers](https://mlsanity.com/paper/2404.09411/) (2024)
- [Generative Distribution Embeddings: Lifting autoencoders to the space of distributions for multiscale representation learning](https://mlsanity.com/paper/2505.18150/) (2025)

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

- [Measure-to-measure Regression with Transformers](https://mlsanity.com/paper/2605.28075/) (2026)
- [Stability of Measure-to-Measure Transformers on Sub-Gaussian Data](https://mlsanity.com/paper/2610.07717/) (2026)

## Abstract (arXiv)

Learning a transport model that maps a source distribution to a target distribution is a canonical problem in machine learning, but scientific applications increasingly require models that can generalize to source and target distributions unseen during training. We introduce distribution-conditioned transport (DCT), a framework that conditions transport maps on learned embeddings of source and target distributions, enabling generalization to unseen distribution pairs. DCT also allows semi-supervised learning for distributional forecasting problems: because it learns from arbitrary distribution pairs, it can leverage distributions observed at only one condition to improve transport prediction. DCT is agnostic to the underlying transport mechanism, supporting models ranging from flow matching to distributional divergence-based models (e.g. Wasserstein, MMD). We demonstrate the practical performance benefits of DCT on synthetic benchmarks and four applications in biology: batch effect transfer in single-cell genomics, perturbation prediction from mass cytometry data, learning clonal transcriptional dynamics in hematopoiesis, and modeling T-cell receptor sequence evolution.