Bi-channel Masked Graph Autoencoders for Spatially Resolved Single-cell Transcriptomics Data Imputation
Anonymous Author(s)
AffiliationAddress email
Abstract
Spatially resolved transcriptomics bring exciting breakthroughs to single-cell analysis by providing physical locations along with gene expression.
However, as acost of the extremely high resolution, the technology also results in much moremissing values in the data, i.e. dropouts.
While a common solution is to performimputation on the missing values, existing imputation methods majorly focus ontranscriptomics data and tend to yield sub-optimal performance on spatial transcriptomics data.
To advance spatial transcriptomics imputation, we propose a newtechnique to adaptively exploit the spatial information of cells and the heterogeneityamong different cell types.
Furthermore, we adopt a mask-then-predict paradigmto explicitly model the recovery of dropouts and enhance the denoising effect.
Compared to existing studies, our work focuses on new large-scale cell-level datainstead of spots or beads.
Preliminary results have demonstrated that our methodoutperforms existing methods for removing dropouts in high-resolution spatiallyresolved transcriptomics data.
Introduction
The rapid advance of single-cell technologies has paved the way to investigate biological processes atan unprecedentedly high resolution.
However, most single-cell technologies such as scRNA-seq [36]and scATAC-seq [6] are unable to capture spatial information of the cells, which is often coupledwith cell identities.
To spatially resolve transcriptomics profiles, spatial transcriptomic technologieshave been recently developed where researchers are able to identify the spatial context of cells andcell clusters within a tissue [7].
Similar to standard single-cell RNA-sequencing, spatially resolvedprofiling 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” [34, 19].
A popular way to address this issue is to perform imputation, which corrects falsezeros by estimating realistic values for those missing values.
Recently, there has been an explosivegrowth of imputation methods for single-cell transcriptomics data, focusing on generative probabilitymodels, matrix factorization [38, 18, 37, 23] and deep learning models [35, 11, 8, 42, 29].
Although promising results have been achieved on single-cell RNA-sequencing data, it remainschallenging to impute spatially resolved single-cell transcriptomics data.
Due to the relatively lowsequencing depth in each unit and the extra experimental processes to obtain spatial locations, spatialtranscriptomics technology suffers from a higher noise level and drop-out rate (zero reads for manygenes).
Furthermore, most existing imputation models do not take advantage of spatial informationand it is still unclear how to advance imputation with spatial information.
To leverage the spatial information, very recently graph-based methods [22, 43] have been designed.
They first builda cell-cell graph based on the spatial position and then apply graph neural networks to the builtgraph.
However, these methods do not consider heterogeneity among different types of cells.
Forexample, tumor cells usually show strong aggregation, and neighboring tumor cells with similarmicro-environment in terms of ligands tend to have higher gene expression similarity than distantcells [5, 21, 26].
By contrast, Treg cells are scarce spatially in many tissues but still tend to havesimilar gene expression profiles [31].
Hence, it is desired to rely on spatial information to impute cellslike 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-channelgraph neural network (GNN).
Specifically, we first construct two channels: a spatial neighboringgraph capturing the spatial information, and a dynamic k-nearest neighbor (kNN) graph capturing thesimilarity 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 [16] which recovers themanually masked node features through a GNN, we introduce a novel mask-then-predict paradigmfor training the proposed imputation method.
Specifically, we randomly mask 60% of genes withzero 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) [15].
Unlike existing spot-level datasets (e.g. 10x Visium [33]) or tiny cell-level datasets (e.g.
MERFISH [44], Fishseq+ [10]), CosMx platform provides unprecedentedly large-scale single-cellspatial transcriptomics datasets, with more than 100,000 cells in each slice.
Our proposed methodachieves state-of-the-art performance on this dataset and significantly outperforms other baselinemodels.
Our contribution can be summarized as follows:
• We propose a novel bi-channel graph neural network to leverage information from both spatialneighboring cells and similar cells to advance spatial transcriptomics data imputation.•
We introduce a mask-then-predict paradigm as the training objective to explicitly model therecovery of dropouts in spatial transcriptomics data, which shows outstanding performance.•
We make the first attempt to evaluate imputation methods on large-scale single-cell spatialtranscriptomics datasets.
Our method improves the data quality and makes downstream tasksmore 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 ourproposed framework in Section 3 and discuss our experimental results in Section 4 .
Finally, weconclude the work with future directions in Section 5 .
Related Work2.1 Spatial Resolved Transcriptomics
While single-cell technologies are able to capture transcriptomic information at the single-cell level,most are unable to capture spatial information of the cells, due to the isolation of single-cells.
However,the relationship between cells and their relative locations within tissue is critical to understandingnormal development and disease pathology.
Spatial transcriptomic technologies are a more recentdevelopment in transcriptomics profiling, which are able to spatially resolve transcriptomics profiles.
With spatial transcriptomics technologies, researchers are able to learn the spatial context of cells,and cell clusters within a tissue [7].
Most commercial spatially resolved profiling technologiesare not at the single-cell resolution that has arisen the problem of cell-type deconvolution andsegmentation [33, 27].
The major technologies/platforms for spatial transcriptomics are Visium by 10x [33], GeoMx Digital
Spatial Profiler (DSP) [27] by NanoString and CosMx Spatial Molecular Imager (SMI) by NanoString,
MERFISH, Vizgen, Resolve, Rebus, and molecular cartography.
10x Visium does not profile at singlecell 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)[15], is able to profile consistently at single cell, and even sub-cellular resolution.
CosMx SMIfollows 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 ofbarcodes (readouts) are UV-cleaved and removed.
These cycles of hybridization and imaging yieldspatial 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 thecell level gene expression, we utilize CellPose [32] 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 datasparsity.
The profiling technology may fail to capture a number of the expressed genes of an individualcell due to low RNA capture rate as well as the stochasticity of mRNA expression.
This results in alarge proportion of false zero counts, defined as “dropout” [34, 19].
A popular way to address dropoutis to perform imputation, which aims to correct false zeros by estimating realistic values for thosegene-cell pairs.
A large number of methods have been developed for the task of scRNA-seq dataimputation, mainly focusing on generative probability models or matrix factorization [38, 30, 14, 18].
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 prevalencein data denoising and missing data applications [28, 12, 13, 3, 4].
For example, as one of the earlyautoencoder methods for scRNA imputation, AutoImpute [35] employs an overcomplete autoencodermodel rather than the usual undercomplete one.
Later, the deep count autoencoder (DCA) [11] isdeveloped for single-cell transcriptomics data imputation.
In contrast to the overcomplete modelof 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 inputgenes.
2.3 Denoising Transciptomics Data via Graph Neural Networks
Despite the huge success achieved by deep generative models such as VAEs in single-cell dataimputation, a new group of methods based on graph neural networks [25] have recently gainedincreasing attention.
For instance, scGNN [42] first builds a cell-cell graph based on gene expressionsimilarity and then utilizes a graph autoencoder together with a standard autoencoder to refine graphstructures and cell representation.
Lastly, an imputation autoencoder is trained with a graph Laplaciansmoothing term added to the reconstruction loss.
GraphSCI [29] constructs a graph from the datawith genes taken as the nodes and the edges between them given by the correlation coefficient of theexpression data.
Given this graph, the autoencoder and graph autoencoder make use of the expressiondata and graph data to reconstruct its input.
More recently, with the development of new spatial transcriptomics technologies, some graph neuralnetworks methods are introduced [17, 22, 9, 43] to address data noise in spatial transcriptomics data.
SpaGCN [17] first constructs an undirected weighted graph of spots from the spatial transcriptomicsdata, where the edge weight is determined by the distance between spots.
A popular GNN model called
GCN [20] is then utilized to aggregate spot gene expression from neighborhoods and update spot geneexpression.
Similarly, CCST [22] constructs a spatial neighboring cell graph, and introduces a selfsupervised graph neural network model, deep graph infomax (DGI) [40], to obtain cell representationsfrom the graph.
Moreover, STAGATE [9] also constructs a spatial cell graph but uses a differentaggregation method, graph attention network [39], to learn low-dimensional latent embeddings withan autoencoder framework.
Notably, all these methods leverage the denoising effect of graph neuralnetworks but they are not specifically designed for imputation.
Meanwhile, most of them are onlydemonstrated on spot-level spatial transcriptomics datasets.
To leverage the spatial information,the latest method, Sprod [43] first projects transcriptomics features to a latent space, then connects
[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]
neighboring cells in the latent space, and prunes it with physical distance.
Then a denoised matrix islearned by minimizing the reconstruction error and a graph Laplacian smoothing term.
Different from the aforementioned methods, we propose to construct two separate graphs for spatialrelations and cell similarity.
In addition, we introduce a novel mask-then-predict paradigm as thetraining 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 typicallyconsists of two parts: the gene expressions and the spatial positions.
We denote the gene expressiondata 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 ofgene-
j captured in corresponding spot-
i or cell- i .
We use another matrix C ∈ R N × 2 to denote thetwo-dimensional coordinates of each spot or cell.
Note that in the case of single-cell level datasets,the source data are preprocessed with a cell segmentation method and the whole space are separatedinto individual cell regions.
The gene expression matrix X is generated by counting RNA moleculesin 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 amask matrix M ∈ { 0 , 1 } N × k , where the value of X i,j can only be observed when M i,j = 1 .
Apartially 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 cellgraph and correspondingly facilitate imputation.
For instance, scGNN [42] constructed a k-nearestneighbor graph using only gene expressions; CCST [22] and STAGATE [9] constructed spatial neighboring graphs that only connect adjacent cells; and Sprod [43] first projects transcriptomicsfeatures to a latent space, then connects neighboring cells in the latent space and prune it with physicaldistance.
In this work, we propose to construct two separate graphs for spatial relations and cellsimilarity respectively.
This will enable our model to impute different types of cells by adaptivelyleveraging 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 ofedges.
Here, the edges describe the spatial relations between cells and can also be represented by anadjacency matrix
A spa ∈ R N × N .
To construct G spa , we first calculate pairwise euclidean distancesbetween 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 weexploit 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 ) ispre-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 adistance 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 neighboralgorithm.
Formally,
[Equation]
where N ( l ) k
( i ) denotes the set containing the k -nearest neighbors of cell i according to the distancematrix ˆD ( l ) .
In practice, we separate cells into fields of view (FOVs), where we only calculatedistances between cells within the same FOV.
This allows our model to have the potential to bescalable 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 thespatial cell graph G spa and the cell similarity graph G ( l ) sim and then updates the node features, i.e., cellrepresentations.
Essentially, it consists of three parts: a spatial graph attention layer, a similaritygraph attention layer, and a gating unit.
The two graph attention layers (GATs) [39] aggregate theinformation 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]
(5) where z ( l ) i denotes the output node features of node i of the l -th layer; σ is an arbitrary non-linearactivation function; N ( i ) is the set of neighboring nodes of node i , W ( l ) is a learnable weight matrixfor linear transformation, [ ·∥· ] indicates the concatenation operation, a ( l ) is a learnable weight vectorfor 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 GATlayers share the same input node features (a.k.a H ( l − 1) ).
The difference between G spa and G ( l ) sim lies inthe 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]
[Equation]
where h ( l − 1) is the input node embedding of any node in layer l , h ( l ) spa is the output node embeddingfrom 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 learnablematrix that project node embedding to a 3-dimensional vector Θ ( l ) , and θ ( l ) i is the i -th element ofvector Θ ( 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 layerfuses information from G spa and G ( [Equation]
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 thereforegradually refined through layers.
3.3 Graph Masked Autoencoders
The idea of graph masked autoencoders is to train an encoder-decoder architecture to recover someartificially added dropouts in the graph.
This serves as a self-supervised pretraining objective and thetrained encoder can be applied to different downstream tasks.
The overall structure is similar to adenoising autoencoder[41], while both the encoder and the decoder can be instantiated as a graphneural network, and both edges and nodes in the input graph can be modified to create noisy input.
Here, we follow the framework of GraphMAE [16] 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 previoussections, we denote the original input gene expression matrix as (cid:101) X .
The node feature (cid:101) X
i,j is maskedwhen 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 ′ perfectlymatches the generation of (cid:101) X , which means the optimization objective is consistent with our imputationgoal.
This is different from previous studies, as we discussed in Section 2.3.
During training, maskedfeatures (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 originalinput (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 anencoder.
Therefore, H ( l ) is the latent code that is the output from the encoder.
We then add a fullyconnected 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 expressionmatrix.
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 alsoexpected to be zero.
To address this kind of unbalanced data, a weighted MSE loss is alternativelyintroduced, defined as:
[Equation]
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 resultingrecovered 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.
Before presenting our experimental results and observations, we first introduce the experimentalsettings.
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 [15].
In this data, 100,149 cells were profiled on 960 pre-selectedmarker 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 andnon-spatial imputation models:
• scGNN [42] first builds a cell-cell graph based on gene expression similarity and then utilizesa graph autoencoder together with a standard autoencoder to refine graph structures and cellrepresentation.
Lastly, an imputation autoencoder is trained with a graph Laplacian smoothingterm added to the reconstruction loss.•
scVI [24] is based on a hierarchical Bayesian model with conditional distributions specified byneural networks.
It models dropouts with a zero-inflated negative binomial distribution, and canestimate the distributional parameters of each gene in each cell.
We intended to run another state-of-the-art model Sprod [43] on our dataset.
However, we failedto get results within 8 hours.
We also lack comparisons with other methods, e.g. scImpute [23],
SAVER [18], etc.
This is still a work-in-progress, we will present a more comprehensive experimentwhen officially published.
4.1.3 Implementation Settings
For baseline models, we follow the default settings on authors’ github.
Regarding preprocessing, weremove cells with no gene expressions.
For scVI, the input is raw read counts of genes.
While for r=
[Table] r= [Table] other models, we use log-normalized expression data.
To create the training data, we randomly set acertain ratio r of gene expressions to be zero, to simulate different degrees of noise.
The experimentsare conducted under two different settings, r = 0 .
3 and r = 0 .
Lastly, all experiments are run 5times, 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 groundtruthlabels and compare them with the imputed values from the imputation model.
The metrics we use areroot 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 outperformsother baselines by showing a higher correlation, and the ablation study shows that both spatial graphand kNN graph benefit the imputation performance.
Conclusion
In this work, we propose a new technique to adaptively exploit the spatial information of cells andthe heterogeneity among different types of cells for imputation.
Furthermore, a mask-then-predictparadigm 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 orbeads.
Preliminary results have demonstrated that our method outperforms existing methods forremoving dropouts in high-resolution spatially resolved transcriptomics data.
This is a work-inprogress and we will continue to add more comprehensive experiments, including clustering results,visualizations and model interpretability analysis.