用于多模态细胞数据集成和插补的模态联结自编码器.

0. TL; DR

Monae (Modal-Nexus Auto-Encoder)是一个能够同时实现高质量数据整合与精准插补的单细胞多组学框架。Monae从两个互补的视角来学习细胞表征:

  1. 图视角 (Graph-linked Embedding): 它构建一个包含所有模态特征的调控关系图,并使用图变分自编码器(Graph VAE)来学习每个特征(基因/peak)的嵌入,捕捉它们之间的模态枢纽关系。
  2. 细胞视角 (Contrastive Cell Embedding): 它为每个模态设计了非对称的自编码器网络,并结合对比学习和自适应聚类,学习一个能最大化区分不同细胞身份的联合细胞嵌入。

在多个大规模、无配对的基准数据集上,Monae在整合和插补两个任务上均取得了与领域专用SOTA方法相当甚至更优的性能,并展现了处理百万级别细胞图谱的强大可扩展性。

1. 背景介绍

随着单细胞技术的发展,我们获得了海量的多组学数据。然而,这些数据往往是不完美的,带来了两大挑战:

  1. 整合: 我们常常需要整合来自不同技术平台的数据,比如scRNA-seqscATAC-seq。它们不仅特征空间完全不同(一个是基因,一个是基因组区域),而且往往是无配对的,即我们不知道哪个RNA细胞对应哪个ATAC细胞。如何跨越这道鸿沟,将它们对齐到一个统一的空间,是数据整合的核心难题。
  2. 插补: 单细胞数据,尤其是scRNA-seq,存在严重的dropout现象——很多本应表达的基因,因为捕获效率低而被错误地记为零。这种技术噪声极大地干扰了我们对细胞真实状态的理解。因此,需要通过数据插补来修复这些缺失的信号,这个过程也常被称为去噪。此外,我们还希望能够进行跨模态插补,即利用一种模态的数据来预测另一种模态,这在缺乏配对数据时尤其有用。

作者认为整合与插补并非孤立的任务。一个好的整合模型,必然深刻理解了细胞的身份和模态间的关联,这种理解正是进行高质量插补的基础。反之,一个能利用多模态信息进行精准插补的模型,其内部必然已经形成了一个对齐良好的细胞表征。Monae是一个旨在通过统一的框架,从根本上协同解决整合与插补两大难题的全新方法。

2. Monae 模型

Monae 是一个无监督的多自编码器框架,其核心在于从图和细胞两个视角进行联合学习,最终实现整合与插补的一体化。

2.1 图链接嵌入学习 (Graph-linked Embedding Learning)

图链接嵌入学习模块的目标是学习一个能捕捉跨模态特征调控关系的嵌入空间。

首先构建一个粗略的、基于先验知识的调控关系图 (Guidance Graph)$G$。图中的每个节点代表一个特征(如一个基因、一个ATAC peak),边代表特征之间的调控关系。例如,根据基因组坐标,可以连接一个基因和它附近的ATAC peak。边的符号可以表示调控的正负效应(如ATAC peak对基因表达通常是正向调控,而甲基化通常是负向调控)。

使用一个图变分自编码器 (Graph VAE)来学习这个调控关系图。编码器是一个图注意力网络 (GAT)GAT通过消息传递机制,让每个特征节点能够聚合其邻居特征(可能来自不同模态)的信息。经过多层GAT,每个特征节点的嵌入 $f_i \in \mathbb{R}^m$ 就编码了其在整个调控网络中的上下文信息。编码器输出每个节点嵌入的均值 $\mu_i$ 和方差 $\sigma_i^2$。

\[q(F|G) = \prod_{i=1}^{|V|} \mathcal{N}(f_i | \text{GAT}_{\mu_i}(G), \text{GAT}_{\sigma_i^2}(G))\]

解码器通过计算任意两个特征嵌入 $f_i$ 和 $f_j$ 的内积,来预测它们之间存在边的概率,从而重构出一个更精细的、密集的调控图 $\hat{G}$。

\[p(\hat{G}|F) = \prod_{i,j} \text{sigmoid}(f_i \cdot f_j^T) = \prod_{i,j} \hat{w}_{ij}\]

通过最小化图重构误差和KL散度,Graph VAE学习到了每个特征的图链接嵌入 $F \in \mathbb{R}^{|V| \times m}$。这个矩阵 $F$ 包含了所有特征在“模态枢纽”中的坐标。

2.2 对比细胞嵌入学习 (Contrastive Cell Embedding Learning)

对比细胞嵌入学习模块的目标是学习一个能最大化细胞身份区分度的联合嵌入空间。

对于每个模态,设计一个非对称的双分支网络(Teacher-Student架构):一个细胞的原始表达谱 $x^{(n)}$ 分别送入Teacher分支和Student分支。Teacher分支和Student分支的网络结构略有不同(如Student分支有额外的Dropout),它们会输出两个略有差异的嵌入向量 $z_{teacher}$ 和 $z_{student}$。这对 $(z_{teacher}, z_{student})$ 就构成了一个正样本对。而batch内所有其他细胞的嵌入,则都构成负样本。

采用标准的对比学习 (Contrastive Learning) InfoNCE损失,目标是拉近正样本对的距离,同时推开所有负样本对。

\[L_{contrastive} = - \frac{1}{2N_{batch}} \sum_{j=1}^{N_{batch}} \log \frac{e^{s_{j, j+}}}{\sum_{r \neq j} (e^{s_{j+, r}} + e^{s_{j, r}})}\]

其中,$s_{j, j+}$ 是正样本对的余弦相似度,$s_{j, r}$ 是与负样本的相似度。

为了进一步增强细胞类型的分离度,Monae引入了一个自适应聚类增强 (Adaptive Clustering)损失。它在对比学习的潜在空间上,动态地进行聚类,并施加一个损失项来惩罚不同簇之间的重叠,使得簇间分离度更大。

\[L_{cluster} = \frac{1}{C} \sum_{a=1}^{C-1} \sum_{b>a} \frac{Y_a S (Y_b)^T}{\sqrt{Y_a S (Y_a)^T \cdot Y_b S (Y_b)^T}}\]

其中 $Y$ 是细胞的软聚类分配概率矩阵,$S$ 是细胞间的相似度矩阵。

通过优化对比损失和聚类损失,模型学习到了每个细胞的对比细胞嵌入 $U_k \in \mathbb{R}^{N_k \times m}$(对于第k个模态)。这个 $U$ 矩阵就是最终用于数据整合的高质量细胞表征。

对比细胞嵌入 $U$ 直接作为最终的整合结果,用于下游的聚类、可视化等任务。通过将对比细胞嵌入 $U$ 和图链接特征嵌入 $F$ 进行矩阵乘法,来生成插补后的表达谱:

这个操作的直觉是:一个细胞(由$U$中的行向量代表)在某个基因(由$F$中的行向量代表)上的表达量,是该细胞的身份状态与该基因在调控网络中的角色相互作用的结果。模型的插补损失 $L_{imputation}$ 就是最小化重构矩阵 $\hat{X}_k$ 与原始矩阵 $X_k$ 之间的差异。

2.3 Monae-E

作者还提出了Monae-E,它在Monae的基础上,直接将细胞节点和TF基序(motif)节点也加入到初始的调控关系图中。这使得Graph VAE能够一步到位地学习到所有实体(细胞、基因、peak、motif)的联合嵌入。这带来了两大好处:

  1. 加速收敛: 丰富的先验知识使得模型训练更快。
  2. 增强可解释性: 可以直接通过计算学习到的TF motif嵌入和基因嵌入之间的距离,来推断TF-靶基因调控网络。

3. 实验分析

作者在多个具有挑战性的大规模数据集上,对MonaeMonae-E的性能进行了全面评估。

3.1 实验一:多模态整合

作者在完全无配对的Muto-2021(人肾脏)数据集和包含批次效应的Ma-2020(小鼠皮肤)数据集上进行了测试。

在衡量整合效果的综合得分上,MonaeMonae-E的表现接近甚至优于GLUE、CoVEL等专门为整合设计的SOTA方法,并远超scButterfly、MultiVI等也试图同时进行整合和插补的多任务方法。

UMAP可视化显示,MonaeMonae-E不仅能很好地混合不同模态,还能非常清晰地分离开不同的细胞类型,生物学结构保留得非常好。

3.2 实验二:模内与跨模态插补

作者评估了Monae在修复dropout(模内插补)和从一种模态预测另一种模态(跨模态插补)上的能力。

通过比较插补前后数据的PCA-UMAP图,可以明显看到,经过Monae插补后,原本因dropout而混杂不清的细胞簇变得界限分明。在定量的聚类指标(ARI, NMI等)上,Monae的性能也优于MAGIC、DCA等专门的单模态插补方法。

在从RNA预测ATAC,以及从ATAC预测RNA的任务中,MonaeMonae-E的性能均名列前茅,与scButterflySOTA方法不相上下。

通过比较预测基因表达的热图,作者发现Monae相比scButterfly,能更好地恢复稀有细胞类型(如Treg, gdT)的特异性表达模式,细节保留得更好。

进一步使用高质量的MultiVI插补结果作为“伪金标准”,来评估MonaescButterfly的预测准确性。结果显示,Monae预测的细胞和基因表达谱与伪金标准的PCC(皮尔逊相关系数)更高,证明其插补结果更逼近真实情况Q。

3.3 实验三:Monae-E的可解释性与生物学发现

作者利用Monae-EMuto-2021肾脏数据集上进行探索性分析。

Monae-E学习到的细胞-基因联合嵌入空间中,观察到细胞类型与其对应的marker基因在空间上是邻近的,这直观地验证了模型学习到了正确的生物学关联。

基于学习到的特征嵌入距离,作者成功地为转录因子HNF4A预测了40个靶基因,其中17个得到了现有文献的支持。对于未被报道的预测,其在RNAATAC两个模态的表达模式也与HNF4A高度一致,为实验验证提供了可靠的候选。

Monae-E不仅能通过特征嵌入距离找到VCAM1基因与其上游~60kb处的已知增强子peak,还能基于插补后的数据,清晰地重建出肾脏近端小管细胞(PT)向受损修复相关的PT_VCAM1细胞分化的伪时间轨迹。