Predicting and interpreting cell-type-specific drug responses in the small-data regime using inductive priors

Abstract

Predicting how small molecules affect diverse cell types phenotypically is central to drug discovery, yet it remains a challenging task. Modelling cell-type-specific transcriptional responses provides a scalable alternative for early candidate identification, enabling broader exploration and lower costs than exhaustive experimental exploration of the chemical space. Here we present PrePR-CT, a graph-based deep learning approach that utilizes cell-type-specific co-expression networks as an inductive bias to predict transcriptional responses to chemical perturbations. Graph attention networks learn biologically meaningful representations that capture cell-type-specific gene interactions, enabling gene-level attributions. Across five single-cell RNA sequencing datasets, including human blood and multiple cancer lines, one bulk transcriptomics dataset and a large-scale small-molecule screen, the method generalizes to unseen perturbations and previously unseen cell types under data-limited settings, achieving higher accuracy for expression variability compared to generative baselines. Attribution analyses identify high-attention genes that complement traditional differential expression analyses, highlighting pathway-specific mechanisms of small-molecule response. By combining scalability, robustness to distribution shifts and interpretability, PrePR-CT enables cell-type-resolved prediction of drug responses, providing a foundation for more precise modelling of cellular perturbations in early drug discovery.

Main

Pharmaceutical research and development is constrained by rising costs and declining productivity, resulting in fewer new drugs reaching the market^1,2. Artificial intelligence (AI) and generative AI hold great potential for drug development^3,4. This includes de novo drug design⁵, repurposing⁶ drugs for alternative indications⁷ and integrating AI techniques into the drug development pipeline⁸. A core step in these efforts is predicting the effects of chemical compounds. This helps identify therapeutic candidates early, reduces costs and improves safety. Chemical perturbations precisely alter cellular processes, revealing the effects of drugs and their mechanisms of action. However, drug responses vary considerably between cell types and the amount of representative data is limited.

The availability of large-scale chemical perturbation datasets enables robust, interpretable predictions across diverse cell types and experimental conditions. For example, in vitro drug screens, exemplified by high-throughput screening (HTS), allow the simultaneous testing of many molecules against specific biological activities or targets. HTS typically analyses bulk RNA sequencing data to assess transcriptional responses of thousands of molecules⁹. However, a significant limitation of this approach is its disregard for cell-type-specific responses to individual perturbations. The transcriptional profiles generated by single-cell RNA sequencing methods provide a powerful resource for capturing the heterogeneity of cellular responses to various perturbations, including chemical compounds and CRISPR-based interventions. Nevertheless, the challenge of scaling HTS to single-cell resolution persists, primarily because of the substantial costs and technical complexities of highly multiplexed chemical experiments. Consequently, these experiments typically allow measurement of the effects of a limited number of drugs, usually fewer than a hundred¹⁰.

Recent experimental studies^11,12,13,14 have developed multiplexed single-cell transcriptional profiling techniques to examine the heterogeneous responses of cancer cell lines to chemical or genetic perturbations. Such datasets motivate machine learning researchers to address the inherent experimental limitations by designing generative in silico models that can produce single-cell gene expression profiles for previously unobserved perturbations. Existing models^15,16 leverage generative approaches to map the data into a lower-dimensional latent space, subsequently employing simple vector arithmetic operations within this space to extrapolate the response of a specific drug to a different cell type. Uncovering the latent space has also proved effective in drug target discovery^17,18. However, a prominent challenge with generative models is their requirement for substantial data quantities to learn distribution parameters and generalize effectively to unseen conditions. One solution¹⁶ involves applying transfer learning in conjunction with bulk RNA sequencing data, such as connectivity map transcriptional data⁹. However, transfer learning restricts the feature space in the training set to the genes shared with the source domain dataset, which may not include the cell-type-specific marker genes.

On the other hand, recent studies have explored using prior information, such as gene–gene interaction networks, to predict cellular responses to unseen perturbations¹⁹. These models demonstrate that leveraging relational graphs yields significant improvements in predicting cellular responses. However, they ignore the cell-type-specific differences between treatment and control groups. Kana et al.²⁰ compared these differences in the principal component analysis space and showed significant variability between cell-type centroids in magnitude and direction, suggesting that the cell-type-specific context is crucial for predicting perturbation responses. Previous methods, such as GEARS¹⁹ and chemCPA¹⁶ demonstrate strong potential for generalizing across perturbations but face challenges in generalizing to out-of-distribution cell types. Introducing a cell-type-specific context as an inductive bias into model design could overcome this limitation, supporting broader generalization across diverse biological contexts²¹. For example, a recent study constructed cell-type-specific protein–protein interaction networks from single-cell data to support drug discovery²², while another integrated protein–protein interactions or gene regulatory networks inferred from transcriptional data with drug target information and disease driver gene sets to identify target genes that optimally shift the cellular state towards a desired post-perturbation profile²³.

To this end, we invoke an inductive prior, such as cell-type-specific co-expression networks, to predict the mean and variance of post-perturbation expression for differentially expressed genes (DEGs) in unseen cell types. We evaluate the performance of our method across five small-scale single-cell perturbation datasets and one large-scale chemical screening experiment. Additionally, we conduct a comparative analysis with two state-of-the-art generative models, focusing exclusively on chemical perturbations. Our results demonstrate the effectiveness of our approach for predicting DEG expression across diverse cellular contexts. We utilize the graph attention network (GAT) algorithm²⁴ to extract meaningful features from the cell-type-specific networks. To enhance interpretability, we analyse the attention values learned by the GAT layers for each gene, revealing their capability to capture non-trivial patterns from the data.

Results

PrePR-CT predicts single-cell perturbation responses using cell-type-specific graphs

Here we introduce PrePR-CT (Predicting Perturbation Responses in Cell Types), a graph-based deep learning method integrating cell-type-specific co-expression networks with single-nucleus RNA sequencing data to predict transcriptional responses to previously unseen chemical perturbations (Fig. 1). PrePR-CT leverages cell-type-specific co-expression patterns as an inductive bias to capture the fundamental cell-type expression patterns, allowing it to generalize to unseen cell types. Such co-expression patterns are modelled as graphs, enabling the generation of cell-type feature vectors using GATs.

Fig. 1: PrePR-CT, a graph-based deep learning method to predict single-cell perturbation response.

PrePR-CT learns cell-type-specific gene co-expression graphs from a batch of training samples using GAT layers, integrates them with unperturbed control gene expression values and pre-defined perturbation embeddings, and processes them through multilayer perceptron (MLP) layers to predict the expression response to a perturbation. Different cell types are colour-coded along the rows of the gene expression value and feature matrices as well as the cell-type graphs (Methods).

Single-cell perturbation datasets often contain sufficient cells in the control state to permit the use of metacells, here SEACells²⁵, to derive more robust expression estimates and to mitigate the imbalance with the less frequent perturbed state cells. As matched pairs of expression measurements before and after perturbation are technically not feasible, we used SEACells to match control and perturbed cells in single-cell datasets (Methods). For bulk RNA sequencing data, we employed either random assignments¹⁹ or selected pairs with minimal Euclidean distance in gene expression between the control and perturbed samples.

Furthermore, we incorporated chemical structure information into the model by generating molecular embeddings from SMILES representations. These chemical embeddings are then combined with the basal cell-type-specific transcriptional profile embeddings. By jointly modelling these complementary representations, the deep neural network learns the relationship between chemical structures and their transcriptional effects, enabling the accurate prediction of gene expression changes following chemical perturbations.

We evaluated the performance of our method across five distinct datasets representing diverse cell types and perturbations: human peripheral blood mononuclear cells (PBMCs) from Kang et al.²⁶, the multi-individual PBMC dataset (NeurIPS) from the Open Problems in Single-cell competition²⁷, a diverse set of multiplexed human cancer cell lines with 18 perturbations in the McFarland dataset¹¹, another set of human cancer cell lines with four perturbations in the Chang dataset¹² and a mouse liver dataset containing 11 cell (sub)types from Nault et al.²⁸ (see Methods for details). Leave-one-out cross-validation was used across cell types, where we removed the test cell-type graph during training and evaluated the predicted perturbation response for that cell type. When multiple perturbations were present in a dataset, compounds were stratified by E-distance¹⁰ and those with maximal and minimal transcriptional effects were selected, thereby ensuring a comprehensive assessment across varying perturbation strengths. We benchmarked our model against scGen¹⁵ and the compositional perturbation autoencoder (CPA)²⁹ for single-perturbation settings and against chemCPA¹⁶ for multi-perturbation settings. In addition to the unseen cell type and unseen perturbation evaluations, we assessed our method under a random splitting protocol on the chemical-protein–protein interaction dataset²³. This setting follows a different evaluation protocol from the out-of-distribution experiments and provides a complementary assessment of model performance (Supplementary Fig. 1).

PrePR-CT predicts perturbation responses in unseen cell types

We tested PrePR-CT’s ability to predict the transcriptional response of an unseen cell type to a known compound in five datasets, using the transcriptional response of another cell type to the testing compound as training data.

PrePR-CT accurately predicts both the mean expression levels (coefficient of determination (R²) > 0.90) and expression variability (R² > 0.70) of PBMC cell types in response to a single dose of IFNβ (Kang dataset) (Fig. 2a,b). Prediction accuracy improves further when focusing on the top 100 DEGs that are often central to biological studies. Accordingly, predicted cell profiles colocalize with observed cells in the Uniform Manifold Approximation and Projection (UMAP) space (Fig. 2c), and PrePR-CT consistently outperforms scGen and CPA in both metrics (Fig. 2d). It also demonstrates lower R² variability across cell types, indicating more stable predictions. To further dissect model performance, genes were stratified by transcriptional change upon perturbation in B cells (Fig. 2e). R² scores were high for all genes (0.91), highly expressed genes (0.93), upregulated DEGs (0.99) and downregulated DEGs (0.84), and only negative for lowly expressed genes (R² = −1.83), which likely indicates only a very weak or absent transcriptional response to the perturbation. Overall, PrePR-CT demonstrates robust modelling of transcriptional responses, especially for DEGs, which are essential for interpreting cell-specific perturbation effects.

Fig. 2: PrePR-CT accurately predicts the effect of a single perturbation in an unseen cell type.

a, UMAP visualization of cell types stimulated with IFNβ (n = 24,249 cells), as demonstrated in Kang et al. b, R² values comparing the predicted versus observed mean (left) and standard deviation (right) of gene expression across seven cell types. Bars show the mean R² over n = 100 bootstrap iterations with ± standard deviation error bars; black dots indicate individual bootstrap iterations. Results are shown for all genes and the top 100 DEGs of the held-out test cell type. c, UMAP visualization comparing predicted (pred) and observed samples of stimulated and control B cells (n = 3,789 cells). d, Comparison of R² scores for the predicted mean and standard deviation of gene expression across CPA, scGen and PrePR-CT under the leave-one-cell-type-out setting (IFNβ perturbation). For each metric and gene set, R² values are reported across the held-out cell types (n = 7). Box plots show the median (centre line) and interquartile range (25th–75th percentiles; box); whiskers extend to 1.5× the interquartile range. e, Scatter plots of predicted versus observed differential expression (mean stimulated − mean control) in B cells across gene subsets: all genes in the B cell graph (n = 3,176), highly expressed genes (n = 1,588), lowly expressed genes (n = 1,588), upregulated DEGs (n = 43) and downregulated DEGs (n = 43). Each point represents one gene. Red lines show linear regression fits and shaded areas indicate the 95% confidence intervals. R² values are shown in each panel.

Source data

To extend the range of perturbations and cell types, we utilized the NeurIPS PBMC dataset, containing 144 small-molecule perturbations of four distinct immune cell types (Fig. 3a). This dataset structure allowed the use of cross-validation and perturbation stratification by transcriptional response strength (Fig. 3b). As illustrated in Fig. 3c, drugs with small transcriptional impact were easier to predict, allowing all models to perform well in estimating mean gene expression and variance across cell types. Nevertheless, PrePR-CT consistently achieved the highest R² scores, particularly among the top 100 DEGs, whereas CPA and scGen showed greater variability, particularly in prediction variance.

Fig. 3: PrePR-CT accurately predicts transcriptional responses to multiple perturbations in an unseen cell type.

a, UMAP visualization of cells coloured by condition (top) and cell type (bottom) across multiple drug perturbations (n = 291,793). b, E-distance between control and perturbed cells for each compound (n = 135 drugs), used to identify perturbations inducing large versus small transcriptional responses. c, R² scores for predicted versus observed gene expression mean (left) and standard deviation (middle) across cell types for penfluridol, flutamide and SCH-58261, evaluated using all genes and the top 100 DEGs of the held-out test cell type. Each point represents one cell-type drug evaluation. The box plot (right) summarizes the performance for penfluridol across available cell types (n = 3; NK cells excluded). The centre line indicates the median and the box bounds indicate the upper and lower quartiles. Whiskers extend to the largest and smallest values within 1.5× the interquartile range. d, The same as in c, but for scriptaid, ketoconazole and dactolisib. The box plot (right) summarizes the performance for dactolisib across evaluated cell types (n = 4). e, UMAP visualization comparing predicted and observed perturbed cells for penfluridol (left; n = 4,914) and dactolisib (right; n = 4,814), shown alongside control cells in B cells.

Source data

When considering drugs that induce the largest observed transcriptional shifts, model performance becomes more variable, with a general decline in predictive accuracy (Fig. 3d). PrePR-CT maintains the highest performance, particularly in capturing gene expression variance, demonstrating its robustness under complex perturbation conditions. Again, UMAP embedding of predicted and observed B cells yields robust co-localization for two representative drugs (Fig. 3e): penfluridol, a mild perturbation (Fig. 3e, left), and dactolisib, a strong perturbation (Fig. 3e, right). For penfluridol, the predicted perturbed cells (blue) closely overlap with both the control population (grey) and the observed perturbed cells (red), indicating that the model effectively captures subtle transcriptional changes. In contrast, dactolisib-treated cells form a distinct cluster in which the predicted cells remain well-aligned with the true responses. Natural killer (NK) cells were excluded from the mild perturbation analysis (Supplementary Fig. 2) due to minimal transcriptional response and low DEG counts, resulting in low or negative R² scores. This exclusion was justified because the absence of a strong biological signal limited the ability to make meaningful predictions. Notably, PrePR-CT did not overfit to noise in these cases, reflecting a desirable inductive bias: the model responds only when a true signal is present, rather than attempting to extract patterns from background variability. This behaviour aids its generalizability and biological relevance.

Supplementary Figs. 3 and 4 present the results for the Nault²⁸ and McFarland¹¹ datasets. In the Nault dataset, cells cluster distinctly by both cell type and treatment condition (Supplementary Fig. 3a), reflecting strong perturbation-driven transcriptional shifts. The obtained model performance varied substantially across cell types but remained consistent for very similar populations, such as central and portal hepatocytes (Supplementary Fig. 3b). As most DEGs in the dataset were downregulated, the model more effectively captured these changes (R²= 0.93 for downregulated genes; Supplementary Fig. 3e). In the McFarland dataset, we found predictions were easier for drugs with minimal effects, with a decline in R² for compounds that induced substantial transcriptional changes (Supplementary Fig. 4).

We further evaluated PrePR-CT generalization by training on the NeurIPS dataset and testing on each tissue type within the Tahoe-100M³⁰ dataset, considering only shared compounds and highly variable common genes (Supplementary Note 1), again predicting transcriptional effects in cell types not seen during training.

To better understand the factors that influence generalization performance, we analyse the results with respect to two key variables: (1) the strength of the transcriptional response induced by each drug measured by E-distance and (2) the transcriptional similarity between training and testing cell types via Euclidean distance. At the cell-type level, a very weakly negative correlation (r = −0.26) between median perturbation E-distance to control over all considered drugs and corresponding median DEG prediction R² indicates that cell lines that respond more strongly to perturbations are not inherently harder to predict (Supplementary Fig. 6). In contrast, at the drug level, stronger median responses for all cell types yield slightly more predictable effects (r = 0.53, and r = 0.38 when removing the outlier). Finally, when examining generalization across cell lines, we found a stronger negative correlation (r = −0.60) between the median Euclidean distance (from each Tahoe cell line to the NeurIPS training cell types) and the corresponding R² scores, illustrating the difficulty of out-of-distribution generalization as dissimilarity from the training data increases.

We benchmarked PrePR-CT against scGen, chemCPA and a mean expression baseline in two settings: <Tahoe controls excluded, strict out-of-distribution> and <Tahoe controls included>. scGen applies only to the latter, as it requires control cells from the test cell types during training, whereas PrePR-CT and chemCPA apply to both. While the mean baseline outperformed PrePR-CT and chemCPA for most cell lines in the first setting, PrePR-CT performed better in the second setting (Supplementary Fig. 7), demonstrating better generalization with the control data.

PrePR-CT predicts the response to unseen perturbations

We also assessed the model’s ability to predict the effect of unseen compounds within a known cell type using the Chang dataset, which includes three structurally similar small molecules (Fig. 4a). We adopted a leave-one-compound-out strategy analogous to the cell-type generalization setting.

Fig. 4: PrePR-CT accurately predicts the transcriptional effects of unseen perturbations in the same cell type (Chang dataset).

a, Molecular structures of the three evaluated compounds: GNE-104, GNE-069 and erlotinib. b, R² scores for the predicted versus observed gene expression mean (left) and standard deviation (right) across compounds, evaluated on all genes and the top 100 DEGs. Bars show the mean R² over n = 100 bootstrap iterations, and error bars denote ± standard deviation. c, Comparison of chemCPA and PrePR-CT performance across four metrics: mean and standard deviation of gene expression for all genes and the top 100 DEGs. Each box plot summarizes the performance across all evaluated compound cell-type test cases (n = 3 compounds). Centre lines indicate the median and boxes the interquartile range (25th–75th percentiles) and whiskers extend to 1.5× interquartile range. d, UMAP visualization of control, observed perturbed and predicted erlotinib-treated cells, showing strong alignment between predicted and observed transcriptional responses (n = 34,154 cells).

Source data

Despite the structural similarity among the compounds, PrePR-CT accurately captured the mean transcriptional response across all genes and the top 100 DEGs, achieving consistently high R² scores (Fig. 4b, left). For gene-wise variance predictions, however, performance differed across compounds: while GNE-104 and GNE-069 showed similar accuracy, erlotinib achieved a markedly higher R² score, especially among the top 100 DEGs (Fig. 4b, right). As scGen does not apply to unseen perturbations, we benchmarked PrePR-CT against chemCPA. PrePR-CT outperformed chemCPA in predicting gene expression variance, particularly for highly responsive genes. UMAP projections of erlotinib-treated cells showed that the predicted profiles closely match the observed perturbed cells and separate from the controls, confirming the ability of PrePR-CT to accurately capture drug-induced transcriptional shifts (Fig. 4c,d).

Next, we extended the evaluation of unseen perturbations to the LINCS dataset, which contains bulk transcriptomic profiles⁹. From this dataset, we constructed a subset of cell lines and perturbations based on the sample size distribution (Methods and Supplementary Fig. 8). In this setting (Supplementary Fig. 8b,c), both PrePR-CT and chemCPA achieved comparable performance in predicting mean DEGs; however, PrePR-CT outperformed chemCPA in predicting the variance of DEGs across all cell lines.

PrePR-CT operates successfully in a small-data regime

A recurring challenge in perturbation-response modelling arises from the different ways experimental setups using limited-size single-cell datasets handle their limited size, requiring a balance between the number of conditions (usually pairs of cell type and treatment) and the resulting number of cells per condition. This typically yields (1) datasets with one cell type and few perturbations, (2) datasets with one perturbation and few cell types, both yielding many cells per perturbation, and (3) datasets containing many cell types and perturbations, yielding a dramatically lower number of cells per pair. These scenarios pose different obstacles to learning perturbation effects, particularly when the variability of the transcriptional response must be estimated from only a few perturbed cells. To systematically evaluate the behaviour of PrePR-CT under conditions similar to (1) and (2), we used subsets of varying sizes from the already small Kang PBMC dataset, consisting of 24,249 cells, for training. The model consistently converged, even when the subsampling rate was aggressively reduced to 20% of the total cells.

Mean expression and variability R² across all genes and the top 100 DEGs show that PrePR-CT consistently outperforms CPA across all subset sizes (Fig. 5a,b). While CPA’s R² scores for mean expression remain relatively stable across dataset sizes, PrePR-CT’s performance declines when the dataset is reduced to below 60% (~15,000 cells). For both methods, R² for mean expression is higher for the top 100 DEGs than for all genes, likely reflecting the lower variability expected at higher gene expression levels.

Fig. 5: PrePR-CT operates successfully in the small-data regime.

a,b, R² scores for predicting the mean (a) and standard deviation (b) of post-perturbation gene expression in B cells as a function of training set size on the Kang dataset. Results are shown for all genes in the B cell graph (n = 3,176) and the top 100 DEGs for both CPA and PrePR-CT. Each point represents the mean performance across n = 3 independent runs with different random seeds; error bars denote ± standard deviation across the three runs. PrePR-CT consistently outperforms CPA across all training set sizes, demonstrating robust performance even in the small-data regime.

Source data

In contrast, both methods exhibit a similar decline in the quality of variability predictions in datasets with fewer than ~15,000 cells. However, CPA’s variability estimates are more precise when considering all genes, while PrePR-CT’s estimates are consistently more precise for the top 100 DEGs. This difference is likely attributable to the use of SEACells to match control and treatment conditions, which enables balanced representation across cellular clusters and thereby captures the full landscape of the control population.

Graph attention network learns meaningful representations

We analysed attention maps to assess the GAT network’s ability to capture meaningful patterns. High-attention genes (HAGs) were identified by summing each gene’s outgoing edge weights and ranking them accordingly. This analysis was performed on cell-type-specific graphs from the Kang dataset. GAT identifies distinct gene sets across different cell types, and the analysis of HAGs suggests that GAT extracts complementary features compared to DEG analysis. To explore the relationship between these two approaches, we compared the union of the top 50 HAGs across all PBMC cell types with that of the top 50 DEGs, yielding 333 and 130 genes for HAGs and DEGs, respectively. Of the genes, 70.5% are uniquely HAGs (Fig. 6a), indicating that a significant number of HAGs are not differentially expressed. Similarly, 24.5% are uniquely DEGs, demonstrating that not all DEGs are required for the prediction. We observed only a 5% overlap between HAGs and DEGs in PBMCs, suggesting that these sets captured complementary signals. HAGs exhibit greater cell-type specificity, while DEGs are more commonly shared, highlighting the ability of GATs to identify unique biological features and to offer complementary insights to traditional differential expression in understanding gene function. Next, we investigated the set of unique HAGs through pathway enrichment analysis and attention scores for ranking. The analysis yielded significant ontology terms, including T cell activation and immune system pathways (Fig. 6b), which align with the known mechanism of action of IFNβ within the immune system, particularly its effects on T cells. Furthermore, the existing literature supports roles for a range of HAGs in the considered perturbations. For example, CD74 is among the top 20 HAGs identified in dendritic cells, and Nguyen-Pham et al.³¹ demonstrated that CD74 negatively regulates dendritic cell migration, underscoring its relevance in the design of protocols aimed at enhancing dendritic cell migratory capacity in immunological research. LCK, LAT and CD3E are among the top 10 HAGs in CD4⁺ T cells and are members of immune system pathways, including T cell receptor signalling and T helper 17 cell differentiation. Subsequently, we examined the associations between HAGs and IFNβ-related pathways. We found that XAF1, ISG20, STAT1 and RSAD2 are involved in multiple interferon-related pathways, including the cellular response to type I interferon (GO:0071357) and the type I interferon signalling pathway (GO:0060337). These genes are among the top DEGs in response to IFNβ in the Kang dataset. This finding suggests that the GAT model effectively attends to the most critical features within the gene set.

Fig. 6: HAG analysis.

a, Venn diagram showing the overlap between HAGs and DEGs. b, Gene set enrichment analysis of HAGs reveals a strong association with immune activation pathways, including T cell activation, cytokine signalling and antigen receptor signalling. c, R² scores before and after random shuffling of cell-type-specific features across multiple cell types, evaluated using all genes (n = 3,176) and the top 100 DEGs. Data are shown as mean ± standard deviation over five independent random shuffling runs (different random seeds). d, R² scores before and after masking HAGs from the input features.

Source data

To further examine the ability of GATs to attend to meaningful biological features, we investigated their capability to learn the correlation structures defining the edges in cell-type-specific graphs. To this end, we computed Spearman’s correlation between the number of neighbours and the attention rank for the genes in each cell type. The results revealed a minimum correlation of 0.68 in B cells and a maximum of 0.91 in FCGR3A⁺ monocytes. These correlation values across different cell types suggest that GATs effectively learn the network structure of correlations within cell-type-specific graphs. To assess the importance of features learned by the GAT model, we performed ablation studies by disrupting these features, measuring R² changes for mean gene expression prediction under different perturbations as a proxy for predictive degradation (Fig. 6c,d). First, we randomized cell type features while preserving the overall graph structure to assess the importance of cell-type-specific information (Fig. 6c). Second, we masked the features of HAGs at test time to evaluate their direct contribution to prediction performance (Fig. 6d). These perturbations allowed us to assess the relative importance of global cell-type features and HAG-derived signals in supporting accurate predictions. Shuffling cell-type features resulted in a pronounced drop in R² values (Fig. 6c), particularly for the top 100 DEGs, compared to the complete set of genes. While the model retained high performance across all genes before shuffling (solid lines), disrupting the alignment between cell-type features and graph structure led to a more significant reduction in prediction accuracy for DEGs (dashed pink line) across most cell types. This finding suggests that DEGs are more sensitive to the integrity of cell-type features, likely due to their stronger responsiveness to perturbations, which becomes obscured when the input feature-cell-type correspondence is disrupted. Masking HAG features at test time also led to a noticeable, albeit smaller, decrease in R² values (Fig. 6d). Again, the top 100 DEGs (dashed pink) exhibited a greater decline in predictive performance than all genes (dashed blue), although overall model performance remained relatively high. This indicates that HAGs contribute disproportionately to the accurate prediction of perturbation-responsive genes, and their removal selectively impairs the model’s ability to reconstruct DEG expression. The model appears more robust to disruption when using only the top 100 DEGs, which likely reflects that DEGs capture more informative and biologically relevant variation, whereas using all genes introduces noise from lowly expressed or uninformative genes, making the model more susceptible to degradation when features are shuffled or masked.

These findings highlight the importance of cell-type features and HAG-derived representations in shaping the GAT model’s performance. DEGs are more sensitive to feature disruptions, supporting their role as key indicators of perturbation effects.

Discussion

We used a gene co-expression network-based cell type representation as an inductive bias to predict cell-type-specific responses to chemical perturbations. This strategy effectively predicted the mean and variance of post-perturbation expression for DEGs across different cell types and chemical perturbations, using leave-one-out cross-validation within each dataset to evaluate out-of-distribution scenarios. By analysing the salient features of GAT-generated cell-type-specific networks and explicitly focusing on genes with HAGs, we identified cell-type markers and perturbation-related genes. This approach offered a complementary perspective by identifying features beyond those examined in differential gene expression analysis. It is essential to note that DEG analysis tends to be more stable than HAG analysis, as HAG genes exhibit lower expression levels. Furthermore, HAG genes are more sensitive to training parameters, such as the number of attention heads in the GAT layers. These findings underscore the need to assess the stability of the model’s interpretations by evaluating their consistency across conditions. Additionally, it is essential to investigate how various model selection criteria impact the interpretability of deep learning models.

The results of our benchmarking against state-of-the-art generative models highlight the advantage of integrating cell-type-specific inductive bias to predict gene expression in response to perturbations. Importantly, our method performs well in the small-data regime, where data scarcity often limits the efficacy of alternative approaches. Moreover, the direct optimization of the Earth Mover’s Distance (EMD; Wasserstein) yields better results than reference generative models, underscoring the robustness and predictive accuracy of our approach. These results demonstrate that the method remains highly effective across diverse single-cell perturbation datasets, emphasizing its broad applicability to single-cell data analysis, particularly in scenarios with limited data availability. In addition, recent benchmarking studies^32,33,34 comparing deep learning methods with linear models have investigated why simple linear baselines often outperform deep learning approaches. These studies highlight the issue of control bias and underscore the need to design deep learning methods, evaluation frameworks and metrics that can effectively distinguish between systematic variation and perturbation-specific responses.

One limitation of the GAT approach is that the graph node features model only the mean and variance of gene expression. A promising direction for future work would be to incorporate additional features, such as cell cycle states, morphological traits and tissue of origin, to improve the biological relevance of the learned representations. Expanding the method to encompass a broader range of cell types would further enhance its generalizability, enabling it to capture diverse cellular contexts more effectively and increase predictive accuracy, particularly in challenging scenarios. Our transfer learning experiment using the NeurIPS and Tahoe datasets highlights several challenges in generalizing across different experiments. We found that the transcriptional distance between training and testing cell types significantly impacts model performance, an observation aligned with current experience in machine learning^35,36. We also observed that the sets of highly variable genes vary across datasets, complicating feature selection. Drugs with strong perturbation effects partly offset this mismatch by generating stronger, more predictable signals. Finally, differences in experimental covariates, such as batch effects or protocol discrepancies, can mask true effects and confound interpretation. Addressing these limitations in future work may lead to the development of more robust and generalizable models for predicting cellular responses to perturbations. The proposed input representation could offer a potential for various drug discovery applications, including predicting drug combination effects, identifying therapeutic targets and exploring interactions between different perturbations.

In summary, we have developed a method that shows strong performance in a small-data regime. This outcome resulted from incorporating an appropriate inductive, cell-specific network bias, which simultaneously enables mechanistic interpretation via the attention layers. This line of research may open the possibility of developing quantitative cell-specific systems biology models^37,38, which could fuel predictive biology³⁹.

Methods

PrePR-CT considers a perturbational single-cell dataset $D=\{({x}_{i},{y}_{i},{g}_{i}){\}}_{i=1}^{N}=\{({x}_{i},{y}_{i},({c}_{i},{p}_{i})){\}}_{i=1}^{N}$, where ${{\bf{x}}}_{i}$ and ${{\bf{y}}}_{i}\in {R}^{m}$ are m-dimensional gene expression vectors representing, respectively, a control cell and its transcriptional response to a perturbation ${g}_{i}$. Each perturbation ${g}_{i}$ comprises a cell-type label ${c}_{i}\in C$ and a $k$-dimensional embedding vector ${{\bf{p}}}_{i}$. We associated each cell type with one of the cell-type-specific graphs $C=\{{G}_{{c}_{1}},{G}_{{c}_{2}},\ldots ,{G}_{{c}_{M}}\}$, where $M$ is the number of cell types in $D$. We predicted the perturbation effect as a function of the cell-type graph, a set of control samples of the same cell type and a perturbation embedding.

Cell-type-specific graphs

In PrePR-CT, each cell type is represented as a graph, with genes as nodes and edges between genes that exhibit sufficiently correlated expression patterns. Smaller disconnected submodules were retained to capture cell-type-specific regulatory patterns while maximizing the overlap of shared genes across cell types. Consistent with the goal of predicting responses to unseen perturbations, we constructed cell-type graphs using expression data from unperturbed (control) cell states, as perturbation-induced measurements were not available during testing. Moreover, these cell-type graphs can contain different gene sets to reflect each cell type’s distinct expression and functions.

The co-expression graphs were constructed by aggregating single-cell expression values of unperturbed cells into metacells using SEACells²⁵ to mitigate the inherent sparsity of single-cell expression data. We calculated the set of 5,000 highly variable genes (HVGs) before the metacell aggregation. To construct the graph for a given cell type ${G}_{{c}_{i}}$, all pairwise correlations between HVGs were computed using only metacells of the selected cell type. Edges with correlation strengths below the 99th percentile of the absolute correlation distribution were removed.

The resulting co-expression graphs included varying numbers of genes based on the defined threshold, with some genes shared across cell types. The selected parameters depended on the processed data and constituted a trade-off between considering the most specific genes per cell type, retaining similar genes between the cell-type graphs and controlling the graph’s size, which affects the time complexity of the training. The graph’s feature matrix was initialized using the mean and standard deviation of the control state to capture the basal expression profile.

Matching control and perturbed cells in single-cell data

Because single-cell transcriptomic assays are destructive, paired measurements from the same cell before and after perturbation were not available, necessitating a strategy to match the control and perturbed cells. Within each cell line, we randomly assigned control-state cells to perturbed cells to avoid bias arising from systematic pairing. As the datasets contained substantially more control than perturbed cells, this procedure inherently subsampled the control population. To ensure that the subsampled controls remained uniformly distributed across the transcriptional manifold, we guided the selection using metacell clusters generated by SEACells. These metacells provided a coarse-grained representation of control-state variation and allowed uniform sampling independent of local differences in cell density. For each cell type and perturbation, we drew an equal number of control cells from the metacells and assigned them to the perturbed cells for training. Gene-expression samples were used as inputs and outputs; SEACell clusters served solely as a tool to align the control and perturbed groups.

Chemical structure representation of the perturbing compound

We used SMILES strings to encode the chemical structures of the perturbing compounds. SMILES provides a compact, machine-readable format for describing molecular structures and serves as input for generating molecular embeddings. We used chemical descriptors computed via the RDKit package, a widely used cheminformatics toolkit, to extract structural and physicochemical features from each SMILES string. This descriptor-based approach captured relevant molecular properties, such as atom types, bond connectivity, partial charges and topological features, which were embedded into fixed-length vectors representing each compound. These embeddings provided structured molecular features that could be integrated with transcriptional response data, allowing our model to predict drug effects across different cell types and conditions. We obtained the SMILES strings for all chemical perturbations from the DrugBank⁴¹ database or the original publications of the corresponding datasets.

Graph attention network

The GAT²⁴ defines an attention mechanism on the graph structure by computing an attention score ${\alpha }_{i,j}$ between a given node i and each node $j$ in its first-order neighbourhood set ${N}_{i}$. This score represents the significance of node $j$’s features to node i and is calculated as:

$${\alpha }_{{ij}}=\frac{\exp \left(\mathrm{LeakyReLU}({{\bf{a}}}^{T}[{{\bf{Wh}}}_{i}\parallel {{\bf{Wh}}}_{j}])\right)}{{\sum }_{s{\in {\mathscr{N}}}_{i}}\exp (\mathrm{LeakyReLU}({{\bf{a}}}^{T}[{\bf{W}}{{\bf{h}}}_{i}\parallel {\bf{W}}{{\bf{h}}}_{s}]))}$$

where ${{h}^{\to }}_{i}$ is a feature representation of node i transformed using the weight matrix $W$, and ${\bf{a}}$ is a learnable weight vector. The self-attention mechanism is then extended to multi-head attention by applying $L$ independent mechanisms, and their outputs are either concatenated or averaged:

$${{\bf{h}}}_{i}^{{\prime} }={\parallel }_{q=1}^{L}\sigma \left({\alpha }_{i,i}^{q}{{\bf{W}}}^{q}{{\bf{h}}}_{i}+\mathop{\sum }\limits_{j\in {{\mathscr{N}}}_{i}}{\alpha }_{i,j}^{q}{{\bf{W}}}^{q}{{\bf{h}}}_{j}\right)$$

The final output ${{h}^{\to }}_{i}^{{\prime} }$ is a high-level feature representation of node i in the graph. We employed GATs in this project because they showed successful examples of generalizing completely unobserved graphs during the training process²⁴.

Loss function

The mean squared error (MSE) is commonly used as a loss function in regression models. However, MSE primarily focuses on the mean and does not consider the overall distribution of the actual and predicted samples. In contrast, the Wasserstein distance, or EMD, is an optimal transport measure for comparing probability distributions with disjoint supports⁴². The Wasserstein distance has achieved significant success in the context of generative adversarial networks, leading to more stable training and improved sample quality by providing a more meaningful and smoother gradient signal compared to traditional generative adversarial networks⁴³. To approximate the EMD, we seek the optimal transportation cost $\gamma$:

$$\begin{array}{rcl}\begin{array}{c}\begin{array}{c}\mathop{\mathrm{min}}\limits_{\gamma }\langle \gamma ,M{\rangle }_{F}\\ \text{s.t.}\gamma 1={\bf{a}},\\ {\gamma }^{T}1={\bf{b}},\\ \gamma \ge 0\end{array}\end{array}\end{array}$$

where $M$ is a cost matrix corresponding to the Euclidean distances between the actual and predicted samples, and ${\mathbf{a}}$ and ${\mathbf{b}}$ are the sample weights of the source and target distributions, respectively^44,45.

For the implementation, we used the mini-batch EMD function implemented by Flamary et al.⁴⁶ to solve the aforementioned optimization problem. The optimal transport library in Python offers several tutorials for using this loss function⁴⁶.

The final loss is defined as:

$$\text{Loss}=\frac{1}{T}\displaystyle \mathop{\sum }\limits_{t=1}^{T}{\rm{EMD}}({Y}_{t},{\hat{Y}}_{t})$$

where $T$ is the number of cell-type–perturbation pairs ${g}_{i}$ in the current batch, and ${Y}_{t}$ and ${\hat{Y}}_{t}$ denote the observed and predicted samples, respectively.

Implementation details

Utilizing a cell-type-specific gene interaction network requires mapping each gene expression sample to the graph of its respective cell type. However, this direct mapping substantially increases the training time complexity. To address this issue, we adopted an alternative formulation. First, we extracted the graphs within each mini-batch and embedded them using the graph attention layers. We then mapped the samples to the features generated by these layers according to their cell-type labels. This avoided treating each cell as an independent graph and thereby reduced computational complexity. Figure 1 demonstrates the training steps.

The model required tuning six hyperparameters during training. Supplementary Table 2 lists these parameters along with the values optimized throughout the training process using the hyperparameter optimization tool Optuna¹⁸ on a fixed random validation set. Optimizer⁴⁷ was used for model training, and no activation function was applied to the output layer.

Evaluation

We evaluated model performance using the coefficient of determination (R²), which measures the agreement between the observed statistics (mean and standard deviation) and their predicted values. Our analysis focused explicitly on gene expression statistics rather than overall gene expression profiles. This distinction is crucial because other methods included in the comparative analysis map all control cells to the perturbed state, resulting in an unequal number of predicted samples compared to the actual perturbed samples in the ground truth data.

Consequently, we concentrated on capturing the distinct shifts from the control state to the perturbed state. The intentional choice of R² as our evaluation metric offers notable advantages over alternative metrics such as Pearson correlation and MSE. Unlike Pearson correlation, which measures only linear relationships and overlooks magnitude considerations, R² provides a more comprehensive assessment of predictive power. This deliberate selection aligned with our objective of capturing nuanced changes in gene expression.

For the gene set selected for evaluation, we assessed all the genes in the graph for the test cell type and the top 100 DEGs. The top 100 DEGs were defined using a combined ranking scheme based on the P value and log fold change statistics computed via t-tests using the scanpy Python package. First, we filtered genes with a P value of 0.05 and then applied the rank product statistic:

$$\mathrm{RP}(\mathrm{gene})={\left(\mathop{\prod }\limits_{i=1}^{s}{r}_{\mathrm{gene},i}\right)}^{\frac{1}{s}}$$

where $s=2$ represents the two ranking criteria: P value and log fold change. This ranking aimed to eliminate noisy genes with extreme log fold change values and very small P-values.

Datasets

We assessed the efficacy and versatility of our method across diverse scenarios by evaluating it on multiple datasets encompassing various cell types and perturbations. Specifically, we employed five distinct datasets: human peripheral PBMCs treated with interferon-beta (IFNβ). IFNβ is known for its potent antiviral activity, growth inhibition and immune modulation. It modulates T cell behaviour by reducing activation and proliferation, limiting infiltration into sensitive tissues and promoting regulatory mechanisms, published by Kang et al.²⁶, the McFarland dataset¹¹, the Chang dataset¹², the single-dose TCDD liver by Nault et al.²⁸ and the cross-individual dataset (NeurIPS) from the Open Problems in Single-Cell competition²⁷.

The McFarland dataset encompasses 13 chemical perturbations, each with a specific dose value, and comprises over 100 cell lines. We focused on the five cell lines with the largest sample counts and excluded drugs lacking SMILES identifiers in the DrugBank database⁴¹.

The Nault dataset entailed single-nucleus RNA sequencing of flash-frozen C57BL6 mouse livers. In this dataset, mice underwent subchronic administration of a specified TCDD dose via oral gavage every 4 days over 28 days. As indicated by Kana et al.²⁰, we excluded all immune cell types, given their propensity to migrate from the lymph to the liver during TCDD administration.

The NeurIPS dataset comprises single-cell gene expression profiles obtained from human peripheral PBMCs treated with 144 compounds from the LINCS Connectivity Map database⁹. Measurements were conducted 24 hours after treatment, with the experiment replicated using cells from three healthy human donors. The full dataset was utilized, except for drug conditions in any cell type with fewer than 30 associated samples.

The LINCS dataset comprises 1,034,271 samples of 122,974 distinct combinations of 83 cell lines treated with 20,544 different drugs. However, the majority of cell line–drug combinations are only sparsely sampled (Supplementary Fig. 7a). We therefore excluded combinations with fewer than 300 samples to reduce potential bias due to sampling, which yielded 65,585 samples from 57 pairs (21 cell lines, 11 perturbations and DMSO as the control). Control–treatment pairs were assigned randomly between DMSO-treated control cells and perturbed cells for a given cell type, a method that has been shown to be successful previously¹⁹.

We processed all the single-cell datasets through a series of steps: removing mitochondrial and ribosomal genes; filtering for cells with at least 1,000 counts and for genes detected in fewer than 50 cells; normalizing the count matrix using scanpy.pp.normalize_total; and applying a $\log (x+1)$ transformation.

For each dataset, we constructed a set of cell-type-specific graphs. The initial gene–gene correlation matrix contained 5,000 × 5,000 edges. After applying a threshold, we retained between 240,000 and 250,000 edges per cell type. Supplementary Table 1 provides additional details for each dataset.

For the transfer learning experiments, we used a subset of the Tahoe dataset, which initially contained 50 unique cell lines from various tissues and 380 distinct compounds. From this dataset, we selected a subset corresponding to the 18 drugs shared with the NeurIPS dataset. To ensure tissue diversity, we identified the cell line with the highest number of samples from each tissue type and subsampled these selected cell lines for testing. To define a shared gene set for transfer learning, we identified the genes common to both datasets and ranked them based on their variability scores computed using Seurat. We then averaged the ranks across the datasets and selected the top 5,000 genes as a joint feature set for downstream analysis. Supplementary Note 1 provides the pre-processing steps in more detail.

Reporting summary

Further information on the research design is available in the Nature Portfolio Reporting Summary linked to this article.