Distribution-Conditioned Transport

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

cs.LGMar 5, 2026

Distribution-Conditioned Transport

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

Organizations: Harvard University · MIT · Beth Israel Deaconess Medical Center · Harvard Medical School ·

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.

Figures & tables

Figure 1: Distribution-conditioned transport. A source distribution (teal) is pushed to a target distribution (purple) by a transport model $\mathcal{T}$ conditioned on source and target embeddings from $\mathcal{E}$ . The same interface supports supervised source-conditioned training, any-to-any source–target conditioning, and zero-shot evaluation on new admissible pairs drawn from $Q_{\text{joint}}$ .

Figure 2

Figure 3: Unsupervised transport generalization. The plot reports $\Delta$ MMD, defined as $K$ -to- $K$ error minus any-to-any DCT error, for IID and OOD targets across MVN/GMM benchmarks, generator families, and values of $K$ . Positive values indicate lower error for any-to-any conditioning.

Generator	$K$ -to- $K$	Any-to-any
scVI	0.6178 $\pm$ 0.0210
Harmony	0.2057 $\pm$ 0.0335
FM	0.2260 $\pm$ 0.0292	0.0839 $\pm$ 0.0089
SWD	0.2835 $\pm$ 0.0217	0.0928 $\pm$ 0.0061
Energy	0.1826 $\pm$ 0.0180	0.0428 $\pm$ 0.0036

Table 1: Batch transfer quality. Values are squared MMD (mean

\pm

std. err.,

\downarrow

) averaged across held-out batch pairs. Any-to-any DCT conditions on encoded source and target batch distributions;

K

-to-

K

uses nearest-neighbor training batches at test time.

Figure 4: Semi-supervised generalization. Supervised models are trained on paired source distributions with $\|\mu\|_{\infty}\leq 2.5$ (shaded) and evaluated across $\|\mu\|_{\infty}\leq 5$ . Semi-supervised DCT uses the any-to-any encoder with a predicted target embedding, while oracle STC uses the true target embedding. Results show MMD distance ( $\downarrow$ ) for multivariate normal (left) and Gaussian mixture (right) targets.

	IID			OOD
	Sup.	Semi	Oracle	Sup.	Semi	Oracle
scGen	.0219 $\pm$ .0057	–	–	.0265 $\pm$ .0022	–	–
CellOT	.0198 $\pm$ .0046	–	–	.0304 $\pm$ .0033	–	–
CellFlow	.0212 $\pm$ .0033	–	–	.0338 $\pm$ .0018	–	–
FM	.0101 $\pm$ .0023	.0106 $\pm$ .0025	.0030 $\pm$ .0003	.0227 $\pm$ .0006	.0208 $\pm$ .0018	.0053 $\pm$ .0013
SWD	.0064 $\pm$ .0022	.0094 $\pm$ .0018	.0010 $\pm$ .0000	.0259 $\pm$ .0008	.0197 $\pm$ .0022	.0026 $\pm$ .0009

Table 2: Trellis PDO perturbation prediction. MMD (mean

\pm

SD,

\downarrow

) on IID replicate-holdout and OOD patient-holdout splits.

	Supervised	Semi-sup.	Oracle	CT MSE
LineageOT [ 13 ]	0.2865 $\pm$ 0.0067	–	–	–
FM	0.3085 $\pm$ 0.0056	0.2369 $\pm$ 0.0052	0.1362 $\pm$ 0.0029	0.0213
SWD	0.2536 $\pm$ 0.0056	0.2249 $\pm$ 0.0053	0.1158 $\pm$ 0.0029	0.0193
Energy (MSE)	0.2524 $\pm$ 0.0055	0.2289 $\pm$ 0.0055	0.1180 $\pm$ 0.0031	0.0187
Energy (mix.)	–	0.1944 $\pm$ 0.0047	–	–

Table 3: Clonal distribution forecasting on Weinreb et al. [57] . MMD (mean

\pm

SE,

\downarrow

) across 628 set pairs from 609 held-out clones. Oracle uses true target embeddings. LineageOT is an external paired-data baseline. CT MSE reports semi-supervised cell-type abundance error across 626 clone–time transitions.

	Supervised	Semi-sup.	Oracle
ProGen2	0.0599 $\pm$ .0211	0.0520 $\pm$ .0372	0.0103 $\pm$ .0041
DFM	0.0587 $\pm$ .0216	0.0499 $\pm$ .0306	0.0110 $\pm$ .0035

Table 4: TCR repertoire forecasting. Values are MMD distance (mean

\pm

standard error,

\downarrow

) over three patient means. Oracle uses true target embeddings.

Appendix figures & tables32 assets

Supplementary material from the paper’s appendix.

Appendix

Target-latent evaluation	IID	OOD
Expectation predictor (MSE)	$0.0103\pm 0.0032$	$0.0191\pm 0.0049$
Target-aided best-of- $32$	$0.0037\pm 0.0009$	$0.0225\pm 0.0024$
Oracle STC	$0.0011\pm 0.0002$	$0.0022\pm 0.0002$

Appendix

Table 5: Trellis oracle-gap diagnostic. The Energy-generator STC model is evaluated with the standard predicted target embedding, a target-aided best-of-

32

diagnostic, and the oracle target embedding. Values are raw MMD-RBF (mean

\pm

standard deviation,

\downarrow

Target-latent evaluation	Energy
Expectation predictor (MSE)	$4.433\pm 0.110$
Mixture predictor	$3.614\pm 0.102$
Target-aided best-of- $32$	$3.380\pm 0.074$
Oracle STC	$2.332\pm 0.046$

Appendix

Table 6: Lineage-tracing oracle-gap diagnostic. The Energy-generator STC model is evaluated with two deployed predictors and two target-aided conditions. Every row uses a fixed 100-cell generated endpoint. Values are energy distance over 628 empirical pairs from 609 held-out clones (mean

\pm

standard error,

\downarrow

Generator	Dominant fitted matrix structure	Interpretation
FM	on diagonal ( $A\approx I$ ), $R^{2}>0.95$	approximately translational
Energy	off diagonal ( $A\approx P_{\mathrm{swap}}$ ), $R^{2}>0.95$	coherent off-diagonal coupling
SWD	off diagonal ( $A\approx P_{\mathrm{swap}}$ ), $R^{2}>0.95$	coherent off-diagonal coupling

Appendix

Table 7: Linear matrix-pairing diagnostic. On the MVN any-to-any benchmark (

K=10{,}000

, 200 epochs), we fit

\hat{x}_{\mathrm{tgt}}=Ax_{\mathrm{src}}+b

from source samples to transported samples. All three deterministic generators have strong affine fits: FM has an approximately identity matrix, while Energy and SWD have off-diagonal matrices. These fits describe source–output association, not optimality of transport cost.

Generator	Regularization	Selected matrix structure	Energy
FM	None	on diagonal	0.019
SWD	None	off diagonal	0.022
SWD	$\lambda=0.1$	on diagonal	0.028
SWD	$\lambda=0.1\to 0$	on diagonal	0.020
Energy	None	off diagonal	0.021
Energy	$\lambda=0.1$	on diagonal	0.042

Appendix

Table 8: Coupling regularization. On the MVN any-to-any benchmark (

K=10{,}000

, 200 epochs), the regularized Energy and SWD rows use the identity target

C_{0}(x)=x

to select an on-diagonal coupling. Energy distance (

\downarrow

, biased empirical estimator) measures marginal distribution quality.

Panel of Figure 5 — Figure 5: Latent-space interpolants on the MVN any-to-any benchmark. Left: STC trajectory obtained by sequentially pushing samples using intermediate latent codes along a linear interpolant between endpoint embeddings. Right: OT displacement interpolation (computed in closed form for Gaussians).

Dataset
Hyperparameter	MVN	GMM
Number of sample sets	50,000	50,000
Samples per set	100	1,000
Data dimension	2	2
Mean prior range	$[0,5]^{2}$	$[0,5]^{2}$
Covariance prior (inv-Wishart $\nu$ )	10	10

Appendix

Table 9: Gaussian benchmark hyperparameters. Dataset, encoder, transport-map, and training settings are reported for the MVN and GMM experiments; the same configurations are used for the unsupervised scaling and semi-supervised Gaussian benchmarks unless otherwise stated.

	$K=10$			$K=100$			$K=1,000$			$K=10,000$
Model	MMD	SWD	Energy	MMD	SWD	Energy	MMD	SWD	Energy	MMD	SWD	Energy
SWD ( $K$ )	0.001	0.036	0.001	0.003	0.041	0.002	0.006	0.083	0.010	0.016	0.104	0.016
SWD (Any)	0.000	0.026	0.000	0.001	0.029	0.000	0.001	0.031	0.000	0.001	0.031	0.000
Energy ( $K$ )	0.001	0.033	0.000	0.003	0.056	0.002	0.005	0.093	0.008	0.016	0.104	0.015
Energy (Any)	0.000	0.026	-0.001	0.001	0.030	0.000	0.001	0.031	0.000	0.001	0.030	0.000
FM ( $K$ )	0.009	0.074	0.006	0.010	0.079	0.008	0.013	0.095	0.011	0.026	0.133	0.021

Appendix

Table 10: Unsupervised MVN scaling across training distributions. We compare

K

-to-

K

baselines and any-to-any DCT models as the number of unique training distributions

K

increases. IID and OOD target sets are reported separately; values are MMD-RBF, SWD, and energy distance (

\downarrow

	$K=10$			$K=100$			$K=1,000$			$K=10,000$
Model	MMD	SWD	Energy	MMD	SWD	Energy	MMD	SWD	Energy	MMD	SWD	Energy
SWD ( $K$ )	0.001	0.108	0.000	0.002	0.125	0.002	0.003	0.149	0.008	0.006	0.189	0.018
SWD (Any)	0.001	0.066	-0.001	0.001	0.093	0.001	0.002	0.111	0.003	0.002	0.121	0.005
Energy ( $K$ )	0.001	0.121	0.001	0.001	0.129	0.001	0.002	0.151	0.005	0.005	0.201	0.016
Energy (Any)	0.000	0.082	-0.002	0.001	0.098	0.000	0.002	0.127	0.002	0.003	0.145	0.005
FM ( $K$ )	0.015	0.295	0.042	0.020	0.362	0.055	0.021	0.357	0.058	0.019	0.346	0.054

Appendix

Table 11: Unsupervised GMM scaling across training distributions. We compare

K

-to-

K

baselines and any-to-any DCT models as the number of unique training distributions

K

increases. IID and OOD target sets are reported separately; values are MMD-RBF, SWD, and energy distance (

\downarrow

Encoder	Generator	$K=10$	$K=100$	$K=1{,}000$	$K=10{,}000$
$K$ -to- $K$	FM	1.0926	0.1250	0.0506	0.0560
	Energy	1.1110	0.1276	0.0409	0.0343
	SWD	1.0310	0.1449	0.0405	0.0325
KME	FM	9.6780	9.2970	6.4750	6.9160
	Energy	7.8860	6.1360	6.1110	6.2600
	SWD	7.2240	5.6940	4.3590	3.7510

Appendix

Table 12: MVN any-to-any encoder ablation. OOD performance is reported for

K

-to-

K

, fixed KME, learned DeepSets, and Set Transformer encoders with matched transport maps. Values are

W_{2}^{2}

distance (

\downarrow

) across

K

unique training distributions.

Encoder	Generator	Sup.	Semi.	Ora.
KME	FM	0.0074	1.5982	1.5982
	Energy	0.0427	3.1122	3.1122
	SWD	0.0308	3.7641	3.7641
DeepSets	FM	0.0010	0.0111	0.0054
	Energy	0.0866	0.0389	0.0135
	SWD	0.0481	0.0304	0.0133

Appendix

Table 13: MVN semi-supervised encoder ablation. Supervised (Sup.), semi-supervised (Semi.), and oracle (Ora.) regimes are evaluated under alternative target-conditioning representations. Values are OOD

W_{2}^{2}

distance (

\downarrow

		IID			OOD
Generator	Method	MMD	SWD	Energy	MMD	SWD	Energy
Flow Matching	Supervised	0.000	0.003	-0.000	0.000	0.006	-0.000
	Semi-sup.	0.001	0.019	0.001	0.001	0.021	0.001
	Oracle	0.001	0.020	0.001	0.002	0.028	0.003
Energy	Supervised	0.000	0.003	-0.000	0.018	0.062	0.020
	Semi-sup.	0.000	0.021	0.000	0.002	0.033	0.002

Appendix

Table 14: Semi-supervised MVN prediction. Supervised rows train only on paired source distributions with

\|\mu\|_{\infty}\leq 2.5

, Semi-sup. rows use the any-to-any encoder with a predicted target embedding, and Oracle rows condition on the true target embedding. Values are distributional distances (

\downarrow

), split by source mean into IID points with

\|\mu\|_{\infty}\leq 2.5

and OOD points with

\|\mu\|_{\infty}>2.5

		IID			OOD
Generator	Method	MMD	SWD	Energy	MMD	SWD	Energy
Flow Matching	Supervised	0.000	0.015	0.000	0.001	0.022	0.001
	Semi-sup.	0.002	0.032	0.002	0.010	0.052	0.011
	Oracle	0.002	0.032	0.002	0.004	0.041	0.005
Energy	Supervised	0.000	0.008	0.000	0.031	0.094	0.037
	Semi-sup.	0.001	0.028	0.001	0.007	0.061	0.008

Appendix

Table 15: Semi-supervised GMM prediction. Supervised rows train only on paired source distributions with

\|\mu\|_{\infty}\leq 2.5

, Semi-sup. rows use the any-to-any encoder with a predicted target embedding, and Oracle rows condition on the true target embedding. Values are distributional distances (

\downarrow

), split by source mean into IID points with

\|\mu\|_{\infty}\leq 2.5

and OOD points with

\|\mu\|_{\infty}>2.5

Generator	Encoder	Energy $\downarrow$	SWD $\downarrow$	MMD-RBF $\downarrow$
scVI		0.9074 $\pm$ 0.0698	0.6643 $\pm$ 0.0566	0.6178 $\pm$ 0.0210
Harmony		0.2584 $\pm$ 0.0402	0.4999 $\pm$ 0.0355	0.2057 $\pm$ 0.0335
SWD	Any-to-any	0.1164 $\pm$ 0.0118	0.3051 $\pm$ 0.0298	0.0928 $\pm$ 0.0061
SWD	$K$ -to- $K$	0.3160 $\pm$ 0.0286	0.5164 $\pm$ 0.0226	0.2835 $\pm$ 0.0217
Energy	Any-to-any	0.0639 $\pm$ 0.0046	0.4248 $\pm$ 0.0257	0.0428 $\pm$ 0.0036
Energy	$K$ -to- $K$	0.4000 $\pm$ 0.0453	0.6222 $\pm$ 0.0333	0.1826 $\pm$ 0.0180

Appendix

Table 16: Full batch effect transfer metrics. Any-to-any DCT and

K

-to-

K

baselines are evaluated on held-out murine pancreas batch pairs against ground-truth target batch distributions. scVI and Harmony are included as external baselines; values are mean

\pm

standard error for energy distance, SWD, and MMD-RBF (

\downarrow

Figure 6: Batch effect transfer on held-out batches. Each row shows one source–target batch pair from the murine pancreas benchmark. Columns compare target cells, any-to-any DCT predictions, and $K$ -to- $K$ predictions over reference density contours from all batches, illustrating that DCT better matches the held-out target distributions.

Model	Energy $\downarrow$
DeepSets	$0.0639\pm 0.0046$
Set Transformer	$0.045\pm 0.005$
KME	$0.691\pm 0.071$

Appendix

Table 17: Batch-transfer encoder comparisons. The any-to-any Energy model is evaluated with the standard learned DeepSets encoder, a Set Transformer encoder, and a fixed KME representation. Values are energy distance (mean

\pm

standard error,

\downarrow

IID (Replicate Holdout)
Generator	Regime	Energy	SWD	MMD-RBF
SWD	Supervised	$0.0897\pm 0.0332$	$0.1266\pm 0.0256$	$0.0064\pm 0.0022$
	Semi-sup.	$0.1305\pm 0.0305$	$0.1524\pm 0.0185$	$0.0094\pm 0.0018$
	Oracle	$0.0154\pm 0.0005$	$0.0633\pm 0.0009$	$0.0010\pm 0.0000$
Energy	Supervised	$0.0779\pm 0.0290$	$0.1199\pm 0.0209$	$0.0054\pm 0.0019$
	Semi-sup.	$0.1442\pm 0.0520$	$0.1588\pm 0.0297$	$0.0103\pm 0.0032$

Appendix

Table 18: Full Trellis PDO perturbation prediction results. Generator families are compared across supervised SC, semi-supervised STC, and oracle STC regimes. Energy, SWD, and MMD-RBF are reported as mean

\pm

standard deviation (

\downarrow

) for IID replicate-holdout and OOD patient-holdout splits; Energy STC (mix.) evaluates the distributional latent predictor by mixing transports over

k=32

sampled target latents.

IID (Replicate Holdout)
Generator	Energy	SWD	MMD-RBF
SWD	$0.1191\pm 0.0297$	$0.1461\pm 0.0199$	$0.0087\pm 0.0019$
Energy	$0.1200\pm 0.0239$	$0.1461\pm 0.0159$	$0.0087\pm 0.0015$
FM	$0.1526\pm 0.0290$	$0.1646\pm 0.0140$	$0.0112\pm 0.0020$
OOD (Patient Holdout)
SWD	$0.6188\pm 0.1697$	$0.3381\pm 0.0484$	$0.0454\pm 0.0111$

Appendix

Table 19: Trellis evaluation excluding test source samples. Values are energy distance, SWD, and MMD-RBF (mean

\pm

standard deviation,

\downarrow

) for IID replicate-holdout and OOD patient-holdout splits.

IID (Replicate Holdout)
Generator	Predictor	Energy	SWD	MMD-RBF
SWD	Ridge	$0.1305\pm 0.0305$	$0.1524\pm 0.0185$	$0.0094\pm 0.0018$
	Random Forest	$0.1530\pm 0.0526$	$0.1598\pm 0.0295$	$0.0106\pm 0.0032$
	MLP	$0.1200\pm 0.0524$	$0.1415\pm 0.0295$	$0.0082\pm 0.0032$
Energy	Ridge	$0.1442\pm 0.0520$	$0.1589\pm 0.0297$	$0.0103\pm 0.0032$
	Random Forest	$0.1515\pm 0.0619$	$0.1619\pm 0.0351$	$0.0107\pm 0.0040$

Appendix

Table 20: Trellis latent point-predictor ablation. Ridge, random forest, and MLP predictors are used to map source embeddings to target embeddings for semi-supervised STC models. Metrics are mean

\pm

standard deviation (

\downarrow

) across SWD, Energy, and FM generators on IID replicate-holdout and OOD patient-holdout splits.

Metric	Regime	$k=1$	$k=2$	$k=4$	$k=8$	$k=16$	$k=32$	Oracle	MSE
IID (Replicate Holdout)
Energy	Best-of-k	$0.1176\pm 0.0310$	$0.0934\pm 0.0258$	$0.0766\pm 0.0209$	$0.0647\pm 0.0163$	$0.0580\pm 0.0158$	$0.0511\pm 0.0135$	$0.0169\pm 0.0026$	$0.1442\pm 0.0520$
Energy	Mixture-of-k	$0.1176\pm 0.0310$	$0.1057\pm 0.0300$	$0.0992\pm 0.0295$	$0.0960\pm 0.0281$	$0.0953\pm 0.0285$	$0.0940\pm 0.0279$	$0.0169\pm 0.0026$	$0.1442\pm 0.0520$
SWD	Best-of-k	$0.1480\pm 0.0203$	$0.1323\pm 0.0187$	$0.1207\pm 0.0167$	$0.1118\pm 0.0143$	$0.1059\pm 0.0141$	$0.1003\pm 0.0127$	$0.0659\pm 0.0037$	$0.1588\pm 0.0297$
SWD	Mixture-of-k	$0.1479\pm 0.0203$	$0.1405\pm 0.0206$	$0.1363\pm 0.0208$	$0.1341\pm 0.0203$	$0.1336\pm 0.0207$	$0.1329\pm 0.0205$	$0.0659\pm 0.0037$	$0.1588\pm 0.0297$
MMD-RBF	Best-of-k	$0.0086\pm 0.0020$	$0.0068\pm 0.0017$	$0.0055\pm 0.0013$	$0.0047\pm 0.0010$	$0.0042\pm 0.0010$	$0.0037\pm 0.0009$	$0.0011\pm 0.0002$	$0.0103\pm 0.0032$

Appendix

Table 21: Trellis distributional latent-predictor diagnostic. For the Energy-generator STC model, best-of-

k

rows use the held-out target only for analysis, while mixture-of-

k

rows are deployable mixtures over sampled target latents. Values are mean

\pm

standard deviation (

\downarrow

) across metrics for IID replicate-holdout and OOD patient-holdout splits; the MSE column is the expectation latent predictor used for the standard semi-supervised row.

Model	Split	SC / baseline	Semi-sup.	Oracle
scGen	IID	$0.0219\pm 0.0057$	–	–
	OOD	$0.0265\pm 0.0022$	–	–
CellOT	IID	$0.0198\pm 0.0046$	–	–
	OOD	$0.0304\pm 0.0033$	–	–
CellFlow	IID	$0.0212\pm 0.0033$	–	–
	OOD	$0.0338\pm 0.0018$	–	–

Appendix

Table 22: Trellis baseline and encoder summary. External perturbation baselines, the MFM-FM source-conditioned baseline, and DCT encoder comparisons are placed on the same raw MMD-RBF scale as the main text. Values are mean

\pm

standard deviation (

\downarrow

) on IID replicate-holdout and OOD patient-holdout splits; for DCT encoder rows, the Energy-generator run is shown. Dashes indicate regimes that are not applicable or were not run.

IID (Replicate Holdout)
Generator	Energy	SWD	MMD-RBF
SWD	$2.1424\pm 0.2856$	$0.6303\pm 0.0473$	$0.1542\pm 0.0199$
Energy	$1.7904\pm 0.1825$	$0.6489\pm 0.0323$	$0.1159\pm 0.0118$
FM	$0.3485\pm 0.0684$	$0.2446\pm 0.0293$	$0.0251\pm 0.0046$
OOD (Patient Holdout)
SWD	$2.3059\pm 0.3605$	$0.6524\pm 0.0532$	$0.1661\pm 0.0232$

Appendix

Table 23: Trellis KME encoder comparison. Semi-supervised STC models use a fixed kernel mean embedding representation in place of the standard learned set encoder. Values are energy distance, SWD, and MMD-RBF (mean

\pm

standard deviation,

\downarrow

) for IID replicate-holdout and OOD patient-holdout splits.

IID (Replicate Holdout)
Generator	Energy	SWD	MMD-RBF
SWD	$0.1220\pm 0.0275$	$0.1479\pm 0.0186$	$0.0089\pm 0.0017$
Energy	$0.1275\pm 0.0298$	$0.1520\pm 0.0200$	$0.0092\pm 0.0019$
FM	$0.1904\pm 0.0772$	$0.1754\pm 0.0336$	$0.0139\pm 0.0051$
OOD (Patient Holdout)
SWD	$0.2567\pm 0.0137$	$0.2080\pm 0.0096$	$0.0195\pm 0.0010$

Appendix

Table 24: Trellis Set Transformer encoder comparison. Semi-supervised STC models use a Set Transformer representation in place of the standard learned DeepSets encoder. Values are energy distance, SWD, and MMD-RBF (mean

\pm

standard deviation,

\downarrow

) for IID replicate-holdout and OOD patient-holdout splits.

IID (Replicate Holdout)
Generator	Energy	SWD	MMD-RBF
SWD	$0.0807\pm 0.0253$	$0.1238\pm 0.0176$	$0.0058\pm 0.0016$
Energy	$0.0836\pm 0.0230$	$0.1254\pm 0.0145$	$0.0059\pm 0.0014$
FM	$0.1235\pm 0.0194$	$0.1477\pm 0.0091$	$0.0088\pm 0.0012$
OOD (Patient Holdout)
SWD	$0.3524\pm 0.0299$	$0.2493\pm 0.0111$	$0.0268\pm 0.0018$

Appendix

Table 25: Trellis source-conditioned graph-encoder comparison. Source-conditioned models use the k-NN graph encoder from Meta Flow Matching rather than the standard DCT set encoder. Metrics compare SWD, Energy, and FM generators on IID replicate-holdout and OOD patient-holdout splits (mean

\pm

standard deviation,

\downarrow

Baseline	Split	Energy	SWD	MMD-RBF
scGen	IID	0.2924 $\pm$ 0.0881	0.2189 $\pm$ 0.0383	0.0219 $\pm$ 0.0057
	OOD	0.3278 $\pm$ 0.0345	0.2272 $\pm$ 0.0162	0.0265 $\pm$ 0.0022
CellOT	IID	0.2795 $\pm$ 0.0760	0.2155 $\pm$ 0.0366	0.0198 $\pm$ 0.0046
	OOD	0.3990 $\pm$ 0.0181	0.2499 $\pm$ 0.0047	0.0304 $\pm$ 0.0033
CellFlow	IID	0.2968 $\pm$ 0.0574	0.2305 $\pm$ 0.0292	0.0212 $\pm$ 0.0033
	OOD	0.4600 $\pm$ 0.0141	0.2828 $\pm$ 0.0102	0.0338 $\pm$ 0.0018

Appendix

Table 26: External Trellis perturbation baselines. scGen, CellOT, and CellFlow are evaluated on the same Trellis PDO perturbation splits used for DCT models. Energy, SWD, and MMD-RBF are reported as mean

\pm

standard deviation (

\downarrow

) for IID replicate-holdout and OOD patient-holdout splits.

IID (Replicate Holdout)
Generator	Energy	SWD	MMD-RBF
Energy	$0.1340\pm 0.0480$	$0.1495\pm 0.0269$	$0.0096\pm 0.0030$
OOD (Patient Holdout)
Energy	$0.2633\pm 0.0167$	$0.2126\pm 0.0086$	$0.0200\pm 0.0016$

Appendix

Table 27: Trellis semi-supervised paired-only ablation. The Energy-generator STC model is trained using only paired source–target examples. Values are energy distance, SWD, and MMD-RBF (mean

\pm

standard deviation,

\downarrow

) for IID replicate-holdout and OOD patient-holdout splits.

IID (Replicate Holdout)
Generator	Energy	SWD	MMD-RBF
Energy	$0.1285\pm 0.0313$	$0.1503\pm 0.0187$	$0.0092\pm 0.0019$
OOD (Patient Holdout)
Energy	$0.6243\pm 0.0747$	$0.3348\pm 0.0198$	$0.0457\pm 0.0058$

Appendix

Table 28: Trellis pairing policy ablation. Performance of semi-supervised models trained with pairing policy such that only control samples are used for source distributions and only perturbed samples are used for target distributions.

Generator	Regime	MMD-RBF $\downarrow$	Energy $\downarrow$	SWD $\downarrow$
LineageOT	External baseline	$0.2865\pm 0.0067$	$6.034\pm 0.164$	$1.932\pm 0.028$
FM	Supervised (SC)	$0.3085\pm 0.0056$	$7.281\pm 0.141$	$2.306\pm 0.025$
	Semi-supervised (STC, MSE)	$0.2369\pm 0.0052$	$5.028\pm 0.123$	$1.778\pm 0.023$
	Semi-supervised (STC, mix.)	$0.1961\pm 0.0042$	$4.061\pm 0.101$	$1.779\pm 0.023$
	Oracle STC	$0.1362\pm 0.0029$	$3.158\pm 0.071$	$1.448\pm 0.020$
SWD	Supervised (SC)	$0.2536\pm 0.0056$	$4.727\pm 0.116$	$1.642\pm 0.022$

Appendix

Table 29: Lineage-tracing forecasting metrics. Held-out clonal populations are evaluated across supervised SC and semi-supervised STC regimes, with LineageOT shown as an external paired-data baseline. Oracle STC uses true target embeddings. Values are mean

\pm

standard error over 628 empirical set pairs (

\downarrow

). MSE rows use the expectation latent predictor; mixture rows draw a fixed 100-cell endpoint from the equal-weight mixture over 32 sampled target embeddings.

Method	MSE $\downarrow$	TV $\downarrow$
DCT (FM)	$0.0213\pm 0.0012$	$0.270\pm 0.010$
DCT (SWD)	$0.0193\pm 0.0013$	$0.248\pm 0.010$
DCT (Energy)	$\mathbf{0.0187\pm 0.0012}$	$\mathbf{0.247\pm 0.009}$

Appendix

Table 30: Lineage cell-type-composition forecasting. Values are mean

\pm

standard error across 626 held-out clone–time transitions. Arrows mark the preferred direction. DCT uses predicted target embeddings.

Target embedding	Access	Energy
Expectation (MSE)	Source only	$4.433\pm 0.110$
Mixture of 32 samples	Source only	$3.614\pm 0.102$
Best of 32 samples	Held-out target	$3.380\pm 0.074$
True target embedding	Held-out target	$2.332\pm 0.046$

Appendix

Table 31: Lineage-tracing target-embedding diagnostics. Deployed predictors use only the source clone. Target-aided diagnostics access the held-out target. Values are energy distance over 628 empirical pairs from 609 held-out clones (mean

\pm

standard error,

\downarrow

Model	Energy
SC, paired only	$4.735\pm 0.118$
STC, paired only	$4.961\pm 0.121$
STC, with orphan marginals	$\mathbf{4.433\pm 0.110}$

Appendix

Table 32: Lineage-tracing orphan-marginal ablation. Energy-distance DCT models trained on paired clones only or with additional orphan marginals. Values are energy distance over 628 empirical pairs from 609 held-out clones (mean

\pm

standard error,

\downarrow

Dataset
Hyperparameter	DFM	ProGen2
Sequences per set (training)	512	256
Maximum sequence length	30	30
Test fraction (multi-TP patients)	0.3	0.3
Encoder
ESM2 backbone	esm2_t6_8M	esm2_t6_8M

Appendix

Table 33: TCR forecasting hyperparameters. Dataset, encoder, generator, training, and evaluation settings for DFM and ProGen2 in the held-out future-timepoint task.

Generator	Regime	Energy	SWD	Sinkhorn	MMD-RBF
DFM	Source-conditioned	$0.0694\pm 0.0267$	$0.0239\pm 0.0037$	$0.4114\pm 0.0140$	$0.0587\pm 0.0216$
	Predicted-target STC	$0.0569\pm 0.0342$	$0.0219\pm 0.0061$	$0.3406\pm 0.0889$	$0.0499\pm 0.0306$
	Oracle STC	$0.0125\pm 0.0037$	$0.0137\pm 0.0021$	$0.2396\pm 0.0482$	$0.0110\pm 0.0035$
ProGen2	Source-conditioned	$0.0703\pm 0.0264$	$0.0244\pm 0.0037$	$0.4125\pm 0.0113$	$0.0599\pm 0.0211$
	Predicted-target STC	$0.0612\pm 0.0439$	$0.0216\pm 0.0069$	$0.3261\pm 0.1017$	$0.0520\pm 0.0372$
	Oracle STC	$0.0121\pm 0.0047$	$0.0135\pm 0.0019$	$0.2192\pm 0.0538$	$0.0103\pm 0.0041$

Appendix

Table 34: TCR repertoire forecasting metrics. Every row uses the same 512 sampled source sequences, 512 target sequences, and 512 outputs generated one per source prompt in each transition. Values are mean

\pm

standard error across three equally weighted patient means (

\downarrow

). Oracle STC embeds the held-out target and is a target-aided diagnostic, not a feasible forecast.

OOD Generalization

Out-of-distribution generalization studies how machine-learning models perform when test data, domains, tasks, or environments differ from training.

See all 392 papers

May 18, 2026

Domain Transfer Becomes Identifiable via a Single Alignment

Domain transfer (DT) maps source to target distributions and supports tasks such as unsupervised image-to-image translation, single-cell analysis, and cross-platform medical imaging. However, DT is fundamentally ill-posed: push-forward mappings are generally non-identifiable, as measure-preserving automorphisms (MPAs) preserve marginals while altering cross-domain correspondences, leading to content-misaligned translation. Recent work shows that MPAs can be eliminated by jointly transferring multiple corresponding source/target conditional distributions, but supervision signals labeling such conditionals are not always available in practice. We develop an alternative route to DT identifiability. Under a structural sparsity condition on the Jacobian support pattern, we show that distribution matching together with a single paired anchor sample suffices to identify the ground-truth transfer -- requiring substantially less supervision than prior approaches. To enable practical high-dimensional learning, we further propose an efficient Jacobian sparsity regularizer based on randomized masked finite differences, yielding a scalable surrogate without explicit Jacobian evaluation. Empirical results on synthetic and real-world DT tasks validate the theory.

Sagar Shrestha, Subash Timilsina, Hoang-Son Nguyen +1

School of Electrical Engineering and Computer Science, Oregon State University, Corvallis, Oregon, USA.

Domain Adaptation

Ask PDF

Oct 4, 2026

TempoBridge: Source-Conditioned Flow Matching with Optimal Transport Couplings for Single-Cell Population Transitions

Destructive single-cell measurements provide unpaired population snapshots rather than observations of the same cells across conditions. Local cell states and transition requests may also be insufficient to distinguish responses across source populations. We introduce TempoBridge, a common source-conditioned transport formulation for temporal, genetic, and chemical population transitions. Source cells initialize latent transport and provide a fixed empirical population summary. The velocity field receives this summary alongside the evolving cell state, flow time, and a structured transition descriptor. Minibatch optimal transport (OT) supplies couplings only for conditional flow-matching training paths; inference requires neither target expression nor OT computation. On held-out donors, TempoBridge achieves an Energy distance of 0.129 versus 0.144 for scGen. Genetic mean-expression

L_2

error is 2.261 versus 3.156 for scGPT-scratch under Seen 2/2. On held-out compounds, condition-averaged drug-effect correlation is 0.598 versus 0.561 for the CellFlow adapter. Temporal ablations show higher mean distributional error after removing source context, replacing optimal transport with random pairing, or replacing flow matching with static residual regression. Together, these results demonstrate the predictive utility of a common source-conditioned transport formulation across held-out donors, gene combinations, and compounds.

Bowen Han, Lingbei Meng, Shihuan Luo +7

Universiti Sains Malaysia · The Chinese University of Hong Kong Shenzhen · Shenzhen Loop Area Institute +3

Single-Cell Perturbation Modeling Conditional Flow Matching Single-Cell RNA Sequencing

Ask PDF

Oct 16, 2026

Dynamic Neural Optimal Transport for Subset Alignment

We study optimal transport between distributions when only a subset of the target population is relevant to the source. The goal is to map the source distribution to a reweighted target distribution whose density ratio with respect to the full target lies in

[0, c]

. This constraint enables the transported distribution to select an appropriate target subpopulation while limiting how strongly any region of the target can be emphasized. We propose a dynamic neural optimal transport framework for this setting using two neural networks. The first parameterizes a time-dependent potential whose spatial gradient defines the velocity field and whose terminal value enforces the density-ratio constraint. The second parameterizes an interpolation between the source distribution and the learned terminal distribution. Unlike prior neural dynamic optimal transport methods, which typically assume a fixed terminal distribution, our formulation learns the terminal distribution within the support of the target. We demonstrate the method on positive-unlabeled learning, unpaired domain translation, and trajectory inference, where meaningful alignments exist between source and target distributions, but direct distribution matching is inappropriate because the target contains additional or irrelevant components

Austin J. Brockmeier, Bilal Riaz

Department of Electrical and Computer Engineering, University of Delaware, Newark, DE, 19711 · Department of Computer and Information Sciences, University of Delaware, Newark, DE, 19711

Optimal Transport

Ask PDF

Abstract

OOD Generalization

Explore similar work

Domain Transfer Becomes Identifiable via a Single Alignment

TempoBridge: Source-Conditioned Flow Matching with Optimal Transport Couplings for Single-Cell Population Transitions

Dynamic Neural Optimal Transport for Subset Alignment