Distribution-Conditioned Transport (DCT)

• DCT is a framework that uses learned embeddings to condition transport maps, enabling effective interpolation across unseen distribution pairs.
• It integrates transport techniques such as Sliced-Wasserstein, MMD, and flow matching to optimize probability measure transformations.
• Empirical results show DCT improves out-of-distribution performance in applications like batch-effect correction, T-cell receptor evolution, and decentralized systems.

Distribution-Conditioned Transport (DCT) is a machine learning and control framework designed for the interpolation and transformation of probability measures, with an emphasis on enabling generalization across previously unseen distribution pairs. Unlike classical one-to-one transport models that require retraining for each new source–target pair, DCT conditions transport mappings explicitly on learned representations (embeddings) of both the source and target distributions. This unifies transport, semi-supervised learning, and distributional forecasting within a single agnostic framework applicable to diverse settings including synthetic data, genomics, and decentralized multi-agent systems (Fishman et al., 5 Mar 2026, Lee, 2 Jan 2026).

1. Formal Problem Definition

Let $X$ denote a data space, typically $\mathbb{R}^d$, and $\mathcal{P}(X)$ the set of Borel probability measures over $X$. Classical optimal transport seeks a map $T: X \rightarrow X$ such that the pushforward $T_\# p = q$ for two distributions $p, q \in \mathcal{P}(X)$. When only empirical samples $S_p = \{x_i\}_{i=1}^m \sim p$ and $S_q = \{y_j\}_{j=1}^n \sim q$ are available, $T$ is commonly trained to minimize statistical divergences or transport losses between $\mathbb{R}^d$0 and $\mathbb{R}^d$1.

DCT overcomes the lack of generalization by introducing source and target embeddings $\mathbb{R}^d$2 via a permutation- and proportion-invariant encoder $\mathbb{R}^d$3. The framework then posits a universal transport map $\mathbb{R}^d$4. For any $\mathbb{R}^d$5, the transformation $\mathbb{R}^d$6 yields samples from the target when $\mathbb{R}^d$7:

$\mathbb{R}^d$8

DCT unifies supervised (paired), unsupervised (any-to-any), and semi-supervised (e.g., with orphan marginals) training regimes by specifying different sampling distributions over source–target pairs.

2. Distribution Embeddings: Theory and Implementation

The encoder $\mathbb{R}^d$9 mapping a sample set $\mathcal{P}(X)$0 to $\mathcal{P}(X)$1 enforces key invariance properties: permutation invariance and invariance to sample duplication. Thus, $\mathcal{P}(X)$2 necessarily factors through the empirical measure $\mathcal{P}(X)$3, $\mathcal{P}(X)$4 for some measurable function $\mathcal{P}(X)$5.

Under mild regularity conditions, a multivariate central limit theorem (CLT) holds for $\mathcal{P}(X)$6:

$\mathcal{P}(X)$7

as $\mathcal{P}(X)$8. This result ensures that mini-batch-based "plug-in" training using empirical samples closely approximates the infinite-sample objective, justifying stochastic optimization [(Fishman et al., 5 Mar 2026), Proposition 2.1]. In practice, $\mathcal{P}(X)$9 is constructed from a per-sample MLP, a sequence of mean-pooled update blocks, and a nonlinear projection (optionally $X$0-normalized). All architecture components maintain permutation equivariance.

3. Conditional Transport Map Formulations

DCT is agnostic to the form of the underlying transport mechanism. Three primary families are specified:

• Sliced-Wasserstein Regression: $X$1 is optimized to minimize the empirical sliced-Wasserstein distance $X$2 between $X$3 and $X$4.
• Energy (MMD) Regression: $X$5 is trained via the energy distance, equivalent to MMD with negative Euclidean kernel.
• Flow Matching (FM): $X$6 parameterizes a time-conditioned vector field $X$7 so that the flow matching loss

$X$8

is minimized for interpolated points $X$9. Sample-level couplings are maintained by integrating $T: X \rightarrow X$0 at inference.

For each mechanism, embeddings $T: X \rightarrow X$1 are fed into the transport generator, enabling efficient adaptation to arbitrary source/target distributions at test time.

4. Training Objectives and Optimization

DCT generalizes several training regimes:

• Supervised (one-to-one): Trained only on observed $T: X \rightarrow X$2 pairs, using losses $T: X \rightarrow X$3.
• Unsupervised (any-to-any): Randomly samples source–target indices $T: X \rightarrow X$4 from a collection $T: X \rightarrow X$5, minimizing $T: X \rightarrow X$6.
• Semi-supervised: Combines losses from both truly paired and unpaired distribution pairs, using a user-defined meta-pairing distribution $T: X \rightarrow X$7.

Plug-in estimation for the population-level loss is justified for general smooth transport divergences by the encoder CLT—minibatch estimates of the loss converge to their population counterparts with negligible bias for moderate $T: X \rightarrow X$8.

5. Algorithmic Structure and Practical Implementation

The core training paradigm is encapsulated by the following steps:

• For each iteration, sample a pair $T: X \rightarrow X$9 using $T_\# p = q$0;
• Subsample mini-batches $T_\# p = q$1;
• Compute $T_\# p = q$2, $T_\# p = q$3;
• Evaluate the loss $T_\# p = q$4 and perform backpropagation.

Unlike explicit K-to-K or domain-ID conditioning, DCT imposes negligible computational overhead—no O($T_\# p = q$5) complexity arises even when training on all $T_\# p = q$6 possible pairs. Hyperparameters (embedding dimension $T_\# p = q$7, number of mean-pool blocks $T_\# p = q$8, network depth, learning rates) are selected as in standard deep architectures [(Fishman et al., 5 Mar 2026), Sec. 5.2].

6. Generalization, Consistency, and Diagnostic Properties

DCT's generalization properties derive from conditioning on continuous distributional embeddings:

• Zero-shot/Out-of-Distribution (OOD) Transfer: Conditioned transport generalizes smoothly to unseen distribution pairs, whereas classical K-to-K (one-hot/domain-ID) baselines yield Voronoi-like failure regions in embedding space [(Fishman et al., 5 Mar 2026), Figs. 2–3].
• CLT-based Asymptotic Consistency: The plug-in loss estimated on minibatches converges to the infinite-sample objective with bias $T_\# p = q$9, variance governed by the embedding CLT (Proposition 2.1).
• Degenerate Conditioning Diagnostics: If a loss is purely distributional and not sample-coupled, the trivial solution $p, q \in \mathcal{P}(X)$0 may arise. Enforcing sample-level coupling, using flow-matching, or incorporating regularizers on $p, q \in \mathcal{P}(X)$1 is necessary to guarantee informative source-target transport [(Fishman et al., 5 Mar 2026), Proposition A.7].

7. Empirical Benchmarks and Applications

DCT has been validated in synthetic and biological settings:

• Synthetic Gaussian and GMM benchmarks: DCT achieves lower OOD interpolation error than K-to-K baselines, especially in any-to-any and semi-supervised settings, effectively leveraging orphan marginals [(Fishman et al., 5 Mar 2026), Figs. 2–5].
• Batch-effect correction in scRNA-seq: On held-out donors, DCT reduces MMD 3–8× compared to scVI, Harmony, and one-hot models, preserving biological structure.
• Mass cytometry perturbation prediction: DCT’s semi-supervised and any-to-any variants improve OOD MMD by 25–35%.
• Lineage tracing and T-cell receptor evolution: DCT improves generalization to unpaired or orphan data, consistently lowering MMD by 12–50% depending on the regime and dataset.

DCT serves as a universal wrapper: distributional embedding conditioning replaces domain-ID or one-hot conditioning, enabling robust any-to-any mapping, semi-supervised learning, and improved OOD performance across diverse transport objective choices (Wasserstein/MMD/flow). In multi-agent optimal transport settings, DCT-inspired decentralized local assignment and memory-corrected communication protocols allow scalable, robust distributional matching under realistic, intermittent communication (Lee, 2 Jan 2026).