{'fraction_non_alphanumeric': 0.02848198391750046, 'fraction_numerical': 0.008593702044073415, 'mean_word_length': 5.049758633494244, 'pattern_counts': {'":': 0, '<': 1, '<?xml version=': 0, '>': 1, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Uncertainty Quantification for Atlas-Level Cell Type Transfer

Single-cell reference atlases are large-scale, cell-level maps that capture cellular heterogeneity within an organ using single cell genomics. Given their size and cellular diversity, these atlases serve as high-quality training data for the transfer of cell type labels to new datasets. Such label transfer, however, must be robust to domain shifts in gene expression due to measurement technique, lab specifics and more general batch effects. This requires methods that provide uncertainty estimates on the cell type predictions to ensure correct interpretation. Here, for the first time, we introduce uncertainty quantification methods for cell type classification on single-cell reference atlases. We benchmark four model classes and show that currently used models lack calibration, robustness, and actionable uncertainty scores. Furthermore, we demonstrate how models that quantify uncertainty are better suited to detect unseen cell types in the setting of atlas-level cell type transfer.


Introduction
Single-cell genomics enables the characterization of cellular states at unprecedented scale and resolution. This has led to the generation of diverse single-cell RNA sequencing (scRNA-seq) datasets that capture cellular heterogeneity at the tissue and organ level. To combine multiple studies into a single representation of a human organ, consortia such as the Human Cell Atlas (Regev et al., 2017) are leading the generation of integrated scRNA-seq atlases. Such atlases are built by integrating multiple datasets across samples, individuals, and technologies, and are deemed to comprehensively capture cellular variation in the given tissue or organ. Recently, the first such comprehensive representation of the human lung was introduced in the Human Lung Cell * Equal contribution 1 Institute of Computational Biology, Helmholtz Center Munich, Germany 2 Department of Mathematics, TU Munich, Germany 3 TUM School of Life Sciences Weihenstephan, Technical University of Munich, Germany. Correspondence to: Fabian Theis <fabian.theis@helmholtz-muenchen.de>.
Next to their potential for discovery, integrated atlases can rapidly accelerate the analysis of new datasets in a harmonised way. Most notably, this is achieved by transferring cell type labels from the reference atlas to new, unannotated data from the same tissue. This task relies on cell type classifiers which are trained on the latent representation and expert cell labels of the integrated atlas. In order to analyze a new dataset, one "projects" the query data to the reference's latent space and applies the classifiers to infer cell type labels. The encoding of the query data is based on heavily regularized batch effect correction approaches to match the reference and is therefore not guaranteed to align perfectly (Lotfollahi et al., 2021). Note that the resulting misalignment naturally becomes greater when the projected data is comprised of samples from a different biological condition, e.g., disease or development, or contains unseen cell states.
The scenarios described above result in variable degrees of dataset "shifts", complicating the transfer of cell type labels. In addition, it is paramount to distinguish shifted latent representations due to unresolved technical batch effects from shifts due to biological differences between the datasets, e.g., new cell types or disease-related cell states. Such a distinction can be accomplished with a classifier that provides a measure of uncertainty, which yields actionable insights for downstream tasks.
In machine learning research data shifts, and how classifiers can overcome them, have been widely studied in uncertainty quantification (Gal, 2016;Gawlikowski et al., 2021). Crucially, model predictions should be reliable, i.e. calibrated, as they are used by data analysts to draw key interpretations from the data which influence further downstream analyses. Cell type classifiers should maintain these properties across dataset shifts while remaining sensitive to subtle biologically driven shifts in gene expression. Finally, they should be transparent to the analyst, by providing an accurate measure of confidence on the prediction task.
In this work, we evaluate uncertainty quantification methods for atlas-level cell type transfer. We find that models that quantify uncertainty are more robust, better calibrated, and provide high quality uncertainty measures that enable them to identify unseen cell types in the projected data.

Dataset and model choice
The HLCA dataset, on which all evaluations are performed, consists of a core of 14 datasets including 166 scRNA-seq healthy tissue samples of the human respiratory system from 107 individuals (Sikkema et al., 2022), see Figure A4 in appendix A for UMAPs. Batch effects between datasets were removed using scANVI with dataset as the batch covariate (Xu et al., 2021). All cells in the integrated atlas were re-annotated using iterative clustering, data-driven marker gene detection, and a consensus manual annotation from 6 lung experts, leading to 58 consensus cell type labels (Sikkema et al., 2022). The resulting training data consists of ∼580 000 cells where each cell is represented by a 30 dimensional embedding and has a unique cell type.
As we focus on the models' ability to provide uncertainty estimates, we consider two baseline models as well as two state-of-the-art classifiers: • Weighted K-nearest neighbor (WKNN) classifier which is used in the scArches method (Lotfollahi et al., 2021) and in the HLCA (Sikkema et al., 2022). • Multi-input multi-output classifier (MIMO) (Havasi et al., 2021), with 3 or 8 subnetworks.
Note that we evaluate the models' uncertainty only on the latent embedding and do not have access to the network used for data integration. This makes methods like scANVI (Xu et al., 2021) for uncertainty prediction inapplicable.
Uncertainty estimates can be divided into predictive, model, and data uncertainty (Gal, 2016;Gawlikowski et al., 2021). Data uncertainty entails label and measurement noise that is inherent to the data, while model uncertainty measures the uncertainty of the model's predictions due to parameter uncertainty. An application of model uncertainty is OOD detection, because it is higher for OOD cells than for ambiguous ID cells, which both have high predictive uncertainty. The uncertainties are calculated using: where H is the entropy and I the mutual information, (x, y) the embedding and cell type of the cell, D the training data and E p(ω|D) the expectation over the posterior (Smith & Gal, 2018). The uncertainties reported throughout this work have been scaled to [0, 1].
In the context of label transfer from reference datasets, data uncertainty would apply to transitioning cell states which cannot be associated clearly with a single cell type. In contrast, model uncertainty highlights how familiar the model is with cells similar to the provided query cell, that is, in an ideal scenario, the model knows when it does not know.
DKL and MIMO belong to the class of single forward pass uncertainty estimation models which we chose for their computational efficiency and performance. Other model classes that perform uncertainty estimation include: Bayesian Neural Networks (Blundell et al., 2015), Bayesian Dropout (Gal & Ghahramani, 2015), or Ensemble models (Lakshminarayanan et al., 2016), but they are not further considered here.
We performed hyperparameter optimisation for all model classes across ∼100 distinct parameter combinations. Furthermore, all reported metrics are computed from model fits across three test sets evaluated on three different trainvalidation splits, resulting in 9 unique train-validation-test splits.

Evaluation setup
As we are interested in both the evaluation of cell type transfer and the ability to provide meaningful uncertainty scores, the evaluation is divided in two parts.
In section 3.1, we evaluate all models in terms of their balanced accuracy, F1-score and expected calibration error (ECE) (Guo et al., 2017). The leave-out set is a dataset integrated in the HLCA but not used during classifier training. Next to the predictive performance, we analyse the calibration of the classifiers. Calibration considers the relation between the frequency of correct predictions and the binned prediction probability and, therefore, indicates how reliable a classifier is. A well calibrated classifier will be correct 7 out of 10 times when providing a confidence score of 0.7. For calibration evaluation we report both ECE and calibration curves.
In section 3.2, we test the models' ability to quantify uncertainty for several out-of-distribution (OOD) scenarios and their ability to discriminate between in-distribution training data (ID) and OOD data. To evaluate the models' uncertainty quality, we exclude three cell types during training and assess their performance through the predictive and model uncertainty. We evaluate OOD data discrimination using the Wasserstein distance (Krishnan & Tickoo, 2020) as well as the area under the precision-recall curve (AUPR).

Label transfer and calibration
We first evaluate the models on their ability to classify cell types on the integrated HLCA dataset using the most finegrained cell type annotation (58 cell types in total). In a typical use case a new, leave-out dataset introduces a data-shift due to differences in experimental design, sample handling, and cell type composition. We use a random traintest split in this analysis to serve as upper bounds on the performance.
We find that all models have comparable accuracy and F1scores ( Figure 1) on the test set. In contrast, the baseline models perform worse on the leave-out dataset, while DKL and MIMO are able to maintain their good performance. These results show that DKL and MIMO are superior in cell type classification to RF and WKNN.
The same is also reflected in the models' calibration, where the ECE for MIMO and DKL is consistently lower (better) compared to the baseline models. The calibration curves in Figure 2 provide a more detailed view on calibration by showing the models' calibration at different confidence levels. It can be seen that WKNN and RF are further away from perfect calibration in the test set than DKL and MIMO. Interestingly, the decline in performance on the leave-out dataset leads to an improved calibration of WKNN, making it comparable to models that quantify uncertainty for the leave-out scenario. Nevertheless, this calibration plot, together with the ECE, should be seen in light of model accuracy, since a model should be robust not only in terms of calibration, but also accuracy under data shifts. We report exact scores in Tables A1,A2 in the appendix A.
Overall, this analysis indicates that models that quantify uncertainty, such as DKL and MIMO, are better suited for the cell type transfer task. They achieve comparable performance in terms of accuracy for both test and leave-out set, while also showing superior calibration compared to the WKNN and RF models.

Uncertainty quantification on unseen cell types
To evaluate the models' ability to detect unseen cell types, we left out three cell types-B cells, Mast cells, and Ionocytes-from the training data and compute the predictive uncertainty for all classes seen during training, as well as the left out class. A classifier trained on K classes should result in high entropy of the output for the unseen, K + 1 class. Figure 3 reports predictive uncertainty for B cells for all model types (see Figure A1 in appendix A for additional cell types). DKL clearly identifies B cells as a new, unseen cell type which is shown through the high uncertainty score for the respective observations. In contrast, predictive uncertainty for WKNN is diffused and does not highlight B cells. Consequently, WKNN's uncertainty estimates can be identified as less reliable which may cause incorrect cell type transfer. We quantitatively evaluate the densities of the Figure 3. Predictive uncertainty on a subset of the HLCA UMAP with B cells left out during classifier training. It can be observed that the MIMO and DKL models predict higher uncertainty for the B cells. RF results are not reported due to poor classification accuracy. Table 1. Area under the precision-recall curve (AUPR) for classification between OOD (unseen cell type) and ID (validation set sampled from reference atlas), using different models and uncertainty types. DKL and MIMO outperform WKNN which is the baseline used in the HLCA. Softmax confidence uses the probability of the predicted class as a measure of uncertainty. predictive uncertainty for seen and unseen cell types via the Wasserstein distance, cf. (Krishnan & Tickoo, 2020), see Figure A2 in appendix A, from which the same conclusions can be drawn.
Finally, we formulate the task of detecting a new cell type as a binary classification problem and quantify the performance through the AUPR. The results are shown in Table 1. Each entry is the AUPR of a model when the given uncertainty type is used for classification. MIMO 8 performs best using model uncertainty for all left out cell types. DKL on the other hand has higher AUPR using predictive uncertainty. Since model uncertainty should be better suited for OOD detection, this suggests that MIMO 8 exhibits better disentanglement of model and data uncertainty. DKL's issue to highlight OOD cells using model uncertainty can also be observed in Figure A3 in appendix A. Nonetheless, DKL has higher AUPR than MIMO 8 for B cells and Ionocytes using predictive uncertainty. When using predictive uncertainty for DKL and model uncertainty for MIMO, both models outperform the WKNN baseline across cell types.
These results indicate that models that quantify uncertainty are superior in detecting unseen cell types, since they produce higher uncertainty on the left out classes. This is a desirable property of the models as it can be used by analysts to discover new cell states which are not present in the integrated scRNA-seq reference atlas.
This analysis requires further investigation and should be carried out on a real OOD scenario, i.e., where a new dataset is projected onto the reference and also excluded form the classifier's optimisation. Nonetheless, one has to appreciate that neither of the baseline models, RF and WKNN, provides such a measure on model uncertainty, and, therefore, an evaluation of models' prediction confidence is not possible. Put differently, it is not possible to state the model doesn't know when it doesn't know. In contrast, models like DKL and MIMO, that provide measures on both data and model uncertainty, should be preferred for building and extending integrated scRNA-seq reference atlases.

Conclusion
In this work, we presented uncertainty quantification on cell type transfer for single-cell reference atlases. We evaluated the considered models on a variety of metrics, including performance, calibration, and uncertainty quantification. We demonstrate that the baseline methods are not well calibrated and lack desirable robustness. In contrast, models that quantify uncertainty provide better confidence and more sensible uncertainty scores, which provide actionable results for downstream interpretations. In fact, such uncertainty measures can be used by analysts to uncover new cell types and cell states not present in the reference integrated atlas, as well as retain a notion on the "confidence" of the prediction, so that more careful conclusions can be drawn. In the future, we seek to extend this benchmark by including real-world query-to-reference mapping as well provide a more extensive evaluation on model and data uncertainty under dataset shifts.   Figure A2. Distribution of predictive uncertainty for the cell types the model was trained on (blue) and the left out cell type (orange). The computed Wasserstein distance, WD, (↑) between both distributions is shown in the title. This quantifies the models' ability to distinguish the OOD cell type from the ID training data. All models outperform the WKNN baseline. While RF shows the largest WD, using this model cannot be recommended due to poor performance in classification accuracy.