Multigrate: 单细胞多组学数据整合.
0. TL; DR
Multigrate 是一个基于多视角神经网络和专家乘积（Product of Experts, PoE）框架的无监督深度生成模型。Multigrate将每种模态的编码器视为一个专家，各自对细胞状态给出一个概率分布（后验）。通过将这些专家的后验分布进行乘积，模型能够计算出一个融合了所有可用模态信息的联合后验分布。
在多个真实的CITE-seq和ATAC-RNA数据集的基准测试中，Multigrate在数据整合和缺失模态插补两项任务上的性能均优于或持平于Seurat v4, totalVI, MOFA+等一系列SOTA方法。
1. 背景介绍
单细胞多组学技术可以同时窥探一个细胞的多个分子层面，如基因表达（RNA）、蛋白质丰度（Protein）和染色质可及性（ATAC）。然而，要将这些来自不同技术、具有不同数据特性完美地拼接起来，面临着三大严峻挑战：
- 数据的不完整性: 在真实的科研项目中几乎不可能为所有样本都获得所有模态的数据。通常的情况是，一部分细胞有RNA和蛋白数据，另一部分可能只有RNA数据。如何有效利用这些部分缺失（partially missing）的数据集进行整合，是现有方法的一大痛点。
- 方法的局限性: 现有的整合方法大多缺乏通用性。像 totalVI 这样的方法，专为CITE-seq（RNA+蛋白）数据设计，难以推广到其他模态组合；像 Seurat v4 和 MOFA+ 这样的方法虽然更通用，但它们或依赖线性模型（如CCA），难以捕捉复杂的非线性关系；或缺乏对缺失模态进行插补（imputation）的机制。
- 缺乏可扩展的参考图谱框架: 随着大型细胞图谱项目的推进，一个理想的整合框架不应只是一次性的整合，而应能构建一个可扩展、可更新的参考图谱。当有新的多组学查询数据产生时，我们希望能够快速地将其映射到这个参考图谱上，而不是每次都从头重新整合所有数据。目前，尚无方法能很好地实现这一点。
因此需要一个通用、灵活、可扩展的解决方案，它必须能够处理任意组合的、包含部分缺失模态的多组学数据；同时具备高质量的整合和插补能力；支持构建参考图谱，并能将新的查询数据高效地映射上来。
2. Multigrate 方法
Multigrate是一个基于多视角VAE的深度生成模型，其核心是专家乘积（PoE）框架和迁移学习。它通过一个统一的神经网络，实现了对异构、不完整多组学数据的整合、插补和参考图谱构建。
图a Multigrate模型结构
Multigrate将每个模态视为一个独立的视角，并为每个视角（模态$i$）构建一个VAE。
对于每个非缺失的模态$i$，其专属的编码器 $f_i$ 会将输入数据 $X_i$ 和批次信息 $S_i$ 编码为一个高斯分布的后验 $q_{\phi_i}(Z_i|X_i, S_i) = \mathcal{N}(Z_i | \mu_i, \sigma_i)$。
为了得到一个融合了所有可用模态信息的联合细胞表征，Multigrate采用了PoE框架来计算联合后验分布 $q_{\phi}(Z_{joint}|X, S)$。其思想非常直观：联合意见是所有专家意见的乘积。
\[q_{\phi}(Z_{joint}|X, S) = \prod_{i=1}^n q_{\phi_i}(Z_i|X_i, S_i)\]由于假设所有后验都是高斯分布，高斯分布的乘积仍然是高斯分布。因此，联合后验的均值 $\mu_{joint}$ 和方差 $\sigma_{joint}$ 可以通过一个简单的解析公式计算得出：
\[\mu_{joint} = \left(\sum_{i=0}^n m_i \mu_i \sigma_i^{-1}\right) \left(\sum_{i=0}^n m_i \sigma_i^{-1}\right)^{-1} \\ \sigma_{joint} = \left(\sum_{i=0}^n m_i \sigma_i^{-1}\right)^{-1}\]其中 $\mu_0, \sigma_0$ 是先验分布（标准正态分布）的参数。$m_i$ 是一个指示变量：如果模态$i$的数据存在，$m_i=1$；如果缺失，$m_i=0$。
PoE框架的巨大优势在于其处理缺失数据的灵活性。当某个模态缺失时，只需将其对应的 $m_i$ 设为0，它就自然地从联合后验的计算中被排除了，而模型依然可以利用其他所有可用的模态信息进行学习。
为了从联合嵌入 $Z_{joint}$ 中重构出原始数据，Multigrate设计了一个两级解码器：
- 首先，$Z_{joint}$ 被送入一个共享的解码器，这个解码器学习恢复出模态间的共性信息，并输出模态特异性的中间表征。
- 然后，每个模态的中间表征再被送入其专属的解码器 $g_i$，用于渲染出该模态的特有风格和特征，最终重构出原始数据 $\tilde{X}_i$。
Multigrate的训练目标是最小化一个由三部分组成的联合损失函数：
\[L_{multigrate} = \sum_{i=1}^d L_{AE_i} + \beta \sum_{i<j} L_{MMD}(Z_{joint_i}, Z_{joint_j})\]$L_{AE}$ 是每个数据集$i$的VAE损失，包含了重构损失和KL散度两部分。
\[L_{AE} = \alpha \mathbb{E}_{q_{\phi}}[\log p_{\theta}(X_i|Z_{joint}, S_i)] - \eta KL(q_{\phi}(Z_{joint}|X_i, S_i) \| p_{\theta}(Z_{joint}|S_i))\]重构损失确保潜在空间编码了足够的信息来恢复原始数据。对于RNA等计数数据，使用负二项损失；对于蛋白等连续数据，使用均方误差（MSE）损失。KL散度是一个正则项，它促使学习到的后验分布接近于先验分布（标准正态分布），保证了潜在空间的良好结构。
$L_{MMD}$ 是最大均值差异损失。为了对齐来自不同数据集（批次）的细胞，Multigrate在它们的联合潜在嵌入 $Z_{joint}$ 之间计算MMD损失。MMD通过计算两个分布在再生核希尔伯特空间中的距离，来衡量它们的差异。通过最小化MMD损失，可以有效地拉近不同批次细胞的分布，实现批次校正。
图b 迁移学习
Multigrate能够与scArches迁移学习框架的无缝结合。我们可以使用多个高质量的多组学数据集，训练一个稳定、全面的Multigrate模型，这个模型就构成了一个多组学参考图谱。当有新的查询多组学数据到来时，无需从头重新训练。利用scArches思想，冻结参考模型的大部分权重，只微调与新数据批次相关的少量条件参数。通过这个高效的微调过程，查询数据被快速、准确地映射到参考图谱的潜在空间中。
一旦映射完成，就可以利用参考图谱的注释，对查询细胞进行细胞类型注释；如果查询数据缺失了某个模态，可以利用训练好的解码器，从其在参考图谱中的联合嵌入出发，插补出缺失的模态数据。
3. 实验分析
作者在多个真实的PBMC（外周血单个核细胞）数据集上，对Multigrate的性能进行了全面的基准测试，涵盖了CITE-seq（RNA+蛋白）和10x Multiome（RNA+ATAC）等多种技术。
3.1 实验一：多组学数据整合
作者比较了Multigrate与MOFA+, Seurat v4, totalVI等SOTA方法的整合性能，评估指标包括生物学信息保留和批次校正两个方面。
在所有测试的数据集上，Multigrate的综合整合得分均名列前茅，与totalVI等为特定数据类型专门设计的顶尖方法性能相当，甚至略优（图c）。
UMAP图和定量指标都显示，Multigrate能够在高效去除批次效应（模态混合良好）和最大化保留生物学异质性（细胞类型分离清晰）之间取得出色的平衡（图a, b）。
3.2 实验二：多组学参考图谱
作者首先用三个健康人的多组学数据集（总计约16万细胞，涵盖RNA、蛋白、ATAC三种模态）构建了一个健康血液细胞参考图谱。然后将一个来自COVID-19患者的CITE-seq数据集作为查询数据，映射到这个参考图谱上。
UMAP图显示，查询数据被平滑且准确地整合到了参考图谱中。不同来源（参考、查询）、不同细胞类型和不同健康状况（健康、COVID-19）的细胞都得到了清晰的展示（图a-c）。
作者利用参考图谱的注释，对查询数据进行细胞类型预测，总体准确率达到了79%。对于主要的细胞类型，预测几乎是完美的。少数预测不佳的细胞类型（如ASDC, Treg），其原因在于它们在参考图谱中的细胞数量本身就极少（<100），难以提供足够的参考信息（图d）。
这个实验有力地证明了Multigrate在构建和利用大规模多组学参考图谱方面的强大能力。
3.3 实验三：缺失模态插补
作者设计了一个实验来测试Multigrate的跨模态插补能力。使用一个CITE-seq数据集，将其中的5000个细胞的蛋白质数据“隐藏”起来，只用它们的RNA数据和另外10000个完整细胞的数据来训练模型。然后用训练好的模型来为那5000个细胞插补出缺失的蛋白质数据。
Multigrate插补出的蛋白质丰度与真实的测量值高度一致。例如，对于T细胞marker CD3蛋白，其插补后的空间表达模式与真实模式几乎无法区分（图a）。
作者计算了所有14个蛋白的插补值与真实值之间的Pearson相关系数。结果显示，Multigrate的平均相关性略高于专门为CITE-seq插补设计的SOTA方法totalVI，以及功能强大的Seurat v4（图b）。
这证明了Multigrate不仅是一个优秀的整合工具，也是一个精准的跨模态数据插补器。
