Towards generalizable single-cell perturbation modeling via the Conditional Monge Gap

TL;DR

This work introduces the Conditional Monge Gap (CMonge), a global neural optimal transport map conditioned on covariates to predict single-cell perturbation responses across seen and unseen drugs and dosages. By jointly learning transport maps for multiple conditions, the method enables cross-task learning and robust out-of-distribution generalization, outperforming condition-specific and some state-of-the-art conditional approaches on scRNA-seq and multiplexed imaging data. The approach leverages a two-stage architecture (latent gene embeddings plus a context-conditioned transport map) and explores conditioning via MoA embeddings or RDKit fingerprints, showing strong performance for both in-distribution and OOD perturbations, including unseen drugs and combinations. The results suggest that conditional OT can bridge structure-based drug representations and observed perturbation effects at scale, offering a promising path for predicting responses to unseen treatments in single-cell data analysis and related domains.

Abstract

Learning the response of single-cells to various treatments offers great potential to enable targeted therapies. In this context, neural optimal transport (OT) has emerged as a principled methodological framework because it inherently accommodates the challenges of unpaired data induced by cell destruction during data acquisition. However, most existing OT approaches are incapable of conditioning on different treatment contexts (e.g., time, drug treatment, drug dosage, or cell type) and we still lack methods that unanimously show promising generalization performance to unseen treatments. Here, we propose the Conditional Monge Gap which learns OT maps conditionally on arbitrary covariates. We demonstrate its value in predicting single-cell perturbation responses conditional to one or multiple drugs, a drug dosage, or combinations thereof. We find that our conditional models achieve results comparable and sometimes even superior to the condition-specific state-of-the-art on scRNA-seq as well as multiplexed protein imaging data. Notably, by aggregating data across conditions we perform cross-task learning which unlocks remarkable generalization abilities to unseen drugs or drug dosages, widely outperforming other conditional models in capturing heterogeneity (i.e., higher moments) in the perturbed population. Finally, by scaling to hundreds of conditions and testing on unseen drugs, we narrow the gap between structure-based and effect-based drug representations, suggesting a promising path to the successful prediction of perturbation effects for unseen treatments.

Figures (16)

• Figure 1: The Conditional Monge Gap (CMonge) for single-cell perturbation response prediction. A) Existing methods learn local maps for each perturbation separately. B) We propose to model perturbation responses via a global estimator that can be conditioned on a potentially unseen condition $c_i$ at inference time. C) In-distribution results for drug-specific models, with or without dose context information. Boxes show the mean Maximum Mean Discrepancy (MMD) per drug on the SciPlex dataset for the highest dose for 9 drug-specific Monge Gap models and 9 Conditional Monge Gap models with conditional dosage information. Lines indicate the average performance of an identity model (lower bound, yellow) and a single Monge model per condition (36 models, upper bound, red). D) Out-of-distribution results for pan-drug models with drug and dose context information. Boxes show the mean MMD per drug on the SciPlex dataset for the highest dose for CMonge and chemCPA. Both conditional models only require drug structure (SMILES) and are trained and evaluated on all 187 drugs in the SciPlex dataset.
• Figure 2: Evaluation of perturbation prediction on the 4i dataset. Each point corresponds to a model trained on one of 35 treatments.
• Figure 3: Comparison of the different conditional and unconditional Monge methods in the ID setting using the $R^2$ metric. For each model it is indicated how many of the 36 drug-dose conditions were in the training set per individual model. A) In-distribution results on 9 selected drugs where CMonge is conditioned on drug and dose. See \ref{['table:drug_dose_id']} for mean and standard deviations. B) Results are grouped by drug and show the $R^2$ metric for both the dose and drug+dose conditioning. Each point represents the mean performance of the model over the four dosages. Error bars represent the 95% confidence interval. Note that CMonge-Dose and Monge-Dose are models per-drug (trained on only one drug), whereas CMonge-DrugDose are pan-condition models where one model is trained on all drugs and dosages. Monge is one model per condition (36 models in total). C) In-distribution results on all drugs in the SciPlex dataset where CMonge is conditioned on drug and dose. See \ref{['table:large_drug_dose_id']} for mean and standard deviations.
• Figure 4: Comparison of different Monge, Conditional Monge (CMonge), and the Identity models using the mean $R^2$ or Maximum Mean Discrepancy (MMD). Panels show overall performance, performance for single treatments, and combinatorial treatments. For each model it is indicated how many of the 33 treatments were in the training set per individual model. A) In distribution results showing $R^2$. The Monge model is trained as one model that sees all conditions, but is not aware of conditional information. The CMonge models are likewise trained on all conditions, but do get conditional information with the Mode-of-Action (MoA) or RDkit embedding. B) Out of distribution (OOD) results showing MMD. All models are trained in a leave-one-drug-out setting. This means that each boxplot shows the result over 33 drugs, from 33 models. The Monge model is also trained in a leave-one-drug-out setting, but does not incorporate conditional information.
• Figure 5: Out-of distribution results on the SciPlex dataset for dose and drug-dose contexts. For each model it is indicated how many of the 36 or 748 drug-dose conditions were in the training set per individual model. A) $R^2$ metric for the dose OOD setting on the selected 9 drugs. Horizontal lines indicate upper and lower bounds of performance. Results are split by dose and shown is the distribution over 9 drugs. B) OOD results for all drugs in the SciPlex dataset. Boxplots show the $R^2$ per drug, shown per dosage. C) Direct comparison for chemCPA and the CMonge-DrugDose-RDkit for the $R^2$, the Wasserstein distance and the MMD. Each dot is the average performance for a drug over all dosages. Points are colored by wether CMonge or RDkit achieves better performance. D) UMAPs for the OOD setting 'Abexinostat-10000' for models trained on the selected 9 drugs. Source, target and transport are taken from one training batch (n=512) and the grey background is the UMAP obtained from all 36 conditions (see Appendix \ref{['fig:extra_umaps']}).