SCTrans: Multi-scale scRNA-seq Sub-vector Completion Transformer for Gene-selective Cell Type Annotation

1 Introduction

The rapid advancements in single-cell RNA sequencing (scRNA-seq) technologies have enabled high-resolution characterization of tissue heterogeneity at cellular-level[Ziegenhain et al., 2017]. As a critical application of scRNA-seq, automatic cell type annotation enables the precise identification of unique and shared genetic expressions across different cell types. This process elucidates the molecular characteristics and functional regulation of cells. Furthermore, cell type annotation facilitates the comparison of cellular typology alterations under pathological states, including cells' stress responses and adaptive mechanisms. This sheds light on cellular functional dysregulation and the mechanisms of disease.

There are a number of attempts made in applying machine learning algorithms to cell type annotation [Wang et al., 2022]. For example, scBERT [Yang et al., 2022] leverages pretrained bidirectional transformer networks to learn the contextual relationships between gene expression vectors. Another representative strategy is to adopt an optimal transport model to identify the known cell types while at the same time discovering new ones, such as scPOT [Zhai et al., 2023]. However, the existing methods mainly focus on learning a nonlinear mapping to infer cell types from individual genes, often overlooking the interdependency between the sub-vectors of gene expression. Compared to individual genes, we consider that gene sub-vectors provide richer information, and performing attention computation among them potentially identities salient ones for downstream biological research and analysis as shown in Figure 1.

More specifically, we propose a multi-scale gene expression Sub-vector Completion Transformer (SCTrans) for cell type annotation over scRNA-seq. SCTrans aims to learn an effective representation of the sub-vectors, such that the interdependency and the association with cell type can be better captured, than dealing with individual genes. As shown in Figure 2, we adopt a multi-scale partitioning strategy to yield gene sub-vectors across different scales. Next, the resulting sub-vectors are tokenized, followed by self-attention computation in each transformer block. One of our training goals is to perform sub-vector completion across different scales: predicting the masked sub-vectors from the unmasked ones. We further impose conditional contrastive regularization on the transformer to learn cell-type-aware representation. By modeling intrinsic contextual relationships, another advantage is to distinguish salient tokens from the others, based on attention scores. As a result, we can perform gene-selective cell type annotation, and provide meaningful and interpretable results. We perform extensive experiments on multiple benchmark dataset to highlight the effectiveness of our design elements and the superior performance over state-of-the-art cell type annotation methods.

In summary, the main contribution of this work are as follows: (a) We design a multi-scale scRNA-seq encoding transformer to learn the representation of gene sub-vectors, whose expressiveness is richer than that of individual genes. (b) Based on the nature of gene co-expression, we perform gene sub-vector completion to capture the intrinsic contextual relationships across different scales. (c) By performing mutual learning between the transformer and a lightweight counterpart, salient tokens of various scales can be distinguished from the others, and lead to significant improvement in terms of annotation accuracy and interpretability.

2 RelatedWork

2.1 Transformer Based Representation Learning

The Transformer architecture [Xu et al., 2019][Dosovitskiy et al., 2021][Wu et al., 2021][Liu et al., 2021] has emerged as a powerful framework, drawing inspiration from transformer research in natural language processing (NLP) [Devlin et al., 2018][Vaswani et al., 2017]. By utilizing the self-attention mechanism, Transformer can model the long time relation, particularly effective in analyzing complex relationships in biological data[Abnar and Zuidema, 2020][Alsaigh et al., 2022]. In single-cell RNA sequencing (scRNA-seq) analysis, Transformer has led to significant advancements. For instance, scBERT [Yang et al., 2022] utilizes the Transformer to annotate cell types in scRNA-seq data, learning from large datasets to understand gene interactions. Similarly, the Exceiver model [Connell et al., 2022], based on the Perceiver IO framework, adapts to new datasets for gene expression analysis. TOSICA [Chen et al., 2023], another Transformer-based tool, excels in integrating diverse datasets for cell type annotation. The DeepMAPS [Ma et al., 2023] model uses a unique approach combining cells and genes to infer biological networks, streamlining the training process. Moreover, scFoundation [Hao et al., 2023] and scAAGA [Meng et al., 2023] models enhance single-cell analysis by transforming gene expression data into informative vectors and employing gene attention modules for feature extraction.

These models are exemplary in their integration of computational techniques with biological data, providing profound insights into genomic analysis. However, a limitation of existing methods is their focus either on single genes, failing to fully reveal the expressions across the gene sub-vectors. This oversight can potentially neglect the deeper implications and nuances within the genomic data.

2.2 Cell Type Annotation

Cell type annotation methods can be broadly categorized into two approaches: cluster-then-annotate and supervised cell type classification. The cluster-then-annotate approach seeks to categorize cells into distinct clusters, and these identified clusters are subsequently manually annotated by experts who examine cluster-specific gene expression patterns [Kiselev et al., 2019]. In this learning paradigm, which falls under the unsupervised learning branch, various clustering approaches have been developed to identify cell types. For example, SC3 accomplishes cell clustering by consistently integrating diverse clustering solutions through a consensus approach [Kiselev et al., 2017]. SIMLR strives to acquire a more accurate metric that better reflects the similarity between samples, capturing the underlying structure of the data through a multiple kernel learning approach [Wang et al., 2017]. Recently, deep learning methods have emerged as powerful tools for clustering scRNA-seq data, overcoming challenges posed by high dropout rates and noise [Tian et al., 2021]. For example, DESC is an unsupervised deep embedding algorithm designed to accurately cluster scRNA-seq data through an iterative self-learning process [Li et al., 2020b]. scziDesk utilizes a convolutional autoencoder to learn representations of the single cell data, followed by a regularized soft kmeans approach to identify cell clusters [Hu et al., 2022]. scNAME method incorporates mask estimation task, in conjunction with a neighborhood contrastive learning framework to achieve effective clustering [Wan et al., 2022]. These deep clustering techniques can characterize complex nonlinear relationships in single cell data to identify the cell groupings.

In contrast, the supervised classification approach involves training a model on data where cell types are preannotated. Subsequently, this model is utilized to predict cell type labels for new, unlabeled single cells [Abdelaal et al., 2019]. While supervised methods can leverage expert knowledge during the training process, their effectiveness may be constrained by the quantity and quality of labeled data available. ScType leverages a comprehensive database of cell type-specific marker genes to facilitate the automated highthroughput annotation of cell types [Ianevski et al., 2022]. scDeepSort leverages a pre-trained weighted graph neural network (GNN) for cell-type annotation [Shao et al., 2021]. MARS, a meta-learning approach, is designed for the identification and annotation of both known and novel cell types [Brbi´c et al., 2020]. scPOT is an end-to-end algorithm specifically designed for the annotation of cell types, including the discovery of novel cell types. This is achieved through the utilization of the optimal transport (OT) framework [Zhai et al., 2023]. scBERT is a pre-trained deep cell-type annotation model that leverages the Bidirectional Encoder Representations from Transformers (BERT [Devlin et al., 2018]).

Significantly different from existing methods, this work proposes a biologically interpretable cell annotation method that considers biological information. The proposed model not only accurately achieves cell annotation but also exploits the contextual relationships that exist between gene expression, providing meaningful biological gene sub-vector selection for downstream biological analysis.

3 Methodology

In this section, we present the details of our proposed SCTrans tailored for cell type annotation. As illustrated in Figure 2, we design an encoding transformer E to learn the representation of gene expression sub-vectors, and adopt an effective sub-vector selection mechanism to identify the salient ones, and facilitate gene-selective cell type annotation.

3.1 Random Sub-vector Completion

Based on the nature of gene co-expression, we adopt a selfsupervised training strategy of random gene sub-vector completion to learn the representation of gene sub-vector for cell type annotation. This process is designed to effectively explore the intrinsic contextual relationships among gene expression patterns at various scales. As illustrated in Figure 2.(a), the gene sequencing data X can be sub-vectorified into multiple gene expression sub-vectors X, and random masks Ω are generated to mask the majority of these sub-vectors, the masked regions are denoted as XΩ. Inspired by Masked Auto-Encoder [He et al., 2022], one of the training goal is to predict the masked sub-vectors, conditioned on the unmasked ones. To achieve this, the unmasked sub-vectors XΨ are fed into the multi-scale encoding transformer E (Ψ denotes the unmasked position index). Specifically, we perform attention computation over multi-scale Queries {Q(i) } and global KeyValue pairs {K, V } to explore the interdependency and learn the representation of gene sub-vectors across different scales, formulated as follows:

where Wq,k,v and bq,k,v are learnable model parameters, N denotes the number of sub-vector sizes, I(·, i) represents the function to produce gene sub-vectors with size i. In the above equation, Q(i) involves the gene sub-vector representations with difference sizes. With the aggregation of multiple transformer blocks, E explores the interdependency of gene subvector with different scales, and the output of multiple blocks are concatenated together to build the final output E(XΨ):

where Λ denotes the number of feature channels. The resulting representation E(XΨ) together with the mask Ω are fed into a simple decoding transformer D to predict the masked gene sub-vectors as follows:

By minimizing the above mean squared error (MSE), the multi-scale encoding transformer is trained to complete the unmasked gene sub-vectors, such that their intrinsic contextual relationships are explored and leveraged.

3.2 Gene-selective Cell Type Annotation

To induce our model to concentrate on salient gene subvectors specific to cell annotation and provide interpretable results, we incorporate a lightweight encoding transformer denoted as ES and perform mutual learning with E. The idea behind our selection module is to utilize the attention scores among multi-scaled Queries and global Keys to distinguish the gene sub-vectors playing more important role than the others, and the formulation is expressed as follows:

where the attention score S(i) is calculated by aggregating the multi-scale Queries {Qi } and the global Key K. To indicate salient gene sub-vectors, a coefficient matrix M(i) is generated in the above equation. The sigmoid function is used in conjunction with a temperature factor T and a learnable threshold θ to modulate the attention score. The value of the elements in M(i) ranges from 0 to 1 to indicate the importance of each gene sub-vector at different positions.

To identify the meaningful gene sub-vectors for cell type annotation, M(i) is further used to guide the lightweight encoding transformer ES. Specifically, we apply M(i) to weight the multi-scale gene sub-vector, and obtain the corresponding Queries as follows:

where ˆQ(i) denotes the weighted queries with sub-vector size i. Different from E, we perform attention computation at each transformer block of ES to mainly explore the relationship of the highlighted gene sub-vectors.

On the other hand, we incorporate a linear classifier C on top of E to infer the cell types: ˆY = C(E(XΨ)), where ˆYj denotes the predicted probability distribution of the jth cell. The predictions are evaluated by using the crossentropy-based loss function Leval as follows:

Minimizing Leval enforces E and C work together to produce the cell type annotation as accurate as possible. To associate gene sub-vector selection with specific cell types, we further perform mutual learning between E and ES by imposing prediction consistency regularization as follows:

where ˆ Yj = C(ES(XΨ)) denotes the cell type prediction from the representation learnt by ES. As will be shown in the experiments, the integration of multi-scale encoding and re-weighting strategies contributes to our cell type annotation performance.

3.3 Model Optimization

Considering that the representation discrepancy among cell types is expected to be significant. For this purpose, we further incorporate a conditional contrastive learning term to regularize the encoding transformer. We can construct semantically congruent and incongruent pairs according to cell type. We aim to push the gene representations of the cells of the same type closer, while moving away from other types of cells. The conditional contrastive loss function is defined as follows:

where Xj (Xk) represents the j (k)-th cell, and τ is a temperature parameter.

By integrating the above four aspects: gene sub-vector completion, cell type prediction, mutual learning and contrastive regularization, the overall optimization of the proposed SCTrans is expressed as follows:

As shown in the above equation, the multi-scale encoding transformer is jointly optimized with other constituent networks to learn cell-type-aware representation for cell type annotation, and the discovered salient gene sub-vectors lead to interpretable results.

4 Experiment

In this work, we conduct experimental studies to evaluate our method on seven representative benchmark datasets, the tissues of datasets including Peripheral Blood Mononuclear Cell (PBMC), Pancreas, Liver and Lung. Table 1 provides concise descriptions of these datasets.We carefully selected these tissues due to their critical roles in health and disease. Our objective in utilizing these diverse datasets is to demonstrate the flexibility and effectiveness of our method.

4.1 Comparison Results

In this section, we compare our model with several representative cell type annotation methods, including scBERT [Yang et al., 2022], Seurat [Hao et al., 2021], SingleR [Aran et al., 2019], CellID [Cortal et al., 2021], scmap [Kiselev et al., 2018], scNym [Kimmel and Kelley, 2021], and Scibet [Li et al., 2020a]. It is worth noting that both CellID and scmap include two variants with different strategies: CellID cell and CellID group for CellID, and scmap cell and scmap cluster for scmap. We evaluate the effectiveness of these cell type annotation methods using classification accuracy and F1 Score, where higher values indicate better performance. For each dataset, we randomly select 80% of the total data for model training and use the remaining 20% as test data. Table 2 provides a quantitative comparison between our model and other competing methods across seven datasets. Based on these results, we make the following observations: 1) PBMC (Zheng68k) presents a challenge in scRNA-seq data, as indicated by the unsatisfactory results obtained by competing methods. In contrast, our method consistently achieves the best results, with an accuracy value exceeding 0.8, significantly outperforming other competing methods. 2) While scBERT achieves the second-best results, our proposed method significantly outperforms all competing methods across all seven datasets in terms of accuracy and F1 score. It consistently attains an accuracy exceeding 0.95 for all datasets except Zheng68k.

Moreover, we conduct a detailed comparison of the performance of our method and scBERT for each cell type. Figure 3 displays the confusion matrix to illustrate the classification results of our model and scBERT across all cell types on the Zheng68k dataset. Our approach outperforms scBERT in classifying every cell type, with particularly notable improvements in classification performance for the'CD4+/CD45RA+/CD25- Naive T' and 'CD4+/CD45RO+Memory' cell types.

In addition, we explore the impact of varying training sample sizes on model performance, comparing our method with scBERT, Scibet, Seurat and CellID cell as depicted in Figure 4. Specifically, we randomly select 10%, 30%, 50%, 70%, and 90% of PBMC cells from the Zheng68K dataset as training data, while the remaining samples serve as testing data. As observed, our model consistently outperforms other methods in terms of accuracy across all settings.

4.2 Model Analysis

Robustness Analysis on Imbalanced Data

In scRNA-seq data, the imbalanced distribution of categories among different cell types poses a challenge. To assess the efficacy of the proposed method in handling imbalanced scRNA-seq data, we specifically chose four distinct PBMC cell populations ('CD19+ B', 'CD8+ cytotoxic T', 'CD34+', and 'CD8+/CD45RA naive cytotoxic cells') from the Zheng68K dataset. In detail, we randomly sampled cells from these types, obtaining counts of 100, 10,000, 100, and 10,000, respectively, to construct the training dataset. For model testing, 100 cells were randomly selected as the testing dataset. As depicted in Figure 5, in the same experimental setup, scBERT exhibits low annotation accuracy for recognizing cell types, whereas our model consistently achieves satisfactory results across all cell types. This underscores the robustness of our method in classifying imbalanced data.

Robustness Analysis on Batch Effects

In the collection of scRNA-seq data, it is common to integrate multiple samples from different sequencing platforms, leading to batch effects. These batch effects can significantly impact the accuracy of experimental results. To assess our model's capability to overcome batch effects, we conduct experiments across datasets, integrating human pancreas datasets generated from diverse sequencing technologies, including Baron, Muraro, Segerstolpe, and Xin. These datasets were harmonized into a unified dataset by aligning common genes and preserving four shared pancreas cell types ('alpha', 'beta', 'delta', and 'gamma').

In this setting, three of the datasets were utilized for training, while the fourth served as the test set. We apply our SCTrans method to this constructed dataset (cross-dataset), and for comparative analysis, we also perform tests on single source data (intra-dataset), where both the training and test sets originated from the same human pancreas dataset, such as Baron or Muraro. As illustrated in Figure 6, our model demonstrates high accuracy and F1 scores on the unified constructed dataset. Notably, it achieves accuracy results that are comparable to, or even better than, those obtained on single source data. These findings indicate that the representation learning from scRNA-seq data facilitated by our model is robust against batch effects.

Biological Meaningful Gene Sub-vector Selection

To assess whether our model can concentrate on the most influential gene segments for cell type recognition task, we conduct experiments involving the selection of a specific ratio of relevant gene sub-vectors selected by SCTrans shown in Figure 7. As scBERT lacks the capability for gene segments selective, we randomly remove a certain ratio of gene subvectors. Subsequently, we compare the resulting accuracy and F1 scores under different settings. Notably, our model demonstrates robust performance, maintaining high accuracy and F1 scores even after the removal of a ratio of gene subvector through selection. In contrast, scBERT experiences a notable decline in classification performance as the ratio of removed gene sub-vectors increases.

Moreover, we anticipate that our SCTrans model can identify meaningful and interpretable gene sub-vectors, which would be highly valuable for downstream biological analyses. To validate this, we undertake a meticulous selection of salient gene sub-vectors and evaluate their biological significance. Specifically, if a marker gene is present within the selected gene sub-vector, this sub-vector is considered biologically meaningful. As depicted in Figure 8, we compare the gene sub-vectors selected by SCTrans against those selected randomly. The results from this comparison indicate that SCTrans is adept at selecting biologically relevant gene sub-vectors. Additionally, SCTrans's selection varies from 0.67 to 0.14, we think this may due to the meaningful tokens's mapping is not average, and the meaningful tokens have weak relation with gene expression in some specific situation. SCTrans's selection may provide insights from other views.

Combining the outcomes from Figures 7 and 8, we can conclude that SCTrans is not only effective in gene selection, which contributes to accurate cell annotation, but also excels in providing biologically meaningful and interpretable results, demonstrating its multifaceted utility in the realm of biological research.

Representation Learning

Figure 9 (right) depicts the UMAP visualization of gene data representations learned by our model, while the left panel illustrates the UMAP visualization based on the original gene data. The visualizations clearly indicate that the representations obtained by our model are more distinguishable than those derived from the original gene data, leading to superior results in terms of ARI scores.

Ablation Study

To further validate the significance of multi-scale attention and contrastive learning in the proposed model, we compare the proposed SCTrans with two variants:

• Without the Multi-scale Attention strategy (SCTrans w/o MA): In this setting, only the single-scale attention will be considered.
• Without Conditional Contrastive Regularization (SCTrans w/o CCR): In this setting, we disable the conditional contrastive regularition termLctrst in Eq.(8).

In detail, we conduct an ablation study on the Zheng68k, MacParland, and Lung datasets to assess the impact of these two important components. As shown in Table 3, the proposed model with multi-scale attention and contrastive learning outperforms the two related variants, validating the effectiveness of these two components in the proposed method.

Additionally, we conduct an evaluation to assess the impact of the multi-scale attention strategy in terms of the reconstruction task. As depicted in Figure 10, the results clearly demonstrate that the implementation of the multi-scale attention strategy significantly enhances the reconstruction process, both in terms of efficiency and accuracy.

5 Conclusion

In this study, we present SCTrans, a novel multi-scale scRNA-seq Sub-vector Completion Transformer designed to overcome the limitations of current cell type annotation methods in scRNA-seq data analysis. SCTrans capitalizes on the richer expressiveness of gene sub-vectors and incorporates self-attention and contrastive regularization techniques. This approach allows SCTrans to excel in identifying crucial genes associated with specific cell types, leading to more accurate cell type annotation compared to current state-of-the-art methods. Furthermore, SCTrans performs attention computation among genes, potentially identifying salient ones, and provides interpretable results for downstream biological research and analysis, enhancing its utility in advancing our understanding of cellular behaviors and functions.

Acknowledgments

This work was supported in part by the National Natural Science Foundation of China (Project No. 62072189, 62106136), in part by the Research Grants Council of the Hong Kong Special Administration Region (Project No. CityU 11206622), in part by the GuangDong Basic and Applied Basic Research Foundation (Project No. 2020A1515010484, 2022A1515011160, 2022A1515010434, 2023A1515030154), and in part by TCL Science and Technology Innovation Fund (Project No. 20231752).

Contribution Statement

L.L. and W.X. are joint first authors, conceived of and designed the computational method; C.L. and S.W are joint corresponding authors. L.L. implemented the main algorithm; L.L. and W.X. did experiments and interpreted the results; X.D.W. and W.J.S. provided an analysis of the biological insights of the experiments; W.X., C.L., S.W. and L.L. wrote the manuscript; W.X., C.L., S.W. and H.S.W revised the manuscript.