Departures: Distributional Transport for Single-Cell Perturbation Prediction with Neural Schr\"odinger Bridges

TL;DR

This paper introduces a novel neural Schr"odinger Bridge approach for predicting single-cell perturbation outcomes, effectively handling unpaired data and capturing cellular heterogeneity with state-of-the-art results.

Contribution

It proposes a scalable Schr"odinger Bridge approximation using Minibatch-OT pairing for precise distribution alignment in single-cell perturbation prediction.

Findings

Achieves state-of-the-art performance on genetic and drug perturbation datasets.

Effectively models heterogeneous single-cell responses.

Provides a scalable method for distributional alignment in unpaired data.

Abstract

Predicting single-cell perturbation outcomes directly advances gene function analysis and facilitates drug candidate selection, making it a key driver of both basic and translational biomedical research. However, a major bottleneck in this task is the unpaired nature of single-cell data, as the same cell cannot be observed both before and after perturbation due to the destructive nature of sequencing. Although some neural generative transport models attempt to tackle unpaired single-cell perturbation data, they either lack explicit conditioning or depend on prior spaces for indirect distribution alignment, limiting precise perturbation modeling. In this work, we approximate Schr\"odinger Bridge (SB), which defines stochastic dynamic mappings recovering the entropy-regularized optimal transport (OT), to directly align the distributions of control and perturbed single-cell populations across different perturbation conditions. Unlike prior SB approximations that rely on bidirectional modeling to infer optimal source-target sample coupling, we leverage Minibatch-OT based pairing to avoid such bidirectional inference and the associated ill-posedness of defining the reverse process. This pairing directly guides bridge learning, yielding a scalable approximation to the SB. We approximate two SB models, one modeling discrete gene activation states and the other continuous expression distributions. Joint training enables accurate perturbation modeling and captures single-cell heterogeneity. Experiments on public genetic and drug perturbation datasets show that our model effectively captures heterogeneous single-cell responses and achieves state-of-the-art performance.