PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Mustapha Bounoua, Giulio Franzese, Pietro Michiardi

arXiv:2609.35106v1cs.LG

TL;DR

Unpaired single-cell measurements prevent direct observation of a cell’s treatment outcome, and existing approaches can either confound perturbation with invariant cell variation or prescribe the perturbation mechanism. DRIFT disentangles invariant and responsive states and transports only the responsive state; it outperforms the strongest published baseline across the reported benchmark settings. The authors identify single-cell-line benchmarks as a current scope boundary.

  • Problem

    Destructive single-cell sequencing leaves control and treatment populations unpaired, while whole-state transport can confound perturbation effects with invariant cell characteristics and prescribed disentanglement mechanisms constrain the modeled dynamics.

    Source: §1 · pp. 1–2

  • Method

    DRIFT uses conditional priors and an information-theoretic invariance constraint to separate invariant and responsive latent blocks, then conditionally transports only the responsive block.

    Source: §1 · p. 2, §4 · p. 9

  • Results

    DRIFT outperforms the strongest published baseline across the Norman additive and holdout splits, Replogle RPE1 screen, and ComboSciPlex dataset, spanning genetic and chemical perturbations and multiple metrics.

    Source: §4 · p. 9

  • Takeaways & Limitations

    The results suggest that separating invariant and responsive components reduces confounding from cell-level variation and yields a more faithful representation of perturbation-induced dynamics.

    Source: §4 · p. 9

    Current benchmarks are limited to single cell lines, so generalization of the invariant–responsive decomposition across diverse cellular backgrounds remains to be tested.

    Source: §4 · p. 9

Abstract

from arXiv · show

Predicting cellular responses to perturbations is a central problem in cellular biology, with broad applications in systems biology and drug discovery. This task is challenging because cellular responses can be complex and cell-state dependent, intrinsic cell-to-cell variability can be confounded with perturbation effects, and destructive single-cell RNA sequencing precludes paired measurements of the same cell before and after treatment. Flow matching transports control cells to perturbed states flexibly, but acting on the full cell state can confound perturbation effects with pre-existing cell-to-cell variability. Disentangled approaches separate responsive from invariant components, but model perturbations through prescribed mechanisms, such as latent shifts or graph edits, limiting their flexibility. We address both limitations in a unified framework. A variational encoder disentangles each cell into an invariant block, capturing state unaffected by the perturbation, and a responsive block, capturing state it changes, through conditional priors and an information-theoretic invariance constraint. Conditional flow matching transports only the responsive block, conditioned on the perturbation and invariant state, yielding a flexible, data-driven model of perturbation effects without confounding pre-existing variability. Across several benchmarks, our method outperforms the strongest published method in settings involving combinatorial and unseen perturbation prediction.

1 INTRODUCTION

Perturbation prediction must bridge unpaired control and treatment populations, while avoiding confounding cell variation with perturbation effects. DRIFT separates invariant and responsive cell components, transports only the responsive component, and reports better benchmark performance than the strongest published method.

1 INTRODUCTION

  • Single-cell sequencing destroys measured cells, leaving unpaired control and treatment populations and preventing direct measurement of per-cell causal effects.Methods therefore rely on assumptions to bridge observed data and the quantity of interest.Source: §1 · p. 1
  • Whole-state transport can confound perturbation effects with invariant cell characteristics, while prescribed disentanglement mechanisms may not generalize to incompatible biological programs.DRIFT instead transports the responsive part while leaving invariant factors fixed.Source: §1 · p. 2
  • DRIFT combines auxiliary-variable conditional priors and an information-theoretic invariance constraint with conditional flow transport of the responsive component.The flow is conditioned on the perturbation and invariant state.Source: §1 · p. 2
  • DRIFT outperforms the strongest published method on every dataset tested.The contribution summary reports state-of-the-art results across multiple benchmarks.Source: §1 · p. 2

2 METHOD

DRIFT learns an invariant and a responsive latent block, then uses a conditional flow to transport only the responsive block. At prediction time, it retains a control cell’s invariant state and decodes it with the generated response.

2.1 PROBLEM STATEMENT AND GENERATIVE MODEL

  • The model predicts a target perturbation’s cell-profile distribution from observed controls, including when the target perturbation may be absent from training.The populations are unpaired, and control and perturbed cells are not assumed to share observed covariates.Source: §2.1 · p. 3
  • The generative model draws invariant state from assay covariates, responsive state from the perturbation and invariant state, and expression from both latent blocks.The invariant block is intended to be perturbation-independent, while the responsive mechanism can depend on basal state.Source: Figure 2 · p. 4, §2.1 · p. 4

2.2 STAGE 1: DISENTANGLING THE VARIANT AND INVARIANT BLOCKS

  • Stage 1 uses auxiliary-variable conditional priors and an information-theoretic penalty on dependence between the invariant block and perturbation to learn the two-block representation.The prior for each block depends on one observed variable; the encoder and decoder are trained to reconstruct the input from both blocks.Source: §2.2 · pp. 4–5
  • Stage 1 delays KL and invariance regularization during warm-up to guard against posterior collapse and a trivial invariant representation.The encoder, decoder, priors, and auxiliary predictor are then trained jointly.Source: §2.2 · p. 5

2.3 STAGE 2: TRANSPORTING THE RESPONSIVE BLOCK

  • With the encoder frozen, Stage 2 learns a conditional flow that maps control responsive latents to perturbed responsive latents while keeping the invariant state fixed.Conditional flow matching fits a velocity field from straight-line source-to-target velocities without simulating the ODE during training.Source: §2.3 · pp. 5–6
  • Because experiments do not reveal control-to-perturbed cell pairs, minibatch entropic optimal transport supplies approximate training pairs for the velocity-field regression.The pairing plan is distinct from the learned flow; its invariant-block cost favors controls with similar invariant states.Source: §2.3 · p. 6

2.4 INFERENCE

  • Inference encodes an observed control cell, retains its invariant state, generates a responsive state conditioned on that state and the target perturbation, and decodes the pair.The prediction is well posed when the control invariant state lies within the support seen by the field under the target perturbation.Source: §2.4 · p. 6

3 EXPERIMENTS

DRIFT is evaluated on combinatorial and zero-shot perturbation benchmarks using correlation, directional-agreement, discriminability, and pointwise-error metrics. It outperforms the strongest baseline across the reported Norman, Replogle RPE1, and ComboSciPlex results, while ablations identify disentanglement as the largest contributor to performance gains.

3 EXPERIMENTS

  • The benchmarks cover combinatorial and zero-shot generalization, including Norman gene pairs and unseen single-gene perturbations in Replogle.The comparisons include foundation, graph-based, generative, and statistical methods; metrics cover correlation, directional agreement, discriminability, and pointwise error.Source: §3 · pp. 6–7, §3.1 · p. 7

3.1 RESULTS: NORMAN

  • On the Norman additive split, DRIFT beats scBIG on all nine metrics; ρ∆D improves by 8.3%, PDS by 7.2%, MSE by 37.2%, and L2 by 20.4%.The ρ∆D gain concerns the 20 most affected genes, while the PDS gain indicates clearer identification of the perturbation condition.Source: Table 1 · p. 7, §3.1 · p. 7
  • On Norman holdout singles, DRIFT beats scBIG on all nine metrics, with ρ∆D rising by 7.8% and MSE improving by 4.6%.Nine of the 26 test conditions include a gene absent from training.Source: §3.1 · p. 7, Table 2 · p. 8
  • On Norman holdout doubles, DRIFT beats scBIG on every metric by wider margins than on singles; ρ∆D improves by 6.9%, ACC∆D reaches 0.9873, and PDS improves by 8.5%.MSE improves by 28.9% and L2 by 15.8%.Source: §3.1 · p. 7, Table 2 · p. 8

3.2 RESULTS: REPLOGLE RPE1

  • On Replogle RPE1, DRIFT leads on all nine metrics, with its largest gain in DES at +36.5%.All 384 test conditions target single genes absent from training, making this a fully zero-shot setting.Source: Table 3 · p. 8, §3.2 · p. 8

3.3 RESULTS: COMBOSCIPLEX

  • On ComboSciPlex, DRIFT leads on five of six metrics, reduces pointwise errors by 23–33% versus scDFM, and raises DE-Spearman ρ from 0.829 to 0.890.It raises ρ∆ from 0.893 to 0.947; in DS, it ranks ahead of scDFM but slightly below CPA.Source: Table 4 · p. 9, §3.3 · p. 9

3.4 ABLATIONS

  • Ablations attribute the largest performance gains to disentanglement, with the invariance penalty adding consistent complementary gains.Conditioning regularization is especially valuable for entirely unseen test perturbations, but less consequential for combinatorial generalization using genes observed during training.Source: §3.4 · p. 9

4 CONCLUSION

DRIFT separates invariant and perturbation-responsive cell state, then transports only the responsive component. It outperforms the strongest published baseline across the reported benchmarks, while current benchmarks remain limited to single cell lines.

4 CONCLUSION

  • DRIFT uses a two-stage model: a variational autoencoder learns invariant and responsive latent blocks, then a conditional flow transports only the responsive block.The encoder uses conditional priors and an information-theoretic invariance constraint; the flow learns a velocity field from data.Source: §4 · p. 9
  • Across the Norman additive and holdout splits, Replogle RPE1, and ComboSciPlex, DRIFT outperforms the strongest published baseline.The authors suggest that separating invariant and responsive components reduces confounding from cell-level variation and yields a more faithful representation of perturbation-induced dynamics.Source: §4 · p. 9
  • Current benchmarks are limited to single cell lines, and testing across diverse cellular backgrounds would assess whether the decomposition generalizes.The authors also note that temporal measurements could help interpret the flow’s time coordinate and connect endpoint mappings into trajectory models.Source: §4 · p. 9

AI USE STATEMENT

  • The authors used generative AI tools for formatting, figure and table generation, code support, GPU-cluster interaction, and artifact management, while taking responsibility for the final content.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

  • The datasets are publicly available through referenced-work repositories, and code and reproduction scripts are available on GitHub.Source: REPRODUCIBILITY STATEMENT · p. 10

M. Lotfollahi, Anna Klimovskaia Susmelj, C. De Donno, Leon Hetzel, Yuge Ji, Ignacio L Ibarra,

The listed references include work on single-cell perturbation prediction, perturbation screens, and computational representations of cells and genes.

References

A METRIC DEFINITIONS

The appendix describes metrics for pointwise prediction, response fidelity, discrimination, and distributional fidelity.

Point-prediction accuracy

  • The point-prediction definitions use the 20 genes with the largest absolute observed expression shifts from the training-control mean as D_p.The appendix states that it follows Ruan et al. (2026) for the definitions and released implementation.Source: A METRIC DEFINITIONS · p. 16

Biological response fidelity

  • DE-Spearman measures rank correlation between predicted and observed log fold changes for statistically significant differentially expressed genes.Source: A METRIC DEFINITIONS · p. 16

Discriminative power

  • PDS ranks observed perturbation profiles by their L1 distance to the predicted profile, excluding genes targeted by the perturbation.Source: A METRIC DEFINITIONS · p. 16

Distributional metrics

  • E-Dist and Wasserstein compare predicted and observed cell populations, while Discr. Cos measures the direction, but not magnitude, of the mean response shift.The population metrics are evaluated on cells projected onto 50 principal components.Source: A METRIC DEFINITIONS · p. 16

ComboSciPlex

B IMPLEMENTATION

The implementation uses benchmark-specific data splits and a two-stage latent model, with perturbation encodings and conditional flow transport used to generate predictions.

Datasets

  • Norman tests combinatorial and zero-shot generalization via 32 held-out gene pairs and 26 held-out conditions; Replogle holds out 384 unseen-gene single-perturbation conditions, while ComboSciPlex evaluates 17 compounds across 31 conditions (24 train, 7 test).For ComboSciPlex, evaluation uses 1,000 test-selected genes and generates 128 cells per condition.Source: B IMPLEMENTATION · p. 16

Stage 1: the disentangling VAE

  • The stage-1 encoder maps expression, perturbation, and covariates to invariant and responsive latent blocks, with separate conditional priors for the two blocks.Covariates include sequencing run, depth, detected genes, and mitochondrial and ribosomal content; ComboSciPlex uses no covariates because its arrayed well label is not independent of perturbation.Source: B IMPLEMENTATION · p. 17

Perturbation features

  • Genetic perturbations use protein embeddings, protein-association neighborhoods, and pathway membership, while compounds use radius-4, 1024-bit Morgan fingerprints reduced to 16 dimensions.The network features describe gene interactions and pathway membership beyond protein sequence alone.Source: B IMPLEMENTATION · p. 17

Objective

  • The KL terms and CLUB penalty are warmed up linearly, and the CLUB critic is trained with its own optimizer at a learning rate of 10^-3.Source: B IMPLEMENTATION · p. 17

Stage 2 and inference

  • The stage-2 conditional flow uses a DiT velocity field, separate tokens for two perturbed genes, and Sinkhorn-based optimal-transport training pairs.Inference integrates the probability-flow ODE with Euler steps using EMA weights; it copies the control cell’s invariant latent while transporting its responsive latent.Source: B IMPLEMENTATION · p. 17

C CONDITIONING REGULARIZATION

Conditioning regularization structures perturbation codes using response supervision, relative geometry, and noise augmentation. These choices encourage codes to reflect response similarity while remaining useful near learned conditions.

C CONDITIONING REGULARIZATION

  • A linear head predicts each condition’s mean training response, tying its code to the responses it must explain.The head is removed after training and kept linear so it cannot absorb the code structure being induced.Source: C CONDITIONING REGULARIZATION · pp. 17–18
  • Relative isometry makes perturbations with similar mean responses receive similar codes while leaving the code scale free to adapt to the prior.It matches the correlation of pairwise distances rather than their absolute values.Source: C CONDITIONING REGULARIZATION · p. 18
  • Adding Gaussian noise to condition embeddings prevents a lookup-table representation and keeps neighborhoods around learned codes valid for decoding nearby unseen genes.Source: C CONDITIONING REGULARIZATION · p. 18

D WHAT EACH LATENT BLOCK ENCODES

The responsive latent block carries perturbation information, while the invariant block carries substantially less, consistent across probes and estimators.

D WHAT EACH LATENT BLOCK ENCODES

E GENERATED POPULATIONS

On held-out Norman additive conditions, generated cell densities match observed distributions, and DRIFT improves on baselines across distributional metrics.

E GENERATED POPULATIONS

  • For the five held-out conditions with the largest shifts from control, generated densities match observed distributions in location and shape, including second modes for CNN1+UBASH3B and ETS2+MAPK1.Source: E GENERATED POPULATIONS · pp. 18–19
  • DRIFT raises Discr. Cos from 0.726 to 0.912 and reduces E-Dist by 22% relative to scBIG across held-out Norman additive conditions.Source: E GENERATED POPULATIONS · p. 19

F ABLATIONS

Ablations identify disentanglement as DRIFT’s largest contributor, with invariance and conditioning regularization adding gains and perturbation features especially helping unseen-gene generalization.

F ABLATIONS

  • The component ablations remove one named component at a time while holding the others and training budget fixed; contribution fractions are reported relative to DRIFT’s margin over Linear.Source: F ABLATIONS · p. 19
  • Disentanglement improves all 20 entries and contributes the most on every dataset, ranging from 11% on the additive split to 35% on holdout double perturbations and 21% overall.Source: F ABLATIONS · p. 20
  • The invariance penalty adds positive mean contribution on every dataset, with a 6% overall contribution and its largest effect on direction of change (ACC∆, 17% overall).Source: F ABLATIONS · p. 20
  • Conditioning regularization raises Replogle PDS from 0.573 to 0.604 when every test gene is unseen, accounting for 43% of DRIFT’s margin over Linear, with a 5% decrease in ρ∆.Its effect is smallest on the additive split, where test pairs combine genes seen individually during training.Source: F ABLATIONS · p. 20
  • STRING and Reactome features improve all nine Replogle metrics, but have limited impact on Norman splits where test conditions recombine individually observed genes.DRIFT with ESM2 alone outperforms Ruan et al. (2026) on every split; these additional features do not account for DRIFT’s advantage over scBIG, which also uses them.Source: F ABLATIONS · p. 20

G FULL RESULTS WITH STANDARD DEVIATIONS

The section reports a genetic-feature comparison and component contributions relative to a Linear baseline, alongside standard deviations across seeds for several datasets and splits.

G FULL RESULTS WITH STANDARD DEVIATIONS

  • Table 9 compares DRIFT’s ESM2, STRING, and Reactome genetic features with an ESM2-only variant, leaving all other components unchanged.Ruan et al. (2026) integrates all three sources; bold and underline mark the best and second-best results.Source: Table 9 · p. 22
  • Table 10 expresses each component’s contribution as a percentage of DRIFT’s margin over the Linear baseline for each dataset and metric.Its legend marks improvement, values within standard error across seeds, and negative values.Source: Table 10 · p. 22
  • Tables 11 and 12 report standard deviations over five seeds for Norman’s additive and holdout splits.Source: Table 11 · p. 22, Table 12 · p. 22
  • Table 13 reports DRIFT’s mean and standard deviation over five seeds on Replogle RPE1.Source: Table 13 · p. 23
  • Table 14 reports DRIFT’s mean and standard deviation over five end-to-end seeds on ComboSciPlex.Source: Table 14 · p. 23

Related papers

Loading 2609.35106v1…