Distribution-Conditioned Transport

Imagine you are a master chef who has learned to cook a specific dish (like a perfect steak) for a specific group of people (say, a group of vegetarians who only eat tofu). You know exactly how to transform their tofu into a steak-like experience.

Now, imagine you are asked to cook for a new group of people you've never met before, or perhaps you need to transform a completely different ingredient (like a carrot) into a steak. Traditional cooking methods (machine learning models) usually fail here because they were only trained on that one specific group and one specific ingredient. They get stuck and say, "I don't know how to do this!"

This paper introduces a new framework called Distribution-Conditioned Transport (DCT). Think of DCT as a universal translator for groups of data.

Here is how it works, broken down with simple analogies:

1. The Problem: The "Orphan" Data

In science (like biology), we often have data that comes in "groups" or "distributions."

The Scenario: Imagine you have photos of cells from 100 different patients. Some patients were scanned at the beginning of a treatment and again at the end (a "paired" set). But many patients were only scanned once (an "orphan" set).
The Old Way: Traditional AI models are like a student who only learns by looking at the "paired" students. If you ask them to predict what happens to an "orphan" patient, they can't do it because they've never seen that specific patient before. They are rigid.

2. The Solution: The "ID Card" for Groups

The authors realized that instead of teaching the AI to memorize every single patient, it should learn to recognize the essence or vibe of a group.

The Analogy: Imagine every group of cells (or data) gets an ID Card (called a "distribution embedding").
- This ID card doesn't describe one specific cell; it describes the whole crowd. It says things like, "This group is mostly young, energetic, and has a lot of red blood cells," or "This group is old, tired, and has high stress markers."
- The AI learns to read these ID cards.

3. The Magic: The "Universal Translator"

Once the AI has these ID cards, it builds a Universal Transport Map.

The Metaphor: Think of the AI as a travel agent.
- Old Travel Agent: "I can only book flights from New York to London because that's the only route I've ever practiced."
- DCT Travel Agent: "I don't care where you are or where you want to go. Just show me your ID card (Source) and the ID card of your destination (Target). I will instantly figure out the best flight path for any two cities, even ones I've never visited before."

Because the AI understands the essence of the groups (via the ID cards), it can generalize. It can take a group of cells from a patient it has never seen and predict how they will change under a new drug, simply by matching the "vibe" of the source group to the "vibe" of the target group.

4. The "Semi-Supervised" Superpower

The paper highlights a special trick called Semi-Supervised Learning.

The Scenario: You have 100 patients. Only 20 have "Before and After" photos. The other 80 only have "Before" photos.
The Old Way: Throw away the 80 "Before" photos because you can't learn from them without an "After" photo.
The DCT Way: The AI uses the 80 "Before" photos to learn more about the nature of the groups. It learns, "Oh, this type of group usually behaves like that type of group." It uses the "orphans" to get smarter, making the predictions for the 20 paired patients much more accurate.

Real-World Examples from the Paper

The authors tested this on four real biological problems:

Batch Effects: Imagine taking photos of cells in Lab A and Lab B. The lighting is different (technical noise). DCT acts like a photo filter that can take a photo from Lab A and make it look exactly like it was taken in Lab B, even if it's a new lab they've never seen.
Drug Prediction: Predicting how a patient's cells will react to a new drug. DCT can guess the reaction for a patient it has never met, based on the "ID card" of their cells.
Cell Evolution: Tracking how stem cells turn into blood cells over time, even when some cells are only observed at one point in time.
Immune System Tracking: Predicting how the immune system evolves after a virus infection.

Summary

Distribution-Conditioned Transport (DCT) is like giving an AI a universal translator that speaks the language of "groups" rather than just individual items.

Old AI: "I know how to turn this specific apple into a pie."
DCT AI: "I understand what makes an apple an apple and what makes a pie a pie. Give me any apple and any pie recipe, and I can tell you how to turn them into each other."

This allows scientists to make predictions about new, unseen situations using data that was previously too messy or incomplete to use.

1. Problem Statement

Modern scientific datasets, particularly in biology (e.g., single-cell genomics), exhibit a hierarchical multiscale structure. Data is often collected from multiple distinct conditions (donors, timepoints, perturbations, or clones), where each condition induces a specific probability distribution over cell states.

The Challenge: Standard transport models (e.g., optimal transport, flow matching) are typically trained on a single pair of source and target distributions. They struggle to generalize to unseen distribution pairs (e.g., a new donor or a new experimental batch) or to leverage "orphan" data (distributions observed at only one timepoint or condition).
The Gap: Existing methods like Multimarginal Stochastic Interpolants (MMSI) or Meta Flow Matching (MFM) either require a fixed set of known distributions, cannot handle unpaired marginals, or fail to generalize to continuous spaces of new distributions.

2. Methodology: Distribution-Conditioned Transport (DCT)

The authors propose DCT, a framework that conditions transport maps on learned embeddings of source and target distributions. This allows the model to learn a universal transport mechanism applicable to any pair of distributions within a metadistribution.

Core Components

Distribution Encoders ( $E$ ):
- A neural network that maps a set of samples $S_i$ (a distribution) to a fixed-dimensional vector embedding $z_i \in \mathbb{R}^d$ .
- Key Property: The encoder is distributionally invariant (permutation and proportionally invariant), meaning it depends only on the empirical measure of the data, not the order or specific sample size.
- Theoretical Guarantee: The encoder satisfies a Central Limit Theorem (CLT). As the sample size $m \to \infty$ , the embedding converges to a population-level embedding with Gaussian fluctuations. This justifies training on mini-batches without biasing the population-level objective.
Conditional Transport Maps ( $T$ ):
- A transport model that takes individual samples $x$ and conditions them on distribution embeddings.
- Source-Conditioned ( $T(x|z_{src})$ ): Used for supervised tasks where the target is implied by the source (e.g., perturbation prediction).
- Source-Target-Conditioned ( $T(x|z_{src}, z_{tgt})$ ): Used for unsupervised "any-to-any" transport, allowing mapping between arbitrary pairs of distributions.
Training Strategy:
- The framework supports Supervised, Unsupervised, and Semi-supervised learning.
- Semi-supervised capability: The model can be trained on "orphan marginals" (unpaired distributions observed at a single timepoint) by using the unsupervised "any-to-any" objective. At test time, a lightweight predictor (e.g., ridge regression) maps the source embedding to the expected target embedding, enabling specific forecasting tasks.
Model Agnosticism:
- DCT is agnostic to the underlying transport mechanism. The authors demonstrate it with:
  - Flow Matching (FM)
  - Sliced Wasserstein Distance (SWD)
  - Energy Distance / MMD
  - Discrete Flow Matching (for sequence data)

3. Key Contributions

Unified Framework: DCT unifies and generalizes previous approaches (like Meta Flow Matching and Multimarginal Interpolants) into a single framework capable of handling continuous spaces of distributions.
Generalization to Unseen Distributions: By conditioning on learned embeddings rather than discrete labels (one-hot encodings), DCT can generalize to source-target pairs never seen during training (Out-of-Distribution or OOD).
Leveraging Unpaired Data: The framework effectively utilizes "orphan marginals" (data observed at only one condition) to improve transport prediction, a common scenario in biological lineage tracing where not all clones are tracked over time.
Theoretical Foundation: The paper provides rigorous proofs regarding the CLT of distribution encoders and the consistency of plug-in losses, ensuring that training on finite samples converges to the correct population-level behavior.

4. Experimental Results

The authors evaluated DCT on synthetic benchmarks and four real-world biological applications.

A. Synthetic Benchmarks (Gaussian & GMM)

Any-to-Any Transport: DCT significantly outperformed "K-to-K" models (which treat distributions as discrete labels) on Out-of-Distribution (OOD) targets. While K-to-K models failed to extrapolate beyond training corners, DCT smoothly interpolated in the embedding space.
Semi-Supervised Learning: In settings with limited paired data, DCT trained with unpaired marginals achieved lower error on OOD targets compared to purely supervised models, approaching the performance of an "Oracle" (which knows the true target embedding).

B. Biological Applications

Batch Effect Transfer (scRNA-seq):
- Task: Transferring cells from a source batch to a target batch (including unseen donors).
- Result: DCT outperformed standard batch correction tools (scVI, Harmony) and K-to-K baselines in preserving biological structure while removing technical noise on held-out donors.
Perturbation Prediction (Mass Cytometry):
- Task: Predicting cell responses to drug perturbations across different patients.
- Result: DCT generalized better to unseen patients (OOD) compared to source-conditioned baselines and existing methods like scGen and CellOT.
Clonal Dynamics (Lineage Traced scRNA-seq):
- Task: Forecasting transcriptional dynamics of hematopoietic clones over time.
- Result: By incorporating "orphan" clones (observed at only one timepoint), the semi-supervised DCT significantly outperformed supervised baselines that could only use fully paired clones.
TCR Repertoire Forecasting:
- Task: Predicting the evolution of T-cell receptor sequences over time.
- Result: Using Discrete Flow Matching (DFM) with DCT, the model reduced prediction error by more than half compared to supervised baselines, demonstrating the ability to learn cross-patient evolutionary patterns.

5. Significance and Impact

Scientific Utility: DCT addresses a critical bottleneck in modern biology: the inability to model complex, hierarchical datasets where data is sparse, unpaired, or involves unseen conditions.
Flexibility: The framework is not tied to a specific generative model, allowing researchers to plug in state-of-the-art transport mechanisms (Flow Matching, GANs, etc.) while gaining the ability to generalize across distributions.
Data Efficiency: By enabling the use of unpaired "orphan" data, DCT maximizes the utility of expensive biological experiments where full longitudinal tracking is often impossible.
Theoretical Rigor: The work bridges the gap between empirical machine learning practices and statistical theory, providing guarantees on how distribution embeddings behave as sample sizes vary.

In summary, Distribution-Conditioned Transport provides a robust, theoretically grounded, and empirically superior framework for learning transport maps that generalize across the space of probability distributions, unlocking new capabilities for forecasting and analysis in complex scientific domains.