PaperlayerView exact revision on arXiv

Source-linked AI summary

Distribution-Conditioned Transport

Nic Fishman, Gokul Gowri, Paolo L. B. Fischer, Marinka Zitnik, Omar Abudayyeh, Jonathan Gootenberg

arXiv:2603.04736v1cs.LG

TL;DR

Transport models often struggle to generalize across the broad range of source and target distributions found in scientific datasets. DCT conditions transport maps on learned source and target distribution embeddings, achieving empirical performance gains on synthetic benchmarks and four biological applications.

  • Problem

    Scientific datasets require transport models that generalize across broad ranges of distributions, beyond methods designed for a single distribution pair.

  • Method

    DCT conditions transport maps on learned source and target distribution embeddings to support supervised, any-to-any, and semi-supervised transport.

  • Results

    DCT achieved empirical performance gains on synthetic benchmarks and four real-world biological transport problems, including held-out batch transfer and unseen-target generalization.

  • Takeaways & Limitations

    DCT offers a unified framework for transport across observed and unseen distributions, including settings with orphan marginals observed at only one timepoint.

  • Takeaways & Limitations

    Source-target-conditioned models can underperform in-distribution, possibly because they underfit under equal compute budgets and may require more training.

Abstract

from arXiv · show

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.

1. Introduction

Distribution-conditioned transport (DCT) addresses transport across multiscale scientific datasets whose condition-specific distributions may be sparse, heterogeneous, and unseen during training. It conditions transport maps on learned source and target distribution embeddings, supporting supervised, unsupervised, and semi-supervised transport settings.

  • Motivation: Scientific datasets contain many condition-specific distributions defined by donors, timepoints, perturbations, or clones, often with incomplete observations across populations.Some populations are observed at multiple timepoints, whereas others are orphan populations observed at only one timepoint.
  • Problem setting: These hierarchical datasets create transport problems involving technical effects, treatment responses, and distributional dynamics, beyond modeling a single fixed distribution pair.Each condition-specific distribution can be viewed as a task sampled from a shared metadistribution over probability measures.
  • Limitations: Existing approaches either solve transport among a fixed set of distributions or learn distribution-space autoencoders without directly providing continuous, unseen-distribution generalization.Multimarginal stochastic interpolants address K-to-K transport but cannot condition on a continuous space of distributions or generalize to unseen distributions.
  • DCT framework: DCT couples generative distribution embeddings with broad transport models and conditions on both source and target embeddings to enable any-to-any transport.The framework extends source-conditioned transport to transport between arbitrary source-target distribution pairs.
  • Problem settings: The framework supports supervised one-to-one forecasting, unsupervised any-to-any transport including unseen batches, and semi-supervised forecasting that uses orphan marginals.These settings cover clonal population forecasting, batch integration, and incomplete lineage-tracing time series.

2. Methods

DCT learns transport maps conditioned on distribution embeddings that are invariant to sample ordering and intended to capture distributional rather than sample-level signal. The framework supports supervised, any-to-any, and partially supervised transport while retaining plug-in consistency and the computational cost of the underlying transport model.

  • Distribution encoders: Distribution encoders map empirical distributions to fixed-dimensional embeddings intended to summarize distributional signal rather than individual samples.The encoder is designed to be distributionally invariant, including permutation invariance, so its output does not depend on sample ordering.
  • Distribution encoders: The encoder CLT enables training downstream objectives on moderate-sized subsamples while recovering population-level behavior up to O(m^-1/2) error.The paper formalizes this plug-in limit theory for DCT in Proposition 2.1 and Appendix A.
  • Supervised transport: In supervised transport, DCT conditions a transport map on source-distribution embeddings and jointly fits the encoder and transport model using the mechanism’s native loss.The loss may be based on flow matching, Sinkhorn/MMD divergence, or another generative transport model.
  • Unsupervised transport: For any-to-any transport, DCT conditions on both source and target distribution embeddings, allowing transport between arbitrary distributions rather than only a fixed set of trained pairs.The unsupervised formulation uses unstructured distributions and applies the same plug-in loss logic to train the chosen transport mechanism.
  • Unsupervised transport: Training on all distribution pairs does not incur quadratic cost per gradient step; computational efficiency matches the base transport model up to a distribution-encoder forward pass.This makes the any-to-any objective practical even when all distribution pairs are considered during training.
  • Partial supervision: In partial supervision, a lightweight model predicts target embeddings from source embeddings using available paired data, after which DCT generates counterfactual samples conditioned on the predicted target embedding.This adaptation leverages unpaired distributions while preserving a source-specific target-prediction objective at test time.

3. Related work

Related work provides distribution embeddings and diverse transport mechanisms that can be incorporated into distribution-conditioned maps. DCT is presented as mechanism-agnostic, while extending prior source-conditioned, multimarginal, and domain-specific transport approaches to broader generalization settings.

  • Distribution embeddings: Distribution embeddings represent probability measures as vectors, using kernel methods or learned representations that preserve Wasserstein distances.
  • Transport mechanisms: Transport maps can be built with adversarial methods, distributional-divergence minimization, normalizing flows, and flow matching.
  • Transport mechanisms: DCT is orthogonal to the transport mechanism: any model mapping source samples to target-matching samples can be conditioned on source or source–target distribution embeddings.The resulting maps satisfy the population-level limits in Eqn. (2) and Eqn. (4).
  • Prior conditional transport: DCT formalizes and generalizes source-conditioning beyond Meta Flow Matching and related cellular perturbation response prediction models.
  • Prior conditional transport: Unlike MMSI’s fixed, training-seen set of K distributions and barycenter-constrained paths, DCT imposes no such transport-path constraint.

4. Unsupervised Transport

This section evaluates source–target-conditioned DCT for unsupervised any-to-any transport on synthetic distributions and held-out scRNA-seq batch conditions. DCT generalizes better to unseen targets and outperforms K-to-K and established batch-integration alternatives in batch effect transfer.

  • Synthetic transport: Any-to-any DCT achieves uniformly lower error than K-to-K models across target means in the bivariate-normal error landscape.K-to-K predictions show Voronoi-like error patterns from nearest-training-distribution assignment.
  • Synthetic transport: The synthetic study compares conventional K-to-K models, which use discrete distribution labels and nearest-training-distribution assignment, with source–target-conditioned models that encode source and target sets.The comparison uses bivariate normal and Gaussian mixture distributions across varying numbers of distinct parameter sets.
  • Synthetic transport: DCT loses its small-K in-distribution disadvantage as K grows and achieves significantly lower error than K-to-K models on out-of-distribution targets.The K-to-K baseline cannot extrapolate beyond its training distributions, whereas DCT embeds targets directly.
  • Batch effect transfer: DCT enables batch effect transfer by predicting how cells from a source batch appear under a target batch’s technical effects.The model is evaluated on murine pancreas cells from 56 mice while holding out all mice from the old-age condition.
  • Batch effect transfer: DCT outperforms K-to-K models, scVI, and Harmony on held-out experimental batches, with predictions that more closely match target cell distributions.This qualitative improvement is reported for held-out donor pairs and summarized in Table 1 and Figure 4.

5. Semi-supervised Transport

DCT improves semi-supervised transport by conditioning on source and target distribution embeddings, enabling extrapolation beyond supervised support and using distributions observed at only one condition. Across Gaussian benchmarks and biological applications, source-target-conditioned models generally improve out-of-distribution or sparse-supervision performance, with some in-distribution tradeoffs.

  • Framework and evaluation: Source-target-conditioned DCTs improve semi-supervised performance across synthetic, mass cytometry, and protein-sequence data using SWD-, Energy-, and flow-matching transport maps.Evaluation uses MMD, SWD, and Energy metrics, with MMD reported in the main text.
  • Framework and evaluation: Semi-supervised models use ridge regression fit on the same data as supervised models, while oracle models condition on true target-distribution embeddings.The source-conditioned FM baseline is identical to MFM except for its population encoder.
  • Gaussian benchmarks: Beyond the supervised training support, semi-supervised models maintain substantially lower Gaussian benchmark error and approach oracle performance, whereas flow matching generalizes very well as an exception.Within support, supervised and semi-supervised methods perform comparably; evaluations span ∥µ∥∞≤5 after training on ∥µ∥∞≤2.5.
  • Mass cytometry: On PDO perturbation prediction, source-conditioned models perform better on held-out replicates from seen patients, while any-to-any models generalize better to unseen patients.The task predicts post-treatment distributions from control populations and treatment identities across 927 matched pairs.
  • Clonal dynamics: For clonal transcriptional dynamics, source-target-conditioned models significantly outperform source-conditioned baselines by incorporating orphan marginals from clones observed at only one timepoint.About 6 · 10^3 clones are measured, but only about 2 · 10^3 are profiled at multiple timepoints.
  • TCR repertoire forecasting: In TCR forecasting, the semi-supervised DFM model reduces energy distance by more than half, while the ProGen model learns nearly identical distribution embeddings.The ProGen embeddings have average cosine similarity 0.998, and the data include only 10 of 37 patients profiled longitudinally.

6. Conclusion … A.1.3. PLUG-IN LOSSES FOR CONDITIONED MODELS

The paper introduces distribution-conditioned transport (DCT) for generalizing transport maps to unseen source–target distributions, while its appendix formalizes the statistical foundations for minibatch training. These foundations establish empirical-measure factorization, encoder central limit theorems, and plug-in loss consistency for conditioned models.

  • 6. Conclusion: DCT conditions transport models on learned source and target distribution embeddings, supporting supervised, any-to-any, and semi-supervised learning with orphan marginals.The framework generalizes transport maps to unseen source and target distributions and formalizes existing approaches under a broader setup.
  • 6. Conclusion: Source–target-conditioned models outperform source-conditioned and K-to-K baselines out-of-distribution but can underperform in-distribution.The paper attributes this pattern potentially to underfitting and different scaling properties under equal compute budgets.
  • A. Theory: The appendix develops empirical-measure factorization and encoder CLTs, including plug-in CLTs for objectives conditioned on encoder states.It adapts plug-in loss theory while addressing DCT-specific joint source–target conditioning and minibatch arguments.
  • A.1.1. EMPIRICAL-MEASURE FACTORIZATION: Permutation-invariant distribution encoders factor through empirical measures, and proportional invariance yields a single functional across sample sizes.These results establish that encoders depend on samples through their empirical distributions and remain consistent under duplicated observations.
  • A.1.2. ENCODER CLT VIA THE FUNCTIONAL DELTA METHOD: If the encoder functional is Hadamard differentiable at the population distribution, the empirical encoder obeys a Gaussian central limit theorem via the functional delta method.The empirical process converges to a Brownian bridge, which the derivative maps to a Gaussian limit with covariance ΣP.
  • A.1.3. PLUG-IN LOSSES FOR CONDITIONED MODELS: The plug-in loss theory covers independent minibatches and paired source–target samples, with joint embedding CLTs accounting for their covariance structure.Independent minibatches produce block-diagonal covariance, whereas paired samples use the coupling-induced cross-covariance.
  • A.1.3. PLUG-IN LOSSES FOR CONDITIONED MODELS: For m/m′ →ρ ∈(0, ∞), the source–target-conditioned plug-in loss has the stated asymptotic limit, and fixed-point reuse does not change the limit.The same result holds when conditioning on any fixed finite subset of source minibatch points because their contribution is oP(1).
  • A.1.3. PLUG-IN LOSSES FOR CONDITIONED MODELS: Under local Lipschitz conditions, plug-in objectives are mean consistent; twice continuous differentiability with bounded Hessians additionally supports asymptotic unbiasedness.These conclusions follow from local expansions combined with the encoder CLTs and moment bounds.

A.2. When do source samples matter? … C.4. Flow matching generator (FM)

The appendices characterize when distribution-conditioned transport can recover source-dependent couplings, how its latent space behaves, and how its standard encoders and generators are constructed. They also define comparisons across unsupervised and supervised transport settings and detail the SWD, Energy, and FM generator families.

  • A.2. When do source samples matter?: Product-coupled distribution-level objectives can admit stochastic generators that match the target distribution while ignoring the source sample.An x-ignoring solution achieves zero loss when the generator can sample the target distribution from its embedding and auxiliary noise.
  • A.2. When do source samples matter?: Meaningful pointwise dependence requires sample-level pairing signals or additional structure selecting a coupling, while purely distributional objectives need not provide it.Possible remedies include deterministic generators, structured couplings such as optimal transport, or coupling- and geometry-selecting regularization.
  • A.3. How is the latent space structured?: Linear interpolation between distribution embeddings produces an STC trajectory that closely tracks closed-form Gaussian OT displacement interpolation in the MVN benchmark.This suggests meaningful latent geometry within the controlled family, but does not establish a guarantee.
  • B.1. Unsupervised Transport: Unsupervised experiments compare DCT source-target-conditioned models with one-hot K-to-K baselines, whose flow-matching versions correspond to MMSI restricted to simplicial edge paths.The K-to-K baseline supports comparison across transport mechanisms and was observed to be more stable than unrestricted MMSI edge-path alternatives.
  • B.2. Supervised Transport: Supervised experiments evaluate source-conditioned models against distributionally invariant-encoder baselines related to Meta Flow Matching.MFM instead uses a heuristic population encoder based on message passing over an input-population k-nearest-neighbor graph.
  • C. Standard distribution-conditioned transport maps: DCT couples a distribution encoder with a conditional generator, with SC generators using zsrc and STC or any-to-any generators using (zsrc, ztgt).Transported sample sets are generated from source and target sample-set embeddings, which need not be paired.

D. Standard evaluation metrics … E.5.1. UNSUPERVISED EVALUATION

The paper evaluates distributional transport with three sample-based distances and tests Gaussian and Gaussian-mixture experiments across supervised, unsupervised, and zero-shot settings. The experimental design varies distribution complexity, model conditioning, training objectives, and evaluation targets to assess generalization.

  • D. Standard evaluation metrics: Standard evaluation uses sliced Wasserstein distance, energy distance, and MMD with an RBF kernel, computed in data space unless otherwise stated.In TCR experiments, the metrics are computed in ESM2 embedding space.
  • D.1. Energy distance (MMD with energy kernel): Energy distance is estimated from finite samples, is equivalent to MMD with kernel k(x, y) = −∥x −y∥, and relates closely to the energy score.The energy score is described as a strictly proper scoring rule for probabilistic prediction.
  • D.2. Sliced 2-Wasserstein distance (SWD): SWD uses 100 random unit-sphere projections, matching sorted scalar projections and averaging squared differences across projections.The procedure estimates the standard sliced 2-Wasserstein distance.
  • D.3. MMD with RBF kernel (MMD-RBF): MMD is estimated empirically as a kernel two-sample statistic with an RBF kernel whose bandwidth uses the median heuristic on pooled samples.The squared MMD combines within-distribution kernel expectations and subtracts twice the cross-distribution expectation.
  • E. Gaussian and Gaussian Mixture Experiments; E.1. Data-Generating Process; E.1.1. MULTIVARIATE NORMAL (MVN) DISTRIBUTIONS: Gaussian experiments include 50,000 bivariate-normal sample sets, each with 100 i.i.d. samples, using randomly sampled means and inverse-Wishart covariances.For density-generalization studies, K takes values in {10, 100, 1,000, 10,000}, creating repeated or unique underlying distributions.
  • E.1.2. GAUSSIAN MIXTURE MODELS (GMM): The Gaussian-mixture setting uses three components, symmetric Dirichlet weights with α = 1.0, inverse-Wishart component covariances, and 1,000 samples per mixture.The larger sample set is intended to capture multimodal structure.
  • E.2. Supervised Gaussian Experiments; E.2.1. SUPERVISED MVN DATASET; E.2.2. SUPERVISED GMM DATASET: Supervised datasets use paired source-target distributions, with MVN shifts and GMM shift-plus-off-axis transformations that create unpaired targets and, for GMMs, bimodal structure.Targets are randomly permuted, removing point-wise correspondence while preserving distributional relatedness.

E.5.2. SEMI-SUPERVISED EVALUATION … F.2. Models

The paper evaluates semi-supervised distributional generalization, showing that an any-to-any encoder with a ridge target-embedding predictor performs well within training support and extrapolates beyond it. Additional Gaussian and MNIST-Colors experiments define the datasets, models, and training configurations used to assess scaling and transport quality.

  • E.5.2. SEMI-SUPERVISED EVALUATION: Semi-supervised training restricts paired data to source means with ∥µ∥∞≤2.5 but evaluates across the full range ∥µ∥∞≤5.Performance is partitioned by source-mean L∞ norm to measure degradation outside supervised support.
  • E.5.2. SEMI-SUPERVISED EVALUATION: The semi-supervised method pretrains an any-to-any encoder on 50,000 distributions, then fits ˆztgt = Wzsrc + b by ridge regression on paired data.At test time, the predicted target embedding conditions the learned transport.
  • E.5.2. SEMI-SUPERVISED EVALUATION: Comparisons include a source-only supervised model trained within restricted support and an oracle using the encoder’s true target embedding.The oracle provides an upper bound for semi-supervised performance.
  • E.6. Hyperparameter Summary: Experiments use shared random-seed and data-split protocols, while Gaussian configurations are consolidated in Table 6.The supervised setup restricts mean priors to [0, 2.5]2 and uses b = (1, 1)⊤, with GMM adding an off-axis displacement of magnitude 0.1.
  • E.7.1. UNSUPERVISED TRANSPORT: SCALING WITH K: The any-to-any encoder maintains consistent IID and OOD performance as K changes, whereas the K-to-K baseline memorizes small in-distribution sets but degrades on OOD targets.Gaussian full results report energy distance averaged across three random seeds, with lower values preferred.
  • E.7.2. SEMI-SUPERVISED GENERALIZATION: Semi-supervised performance matches or exceeds supervised performance on IID targets and has substantially lower OOD error outside restricted support.This indicates that unsupervised distributional structure supports extrapolation even when the latent predictor generalizes beyond paired training data.
  • E.7.2. SEMI-SUPERVISED GENERALIZATION: Supervised flow matching also shows low OOD error in the synthetic setting, but this anomaly may reflect simple linear transformations and low-dimensional extrapolation.The behavior is not replicated in real-world experiments.
  • F. MNIST-Colors image benchmark: MNIST-Colors tests whether Gaussian scaling behavior persists on image distributions, using colored 64-image sets, held-out colors, and context-conditioned U-Net transport models.Colors are applied to grayscale MNIST digits; models use 64-dimensional embeddings, concatenated source and target conditioning, and 200 training epochs with learning rate 2 × 10−4.

F.3. Metrics and results · G. Experimental details for scRNA-seq batch effect transfer

The section defines distributional and color metrics for evaluation and summarizes MNIST-Colors behavior across IID and OOD targets. It also introduces the scRNA-seq batch effect transfer experiment’s preprocessing, configurations, and evaluation protocol.

  • F.3. Metrics and results: Evaluation reports SWD and MMD-RBF on flattened pixels, plus Color MSE between generated and target samples’ per-image mean RGB vectors.Full results are provided in Table 11.
  • F.3. Metrics and results: GMM experiments compare MMD, SWD, and Energy distance across K unique distributions.These distributional distances are reported in Table 8.
  • F.3. Metrics and results: MVN semi-supervised experiments report distributional distances for IID targets with ∥µ∥∞≤2.5 and OOD targets with ∥µ∥∞> 2.5.The IID and OOD regimes are defined in Table 9.
  • F.3. Metrics and results: At small K, K-to-K baselines can compete on IID MNIST-Colors targets but degrade on OOD targets.This pattern mirrors the Gaussian experiments.
  • F.3. Metrics and results: Any-to-any models consistently improve OOD MNIST-Colors performance across generator families.The result is stated alongside the contrast with K-to-K baselines.
  • F.3. Metrics and results: As K increases, the performance gap between IID and OOD regimes narrows because less interpolation beyond the training set is needed.This interpretation is given for the MNIST-Colors results.
  • G. Experimental details for scRNA-seq batch effect transfer: The scRNA-seq batch effect transfer appendix describes its preprocessing pipeline, model configurations, and evaluation protocol.Generator families and metrics are defined in Appendices C and D.

G.1. Data preprocessing … H.4. Cotraining and stratified sampling

The appendix details preprocessing, sampling, model configurations, evaluation protocols, and Trellis drug-perturbation experiments, including scGen and CellOT baselines. It also examines cotraining and stratified sampling for improving latent prediction in any-to-any models.

  • G.1. Data preprocessing: The pancreas atlas uses the top 10 standardized principal components, with three two-year donor mice held out and 53 donors retained for training.Each sampled donor set contains n = 128 cells, and training forms random source-target donor pairs.
  • G.2. Model configurations: The study evaluates flow matching, energy/MMD regression, and sliced Wasserstein regression with distribution encoders versus one-hot donor-conditioned baselines.Held-out donors are assigned to the nearest training donor for one-hot encoding, while scVI and Harmony provide additional comparisons.
  • G.3. Full results across metrics: Predictions are evaluated over 10 independent source-target set pairs per test donor pair using energy distance, sliced Wasserstein distance, and RBF-kernel MMD.Metrics compare generated and ground-truth target distributions and are aggregated as mean ± standard error; results are reported in Tables 1 and 12.
  • H.1. Full results across metrics: The Trellis experiments report standard distributional metrics for IID replicate holdouts and OOD patient holdouts across the tested conditioning regimes.Tables 13 and 14 provide the main model and scGen metric summaries, respectively.
  • H.2. scGen baseline: scGen is trained in two stages: a VAE learns representations, then drug-specific latent difference vectors are computed from training profiles.The comparison uses the same IID and OOD holdouts as the main Trellis results.
  • H.3. CellOT baseline: CellOT performs comparably to scGen, and both baselines show a smaller IID–OOD performance gap than DCTs, likely because neither conditions at the patient level.CellOT trains a separate model for each drug perturbation; its results are reported in Table 15.
  • H.4. Cotraining and stratified sampling: Cotraining adds a latent-difference prediction loss, while stratified sampling ensures corresponding pairs comprise fraction η of training pairs; the selected parameters are α = 0.001 and η = 0.5.The distribution encoder otherwise leaves only about 10 −1000 datapoints for biological latent predictors, motivating this intervention.

H.5. Predictor ablation · H.6. Source-conditioned transport: encoder analysis

The predictor ablation benchmarks ridge regression, random forest regression, and a small MLP, finding that MLP generally performs best with modest gains over ridge. The encoder analysis compares the GDE-based source encoder with a k-NN graph-based encoder and finds similar performance.

  • H.5. Predictor ablation: H.5 Predictor ablation evaluates ridge regression, random forest regression, and a small MLP as latent predictor architectures.Ridge regression was used in Table 2, while the alternatives were random forest regression and a two-layer MLP.
  • H.5. Predictor ablation: The ablation is motivated by strong oracle performance, which indicates that the unsupervised any-to-any encoder and generator can support higher performance with a more accurate latent predictor.This motivation applies to both SC and STC models.
  • H.5. Predictor ablation: Table 17 reports distributional metrics for IID replicate holdout and OOD patient holdout settings.Results are reported as mean ± standard deviation.
  • H.5. Predictor ablation: The predictor search cross-validates ridge regularization, random forest estimator count and depth, and MLP hidden dimension and weight decay.The MLP uses ReLU activations and dropout regularization.
  • H.5. Predictor ablation: The MLP predictor consistently outperforms ridge regression and random forest regression across metrics and different generators.The gains relative to ridge regression are overall modest.
  • H.6. Source-conditioned transport: encoder analysis: H.6 Source-conditioned transport compares SC DCT using its GDE-based source encoder with the k-NN graph-based encoder from the original Meta Flow Matching work.The comparison tests dependence on the encoder used to embed the source.
  • H.6. Source-conditioned transport: encoder analysis: The k-NN-based and GDE-based encoders produce very similar performance for SC DCT.The results are reported in Table 18 and compared with Table 13.

I. Experimental details for lineage-traced scRNA-seq … J.6. Hyperparameter Summary

The appendices detail lineage-traced scRNA-seq and T-cell receptor forecasting experiments, including data construction, distribution-conditioned model configurations, evaluation protocols, and reproducibility settings. Results indicate qualitatively consistent scRNA-seq performance across metrics and substantially stronger DFM than ProGen2 performance for TCR forecasting, especially semi-supervisedly.

  • I. Experimental details for lineage-traced scRNA-seq: The lineage-traced scRNA-seq experiment uses hematopoietic differentiation data at days 2, 4, and 6, with clone-level consecutive-timepoint pairs and 1:1 clone train/test splitting.Clone splits are performed at the clone rather than cell level; 1256 clones are observed at consecutive timepoints.
  • I.1. Data preprocessing: ScRNA-seq preprocessing normalizes counts, applies log(x + 1), selects the top 10,000 highly variable genes, rescales genes, and performs downstream analysis in 50 principal components.Sets contain n = 100 cells with minimum nmin = 3, padding smaller clone-timepoint groups by replacement sampling.
  • I.2. Model configurations: The scRNA-seq study compares FM, Energy, and SWD generators under supervised SC, semi-supervised STC, and oracle STC regimes, using Energy, SWD, and MMD-RBF metrics.Semi-supervised STC predicts target embeddings with post-hoc ridge regression, while oracle STC uses the true target embedding.
  • I.3. Full results across metrics: Across the evaluated scRNA-seq models, results remain qualitatively similar regardless of whether Energy, SWD, or MMD-RBF is used.MMD-RBF is shown in Table 5.3, while all three metrics are reported in Table 19.
  • J. TCR Repertoire Forecasting Experiments: The TCR experiments forecast each held-out patient’s repertoire from time t to t+1 using COVID-19 TRB CDR3 data, with patient-level testing and unpaired single-timepoint training data for semi-supervised learning.The dataset includes 37 patients, including 10 profiled longitudinally; 3 of the 10 multi-timepoint patients are held out.
  • J.2. Embedding Architecture: TCR repertoires are encoded by frozen 8-million-parameter ESM2, producing 320-dimensional sequence embeddings that a transformer aggregator converts into 128-dimensional distribution embeddings.The aggregator uses two four-head self-attention layers, mean pooling, and a latent projection.
  • J.3. Generator Architectures: The TCR comparison uses ESM2 discrete flow matching with iterative refinement and ProGen2 with prefix-based autoregressive conditioning, while both receive source and target distribution embeddings.DFM conditions at input and output levels throughout the transformer and learns transitions through masked-token refinement; ProGen2 inserts one learned prefix token.
  • J.4. Evaluation Protocol: DFM substantially outperforms ProGen2, particularly semi-supervisedly, plausibly because simultaneous iterative refinement preserves source structure and distributes conditioning more thoroughly than causal prefix propagation.Evaluation generates 2048 sequences, compares them with actual next-timepoint sequences in ESM2 space, and reports Energy, SWD, and MMD-RBF with standard errors across test pairs.

J.6.1. RESULTS

Table 21 reports TCR repertoire prediction performance across conditioning strategies, showing meaningful conditioning for DFM but little conditioning effect for ProGen2 because of encoder collapse.

  • J.6.1. RESULTS: Table 21 reports TCR repertoire prediction performance across conditioning strategies as mean ± standard error.The results span all reported metrics.
  • J.6.1. RESULTS: DFM shows meaningful conditioning, with Oracle ≫ Predictor ≫ Within-patient performance.The ordering is reported across the conditioning strategies.
  • J.6.1. RESULTS: ProGen2 shows little conditioning effect because its encoder collapses, with all embeddings having cosine similarity > 0.998.The high embedding similarity is given as the stated explanation for the weak conditioning effect.
Loading 2603.04736v1…