GRAPE: Heterogeneous Graph Representation Learning for Genetic Perturbation with Coding and Non-Coding Biotype

1 Introduction

Single-cell genetic perturbation refers to the process of knocking out specific genes in single cells using CRISPR technology [Barrangou and Doudna, 2016; Lino et al., 2018], thereby altering their normal physiological state or behavior. Despite significant advances in single-cell technology and high-throughput screening, it remains impractical to perform perturbation experiments on all genes within the vast genome of a cell. Accurate prediction of the effects of untested genetic perturbations helps reveal the specific roles of target genes in cellular biological processes and uncovers the regulatory relationships and network structures between genes. Using predictive models, potentially important genes can be identified before experiments, improving the efficiency of the utilization of experimental resources.

Generative models have been widely applied in perturbation modeling, with existing approaches focusing on modeling the latent distribution of data to generate gene expression changes under various perturbations. SAMS-VAE [Bereket and Karaletsos, 2024] enhances model interpretability by using sparse global latent variables to disentangle the specific effects of perturbations. GraphVCI [Wu et al., 2022], on the other hand, leverages graph convolutional networks to model gene regulatory networks (GRN) and predict gene expression under counterfactual perturbations. In contrast, GEARS [Roohani et al., 2022] constructs a graph representation learning model based on gene pathways, ultimately producing prediction based on gene embeddings.

However, these methods have several limitations. 1. Lack of Biological Information: They do not fully leverage the rich biological information available in genes, such as detailed gene descriptions and DNA sequences, which provide valuable insights into the functional roles of genes. 2. Failure to Capture Implicit Gene Interactions: Relying on simplistic evaluation metrics, these methods fail to capture the complex nature of gene regulation, where genes interact in diverse and multifaceted ways, hindering the accurate modeling and prediction of gene behavior across various biological contexts. 3. Ignoring Biotype-Specific Gene Functions: Coding genes (c-genes) are transcribed into mRNA and translated into proteins, while non-coding genes (nc-genes) are transcribed into RNA but do not produce proteins. Biotype refers to the classification of genes into coding and non-coding categories based on their functional roles, with the functional dimensions of genes varying across biotypes [Dykes and Emanueli, 2017; B¨ohmdorfer and Wierzbicki, 2015; Magistri et al., 2012; Pankiewicz et al., 2021]. These methods overlook these differences, limiting their ability to model gene interactions effectively.

To address these issues, we propose GRAPE (Heterogeneous Graph Representation LeArning for Genetic PErturbation) to predict genetic perturbation, which utilizes pre-trained large language model and DNA sequence model for gene feature initialization. Furthermore, GRAPE constructs a GRN as heterogeneous graph to capture the distinct roles of different biotypes and automatically learn the interactions between genes via graph structure learning. Experiments on publicly available datasets show that GRAPE outperforms existing methods.

The main contributions of our work are as follows:

• Perspective: We introduce GRAPE, the first model to incorporate cellular gene biotype information (coding and non-coding genes), to model gene regulatory networks as heterogeneous graphs. This approach captures the distinct functional roles of genes from different biotypes through heterogeneous graph neural networks.
• Algorithm: We initialize gene representations by leveraging pre-trained models [Kenton and Toutanova, 2019; Nguyen et al., 2024] to extract multi-modal features of genes (textual descriptions and DNA sequences), which are then semantically aligned and fused to generate richer and more comprehensive representations. Besides, in constructing the gene interaction network, we propose a novel graph structure learning (GSL) method to capture implicit gene relationships.
• Experiments: We demonstrate the superiority of GRAPE over existing methods on publicly datasets, using a range of comprehensive evaluation metrics.

2 RelatedWork

2.1 Graph Structure Learning (GSL) and Heterogeneoous Graph Neural Network (HGNN)

GSL optimizes the graph's structure to enhance information flow, improve node representations.[Franceschi et al., 2019] jointly learn the graph structure and GCN parameters; SLAPS ([Fatemi et al., 2021]) uses self-supervision to better infer graph structures for GNNs; HGSL([Zhao et al., 2021]) generating relation subgraphs based on feature similarity, propagation, and semantics. HGNNs learn relationships between different node types and generate node representations. HetGNN [Zhang et al., 2019] uses a random walk strategy and two aggregation modules to generate embeddings; HAN [Wang et al., 2019] uses hierarchical attention mechanisms at both node and semantic levels; SeHGNN

[Shi, 2022] reduces complexity by using a lightweight mean aggregator and avoids excessive attention mechanisms. Despite their individual advantages, there remains room for improvement in capturing more complex and implicit relationships within heterogeneous graphs.

2.2 Genetic PerturbationModel

Gene perturbation is a essential direction in single-cell perturbation research. The Variational Autoencoder (VAE) framework and Graph Neural Networks (GNNs) have emerged as the dominant approaches in this filed of study. scGen [Lotfollahi et al., 2019] is the first to use deep learning methods to model single-cell perturbations, utilizing variational autoencoders to capture cellular responses to perturbations in single-cell gene expression; GEARS [Roohani et al., 2022] introduces ontology modality to model gene-gene interaction relationships; CellOracle [Bunne et al., 2023] leverages generegulatory networks from single-cell multi-omics data to simulate transcription factor perturbations; SAMS-VAE [Bereket and Karaletsos, 2024] incorporates sparse global latent variables to enhance model interpretability; GraphVCI [Wu et al., 2022] propose a novel graph-based Bayesian causal inference framework to predict gene expression responses. These methods have several limitations. They fail to fully leverage available biological information, which provide valuable insights into gene functions. Moreover, relying on simplistic evaluation metrics, they cannot capture the complex nature of gene regulation, hindering accurate predictions. Additionally, overlooking the functional differences between genes of different biotypes hinders their ability to effectively model gene interactions.

3 Preliminaries

In this section, we expound some basic concepts and give some definitions of our model.

Definition 1. Heterogeneous Graph We introduce an undirected heterogeneous graph G = (N,R,X), where N is node set, along with node type mapping function φ : N → T = {c-gene, nc-gene}, R is edge set, along with edge type mapping function ϕ : R → E = {rcc, rcn, rnn}, X is node feature matrix.

Definition 2. Model Input and Output The model takes as input gene textual descriptions and DNA sequence information, along with the perturbation conditions. As illustrated in Fig.1, the task of our model is to predict the change in single-cell expression levels under a given perturbation condition (e.g., knockdown of gene k).

4 Method

In this section, we introduce the proposed model GRAPE. Figure.2.A illustrates the workflow for initializing gene features using sequence and textual descriptions by DNA sequence model [Nguyen et al., 2024] nad LLM [Kenton and Toutanova, 2019], respectively. Figure.2.B presents the methodology for constructing gene interaction networks and performing heterogeneous graph representation learning.

Figure 2: Overview of the GRAPE.

4.1 Gene Representation Initialization

Existing methods often initialize gene representations based on gene indices, which lack biological prior information and interpretability, potentially increasing the cost of model convergence. Therefore, we propose multi-modal biological prior information for representation initialization, enabling a more effective capture of the intrinsic gene characteristics. Building upon the pre-trained models [Nguyen et al., 2024; Kenton and Toutanova, 2019], we align gene sequences with their textual descriptions to effectively integrate and capture the shared semantic information between them.

Formally, we extract DNA sequence features using [Nguyen et al., 2024] and textual description features using [Kenton and Toutanova, 2019], resulting in the multi-modal features matrix Xm = [x1m, x2m, ..., xNm] ∈ RN×dm , where m ∈ {seq, text} denotes either the sequence or the textual description modality, with a dimensionality of dm. Our goal is to ensure that the semantic integrity of each modality is preserved while achieving a unified feature representation Zm = f(Xm) = [z1m, z2m, ..., zNm ]T ∈ RN×qm after the alignment, where f represents the transformation process.

Specifically, we aim to align the features extracted from different modalities for the same gene, while ensuring that cross-modal features from different genes are pushed apart. We achieve this using contrastive learning [Oord et al., 2018], and the corresponding formula is as follows:

where D is the discriminator function that estimates the cosine similarity between the features after transformation. This indirect contrastive learning approach helps mitigate the disruption of the original semantics of each modality during alignment. After modality alignment, we obtain the multi-modal fusion features through a multi-layer perceptron (MLP) Ψ:

4.2 Gene Regulation Network Construction

Based on the multi-modal features previously extracted, we can construct a graph for each modality m, where m ∈ M = {seq, text, fus}, according to the similarity of the characteristics between genes, thus revealing gene associations from different perspectives. Specifically, each gene is treated as a node in the graph, with multiple modality-specific representations. The edge between node zim and node zjm is obtained by:

where ϵm is used to control the sparsity of the graph corresponding to the m modality. Sm is a multi-head similarity module that computes multiple similarity scores using different heads with independent parameters. The definition is as follow:

where W tm denotes the learnable parameters of the t-th similarity head in the m modality. After the above graph construction steps and edge filtering, we obtain the gene interaction network structures Gm from the perspective of each modality. To aggregate the Gene Regulatory Networks (GRNs) Gagg from different modalities, we can express the aggregation as the mean across all modalities:

However, graphs constructed using simple similarity-based methods tend to capture only coarse-grained neighbor relationships, making it challenging to comprehensively model the complex and latent interactions between nodes. In realworld scenarios, genes often exhibit latent interactions across multiple scales, necessitating a more flexible and holistic mechanism to model their associations. To address this, we introduce the K-hop GCN aggregation mechanism, which aggregates the results from different K-hop convolutional layers to integrate multi-scale structural information, enabling a more effective representation of the intricate gene interactions. In practice, we aggregate the node relationships across different hop distances as follows:

when l = 0, the feature matrix H(l) is typically initialized as the initial node fusion representations Zfus (Eq. (2)), with GCN serving as the backbone for feature aggregation [Kipf and Welling, 2016].

Then the K-hop structure is defined as:

where Φ denotes a channel attention layer. The edge set E are selected to capture the most significant relationships between nodes. To maintain the structural consistency between the original graph Gagg and the k-hop graph G, the regularization (Eq. (10)) is proposed, ensuring that the model does not deviate from the local structural information of the original graph while learning higher-order graph structures.

4.3 Heterogeneous Graph Representation Learning

As mentioned earlier, c-genes and nc-genes regulate the cellular state at different level [Dykes and Emanueli, 2017; B¨ohmdorfer and Wierzbicki, 2015; Magistri et al., 2012; Pankiewicz et al., 2021]. Therefore, it is necessary to construct a gene interaction network using a heterogeneous graph based on the gene biotype. To simplify the problem, we regard the heterogeneous graph as an undirected graph.

Based on the graph construction before, we map the nodes using φ and the edges using ϕ as defined in Definition.1. Given the different types of associations between genes, we design a multi-head heterogeneous graph attention network (HGAT) to aggregate these features:

where H represents number of head, and P (0) = Zfus is initialized as the fusion gene embeddings with perturbation information, which will be elaborated on in the following section. The GAT used for feature aggregation can be found in [Brody et al., 2021; Veliˇckovi´c et al., 2017].This framework enables the model to actively learn the weights of different types of edges by leveraging the attention mechanism. In the context of a heterogeneous graph, where various edge types represent distinct gene interactions, the model can assign independent parameters to each type of relationship. By doing so, it allows the model to prioritize the most relevant interactions for each specific task, rather than treating all the edges equally.

4.4 Genetic Perturbation Prediction

During training process, GRAPE initializes gene embeddings Zm and constructs gene regulation network G at first. Given a gene perturbation condition c, which biologically corresponds to the knockout of a specific gene , we incorporate this perturbation information into the original input in multi-head HGAT to simulate its effects. Specifically, this involves masking the representations of the corresponding knockout genes in Zfus, ultimately resulting in the final output Pc = {p1c , p1c , . . . , pNc } ∈ RN×d of Eq. (11).

Considering that each gene has its own unique expression pattern [Roohani et al., 2022], we apply an independent linear transformation to each gene to obtain the final output under perturbation c:

where ⊙ denotes Hadamard Product. Ultimately, the prediction of gene expression can be obtained by:

where ctrl represents the expectation of gene expression levels from the control group cells (non-knockout).

To optimize the GRAPE's performance, we compute the mean squared error (MSE) between the predicted and true gene expression values under condition c:

where gc ∈ RN×1 represents the observed gene expression value sampled under perturbation condition c. Moreover, to control the direction of gene expression changes, we use the Huber loss function, which combines MSE's sensitivity to small errors with MAE's robustness to large errors, reducing the impact of outliers and ensuring more stable training in noisy data.

The reconstruction objective of the model for gene expression values is represented as:

4.5 Training Steps

During model training, we optimize the above objectives simultaneously, and the final training objective is defined as:

here, Lalignis specifically optimized during the pre-training phase before the main model training.

4.6 Analysis of Time Complexity

To simplify the analysis, we assume the feature dimension has no impact and treat computations as O(1). The computational cost of the model consists of four main components: multihead attention O(H1 · N 2 ), K-hop similarity (O(K · m)), channel attention O(N 2 ), and feature extractions from multihead HGAT O(L · H2 · (|E| + N)). Here, N denotes the number of genes, H1 represents head number in Eq. (4), m stands for the number of edges in the result of Eq. (3), with L and H2 referring to the layer depth and the number of attention heads in Eq. (11), respectively.

5 Results and Discussion

To demonstrate the effectiveness of our method, we conduct extensive experiments on public available datasets. Adamson dataset [Adamson et al., 2016] contains data from 87 types of single-gene perturbations, with a total of 68, 603 singlecell RNA sequencing results. Norman dataset [Norman et al., 2019] includes results from both single-gene and double-gene perturbations. We only select experimental data for singlegene perturbations, with over 100 types of single-gene perturbations in total. In each dataset, hundreds of single-cell RNA sequencing data can be obtained for each type of genetic perturbation.

5.1 Experiment Settings

Gene text descriptions, sequences, and biotypes can all be accessed through the API provided by Ensemble (Click here to visit the tutorial). In the training process of GRAPE, we randomly select 70% of gene perturbation conditions for the training set and use the remaining for testing, then select single-cell data according to these conditions. The number of heads in both the multi-head similarity module (Eq. 4) and the multi-head HGAT module (Eq. 11) is set to 4. The model is trained using the AdamW [Loshchilov, 2017] optimizer with a learning rate of 0.001 and a batch size of 64. For evaluation, we adopt multiple metrics, including Direction Accuracy Coefficient (DAC), Root Mean Square Error (RMSE), and Pearson Correlation Coefficient delta (PCC delta), to assess the model's prediction accuracy and biological relevance (refer to the github for further details). All our method and its competitors are conducted using four Nvidia A100 GPUs.

5.2 GRAPE outperform existing methods

Table 1 demonstrates that GRAPE significantly outperforms existing gene perturbation methods ([Roohani et al., 2022; Wu et al., 2022; Bereket and Karaletsos, 2024]) across all evaluated metrics. While Root Mean Square Error (RMSE) measures the overall prediction error, DAC and PCC delta provide deeper biological insights. DAC evaluates the consistency between predicted and true directions of gene expression changes, which is crucial for understanding gene regulatory effects and ensuring accurate biological interpretations, as incorrect directional predictions can lead to misleading conclusions. PCC delta, on the other hand, quantifies the consistency in the change of gene expression values between the predicted and true values under perturbation conditions. Together, these metrics offer a comprehensive assessment of model accuracy and relevance in biological contexts.

GRAPE showcases superior performance compared to GEARS [Roohani et al., 2022], primarily because it goes beyond GEARS' simplified gene interaction network by accounting for the diverse relationships across different biotype genes. GRAPE actively uncovers latent connections between genes, providing a deeper understanding of gene regulatory dynamics. In contrast, generative models [Wu et al., 2022; Bereket and Karaletsos, 2024] perform poorly due to their inability to effectively capture the semantic meaning of gene perturbation conditions. This limitation hinders their ability to model biologically relevant features, resulting in less accurate predictions. Moreover, GRAPE leverages multimodal information to initialize gene representations, providing a more comprehensive and biologically meaningful characterization of genes.

To further evaluate the model's performance, we visualize the distribution of DAC and PCC values across samples (Fig.4). These plots provide a detailed view of the prediction accuracy at the sample level, highlighting the variability and consistency of the model's predictions. By focusing on these key metrics, the violin plots offer an intuitive representation of how well the model captures the directionality and correlation of gene expression changes under different perturbations.

5.3 GRAPE Enhances Gene Representation with Biotype-Aware Heterogeneous Graphs

To evaluate GRAPE leveraging biotype knowledge, we first visualize gene features extracted from pre-trained LLM [Kenton and Toutanova, 2019] and DNA sequence model [Nguyen et al., 2024], focusing on two modalities: gene textual descriptions (Fig.3.A) and gene sequences (Fig.3.B). For this, we employe Uniform Manifold Approximation and Projection (UMAP) to visualize the high-dimensional gene features in a lower-dimensional space. The visualizations of Fig.3.A and Fig.3.B reveal that the clustering of gene representations did not align with the biotype labels. Despite the inherent complexity and richness of these features, the unsupervised clustering approach failed to exhibit any clear, biotypespecific grouping. This suggests that the pre-trained models, although capable of extracting meaningful features, does not sufficiently capture the underlying biological relevance of gene biotypes on its own.

In contrast, after training the GRAPE, which incorporates biotype-aware heterogeneous graphs into the representation learning process (Eq. (11)), the extracted gene features exhibited a distinct biotype-based clustering pattern (Fig.3.C). The representations of genes that belong to the same biotype formed tight and well-separated clusters, indicating that the model successfully learned to embed gene biotype information into the feature space. This outcome highlights the efficacy of GRAPE in improving gene representation by dynamically adjusting the gene regulatory network and then extracting features through HGAT that incorporates biotype-aware information.

Constructing a heterogeneous graph network effectively represents the complex relationships between genes and simulates interactions between different biotype genes, thereby enhancing gene representation through HGAT. On the other hand, when biotype is directly concatenated as a feature to the gene representation, the model may fail to fully explore the underlying connections between genes based on their biotype. This approach only provides a singular biotype feature, lacking interactions with other gene attributes and fails to leverage the diversity of nodes and edges in the GRN. As a result, the model may not capture the biological significance of genes in their multidimensional relationships. Additionally, directly incorporating the biotype label as a feature could lead to over-reliance on the biotype information, potentially neglecting other crucial biological features, which could in turn reduce the model's generalization ability.

5.4 Ablation Study

To further evaluate the effectiveness of GRAPE, we compare it with the following methods through an ablation study. 1)only seq: Excludes the textual description modality features, using only the DNA sequence features in the model. 2)only desc: Excludes the DNA sequence modality features, using only the textual description features in the model. 3)only fusion: Uses the fusion modality, combining both desc and seq features into a unified representation, excluding the individual modalities. 4)w/o GSL: Excludes the K-hop similarity graph structure learning process. 5)w/o biotype: Ignores biotype information, converting the GRN into a homogeneous graph, and uses a homogeneous GNN instead of HGNN. The results are shown in Figure.5.

The experimental results show that GRAPE outperforms methods relying on individual modalities. This indicates that integrating multiple modalities allows the model to capture a broader range of information, leveraging complementary insights from each modality for improved performance and more robust feature extraction. Additionally, the performance of no GSL is less than ideal, highlighting the significant role that graph structure learning plays in capturing the hidden relationships between genes. Without graph structure learning, the model fails to fully leverage the intricate connections and dependencies inherent in the gene interactions, which are essential for accurate representation and prediction.

More importantly, the results clearly demonstrate the effectiveness of introducing biotype information. By leveraging HGNN, the model is able to capture and simulate the complex, multi-dimensional regulatory relationships between genes. This approach not only allows for a more nuanced understanding of gene interactions across different biological categories, but also enhances the model's ability to predict gene behavior with greater accuracy.

6 Conclusion

In this work, we present GRAPE, a heterogeneous graph neural network (HGNN) that integrates gene biotype information for predicting genetic perturbations. By leveraging pre-trained large language model (LLM) and DNA sequence model to extract features from gene textual descriptions and DNA sequences, we initialize gene representations in a way that captures both sequence-based and textual semantic information. Furthermore, for the first time, we incorporate gene biotype data to simulate the distinct roles of different biotypes in regulating cellular processes. Through graph structure learning (GSL), GRAPE dynamically refines the gene regulatory network (GRN), effectively capturing implicit gene relationships. Our method outperforms existing approaches on publicly available datasets, achieving state-of-the-art performance in predicting genetic perturbations. If sufficient data is available, it can enable a zero-shot model on top of GRAPE, automatically constructing gene regulatory networks (GRN) and inferring gene interactions to predict genetic perturbation results without human annotations.

Acknowledgments

This work was supported by National Science and Technology Major Project (No. 2022ZD0115101), National Natural Science Foundation of China Project (No. 623B2086), National Natural Science Foundation of China Project (No. U21A20427), Project (No. WU2022A009) from the Center of Synthetic Biology and Integrated Bioengineering of Westlake University, Project (No. WU2023C019) from the Westlake University Industries of the Future Research Funding.

Contribution Statement

Changxi Chi and Jun Xia contributed equally to this work. Stan Z. Li supervised the project.