Departures: Distributional Transport for Single-Cell Perturbation Prediction with Neural Schrödinger Bridges

2025-11-17

AI Summary

Single-cell perturbation prediction faces a fundamental challenge: sequencing is destructive, yielding inherently unpaired data; existing methods either lack explicit conditional control or rely on prior spatial assumptions for indirect distribution alignment, limiting accurate modeling of heterogeneous cellular responses. To address this, we propose a conditional distribution transport framework based on the Neural Schrödinger Bridge (NSB), the first to jointly model discrete gene activation states and continuous expression-level changes. We introduce a Minibatch Optimal Transport (Minibatch-OT) pairing strategy to guide bridge learning, circumventing ill-posedness in bidirectional modeling while enabling scalable, high-fidelity probabilistic mapping at the population level. Evaluated across diverse genetic and pharmacological perturbation datasets, our method achieves state-of-the-art performance, significantly improving both predictive accuracy and biological interpretability of single-cell perturbation responses.

Technology Category

• Reasoning under Uncertainty: Relational Probabilistic Models
• Search and Optimization: Learning to Search
• Machine Learning: Probabilistic Circuits and Graphical Models

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• User Modeling, Personalization and Recommendation: Explainable and interpretable methods for personalization

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ödinger 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.

Problem

• Predicting single-cell perturbation outcomes for gene function analysis and drug selection
• Addressing unpaired data limitations in single-cell perturbation modeling
• Developing scalable Schrödinger Bridge approximations for distribution alignment

Innovation

• Schrödinger Bridge approximates optimal transport directly
• Minibatch-OT pairing avoids bidirectional inference complexity
• Joint discrete-continuous models capture single-cell heterogeneity