Impeller: a path-based heterogeneous graph learning method for spatial transcriptomic data imputation

Abstract

Motivation

Recent advances in spatial transcriptomics allow spatially resolved gene expression measurements with cellular or even sub-cellular resolution, directly characterizing the complex spatiotemporal gene expression landscape and cell-to-cell interactions in their native microenvironments. Due to technology limitations, most spatial transcriptomic technologies still yield incomplete expression measurements with excessive missing values. Therefore, gene imputation is critical to filling in missing data, enhancing resolution, and improving overall interpretability. However, existing methods either require additional matched single-cell RNA-seq data, which is rarely available, or ignore spatial proximity or expression similarity information.

Results

To address these issues, we introduce Impeller, a path-based heterogeneous graph learning method for spatial transcriptomic data imputation. Impeller has two unique characteristics distinct from existing approaches. First, it builds a heterogeneous graph with two types of edges representing spatial proximity and expression similarity. Therefore, Impeller can simultaneously model smooth gene expression changes across spatial dimensions and capture similar gene expression signatures of faraway cells from the same type. Moreover, Impeller incorporates both short- and long-range cell-to-cell interactions (e.g. via paracrine and endocrine) by stacking multiple GNN layers. We use a learnable path operator in Impeller to avoid the over-smoothing issue of the traditional Laplacian matrices. Extensive experiments on diverse datasets from three popular platforms and two species demonstrate the superiority of Impeller over various state-of-the-art imputation methods.

Availability and implementation

The code and preprocessed data used in this study are available at https://github.com/aicb-ZhangLabs/Impeller and https://zenodo.org/records/11212604.

1 Introduction

The orchestration of cellular life hinges on the precise control of when and where genes are activated or silenced. Characterizing such spatiotemporal gene expression patterns is crucial for a better understanding of life, from development to disease to adaptation (Mantri et al. 2021). While single-cell RNA sequencing (scRNA-seq) is a revolutionary and widely available technology that enables simultaneous gene expression profiling over thousands of cells, it usually needs to dissociate cells from their native tissue and thus loses the spatial context (Lähnemann et al. 2020). Recent advances in spatial transcriptomics (Ståhl et al. 2016) allow spatially resolved gene expression measurements at a single-cell or even sub-cellular resolution, providing unprecedented opportunities to characterize the complex landscape of spatiotemporal gene expression and understand the intricate interplay between cells in their native microenvironments (Strell et al. 2019). However, due to technical and biological limitations, most spatial transcriptomic profiling technologies still yield incomplete datasets with excessive missing gene expression values, hindering our biological interpretation of such valuable datasets (Choe et al. 2023). Therefore, gene imputation is a critical task to enrich spatial transcriptomics by filling in missing data, enhancing resolution, and improving the overall quality and interpretability of the datasets.

Several methods have been successfully developed for gene imputation in spatial transcriptomics, which can be broadly summarized into two categories—reference-based and reference-free approaches. Since scRNA-seq data usually offer a deeper dive into transcriptome profiling, reference-based methods integrated spatial transcriptomic data with matched scRNA-seq data from the same sample for accurate imputation. While promising, these referenced-based methods usually suffer from two limitations. First, most studies do not always have matched scRNA-seq data, especially those using valuable and rare samples. Second, even with matched data, there can be significant gene expression distribution shifts due to sequencing protocol differences (e.g. single nuclei RNA-seq versus whole cell spatial transcriptomics) (Zeng et al. 2022).

Researchers also used reference-free methods for direct gene expression imputation. For instance, traditional gene imputation methods designed for scRNA-seq data, such as scVI (Lopez et al. 2018), ALRA (Linderman et al. 2018), Magic (van Dijk et al. 2018), and scGNN (Wang et al. 2021), have been adapted for spatial transcriptomic data imputation. While effectively capturing cell-type-specific gene expression signatures, these methods completely ignored the rich spatial information, resulting in suboptimal results. Later, scientists emphasized the importance of spatial context for cell-to-cell interaction (CCI) in modulating expression changes in response to external stimuli (Armingol et al. 2021). Therefore, Graph Neural Network (GNN) based methods have been developed to mimic CCIs for imputation tasks with improved performance. However, different types of CCI involve distinct cell signaling mechanisms with varying interaction ranges. Existing GNN-based methods used very shallow convolutional layers for computational convenience, successfully modeling short-range CCI (e.g. via autocrine and juxtacrine) but ignoring long-range interactions (e.g. via paracrine and endocrine). As a result, they cannot fully exploit the spatial information for gene expression imputation.

To address the abovementioned issues, we propose Impeller, a path-based heterogeneous graph learning method for accurate spatial transcriptomic data imputation. Impeller contains two unique components to exploit both transcriptomic and spatial information. First, it builds a heterogeneous graph with nodes representing cells and two types of edges describing expression similarity and spatial proximity. Therefore, the expression-based edges allow it to capture cell-type-specific expression signatures of faraway cells from the same type, and the proximity-based edges incorporate CCI effects in the spatial context. Second, Impeller models long-range CCI by stacking multiple GNN layers and uses a learnable path operator instead of the traditional Laplacian matrices to avoid the over-smoothing problem. Extensive experiments on diverse datasets from three popular platforms and two species demonstrate the superiority of Impeller over various state-of-the-art imputation methods.

Our main contributions are summarized below:

• We propose a graph neural network, Impeller, for reference-free spatial transcriptomic data imputation. Impeller incorporates cell-type-specific expression signatures and CCI via a heterogeneous graph with edges representing transcriptomic similarity and spatial proximity.

• Impeller stacks multiple GNN layers to include both short- and long-range cell-to-cell interactions in the spatial context. Moreover, it uses a learnable path-based operator to avoid over-smoothing.

• To the best of our knowledge, this is the first paper to combine cell-type-specific expression signatures with spatial short- and long-range CCI for gene expression imputation.

• We extensively evaluate Impeller alongside state-of-the-art competitive methods on datasets from three sequencing platforms and two species. The results demonstrate that Impeller outperforms all of the baselines.

2 Related work

2.1 Imputation methods ignoring spatial information

Earlier spatial transcriptomic data imputation methods adapted the computational strategies originally developed for scRNA-seq data, overlooking the spatial coordinate information of each spot. For instance, eKNN (expression-based K nearest neighbor), and eSNN (expression-based Shared nearest neighbor) are methods implemented using the Seurat R-package that rely on gene expressions of nearest neighbors. MAGIC adopted data diffusion across similar cells to impute missing transcriptomic data. ALARA used low-rank approximation to distinguish genuine nonexpression from technical dropouts, thus preserving true gene absence in samples. scVI used a deep variational autoencoder for gene imputation by assuming the read counts per gene follow a zero-inflated negative binomial distribution. However, these methods completely ignored the rich spatial information, resulting in sub-optimal performance.

2.1 Imputation methods utilizing spatial information

Later on, several methods were developed to exploit the spatial coordinate information to improve imputation accuracy. Since scRNA-seq data are usually sequenced deeper to provide more accurate expression measurements, several methods incorporated additional scRNA-seq data during the imputation process. For instance, gimVI used a low-rank approximation and included scRNA reference (Lopez et al. 2019). Tangram mapped scRNA-seq data onto spatial transcriptomics data to facilitate imputation by fitting expression values on the shared genes (Biancalani et al. 2021). STLearn used gene expression data, spatial distance, and tissue morphology data for imputing absent gene reads (Pham et al. 2020). However, additional scRNA-seq data are not always available and there can be large gene expression distribution shifts between these datasets due to differences in sequencing protocols (e.g. single-cell versus single-nuclei), resulting in limited applications for reference-based methods.

On the other hand, several reference-free methods have been developed for more generalized settings. For example, the seKNN (spatial-expression-based K nearest neighbor) and seSNN (spatial-expression-based shared nearest neighbor) models (Satija et al. 2015, Butler et al. 2018, Stuart et al. 2019, Hao et al. 2021) incorporated cell-to-cell distance when defining the KNN for imputation tasks. Recently, STAGATE (Dong and Zhang 2022) is a graph attention auto-encoder framework that effectively imputes genes by integrating spatial data and cell type labels. Overall, these methods did not deeply integrate and exploit the full potential of combining expression and spatial data.

3 Materials and methods

3.1 Problem definition

Here, we aim to impute the excessive missing gene expression values in spatial transcriptomics data without matched reference scRNA-seq data. Formally, given a sparse cell-by-gene count matrix $Xobs∈Rn×m$ which represents observations for n cells across m genes, and the spatial coordinates $C∈Rn×2$ of these cells, our goal is to impute the gene expression matrix $X^∈Rn×m$⁠. $Xobs$ is derived from the ground truth matrix $Xgt∈Rn×m$⁠, which contains the observed nonzero entries pre-masking. To simulate real-world data conditions, 10% of the nonzero entries in $Xgt$ are masked to form a test set and another 10% for validation, thus creating $Xobs$⁠. This matrix serves as the input for our imputation model. The major challenge is to generate $X^$ that is as close as possible to the ground truth gene expression $Xgt$⁠, using both the observed gene expressions in $Xobs$ and the spatial information in $C$⁠.

3.2 Heterogeneous graph construction

As shown in Fig. 1, we build our Impeller model based on two widely accepted biological insights–(i) gene expression can be modulated by surrounding cells via CCI; (ii) faraway cells of the same cell type may share stable gene expression signatures. Therefore, Impeller first builds a heterogeneous graph G to fully exploit both spatial and cell-type information, with nodes and edges representing cells and their relationships.

Specifically, G contains two complementary graphs: a spatial graph (⁠$Gs$⁠) and a gene similarity graph (⁠$Gg$⁠). Edges in $Gs$ represent the cell’s spatial proximity to model CCI, while edges in $Gg$ denote the cell’s transcriptomic similarity to capture the cell-type-specific expression signatures.

3.2.1 Spatial graph construction

3.2.2 Gene similarity graph construction

3.3 GNN model on heterogeneous graph

With the heterogeneous graph built, Impeller uses a path-based heterogeneous GNN to synthesize the impacts of spatial CCI (⁠$Gs$⁠) and cell-type-specific expression signatures (⁠$Gg$⁠) for the imputation task. We introduce the problem of traditional GNN, our learnable path operator, and the overall architecture of Impeller as follows.

3.3.1 Problem of traditional GNN

In order to capture long-range CCI interaction, we have to include relatively far away cells by stacking multiple GNN layers via a larger L. Traditional Laplacian matrices-based GNN suffers from over-smoothing, resulting in deteriorated performance as L increases (Eliasof et al. 2022). Therefore, we introduce a learnable path operator to overcome this issue and better capture the long-range CCI.

3.3.2 Learnable path operator

3.3.3 The overall architecture of impeller

3.4 Computational complexity analysis

3.4.1 k-hop complexity analysis

Traditional GNNs need to gather information from k-hop neighbor nodes after stacking of k layers. Given the complexity of each layer as $O(n×dt)$⁠, where n is the number of nodes and d_t is the average node degree, the overall complexity becomes $O(n×dt×k)$⁠. In contrast, Impeller can directly access neighbors up to k-hop distance via a single layer by setting $ks=kg=k$⁠. The computational complexity per layer for Impeller is $O(n×(Ts×ks+Tg×kg))$⁠, with T_s and T_g representing the number of paths in $Gs$ and $Gg$⁠, k_s and k_g denoting path lengths. As a result, when $Ts<dt$ (a condition satisfied in our task), Impeller offers superior computational efficiency.

3.4.2 The number of parameters

Traditional GNNs have $O(demb(l)×demb(l+1))$ parameters per layer while the path operator of Impeller adds $(ks+kg)*demb(l)$ parameters. Since $ks+kg$ is typically much smaller than $demb(l+1)$⁠, Impeller’s number of parameters remains on par with traditional GNNs.

4 Experiments

4.1 Detailed experimental setup

4.1.1 Data sources and preprocessing

In our study, we tested Impeller using diverse datasets from three popular sequencing platforms and two organisms. Specifically, we included 10X Visium datasets from the human dorsolateral prefrontal cortex (DLPFC), (Maynard et al. 2021), Steroseq datasets from mouse olfactory bulb (Chen et al. 2021), and Slide-seqV2 from mouse olfactory bulb (Stickels et al. 2021) in our analyses. Detailed attributes of these datasets are summarized in Table 1 (for filter details and visualizations, see the Supplementary Appendix). After standard pre-processing and normalization procedures, we downsampled the data according to scGNN, where 10% of nonzero entries in the dataset were used as a test set, and another 10% of nonzero entries were reserved for validation. For a fair comparison, we repeat ten times with different mask configurations.

4.1.2 Baseline methods for benchmarking

We conducted a comparative study utilizing 12 state-of-the-art methods, including reference-free and reference-based methods that originally required additional scRNA-seq data. However, in our analysis, we did not use any additional scRNA-seq data for a fair comparison.

First, we included methods directly adapted from scRNA-seq data imputation and completely ignored the rich spatial information, including a deep generative model scVI, a low-rank approximation model ALRA, nearest neighbors-based models eKNN and eSNN, a diffusion-based model MAGIC, and a GNN-based model scGNN. Furthermore, we used several imputation methods specifically designed for spatial transcriptomic data, such as seKNN (spatial-expression-based K nearest neighbor), and seSNN (spatial-expression-based shared nearest neighbor). gimVI and Tangram need additional scRNA-seq from matched samples, so we used a reference-free implementation available through their website for a fair comparison. Lastly, we included STAGATE a graph attention auto-encoder framework by amalgamating spatial data and gene expression profiles. We use default parameters in most baseline methods (for details, see the Supplementary Appendix).

4.1.3 Evaluation metrics

4.2 Experimental results

4.2.1 Improved imputation accuracy

We benchmarked our performance against 12 leading methods by assessing imputation accuracy across 14 datasets. These datasets span three prominent sequencing platforms (10x Visium, Stereoseq, and Slideseq) and two species (human and mouse). Table 2 summarizes the performance of Impellers and other baselines (for results of the other six samples of the DLPFC dataset, please see the Supplementary Appendix). For a fair comparison, we did not include any additional scRNA-seq data to facilitate the imputation task. Overall, Impeller consistently outperforms others in all datasets using L1 distance, Cosine Similarity, and RMSE, indicating the effectiveness and robustness of our strategy.

In addition, we found that most methods utilizing spatial information (w* group in Table 2) demonstrated higher imputation accuracy than those ignoring spatial information (wo* group in Table 2), validating the presence of rich information in the spatial context. Notably, Impeller surpasses even the best gene expression-only method, eKNN, with improvements of 11.32% on 10xVisium DLPFC, 31.09% on Stereoseq, and 6.01% on SlideseqV2 Mouse. Furthermore, compared to uniform averaging using KNN, GNN allows for more flexible neighbor information aggregation for better imputation accuracy, as reflected by the noticeably improved performance of Impeller and STAGATE.

4.2.2 Impact of long-range CCI

To probe disparities between Impeller and traditional GNNs in capturing long-range cell dependencies, we examined several models–Impeller, GCN (Kipf and Welling 2016), GraphSAGE (Hamilton et al. 2017), GAT (Veličković et al. 2017), and GraphTransformer (Shi et al. 2020)–across varying receptive fields in the Stereoseq dataset.

In Table 3, GAT and GraphSAGE suffer from gradient vanishing/exploding issues as more layers are added to capture long-range CCI, resulting in quickly degraded performance. GCN works best initially, but its performance drops with more layers added. This could be because the number of neighbors grows fast as we increase the receptive field, leaving it difficult for the target cell to understand the influence of each neighbor. Furthermore, GraphTransformer starts with high errors at a receptive field of 2. It works best at a receptive field of 8, but the error goes up again at 32. This increase in error is similar to the problem of GCN, as all cells start to look too similar to make useful representations. On the other hand, Impeller effectively tackles these challenges by the path operator, as reflected by the consistently improved results until the receptive field of 32. As the receptive field continues to grow, Impeller’s performance slightly declines, likely because distant information becomes less relevant for the target cell’s gene imputation. An additional perturbation study, demonstrating the effectiveness of Impeller in capturing CCI, is shown in the Supplementary Appendix.

4.2.3 Advantage of heterogeneous graph

In our study, we explored the influence of graph modalities on imputation accuracy by assessing three key variants: var_s, using solely the spatial graph; var_g, utilizing only the gene similarity graph; var_h, integrating both graphs. We then calculated the performance improvement from adding $Gg$ by comparing var_h with var_s, and the improvement from adding $Gs$ by comparing var_h with var_g. As shown in Fig. 2, the majority of the cases (22 out of 24) exhibit positive improvements. Specifically, in the DLPFC sample 151674, the inclusion of the gene similarity graph yields a 17.3% improvement, and the 13.6% enhancement is achieved by adding the spatial graph alone. Similarly, in sample 151508, the gene similarity graph and the spatial graph contribute to improvements of 3.6% and 9.9%, respectively. These results underscore the efficacy of our approach, particularly in scenarios where the complex interaction between spatial and gene expression data is pivotal for enhancing gene imputation accuracy.

4.3 Ablation study

We conducted an ablation study to evaluate the performance of four primary path operator variants of Impeller: $opglo$⁠, where all Impeller layers and channels (each channel representing 1D of $f(l)$⁠) share one path operator; $opcha$⁠, where channels share an operator but layers have distinct ones; $oplay$⁠, where all layers share one, but channels have individual operators; and $opind$ where every layer and channel possesses an independent path operator. As depicted in Fig. 3, both $opglo$ and $opcha$ performed poorly on the DLPFC dataset, indicating the importance of distinct operators for each channel. Notably, $oplay$ and $opind$ showed comparable results, suggesting that layer-specific operators might be optional, depending on the specific application. Another ablation study regarding different path construction and graph construction methods is shown in the Supplementary Appendix.

4.4 Parameter analysis

To investigate the influence of Impeller’s various hyper-parameters, we conducted extensive experiments using the DLPFC dataset (Sample ID: 151507) and reported the mean and standard deviation of the imputation accuracy over ten repetitions.

First, we studied the impact of q_s and q_g on the RMSE of a random walk on $Gs$ and $Gg$ following the Node2Vec mechanism. Higher values of q (i.e. q_s and q_g) encourage the walk to sample more distant nodes, enhancing the exploration of the global graph structure, while lower values bias the walk toward neighboring nodes, facilitating local exploration. As shown in Fig. 4A, Impeller exhibits strong robustness with RMSE from 0.33 to 0.36 when q_s and q_g varied from 0.1 to 5. However, higher values of q_s and q_g tend to induce larger errors. For generality, we selected 1 as the default value for q_s and q_g.

We investigated the impact of random walk length (k_s and k_g) and layer number (L), shown in Fig. 4B. A path length of 2 with 10 layers results in maximum errors, reducing our model to a standard ten-layer GCN. This is because, at this path length, the model focuses on immediate neighbors, akin to how traditional GCNs operate. Such a setup, while deep, limits neighborhood exploration and increases over-smoothing risk. Conversely, a path length of 8 with 4 layers allows for capturing broader interactions (up to 28 hops), balancing extended reach and computational efficiency, thus avoiding over-smoothing and optimizing long-range CCI capture.

Next, we examined the impact of the number of random walks (T_s and T_g). As shown in Fig. 4C, T_s and T_g appeared to have a minimal effect on results, due to the robustness of Impeller which resamples at each epoch during training. We chose 8 as the default number of random walks.

Lastly, we evaluated how the embedding dimension $demb(l)$ affects Impeller’s performance. As shown in Fig. 4D, smaller $demb(l)$ (such as 2, 4, 8) leads to limited expressive power and larger imputation errors. As $demb(l)$ increases to 16, 32, 64, or 128, Impeller’s expressive power improves and operations converge well in each run. Due to our early stopping criterion, we cease training if the validation RMSE does not improve for 50 consecutive epochs. When $demb(l)$ was set to 256 or 512, it’s hard for Impeller to converge quickly at these dimensions. To strike a balance between complexity and representational power, we opted for $demb(l)$ of 64.

In summary, these comprehensive parameter analyses reveal that Impeller is robust across a wide range of parameter settings, while still providing tunable options for balancing computational efficiency and prediction accuracy. These results further substantiate the effectiveness and practicality of our proposed model for gene imputation tasks.

4.5 Neighbor visualization

To better understand the differences between traditional GNNs and our path-based GNN, Impeller, we turned to a visual example (sample 151507 from the DLPFC dataset). Figure 1B shows how the typical GNN gathers information from far-away neighbors. The center node (red sphere) stacks five GNN layers to gather information from distant nodes like the one shown in yellow. But this method sometimes pulls in extra information from different tissue layers that is not needed. On the other hand, Fig. 1C shows our Impeller model. Instead of stacking GNN layers, Impeller samples a direct path from the center node to the target node. While using this direct path method, Impeller offers better gene imputation performance by capturing the relevant long-range CCI.

4.6 Running time analysis

As shown in Table 4, we conducted a comparative model parameter and runtime analysis with popular graph-based models (GCN, GAT, GraphSAGE, and Transformer) on the DLPFC dataset. As discussed in Section 3.4.2, our model maintains a parameter count comparable to traditional GNNs, with the complexity per layer defined as $O(demb(l)×demb(l+1))$⁠. Specifically, our model introduces only a 3.5% increase in parameters for GCN and a 2.7% increase for GAT. In contrast, it achieves a 48.0% reduction in parameters for GraphSAGE and a 74.1% reduction for GraphTransformer (Table 4). Despite its additional path sampling step, Impeller remarkably outperformed the others in training and inference efficiency. This can be partially credited to leveraging the DGL library’s optimized implementation for path sampling (https://docs.dgl.ai/en/0.8.x/api/python/dgl.sampling.html) and the inherently faster multiplication process used in path-based convolution compared to edge-wise information aggregation in traditional GNNs. In addition, Impeller showed the lowest RMSE, indicating superior prediction accuracy. Hence, Impeller offers a balanced blend of efficiency and precision for spatial transcriptomic data imputation, outperforming other graph-based models.

5 Conclusion

In this study, we introduced Impeller, a path-based heterogeneous graph learning approach tailored for spatial transcriptomic data imputation. By constructing a heterogeneous graph capturing both spatial proximity and gene expression similarity, Impeller offers a refined representation of cellular landscapes. Further, its integration of multiple GNN layers, coupled with a learnable path operator, ensures comprehensive modeling of both short and long-range cellular interactions while effectively averting over-smoothing issues. Benchmark tests across diverse datasets spanning various platforms and species underscore Impeller’s superior performance compared to state-of-the-art imputation methods. This work not only establishes Impeller’s prowess in spatial transcriptomic imputation but also underscores its potential to model both short- and long-range cell-cell interactions.

Acknowledgements

The authors thank the anonymous reviewers for their valuable suggestions and the UCI OIT department for GPU resources.

Supplementary data

Supplementary data are available at Bioinformatics online.

Conflict of interest

None declared.

Funding

This work was supported by the National Institutes of Health [R01HG012572, R01NS128523].

Data availability

The code and preprocessed data used in this study are available at https://github.com/aicb-ZhangLabs/Impeller and https://zenodo.org/records/11212604.

References