GeneMamba: An Efficient and Effective Foundation Model on Single Cell Data

1. Problem Statement

Single-cell RNA sequencing (scRNA-seq) has revolutionized the study of cellular heterogeneity but presents significant computational challenges due to the high dimensionality, sparsity, and massive scale of the data (often involving millions of cells and thousands of genes).

Existing foundation models for single-cell data, primarily based on Transformer architectures (e.g., scBERT, scGPT, Geneformer), face critical limitations:

• Quadratic Complexity: The self-attention mechanism in Transformers scales quadratically with sequence length ($O(N^2)$), making them computationally expensive and memory-intensive for the long sequences typical of transcriptomes.
• Long-Range Dependency: Transformers often struggle to capture long-range dependencies essential for modeling complex gene regulatory networks and cell state transitions.
• Scalability: Training on datasets with tens of millions of cells is resource-prohibitive for standard Transformer models.

2. Methodology

The authors propose GeneMamba, a novel foundation model that replaces the Transformer backbone with a State Space Model (SSM) architecture, specifically a Bidirectional Mamba (Bi-Mamba) module, combined with biologically informed tokenization and loss functions.

A. Data Processing and Tokenization

• Input Representation: Instead of raw counts, the model processes gene expression matrices.
• Rank-Based Discretization: The authors adopt a rank-based strategy (similar to Geneformer). Genes within each cell are normalized and ranked in descending order of expression.
  • Rationale: This approach captures relative expression levels, is robust to batch effects and noise, and aligns with biological regulatory interactions where relative abundance matters more than absolute values.
  • Output: Each cell is represented as a sequence of gene tokens (ranks), with a maximum sequence length of 2,048 or 4,096 genes.

B. Architecture: Bi-Mamba

GeneMamba utilizes the Bi-Mamba block, which addresses the unidirectional nature of standard Mamba models to capture full contextual dependencies.

• Bidirectional Processing: The input sequence is processed in both forward and reverse directions using shared-weight Mamba layers.
• Gating Mechanism: The outputs from the forward and reverse passes are combined using a learnable gating mechanism (sigmoid-based) to balance upstream and downstream context.
• Efficiency: SSMs offer linear computational complexity ($O(N)$) with respect to sequence length, enabling the processing of ultra-long sequences with significantly reduced memory and compute requirements compared to Transformers.

C. Pretraining Objectives

The model is pretrained using a multi-task loss function designed to capture both sequential patterns and biological structure:

1. Next Gene Prediction Loss ($L_{lang}$): A standard causal language modeling objective (Next Token Prediction) where the model predicts the next gene token based on the preceding context within a cell.
2. Gene Pathway Loss ($L_{pathway}$): An InfoNCE contrastive loss that enforces similarity between embeddings of genes belonging to the same biological pathway and dissimilarity between genes from different pathways. This injects biological prior knowledge into the embedding space.
3. Total Loss: $L = L_{lang} + \gamma L_{pathway}$ (where $\gamma = 0.1$).

3. Key Contributions

1. Scalable Foundation Model: Introduction of GeneMamba, the first SSM-based foundation model tailored for single-cell data, capable of training on ~30 million cells (filtered from 50M+ raw cells) with linear scaling.
2. Biologically Informed Architecture: Integration of Bi-Mamba for bidirectional context and Rank-based tokenization to handle sparsity and noise effectively.
3. Biological Interpretability: The use of pathway-guided contrastive learning ensures that the learned embeddings reflect known biological relationships (gene pathways), enhancing interpretability.
4. Comprehensive Validation: Extensive benchmarking across diverse downstream tasks, demonstrating superiority over Transformer-based baselines in efficiency and performance.

4. Experimental Results

The model was evaluated on multiple downstream tasks using datasets like PBMC12k, hPancreas, Myeloid, and COVID-19.

• Multi-Batch Integration:
  • GeneMamba outperformed baselines (Harmony, Geneformer, scGPT, scFoundation) in balancing batch correction (Avg_batch) and biological preservation (Avg_bio).
  • It successfully mitigated batch effects while retaining distinct cell type clusters, a common trade-off in integration tasks.
• Cell Type Annotation:
  • Achieved state-of-the-art accuracy and Macro-F1 scores on datasets like hPancreas (Acc: 0.9713) and Myeloid_b (Acc: 0.9603).
  • Demonstrated robustness in distinguishing nuanced cell types, even in imbalanced datasets.
• Gene Rank Reconstruction:
  • GeneMamba showed superior ability to reconstruct the original gene ranking order compared to Geneformer and unidirectional Mamba variants.
  • Metrics: Achieved a BLEU score of 0.987 and Spearman correlation of 0.711 on the PBMC12k dataset, significantly outperforming Geneformer (BLEU 0.968).
• Computational Efficiency:
  • GeneMamba reduced training time per sample by approximately 42% compared to Transformers at sequence length 2048.
  • At longer sequence lengths (8192), the efficiency gap widened significantly, with GeneMamba maintaining linear scaling while Transformers became prohibitively slow.
• Perturbation Analysis:
  • In gene perturbation tasks (predicting expression changes after genetic intervention), GeneMamba achieved lower Mean Squared Error (MSE) and higher correlation with ground truth compared to baseline GEARS models.

5. Significance and Impact

• Scalability: GeneMamba provides a viable path for training foundation models on tens of millions of cells, a scale previously unattainable with Transformers due to memory constraints.
• Efficiency: By leveraging SSMs, it drastically reduces the computational cost (FLOPs) and training time, making large-scale single-cell analysis more accessible.
• Biological Relevance: The model's ability to reconstruct gene rankings and preserve pathway relationships suggests it captures the underlying biological logic of gene regulation better than previous models.
• Future Foundation: GeneMamba establishes a robust, efficient, and interpretable backbone for future single-cell omics research, facilitating advancements in disease modeling, drug discovery, and developmental biology.

In summary, GeneMamba represents a paradigm shift in single-cell deep learning, moving from quadratic-complexity Transformers to linear-complexity State Space Models, while simultaneously improving biological interpretability through rank-based tokenization and pathway-aware training.