Single-cell mosaic integration and cell state transfer with auto-scaling self-attention mechanism

Abstract

The integration of data from multiple modalities generated by single-cell omics technologies is crucial for accurately identifying cell states. One challenge in comprehending multi-omics data resides in mosaic integration, in which different data modalities are profiled in different subsets of cells, as it requires simultaneous batch effect removal and modality alignment. Here, we develop Multi-omics Mosaic Auto-scaling Attention Variational Inference (mmAAVI), a scalable deep generative model for single-cell mosaic integration. Leveraging auto-scaling self-attention mechanisms, mmAAVI can map arbitrary combinations of omics to the common embedding space. If existing well-annotated cell states, the model can perform semisupervised learning to utilize existing these annotations. We validated the performance of mmAAVI and five other commonly used methods on four benchmark datasets, which vary in cell numbers, omics types, and missing patterns. mmAAVI consistently demonstrated its superiority. We also validated mmAAVI's ability for cell state knowledge transfer, achieving balanced accuracies of 0.82 and 0.97 with less 1% labeled cells between batches with completely different omics. The full package is available at https://github.com/luyiyun/mmAAVI.

Keywords: mosaic integration; self-attention; semi-supervised learning; single-cell; variational inference.

Figure 1

Schematic of mosaic integration and mmAAVI. A, graph illustration of mosaic integration data with three batches and three modalities. Batch effects may arise from differences in individuals, experimental conditions, or tissue sources. Potential modalities include chromatin accessibility (DNA level), gene expression (RNA level), and epitope (protein level), among others. Due to the combined impact of batch effects and modalities missingness, biological variations (such as cell types) are obscured by systematic errors, making direct integration analysis challenging. (b, c) The workflow for unsupervised (b) and semisupervised (c) analysis using mmAAVI. (d) Schematic of mmAAVI model. Multimodal data for each cell are transformed into embeddings by modality-specific encoders and then fused into a global feature , a low-dimensional representation of the cell state following mixture distribution parameterized by discrete and continue . A discriminator is used to harmonize the distribution of across different batches . Meanwhile, a guidance graph with prior knowledge is transformed into feature embeddings by a graph encoder. Next, modality-specific hybrid decoders map samples from the posterior distribution of and , along with the batch, , to parameters of the distribution for each feature of existed modalities. The posterior mean of can be used as input to clustering and visualization algorithms.

Figure 2

Benchmark results on the human PBMC dataset and mouse brain cortex dataset. (a) Layout of input data matrices in the human PBMC dataset. (b, c) UMAP visualization of cell embeddings learned by mmAAVI for the human PBMC dataset, where cells are colored by cell batches (b) and cell types (c). (d) The means of performance scores of mmAAVI and other methods on the PBMC dataset using the coarse annotation (four categories) for cell type labels. (e) Layout of input data matrices in the mouse brain cortex dataset. (f, g) UMAP visualization of cell embeddings learned by mmAAVI for the mouse brain cortex dataset, where cells are colored by cell batches (f) and cell types (g). (h) The means of performance scores of mmAAVI and other methods on the mouse brain cortex dataset.

Figure 3

Benchmark results on the human nephron dataset (Nephron) and adult mouse cortical neuron cell (Triple) dataset. (a) layout of input data matrices in the Nephron dataset. (b–d) UMAP visualization of cell embeddings learned by mmAAVI for the Nephron dataset, where cells are colored by samples (b), omics (c), and cell types (d). (e) The means of performance scores of mmAAVI and other methods on the Nephron dataset. (f–i) UMAP visualization of cell embeddings learned by mmAAVI for the Triple dataset, where cells are colored by cell batches (f) and cell types (g–i) for each omics layer. (k) Layout of input data matrices in the Triple dataset. (j–l) The means of performance scores of mmAAVI and other methods on the Triple dataset. (m) The total score of all methods on the Triple dataset.

Figure 4

Results of semisupervised learning on the human nephron dataset (Nephron) and adult mouse cortical neuron cell (Triple) dataset. (a, b) The test balanced accuracy (bACC) of mmAAVI and benchmark approaches on each batch of the Nephron (a) and Triple (b) datasets, where "global" represents the bACC on all test cells, and "average" represents the average values across bACCs of all batches. (c, d) UMAP visualization of cell embeddings learned by semisupervised version of mmAAVI for the Nephron (c) and Triple (d) datasets, where cells are colored by annotation of “seed” cells, prediction of semisupervised version of mmAAVI, and true cell types. (e, f) The performance scores of semisupervised version and unsupervised version of mmAAVI on the Nephron (e) and Triple (f) datasets.