Bi-channel Masked Graph Autoencoders for Spatially Resolved Single-cell Transcriptomics Data Imputation
Abstract
Spatially resolved transcriptomics bring exciting breakthroughs to single-cell analysis by providing physical locations along with gene expression.
However, as a cost of the extremely high resolution, the technology also results in much more missing values in the data, i.e. dropouts.
While a common solution is to perform imputation on the missing values, existing imputation methods majorly focus on transcriptomics data and tend to yield sub-optimal performance on spatial transcriptomics data.
To advance spatial transcriptomics imputation, we propose a new technique to adaptively exploit the spatial information of cells and the heterogeneity among different cell types.
Furthermore, we adopt a mask-then-predict paradigm to explicitly model the recovery of dropouts and enhance the denoising effect.
Compared to existing studies, our work focuses on new large-scale cell-level data instead of spots or beads.
Preliminary results have demonstrated that our method outperforms existing methods for removing dropouts in high-resolution spatially resolved transcriptomics data.
Introduction
The rapid advance of single-cell technologies has paved the way to investigate biological processes at an unprecedentedly high resolution.
However, most single-cell technologies such as scRNA-seq [38] and scATAC-seq [6] are unable to capture spatial information of the cells, which is often coupled with cell identities.
To spatially resolve transcriptomics profiles, spatial transcriptomic technologies have been recently developed where researchers are able to identify the spatial context of cells and cell clusters within a tissue [7].
Similar to standard single-cell RNA-sequencing, spatially resolved profiling technology may fail to capture a significant number of the expressed genes due to the low
RNA capture rate.
This leads to an artificially large proportion of false zero counts, i e., the so-called “dropout” [36, 20].
A popular way to address this issue is to perform imputation, which corrects false zeros by estimating realistic values for those missing values.
Recently, there has been an explosive growth of imputation methods for single-cell transcriptomics data, focusing on generative probability models, matrix factorization [40, 19, 39, 24] and deep learning models [37, 12, 8, 44, 31, 29].
Although promising results have been achieved on single-cell RNA-sequencing data, it remains challenging to impute spatially resolved single-cell transcriptomics data.
Due to the relatively small library size in each unit and the extra experimental processes to obtain spatial locations, spatial transcriptomics technology suffers from a higher noise level and drop-out rate (zero reads for many genes).
Furthermore, most existing imputation models do not take advantage of spatial information, and it is still unclear how to advance imputation with spatial information.
To leverage the spatial information, a very recently graph-based methods [23, 45] have been designed.
They first build a cell-cell graph based on the spatial position and then apply graph neural networks to the constructed graph.
However, these methods do not consider heterogeneity among different types of cells.
For example, tumor cells usually show strong aggregation, and neighboring tumor cells with similar micro-environment in terms of ligands tend to have higher gene expression similarity than distant cells [5, 22, 27].
By contrast, Treg cells are scarce spatially in many tissues but still tend to have similar gene expression profiles [33].
Hence, it is desired to rely on spatial information to impute cells like tumors, while leveraging information from the same type of cells to impute cells like Treg cells.
To capture such heterogeneity while exploiting spatial information, we propose a novel bi-channel graph neural network (GNN).
Specifically, we first construct two channels: a spatial neighboring graph capturing the spatial information and a dynamic k-nearest neighbor (kNN) graph capturing the similarity of cells.
Then we adopt a gating unit to adaptively fuse information from the two graphs, which adjusts the exploitation of spatial information and heterogeneity among different types of cells.
Furthermore, inspired by the denoising effect of masked graph autoencoders [17] which recovers the manually masked node features through a GNN, we introduce a novel mask-then-predict paradigm for training the proposed imputation method.
Specifically, we randomly mask 60% of genes with zero in each cell and train our bi-channel graph neural network to recover the input expression.
This training paradigm explicitly models the recovery of dropouts with graph neural networks.
We evaluate our method on an unpublished dataset produced by CosMx Spatial Molecular Imager (SMI) [16].
Unlike existing spot-level datasets (e.g. 10x Visium [35]) or tiny cell-level datasets (e.g.
MERFISH [46], Fishseq+ [11]), CosMx platform provides unprecedentedly large-scale single-cell spatial transcriptomics datasets, with more than 100,000 cells in each slice.
Our proposed method achieves state-of-the-art performance on this dataset and outperforms other baseline models.
Our contribution can be summarized as follows:
• We propose a novel bi-channel graph neural network to leverage information from both spatial neighboring cells and similar cells to advance spatial transcriptomics data imputation.
• We introduce a mask-
then-predict paradigm as the training objective to explicitly model the recovery of dropouts in spatial transcriptomics data, which shows outstanding performance.
We make the first attempt to evaluate imputation methods on large-scale single-cell spatial transcriptomics datasets.
Our method improves the data quality and makes downstream tasks more viable while providing a strong baseline for future spatial imputation models.
The rest of the paper is organized as follows.
In Section 2 , we review related work.
We detail our proposed framework in Section 3 and discuss our experimental results in Section 4 .
Finally, we conclude the work with future directions in Section 5 .
Related Work 2.1 Spatial Resolved Transcriptomics
While single-cell technologies can capture transcriptomic information at the single-cell level, most are unable to capture spatial information of the cells due to the isolation procedure.
However, the relationship between cells and their relative locations within tissue is critical to understanding normal development and disease pathology.
Recently, spatial transcriptomic technologies are developed to spatially resolve transcriptomics profiles [35, 28].
With spatial transcriptomics data, researchers can learn the spatial context of cells and cell clusters within a tissue [7].
Most commercial spatially resolved profiling technologies are not at the single-cell resolution that has arisen the problem of cell-type deconvolution and segmentation [35, 28].
The major technologies/platforms for spatial transcriptomics are Visium by 10x [35], GeoMx Digital
Spatial Profiler (DSP) [28] by NanoString and CosMx Spatial Molecular Imager (SMI) by NanoString,
MERFISH, Vizgen, Resolve, Rebus, and molecular cartography.
10x Visium does not profile at single-cell resolution, and while GeoMx DSP is capable of single-cell resolution through user-drawn profiling regions, the scalability is limited.
The most recent platform, CosMx Spatial Molecular
Imager (SMI) [16], can profile consistently at single cell and even sub-cellular resolution.
CosMx
SMI follows much of the initial protocol as GeoMx DSP, with barcoding and ISH hybridization.
However, the SMI instrument performs 16 cycles of automated cyclic readout, and in each cycle the set of barcodes (readouts) are UV-cleaved and removed.
These cycles of hybridization and imaging yield spatial resolved profiling of RNA and protein at single-cell ( ∼ 10 µm ) and subcellular ( ∼ 1 µm ) resolution.
In this work, we use an unpublished dataset produced by the CosMx platform.
In order to obtain the cell level gene expression, we utilize CellPose [34] software to conduct cell segmentation.
2.2 Single-cell Transcriptomics Data Imputation
The increased resolution of transcriptomics profiling methods comes at a cost of increasing data sparsity.
The profiling technology may fail to capture a number of the expressed genes of an individual cell due to low amounts of mRNA in individual cells and low capture rate.
This results in a large proportion of false zero counts, defined as “dropout” [20, 36].
A popular way to address dropout is to perform imputation, which aims to correct false zeros by estimating realistic values for those gene-cell pairs.
A large number of methods have been developed for the task of scRNA-seq data imputation, mainly focusing on generative probability models or matrix factorization [15, 19, 32, 40].
Aside from these methods, deep learning models have gained immense popularity over recent years.
A natural deep learning architecture for the imputation task is the autoencoder, due to its prevalence in data denoising and missing data applications [3, 4, 13, 14, 30].
For example, as one of the early autoencoder methods for scRNA imputation, AutoImpute [37] employs an overcomplete autoencoder model rather than the usual undercomplete one.
Later, the deep count autoencoder (DCA) [12] is developed for single-cell transcriptomics data imputation.
In contrast to the overcomplete model of AutoImpute, DCA takes a form similar to a standard variational autoencoder (VAE).
Moreover,
DeepImpute [2] builds multiple neural networks in parallel to impute target genes using a set of input genes.
2.3 Denoising Transciptomics Data via Graph Neural Networks
Despite the huge success achieved by deep generative models such as VAEs in single-cell data imputation, a new series of methods based on graph neural networks [26] have recently gained increasing attention.
For instance, scGNN [44] first builds a cell-cell graph based on gene expression similarity and then utilizes a graph autoencoder together with a standard autoencoder to refine graph structures and cell representation.
Lastly, an imputation autoencoder is trained with a graph Laplacian smoothing term added to the reconstruction loss.
GraphSCI [31] constructs a graph from the data with genes taken as the nodes and the edges between them given by the correlation coefficient of the expression data.
Given this graph, the autoencoder and graph autoencoder make use of the expression data and graph data to reconstruct its input.
More recently, with the development of new spatial transcriptomics technologies, some graph neural networks methods are introduced [18, 23, 10, 45] to address data noise in spatial transcriptomics data.
SpaGCN [18] first constructs an undirected weighted graph of spots from the spatial transcriptomics data, where the edge weight is determined by the distance between spots.
A popular GNN model called
GCN [21] is then utilized to aggregate spot gene expression from neighborhoods and update spot gene expression.
Similarly, CCST [23] constructs a spatial neighboring cell graph, and introduces a selfsupervised graph neural network model, deep graph infomax (DGI) [42], to obtain cell representations from the graph.
Moreover, STAGATE [10] also constructs a spatial cell graph but uses a different aggregation method, graph attention network [41], to learn low-dimensional latent embeddings with an autoencoder framework.
Notably, all these methods leverage the denoising effect of graph neural networks but they are not specifically designed for imputation.
Meanwhile, most of them are only demonstrated on spot-level spatial transcriptomics datasets.
To leverage the spatial information, the latest method, Sprod [45] first projects transcriptomics features to a latent space, then connects neighboring cells in the latent space, and prunes it with physical distance.
Then a denoised matrix is learned by minimizing the reconstruction error and a graph Laplacian smoothing term.
Different from the aforementioned methods, we propose to construct two separate graphs for spatial relations and cell similarity.
In addition, we introduce a novel mask-then-predict paradigm as the training objective to explicitly model the recovery of dropouts in spatial transcriptomics data.
The Proposed Framework
We first introduce some notations and basic concepts.
A spatial transcriptomics dataset typically consists of two parts: the gene expressions and the spatial positions.
We denote the gene expression data as a matrix X ∈ R N × k , where N is the number of spots or cells (depending on the dataset) and k is the number of gene types measured in the dataset.
Hereby, X i,j represents the count of gene- j captured in corresponding spot-
i or cell- i .
We use another matrix C ∈ R N × 2 to denote the two-dimensional coordinates of each spot or cell.
Note that in the case of single-cell level datasets,
[0, mask , 1.1, mask ] [0, 0, …, mask, 1] [1, 0.5, …, 2.1, 0] [0.7, mask, …, 1, 0] [0.6, 1.2, …, 1, 0]
Spatial Cell Graph
Cell Similarity Graph
Gating
Cell Representation
Dynamic k-NN (rebuild every layer)
[Figure]
Representation
[Figure]
mask : to be imputed
[Figure]
the source data are preprocessed with a cell segmentation method and the whole space are separated into individual cell regions.
The gene expression matrix X is generated by counting RNA molecules in each cell region, and the coordinates C is the center position of each cell region.
In the spatial transcriptomics imputation problem, a part of values in X are missing, denoted as a mask matrix M ∈ { 0 , 1 } N × k , where the value of X i,j can only be observed when M i,j = 1 .
A partially observed data matrix (cid:101) X is hereby defined as:
[Equation]
Our goal is to predict the missing values X i,j at M i,j = 0 , given the partial gene expression data (cid:101)
X i,j and the spatial positions C .
In the following, we present the details of our proposed framework.
3.1 Cell Graph Construction with Spatial Distance and Gene Expression Similarity
Based on gene spatial distance and expression similarity, there are multiple ways to construct a cell graph and correspondingly facilitate imputation.
For instance, scGNN [44] constructed a k-nearestneighbor graph using only gene expressions; CCST [23] and STAGATE [10] constructed spatial neighboring graphs that only connect adjacent cells; and Sprod [45] first projects transcriptomics features to a latent space, then connects neighboring cells in the latent space and prune it with physical distance.
In this work, we propose to construct two separate graphs for spatial relations and cell similarity respectively.
This will facilitate our model to impute different types of cells by adaptively leveraging both spatial and similarity information.
A spatial cell graph is denoted as G spa =
( V , E spa ) , where V is the set of cells, and E spa is the set of edges.
Here, the edges describe the spatial relations between cells and can also be represented by an adjacency matrix A spa ∈ R N × N .
To construct G spa , we first calculate pairwise euclidean distances between cells.
Concretely, a distance matrix D ∈ R N × N is calculated as:
[Equation]
We then build the adjacency matrix A spa by setting a threshold p on the distance matrix:
[Equation]
In addition, we construct a dynamic cell similarity graph, denoted as G ( l ) sim
=
( V , E ( l ) sim ) , where we exploit the low-dimensional hidden states from our model.
To be concrete, let H ( l − 1) ∈ R N × d ( l − 1) and H ( l ) ∈ R N × d ( l ) be the input and output cell representations of the l -th layer, where d ( l ) is pre-defined hidden size.
H (0) can be initialized by a low-dimensional representation generated from ˆX with unsupervised algorithms such as principle component analysis (PCA) [1].
We construct a distance matrix ˆD ( l ) in each layer, using a formula similar to Eq 2, while C is replaced with H ( l − 1) .
After that, a similarity-based adjacency matrix A ( l ) sim ∈ R N × N is built with a k-nearest neighbor algorithm.
Formally,
[Equation]
where N ( l ) k
( i ) denotes the set containing the k -nearest neighbors of cell i according to the distance matrix ˆD ( l ) .
In practice, we separate cells into fields of view (FOVs), where we only calculate distances between cells within the same FOV.
This allows our model to have the potential to be scalable to millions of cells.
3.2 Bi-channel Graph Neural Network
Our bi-channel graph neural network layer as shown in Figure 3 performs aggregation on both the spatial cell graph G spa and the cell similarity graph G ( l ) sim and then updates the node features, i.e., cell representations.
Essentially, it consists of three parts: a spatial graph attention layer, a similarity graph attention layer, and a gating unit.
The two graph attention layers (GATs) [41] aggregate the information from G spa and G ( l ) sim parallelly.
Then a gating unit fuses information from the two GATs, and generates the output of the whole layer.
Formally, a GAT layer can be formulated as:
[Equation]
where z ( l ) i denotes the output node features of node i of the l -th layer; σ is an arbitrary non-linear activation function; N ( i ) is the set of neighboring nodes of node i , W ( l ) is a learnable weight matrix for linear transformation, [ ·∥· ] indicates the concatenation operation, a ( l ) is a learnable weight vector for attention and α ( l ) ij is the attention score between node i and node j .
In a bi-channel graph neural network layer, we have two GAT layers for G spa and G ( l ) sim , respectively.
One is parameterized by W ( l ) spa and a ( l ) spa , denoted as GAT ( l ) spa .
The other is parameterized by W ( l ) sim and a ( l ) sim , denoted as GAT ( l ) sim .
Formally, the input and output of GAT ( l ) spa and GAT ( l ) sim can be written as:
[Equation]
Note that G spa and G ( l ) sim share the same set of nodes V (i.e. the same set of cells), and the two GAT layers share the same input node features (a.k.a H ( l − 1) ).
The difference between G spa and G ( l ) sim lies in the adjacency matrix A spa and A ( l ) sim .
Hence, N ( i ) in Eq 5 is different in GAT ( l ) spa and GAT ( l ) sim .
Next, we define a gating unit as follows:
[Equation]
) where h ( l − 1) is the input node embedding of any node in layer l , h ( l ) spa is the output node embedding from GAT ( l ) spa , h ( l ) sim is the output node embedding from GAT ( l ) sim , W ( l ) gate ∈ R 3 × d ( l − 1) is a learnable matrix that project node embedding to a 3-dimensional vector Θ ( l ) , and θ ( l ) i is the i -th element of vector Θ ( l ) .
As a result, the new embedding of each node is a weighted sum of the output of GAT ( l ) spa , the output of GAT ( l ) sim , and the old embedding.
In this way, our bi-channel graph neural network layer fuses information from G spa and G ( l ) sim .
Overall, a bi-channel graph neural network layer ϕ
( l ) can be summarized as:
[Equation]
where H ( l ) will be used to build A ( l +1) sim and fed to layer l + 1 .
The similarity graph G ( l ) sim is therefore gradually refined through layers.
3.3 Graph Masked Autoencoders
The idea of graph masked autoencoders is to train an encoder-decoder architecture to recover some artificially added dropouts in the graph.
This serves as a self-supervised pretraining objective and the trained encoder can be applied to different downstream tasks.
The overall structure is similar to a denoising autoencoder[43], while both the encoder and the decoder can be instantiated as a graph neural network, and both edges and nodes in the input graph can be modified to create noisy input.
Here, we follow the framework of GraphMAE [17] to only mask node features as shown in Figure 3.
More specifically, we uniformly sample a mask M ′ ∈ { 0 , 1 } N × k .
To be consistent with previous sections, we denote the original input gene expression matrix as (cid:101) X .
The node feature (cid:101) X
i,j is masked when M ′ i,j = 0 , resulting in a new node feature matrix (cid:101)
X
′ ∈ R N × k written as:
[Equation]
where ⊙ indicates element-wise multiplication.
Note that the generation process of (cid:101) X ′ perfectly matches the generation of (cid:101) X , which means the optimization objective is consistent with our imputation goal.
This is different from previous studies, as we discussed in Section 2.3.
During training, masked features (cid:101)
X ′ are used as initial node features for ϕ (1) , i.e. H (0)
= (cid:101) X ′ .
As we mentioned in Section 3.1, we can also add dimension reduction methods, such as H (0) =
PCA (cid:16) (cid:101) X ′ (cid:17) .
For testing, the original input (cid:101) X is enabled, which provides extra information for imputation.
Finally, let L be the total number of graph neural network layers, we consider all those L layers as an encoder.
Therefore, H ( L ) is the latent code that is the output from the encoder.
We then add a fully connected decoder to predict the unmasked features, the prediction is thus denoted as:
[Equation]
(11) where W o ∈ R d ( L )
× k is the learnble weight matrix, and Z ∈ R N × k is the recovered gene expression matrix.
3.4 The Objective Function
As a masked autoencoder, the objective function for the whole framework can be defined as:
[Equation]
where we only calculate mean squared error (MSE) for the prediction of masked values.
In addition, due to the data sparsity, most values in (cid:101) X are zero, therefore most predictions are also expected to be zero.
To address this kind of unbalanced data, a weighted MSE loss is alternatively introduced, defined as:
[Equation]
(13) where we use the ground truth value in (cid:101) X to assign a weight encouraging non-zero expressions, and γ is a scaling factor that controls this effect.
After the model is trained, we run our model on the original input matrix (cid:101) X , and the resulting recovered gene expression matrix Z will be the final imputation result.
Experiment
In this section, we evaluate the effectiveness of our methodon spatial transcriptomics data imputation.
Our source codes will be integrated into the DANCE package [9].
Before presenting our experimental results and observations, we first introduce the experimental settings.
4.1 Experimental settings 4.1.1 Datasets
We illustrate the effectiveness of our algorithm on a spatial transcriptomics dataset over human
Non-Small Cell Lung Cancer (NSCLC)
Formalin-Fixed Paraffin-Embedded (FFPE) slides using the
CosMX platform from Nanostring [16].
In this data, 100,149 cells were profiled on 960 pre-selected marker genes and 20 negative probe controls with an average of 300 transcripts per cell and an average dropout rate of 87%.
4.1.2 Baselines
To evaluate the effectiveness of our method, we compare it with the state-of-the-art spatial and non-spatial imputation models:
• scGNN [44] first builds a cell-cell graph based on gene expression similarity and then utilizes a graph autoencoder together with a standard autoencoder to refine graph structures and cell representation.
Lastly, an imputation autoencoder is trained with a graph Laplacian smoothing term added to the reconstruction loss.
• scVI [25] is based on a hierarchical Bayesian model with conditional distributions specified by neural networks.
It models dropouts with a zero-inflated negative binomial distribution, and can estimate the distributional parameters of each gene in each cell.
We intended to run another state-of-the-art model Sprod [45] on our dataset.
However, we failed to get results within 8 hours.
We also lack comparisons with other methods, e.g. scImpute [24],
SAVER [19], etc.
This is still a work-in-progress, we will present a more comprehensive experiment when officially published.
4.1.3 Implementation Settings
For baseline models, we follow the default settings on authors’ github.
Regarding preprocessing, we remove cells with no gene expressions.
For scVI, the input is raw read counts of genes.
While for other models, we use log-normalized expression data.
To create the training data, we randomly set a certain ratio r of gene expressions to be zero, to simulate different degrees of noise.
The experiments are conducted under two different settings, r = 0 .
3 and r = 0 .
Lastly, all experiments are runtimes, and the average performance is reported.
4.2 Imputation Performance
To evaluate the imputation performance, we use values that are dropped in training data as groundtruth labels and compare them with the imputed values from the imputation model.
The metrics we use are root mean squared error (RMSE, the lower, the better) and Pearson correlation coefficient (Correlation, the higher the better).
All evaluations are conducted based on log-normalized values.
In Table 1, we present the experiment results.
It is shown that our method significantly outperforms other baselines by showing a higher correlation, and the ablation study shows that both spatial graph and kNN graph benefit the imputation performance.
Conclusion
In this work, we propose a new technique to adaptively exploit the spatial information of cells and the heterogeneity among different types of cells for imputation.
Furthermore, a mask-then-predict paradigm is adopted to explicitly model the recovery of dropouts and enhance the denoising effect.
Compared to existing studies, our work focus on new large-scale cell-level data instead of spots or r=
[Table] [Table]
r=0.
[Table]