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 a costof the extremely high resolution, the technology also results in much more missingvalues in the data, i.e. dropouts.
While a common solution is to perform imputationon the missing values, existing imputation methods majorly focus on transcriptomics data and tend to yield sub-optimal performance on spatial transcriptomicsdata.
To advance spatial transcriptomics imputation, we propose a new techniqueto adaptively exploit the spatial information of cells and the heterogeneity amongdifferent types of cells.
Furthermore, we adopt a mask-then-predict paradigmto explicitly model the recovery of dropouts and enhance the denoising effect.
Compared to previous studies, our work focus on new large-scale cell-level datainstead of spots or beads.
Preliminary results have demonstrated that our methodoutperforms previous methods for removing dropouts in high-resolution spatiallyresolved transcriptomics data.
Introduction
The rapid advance of single-cell technologies have paved the way for investigating biological processes at an unprecedentedly high resolution.
However, most single-cell technologies such asscRNA-seq [35] and scATAC-seq [6] are unable to capture spatial information of the cells, which isoften coupled with cell identities.
To spatially resolve transcriptomics profiles, spatial transcriptomictechnologies are recently developed with which researchers are able to identify the spatial context ofcells and cell clusters within a tissue [7].
Similar to standard single-cell RNA-sequencing, spatiallyresolved profiling technology may fail to capture a significant number of the expressed genes dueto low RNA capture rate.
This results in an artificially large proportion of false zero counts, i e.,the so-called “dropout” [33, 19].
A direct way to address this issue is to perform imputation, whichcorrects false zeros by estimating realistic values for those missing values.
Recently, there has been anexplosive growth of imputation methods for single-cell transcriptomics data, focusing on generativeprobability models, matrix factorization [37, 18, 36, 23] and deep learning models [34, 11, 8, 41, 28].
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, few graph-based methods [22, 42] propose to build a cell-cell graph based on thespatial position and apply graph neural networks on built graph.
However, these methods do notconsider heterogeneity among different types of cells.
For example, tumor cells usually show strongaggregation, and neighboring tumor cells with similar micro-environment in terms of ligands tendto have higher gene expression similarity than distant cells [5, 21, 25].
By contrast, Treg cells arescarce spatially in many tissues but still tend to have similar gene expression profile [30].
Hence, it isdesired to rely on spatial information to impute cells like tumors, while leveraging information fromthe same type of cells to impute cells like Treg cells.
To capture such heterogeneity while exploitingspatial information, we hereby propose a novel bi-channel graph neural network (GNN).
Specifically,we first construct two channels: a spatial neighboring graph capturing the spatial information, and adynamic k-nearest neighbor (kNN) graph capturing the similarity of cells.
Then we adopt a gatingunit to adaptively fuse the information from the two graphs, which adjusts the exploitation of spatialinformation 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 networks 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 [32]) or tiny cell-level datasets (e.g.
MERFISH [43], Fishseq+ [10]), CosMx platform provided unprecedentedly large-scale single-cellspatial transcriptomics datasets, with more than 100,000 cells in each slice.
Our proposed methodachieved 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 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 therecovery of dropouts in spatial transcriptomics data, which shows outstanding performance onour
dataset.•
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 some of the related work.
Weexplain our proposed framework in Section 3 and report 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 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].
Though, most spatially resolved profiling technologies are not atthe single-cell resolution, so arises the problem of cell-type deconvolution and segmentation.
The major technologies/platforms for spatial transcriptomics are Visium by 10x [32], GeoMx Digital
Spatial Profiler (DSP) [26] 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 CosMx platform.
In order to obtain celllevel gene expression, we utilize CellPose [31] 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” [33, 19].
A direct way to address dropout isto perform imputation, a method to correct false zeros by estimating realistic values for those gene-cellpairs.
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 [37, 29, 14, 18].
Asidefrom these methods, deep learning models have gained immense popularity over the recent years.
Anatural deep learning architecture for the imputation task is the autoencoder, due to its prevalencein data denoising and missing data applications [27, 12, 13, 3, 4].
For example, as one of the earlyautoencoder methods for scRNA imputation, AutoImpute [34] 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).
Moreoever,another model, DeepImpute [2] builds multiple neural networks in parallel to impute target genesusing a set of input genes.
2.3 Denoising Transciptomics Data with 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 have recently gained increasingattention.
For instance, scGNN [41] first builds a cell-cell graph based on gene expression similarityand then utilizes a graph autoencoder together with a standard autoencoder to refine graph structuresand cell representation.
Lastly, an imputation autoencoder is trained with a graph Laplacian smoothingterm added to the reconstruction loss.
Another method called GraphSCI [28] constructs a graph fromthe data with genes taken as the nodes and the edges between them given by correlation coefficientof the expression data.
Given this graph, the autoencoder and graph autoencoder make use of theexpression data 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, 42] 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) [39], to obtain cell representationsfrom the graph.
Moreover, STAGATE [9] also constructs a spatial cell graph but uses a differentaggregation method, graph attention network [38], to learn low-dimensional latent embeddings withan autoencoder framework.
Notably, all these methods leverage the deonising 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 [42] 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 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 case of single-cell level datasets, thesource data are preprocessed with a cell segmentation method and the whole space are separated intoindividual cell regions.
The gene expression matrix X is generated by counting RNA molecules ineach cell region, and the coordinates C is the center position of each cell region.
In the spatial transcriptomics imputation problem, part of values in X are missing, denoted as a maskmatrix M ∈ { 0 , 1 } N × k , where the value of X i,j can only be observed when M i,j = 1 .
A partiallyobserved 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 facilitate imputation.
For instance, scGNN [41] constructed a k-nearest-neighbor graphusing only gene expressions; CCST [22] and STAGATE [9] constructed spatial neighboring graphsthat only connect adjacent cells; Sprod [42] 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 instead propose to construct two separate graphs for spatial relations and similarity.
A spatial cell graph is denoted as G spa =
( V , E spa ) , where V is the set of all cells, 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 equation 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-nearestneighbor algorithm.
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 be scalable to millions ofcells in the future.
3.2 Bi-channel Graph Neural Network
Our bi-channel graph neural network layer performs aggregation on both the spatial cell graph
G spa and the cell similarity graph G ( l ) sim , and then update 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) [38] aggregate the information from G spaand G ( l ) sim parallelly.
Then a gating unit fuses information from the two GATs, and generate the outputof 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-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 anda ( 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 equation 5 is different in GAT ( l ) spa and GAT ( l ) sim .
Next, we define a gating unit as follows:
[Equation]
[Equation]
(8) 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 ( 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 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[40], 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.
More specifically, we uniformly sample a mask M ′ ∈ { 0 , 1 } N × k .
To be consistent with section ?
,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 ′ 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 dimensionality reduction methods, such as H (0) =
PCA (cid:16) (cid:101) X ′ (cid:17) .
For testing, theoriginal 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 asan encoder.
Therefore, H ( l ) is the latent code that output from the encoder.
We then add an 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, Z ∈ R N × k is the recovered gene expressionmatrix.
3.4 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.
Besides, 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.
Experiments
In this section, we evaluate the effectiveness of our method on 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 potential usefulness of our algorithm on a spatial transcriptomics dataset over human
Non-Small Cell Lung Cancer (NSCLC)
Formalin-Fixed Paraffin-Embedded (FFPE) slides using r=
[Table]
[Table]
r=0.
[Table]
CosMX platform [15].
In this data, 100,149 cells were profiling on 960 pre-selected marker genesand 20 negative probe controls with an average of 300 transcripts per cell and an average dropout rateof 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 [41] 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 [42] 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 Experimental 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 forother 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
The imputation performance is defined as the accuracy of prediction over missing values.
To evaluatethe imputation performance, we use values that are dropped in training data as labels, and comparethem with the imputed values from the imputation model.
The metrics we use are rooted meansquared error (RMSE, the lower the better) and Pearson correlation coefficient (Correlation, thehigher the better).
All evaluations are conducted based on log-normalized values.
In table 1, we present the experiment results.
It is shown from the experiments that out methodsignificantly outperforms other baselines by showing a higher correlation, and the ablation studyshows that both spatial graph and kNN graph benefit the imputation performance.
Conclusion
In this study, we propose a new technique to adaptively exploit the spatial information of cells and theheterogeneity among different types of cells.
Furthermore, a mask-then-predict paradigm is adoptedto explicitly model the recovery of dropouts and enhance the denoising effect.
Compared to previousstudies, our work focus on new large-scale cell-level data instead of spots or beads.
Preliminaryresults have demonstrated that our method outperforms previous methods for removing dropoutsin high-resolution spatially resolved transcriptomics data.
This is a work-in-progress and we willcontinue to polish our ideas and experiments.