CloneOT: predicting clone-specific perturbation responses with lineage-supervised neural inverse optimal transport
Authors: Yanshuo Chen, Wei Chen, Heng Huang
Organizations: Department of Computer Science, University of Maryland, College Park, MD, 20742, USA. · Center of Bioinformatics and Computational Biology, College Park, MD, 20740, USA. · Department of Biostatistics, University of Pittsburgh, Pittsburgh, PA, 15261, USA. · Department of Pediatrics, UPMC Children’s Hospital of Pittsburgh, Pittsburgh, PA, 15224, USA. · Institute of Health Computing, University of Maryland, North Bethesda, MD, 20852, USA.
Predicting single-cell perturbation responses requires mapping between two unpaired populations of cells measured before and after the perturbation. Optimal transport (OT) provides a principled framework for constructing such mappings by minimizing the transport cost, and neural OT methods (e.g., CellOT) have been employed for this prediction task. However, these methods rely on a fixed transport cost, typically the squared Euclidean distance. This cost may connect cells with the wrong responses. Lineage tracing provides the information needed to learn the cost, as heritable barcodes link the cells of the same clone before and after the perturbation. Here, we propose \CloneOT, a neural inverse OT method that learns both the transport cost and the transport map from partially lineage-traced single-cell datasets. \CloneOT\ represents the cost as a neural network. It learns the cost by minimizing the extra cost that the lineage adds to the optimal transport, so that the cheapest transport keeps the cells of each clone together. Our experiments demonstrate that \CloneOT\ recovers the correct transport on 2D datasets, while the squared Euclidean cost fails. Moreover, \CloneOT\ outperforms five baselines on 14 lineage-tracing perturbation scenarios of cancer cells, for both seen and unseen clones. Furthermore, we show that \CloneOT\ reveals the clone-specific transcriptional programs of drug-tolerant persister cells and forecasts the fate of each clone from its cells before treatment. Our implementation is available at https://github.com/poseidonchan/CloneOT.
Figures & tables
Figure 1: Overview of CloneOT. a. Lineage barcodes identify, for a fraction of the clones (colored), which source and target cells belong to the same clone. b. The transport cost cθ(x,y)=hθ(φθ(x)−ψθ(y)) is a neural network with neural ODE flows φθ,ψθ and a symmetric input convex network hθ . Training alternates between solving the entropic OT problems without and with the lineage constraint by neural dual potentials and descending L(θ)=OTθA−OTθ , the extra cost of respecting the lineage. c. A neural map Tη amortizes the learned coupling and predicts the response of a new cell.
Figure 2: CloneOT on 2D datasets. Source points (dark blue), target points (pink) and predictions for new source points (light blue) are shown with transport rays for a quarter of the points. Red rays end on the wrong class. a. The T-shape dataset [ 23 ] . For 20% of the source points, the corresponding target points are given. b. The rotated-moons dataset [ 29 ] . The only correspondence is between each source moon and its rotated moon. Neural W2OT uses the squared Euclidean cost, and all methods are shown for training seed 0.
Figure 3: Performance benchmark with 80% labeled clones. a. Seen clones (IID setting). b. Unseen clones (OOD setting). Bars show means over the test clones of each study with 95% bootstrap confidence intervals. Stars denote two-sided Wilcoxon signed-rank tests against CloneOT over clones (* p<0.05 , ** p<0.01 , *** p<0.001 , **** p<10−4 , ns not significant), and red arrows mark methods significantly better than CloneOT. c. Clone-specific correlation for seen clones, CloneOT against each baseline with one point per scenario and p from a paired Wilcoxon signed-rank test over scenarios.
Figure 4: Lineage supervision improves CloneOT. a. r2 , ℓ2 , energy score and clone-specific correlation (Fig. 3 c) for seen clones against the fraction of labeled training clones. The value at 0% is CloneOT without lineage (open circles). b. The same for unseen clones. Points are means over the 14 scenarios with 95% confidence intervals over clones (over scenarios for the clone-specific correlation).
Figure 5: CloneOT predicts the clone-specific programs and fates of drug-tolerant persister cells (Watermelon, 80% labeled clones). a. Correlation across seen test clones between the predicted and the observed program score, pooled over the 12 scenarios (1356 clones summed over scenarios) with 95% bootstrap confidence intervals. Programs are the leading principal components of the post-treatment cells and are labeled by the Hallmark sets enriched at their two ends. Asterisks mark a gain from lineage (paired bootstrap, Benjamini–Hochberg q<0.05 ), and gray marks no gain. b. Forecast of the day-14 fate (cycling or non-cycling persisters) of day-0 clones from nine gene set scores by cross-validated logistic regression, with AUC averaged over 20 repetitions. c. Observed program score of the seen test clones in each third of the predicted score, with 95% confidence intervals.
Inferring dynamics from population snapshots is a fundamental challenge in machine learning and biology. In scRNA-sequencing (scRNA-seq), destructive measurements preclude direct tracking of individual cells across time, making trajectory inference underdetermined. Optimal Transport (OT) provides a principled framework for snapshot alignment, but a long-standing modeling question is which cost functions yield biologically meaningful couplings. Standard OT approaches rely on gene-expression distances, implicitly treating cells as independent points and neglecting structured cell-cell communication mediated by ligand-receptor signaling. We introduce CellBRIDGE (Cell-Based Regularized Interaction-Driven Gene Expression), which augments feature-based OT with a directed, typed interaction cost derived from ligand-receptor activity. By explicitly modeling cell-cell communication, CellBRIDGE improves cross-snapshot couplings and downstream trajectory estimates across synthetic and real scRNA-seq datasets relative to feature-only baselines. Notably, CellBRIDGE enables mechanistically interpretable in silico perturbations: on lung cancer data, silencing specific ligand-receptor pairs induces trajectory shifts that recapitulate expected effects of targeted pathway inhibition.
Silas Ruhrberg Estévez, Nicolas Huynh, Tennison Liu +4
DAMTP, University of Cambridge · Francis Crick Institute · University of Colorado Anschutz Medical Campus
Unbalanced optimal transport (UOT) provides a principled framework for modeling single-cell transitions and birth-death dynamics, but its high computational cost limits scalability to large-scale datasets. Although single-cell data often contain hierarchical annotations and known transition priors, existing UOT approximations rarely exploit this multiscale structure or prior knowledge. We introduce Multiscale Supervised Unbalanced Optimal Transport Flow Matching (MUST-FM), a simulation-free framework that scales UOT by leveraging hierarchical data structure. MUST-FM further supports an optional supervised formulation that incorporates transition priors, such as cell lineages, to guide the learning of displacement fields and mass variations. Experiments show that MUST-FM reduces computational overhead while achieving robust and biologically meaningful trajectory inference, enabling dynamic modeling of atlas-scale single-cell datasets.
Qiangwei Peng, Lezhi Chen, Peijie Zhou
College of Electrical Engineering Sichuan University, Sichuan, China · AI for Science Institute, Beijing, China
Single-cell perturbation atlases rarely measure every intervention in every cellular context: a query perturbation is often observed in one or more source contexts but missing in the recipient context where its effect is needed. Ignoring those measured responses discards query-specific experimental evidence, whereas copying or weakly calibrating them across contexts risks transferring the wrong signal. We propose PerturbMap, which predicts a missing recipient-context effect by combining a recipient-local low-rank base with accepted proposals that transport the same perturbation's measured source responses through source-to-recipient ridge experts fit on paired training perturbations, with proposal weights determined by route reliability estimated on validation anchors. On the Perturb-CITE-seq melanoma cohort, PerturbMap improves full-effect MSE by 4.1% over a recipient-local low-rank base and achieves lower MSE than FedAvg, zero-response, raw-copy, calibrated-copy, and identity-shuffled affine controls. It remains within 2.82×10−6 MSE of our centralized token-matched pooled reference, which uses a stronger training interface. A condition-mean specificity diagnostic shows the same direction: same-recipient top-10 counterpart retrieval by cosine increases from 74.5% for the low-rank base to 80.5% for PerturbMap.
Panpan Cui, Yiqi Liu, Wenhao Sun
School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences · Institute of Computing Technology, CAS; University of Chinese Academy of Sciences · 3Hong Kong University of Science and Technology