CellMemory: hierarchical interpretation of out-of-distribution cells using bottlenecked transformer

Abstract

Machine learning methods, especially Transformer architectures, have been widely employed in single-cell omics studies. However, interpretability and accurate representation of out-of-distribution (OOD) cells remains challenging. Inspired by the global workspace theory in cognitive neuroscience, we introduce CellMemory, a bottlenecked Transformer with improved generalizability designed for the hierarchical interpretation of OOD cells. Without pre-training, CellMemory outperforms existing single-cell foundation models and accurately deciphers spatial transcriptomics at high resolution. Leveraging its robust representations, we further elucidate malignant cells and their founder cells across patients, providing reliable characterizations of the cellular changes caused by the disease.

Supplementary Information

The online version contains supplementary material available at 10.1186/s13059-025-03638-y.

Background

In recent years, single-cell sequencing has revolutionized genomics by providing unprecedented insights into cellular heterogeneity. International consortia like Human Cell Atlas and Tabula Sapiens are utilizing population-scale single-cell resources to create a consensus reference [1, 2] reminiscent of the Human Genome Project [3] but at a cellular level [4]. However, cells from different individuals, technologies, or species do not exhibit similar distribution patterns. For example, there are variations in the states of malignant and healthy cells, as well as heterogeneity among malignant cells across different patients. Moreover, significant technical variations are present between single-cell and spatial omics data. These cells that deviate substantially from the established paradigms are defined as out-of-distribution (OOD) cells [5].

Deciphering of OOD cells requires concurrent identity inference and data integration, a process known as reference mapping [5]. Currently, only a limited number of methods can perform both [6–9] simultaneously. Furthermore, these methods often fall short of adequately considering gene interactions, which are crucial for characterizing cell representations. For biological researchers, the interpretation provided by explainable artificial intelligence (xAI) may be important as these explanations of model behavior could offer new biological perspectives [10]. But the application of xAI to cells outside the reference scope is an area that has not been thoroughly explored [10, 11].

Pre-trained models based on Transformer architectures, capable of understanding the interrelationships among features, present a novel strategy for exploring data across various fields [12–15]. However, current single-cell Foundation Models (scFMs) [16–18] struggle with handling long token inputs due to the computational complexity of self-attention mechanisms. This limits their capacity to comprehend the regulatory relationships among genes and to elucidate comprehensive cell representations. Moreover, the sparsity and high noise levels in single-cell data necessitate models with robust generalization to process such information effectively. In this study, we developed a neuroscience-inspired Transformer architecture to solve the mentioned problems.

The Global Workspace Theory (GWT) in neuroscience suggests that consciousness emerges from the selective dissemination of information via a shared “global workspace” [19–21] or “memory space”. This workspace, a network of densely interconnected neurons, facilitates competition between neuron modules to write information to this limited-capacity space [22]. Inspired by this concept, we propose that combining GWT with a Transformer architecture can efficiently manage informative and sparse single-cell data. We have developed CellMemory, which is a bottlenecked architecture within the Transformer that uses a cross-attention mechanism. This unique design allows CellMemory to learn generalized representations from standardized paradigms and perform interpretable inferences for OOD cells. By incorporating various specialist modules into the global workspace transparently, CellMemory allows for a hierarchical interpretation. Through the joint analysis of model parameters and biological information, we observed that the model’s internal decision-making rules align with some established biological patterns, enhancing our understanding of how the model interprets cell representations.

A comprehensive benchmarking study compared 3 scFMs and 16 task-specific methods to CellMemory utilizing data from over 15 million cells. CellMemory demonstrated harmonious integration and accurate label transfer, even outperforming scFMs [16, 17] regarding generalization ability and computational efficiency [23]. Remarkably, CellMemory provide interpretable inference for single-cell spatial transcriptomics at a granular level.

Finally, we used CellMemory and healthy references to delineate OOD malignant cells. CellMemory helps us contextualize malignant cells in harmonious embeddings and explain their origins. Our study focused on a transitional stage in the development of lung cancer, revealing that tumor cells in some lung cancer patients originated from different founder cells. This finding indicates that understanding the cell of origin of a tumor in some patients is crucial for comprehending drug resistance. It underscores the importance of leveraging xAI and existing references to better understand complex diseases by offering insights into OOD disease states.

Results

The inspiration for CellMemory

The Global Workspace Theory in neuroscience suggests that consciousness arises from the selective dissemination of information through a shared “global workspace” [19–21]. The brain consists of numerous specialized modules that handle various tasks like vision, memory, and emotions. The “global workspace” or “memory space” is a network of densely interconnected neurons that enables these specialized modules to share and communicate information. Due to its limited capacity, not all information can be retained in the global workspace. Specialized modules compete to write information to the global workspace, which then becomes part of conscious awareness. Once stored, the information is broadcasted to diverse specialized neural modules via synchronized neural firing and rapid signaling. This ensures that various brain regions process the information concurrently, contributing to a unified conscious experience.

Inspired by GWT and related works [12, 21, 22], we propose that an architecture comprising specialists trained to communicate effectively within a constraint global workspace can substantially enhance computational efficiency and improve generalization in understanding cell representations. To achieve this, we have developed CellMemory, which utilizes a bottlenecked architecture within a Transformer-based framework (Fig. 1a). CellMemory is tailored to derive cellular representations from single-cell datasets and to address the challenge of making inferences about OOD cells in an explainable manner. Furthermore, it can integrate population-scale single-cell datasets at a granular level (high-resolution annotation), adeptly managing variations in cell state, technical platform, population, omics, and species (Fig. 1b).

Fig. 1.

The illustration of CellMemory. a CellMemory learns cell representations from references in a supervised manner. The densely interconnected neurons in CellMemory enable specialized modules to share and communicate information. The “Specialists” (specialized modules) compete with each other to write information to the “Memory”. The information is then broadcasted to “Specialists” through synchronized neural firing and rapid signaling. The processes of writing and broadcasting between “Specialists” and “Memory” are executed by the cross-attention mechanism. The CLS token of CellMemory is used for training the model by comparing predictions with the real labels. b The application of CellMemory. CellMemory is capable of training models using single-cell data from any species without being constrained by a predefined feature set. Upon analysis of the query set with the trained CellMemory, the generated cell embeddings [CLS] enable population-scale integration (in the illustration, shapes represent cell types, and colors represent batch information), granule-level annotation (both shapes and colors represent cellular identity, and similar colors represent similarities in cell identities), and inference of out-of-distribution cells (both shapes and colors represent cellular identity). While performing these analyses with cell embeddings, CellMemory also generates the attention/memory score, which allows for the interpretation of model decisions at the level of specific genes and gene programs. TAG: Top Attention Gene

CellMemory employs a strategy of natural language processing (NLP) models for generating token and position embeddings [13]. The position embedding is employed to characterize features such as genes, while gene expression is processed in a bag-of-words manner and characterized by token embedding (Methods). The feature matrix is then transformed by the embedding layer into the “Specialists” module. Concurrently, the CLS (classification) token is incorporated into the “Specialists” module as the cell representation (Fig. 1a).

The conventional self-attention mechanism [12, 13] used by BERT has a quadratic complexity concerning sequence length, resulting in increased time and memory complexity [24] as sequences become larger. This presents a challenge when dealing with large-scale single-cell omics data. In contrast, CellMemory uses the “Specialists” module (length=$M$) to perform cross-attention with the “Memory” (length=$H$, $H\ll M$) (Fig. 1). In this process, a significant amount of biological information competes for the limited “memory space” (Additional file 1: Fig. S1a). The refined information is then broadcast to all modules by the “Memory” (Additional file 1: Fig. S1b) (Methods). The bottleneck space acts as a filter, prioritizing the most significant information and ensuring that the most important biological details are efficiently communicated. As a result, the CellMemory architecture has impressive generalization and substantially reduces computational costs (Additional file 1: Fig. S1c).

By supervised learning on reference data, CellMemory performs interpretable inference on query data. The bottlenecked architecture and its perception of cellular context ensure the generalized predictions, including cell embeddings and annotations. Furthermore, CellMemory’s unique architecture provides a hierarchical interpretation at two levels (Fig. 1b, Additional file 1: Fig. S1d).

Level 1

Each cell is assigned an attention score across all features, highlighting those that are most critical for the model’s understanding of that cell.

Level 2

For each cell, every memory slot within the model focuses on the subset of features, with the importance of these features represented by the memory score. This level reveals that different memory slots aggregate information according to different patterns (gene programs), thereby providing a more detailed characterization of model behavior (Methods).

CellMemory enables one-stop reference mapping of single-cell atlas across diverse scenarios

To validate whether CellMemory can generate accurate cell representations, we evaluated its annotation performance against 3 scFMs [16, 17, 25] and several advanced annotation methods [6, 7, 18, 26, 27] using over 4.6 million cells [27–35] with diverse biological and technological attributes (Additional file 1: Fig. S2a). We set up benchmarking scenarios. For example, the hLung data comprised cells from 5 sequencing platforms (10x, Smart-seq2, Singleron, etc.) spanning both diseased and normal human lung tissues; the mHypoMap integrated 17 published mouse hypothalamus datasets with significant heterogeneity; and the Immune dataset consists of 17 tissues, with the training set covering 9 tissues and the test set 8 (more details see Methods). Our design intentionally leveraged CellMemory’s bottlenecked Transformer architecture to facilitate robust label transfer. Specifically, the model was trained on cells from one condition and subsequently annotated OOD cells from distinct states (varying in platform, tissue, or species), with performance assessed by F1-score (macro) and accuracy. These metrics, particularly the F1-score’s nuanced evaluation of rare cell types, directly reflect the method’s capacity to learn compact yet informative representations. Remarkably, CellMemory outperforms scFMs [16, 17, 25] in various datasets, including identifying rare cell types on most benchmark datasets [36] (F1-score in Fig. 2a, Additional file 1: Figs. S2-S4). In the hPancreas dataset, the smallest dataset in our benchmark, where the query set contained a rare cell type (beta_minor) accounting for only 0.3% (Additional file 1: Fig. S2c). In this instance, Geneformer annotated only 11% of these cells, and Seurat failed to annotate any. In contrast, CellMemory achieved an annotation accuracy of 81% for this rare cell type. It is noteworthy that although scFMs demonstrate improved overall accuracy, especially on small-scale datasets, they tend to neglect low-abundance cell types, likely due to data imbalance during the pre-training stage or insufficient model generalization (Fig. 2a, Additional file 1: Fig. S2). Despite not requiring pre-training, CellMemory still achieves an outstanding F1-score and scalability, demonstrating its excellent generalization. Furthermore, with its smaller model size and bottlenecked architecture, CellMemory attains higher computational efficiency than self-attention-based Transformers (Fig. 2a, b, Additional file 1: Figs. S2-S4).

Fig. 2.

Annotation and integration benchmark. a F1-score (macro) is employed to assess the annotation performance of single-cell datasets. Each circle represents the mean of the five-fold cross-validation for the respective method. The overall performance across all datasets corresponds to the Annotation Score. In the term, “h” represents human, and “m” represents mouse. hBreast: Prediction of human breast cancer samples using normal breast samples. hLung: Prediction of human LUSC samples using normal lung samples. hPancreas: Prediction of human Type 1 diabetes samples using pancreatic islets normal control samples. Immune: Immune cells from human tissue group 1 were used as the training set, and immune cells from human tissue group 2 were used as the test set. mHypoMap: Prediction of mouse hypothalamus samples produced by the Drop-seq platform, using samples produced by the 10x platform. mPancreas: Prediction of mouse type 1/2 diabetes samples using normal pancreatic islet samples. crossSpecies (cortex): The human cortex samples were used as the training set, while one human and four non-human primate cortex samples were used as the test set. CellMemory: The model was trained using all input genes, considering genes with zero expression in each cell. CellMemory-fast: Genes with zero expression in each cell were filtered out, and the model was trained using the Mixed Precision Training strategy. b The average time taken to train one epoch is recorded for each dataset for CellMemory, scGPT, and Geneformer. We did not include UCE’s computation time because most of its time is used for encoding the embedding matrix. c The integration benchmark for the dataset includes popular integration methods such scPoli, scVI, and scGPT. The Integration Score is comprehensively evaluated based on the retention of biological information and the removal of batch effects. The overall score is calculated as a weighted average of the batch correction score and the biological conservation score, with weights of 0.4 and 0.6, respectively. d, e Annotation and integration benchmarks for the SEA-AD and AIDA datasets. f UMAPs are generated from SEA-AD cell embeddings obtained by four integration methods. Even the UMAP structure generated by scPoli or scFM does not appear particularly clear. Cell colored by original annotation. A total of 131 cell types were labeled, for detailed annotation, see Additional file 1: Fig. S8

We further evaluated CellMemory’s ability to integrate OOD cells during reference mapping using cell embeddings from test sets (The same CLS embedding used as in the annotation evaluation, see Methods). In comparison with advanced integration tools [37] such as scVI [38], scGPT [16], and scPoli [6], which is tailored for the integration of population-scale datasets. CellMemory consistently preserved biological representations while minimizing batch effects (Fig. 2c, Additional file 1: Fig. S5). We attribute these performances to CellMemory’s focus on learning cell representations combined with the generalization conferred by its bottlenecked Transformer architecture, which enables effective integration. Our experiments also demonstrated that specific token preprocessing strategies, such as scaling expression values into defined bins (Additional file 1: Fig. S6), contribute to enhanced integration performance by reducing batch effects. This suggests that a scale-to-bin strategy is more effective for batch correction when integrating data with significant batch effects. Furthermore, incorporating additional tokens may help the model generate a more comprehensive understanding of cell representations, emphasizing that higher computational efficiency will enable researchers to consider more biological information (Additional file 1: Fig. S6).

As sequencing technologies advance and generate larger datasets, researchers increasingly rely on iterative data integration and annotation strategies [33, 34] (beginning with distinguishing cell types at a broad level and then progressing to finer distinctions between subtypes or even supertypes within each major category). In this investigation, we demonstrated CellMemory’s one-stop reference mapping capability on two population-scale datasets characterized by high-resolution cell identities (Fig. 2d). The first atlas—the brain cortex of Alzheimer’s patients (SEA-AD), comprises 1,395,601 cells from 84 donors with 131 categories of granule-level annotation [33], where normal samples served as the reference for disease sample inference. The second atlas, the Asian Immune Diversity Atlas (AIDA) [39], encompasses 1,058,909 peripheral blood mononuclear cells from 503 healthy donors across five Asian ancestries and a European population. We trained the model using cells from the Korean population and subsequently integrated data from the remaining five populations (Methods). In the annotation benchmark, CellMemory demonstrates superior performance in both accuracy and F1-score (Fig. 2d). Specifically, for the SEA-AD dataset containing 131 cell types, CellMemory achieves an accuracy of nearly 90%, compared to only 60% for scPoli. Integration results further confirmed that CellMemory not only removes batch effects but also preserves comprehensive biological information, ranking first in 7 out of 8 metrics (Fig. 2e, Additional file 1: Figs. S7-S8). Compared to advanced methods such as scGPT and scPoli, the cell embeddings generated by CellMemory exhibited clear consistency with high-resolution cell states (Fig. 2f). These findings illustrate that the enhanced understanding and generalization achieved through its bottlenecked Transformer architecture underpin CellMemory’s superior data comprehension and high-resolution knowledge transfer in population-scale datasets.

CellMemory facilitates the interpretable characterization of single-cell spatial omics

Deciphering single-cell spatial omics remains challenging due to significant technical bias and variations of information richness compared to single-cell omics data. To characterize these OOD cells, we applied CellMemory to single-cell spatial transcriptomics data (Fig. 3a). A benchmarking was performed to assess the efficacy of spatial data cell type annotation using data from CosMx [34, 40], MERFISH [35, 41], and Slide-seq [42]. These spatial technology platforms encompass both high-throughput sequencing and in situ hybridization approaches, comprising a total of 645k cells with each dataset containing three replicate samples. CellMemory was compared with 5 state-of-the-art spatial data annotation methods [43–47] and 2 top-performance scFMs (Methods), where single-cell data served as the reference for model training and spatial transcriptomics data were used as the query set for prediction. Accuracy was employed as the primary metric for evaluation due to the complexity of spatial data representation. Our results demonstrate that CellMemory achieves higher accuracy across datasets generated by all three sequencing platforms, thereby exhibiting robust generalization on OOD cells, even without incorporating spatial coordinate information (Fig. 3b, Additional file 1: Fig. S9). Additionally, the performance of scGPT and Geneformer indicates that pre-training strategies may further enhance the comprehension of spatial data across species. Next, we used CellMemory to explore its reliability and interpretability on data produced by additional spatial technology platforms.

Fig. 3.

CellMemory interprets single-cell spatial data across omics. a CellMemory can characterize single-cell spatial omics from various sequencing platforms, including CosMx, MERFISH, Slide-seq, Stereo-seq, Xenium, and Slide-tags. b The accuracy of spatial annotation tools is evaluated using datasets from mHypo (mouse hypothalamus measured by MERFISH), hNSCLC (human NSCLC measured by CosMx), and msSermato (mouse spermatogenesis measured by Slide-seq). Each dataset comprised three samples for replication. c CellMemory was trained using single-cell data at the L2 cell type resolution, to generate CLS embeddings and annotations for Slide-tags cells. The right part is plotted by the L1 (original) and L2 (CellMemory) cell types in spatial coordinates. d CellMemory model was built using single-cell data at the L3 resolution, to integrate single-cell and Slide-tags data. e The cells highlighted in the co-embedding are Slide-tags cells, labeled with Slide-tags cell, L1 (original), L2 (CellMemory), and L3 (CellMemory) cell identity. f The identification of Slide-tags cells at L3 resolution by CellMemory (L4 IT_2 and Micro-PVM_1) is displayed (the first line), along with the expression (the second line) and attention score (the third line) of TAGs (VWC2L, F13A1). The left half represents the spatial coordinate, and the right half is the UMAP coordinate. L4 IT_2 (VWC2L: TAGs top 2); Micro-PVM_1 (F13A1: TAGs top 1). g UMAP representation of 4 million mouse whole brain cells from MERFISH, colored by subclasses. h Annotation benchmark comparison of CellMemory with other state-of-art methods, including scGPT, Geneformer, and CellTypist. The reference dataset consists of mouse whole-brain 10x single-cell data, the query set comprises MERFISH cells derived from 59 coronal sections. i Visualization of the 5 (total 59) mouse whole-brain sections, with cells colored by predictions of CellMemory. Cell type labels are provided in Additional file 1: Fig. S17. j Spatial coordinates of mouse brain MERFISH section 36, colored by CellMemory predictions. k Heatmap displaying the attention scores of TAGs for all subclasses of IT-ET Glut in section 36. l Distribution of subclasses and their corresponding TAGs’ attention scores within the spatial coordinates of section 36. PF Fzd5 Glut (Ctxn3: TAGs top 2); L2/3 IT RSP Glut (Igfbp6: top 21); L2/3 IT CTX Glut (Chrm2: top 4); ARH-PVp Tbx3 (Tbx3: top 1); L4 RSP-ACA Glut (Cbln1: top 1); L6 IT CTX Glut (Syt6: top 2); CA1-ProS Glut (Fibcd1: top 1); CA3 Glut (Frzb: top 2); L6 IT CTX Glut (Sema3e: top 5)

Ductal carcinoma in situ (DCIS) is a non-obligate precursor of invasive ductal carcinoma [48]. Characterization of key markers at invasive boundaries is crucial for understanding the molecular features associated with cancer heterogeneity. The inherent complexities and similarities among cells within DCIS regions pose challenges in achieving accurate interpretation. To overcome this, we utilized CellMemory to model a single-cell breast cancer dataset [49] and integrate subcellular resolution breast cancer data from 10x Xenium [48] (Additional file 1: Fig. S10). The visualization of the attention score demonstrated CellMemory’s interpretable representation of cell groups (Additional file 1: Fig. S10c, e). In a particular section of breast cancer tissue, CellMemory identifies cancer-associated epithelial cells (Epi-Cancer) that are different from the inner layer of ductal cells. Notably, two ductal regions show distinct cellular compositions (Additional file 1: Fig. S10d). The first region (DCIS 1) contains a greater proportion of normal epithelial cells (Epi-Normal), while the second region (DCIS 2) has significantly fewer of these cells. The spatial distribution of Top Attention Genes (TAGs) like KRT14 [50] and SERHL2, which are used to identify Epi-Normal and Epi-Cancer cells respectively, exhibit variations (Additional file 1: Fig. S10d). Previous reports have indicated the presence of KRT15⁺ myoepithelial cells around DCIS, which are notably absent in highly invasive tumor regions. The absence of these cells is linked to a poor prognosis [51]. When CellMemory assigns weights to KRT15 in spatial cells, a significant reduction of KRT15⁺ epithelial cells in DCIS 2 was observed (Additional file 1: Fig. S10f-g), indicating a higher degree of tumor invasion in this region. The confirmation of previous identifications of invasive tumor cells validated the insights provided by CellMemory. By providing spatial cell characterization and identifying critical TAGs associated with invasion, CellMemory contributes to an interpretable inference of cancer biology, enabling the dissection of tumor ecosystems in spatial cancer data.

CellMemory deciphers single-cell spatial data at the granule level

Characterizing single-cell spatial transcriptomics data at the granule level remains underexplored, especially for OOD cells. We used human cortex data [52] generated by Slide-tags as an example. Initially, CellMemory utilized a single-cell cortex dataset [33] from normal samples to build a model at level 2 (L2, 24 cell types). As illustrated in Fig. 3c, CellMemory produced robust embeddings and advanced annotations for the Slide-tags data, showcasing a harmonious integration of two distinct omics datasets (Additional file 1: Fig. S11a). Moreover, the attention score of TAGs offered interpretable insights into the decision-making process for integrating spatial omics (Additional file 1: Fig. S11b). Building on this foundation, we constructed a level 3 model that incorporated granule-level cell information (131 cell types) by integrating single-cell and Slide-tags data through CLS embedding (Fig. 3d, Additional file 1: Fig. S12).

To further assess integration performance and prediction reliability, we incorporated Slide-tags information, and L1/2/3 annotations (Fig. 3e). The consistent cell representations and coherent hierarchical annotations demonstrate that CellMemory can achieve granule-level characterization of single-cell spatial data. The effectiveness of the annotation is visually supported in the Sankey diagram, which depicts a transition from coarse to granular categories (Additional file 1: Fig. S13). Moreover, the consistency of marker gene expression between Slide-tags data and the 10x reference further validates the reliability of CellMemory’s predictions (Additional file 1: Fig. S14). To validate the model’s interpretation for specific cell states, we first examined the TAGs (top 1–12) in the L4 IT_2 and observed that the attention scores of these genes correlated with their expression levels, with the scores resembling a denoised version of the gene expression. This consistent trend and denoising-like behavior suggest that CellMemory leverages these features to achieve a comprehensive understanding of cell representations. Subsequently, we noted that the well-established expression of LHFPL3 in oligodendrocyte precursor cells (OPC) was emphasized [53], and CellMemory appropriately assigned a greater weight to LHFPL3 in OPC_2. While VWC2L is expressed across several clusters, CellMemory uniquely assigns it higher weights in L4 IT_2 cells after considering all gene correlations. Remarkably, CellMemory accurately characterizes rare cell states, such as Micro-PVM_1 (F13A1⁺), which exhibit a very low proportion in the reference set (only 0.43‰) and in the spatial data (3.4‰), providing reasonable explanations for cell representation (Fig. 3f, Additional file 1: Figs. S14-S15). In addition, CellMemory’s performance has been validated using Stereo-seq spatial data from mouse hemi-brain [54, 55], achieving biologically meaningful interpretations (Additional file 1: Fig. S16).

The mammalian whole brain is incredibly complex. Understanding the cellular distribution within its spatial coordinates is crucial for comprehending consciousness, memory, and disease. To address this challenge, CellMemory was trained using 781 scRNA-seq libraries, encompassing 4 million single-cell transcriptomes from the mouse brain, to analyze the Allen Institute for Brain Science (AIBS) MERFISH dataset [56]. This dataset includes 59 serial full coronal sections at 200-µm intervals spanning the entire mouse brain, totaling around 4 million cells (Fig. 3g–i, Additional file 1: Fig. S17). In challenging a scenario involving up to 338 subclasses in cross-omics representation, CellMemory outperformed scFMs in generalization (Fig. 3h). Compared to a well-known annotation method, CellMemory demonstrated up to 15% improved representation accuracy. More specifically, among these 59 sections, CellMemory exhibited lower annotation accuracy than scGPT in only 3 sections, whereas in the remaining 56 sections its performance was superior. This implies that CellMemory accomplishes high-resolution knowledge transfer in a single step, unlike previous methods that relied on hierarchical annotation processes. As an example, we present CellMemory’s interpretation of the IT-ET Glut neuronal subclasses in section 36 (Fig. 3j, k). The attention scores of TAGs displayed spatial distributions highly convergent with cell types, partially elucidating the logic behind CellMemory’s representation of OOD cells from MERFISH platform (Fig. 3l).

Collectively, our validation with Slide-tags, Stereo-seq, and MERFISH data demonstrated that CellMemory delivers high resolution in the characterization of single-cell spatial omics. This is essential for dissecting spatial cellular heterogeneity in complex biological systems and will significantly enhance our perception of spatial cell states.

Integrating the Asian Immune Diversity Atlas (AIDA) with hierarchical interpretation

Integrating large-scale single-cell data across diverse populations and technological platforms presents a challenge, often resulting in the loss of low-abundance cells or accurate information due to batch effects between datasets and insufficient model generalization. To validate the value of CellMemory’s results for analysis, we used PBMC CITE-seq data [7] to construct a model that integrates AIDA data from 10x 5’ sequencing. The AIDA dataset includes 503 healthy donors from five Asian ancestries across Japan, Singapore, and South Korea, as well as European controls (Fig. 4a, b). The cell embeddings generated by CellMemory demonstrated harmonious integration and high consistency with the original manually curated annotations (Additional file 1: Fig. S18). Notably, CellMemory identified a rare cell type called ASDC, which was not previously recognized in the AIDA dataset and constitutes only about 0.1‰ of the total cell population (Fig. 4b). The distinct expression profiles of AXL and PPP1R14A confirmed the characterization [57] (Additional file 1: Fig. S19).

Fig. 4.

Integration of the Asian Immune Diversity Atlas with hierarchical interpretation. a The Asian Immune Diversity Atlas (AIDA) comprises 985,470 cells from 503 healthy donors. The samples are obtained from Korea, Japan, and Singapore, including five Asian ethnicities, and five European donors as a control group. b The integration of CellMemory with the AIDA dataset generated cell embeddings, which are colored by cell type annotations (left) from CellMemory and ethnicity (right) (randomly selecting 30 donors from each Asian population for fair observation of cell composition and state). c Embeddings of CD8 TCM cells, annotated with CellMemory and original annotations. d Attention scores of TAGs (CD8A) for CD8 TCM cells are displayed. e The expression of IL7R, CD4, and CD8A/B in CD4 TCM, CD8 TCM, and the MissMatch group. f Memory scores for the top five genes of each memory slot in the CD8 TCM group. Notably, the information attended to by different slots is not entirely distinct. While each slot tends to process information associated with some patterns, certain features remain critical across multiple slots. g Enrichment terms for the top 1–100 TAGs of CD8 TCM. h For the CD8 TCM, the most significant terms in each memory slot (enrichment results for the top 1–50 gene sets in each memory slot)

Notably, CellMemory’s prediction revealed a subgroup within CD8 TCM cells that had originally been mislabeled as CD4+_T_cm cells (Fig. 4c). Upon examining the attention scores of TAGs (CD8A) (Fig. 4d) and marker expression for both CD8 and CD4 T cells, we noticed that this mismatch subgroup closely resembled the CD8 TCM cell phenotype (Fig. 4e). This observation underscores the complexity involved in characterizing immune cells in large-scale, multi-population studies and highlights the pivotal role of model’s generalization.

Then, we investigated the decision-making process of CellMemory, particularly focusing on the misclassified CD8 TCM cells. Specifically, we examined how each memory slot processes information within the memory space (Methods). Each slot tends to assign different weights to specific features, as illuminated by the memory scores of TAGs across all slots (Fig. 4f). And the most significant enrichment terms within each slot generally correspond to a subset of the enrichment results derived from cell type-specific TAGs (Fig. 4g, h, Additional file 1: Figs. S20-S21). For instance, slot 8 prioritizes the term “Regulation of lymphocyte activation,” while slot 6 focuses on “Adaptive immune response”—both of which are prominent terms in the CD8 TCM TAGs enrichment results. This discovery indicates that within CellMemory, multiple specialist modules requiring coordinated management are organized into different memory slots based on identifiable patterns, such as biological functions (Additional file 1: Fig. S1d). Consequently, each memory slot becomes specialized in comprehending and processing an aspect of information, facilitating efficient collaborative processing. This mechanism is analogous to the cognitive process of driving, wherein the driver must synthesize and decipher a multitude of informational cues, such as road markings, traffic signs, and auditory cues, to navigate effectively. The bottlenecked architecture of CellMemory substantially enhances our understanding of the model’s dynamic approach to processing biological information, imparting an additional dimension of interpretability to its role in orchestrating complex biological processes.

Studying malignant cells and the heterogeneity among patients using normal cells as a baseline

While mapping datasets from healthy individuals to healthy reference atlases is feasible, interpreting diseased samples using these atlases, although highly desirable, presents distinct challenges [58]. Malignant cells exhibit notable complexity, characterized by aneuploidy and substantial heterogeneity among patients, which compounds the intrinsic challenges of their analysis. With CellMemory, we address this issue by leveraging its robust generalization capability, to characterize malignant cell states against an established healthy reference.

Mixed phenotype acute leukemia (MPAL) is a high-risk subtype of leukemia characterized by the co-expression of both myeloid and lymphoid features. Recent reports have indicated that MPAL is associated with abnormalities in primitive hematopoietic progenitors [59]. Accurate characterization of the malignant cell components and molecular signatures in MPAL patients is crucial for elucidating the potential mechanisms underlying this complex disease. In this study, we trained CellMemory using bone marrow and peripheral blood mononuclear cells (BMMC, PBMC) from healthy individuals. Subsequently, we applied this model to analyze malignant cells from two MPAL patients [60], one with B myeloid MPAL (MPAL-a), and the other with T-myeloid MPAL (MPAL-b) (Fig. 5a, b). CellMemory was utilized to define the lineage proportions in these patients (Methods), identifying the progenitor-like state as the predominant cellular origin (Fig. 5c).

Fig. 5.

CellMemory interprets the individual heterogeneity of complex diseases. a CellMemory integrated malignant blood cells from MPAL patients with bone marrow and peripheral blood mononuclear cells from healthy samples. b Colored points show the positions of MPAL patient cells (MPAL-a: 3435 cells; MPAL-b: 2947 cells) in the co-embedding space and their inferred cell identities. c Classify the lineage proportions based on cell types. d The expression of lineage-specific genes in lymphoid and myeloid cells is inferred from MPAL-a patient cells. e The heatmap showing the memory score of progenitor-like cells from MPAL-a patient in memory space. Eight memory slots are displayed, each showing the top five TAGs that it focuses on. The top 1–50 TAGs for each slot were subjected to enrichment analysis, with significant terms displayed above the heatmap. f Expression and attention score of EGR1 in healthy cells and MPAL-a cells. g Cell embedding generated by CellMemory for MB patients, colored by the inferred potential origin types of MB cells. h The centroids of RL, GCP, and eCN/UBC cell populations are compared to centroids of cells from three subgroups of MB patients (each subgroup displaying three patients). i UMAP shows the expression and attention scores of MKI67 and LMX1A. j Heatmap displaying the expression of cell type-specific genes in fetal cerebellum cells and malignant cells from MB patients

After examining the specific features of lymphoid and myeloid cells in MPAL-a, we consistently observed characteristics related to the normal lineage and developmental stage of these malignant cells (Fig. 5d). Information at both the gene and gene program levels provides biologically meaningful interpretations of these malignant cells with progenitor-like properties. The GO term “Hematopoietic cell lineage,” which is enriched among the top attention genes in memory slot 3, suggests the cellular origin, whereas the term “Response to oxidative stress” in slot 4 reflects the cells’ stress state [61]. Additionally, the term in slot 6 indicates immune function dysregulation [62] (Fig. 5e, Additional file 1: Fig. S22a). EGR1, a key feature of the progenitor lineage [63], was identified as a TAG by CellMemory, suggesting that it may play an important role in maintaining the stem cell characteristics (Fig. 5f). Furthermore, CellMemory provided distinct hierarchical interpretations for MPAL-b, with increased attention to DNTT, implying potential differences in lineage or stages for the cell of origin among patients [64] (Additional file 1: Fig. S22b-e).

Medulloblastoma (MB) is a malignant pediatric cerebellar tumor that can be molecularly divided into distinct subgroups [65]: WNT, SHH, Group 3, and Group 4. Accurately identifying the origins of each subgroup is crucial for a better understanding of the disease and for developing targeted interventions. While it is currently believed that WNT MB originates from mossy fiber neurons in the embryonic dorsal brainstem [66], the origins of Group 3 and Group 4 MB are still unclear. A recent report suggested a potential association between the development of embryonic cerebellar glutamatergic neuronal lineages and MB [67]. To explore this, we employed single-cell data from SPLIT-seq [68], which was generated from developing human embryonic cerebellar glutamatergic neuronal lineages, to construct a model. This model was then used to characterize individuals from various MB subgroups (SHH, Group 3, Group 4) based on data obtained through the Smart-seq2 [69] method, aiming to uncover the diverse origins of MB.

These findings indicate that cell types in MB can be confidently categorized into three lineages: eCN/UBC (excitatory cerebellar interneuron/unipolar brush cell), RL (rhombic lip), and GCP (granule cell progenitor) (Fig. 5g). The centroids of these cell identities (Methods) demonstrated high consistency with the centroids of MB subgroup donor cells (Fig. 5h, Additional file 1: Fig. S23). The representation suggested that Group 3 MB may originate from RL, Group 4 from eCN/UBC, and SHH from GCP. Additionally, the distribution of MKI67 attention scores in cell embeddings shows clear similarities between MB cells and RL representations during embryonic cerebellar development [67], as seen with LMX1A (Fig. 5i). Finally, the expression of markers for RL, GCP, and UBC in MB cells, along with their attention scores, supports the proposed origins of different MB subgroups, in agreement with previous reports [67] (Fig. 5j, Additional file 1: Fig. S23).

In conclusion, precise characterization of OOD cells using CellMemory reveals the heterogeneous origins of complex tumors. The clear and interpretable representations generated provide a robust framework for elucidating the relationship between normal and disease states, underscoring the potential of CellMemory as an invaluable tool for identifying tumor founder cells.

Investigating the intermediate stages in the development of lung cancer

Cancer arises from disruptions in normal cellular function and genetic alterations that enable tumor cell proliferation [70]. However, accurately characterizing malignant cells remains challenging for computational methods due to the similarities between tumor cells and their adjacent normal cells, as well as the considerable heterogeneity among tumors from different individuals. In this study, we demonstrate the effectiveness of CellMemory by constructing a model based on a lung cancer atlas [34] (which includes tumor cells), and subsequently integrating another dataset consisting of 293,432 cells from 52 lung cancer patients [71]. The query dataset covers lung adenocarcinoma (LUAD) and lung squamous cell carcinoma (LUSC), with a focus on LUAD results.

Accurately characterizing tumor cells and other epithelial cells is an error-intolerant task. Therefore, we first compared CellMemory with scPoli, a method capable of reference mapping, on the validation set. CellMemory demonstrated robust performance (Fig. 6a) (Methods). The trained CellMemory was applied to integrate data from LUAD patients, resulting in the annotation of 24 cell types, including tumor cells (Fig. 6b, Additional file 1: Fig. S24a, b). The cell embedding revealed that transitional cells (SCGB3A2) within the epithelial cell population are most closely related to tumor cells. Moreover, a subset of AT2 cells was embedded within transitional cells (Fig. 6b), suggesting a similarity between AT2 and transitional cells. These findings support the reports that LUAD may exploit lung plasticity and regeneration mechanisms to generate tumors [71–73]. We named the AT2 cells embedded within transitional cells as AT2-2 (Methods), while those positioned further away were designated as AT2-1 (Fig. 6c, Additional file 1: Fig. S24c). Analysis of intercellular interactions [74] indicates that tumor cells engage more significant interactions with cells like macrophages compared to transitional cells. Furthermore, the copy number variation (CNV) results [75] and hierarchical interpretation of tumor cells emphasize their malignant nature (Additional file 1: Figs. S24-S25), implying that the tumor cells identified by CellMemory display distinct characteristics associated with malignancy [76].

Fig. 6.

CellMemory accurately characterizes the lung cancer atlas. a The accuracy of malignant cell identification by CellMemory and scPoli was validated through five-fold cross-validation. Filter: filter zero expression genes in cells. Full: consider all input genes. b UMAP generated by CellMemory, integrating and annotating 149,595 cells from LUAD donors. c UMAP of all epithelial cells in the LUAD dataset. d Interactions between tumor-associated epithelial cells and monocytes, macrophages, stromal cells, etc. e Dop plot demonstrating the specific gene expression in epithelial cells. f UMAP displaying the expression of SFTPC, SFTPB, SCGB1A1, and SCGB3A2 in epithelial cells. g Proportion of stages identified as AT2-2 cells. h, i UMAP displays epithelial cells from PA03, with color indicating cell type (h) and sample tissue (i). j CNV levels of tumor-associated epithelial cells in PA03. k CNV of tumor-associated epithelial cells in PA03 inferred using inferCNV. The CNV distances between each cell type were generated by inferCNV, which corresponds to the alignment of rows. Tumor cells and Transitional cells are down-sampled to 400 cells

Compared to AT2-1 cells, AT2-2 cells show closer associations with tumor cells (Fig. 6d). The gene expression profiles of AT2-2 cells were found to be intermediate between those of AT2-1 and transitional cells. Specifically, AT2-2 cells lose characteristics associated with AT2-1 (SFTPC, SFTPA1/2) while acquiring features typically associated with transitional cells (SCGB3A2, SPINK) (Fig. 6e, f, Additional file 1: Fig. S24d-f). Examination of the memory space revealed that, in memory slot 2, the gene SFTPB and the term “Surfactant metabolism” reflected the identity of AT2 cells, while in slots 1, 7, and 8, the genes and terms indicate that AT2-2 cells are acquiring tumor-related cell states (Additional file 1: Fig. S25c, d). Furthermore, the interactions, stages, and CNV patterns of these cells suggest that AT2-2 cells represent an intermediate state in the progression of LUAD from AT2 cells (Fig. 6d, g, Additional file 1: Fig. S24g-i). Genes specifically expressed in AT2-2, including MSLN, CAPN8, NAMPT, and TC2N, have been identified as potential therapeutic targets, with the knockout of these genes significantly suppressing LUAD [77–80]. HLA-E, known in pancreatic cancer to be exploited by circulating tumor cells (CTCs) to evade immune checkpoints and promote metastasis [81], is closely related to the most focused GO term “Antigen processing and presentation of exogenous peptide antigen” in memory slot 1 of AT2-2 (Additional file 1: Fig. S25d). These findings imply that the AT2-2 state represents a notable intermediate stage in LUAD progression, potentially offering new insights for therapeutic intervention.

CellMemory generates cell embeddings that comprehensively incorporate all input gene interactions without being constrained by prior information. This effectively characterizes the cell states for interpretable integration and offers a robust method for exploring the continuum of cell states within tumor contexts.

Dissecting the heterogeneity of founder cells among lung cancer patients

Single-cell transcriptomics enables a detailed understanding of patient-specific disease states. In the LUAD dataset, the cell embeddings generated by CellMemory revealed the similarity between AT2 cells and tumor cells. For example, in patient PA11, the cell representations and CNV patterns of AT2 cells closely resembled those of tumor cells (Additional file 1: Fig. S26), consistent with the common view of founder cells in LUAD. Additionally, the cell embeddings indicated that club cells and their neighboring transitional cells follow a trajectory that converges toward tumor cells. In a subset of patients, tumor tissues displayed CNV patterns in club cells similar to those observed in tumor cells (Additional file 1: Fig. S24g). We provide a personalized characterization for patient PA03 (Fig. 6h, i, Additional file 1: Fig. S27). The embedding of epithelial cells in PA03 revealed a distinct cluster that had transitioned from club cells to transitional cells. Previous reports have indicated that bronchial cells can use intermediate state cells with differentiation potential to replenish alveolar cells [82]. Accordingly, we have designated this cluster as ClubtoTrans cells (SCGB1A1, SCGB3A2, SFTPB) (Methods).

In the tumor tissue of patient PA03, AT2-2 cells exhibited CNV levels and patterns akin to those observed in tumor cells (Fig. 6j, k). Interestingly, a small number of club and ClubtoTrans cells were identified within the tumor tissue of the PA03 patient. The CNV patterns observed in tumor cells were also evident within ClubtoTrans cells. Based on clustering results from InferCNV, the CNV profile of ClubtoTrans cells is more closely aligned with that of tumor cells (Fig. 6k). Moreover, in key oncogenic regions of LUAD (e.g., KRT8, CEACAM6), ClubtoTrans cells exhibited significant CNV amplification (Fig. 6k). Recent reports have pointed out that club cells can develop tumors by leveraging the plasticity and regenerative capacity of lung [72]. During the transition, club cells acquire an AT2-like epigenetic pattern, consistent with CellMemory’s characterization of the single-cell transcriptome (Fig. 6h). Importantly, the heterogeneity of founder cells among LUAD patients is closely associated with tumor state transdifferentiation [83]. In the club and ClubToTrans cells of patient PA03, expression of LUSC markers was observed (Additional file 1: Figs. S27e-S28). This suggests that the patient may be in the early stages of transitioning to lung adenosquamous carcinoma (LUAS) [83]. However, the potential impact of this transition on drug resistance and prognosis was not recognized in the original study. The integration results generated by CellMemory revealed similar state changes in another cohort of LUAD data for certain donors (LUAD dataset 2, Methods) (Additional file 1: Figs. S29-S31).

It can be challenging to focus on the interrelationships among subpopulations when integrating population-scale data, due to confounding factors such as background noise and batch effects. However, CellMemory’s analysis of lung cancer data shows that it can effectively transfer important knowledge without bias (Additional file 1: Figs. S32-S34). Identifying potential founder cells specific to patients will provide valuable insights into individual drug resistance and guide personalized targeted interventions for tumor progression.

Discussion

To leverage existing reference data for characterizing out-of-distribution (OOD) cells [5], we developed CellMemory, a Transformer-based model that employs a cross-attention mechanism inspired by global workspace theory. We conducted a comprehensive comparison of 3 scFMs and 16 task-specific methods for single-cell data integration and annotation, utilizing data from over 15 million cells. In these experiments, CellMemory demonstrated three advantages: (1) exceptional generalization performance; (2) greater computational efficiency; and (3) the ability to handle longer token sequences.

Owing to its generalized representations, CellMemory outperforms other methods, including scGPT and Geneformer [23], particularly in cross-omics (single-cell and single-cell spatial omics) situations. We attribute this performance to the bottlenecked architecture, which incorporates attention and competitive mechanisms that enable the retention of essential information [20]. We also evaluated different strategies for processing tokens in single-cell data and observed that these strategies have varying impacts on model performance. When integrating datasets that contain technical variability, such as single-cell and single-cell spatial data, we recommend scaling gene token values into bin ranges. Furthermore, self-attention-based scFMs often exclude genes with zero expression to reduce computational cost [16, 17]. Our analyses revealed that retaining zero expression genes in cells facilitates a more comprehensive understanding of cell representations. These findings emphasize the importance of devising efficient model architectures that capture a broader range of gene regulatory interactions. Notably, although CellMemory achieved good integration performance, recent studies have cautioned that some integration metrics may be misleading, and more holistic evaluation strategies could be beneficial [84].

We then focused on scenarios involving cross-omics, ethnic populations, and comparisons between disease and healthy states to explore interpretable inference for OOD cells. Notably, by analyzing the genes that received the highest attention in cell type specific memory slots, we identified significant differences in the biological functions of gene programs targeted by different memory slots. Various specialist modules were organized into distinct shared global workspaces within a bottleneck architecture, ensuring consistency in task completion. This process is analogous to how the brain coordinates and integrates driving-related visual and auditory information in the consciousness space for safe driving. This hierarchical interpretation facilitates the analysis of OOD objects related to cellular states, further enhancing our understanding of the dynamic manner in which the model processes biological information.

The identification of malignant cells and the characterization of disease states have posed significant challenges in analyzing single-cell datasets. Previous studies have shown that accurately characterizing disease state cells typically requires the inclusion of control groups (usually derived from healthy samples) during model training [58]. CellMemory, without relying on control samples, exhibits exceptional accuracy in distinguishing between normal and malignant cells in lung cancer. It effectively illustrates the cell representations of disease samples, investigates the intermediate states associated with LUAD origin, and reveals the heterogeneous founder cells in certain individuals. These findings emphasize the crucial role of employing a model capable of discerning the subtle details with utmost clarity to comprehensively interpret disease states. By assisting researchers in delineating both the broader context and intricate specifics, CellMemory advances the practice of precision medicine at an individualized level.

As advancements in single-cell references continue, CellMemory is expected to play a pivotal role in facilitating reference mapping without the need for repeated training. Furthermore, its bottlenecked architecture is ideally suited for integrating multimodal data, thus enabling comprehensive information exchange within the low-dimensional memory space. Its capability for processing long tokens makes it particularly well-suited for handling sparse, high-dimensional single-cell multi-omics data, and reducing pre-training costs. In the forthcoming phase, we aim to extend the applications of bottlenecked Transformer across these domains.

Conclusions

We present CellMemory, an interpretable and accurate approach for representing OOD cells. The low-dimensional memory space within the CellMemory architecture enables hierarchical interpretation of OOD cells. We extensively validated CellMemory on more than 15 million cells from single-cell and spatial transcriptomic profiles, demonstrating its superior generalization and computational efficiency compared to existing single-cell foundation models. Furthermore, its interpretable cell representations provide valuable insights into the characterization of malignant cells and their founder cell populations. As high-quality single-cell reference atlases continue to advance, CellMemory is poised to enable robust and scalable mapping of OOD cells onto these references.

Methods

The Global Workspace Theory

Our project builds on the Global Workspace Theory (GWT) in neuroscience, which suggests that consciousness arises from the selective dissemination of information through a shared “global workspace” [19–22].

The brain consists of numerous specialized modules and processes that handle various tasks, such as vision, hearing, memory, and emotions. The “global workspace,” a network of densely interconnected neurons, enables these specialized modules to share and communicate information. However, not all pieces of information can reach the global workspace due to its limited capacity. Instead, the specialized modules may compete with each other to write information to the global workspace. When certain information becomes dominant in the global workspace, it enters our conscious awareness, while other, non-dominant information remains in the unconscious realm. Once information occupies the global workspace, it is broadcasted to diverse specialized neural modules through synchronized neural firing and rapid signaling. This broadcast mechanism ensures that various brain regions concurrently process the information, contributing to a unified conscious experience.

Recently, GWT was used to model complex environments in deep learning [22]. Their approach leveraged the shared global workspace as a communication bottleneck among specialists, significantly outperforming the conventional self-attention frameworks across various visual reasoning benchmarks.

Inspired by the GWT in neuroscience and the promising results in deep learning, our project tailored GWT for analyzing single-cell omics data. Single-cell data presents a unique challenge due to its notoriously high dimensionality and intricate underlying patterns. We found that the global workspace model is ideally suited to address this issue. Its bottleneck architecture acts as a filter, prioritizing and highlighting the most significant information, ensuring that the most pertinent biological details are efficiently communicated and processed. Furthermore, the conventional self-attention mechanism used by BERT presents a quadratic complexity concerning sequence length, leading to increasing time and memory complexity as sequences expand [12, 13, 24]. This becomes drastically problematic when dealing with lengthy single-cell omics data. Alternatively, we employed cross-attention to implement the shared global workspace, substantially reducing time and space complexity and resulting in a more efficient single-cell omics data analysis.

The methodology of CellMemory

Assuming we have an input matrix X $\in$ ${\mathbb{R}}^{N\times M}$, where $N$ represents the number of cells and $M$ stands for the total number of genes.

Embedding

During the embedding process, we encoded each gene of each cell using a word embedding strategy. The processed expression values were transformed into discrete values through a bag-of-words approach (norm: rounding up, or bin: scaling to bin ranges), generating token embeddings. The positions of genes were utilized to create position embeddings.

Construct specialist matrix

For each cell, $n\in N$ of the input matrix X, an extra CLS token is appended to the front of genes. This extended representation is further converted into a $K$-dimensional embedding space, resulting in the specialist matrix R $\in {\mathbb{R}}^{N\times (1+M)\times K}$, where K is a hyperparameter (please refer to “Hyperparameter settings” for details) identifying the number of dimensions essential for processing complex data relationships, and M represents the number of specialists (genes).

Compete to write

The specialist modules $\mathbf{R}$ compete to write the information to the shared global workspace (Additional file 1: Fig. S1a), resulting in the memory matrix D $\in {\mathbb{R}}^{N\times H\times K}$, with $H$ distinct memory slots serving as a hyperparameter. Typically, $H$ is much smaller than the original input dimension $M$, reflecting the bottleneck of the global workspace. The writing procedure involves a cross-attention mechanism between the specialist matrix $\mathbf{R}$ and the memory matrix $\mathbf{D}$, detailed as follows.

1. Initially, we construct the Query, Key, and Value matrices. Here,${\tilde{\mathbf{W}}}_q$, ${\tilde{\mathbf{W}}}_k$, and ${\tilde{\mathbf{W}}}_v$ represent the respective weight matrices for the Query, Key, and Value when calculating the cross-attention matrix.

   $$\begin{array}{c}\mathrm{Query}:\tilde{\mathbf{Q}}=\mathbf{D}{\stackrel{\mathbf{\sim }}{\mathbf{W}}}_{\mathbf{q}}\\ \mathrm{Key}:\stackrel{\mathbf{\sim }}{\mathbf{K}}=\mathbf{R}{\stackrel{\mathbf{\sim }}{\mathbf{W}}}_{\mathrm{k}}\\ \mathrm{Value}:\stackrel{\mathbf{\sim }}{\mathbf{V}}=\mathbf{R}{\stackrel{\mathbf{\sim }}{\mathbf{W}}}_{\nu }\end{array}$$ 1
2. Then, the attention matrix $\tilde{\mathbf{A}}$ is calculated by multiplying the Query and Key matrices, scaled by the last dimension of attention weights $d$, and applying a softmax function for normalization. We select the top $k$ largest attention scores along the $1+M$ dimension, where specialized modules compete to write relevant information to the global workspace. This algorithm aligns with the principles postulated in the global workspace theory in neuroscience, wherein various cognitive processes compete to access a centralized, conscious processing stage. It is worth noting that, in our experiments, we observed a slight performance improvement when using the top $k$ mechanism, but it impacted the interpretability of the model results. Therefore, in this article application, we did not employ the top $k$. If the focus is solely on cell type prediction and the input tokens are relatively long, we recommend considering the use of the top $k$.

   $$\stackrel{\mathbf{\sim }}{\mathbf{A}}={\mathrm{Top}}_{\mathrm{k}}\left(\mathrm{softmax},\left(\frac{\stackrel{\mathbf{\sim }}{\mathbf{Q}}{\stackrel{\mathbf{\sim }}{\mathbf{K}}}^{\mathrm{T}}}{\sqrt{\mathrm{d}}}\right)\right)$$ 2
3. Finally, the attention matrix $\tilde{\mathbf{A}}$ is multiplied by the Value matrix $\tilde{\mathbf{V}}$, resulting in the updated memory $\mathbf{D}$.

   $$\mathbf{D}=\stackrel{\mathbf{\sim }}{\mathbf{A}}\stackrel{\mathbf{\sim }}{\mathbf{V}}$$ 3

Broadcast to specialists

The memory matrix $\mathbf{D}$ is broadcasted back to the specialists (Additional file 1: Fig. S1b), leading to the creation of an updated specialist matrix $\mathbf{R}$. This broadcasting action also utilizes a cross-attention mechanism, which is demonstrated as follows.

1. Similar to the writing mechanism, we formulate the Query, Key, and Value matrices when computing the cross-attention for broadcasting. Specifically, the weight matrices associated with the Query, Key, and Value are represented by ${\widehat{\mathbf{W}}}_q$, ${\widehat{\mathbf{W}}}_k$, and ${\widehat{\mathbf{W}}}_v$, respectively.

   $$\begin{array}{c}\mathrm{Query}:\widehat{\mathbf{Q}}=\mathbf{R}{\widehat{\mathbf{W}}}_q\hfill \\ \mathrm{Key}:\widehat{\mathbf{K}}=\mathbf{D}{\widehat{\mathbf{W}}}_k\hfill \\ \mathrm{Value}:\widehat{\mathbf{V}}=\mathbf{D}{\widehat{\mathbf{W}}}_v\hfill \end{array}$$ 4
2. Then, we calculate the attention matrix $\widehat{\mathbf{A}}$ by applying softmax over the slots.

   $$\widehat{\mathbf{A}}=\mathrm{softmax}\left(\frac{\widehat{\mathbf{Q}}{\widehat{\mathbf{K}}}^{\mathrm{T}}}{\sqrt{d}}\right)$$ 5
3. Eventually, we construct the updated specialist matrix $\mathbf{R}$.

   $$\mathbf{R}=\widehat{\mathbf{A}}\widehat{\mathbf{V}}$$ 6

The processes of “Compete to Write” and “Broadcast to Specialists” are not confined to singular iterations. Instead, they can be orchestrated in a recurrent sequence: write → broadcast → write → broadcast, and so forth. By repeatedly updating the global workspace and broadcasting for specialized feedback, the system ensures a comprehensive integration of information, as each iteration enhances the clarity and resolution of the analysis. Therefore, when applied to the single-cell omics data, this repeated cycle promotes a thorough understanding and in-depth analysis of the data.

Model training and prediction

In the training process of CellMemory, the CLS token matrix S $\in {\mathbb{R}}^{N\times K}$ is extracted from the first dimension of the revised specialist matrix R at index 0: $S_{n,k}=R_{n,0,k}$. Through the iterative process of writing and broadcasting, the CLS token acquires substantial intracellular information. The CLS token S is then projected through a fully connected layer f and processed to the probability for each class c $\in \mathbf{C}$ using softmax function. The predictions are compared to the true labels, and the cross-entropy loss ${\mathcal{L}}_{\mathit{CE}}$ is calculated as follows.

$$\widehat{\mathit{y}}=softmax\left(f,\left(\mathit{S}\right)\right)$$ 7

$${\mathcal{L}}_{\mathit{CE}}=-\frac{1}{N}{\sum }_{n=1}^N{\sum }_{C=1}^C{y}_{n,c}log{\widehat{y}}_{n,c}$$ 8

CellMemory is guided to focus on acquiring information related to the identity representation of cells. During the prediction, the softmax function assigns each cell’s prediction results to the probability of each category, with the category having the highest probability (or confidence score) being considered as the cell’s identity.

Hyperparameter settings

The gene expression matrix, after undergoing embedding, is input into a 4-layered transformer along with the CLS token for classification. We set the embedding dimension (K) to 256 and the feed-forward network dimension to 512. The Query and Key sizes are set to 32, and the Value size is set to 64. We employ 4 heads during reading from and writing into the global workspace, which comprises 8 memory slots (H). We train the model using an early stopping strategy, stopping training when the validation set’s loss has not decreased for 5 epochs. We use Adam optimizer with a learning rate of 0.0003 and anneal the learning rate using cosine annealing [22].

Performing downstream tasks

1. CLS token: The CLS token S $\in {\mathbb{R}}^{N\times K}$ effectively communicates with gene tokens inside the cell through attention mechanisms and the bottleneck architecture in a supervised manner, obtaining a complete cellular representation. It is employed for single-cell and single-cell spatial data reference mapping.

2. Attention/Memory score matrix: The interpretability of CellMemory originates from the cross-attention mechanism’s explicit demonstration of the importance of input features. The gene expression matrix undergoes a writing and broadcasting process with the memory, leading to the generation of two attention score matrices ($\tilde{\mathbf{A}}$ $\in {\mathbb{R}}^{N\times H\times (1+M)}$ and $\widehat{\mathbf{A}}\in {\mathbb{R}}^{N\times (1+M)\times H}$). Matrix multiplication is performed on the two attention matrices derived from the writing and broadcasting steps, yielding a resultant matrix Z = $\widehat{\mathbf{A}}\tilde{\mathbf{A}}$ $\in {\mathbb{R}}^{N\times (1+M)\times (1+M)}$. From matrix Z, we extract matrix Z $\in {\mathbb{R}}^{N\times M}$, where ${\mathbf{Z}\prime }_{n,m}={\mathbf{Z}}_{n,1,m+1}$, capturing the attention of each CLS token concerning each of the M genes across N cells.

   • (i)
     Interpretation Level 1 (gene level): To generate cell-type-specific TAGs, we compute the average attention score for specific groups of cells based on the global attention score matrix Z. Then, the scores are sorted to obtain the top attention genes for each group. In this manner, we obtain the feature attribution score for each cell. Typically, we observe that, within a given cell type, the attention scores of the top 20 genes exhibit clear consistency with gene expression levels. When systematically analyzing the functional pathways associated with these top attention genes, we recommend considering a larger gene set, such as the top 50 or 100 genes.
   • (ii)
     Interpretation Level 2 (gene program level): To generate memory slot specific TAGs, we utilize the memory matrix $\widehat{\mathbf{A}}$ produced during broadcasting. The mean is calculated across selected cells, then, based on averaged memory scores, the features are sorted to obtain TAGs for each slot. At this level, it becomes possible to determine, for each cell, which features each slot within the model focuses on the most. When analyzing the functional pathway enrichment of genes associated with a memory slot, we suggest using a gene set comprising the top 50 genes.

Theoretical analysis

CellMemory is more computationally efficient than the conventional self-attention mechanisms, especially in single-cell omics data analysis. Specifically, the traditional self-attention approach used in BERT is characterized by quadratic complexity as a function of sequence length. This inherent complexity becomes a substantial impediment as sequences extend, particularly when dealing with voluminous single-cell omics datasets. Using traditional self-attention for such extensive data requires substantial computational resources and elongated processing time, making it unsuitable for real-time applications. In contrast, CellMemory harnesses a cross-attention mechanism to implement the shared global workspace, significantly reducing time and space complexity.

Time complexity

Our proposed method, CellMemory, outperforms traditional self-attention mechanisms in terms of computational time complexity. Consider the “Compete to Write” procedure. For each cell $n\in N$, we have Query ${\tilde{\mathbf{Q}}}_n$ $\in$ ${\mathbb{R}}^{H\times K}$ and Key ${\tilde{\mathbf{K}}}_n$ $\in$ ${\mathbb{R}}^{\left(1,+,M\right)\times K}$, and the total time complexity to calculate cross-attention matrix ${\tilde{\mathbf{A}}}_n$ is $\mathcal{O}(MHK)$ in our proposed method. Similarly, due to the consistent use of cross-attention, the time complexity to calculate ${\tilde{\mathbf{A}}}_n$ is also $\mathcal{O}(MHK)$ in the “Broadcast to Specialists” procedure. In comparison, the traditional self-attention mechanism exhibits a total time complexity of $\mathcal{O}(M^{2}K)$. Specifically, $H$, the hyperparameter representing memory slots, is typically chosen to be much smaller than the input dimension $M$ to reflect the bottleneck of the shared global workspace.

Space complexity

Furthermore, CellMemory also surpasses the traditional self-attention mechanisms regarding space complexity. For each cell $n\in N$, the “Compete to Write” procedure yields a cross-attention matrix ${\tilde{\mathbf{A}}}_n$ with dimensions $H\times (1+M)$, while the “Broadcast to Specialists” procedure results in ${\tilde{\mathbf{A}}}_n$ with dimensions $(1+M)\times H$. Consequently, the space complexity of our method is $\mathcal{O}(MH)$. In contrast, the traditional self-attention mechanisms operate with a quadratic space complexity of $\mathcal{O}(M^2)$. Given $H\ll M$, it is evident that CellMemory offers a more memory-efficient alternative, saving valuable computational resources compared to the conventional self-attention techniques.

Relation with other transformer variants

Transformers, with their attention-based architecture, have become a cornerstone in deep learning. However, traditional transformers face challenges like quadratic time and memory complexity and a large parameter count, impeding scalability. To address these issues, numerous transformer variants have emerged in recent years. Specifically, several works employ sparse attention to reduce computation time and memory usage [85, 86]. Yet, sparse attention requires sparse matrix multiplication, which is not ubiquitously supported on all accelerators [87]. To address these challenges, linear attention has emerged, where softmax attention is approximated using random feature kernels [87, 88]. However, this approximation often underperforms relative to the vanilla self-attention mechanism, especially when scaled to large models [89]. Additionally, the Funnel Transformer progressively reduces the length of hidden representations to handle long sequences [90]. Although it efficaciously alleviates the computational and memory cost, the complexity still remains quadratic [91]. Similarly, the Compressive Transformer compresses memories for long-range sequences [92], but demands extra memory to store compressed hidden states compared to the vanilla Transformer. Moreover, the Linformer addresses long sequences using low-rank approximation methods; its projection operation, however, introduces additional computational costs [93].

Given these constraints, we adopt a bottlenecked architecture to enhance efficiency by reducing both computational and memory overhead. Our method aligns with variants incorporating a “Neural Memory” and implementing down-sampling within the attention mechanism [94]. This approach is directly inspired by the transformer shared global workspace [22] and bears similarities to the Perceiver model [95]. Both strategies refine the pairwise self-attention mechanism by introducing global coordination and a unified memory. This integrated memory acts as a constrained communication channel, distinct from other transformer bottleneck techniques like discrete bottleneck [96] and restricted cross-modal attention [97]. In our model, bottleneck effects are achieved through a small, integrated memory, facilitating information integration and dimension reduction. These features are particularly advantageous for processing complex and noisy biological data, as in our study.

Data preprocessing

During the data preprocessing stage, we prepared multiple formats of the data to accommodate the input requirements of different methods. First, single-cell gene expression data were standardized using a log-transformation, followed by the identification of highly variable genes (HVGs). All models were trained on a cell-by-HVGs matrix (normalized data or raw counts data) of consistent shape. For example, methods such as TOSICA, Seurat require a standardized matrix, whereas scVI utilizes raw count matrices for input. Geneformer requires raw count matrices and converts gene expression signals into rank information, while scGPT scales each cell’s gene expression signals into a bin range of 1–50. All methods were processed using their default recommended data input pipelines. In the benchmark (annotation and integration) comparison, CellMemory maintained a consistent preprocessing approach by using normalized gene expression data for model training. Specifically, for CellMemory, all input genes, including zero-expression genes were considered, and the gene expression token embeddings were processed in the norm manner (see the methodology of CellMemory). CellMemory-fast filters out genes with zero expression in each cell, with the token embeddings processed in the norm manner. The preprocessing strategies for scGPT, Geneformer, and UCE all involve filtering out genes with zero expression in each cell.

Single-cell dataset

hBreast

We collected a single cell atlas of the adult human breast from the CELLxGENE database. This dataset includes 244,285 cells (100,000 cells from normal samples by down-sampling; 144,285 cells from breast cancer samples; 123 batches), identifying 58 biological cell states. The data was subset to the 3000 highly variable genes (HVGs) to benchmark. We established an annotation scenario using normal cells as the reference for annotating cells from breast cancer samples. The CLS embedding of the query set was evaluated in the integration benchmark. Batch information was denoted using the “library_uuid.”

hLung

The Lung atlas is available in the CELLxGENE database, comprising both normal and LUSC samples. It contains 305,319 cells (212,889 cells from normal samples and 92,430 cells from LUSC samples) and identifies 24 cell types. We restricted the reference dataset to data exclusively from the 10x platform to predict the cell type of the query set of LUSC. The query set incorporated data from 5 platforms (SMART-seq3, Singleron, DB, etc.). In the integration benchmark, the term “dataset” was utilized as the batch label. A total of 3000 HVGs were used for benchmarking.

hPancreas

The human pancreas dataset includes 47,644 cells (26,197 cells from control samples; 21,447 cells from T1D patients, and 13 cell types). We removed the hybrid cell type from the dataset, and cells from control donors were used as the reference while cells from T1D patients were analyzed as the query set. A total of 3000 HVGs were used for benchmarking.

Immune

The training set was composed of 141,524 immune cells derived from 8 tissues, including the lung and liver. The test set consisted of 188,238 immune cells from 9 tissues, such as the spleen and thymus. A total of 4000 HVGs were used for benchmarking.

mPancreas

The mPancreas contains scRNA-seq data from 239,679 mouse pancreatic islet cells (with the reference down-sampled to 100,000). The normal cells within the dataset served as a reference to train the model, and the T1/2D samples were utilized as the query set (139,679 cells). This dataset comprises 20 cell types, consisting of 44 batches. The term “batch_integration” was utilized as the batch label. A total of 3000 HVGs were used for benchmarking.

mHypoMap

The mouse hypothalamic atlas from HypoMap contains 66 cell populations derived from the transcriptional profiling of 161,903 cells, consisting of 29 batches. We used the data from the Drop-seq platform as the query set (61,903 cells) and the data from the 10x platform as the train set (downsampled to 100,000 cells). The term “Sample_ID” was utilized as the batch label. A total of 4000 HVGs were used for benchmarking.

crossSpecies (cortex)

In the cross-species benchmark, we utilized human cerebral cortex data, comprising 200,000 cells, as the training set. The query set was constituted by a dataset that includes one human and four non-human cerebral cortex datasets (human: 156,285 cells, Gorilla: 139,945 cells; Chimpanzee: 112,929 cells; Macaque: 89,136 cells; Marmoset: 75,861 cells). This part ensured the retention of a set of homologous genes across all species. The performance of the integration was evaluated using chimpanzee cortex data. A total of 3000 HVGs were used for benchmarking.

Drosophila

The Drosophila dataset (body) comprises 265,979 cells (17 cell types), categorized into four age groups (5, 30, 50, 70 days), aiming to investigate aging modifications in Drosophila at the single-cell level. The younger samples of 5 days and 30 days served as the reference (173,181 cells), and the older samples of 50 days and 70 days were used as the query (92,798 cells). In the benchmark comparison, we omitted Geneformer and scGPT from the Drosophila dataset due to substantial disparities between the gene features used for their pre-training and those of Drosophila’s gene features. A total of 3000 HVGs were used for benchmarking.

Breast cancer atlas

Using the breast cancer single-cell data as a reference, we constructed a model to integrate the 10x Xenium data. We retained the overlapping gene set (301 genes) with the Xenium data and included cells with counts more than 20, and genes more than 10, totaling 94,535 cells. The model was built at the “celltype_major” resolution, integrating the Xenium and 10x 3’/5’ single-cell data (from the same sample). Gene expression values were scaled to the bin range.

Brain cell atlas of Alzheimer’s disease (SEA-AD)

This atlas was sampled from the dorsolateral prefrontal cortex of Alzheimer’s disease patients, comprising 1,395,601 cells. Among them, 796,768 cells originated from normal samples and 598,833 cells from disease samples. It includes annotations at both subclass (24) and supertype (131) resolutions. In the integration assessment, cells from normal samples were used to construct a reference, which was then integrated with cells from disease samples, while retaining the 3000 highly variable genes.

The Asian Immune Diversity Atlas (AIDA)

AIDA is an immune dataset comprising 985,470 PBMC cells from 503 healthy donors across five Asian ancestries from Japan, Singapore, and South Korea, along with European controls (5 donors, 73,439 cells) via 10x 5’ single-cell sequencing.

Adolescent mouse nervous system

The single-cell mouse brain atlas of cell types from the Linnarsson Lab is employed as the reference dataset to characterize Stereo-seq adult mouse coronal hemibrain dataset. Cell types consistent with the organization of the query dataset were retained, and the model was trained using over 120 cell types. This work provided identified marker genes, which were used to train the model for reference mapping.

Mouse whole-brain scRNA-seq atlas

A scRNA-seq whole adult mouse brain dataset of 4,057,701 cells was generated by the 10x platform (v2, v3) with 338 subclasses. This dataset was used to build a model to infer mouse whole-brain MERFISH data. The gene set was intersected with MERFISH data to retain 500 genes, and cells with more than 30 captured genes were kept.

CITE-seq PBMC

We used PBMC cells produced by CITE-seq (161,764 cells, 36 classes) to integrate and annotate the AIDA dataset. A total of 4000 highly variable genes were utilized.

MPAL lineage inference

Bone marrow and peripheral blood mononuclear cells from healthy samples were used, totaling 35,582 cells (MBBCs=12,602; CD34+-enriched BMMCs=8176; PBMCs=14,804). A model was constructed using 1000 highly variable genes retained from the reference. MPAL-a (3435 cells) corresponds to B-myeloid MPAL (MPAL4) in the original study, and MPAL-b (2,947 cells) corresponds to T-myeloid MPAL (MPAL5). Cells labeled as healthy-like in the original study were excluded. We classified lineages by cell type. Erythroid_like: Late.Eryth. Lymphoid_like: B; Pre.B; CLP.1; CLP.2; Plasma; Unk_1. Myeloid_like: GMP.Neut; GMP; pDC; cDC; Unk_2; CD14.Mono.1; CD14.Mono.2; CD16.Mono. Progenitor_like: Early.Eryth; Early.Baso; HSC; CMP.LMPP. TNK_like: NK; CD8.CM; CD8.N; CD8.EM; CD4.M; CD4.N1; CD4.N2.

Origin inference for different MB subgroups

Glutamatergic neuronal lineages from the fetal cerebellum were used as the reference, with 500 highly variable genes. Malignant MB cells produced by Smart-seq4 (including SHH, Group 3, and Group 4 subgroups from 31 donors) were used as the query set.

Integration of lung cancer dataset

We used all healthy samples, LUSC samples, and LUAD samples (downsampled to 200k cells) from hLung dataset, and retained 3000 highly variable genes. We then integrated a lung cancer dataset comprising 293,432 cells from 52 lung cancer patients (including LUAD and LUSC) using the trained CellMemory. For tumor cell similarity inference, all tumor cell types were removed from the reference, retaining 5000 highly variable genes to train the model.

Integration of LUAD dataset 2

A model was constructed using the same reference as the “Integration of lung cancer dataset.” This model was then integrated into the LUAD dataset comprising 488,236 cells from 58 LUAD donors.

Single-cell spatial dataset

mHypo

The mouse hypothalamic spatial dataset measured by MERFISH, comprising three samples with naïve behavior, totaling 180,754 cells. The single-cell data from mHypoMap was used as the reference. Aligning genes (154) were captured by MERFISH resulting in 8 categories.

hNSCLC

The scRNA-seq dataset of hLung served as the reference to predict NSCLS spatial dataset. The NSCLC spatial dataset includes Lung-6, Lung-9, and Lung-13, consisting of 256,581 cells from 11 cell types. These samples were generated from a Formalin-Fixed Paraffin-Embedded (FFPE) sample by the CosMx SMI platform. A shared gene set of 952 genes from the reference dataset was selected to align with the query dataset.

mSpermato

The spatial dataset of mouse spermatogenesis (measured by Slide-seq) was acquired from three wild-type (WT) mice and three leptin-deficient diabetic (ob/ob) mice, including 207,335 cells and 9 cell types. The WT dataset served as the reference, and the ob/ob dataset served as the query set.

Xenium breast cancer

This dataset consists of two replicate samples from a breast cancer patient, as well as single-cell 3’ and 5’ data. Cells with counts of more than 20 and more than 10 genes were retained, totaling 305,631 cells and 301 genes.

Slide-tags human cortex

The spatial transcriptome data generated by Slide-tags sequencing consists of 4065 cells belonging to 7 cell types. Cells from normal samples in the brain cell atlas of Alzheimer’s disease were utilized as a reference for integrating the Slide-tags spatial data.

Stereo-seq adult mouse coronal hemibrain

The single-cell spatial dataset of the mouse brain hemisphere (including 50,140 cells) was measured by Stereo-seq. We utilized segmented data as the query set. It maintained the subset of genes that align with the reference (“Adolescent mouse nervous system”).

Mouse whole-brain MERFISH atlas

The MERFISH dataset, encompassing 4,334,174 cells across 59 sections from one whole male mouse brain with a 500-gene panel, includes a hierarchical cell type annotation, organized into four levels of cellular identities. To evaluate the performance of benchmark models, we used the subclass information. Cells with more than 30 captured genes were kept.

Comparison with annotation methods

In the annotation benchmark, we trained models using reference data while predicting the query data. All methods generated cell type predictions for each cell in the query data. We ensured consistency between the reference and query data, employing the same cell type resolution. Additionally, the inputs for all methods were provided according to their default requirements. CellTypist (1.5.3, python) emerged as an efficient annotation tool that outperforms most methods. Seurat (4.3.0.1, R) is a widely employed method in single-cell data analysis. In this study, we utilized the implementation of the cell type annotation pipeline as described in the “Mapping and annotating query datasets” chapter. TOSICA (1.0.0, python) is a self-attention Transformer-based model that computes attention at the pathway level and performs cell type annotation. scPoli (python) used the default set.

Comparison with the pre-trained Transformer model

Geneformer

Geneformer is a scFM that leverages pretraining with a dataset of 30 million cells, allowing it to perform tasks such as cell type annotation and the inference of key network regulators. The cells used for pre-training are sourced from human data. Geneformer requires pre-trained weights as input. We implemented the cell type annotation task using the annotation process provided by Geneformer. In the benchmark tasks, Geneformer was fine-tuned.

scGPT

scGPT is a scFM that leverages pre-training with a dataset of 33 million cells. The cells used for pre-training are sourced from human data. The Flash-attention version is used to benchmark. We used the annotation and integration pipeline from scGPT to separately perform the annotation and integration tasks. In the benchmark comparisons, scGPT was fine-tuned.

UCE

A scFM pre-trained on over 30 million cells, comprises 33 Transformer layers and 650 million parameters. We used its annotation pipeline to perform the annotation task. We followed UCE’s default usage, wherein its annotation strategy involves obtaining latent embeddings via a zero-shot approach, followed by a classification head.

scBERT

scBERT adopts the advanced strategy derived from BERT, tailored to implement single-cell data analysis. scBERT is exclusively pre-trained on a human dataset (Panglao dataset, 1,126,580 cells), there are no pre-trained models available for other species.

Comparison with spatial annotation methods

In the spatial data annotation task, we trained the model using single-cell reference data and then used it to annotate cell types in single-cell spatial transcriptomics datasets. In this benchmark scenario, the training and test sets for each model were kept consistent. CellMemory seeks to label individual cells and cell-like objects, such as Stereo-seq spots or Slide-seq beads, through discrete annotations. To enhance the interpretability of spatial single-cell data, understanding the most critical features for each cell is essential. Compositional annotation refers to probabilities associated with discrete categories, indicating objects closely related to the assigned category (the class with the maximum compositions). Default parameters were utilized for all other spatial annotation methods (Tangram (“clusters” level), TACCO, RCTD, NMFreg, NNLS).

Comparison with integration methods

For methods that perform reference mapping, we trained the model on the reference data and generated cell embeddings and label predictions for the query data. These cell embeddings were then used for further comparison. For unsupervised dimensionality reduction methods, we directly input the query data to obtain cell embeddings. In benchmark integration tasks, aside from scGPT and scPoli, all other integration methods were performed using scIB [37]. scIB is an integration tool that consolidates many advanced integration methods and incorporates a comprehensive set of evaluation metrics. This tool enables users to efficiently compare multiple integration approaches. scGPT performs the integration process, and scPoli performs integration and annotation using default steps. The CellMemory CLS embedding used by the integration benchmark is the same as that used by the annotation. The original cell types are applied to biological conservation evaluation.

Training strategy of the CellMemory

In the process of model training, the reference dataset is divided into a training set and a validation set with an 80:20 ratio. In the annotation benchmark validation, we employed a five-fold cross-validation to ensure the reliability of the results.

Evaluation metrics

We employed accuracy and macro f1-score as evaluation metrics to assess the classification performance. These metrics were computed using the Python library sklearn. We assessed the data integration using metrics from scIB package. The overall integration score is a weighted average of the average batch mixing score and the average biological conservation score, with weights of 0.4 and 0.6 respectively.

Enrichment analysis

The enrichment analysis was performed using the Metascape [98] web-based platform. The analysis was conducted using the TAGs (Interpretation level 1: top 1–100 TAGs; Interpretation level 2: top 1–50 TAGs), with enriched terms including GO terms and pathways. In the enrichment results, the length of each bar represents statistical significance (log10-transformed p-value), and only pathways meeting a significance threshold (p<0.01) were retained. To facilitate the comparison between the two interpretation levels, we retained only the top eight most significant enrichment results, corresponding to the number of slots.

Centroid determination for MB

When inferring the identity of MB origin cells, we first applied Leiden clustering to all MB cells, resulting in 13 clusters. We identified the cluster with the highest mean confidence score for each cell type, then averaged the embedding coordinates of all cells in that cluster to determine its centroid. The embedding coordinates of all cells from a specific donor were used to determine the donor centroid, while the embedding coordinates of all cells from a specific subgroup were used to determine the subgroup centroid.

Identification of tumor cells

In the decipher of lung cancer data (LUAD+LUSC), characterizing tumor cells relative to other epithelial cells is a critical task due to their high degree of similarity. During model training with a lung cancer atlas that includes tumor cell type, we performed five-fold cross-validation to quantify the accuracy of tumor cell annotations. We compared with scPoli, another method capable of reference mapping, using default parameters.

Cell-cell interaction

In the analysis of LUAD data, CellPhoneDB [74] was used to analyze the interactions among cell types associated with tumor origins (e.g., AT2, Club, Transitional cells) and those that potentially interact with them (e.g., Monocyte, Macrophage, Stromal cells), employing default parameter setting.

CNV inference

In the inference of LUAD data, inferCNV [75] was employed to infer CNV patterns and levels in lung cancer patients. Our focus was on tumor cells and epithelial cells that might be associated to tumor origins, with immune cells serving as the normal control. We ran the inferCNV with default parameters.

Definition of AT2 subgroups

The Leiden clustering algorithm was employed to analyze the epithelial cell embeddings generated by CellMemory in LUAD analysis. AT2 cells that overlapped with cluster 5 were designated as AT2-1, while AT2 cells that were not included in cluster 5 were labeled as AT2-2.

Authors’ contributions

D.L., L.J., and Q.W. designed the project. Q.W. and H.Z. built the model. Q.W., Y.H., and Y.C. conducted the model application, benchmark, and analysis. Q.W., L.J., D.L., Y.W., and H.Z. wrote the manuscript. Y.L., G.L., J.C., Y.L., J.Z., X.Z., and M.K. provided comments and application suggestions for the project. All authors discussed the results and commented on the manuscript.

Funding

This research was supported by grants from the National Key Research and Development Program of China (grant no. 2024YFA1802101 to L.J.), the Strategic Priority Research Program of the Chinese Academy of Sciences (grant no. XDA0460305 to J.C., XDA0460402 to G.L. and XDB38020500 to L.J.), the National Key Research and Development Program of China (grant nos. 2023YFC3402703 and 2019YFA0801702 to L.J.), the National Natural Science Foundation of China (grant nos. 92374104 and 31970760 to L.J., 32400484 to Y.L. and 32100479 to G.L.), the International Partnership Program of the Chinese Academy of Sciences (grant no. 153F11KYSB20210006 to L.J.), the CAS Youth Interdisciplinary Team to L.J., the Program of the Beijing Natural Science Foundation (grant nos. L246036 and L246006 to L.J.), and the National Funded Postdoctoral Researcher Program (grant no. GZC20232568 to Y.L.). D.L. was supported by the NUHS seed fund (25-0381-A0001).

Data availability

All single-cell and single-cell spatial omics datasets analyzed in the current study are publicly available and can be downloaded from their public repository. Detailed information and download links for each dataset are provided in Additional file 3 Table S2. Preprocessed versions of the data can be accessed via 10.5281/zenodo.15511819 [99]. A comprehensive description of the data preprocessing strategies is available in the Methods section. Our Python implementation along with the versions of the packages used can be found at https://github.com/QifeiWang4Bio/CellMemory and 10.5281/zenodo.15511926 [100] under an MIT license.

Declarations

Ethics approval and consent to participate

Not applicable.

Consent for publication

Not applicable.

Competing interests

The authors declare no competing interests.