Distribution-Conditioned Transport

Abstract: 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.

• The paper introduces a novel framework that leverages distribution encoders with conditional generative models to achieve robust transport between probability distributions.
• The methodology decouples embedding learning from the transport mapping, enabling efficient scaling and unbiased minibatch convergence with strong OOD generalization.
• Empirical results demonstrate lower out-of-distribution errors and improved batch effect correction, with applications in scRNA-seq, mass cytometry, and TCR evolution.

Distribution-Conditioned Transport: Framework, Numerical Results, and Implications

Formal Framework and Motivation

"Distribution-Conditioned Transport" (2603.04736) establishes a unifying framework for transport modeling between probability distributions that is explicitly designed for generalization to unseen source and target distributions. The approach leverages distribution encoders that generate latent representations of empirical sample sets, ensuring invariance with respect to both sample permutations and sizes. By auguring these embeddings into conditional generative models (including flow matching, divergence-based objectives, normalizing flows, and adversarial generators), the proposed framework enables three major operating modes: supervised (one-to-one), unsupervised (any-to-any), and semi-supervised transport, the latter incorporating orphan marginals—marginal distributions observed without paired counterparts.

The hierarchical structure of modern scientific datasets—such as those found in single-cell genomics, mass cytometry, lineage tracing, and TCR repertoire sequencing—often presents sparsity and incomplete pairing, making traditional paired distribution transport insufficient. The DCT framework is agnostic to the chosen transport mechanism and decouples embedding learning from the generative mapping, thus permitting broad applicability and efficient exploitation of available data.

Methodology

Distribution encoders $\mathcal{E}: S_i \to z_i \in \mathbb{R}^d$ are trained to ensure that $z_i$ captures the distributional structure of sample sets $S_i$ without being influenced by sampling noise. Invariance properties guarantee reliance only on empirical measures. For downstream transport objectives with sufficient smoothness, the encoder admits a central limit theorem, enabling unbiased minibatch training with convergence guarantees—critical for scalability to large datasets.

Two main forms of conditional transport maps are considered:

• Source-conditioned: $T(x | z_{\text{src}})$, used for supervised tasks with explicit source-target pairs.
• Source-target-conditioned: $T(x | z_{\text{src}}, z_{\text{tgt}})$, used for any-to-any and semi-supervised tasks, capitalizing on the full space of learned distributional embeddings.

Semi-supervised learning leverages regression models to predict target embeddings from source embeddings within the set of available paired distributions, supporting generalization beyond training pairs.

Empirical Results and Numerical Evidence

The paper provides comprehensive numerical evaluation across synthetic and real datasets. Notable findings include:

• Synthetic Gaussian and GMM Benchmarks: When scaling the number of distributions $K$, source-target-conditioned models consistently achieve lower out-of-distribution (OOD) error compared to K-to-K baselines, especially as $K$ increases. These baselines, which memorize discrete distribution identities, perform adequately on in-distribution targets for small $K$, but fail to interpolate for unseen targets. DCTs retain low error for both in-distribution and OOD targets, demonstrating superior generalization.
• scRNA-seq Batch Effect Transfer: On murine pancreas scRNA-seq data, DCT models outperform both traditional batch correction methods (e.g., scVI, Harmony) and K-to-K models in transferring cells across held-out experimental batches. The MMD distance for DCT is substantially lower (e.g., 0.1164 vs. 0.3160 for SWD) on held-out test distributions.
• Drug Perturbation Prediction (Mass Cytometry): For patient-specific prediction of chemotherapy response in colorectal organoids, DCT semi-supervised models show improved generalization to unseen patients. OOD patient-level MMD distance is reduced by more than 35% relative to source-conditioned baselines.
• Modeling Clonal Population Dynamics (scRNA-seq Lineage Tracing): Incorporation of orphan marginals via any-to-any pairing in DCT leads to significantly lower forecasting errors; e.g., semi-supervised models achieve MMD distances of 8.67 compared to 9.87 for supervised competitors.
• TCR Sequence Evolution Forecasting: Using sequence embeddings from pretrained ESM2, the discrete flow matching DCT substantially improves energy distance and MMD-RBF on longitudinal immune repertoire forecasting; semi-supervised models halve error rates relative to supervised approaches.

The numerical results consistently document that DCT's generalization properties yield robust performance across application domains, particularly in OOD settings where conventional methods struggle.

Theoretical and Practical Implications

On the theoretical front, DCT extends meta-learning concepts to the space of probability measures, providing convergence proofs, mean consistency, and unbiased plug-in loss estimates—all validated through the functional delta method and empirical measure factorization. The use of permutation/proportional invariance for encoders aligns with recent efforts in kernel mean embeddings, deep set architectures, and generative distribution embeddings.

Practically, the flexibility to condition transport on arbitrary distribution embeddings enables:

• Zero-shot generalization to unseen conditions, donors, or experimental batches
• Utilization of unpaired marginals for enhanced predictive modeling, reducing data waste
• Modular combination of transport mechanisms (flow matching, sliced Wasserstein regression, MMD, etc.) with learned distributional context
• Efficient scaling without quadratic computational overhead even when training on all pairwise combinations

For biological applications, this provides notable improvements in batch integration, perturbation forecasting, dynamic inference, and evolutionary modeling, often with direct transferability to other domains with hierarchical, sparsely paired distributional datasets.

Limitations and Prospects

While DCT achieves strong OOD performance, it may underperform on in-distribution targets relative to memorization-based K-to-K models under limited training regimes. The scaling properties and optimization landscape of source-target-conditioned models could require further adaptation (e.g., additional training iterations or regularization) for maximum efficacy.

Future directions include:

• Extension to higher-order multi-marginal transport and barycentric interpolation
• Incorporation of additional structure or regularization in the generative models to enforce desirable coupling and avoid degenerate solutions (x-ignoring modes)
• Application to trajectory inference, evolutionary modeling, and time-series data with complex dependence structures
• Exploration of alternative distribution embeddings (e.g., kernel methods, transformer-based embeddings) for improved representation learning

Conclusion

Distribution-Conditioned Transport (2603.04736) formalizes a robust, theoretically grounded framework for learning transport maps between probability distributions with strong out-of-distribution generalization. Through the integration of distributional encoders and conditional generative models, DCT supports supervised, unsupervised, and semi-supervised transport, efficiently utilizing all available data, including orphan marginals. The empirical results demonstrate substantial performance gains on synthetic and real biological datasets, particularly in tasks where generalization to unseen contexts is critical. This approach represents a significant expansion of meta-learning principles in generative modeling, with broad implications for both theoretical and applied machine learning.