CloneOT: predicting clone-specific perturbation responses with lineage-supervised neural inverse optimal transport

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 (textite.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.