Distribution-Independent Cell Type Identification for Single-Cell RNA-seq Data

1 Introduction

Since being recognized as the yearly technology by Nature Methods in 2013, single-cell RNA sequencing (scRNA-seq) technology has seen substantial and swift development. The scale of the sequencing data has expanded, encompassing from a small group of dozens to hundreds of thousands, even millions [Regev et al., 2017]. Using gene expression profiling, scRNA-seq enables researchers to inspect the individual cellular-level variability of disease tumors. The crucial process of identifying cell types in the analysis of scRNA-seq data aids in comprehending the source of tissue heterogeneity. The standard cell type annotation method initially clusters the cell population before identifying cluster-specific marker genes. The cells are then classified according to the ontological functions of their genes. However, as the scale of sequencing data rapidly expands, identifying marker genes to annotate cells turns into an increasingly daunting and timeintensive task [Kiselev et al., 2019].

Given the abundance of extensively annotated scRNA-seq datasets, researchers have started to leverage classification and retrieval machine learning approaches to automate the process of transferring cell type labels from reference data to target data [Cao et al., 2020]. During the early stages, the majority of automated annotation methods were implemented within closed domains, meaning that every cell type present in the target data also existed in the reference data. However, in realistic scenarios, this limitation can be excessively stringent. Thus, researchers eventually introduced automated annotation tasks applicable to open domains [Xu et al., 2021]. To simplify terminologies, we shall denote cell types that are shared with both the reference data and the target data as seen cell types. Conversely, cell types absent in the reference data but present in the target data will be referred to as novel cell types. For tasks operating in open domains, it's a common practice among most annotation methods to classify these novel cell types under a category named "unassigned".

Given the demand for more detailed cellular and gene-level scrutiny within the "unassigned" group, the recent proposition of scGAD highlights the need to integrate cell classification and cell clustering within a unified framework [Zhai et al., 2023]. However, certain limitations reside in scGAD's algorithm. Firstly, its mutual nearest neighbor retrieval process, an essential component of model training, is noted for its time-intensiveness which may significantly hinder the algorithm's computational efficiency. Secondly, scGAD overlooks a crucial characteristic of scRNA-seq data: the imbalance of cell types [L¨ahnemann et al., 2020]. The differing frequencies of diverse cell types follow a long-tailed distribution where a minor fraction of classes heavily influence the overall data distribution, and many classes only link with a minuscule amount of instances. If we avoid formulating an algorithm specifically tailored to the peculiarities of these long-tailed distribution datasets, any derived solution would unquestionably be subpar.

To address the challenges discussed earlier, this study introduces a cutting-edge, distribution-independent framework for universal cell type identification termed scDET, premised on the concept of autonomously equilibrated dualconsultative contrastive supervision. The scDET framework encompasses two synergistic components: a contrastive learning branch and a pseudo-labeling branch. These branches work in tandem to offer mutual guidance, particularly for the imbalanced cell type annotation challenge within open-world settings. More precisely, the contrastive learning branch undertakes the role of distributional estimation throughout the model's training phase. This estimation acts as a means of regularizing the output of a linear classifier, thereby optimizing the generation of pseudo-labels. Conversely, the pseudo-labels thus obtained are selectively resampled and adjusted to counterbalance and augment the supervision provided to the contrastive learning branch. To amalgamate the insights from both branches and cultivate an improved representational space, we have engineered an innovative contrastive loss function predicated on pseudolabels. This function is designed to coalesce samples in the feature space by leveraging their respective positiveness scores, thereby enhancing the clustering process.

We highlight the main contributions as follows:

• We integrate the paradigms of long-tail and open-world learning within the context of scRNA-seq data annotation, pinpointing distinct challenges in this domain that remain insufficiently tackled by current methodologies.
• We design a novel distribution-independent cell type annotation paradigm that comprises a contrastive learning branch and pseudo-labeling branch, which work collaboratively to tackle the data imbalance and classify novel cell types simultaneously.
• We carry out comprehensive experimental evaluations across diverse datasets and engage in detailed comparative analyses with leading-edge benchmarks, thereby affirming the effectiveness of scDET.

2 RelatedWork

2.1 Cell Type Identification for scRNA-Seq Data

Unsupervised clustering and supervised classification represent the two primary methodological approaches for determining cell types within scRNA-seq data [Petegrosso et al., 2020]. With the accelerated advancement of deep learning technologies, a plethora of statistical learning approaches have been developed in both directions. In the realm of clustering, scziDesk identifies distinct cell populations by employing a soft self-training k-means algorithm within a lowdimensional feature space [Chen et al., 2020a]. scNAME enhances clustering accuracy by integrating a mask estimation strategy alongside a neighborhood contrastive learning framework [Wan et al., 2022]. As a semi-supervised clustering technique, scCNC embeds expert knowledge into the process using a capsule network framework [Wang et al., 2022]. On the classification front, scNym leverages a combination of semi-supervised and adversarial learning approaches to effectively incorporate gene expression information [Kimmel and Kelley, 2021]. scArches adopt transfer learning accompanied by fine-tuned parameter optimization to contextualize target datasets appropriately [Lotfollahi et al., 2022]. MARS utilizes meta-learning principles to reinforce the feature similarity amongst identical cell types [Brbi´c et al., 2020]. To streamline and integrate the processes of cell clustering and classification, scGAD has been introduced [Zhai et al., 2023]. This self-supervised learning architecture establishes a coherent linkage between the reference and target data, thus creating a unified framework for cell type identification.

2.2 Long-Tail Learning and Contrastive Learning

Long-tail learning is a typical machine learning scenario where the majority of samples belong to a few classes, while the minority classes have only a small number of samples [Zhang et al., 2023]. Existing research on long-tail learning mainly focused on fully/semi/self-supervised scenarios, with representative methods that highlight the minority samples via re-sampling [Shen et al., 2016], re-weighting [Cao et al., 2019], or knowledge transferring [Wang et al., 2017]. However, the requirement of the data distribution obstructs the application of fully/semi-supervised methods in our task, while self-supervised methods would disregard the available label information for known samples. Contrastive learning provides distinctive representations by controlling the instance similarity in feature space, achieving significant progress in recent years [He et al., 2020]. The existing literature generalizes the idea to different scenarios [Khosla et al., 2020], implements it for solving various downstream tasks [Chaitanya et al., 2020], and provides theoretical support [Wang and Isola, 2020]. However, few attempts have been made at contrastive-based open-world cell type identification. scGAD generates pseudo-positive pairs for closely aligned representations, but this paradigm could lead to another dilemma: the representations and pseudo-labels are interdependent, which means an inferior feature space could lead to false positive pairs, and in turn deteriorate the learning of feature space.

3 Method

3.1 ProblemFormulation

To commence, we'll clarify some notations. Our examination of a distribution-independent cell type identification task provides us with certain labeled reference data, denoted as Dr = {(xri , yr i )nr i=1 }, and unlabeled target data represented as Dt = {(xti )nt i=1 }. Both datasets may originate from either the same or different scRNA-seq datasets. The labels for Dr and Dt are represented respectively as Cr and Ct. In the problem under consideration, we assume that Cr is a subset of Ct. Moreover, the set of labels we have seen, or the seen label set, is defined as Cs = Cr ∩ Ct. Additionally, the novel label set is demarcated as Cn = Ct\Cr. The purpose of our work is to assign either seen cell type labels or clustering labels to cells within the target data. It is generally accepted that the total number of cell types within the whole dataset represented as Dr ∪ Dt, is known prior to the assignment, since we can estimate it efficaciously using existing methods.

Given the high-dimensional and sparse characteristics inherent in scRNA-seq data, we have developed a model composed of three primary components: a denoising autoencoder, a contrastive branch, and a pseudo-labeling branch (see Figure 1). The autoencoder compresses the input, represented as x, into an embedded feature that we signify as z. Following the compression, the input is then reconstructed through the use of z. The contrastive branch transforms the embedding, denoted as z, into a renewed representation symbolized as υ via a projection layer. Simultaneously, the pseudo-labeling branch categorizes the same embedding z into one of the classes within the set |Cr ∪ Ct|. This assignment is done with a probability p through classifier layer implementation. Motivated by the advancements in self-supervised learning [Liu et al., 2021], we have employed a data augmentation approach to produce an alternate view, ˜x, of gene expression, enhancing the ability to capture inter-gene correlations. Comprehensive details are available in the supplementary.

3.2 Revisiting the Contrastive Learning

Contrastive Learning provides distinctive and transferable representations by controlling the instance similarity in feature space, achieving significant progress in recent years [Chen et al., 2020b]. Specifically, it aims to find a projection function to acquire optimal feature representation υ of the gene expression input x, such that υ retains the discriminative semantic information of the input cell. The general contrastive learning loss function can be defined as,

where D is the training set and A(i) = D \ i. P(i) is the set of indices of positive pairs in A(i), | · | denotes the operation to compute the inner product similarity, and τ is the temperature parameter. For the unsupervised learning scenario, the positive set is formulated as the two views of the same cell. For the supervised learning scenario, the positive set can be the cells that share the same ground truth labels.

3.3 Dynamic Distribution Estimation

To enhance and re-balance the contrastive representation learning, we propose to learn an auxiliary classifier for pseudo-labeling. Our pseudo-labeling branch includes a cross-entropy loss Lr on reference data, a self-training objective Lt on target data, and a regularization term Lreg on the whole data. Formally, we use the prototype parameterized classifier and randomly initialize a set of prototypes {µi}|Cr∪Ct| i=1 , each standing for one cell type. Then we can calculate the assignment probability pi for each cell xi by softmax on cosine similarity between the hidden feature hi and the prototypes {µi}|Cr∪Ct| i=1 , i.e.,

and the soft pseudo-label qi is produced with a sharper temperature ˆ τ in a similar fashion. Considering that we have two different views of cells, the classification objectives are then simply cross-entropy (CE) loss between the predictions and pseudo-labels or ground-truth labels,

However, these two objectives do not guarantee the nontrivial solution of the model due to the prediction bias induced by weak supervision of novel cell types. So we propose to align the predictions with the data distribution to avoid nonactivated classifiers. The difficulty lies in the fact that the distribution of the training set is independent in our task, and simply using the distribution of reference set πDr or a balance prior results in inferior performance. Furthermore, we observed that estimating the distribution from the pseudolabeling branch itself could accumulate the estimation error and deteriorate model performance.

To this end, we perform estimation on the contrastivelearning branch as an alternative to avoid bias accumulation. Concretely, we perform k-means clustering on all samples in Dr ∪ Dt to obtain |Cr ∪ Ct| clusters with size {si}|Cr∪Ct| i=1 , and normalize the sample number {si}|Cr∪Ct| i=1 to frequency {πi}|Cr∪Ct| i=1 , i.e., πi = si/ ∑|Cr∪Ct| j=1 sj. Moreover, we need to determine the corresponding relationship between clusters and cell types. So we use the Hungarian optimal assignment algorithm [Kuhn, 1955] to map |Cs| clusters to each known cell type. For the remaining |Cn| clusters, we sort them by cluster size and assign them sequentially to novel cell types. Finally, we regularize the mean prediction of the pseudolabeling branch with the aligned distribution {πi}|Cr∪Ct| i=1 ,

where KL(·||·) is the Kullback-Leibler (KL) divergence between the two distributions, and ρ ∈ [0, 1] is a hyperparameter used to smooth the target long-tailed distribution. To balance estimation accuracy and computation overhead, we re-estimate the dataset distribution every m epoch. The training objective of the pseudo-labeling branch is as follows,

3.4 Debiasing and Sampling Process

In the last section, we regularize the pseudo-labeling branch with the estimated data distribution {πi}|Cr∪Ct| i=1 acquired from the contrastive-learning branch. In this step, we aim to transfer the knowledge from the pseudo-labeling branch to the contrastive-learning branch in turn and further enhance the representation learning, which is also beneficial to the estimation of data distribution. However, the long-tail distribution of the underlying dataset places additional requirements on knowledge transferring. On one hand, we should ensure that the knowledge to be transmitted is not affected by the long-tailed distribution. Meanwhile, this process should help the contrastive-learning branch cope with the issue of data imbalance and the lack of supervision for novel cell types.

We design a debiasing and sampling step of the pseudolabels to meet the two aforementioned requirements for knowledge transfer. First, similar to the previous work [Menon et al., 2020], we apply post-hoc logits adjustment based on the estimated distribution to the predicted logits to eliminate the bias caused by long-tailed distribution,

where ˆ pi denotes the rectified cell type probability prediction of cell xi and γ is a hyper-parameter. Next, to re-balance the learning process and filter low-precision pseudo-labels, we propose to sample the unlabeled cells. Specifically, we sample the pseudo-labels of unlabeled cells in a training batch ˆ pB = {ˆ pi : i = 1, 2, ...,B} according to their prediction cell type c, where ci = argmax ˆ pi. Formally, the cell type-wise sampling rate SRc can be defined as,

where CB ⊂ Cr denotes the set of labels that are involved in the current batch, and α, β ∈ [0, 1] are two hyperparameters. Setting α = β = 0 means sampling all pseudo-labels, while α = β = 1 means setting the sampling rate to be inversely proportional to the estimated number of samples in its cell types. Note that we prioritize selecting samples with higher prediction confidence in each cell type to remove the potentially false pseudo-labels. The sampled M(Dt) complements the original long-tailed distribution to serve as a re-balance term for Dr and also provides additional supervision for novel cell types. With this debiasing and sampling step, the sampled cells {Dr ∪ M(Dt)} and their corresponding rectified pseudo-labels mitigate the impact of the imbalance issue and compensate the unlabeled open-set cells simultaneously.

3.5 Soft Contrastive Learning

In the above two sections, we generate pseudo-labels via the pseudo-labeling branch with help from the contrastivelearning branch. Those pseudo-labels could provide additional information to improve the contrastive learning branch as well. In particular, we adopt a soft contrastive strategy to fully leverage the probabilistic information in pseudo-labels. We design a pair-wise positiveness score to adjust the contribution of different samples to the anchor cell. For cell pair (xi, ˜xj), the positiveness score wij = Sim(ˆ pi, ˆ˜ pj). wij is obtained by calculating the similarity of the rectified cell type assignment distribution between xi and ˜xj. In practice, we implement the similarity metric with the dot product operation, which can be interpreted as the probability of cell xi and ˜xj belonging to the same cell type. Finally, we formulate the soft contrastive loss by incorporating w into Equation (1),

Minimizing the above equation encourages the similarity of features between two samples to be proportional to the corresponding positiveness score. In this way, we effectively transfer knowledge from the pseudo-labeling branch to contrastive learning. The training objective of the contrastive learning (CL) branch consists of an unsupervised CL loss on all cells, a supervised CL loss on labeled cells, and the proposed soft CL loss on both labeled and sampled unlabeled subset,

It is worth noting that the backbone of the two branches is shared. So the two-branch structure only slightly increases computational overheads. During inference, we utilize the pseudo-label learning branch and obtain predictions by finding the maximum component of the classification probability. Overall loss. Together with the reconstruction loss Lrec (see supplementary), we give the overall training objective as,

where λ1 and λ2 are two weight hyper-parameters.

4 Experiment

4.1 Setup

Data preparation. Our experiments encompass two types of annotation scenarios: intra-data annotation and inter-data annotation. In the process of intra-data annotation, we meticulously assembled a collection of 10 distinct datasets, each obtained from various organisms. The number of cells within these datasets exhibits a considerable range, with a minimum of 6,462 and a maximum of 110,704 cells. Additionally, the diversity of cell types is notable, fluctuating between 9 and 45 different types. In the absence of specific indications to the contrary, our standard procedure involves bifurcating all the cell types into two equal cohorts: 50% are classified as

"seen" and the remaining 50% as "novel". After this classification, we randomly select 50% of the samples from the seen cell types to constitute the reference set, denoted as Dr, while the remainder of the samples is designated to form the testing set, Dt. Regarding inter-data annotation, our approach entails the selection of 10 paired groups of datasets. Each pair is composed of a reference dataset and a corresponding target dataset, between which batch effects are observed. The foundational attributes and details about these datasets are comprehensively enumerated in the supplementary.

Comparison baseline. Our research investigates a new task, wherein scGAD emerges as the optimal baseline method for our comparative analysis. Furthermore, we evaluate the performance of our method against an array of specialized techniques, including three clustering algorithms, i.e., scziDesk, scNAME, and scCNC, and three annotation methodologies, i.e., MARS, scNym, and scArches, each tailored for scRNAseq data. In the context of the clustering approaches, it is notable that only scCNC is trained using both datasets Dr and Dt, whereas scziDesk and scNAME are exposed solely to Dt during their training phase. We present their clustering efficacy for both known and novel cell types, offering a comprehensive assessment of their performance. Turning to the annotation strategies, we initially deploy these methods to categorize the target cells into established cell types and segregate an "unassigned" category. Subsequently, we implement k-means clustering on the cells within this "unassigned" classification to delineate novel cellular clusters. Lastly, we run each compared method in their default settings.

Evaluation metrics. Similar to scGAD, our study presents the classification accuracy for recognized cell types, denoted as Cs, alongside the clustering accuracy for previously unidentified cell types, symbolized as Cn. These metrics are used for evaluating the performance of annotation baselines. Conversely, the performance of clustering baselines is assessed by examining clustering accuracy across both known and novel cell types. To determine the clustering accuracy more precisely, the Hungarian algorithm is employed to address the optimal assignment problem, as detailed in [Kuhn, 1955]. In scenarios where accuracy across the collective set of cell types, Cs ∪ Cn, is being evaluated, the Hungarian algorithm is applied to ascertain the best assignment strategy for both seen and novel cell categories. It is important to note that the accuracies reported are not singular measurements but rather the mean of three independent experimental runs, ensuring the reliability of the results.

Implementation details. We implement all our techniques using PyTorch and conduct the experiments using 2 Tesla A100 GPUs. Following scGAD, the encoder has two layers with sizes 512 and 256, the decoder has the reverse structure of the encoder, and the latent space has a dimension of 128. The pseudo-labeling branch is implemented by two fully connected layers with size 128 and a number of the whole cell types. The contrastive-learning branch also consists of two linear layers sized 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 1.0, and the loss weight parameters λ1 and λ2 are both set to 1.0. The other hyperparameters ρ, γ, α and β are all set to 0.5. Lastly, the whole model is trained by 500 epochs for each dataset and the data distribution is re-estimated every 10 epochs.

4.2 Results

Intra-data experiments. We first explore the performance of scDET under the intra-data annotation setting without access to the batch effect. From the results in Table 1, scDET gives consistently superior performance than other methods under the overall accuracy on almost all datasets. Specially, compared with scGAD, our method achieves higher novel accuracy and better trade-off between classifying the seen cell types and clustering the novel cell types. It is not surprising that scDET gets such an excellent performance since the autonomously equilibrated dual-consultative contrastive learning framework can collaboratively tackle the scRNA-seq data imbalance issue and discover the out-of-distribution novel cell types. Besides, although scNym can obtain relatively high annotation accuracy on seen cell types, it has a sharp drop in clustering accuracy on novel cell types. MARS and scArches occasionally achieve competitive novel accuracy on some datasets, but they can only provide sub-optimal results for both classification and clustering accuracy. This evidence fully shows that the two-step strategy that first detects novel cells with "unassigned" labels and then clusters them is not an appropriate solution to this open-world task. 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 general, scDET outperforms other baselines and achieves remarkable progress in this setting.

Inter-data experiments. Then we turn to study a more challenging setting with the batch effect, i.e., inter-data annotation. From the results in Table 2, scDET still achieves better results than other algorithms on most mixed datasets, especially for novel accuracy and overall accuracy. For those groups where scGAD performs well, scDET can also perform well, and when the scGAD's performance is less than satisfactory on some datasets, scDET can show relatively stable and excellent performance. Moreover, compared with the intra-data setting, there is no significant decline in the performance of scDET, indicating that it could resist the effect of batch effect to some extent. In comparison, MARS and scArches are easily susceptible to batch effect and result in inferior precision on cell type identification, because they separate the learning process on reference data from that on target data. Similarly, scNym is easy to overfit to seen cell types and lacks strong discriminative power on novel cell types. Although the accuracy of scziDesk and scNAME seems competitive, that is because they do not use reference sets, thereby avoiding the batch effect. However, the high clustering accuracy on seen cell types does not facilitate the annotation process in practicality. In summary, scDET can work well under the challenging inter-data annotation scenario.

Feature visualization. To observe the cell type identification results more intuitively, we extract the low-dimensional embedding features of four methods and use the UMAP approach to visualize them in Figure 2. We can see that MARS confuses known cell types with novel cell types, and scNym fails to separate groups of novel cell types. Although scGAD performs better than them, it still does not recognize the known pancreatic D cells well and mixes some novel pancreatic ductal cells with pancreatic acinar cells. By contrast, our method does a good job of separating each known cell type from each novel cell type.

4.3 Ablation Study

Robustness analysis. We first investigate the effect of novel cell type numbers on the performance of each method and conduct control experiments on Quake 10x and Quake Smartseq2 datasets that hold massive cell types. From the results in Figure 3, we can see that the overall accuracy of almost all tested methods declines with the increase in the number of novel cell types. This is reasonable because the number of new cell types determines the difficulty of discovering them.

However, scDET always outperformed other methods regardless of the number of novel cell types, demonstrating the stability of our method. Then we vary the ratio of labeled data to study its impact on the results of five methods. Figure 3 shows the variation in overall accuracy with the changing of labeled ratio on Quake 10x and Quake Smart-seq2 datasets. We find that scDET still achieves consistently better results than the other baselines and maintains its superior performance without being affected by the ratio of labeled data.

Effectiveness of Lreg. Recall in Section 3.3, we propose to regularize the predictions of the pseudo-labeling branch by the estimated train set distribution. In Table 3, we show the performance of the pseudo-labeling branch on Quake 10x and Quake Smart-seq2 with different estimation strategies. "Oracle" denotes we use the true distribution π of Dr ∪ Dt (unknown in practice) as the target distribution in Equation 5, and it serves as an upper bound of the performance. Compared to using a balance prior, regularizing the predictions with oracle distribution significantly improves the performance on both seen and novel cell types, showing the importance of the distribution estimation. Meanwhile, the similar results achieved by our estimation strategy imply it could be a reliable proxy to π. Furthermore, we investigate whether two alternative estimation strategies could help the pseudo-labeling branch: 1) only regularize seen cell types prediction with πDr , 2) perform k-means clustering on the feature of the pseudo-labeling branch. Both of them result in inferior accuracy and could in turn deteriorate the contrastive learning process.

Effectiveness of Lsoft CL . In Section 3.5, we design a novel soft contrastive loss based on pseudo-labels to transfer the knowledge of the pseudo-labeling branch into the contrastive learning branch. As an opposite, we could also construct the loss in a hard manner where we formulate the positive pairs on top of the predictive cell type with the largest logit and further perform the supervised contrastive loss. Intuitively, the hard design discards the probability distribution information and is more susceptible to the false pseudo-labels, while the soft contrastive loss utilized in our method could help alleviate the influence of erroneous pseudo-labels. The results also support the intuition, as shown in Table 4, that the proposed Lsof t CL outperforms the supervised CL loss by a large margin. This phenomenon is also observed in knowledge distillation that transferring knowledge by using soft labels rather than one-hot predictions can achieve better performance.

5 Conclusion

In this paper, we formulate a realistic distributionindependent cell type identification task that unifies longtailed and open-set learning and design a novel framework scDET for this task to fight against scRNA-seq data imbalance and classify novel cell types simultaneously. Extensive experiments on various datasets verify the significance of scDET, and deeper analyses show the effectiveness of its proposed individual components.

Contribution Statement

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