通过scVAEIT实现单细胞多模态混合整合与插补的鲁棒概率建模.

0. TL; DR

scVAEIT是一个用于马赛克式测量数据的多模态整合与插补的概率性变分自编码器模型。它引入了一个缺失掩码（missing mask），在训练过程中随机地遮盖掉一部分观测到的特征。模型被训练来根据未被遮盖的特征，去预测被遮盖的特征。

scVAEIT在插补缺失模态和特征方面表现出极高的准确性，并且能够鲁棒地处理与训练数据生物学上不同的新细胞类型。它在调整批次效应的同时，能很好地保留生物学变异，从而为整合后的数据集提供了更高质量的潜在表征。

1. 背景介绍

随着单细胞多组学技术的发展，我们手中的数据集变得越来越多样和复杂。一个日益普遍的场景是马赛克整合（mosaic integration）：作者希望整合多个数据集，但这些数据集的测量面板并不完全相同。例如可能有：

在这种马赛克式的场景下，数据整合的目标也变得更加复杂和宏大：

  1. 构建统一的潜在空间： 如何利用所有共享和非共享的特征，为所有细胞学习一个统一的、高质量的低维表示？
  2. 精准插补缺失模态/特征： 如何为CITE-seq数据补上缺失的ATAC信息？如何为DOGMA-seq中某个细胞碰巧未测到的某个蛋白质，进行准确的填空？
  3. 鲁棒的迁移学习： 如何将从一个大规模参考图谱上学到的知识，迁移到一个只测量了部分特征的新数据集上？

现有的整合方法在面对马赛克世界时，往往捉襟见肘：

为了攻克马赛克整合的难题，作者需要一个能够显式地建模和利用缺失性信息的全新框架。scVAEIT正是为此而生。

2. scVAEIT 方法

scVAEIT的核心思想是，通过引入一个随机掩码（missing mask），将无监督的VAE重构任务，巧妙地转化为一个有监督的条件生成（conditional generation）任务，从而学习到特征与模态间复杂的依赖关系。

scVAEIT的基础是一个变分自编码器。对于一个细胞$n$的完整多组学数据 $X_n$，scVAEIT首先引入一个二进制的掩码向量 $M_n$。$M_n$ 与 $X_n$ 维度相同。如果$M_n$的第$j$个元素为1，表示特征$j$是可见的（observed）；如果为0，表示是被遮盖的（masked）。这样原始数据被分为了两部分：$X_{M_n^c}$（可见部分）和 $X_{M_n}$（被遮盖部分）。

scVAEIT的编码器接收可见数据 $X_{M_n^c}$ 和掩码 $M_n$ 作为输入。它的任务是学习一个条件后验分布 $q_\psi(Z_n | X_n, M_n)$，即在给定可见数据和缺失模式的条件下，推断出细胞的低维潜在表征 $Z_n$。

解码器则接收潜变量 $Z_n$、可见数据 $X_{M_n^c}$ 和掩码 $M_n$ 作为输入。它的任务是预测被遮盖的数据 $X_{M_n}$ 的条件分布 $p_\theta(X_{M_n} | Z_n, X_{M_n^c}, M_n)$。

真实掩码 ($M_a$) 代表数据中真实存在的缺失模式（例如，CITE-seq数据天然缺失ATAC模态）。在训练过程中，即使对于一个完整的数据点，作者也随机生成一个掩码 $M$，人为地遮盖掉一部分特征。随机掩码的优势：

模型的训练目标是最大化一个加权的证据下界（ELBO）。这个ELBO巧妙地结合了插补（imputation）和重构（reconstruction）两个任务。

\[L := \beta L_M + (1-\beta) L_{recon}\]

插补损失 ($L_M$) 最大化模型预测被遮盖部分的对数似然。这相当于一个有监督的学习任务，模型必须学会根据可见部分来推断不可见部分。

\[L_M \ge \mathbb{E}_{q_\psi(Z|X,M)}[\log p_{\theta_2}(X_M | Z, X_{M^c}, M)] - KL(q_\psi(Z|X,M) \| p_{\theta_1}(Z|X_{M^c}, M))\]

重构损失 ($L_{recon}$) 最大化模型重构可见部分的对数似然。这是一个标准的VAE去噪任务。

scVAEIT为不同模态定制了不同的生成分布，如为RNA和蛋白质使用负二项分布，为ATAC-seq使用伯努利分布。

为了降低计算复杂度，scVAEIT在编码器和解码器的第一层和最后一层采用了子连接（subconnections），即不同模态（甚至不同染色体上的peaks）在进入网络时，先经过独立的子网络处理，再进行整合。

3. 实验分析

作者通过一系列精心设计的实验，全面评估了scVAEIT在跨模态翻译、三模态整合和马赛克整合等任务中的性能。

3.1 实验一：跨模态翻译与迁移学习

作者在一个CITE-seq PBMC数据集上，通过留一细胞类型的交叉验证，测试了模型的跨模态翻译能力，并进一步测试了其向两个外部数据集（一个来自不同组织CBMC，一个来自不同技术REAP-seq）的迁移学习能力。

3.2 实验二：三模态整合与插补

作者在一个同时测量了RNA、蛋白质和ATAC的DOGMA-seq数据集上，评估了scVAEIT处理三种模态的能力。

3.3 实验三：对缺失特征的鲁棒性

作者通过在测试时随机移除不同比例的输入特征，来评估模型在面对模型不匹配（model misspecification）时的鲁棒性。

结果显示，随着缺失特征比例的增加，totalVI和MultiVI的性能急剧下降，稳定性也变差。相比之下，scVAEIT的性能则异常稳定，几乎不受缺失特征的影响。Seurat作为非参数方法，也表现出较好的稳定性，但其整体精度低于scVAEIT。

这充分证明了scVAEIT的随机掩码训练策略为其带来了鲁棒性，使其能够从容应对各种不完整的输入数据。

3.4 实验四：多源、多模态的马赛克整合

作者整合了三个来源不同、模态组合也不同的PBMC数据集（DOGMA-seq, CITE-seq, ASAP-seq）。