sc2Flow: Mitigating Mean Prediction Bias in Single-Cell Perturbation with Dual-Stage Flow Matching

1 Introduction

High-throughput screening is a critical technology for advancing personalized cancer therapy [Carnero, 2006], enabling the identification of effective drugs and key genes through the analysis of post-perturbation gene expression profiles [Szalai and Veres, 2023]. However, testing all possible combinations of approved drugs would require astronomical resources and time [Menden et al., 2019]. Although numerous deep learning-based approaches have been proposed to predict unobserved perturbation combinations [Aliper et al., 2016], significant challenges remain.

Single-cell drug response exhibits a high degree of heterogeneity. As illustrated in Fig. 1a and Fig. 1b, when a minimal fraction of a cellular subtype (samples in the red box) responds to drug perturbation, other subtypes of the same cell population remain unresponsive. Traditional deep learning approaches for perturbation prediction, such as Variational Autoencoders (VAEs [Kingma et al., 2013]), typically rely on regression objectives like Mean Squared Error (MSE). From a statistical perspective, minimizing MSE converges to the conditional expectation of the target distribution [Bishop and Nasrabadi, 2006]. Thus, these models tend to predict the average of the two modes (response/non-response), resulting in universal non-response for all samples or biologically nonexistent intermediate states. At the gene level, the same issue persists: models tend to predict zero or underestimate its expression value.

To overcome this, Flow Matching (FM) has emerged as a powerful generative paradigm [Lipman et al., 2022]. This approach aims to reconstruct the full target distribution rather than merely its expectation. Even when a gene is unexpressed in the majority of cells, Flow Matching can learn to "push" the appropriate probability mass toward highexpression states for specific subpopulations.

However, existing Flow Matching methods (e.g., CellFlow [Klein et al., 2025]) predominantly rely on MLP architectures, which require fixed-dimensional inputs and outputs, forcing the model to process the entire genome. However, in most samples, fewer than 5% of genes are expressed (Fig. 1d left), even after HVGs selection (Fig. 1d right). This leads to substantial parameter redundancy, as a vast number of weights are wasted computing invalid "zero" values. Moreover, their parameter count grows linearly with gene set size. In contrast, Transformer [Vaswani et al., 2017] is inherently suited for processing variable-length sequences. Its number of parameters is independent of the input sequence length, offering superior scalability. Although the attention mechanism scales quadratically, the input length is very short if only expressed genes are processed.

Nevertheless, implementing Transformer-based FM for perturbation prediction presents a unique technical challenge: FM typically requires the source and target spaces to share consistent dimensionality. However, gene set before and after perturbation are different. This creates a dimensionality mismatch, preventing the direct application of Controlto-Perturbed flow generation.

To tackle these challenges, we introduce sc2Flow, a dualstage flow matching framework designed to predict singlecell drug perturbations. Our method decomposes the problem into two distinct phases: (1) Discrete Value Prediction: The model takes the discretized expression state of control cells (binary indicators of gene expression) and perturbation embeddings as input to predict the post-perturbation gene expression mask using discrete flow matching [Gat et al., 2024].

To investigate potential biological information loss caused by disregarding continuous expression values, we compared UMAP visualizations of discretized (binary) data against the original continuous data. The results demonstrate that the discrete representation preserves significant perturbation feature (Fig. 1c). Given the enormous input dimensionality in this stage, we employ a multi-layer residual network based on MLPs for prediction to reduce computational costs. (2) Continuous Value Prediction: A Transformer model utilizes tokens from these filtered genes to predict their quantitative expression level. Since the first stage has already predicted the target post-perturbation genes, we can construct a vector field that maps Gaussian noise to the true expression values for this gene subset, enabling training via flow matching while avoiding the dimensionality mismatch problem.

Extensive experiments on Sci-Plex3 [Srivatsan et al., 2020] and two other single-cell datasets [McFarland et al., 2020; Zhao et al., 2021b] demonstrate that sc2Flow significantly outperforms fully connected baselines (chemCPA [Hetzel et al., 2022], biolord [Piran et al., 2024]) and single-stage flow matching models CellFlow [Klein et al., 2025]. Our contributions can be summarized as follows:

• We present sc2Flow, a Transformer-based framework that unifies discrete and continuous flow matching for precise drug perturbation prediction.
• sc2Flow is the first method that effectively mitigates mean prediction bias while maintaining both scalability and high parameter efficiency.
• Extensive benchmarking across multiple datasets confirms that sc2Flow not only predicts expression values accurately but also significantly outperforms existing state-of-the-art methods in distribution matching.

2 RelatedWork

2.1 Encoder-Decoder based Prediction

Previous methods primarily rely on encoder-decoder frameworks. Early approaches used VAEs to encode gene expression before and after perturbation into a latent space to predict responses on novel cell types (e.g., scGen [Lotfollahi et al., 2019], CPA [Lotfollahi et al., 2023]). Later works incorporated drug structure encodings to generalize to novel compounds (e.g., PRnet [Qi et al., 2024], ChemCPA [Hetzel et al., 2022], Biolord [Piran et al., 2024]). Although these methods can achieve high scores on numerical prediction metrics (such as R2 ), due to limitations in sequencing technology, researchers cannot obtain paired samples before and after perturbation. Consequently, they can only evaluate the difference between the average values of all predicted and target samples under a specific perturbation condition, thereby concealing the aforementioned mean prediction bias problem. sc2Flow employs flow matching methods to mitigate this issue.

2.2 FlowMatching

Flow matching aims to learn a velocity field defined by Ordinary Differential Equations (ODEs) to construct a deterministic mapping between two distributions [Lipman et al., 2022]. By optimizing the probability flow field of the perturbation response rather than relying on traditional Mean Squared Error (MSE) objectives, this approach explicitly models the continuous trajectory of cell state transitions, thereby mitigating mean prediction bias. Recently, many flow matching-based models for drug perturbation prediction (such as CellFlow) have emerged. Since flow matching requires consistent input and output dimensions, even when using Transformer architectures, these single-stage models must include expression values for all genes. In contrast, while sc2Flow faces the same constraint, it avoids interference from non-expressed genes because the expressed genes have already been filtered out during the discrete prediction stage.

3 Methods

This section outlines the architecture and computational methodology of sc2Flow. As illustrated in Fig. 2, the framework decomposes the drug perturbation prediction task into two independent components: (1) Discrete Flow Matching (DFM) for predicting the binary gene expression states (i.e., identifying which genes are expressed), and (2) Continuous FlowMatching (CFM) for predicting the precise expression values of these genes.

This decoupling addresses the zero-inflation and high dimensionality of single-cell data, which make direct vector prediction computationally intensive and difficult to optimize. By filtering genes via DFM, the Transformer-based CFM can focus on a compact sequence of relevant genes. We employ Uni-Mol2 for 3D molecular encoding and incorporate dosage via explicit embeddings. At inference, DFM identifies active genes to guide CFM in generating the final perturbed profile.

3.1 ProblemDefinition

Traditionally, drug perturbation prediction is formulated as a multivariate regression task. In this work, we reframe it as a conditional generative modeling problem using the Flow Matching framework. Let X ∈ RN denote the gene expression space for N genes. We define a source distribution p0 representing the control group profiles x0, and a target distribution p1 representing the perturbed group profiles x1. The generation is conditioned on the perturbation context c = {cdrug, cdose}. The goal is to learn a time-dependent vector field vt(x, c) that pushes samples from p0 to p1 over the time interval t ∈ [0, 1].

3.2 Conditional Encoding

Both DFM and CFM components share a unified conditional encoding scheme to process gene identities and drug perturbations.

Gene and Drug Embeddings

For gene identification, we assign a learnable embedding vector to each gene. Let Egene ∈ RN×d be the embedding matrix where the i-th row corresponds to the i-th gene. The process to obtain gene identity embeddings hgene ∈ RN×d for sample x ∈ RN is defined as:hgene = Egene(x) (1)

For drug representation, the SMILES string of the drug is processed by the frozen Uni-Mol2 model. The resulting representation is projected to the model dimension d via a linear layer:

Dosage Encoding

To incorporate dosage information, we normalize the scalar dosage values to the range [0, 1]. This normalized scalar sdose scales a learnable dosage embedding vector edose ∈ Rd :

This multiplicative interaction ensures that the dosage magnitude directly modulates the embedding intensity.

3.3 Discrete FlowMatching (DFM)

The DFM module models the transition of gene states from the control group to the perturbed group. We treat the expression state as a binary variable (0 for unexpressed, 1 for expressed).

Conditional Probability Path

Let x ∈ {0, 1}N represent the binary expression vector. To apply Flow Matching on discrete data, we operate on the probability simplex. We map the discrete states to one-hot distributions. Let p0 = δ(x0) and p1 = δ(x1) be the probability mass functions of the source and target states, respectively. We define a linear conditional probability path pt(x) at time t as:

where pt ∈ RN×2 represents the probability of each gene being expressed or not at time t. A noisy sample xt can be drawn from this intermediate distribution:

Velocity Prediction via ResNet

The core of DFM is to predict the flow velocity, which drives the probability mass from t = 0 to t = 1.

We employ a Residual Network (ResNet) composed of Nd = 12 MLP Residual blocks. First, the model sets two learnable embeddings for expressed and unexpressed genes. Let ht ∈ RN×d be the embedding of xt. Then, the input to the network is the concatenation of ht, gene embeddings, condition embeddings, and the time step embedding temb:

The network predicts the logits ˆp1 for the target state x1:

Although the model predicts the final state distribution ˆp1, this formulation implicitly defines the vector field necessary to reach the target.

Training Objective

Since the target x1 is binary, the model is trained using the binary cross-entropy loss averaged over all genes:

Inference Process

During inference, we solve the ODE defined by the predicted vector field. The update rule for the probability pt moving from t to t + h is derived from the linear path assumption. The velocity field ut required to reach ˆp1 from pt is:

Applying an Euler step:

Finally, the discrete vector is sampled: xt+h = Categorical(pt+h). During inference, we generate the values by numerically integrating the ODE from t = 0 (control) to t = 1 (perturbation) using the predicted vector field ut.

3.4 Continuous FlowMatching (CFM)

The CFMmodule predicts the precise expression values. To enhance computational efficiency, the CFM module restricts its input exclusively to the subset of expressed genes. During training, we select this subset based on ground truth labels to avoid being affected by the overall performance of the DFM, thereby ensuring training stability. Conversely, during inference, the subset is dynamically determined by the binary states predicted by the DFM stage.

Sequence Construction

Let M be the number of expressed genes in x1. We construct a sequence of tokens Y ∈ RM×d from x1. Each token combines the gene's identity embedding and its continuous expression value. The conditional flow for the continuous values is defined as:yt = (1 - t)z + ty1 (11)

where z ∼ N(0, I) is the initial Gaussian noise and y1 is the target expression.

Transformer with Single-Stream Fusion

We employ a Transformer architecture to capture the interactions between genes and the perturbation conditions. To efficiently fuse the multimodal inputs (gene tokens, drug, dose), we adopt a Single-Stream architecture proposed in [Cai et al.,

2025]. We also use Sinusoidal Positional Encoding SinCos(·) [Su et al., 2024] to map the scalar expression value y ∈ R to a vector:

Then, this sequence is processed by Nc = 5 layers of Transformer blocks. Each block consists of Multi-Head SelfAttention (MSA) and a Feed-Forward Network (FFN), with Adaptive Layer Norm (AdaLN) conditioning on the time t. Additionally, to stabilize training directly in the original data space, we adopt the method of [Li and He, 2025], which directly predicts the final sample y1:

This allows every gene token to directly attend to the drug and dosage information within the self-attention mechanism.

Prediction and Loss

Similar to the DFM stage, the CFM model is trained to predict the clean target data y1 from the noisy input yt. The output corresponding to the gene tokens is projected back to the scalar expression space to obtain ˆy1. The training objective is the Mean Squared Error (MSE):

During inference, we generate the values by numerically integrating the ODE from t = 0 (noise) to t = 1 (prediction) using the predicted vector field vt = (ˆy1 - yt)/(1 - t).

4 Experiment

4.1 Setup

Baseline Methods

We trained all models using their official open-source implementations with optimal hyperparameters. Detailed training configurations are provided in the Supplementary Materials. The baseline methods are listed as follows:

• chemCPA [Hetzel et al., 2022] is an adversarial learning method based on VAE. It demonstrates state-ofthe-art performance on single-cell drug perturbation response benchmarks for single-drug treatments [Wei et al., 2025].
• biolord [Piran et al., 2024] is based on VAE and enables counterfactual predictions by disentangling the biological state of cells from external perturbations. It has since become a widely adopted state-of-the-art baseline in the field.
• CellFlow [Klein et al., 2025] is a recent drug perturbation response generation model based on Optimal Transport Conditional Flow Matching, which predicts velocity fields via MLPs.
• Unpredicted is a reference method. It makes no predictions and directly outputs the control group, serving to illustrate the inherent difficulty of predicting perturbation effects in the dataset.

Metrics

For the evaluation of generated molecules, we consider the following metrics.

• R2 (Coefficient of Determination) Measures the proportion of variance in the true gene expression data that is predictable from the generated data. It quantifies the overall fidelity of the predictions, with a higher R2 indicating better agreement.
• Energy Distance (ED) This parameter-free metric provides a robust, objective measure of global distribution similarity. By capturing all statistical moments without hyperparameter tuning, it ensures the generated population faithfully mirrors the shape and spread of real data.
• Maximum Mean Discrepancy (MMD) A kernel-based metric (using Gaussian RBF) that is particularly advantageous for high-dimensional data, allowing it to capture fine-grained, local discrepancies and subtle structural differences in complex single-cell phenotypes.

Dataset

We conduct experiments on the following single-cell drug perturbation datasets.

• Sci-Plex3 [Srivatsan et al., 2020] comprises approximately 572,853 samples, involving 3 cell lines and 188 drug perturbations. Following the practice of prior work, we allocate the 9 drugs with the most significant perturbation effects as the test set.
• McFarland [McFarland et al., 2020] contains 127,973 samples across 172 cell types treated with 11 drugs. We create partitions based on drugs. Specifically, we generate 5 different data splits, each with 2 to 3 drugs held out in the test set.
• Zhao [Zhao et al., 2021a] includes roughly 155,980 samples, recording the perturbation effects of 4 drugs on Glioma. We adopt a leave-one-drug-out strategy, resulting in 4 data splits where each drug serves as the test set once.

Gene Selection

All models were trained on an identical gene panel to ensure a fair comparison.

• HVGs: Since all baseline methods predict only a preselected set of 2000 HVGs, for a fair comparison, our method (sc2Flow) formally operates on the same full gene set of size 2000 after normalization using log(1 + x).
• Differentially Expressed Genes (DEGs): Since most genes remain unchanged post-perturbation, we focus evaluation on the top 50 DEGs per drug (identified by Seurat). This provides a more rigorous assessment of the model's predictive power for biologically relevant responses.

4.2 Results and Analysis

Robustness Across Dosage and Datasets

We evaluated sc2Flow against state-of-the-art baselines across three datasets (Sci-Plex3, McFarland, and Zhao). The quantitative results (Table 1 and Table 2) and visual distributions (Fig. 3) demonstrate that sc2Flow establishes a new standard for perturbation prediction.

Superiority in Distributional Alignment

The most significant advantage of sc2Flow lies in its ability to model the high-dimensional data distribution rather than merely predicting conditional means. As shown in Table 1 and Table 2, sc2Flow achieves substantial improvements in ED and MMD across all datasets. On the McFarland dataset, sc2Flow reduces the MMD by 89.5% compared to biolord (0.639 → 0.067) and 85% compared to chemCPA. Similarly, the Energy Distance is reduced by over 68% compared to the strongest baseline. On the Zhao dataset, the performance gap is even more pronounced, with sc2Flow achieving an MMD of 0.012, an order of magnitude lower than chemCPA (0.139) and biolord (0.633). These metrics confirm that while VAEbased methods (chemCPA, biolord) often suffer from or averaging effects-resulting in generated cells that cluster around the mean-sc2Flow's flow matching paradigm successfully captures the heterogeneity of the perturbed cell populations.

Precision in Gene Expression Recovery

Beyond distribution matching, sc2Flow excels in recovering precise gene expression levels, measured by R2 . In the McFarland dataset (Table 1), our model achieves an R2 of 0.901, comparable to the state-of-the-art biolord (0.907) and significantly outperforming chemCPA (0.785). However, in the more challenging Zhao dataset (Table 2), sc2Flow demonstrates superior robustness, achieving the highest R2 of 0.816, whereas baseline performance degrades significantly (chemCPA: 0.700). Single-stage flow matching models like CellFlow compress high-dimensional gene space into lowdimensional latent vectors (e.g., 100 principal components)

to manage computational cost. This compression inherently causes information loss, leading to larger reconstruction errors (Fig. 3, R2 panel). In contrast, sc2Flow's first stage (DFM) explicitly filters non-expressed genes, allowing the second stage (CFM) to operate on the full raw dimensionality of the active genes without information loss.

A critical requirement for drug discovery is the ability to predict changes in DEGs, which represent the biologically relevant response to treatment. When evaluating specifically on DEGs, sc2Flow maintains high R2 on DEGs, whereas baselines often show performance drops when zoomed in on DEGs (Table 1). By explicitly modeling the "zero-inflation" via the Discrete Flow Matching stage, sc2Flow avoids the common pitfall of predicting small, non-zero noise values for unexpressed genes. This ensures that the model focuses its generative capacity on genes that are actually biologically active, resulting in generated profiles that are biologically indistinguishable from real data (Fig. 3).

We observed distinct behavioral differences between model architectures across different perturbation intensities. As shown in (Fig. 3), traditional models (chemCPA, biolord) tend to default to the "unpredicted" (control) state when facing high-dose perturbations, essentially failing to capture the strong drug effect. This demonstrates that traditional methods are affected by the mean prediction bias problem.

In contrast, sc2Flow exhibits active predictive capabilities at high doses. By learning the continuous velocity field from control to perturbed states, our model effectively "pushes" the cell representation along the correct trajectory. This is evidenced by the consistent superiority in distribution metrics (ED and MMD) across the comprehensive 9-split crossvalidation on the experimental datasets, where sc2Flow outperformed competitors in nearly all scenarios.

4.3 Visual Comparison of Generated Samples

To visualize the superior distribution matching capabilities of sc2Flow, we analyzed the response of A549 cells to Alvespimycin (0.001 µM). Fig. 4 displays the ground truth sample alongside the predicted samples from each model that achieved the minimum Mean Squared Error (MSE) relative to the real data.

Results show that sc2Flow best aligns with the ground truth, capturing the high sparsity and sharp peaks characteristic of real gene expression. Unlike baselines that tend toward mean-averaging, sc2Flow's two-stage architecture first identifies expressed genes and then generates precise expression levels. This success is driven by the FlowMatching objective, which learns deterministic trajectories to the target distribution rather than approximating a static mean. Consequently, sc2Flow produces high-fidelity peaks that faithfully reflect biological signals, overcoming the "mean prediction bias" of traditional methods.

5 Conclusion

In this work, we propose sc2Flow, a Transformer-based TwoStage FlowMatching framework for efficient single-cell drug perturbation prediction. Extensive experiments demonstrate that sc2Flow significantly surpasses state-of-the-art autoencoder and single-stage baselines, particularly in preserving the distributional heterogeneity of cell responses. The reasons for sc2Flow's strong performance include the decoupling of discrete and continuous prediction which effectively handles data sparsity, the explicit exclusion of non-expressed genes to isolate the Transformer from noise, and the use of flow matching to capture complex distributional shifts rather than simple average effects. Moreover, sc2Flow establishes a new efficient paradigm for in silico drug screening under resourceconstrained settings.

6 Data and Code Availability

The code and datasets for this work are available at https: //github.com/hanwenlv-cmd/sc2Flow.

Acknowledgments

This work was supported by the National Natural Science Foundation of China [62372165].