Realistic Cell Type Annotation and Discovery for Single-cell RNA-seq Data

1 Introduction

Single-cell RNA sequencing (scRNA-seq) technologies allow us to measure gene expressions in millions of single cells and promise to provide high-resolution insights into the complex cellular ecosystem [Ziegenhain et al., 2017]. Cell annotation is the fundamental step in analyzing scRNA-seq data [Luecken and Theis, 2019]. In recent years, with more and more well-annotated scRNA-seq data becoming available, researchers turn to use neural networks to achieve automatic annotation of cell types [Cao et al., 2019]. Naturally, suppose Cr and Ct represent the label sets of reference and target data, respectively. Earlier developed methods are based on a close-set assumption, which can be expressed as Ct ⊆ Cr. However, this assumption is difficult to satisfy for data in the wild[Xu et al., 2021]. Therefore, to take into account a more realistic situation, the open-set scenario is introduced, that is, Cr ⊂ Ct, and several methods are proposed to settle this task, aiming to annotate cells with cell type labels in reference data or a unified "unassigned" label [Kimmel and Kelley, 2020].

Although the existing methods have achieved remarkable progress, they can not carry out further fine-grained analysis for them, which is not conducive to the subsequent downstream analysis [Brbi´c et al., 2020]. To address this need, we propose a more practical and challenging annotation task called realistic cell type annotation and discovery for scRNAseq data, whereby cells from novel cell types are given cluster labels instead of "unassigned" label and cells from seen cell types are given cell type label. One may argue that we can first use the annotation methods for the open-set scenario to find cells with "unassigned" label and then further cluster them into different groups. However, since the annotation results have a significant effect on the subsequent clustering process [Chen et al., 2020b], completely separating two processes is not conducive to problem-solving. Furthermore, we can also prove experimentally that the two-step approach only provides sub-optimal results.

However, settling this new task may face some challenges. First, the lack of label supervision for novel cell types will cause the model to be biased towards the seen cell types, thus further generating an imbalanced prediction state. Second, removing batch effects and other confounding noises, while maintaining biological signals of interest, is also essential but challenging for exploiting the clustering structure in target data [L¨ahnemann et al., 2020]. In this paper, we propose an end-to-end algorithm called scPOT based on a unified optimal transport (OT) framework to address these issues. For the cell type discovery, we introduce an OT-based prototypical selfsupervised learning paradigm to facilitate both global discrimination of clusters and local consistency of cells, which can avoid over-reliance on reference supervision and help recognize the whole cell types automatically. With regard to seen cell type annotation, we design an unbalanced OT-based partial alignment method to detect the common cells in the target data, which can uncover the intrinsic difference between reference and target data based on the statistical information of the assignment matrix. By leveraging the clustering accuracy on reference data, we propose a solution to estimate the number of cell types in target data, which is a challenging and poorly investigated problem in single-cell annotation. Lastly, to evaluate the performance of scPOT comprehensively, we choose various comparison baselines and build the intra-data and cross-data benchmarks on the basis of massive, highly imbalanced scRNA-seq data.

We highlight the main contribution as follows:

• We propose a new, practical, and challenging task called realistic cell type annotation and discovery in the singlecell annotation field. To solve this problem effectively, we further propose a novel method named scPOT.
• We introduce a unified OT framework based on a cellprototype alignment schedule to achieve seen cell type annotation and novel cell type clustering simultaneously.
• An easy yet effective solution is designed for the challenging problem of estimating the overall cell type number in target data.
• Comprehensive evaluation benchmarks are constructed to validate the practicality of scPOT, and deeper analyses show the effectiveness of its individual components.

2 RelatedWork

Single-Cell RNA-Seq Data Clustering. As an unsupervised learning branch, clustering is widely used for identifying cell types [Lakkis et al., 2021]. Early efforts are mostly based on traditional dimension reduction and hard clustering methods [Satija et al., 2015]. However, since the scRNA-seq data possess the characteristics of high dimension, sparseness, and complex nonlinear relationships, the traditional clustering methods might obtain unsatisfactory results. Recently, with the breakthrough of deep learning, several deep clustering methods have emerged to serve scRNAseq data. scziDesk clusters the cell population in the learned latent space by a soft self-training K-means algorithm [Chen et al., 2020a]. scNAME incorporates a mask estimation task and a neighborhood contrastive learning framework to cluster cells [Wan et al., 2022]. As a semi-supervised clustering method based on the capsule network, scCNC integrates domain knowledge into the clustering process [Wang et al., 2022]. However, although these methods can discover novel cell types in target data, they cannot recognize the seen cell types that previously existed in reference data.

Single-Cell RNA-Seq Data Annotation. The traditional cell annotation methods usually first cluster cells and then manually annotate these clusters to different cell types based on marker genes, which is time-consuming and subjective [Kiselev et al., 2019]. With the tremendous increase of wellannotated scRNA-seq datasets, more and more studies turn to exploring automatic annotation methods [Shao et al., 2020]. ItClust is a transfer learning-based method that takes advantage of cell-type-specific gene expression information learned from reference data [Hu et al., 2020]. MARS applies metalearning to encourage the same cell types to have similar features while those of different cell types are far apart [Brbi´c et al., 2020]. scNym is a gene expression knowledge integration framework that uses semi-supervised and adversarial learning techniques [Kimmel and Kelley, 2020]. scArches uses transfer learning and parameter optimization to enable reference building and contextualization of target data [Lotfollahi et al., 2022]. Overall, these methods can only roughly annotate cells from novel cell types with "unassigned" label, which is not conducive to subsequent downstream analysis.

3 Method

We first give some notations. In realistic cell type annotation and discovery task, we are given some labeled reference data Dr = {(xri , yr i )nr i=1 } and unlabeled target data Dt = {(xti )nt i=1 }, which can come from the same scRNA-seq dataset or different scRNA-seq datasets. The label sets of reference and target data are denoted as Cr and Ct, respectively. In our problem, we assume that Cr ⊂ Ct; furthermore, the seen label set is defined as Cs = Cr ∩ Ct, and the novel label set is defined as Cn = Ct\Cr. The goal is to assign either seen cell type labels or clustering labels to cells in the target data.

Considering the traits of scRNA-seq data, we assume that {xi}nr+nt i=1 follows a zero-inflated negative binomial distribution and use a denoising autoencoder model to reconstruct data [Eraslan et al., 2019]. Inspired by the recent progress in self-supervised learning [He et al., 2020], we use a data augmentation strategy to generate different views of gene expression, which can capture the correlations across genes better. The detailed information can be seen in the supplementary. In order to assign an annotation or clustering label for each cell, we attach two prototype-based classifiers Φr and Φt to the latent layer. Φr projects the l2 normalized embedding zi into one of the |Cr| seen cell types together with a similarity vector sri , where sri = V r zi and V r = [vr 1 , vr 2 , ..., v|Cr r |]T is the l2 normalized reference prototype matrix. Similarly, Φt maps zi to one of the |Cr ∪ Ct| clusters together with a similarity vector sti , where sti = V t zi and V t = [vt1 , vt2 , ..., vt |Cr∪Ct | ]T is the l2 normalized target prototype matrix. The value of |Cn| can be estimated and entered into the model as a prior. The specific estimation method will be introduced later. Since we also take the augmented data as input, its embedding representation and predictive probability can be written as ˜ zi and ˜ si, respectively. In the testing phase, we match each reference prototype with the nearest target prototype by the cosine similarity. Then for each target cell, we calculate the maximum component index of sri ∪sti and take it as the prediction label.

3.1 OT-Based Prototypical Representation Learning for Novel Cell Type Discovery

To exploit the cell type structure and facilitate the feature representation learning for the target data, we propose a prototypical self-supervised clustering technique from the perspective of OT [Villani, 2009], which solves a full mapping problem between target cells and target prototypes. Specifically, we briefly recall the well-known OT problem. Let Ωd = {µ : µT 1d = 1} denote a simplex set, where 1d refers to a d-dimensional vector of all one. Given two simplex distribution vectors α ∈ Ωm and β ∈ Ωn, we can define the transport polytope of α and β as follows,

The transport polytope A(α, β) can also be interpreted as a set of all possible joint probabilities of (F, G), where F and G are two d-dimensional random variables with marginal distribution α and β, respectively. When we have a similarity matrix M, the joint probability matrix Q∗ , also called as coupling matrix mapping α to β, can be quantified by optimizing the following maximization problem,

where ϵ > 0 and H(Q) = - ∑ij Qij logQij is the entropy regularization. The optimal Q∗ is shown to be unique with the form Q∗ = Diag(u) exp(M/ϵ)Diag(v), where u and v can be solved by sinkhorn-knopp algorithm [Cuturi, 2013].

Given B normalized target feature vectors Z t = [zt1 , zt2 , ..., z tB ]T , we are interested in mapping them to the target prototypes V t . We denote this mapping by Qtt = [qtt 1 , qtt 2 , ..., qtt B ]T , and optimize Qtt to maximize the similarity matrix Stt = [stt 1 , stt 2 , ..., sttB ]T between the target features and target prototypes by solving the following OT problem,

The constraint of marginal uniform distribution enforces that on average each prototype is selected at least B |Cr∪Ct| times in the batch. The solution Qtt satisfies the condition that the sum of each row equals to 1 |Cr∪Ct| strictly. To further obtain the prototype assignment distribution for each target cell, we need tomultiplyQtt by |Cr∪Ct|, i.e., Qtt ←Qtt ×|Cr∪Ct|, to ensure each row of Qtt is a probability vector. Similarly, for another augmented branch, we can also obtain the soft assignment matrix ˜Qtt . To encourage the separability of clusters and guarantee the discrimination of features, we propose a swapped prototypical representation learning loss as follows,

rameter. By replacing stt ik with ˜ stt ik , the ˜ ptt ik can be obtained as same way as ptt ik . We observe that a strong entropy regularization (i.e., using a high ϵ) in OT generally leads to a trivial solution where all cells collapse into a unique representation and are all assigned uniformly to all prototypes [Caron et al., 2020]. Hence, in practice, we keep ϵ low.

Note that Lpro does not involve the cells in the reference data and the local structure of the whole data cannot be captured. Therefore, we propose to encourage the prediction consistency of similar cells and utilize the pairwise constraint to group the cells from the same cell types. To achieve this, we rely on the ground-truth annotations from the reference data and pseudo-labels generated on the target data. Specifically, for the reference data, we already know which pairs should belong to the same cell types according to groundtruth labels. To obtain the pseudo-labels for the target data, we compute the cosine distance between all pairs of normalized feature embeddings in the reference and target batches. We then rank the computed distances and for each target cell generate the pseudo-label for its most similar neighbor. Given two mini-batch feature embeddings {zr i }Bi=1 ∪{zt j }Bj=1 , we denote its closest set as {z′i r }Bi=1 ∪ {z′j t }Bj=1 . Note that {z′i r }Bi=1 is always correct since it is generated using the ground-truth labels. Similarly, for the augmented branch {˜ zr i }Bi=1 ∪ {˜ zt j }Bj=1 , we also can obtain the corresponding closest set {˜ z′i r }Bi=1 ∪ {˜ z′j t }Bj=1 . Then our pairwise objective is defined as a modified binary cross-entropy loss,

where σ is the sigmoid function and ⟨.⟩ refers to the vector inner product operation. We update similarities and positive pairs in an online fashion and thus benefit from improved feature representation during training. Eventually, to encourage both global discrimination of clusters and local consistency of cells, we unify Lpro with Lpair as follows,

3.2 OT-Based Prototypical Partial Alignment for Seen Cell Type Annotation

For seen cell type classifier Φr, we can use the known label information of reference data to train it based on the standard cross-entropy loss. Given B reference cells, we have

where ϕ is the softmax function. Since the target data possess some novel cell types, we cannot match all target cells with reference prototypes and we should consider partial alignment to avoid misalignment between them. To extract the shared knowledge across data, we propose an unbalanced OTbased seen cell type annotation method, which solves a partial mapping problem between target cells and reference prototypes. Specifically, OT ϵ (M, α, β) is not suitable for the partial alignment problem, because its solution Q∗ satisfies the condition of A(α, β) strictly. So unbalanced OT is designed to relax the conservation of marginal constraints by allowing the system to use soft penalties, which can be formulated as,

where KL is the Kullback-Leibler Divergence. This optimization problem can be solved by the generalized sinkhornknopp algorithm [Chizat et al., 2018]. Given B normalized target features Z t = [zt1 , zt2 , ..., z tB ]T , to achieve mapping them to the reference prototypes V r , we can obtain the optimal assignment matrix Qrt by optimizing the following unbalanced OT objective,

where Srt = [srt 1 , srt 2 , ..., srt B ]T . Note that the target cells from novel cell types will be assigned with relatively low weights in Qrt . Based on this observation, we select target cells with top confidence as seen cell types by mining the statistical information of the assignment matrix Qrt .

We first normalize the assignment matrix, i.e., Qrt ← Qrt / ∑ Qrt . Then we generate scores based on the statistical property of Qrt , which depicts the geometrical relationship between target cells and reference prototypes. For i-th row of Qrt , it can be seen as a prediction vector of i-th target cell and we can obtain its pseudo-label ˆ yt i by argmax operation. And we assign it confidence score ξ rt i with the maximum value of i-th row in Qrt , i.e., ξ rt i = max({Qrt i1 , Qrt i2 , ..., Qrt i|Cr| }). A higher score ξ rt i implies that zt i is relatively closer to a source prototype than any other cells and is more likely from a seen cell type. Meanwhile, to select target cells with top confidence, we also evaluate the confidence score of j-th reference prototype with the sum of j-th column, i.e., ζ rt j = ∑Bi=1 Qrt ij . Analogously, a higher score ζ rt j means vr j is a more reliable reference prototype, which is assigned to target cells more frequently. Then the target cells from seen cell types can be detected by statistic values, i.e.,

where δt i = 1 indicates the i-th target cell is regarded as from the seen cell types with top confidence, which can be assigned with pseudo-label ˆ yt i . Similarly, the same procedure can be implemented on the augmented branch to obtain ˜ δt i and ˆ˜ yt i . Naturally, we can use the pseudo-labels of selected target cells to compute the swapped partial alignment loss below,

where ϕ is the softmax function. Combining Lce with Lalign, we give the training loss on the reference prototypes,

Overall loss. Together with the data denoising loss Lden (see supplementary), we give the overall training objective as

where λ and γ are two weight hyper-parameters.

3.3 Estimation of the |Ct| Value

Here, we propose a solution to the challenging and underinvestigated problem in cell annotation: estimating the cell type number |Ct| in target data. Almost all annotation methods assume that the number of |Ct| is prior. However, this assumption is unrealistic in the real world. This calls on the community to develop a method for estimating |Ct|. Our main idea derives from the information available in Dr. Specifically, we perform sphere k-means clustering on the whole dataset Dr ⋃ Dt and then evaluate clustering accuracy only on the reference data Dr. Let | ˆ Ct| represent the estimated value. If | ˆ Ct| > |Ct|, then ˆ Ct - Ct can be called the extra cell types, and all cells assigned to extra cell types are mispredicted. Similarly, if | ˆ Ct| < |Ct|, then Ct - ˆ Ct can be called the extra true cell types, and all cells with those cell types are predicted incorrectly. Based on this analysis, whether | ˆ Ct| is higher or lower than |Ct| will have a negative impact on the clustering accuracy on Dr. In other words, the clustering accuracy on Dr will be maximized when | ˆ Ct| = |Ct|. According to this intuition, we use AC = f(| ˆ Ct|, Dr) to measure the clustering accuracy on Dr, which we optimize with Brent's algorithm to find the optimal | ˆ Ct| [Brent, 2013].

4 Experiment

4.1 Setup

Dataset. Our experiments consist of intra-data annotation and cross-data annotation. For the former, we collect 10 datasets sequenced from different organisms. The cell numbers range from 6462 to 110704, and the cell type numbers vary from 9 to 45. Unless otherwise noted, we first divide all cell types into 50% seen and 50% novel. Then we select 50% samples in seen cell types as Dr and the rest as Dt. For the latter, we select 10 groups of datasets. Each group consists of a reference dataset and a target dataset, and the batch effect exists between them. The basic information of these datasets can be seen in the supplementary.

Baselines. Our task is to establish a new cell annotation setting for which no ready-to-use baselines exist. Thus, we compare scPOT with recently developed scRNA-seq clustering and annotation algorithms, including three clustering methods (scziDesk, scCNC, and scNAME) and four annotation methods (MARS, ItClust, scNym, and scArches). For clustering methods, only scCNC participates in training with both Dr and Dt, while the other two trains only on Dt. We report their clustering performance on seen and novel cell types. For annotation methods, we first use them to classify target cells into seen cell types and identify the "unassigned" group. Next, we apply k-means clustering on the "unassigned" group to obtain novel clusters. Detailed information on these baselines can be seen in the supplementary.

Evaluation Protocols. We report the classification accuracy on seen cell types and clustering accuracy on novel cell types for scPOT and annotation baselines while reporting clustering accuracy on both seen and novel cell types for clustering baselines. Specifically, to compute the clustering accuracy, we apply the Hungarian algorithm to solve the optimal assignment problem [Kuhn, 1955]. When reporting accuracy on all cell types, we solve the optimal assignment problem on both seen and novel cell types. The reported accuracies are the mean values of three runs.

Implementation Details. Our algorithm is implemented by PyTorch, and we conduct the experiments with 2 Tesla A100 GPUs. The two layers of the encoder are sized 512 and 256, respectively, and the decoder has the reverse structure of the encoder. The bottleneck layer has a size of 128. The training mini-batch size is set to 256, and the optimizer is Adam with a learning rate of 1e-4. The temperature τ is set to 0.1, and the loss weight parameters λ and γ are both set to 1.0. The parameters ϵ and κ in OT are set slightly differently in various datasets and the details can be seen in the supplementary. We first train the whole model using Lden loss with 600 epochs. Then, we apply the sphere k-means algorithm on target embeddings to obtain cluster centers as the initial values of target prototypes. The initialization of reference prototypes can be obtained by the mean values of reference embeddings based on ground-truth labels. Finally, we train the model with the overall loss Ltol until the predictions no longer change.

4.2 Results

Intra-data Annotation. To begin, we explore the performance of scPOT under the intra-data setting without batch effect. From the results in Tabel 1, we can conclude that scPOT consistently achieves stable and the best performance under three types of accuracy on most datasets. It is not worth surprising that scPOT gets such an excellent performance since the unified OT framework based on a cell-prototype alignment schedule can effectively realize label transferring of seen cell types and cell clustering of novel cell types by solving the optimal assignment problem. The results in Tabel 1 also fully show that the two-step strategy that first annotates cells of novel cell types with "unassigned" label and then further clusters them is not an appropriate solution to this new task. By comparison, although scNym can obtain relatively high annotation accuracy on seen cell types, and even higher than our method on individual datasets, it has a sharp drop in clustering accuracy on novel cell types. As an unsupervised annotation method, MARS allows for fine-grained analysis of novel cell types and has the ability to assign clustering labels for them. However, it can only provide sub-optimal results for both annotation and clustering accuracy. scCNC and ItClust are not excellent in both annotation and clustering accuracy, and ItClust even fails catastrophically in annotating seen cell types. As clustering methods, scziDesk and scNAME cannot provide satisfactory results for the reason that they do not utilize the information in reference datasets, which makes them less competitive. In summary, scPOT outperforms other competing methods under three kinds of accuracy in the intradata setting. This superior performance shows that scPOT achieves remarkable progress for this novel task.

Cross-Data Annotation. In this section, we turn to explore a more challenging setting, that is, the cross-data setting with batch effect. We compare scPOT with other seven competitive methods on ten groups of mixed datasets. As shown in Table 2, scPOT achieves consistently better results than other methods on most datasets, which demonstrates that scPOT can also deliver excellent performance for the crossdata situation. Moreover, compared with the intra-data setting, there is no significant decline in the performance of scPOT, indicating that scPOT could resist the effect of batch effect to some extent. scziDesk and scNAME are unsupervised clustering methods and do not use reference data, they do not have to deal with batch effects, but at the same time, they can not get the cell type knowledge from reference data. In comparison, since MARS and ItClust separate the learning process on reference data from the training process on target data, elevating susceptibility to batch effect, and in turn, leading to model overfitting and false cell type annotations.

4.3 Ablation Study

Robustness Analysis. Since the novel cell type number |Cn| determines the difficulty for methods to discover and cluster novel cells, it is imperative to explore the influence of the variation of |Cn| on the methods. We evaluate the performance of all eight methods on Quake 10x and Quake Smart-seq2 with 36 and 45 total cell type numbers, respectively. Here, |Cn| varies in the range of [4, 11, 18, 25, 32] for Quake 10x and [5, 14, 23, 32, 41] for Quake Smart-seq2. Figure 2(a) and Figure 2(b) are the line graphs that show the results intuitively. It is easy to see that the change of |Cn| has a huge effect on the overall accuracy of all methods and it is reasonable because |Cn| affects the specific gravity of different cell types and can therefore influence the performance of methods. From the Figure, we can conclude that no matter what value |Cn| takes, scPOT beats other methods with clear margins. Moreover, with the change of |Cn|, the overall accuracy of scPOT varies only slightly, suggesting the robustness of scPOT. On the contrary, other methods are affected by the variation of |Cn| to varying degrees. The overall accuracy of scNym, scArches and scCNC drops catastrophically with increasing |Cn|. It is noteworthy that the result for ItClust rises dramatically on Quake 10x and drops significantly on Quake Smart-seq2, which suggests its instability. Although the results of MARS, scziDesk and scNAME are relatively stable, they tend to provide sub-optimal results. Therefore, we can conclude that scPOT is robust for the variation of |Cn|.

Besides, since the ratio of labeled data determines how much information can be used for annotation and clustering, we explore its impact by conducting experiments on Quake 10x and Quake Smart-seq2 datasets. The ratio of labeled data varies in the range of [0.1, 0.3, 0.5, 0.7, 0.9], and Figure 2(c) and 2(d) describe the variation tendency overall accuracy of eight methods. We can find that scPOT still performs best, no matter what value the ratio of labeled data takes. However, the other seven methods all show an obvious downward trend. This result is in line with our speculation for the reason that these three methods are all to make the model learn the knowledge from the reference data first and then transfer the learned knowledge or model to the target data to make predictions, which is sensitive to the size of the reference dataset. In conclusion, this experiment provides intuitive evidence to confirm that scPOT can provide reliable and remarkable performance, even with a few labeled data.

Validity of the |Ct| Value Estimation Method. The estimated value of |Ct| determines the performance of methods for discovering novel cell types. Thus, it is imperative to conduct experiments to validate the validity of the |Ct| value estimation method. Specifically, Quake 10x and Quake Smartseq2 are used as our experimental data, whose total cell type numbers are 36 and 45, respectively. We study the case when |Ct| varies in the range of [-15,-10,-5, 0, 5, 10, 15] and "increment=0" means that the estimated value of |Ct| is equal to the true value. The results are shown in Table 3. For these two datasets, we can clearly see that the clustering accuracy gets the maximum value when the increment is 0, indicating the validity of our estimation method.

Effect of Lpro, Lpair and Lalign. Here, we carry out an ablation study on 10 real datasets to learn about the effect and performance gain of introducing Lpro,Lpair and Lalign in scPOT, respectively. The results are shown in Table 4. We can clearly see that removing Lpro results in the most significant effect on the overall accuracy, mainly because it plays a key role in the discovery and separation of novel cell types. It can also be seen that when the number of cell types in target data is relatively large, adding the local structure constraint Lpair has a greater effect on overall accuracy, in contrast, the effect of prototype-oriented partial alignment loss Lalign is more significant. Overall, we can conclude that the strategies we propose are of significant value to address this new task.

5 Conclusion

In this article, we propose a new, practical, and challenging task called realistic cell type annotation and discovery in the single-cell field and design a unified OT framework called scPOT to address it. scPOT mainly consists of two main parts, which are an OT-based prototypical learning paradigm for the novel cell type discovery and an OT-based partial alignment strategy to realize seen label transferring. We also introduce a solution for the estimation problem of the total number of cell types in target data. To evaluate the algorithm's performance, we carefully construct comprehensive baselines and benchmarks. The results on massive real datasets verify the superiority and robustness of scGAD compared to several annotation and clustering methods.

Contribution Statement

Yuyao Zhai and Liang Chen made the same contribution to this paper, and Minghua Deng is the corresponding author.