DRIFT: Disentangled Responsive-Invariant Flow Transport for Single-Cell Perturbation Prediction

Predicting how a cell responds to a perturbation is a central problem in modeling cellular biology, yet the response can be complex and the data are structurally incomplete: profiling destroys the cell, so the treated and untreated states of the same cell can never be jointly observed. Flow matching has been used to transport control cells to perturbed states, learning complex responses without prescribing their functional form, but the learned flow acts on the entire cell state, including the basal state, which confounds the perturbation's effect with pre-existing cell-to-cell variability. Disentangled and causal approaches avoid this confound by fixing a form for the perturbation's mechanism, e.g., an additive latent shift, a sparse mechanism shift, or an edit to a postulated structural graph, but assumed mechanism may not match the true biology. We address both limitations within a single framework: a variational encoder separates a cell's representation into an invariant block, capturing state the perturbation leaves untouched, and a responsive block, capturing state it changes, with the split enforced by information-theoretic invariance and conditional priors driven by distinct auxiliary variables; only the responsive block is then transported, via a conditional flow matching, jointly conditioned on the perturbation and the invariant state. Our method outperforms state-of-the-art methods on several benchmarks, demonstrating its effectiveness in predicting cell responses to perturbations.