scMamba：超越高可变特征选择的单细胞多组学整合的可扩展基础模型.

0. TL; DR

scMamba是一个专为单细胞多组学数据设计的可扩展基础模型。它直接处理包含数万个基因和染色质开放区域（peaks）的原始数据，无需预先筛选高可变特征，从而最大程度地保留了生物信息的完整性。

scMamba采用基于Patch的细胞分词策略，将基因组区域视为token，细胞视为句子，并保留了基因组的位置信息。模型基于Mamba结构，能够高效地从高维、稀疏的单细胞数据中学习深层生物学规律。作者提出一种结合了余弦相似度正则化的对比学习方法，实现了跨组学数据更优越的对齐效果。

在多个基准测试中，scMamba 在生物变异保留、组学对齐、细胞聚类、类型注释和轨迹推断等关键任务上，全面超越了现有的 SOTA 方法，并展现了处理数十万细胞级别图谱数据的强大扩展性。

1. 背景介绍

单细胞测序技术能够以前所未有的分辨率窥探单个细胞内的分子世界，例如基因表达（scRNA-seq）、染色质可及性（scATAC-seq）和表面蛋白丰度（CITE-seq）等。而单细胞多组学技术进一步允许在同一个细胞内同时测量多种分子信息（如 RNA + ATAC）。如何有效地拼合这些来自不同分子层面（模态）的信息是一个巨大的计算挑战。

现有的端到端的多模态整合方法如totalVI、GLUE、scCLIP 等，它们直接对多模态数据进行联合建模。这些方法面临一个共同的问题：为了降低计算复杂性，它们要求用户在预处理阶段筛选出 3000-5000个高可变基因（HVGs）或高可变开放区域（HVPs）。这种做法基于一个核心假设：最具生物学意义的信号蕴含在那些变异最大的特征中。但这个假设并不总是成立，许多在细胞身份或功能中起着关键作用的基因（如管家基因、某些发育调控因子）其表达水平可能相对稳定，但在特定上下文中至关重要。

scMamba的目标是构建一个能够直接处理全部原始特征的基础模型，彻底摆脱对高可变特征筛选的依赖，从而更全面地捕捉单细胞数据中蕴含的丰富生物学信号。

2. scMamba 模型

scMamba 的整体架构如下图所示，其核心流程包含三大关键组件：(1) 基于 Patch 的细胞分词，(2) 用于特征提取的 scMamba 编码器，以及 (3) 用于对齐不同组学的对比学习策略。

2.1 基于 Patch 的细胞分词 (Patch-based Cell Tokenization)

传统的单细胞基础模型（如 scGPT）借鉴自然语言处理（NLP），将基因视为token，将细胞视为句子。但这种策略在处理数万个基因时会因序列过长而导致计算量爆炸。

scMamba 提出了一种基于** Patch** 的细胞分策略：

  1. 排序 (Ordering): 首先，将所有基因（或 peaks）按照它们在参考基因组上的物理坐标进行排序。这步操作保留了基因组的空间邻近信息。
  2. 分块 (Patching): 对于每个细胞，将其完整的特征向量（如 20000 个基因的表达值）分割成连续的、不重叠的“Patch”。例如，每 100 个基因作为一个 Patch，一个 20000 维的向量就被分成了 200 个 Patch。
  3. 嵌入 (Embedding): 将每个 Patch 通过一个可训练的线性投影矩阵，映射到一个固定维度的嵌入向量（Embedding）。
  4. 加入位置信息: 为了让模型知道每个 Patch 的原始位置，为每个 Patch 嵌入添加一个可学习的“位置嵌入 (Positional Embedding)”。

最终，对于一个细胞 $i$，其输入嵌入 $T^{(i)} \in \mathbb{R}^{C \times D}$ (其中 C 是 Patch 数量，D 是嵌入维度) 定义为：

\[T = [\mathbf{t}^{(1)}W; \mathbf{t}^{(2)}W; \dots; \mathbf{t}^{(C)}W] + E_{pos}\]

其中$\mathbf{t}^{(j)}$ 是第 $j$ 个 Patch，$W$ 是可学习的投影矩阵，$E_{pos}$ 是位置嵌入矩阵。

这种策略将一个超高维的细胞特征向量抽象成了一个较短的 Patch 序列，既保留了基因组的局部上下文信息，又极大地提高了计算效率，使得处理全量特征成为可能。

2.2 基于 Mamba 的高效编码器

获得细胞的 Patch 序列后，使用 scMamba 编码器 对其进行深度特征提取。该编码器由多个 scMamba Block 堆叠而成，每个 Block 的核心是 Mamba2 模块。

Mamba 是一种状态空间模型（SSM），相较于 Transformer 的自注意力机制，它具有两大优势：

Mamba2 的核心是其状态空间对偶（State Space Duality, SSD）算法。它将复杂的 SSM 运算简化为一个与注意力机制非常相似但更高效的形式。其计算可以概括为：

\[y = M x \quad \text{with} \quad M = L \circ (CB^T)\]

其中 $x$ 是输入序列，$y$ 是输出。矩阵 $L$ 是一个下三角矩阵，它的元素值可以动态变化，赋予了模型根据上下文优先关注或忽略某些输入 Patch 的能力。$L$ 的结构如下：

\[L = \mathbf{1SS}(a_{0:C}) := \begin{bmatrix} 1 & & & \\ a_1 & 1 & & \\ a_2a_1 & a_2 & 1 & \\ \vdots & \vdots & \ddots & \\ a_{C-1}...a_1 & a_{C-1}...a_2 & \dots & 1 \end{bmatrix}\]

通过这个高效的编码器，scMamba 能够从原始的全量特征中提取出富有生物学意义的细胞表征。

2.3 增强的对比学习与多组学整合

对于多组学数据（如 scRNA-seq 和 scATAC-seq），分别为它们构建独立的 scMamba 编码器，并得到每个细胞的两种模态的嵌入向量 $x_i$ (RNA) 和 $y_i$ (ATAC)。接下来的目标是让这两个来自同一个细胞的嵌入向量在共享的潜在空间中尽可能地靠近。

为此作者设计了一种新颖的对比学习策略。

  1. 定义正负样本：
    • 正样本对 (Positive Pair): 来自同一个细胞的 RNA 嵌入 $x_i$ 和 ATAC 嵌入 $y_i$。
    • 负样本对 (Negative Pairs): 来自不同细胞的 RNA 和 ATAC 嵌入（如 $x_i$ 和 $y_j$, 其中 $i \neq j$）。
  2. 设计组合损失函数： 总损失函数 $L$ 由两部分构成：
    • 对比损失 (Contrastive Loss, $L_{con}$): 标准的 InfoNCE 损失。其目标是拉近正样本对的相似度，同时推开负样本对的相似度。
\[L_{con} = -\frac{1}{N} \left( \sum_{i} \log \frac{\exp(x_i^T y_i / \sigma)}{\sum_{j=1}^{N} \exp(x_i^T y_j / \sigma)} + \sum_{i} \log \frac{\exp(y_i^T x_i / \sigma)}{\sum_{j=1}^{N} \exp(y_i^T x_j / \sigma)} \right)\] \[L_{sim} = 1 - \frac{1}{N} \sum_{i} \cos(x_i, y_i)\]

最终的总损失函数为： \(L = \lambda_{con} \cdot L_{con} + \lambda_{sim} \cdot L_{sim}\) 通过优化这个目标函数，scMamba 能够学习到一个统一的、生物学意义明确的细胞嵌入空间。

3. 实验分析

作者通过一系列详尽的实验，将 scMamba 与 7 个当前最先进的（SOTA）整合方法（包括 scCLIP, CVQVAE, GLUE, SCALEX, scVI, Harmony, Scanorama）进行了全方位的比较。

3.1 实验一：多组学整合性能基准测试

作者使用三组公开的 scRNA+scATAC 数据集（SNARE-seq, 10x Multiome）进行评估。评估指标分为两类：生物变异保留（评估聚类准确度，如 ARI, NMI）和组学对齐（评估不同模态混合程度，如 OEMS, SAS）。

在所有数据集上，scMamba 的综合整合得分（生物变异保留和组学对齐的加权平均）始终排名第一，平均提升超过 10%。

UMAP 可视化结果显示，scMamba 能够在完美混合不同组学（图中左侧，不同颜色点均匀混合）的同时，清晰地分离不同的细胞类型（图中右侧，相同颜色点紧密聚集）。相比之下，其他方法或多或少存在组学分离（如 SCALEX）或细胞类型混淆（如 GLUE）的问题。

作者使用 FOSCTTM 指标（越低越好）来衡量模型能否精确地找到跨模态的同一个细胞。结果显示，scMamba 的匹配误差远低于其他所有方法，证明了其卓越的对齐精度。

3.2 实验二：图谱级别的大规模整合能力

作者使用了一个包含 37.7万个细胞 的人类胎儿发育图谱数据集，以测试模型的可扩展性。

随着数据规模增大，scMamba 始终保持着顶尖的整合性能。相比之下，一些方法（如 GLUE）虽然效果尚可，但计算开销巨大（耗时55小时，占用268GB内存）；另一些方法（如 SCALEX, scVI）则在规模增大后性能明显下降。

scMamba 的运行时间和内存占用随细胞数量近乎线性增长，表现出优异的计算效率。在处理 30 万细胞的数据时，scMamba 兼顾了高性能和高效率。

3.3 实验三：下游任务

⚪ 细胞类型注释

作者利用 scMamba 生成的细胞嵌入来训练一个简单的分类器，用于细胞类型注释。

在 CITE-seq 数据集上，scMamba 取得了极高的注释准确率（大部分细胞类型 > 0.9），甚至能够准确区分一些高度相似的 B 细胞亚型。scMamba 成功地利用一个多组学参考数据集，对一个纯 scRNA-seq 数据集进行了高精度注释，展现了其作为参考图谱的强大潜力。

在与 scGPT 和 CellPLM 等其他单细胞基础模型的直接比较中，scMamba 在准确率、F1-score 等所有分类指标上均取得了最优表现。

⚪ 细胞分化轨迹重建

作者使用一个包含人类造血分化过程的 BMMC 数据集，来评估 scMamba 重建细胞发育轨迹的能力。

在 scMamba 整合后的嵌入空间中，能够成功地重建出了一条从造血干细胞（HSC）到各类成熟细胞的、清晰连续的分化轨迹，而这在原始的单一组学数据中是模糊不清的。

通过一系列定量指标（如邻域重叠度、伪时间一致性、轨迹保守性得分），证明 scMamba 在整合过程中最大程度地保留了原始数据中蕴含的细胞分化动态信息，再次全面超越所有对比方法。