0. TL; DR

JAMIE (Joint variational Autoencoders for Multimodal Imputation and Embedding)是一个专为处理部分配对单细胞多组学数据而设计的深度学习框架。它通过关联性潜在空间聚合（correlation-based latent aggregation）机制，灵活地利用任何可用的细胞配对信息，来指导不同模态潜在空间的对齐。JAMIE为每个模态学习一个既相互关联又保持独立的潜在空间，其编码器-解码器架构天然支持跨模态插补。

在多个模拟和真实的单细胞多组学数据集上（涵盖基因表达、染色质可及性、电生理学等），JAMIE在数据整合和跨模态插补两项任务上的性能均显著优于现有的SOTA方法。

1. 背景介绍

对基因表达（RNA）和染色质开放状态（ATAC）的配队场景中，经常遇到棘手的局面：

部分配对 (Partially Matched): 只有一小部分细胞是同时测了RNA和ATAC的，大部分细胞都只有其中一种信息。
完全不配对 (Unpaired): 有两批细胞，一批测了RNA，另一批测了ATAC，它们来自同一个组织，但完全不知道细胞间的一一对应关系。

现有的一些无监督对齐方法（如UnionCom, MMD-MA）虽然不要求配对，但它们在对齐流形时，可能会忽略模态特异性的信息，或者在数据分布复杂时对齐失败。一些方法为了实现整合，倾向于将所有模态都压缩到一个单一的、共享的潜在空间中。这虽然能实现对齐，但代价是牺牲了每个模态独有的、宝贵的特异性信息。

JAMIE是一个更灵活、更强大的框架，它有效利用任何可用的配对信息，在整合的同时保留每个模态的特异性，还能精准地插补出缺失的模态信息。

2. JAMIE 模型

JAMIE是一个基于双变分自编码器 (Dual VAEs)的架构，为每种待整合的模态（假设为$X$和$Y$）都构建一个独立的变分自编码器。

编码器 $f_E^X$ 将输入数据 $X$ 编码为一个连续的潜在空间分布，通常是高斯分布。它输出每个细胞 $i$ 的潜在表征的均值 $\mu_i^X$ 和对数方差 $\log(\sigma_i^X)^2$。然后通过重参数化技巧从这个分布中采样得到潜在向量 $L_i^X \sim \mathcal{N}(\mu_i^X, \text{diag}(\sigma_i^X))$。对模态Y的编码器 $f_E^Y$ 也进行同样的操作，得到 $L_j^Y$。编码器为每个模态保留了一个独立的、特异性的潜在空间（$L^X$ 和 $L^Y$），而不是强行将它们压缩到同一个空间。

解码器 $f_D^X$ 和 $f_D^Y$ 负责从聚合后的潜在空间中，将向量重构回原始的高维数据空间，得到 $\tilde{X}$ 和 $\tilde{Y}$。

JAMIE通过关联性潜在空间聚合 (Correlation-based Latent Aggregation) 将两个独立的潜在空间软性地连接起来。模型的输入除了数据矩阵 $X$ 和 $Y$，还有一个可选的配对矩阵 $F$ ($F \in \mathbb{R}^{n_X \times n_Y}$)。$F$ 编码了细胞间的配对信息：如果第 $i$ 个X细胞和第 $j$ 个Y细胞是同一个细胞，$F_{ij}=1$；如果它们不是同一个细胞，$F_{ij}=0$；$F$ 也可以是0到1之间的概率值，表示配对的置信度。

在得到模态特异性的潜在空间 $L^X$ 和 $L^Y$ 后，通过以下公式计算聚合后的潜在空间 $M^X$ 和 $M^Y$：

\[M^X = (L^X + \delta L^Y F^T) (\text{diag}(1 + \delta F \mathbf{1}_{n_Y}))^{-1} \\ M^Y = (L^Y + \delta^{-1} L^X F) (\text{diag}(1 + \delta^{-1} F^T \mathbf{1}_{n_X}))^{-1}\]

这个公式表明，对于一个X模态的细胞 $i$，其聚合后的嵌入 $M_i^X$，是它自身的嵌入 $L_i^X$ 与所有与它配对的Y模态细胞的嵌入 $L_j^Y$ 的加权平均。$F$ 矩阵决定了加权的模式。$\delta$ 是一个可学习的参数，它自动平衡了两种模态在聚合过程中的相对权重。

通过这种方式，配对信息将两个独立的潜在空间中对应的细胞点拉近，从而实现了对齐。而对于那些没有配对信息的细胞，它们的聚合嵌入就约等于其自身嵌入，其对齐则依赖于与其他已配对细胞的流形关系。

JAMIE的训练目标是最小化一个由四部分组成的联合损失函数：

\[\ell_{total} = \ell_{kl} + \alpha\ell_{combination} + \beta\ell_{reconstruction} + \gamma\ell_{alignment}\]

$\ell_{kl}$ (KL散度损失): 这是标准VAE的损失项，它促使编码器生成的潜在分布趋向于标准正态分布，保证了潜在空间的连续性和良好的结构。作者使用KL退火（KL annealing）策略，在训练初期给KL损失一个较小的权重，让模型先专注于重构，再逐渐加强对潜在空间结构的约束。
$\ell_{combination}$ (组合损失): 这个损失项要求聚合后的潜在空间 $M$ 与原始的潜在空间 $L$ 尽可能相似。它起到一个正则化的作用，防止聚合过程过度扭曲每个模态原有的结构。 $\ell_{combination} = \|L^X - M^X\|_F^2 + \|L^Y - M^Y\|_F^2$
$\ell_{reconstruction}$ (重构损失): 这是标准的自编码器损失，即最小化重构数据与原始数据之间的均方误差（MSE）。 $\ell_{reconstruction} = \|X - \tilde{X}\|_F^2 + \|Y - \tilde{Y}\|_F^2$
$\ell_{alignment}$ (对齐损失): 这是一个可选的损失项，在配对信息非常稀疏时特别有用。它会首先通过无监督流形对齐的方法，推断一个更密集的伪配对矩阵 $\tilde{F}$，然后强制要求聚合后的两个潜在空间 $M^X$ 和 $M^Y$ 在这个 $\tilde{F}$ 的指导下更加接近。 $\ell_{alignment} = \|M^X - M^Y \tilde{F}^T\|_F^2$

通过联合优化这四个损失，JAMIE得以在保留模态特异性的前提下，利用部分配对信息高效地实现潜在空间的对齐。

3. 实验分析

作者在多个模拟数据集和真实的、包含多种模态（基因表达、染色质可及性、电生理学）的单细胞数据集上，对JAMIE的性能进行了全面评估。

3.1 实验一：数据整合

作者评估了JAMIE学习到的联合嵌入空间在对齐不同模态和分离细胞类型方面的能力。

在所有测试的数据集上（模拟数据、小鼠大脑Patch-seq、人脑Multiome），JAMIE在标签迁移准确率（LTA）和真实配对距离（FOSCTTM）两个核心指标上，均显著优于或持平于NLMA, UnionCom, CCA等一系列SOTA**对齐方法。

UMAP可视化清晰地显示，JAMIE能够生成高质量的联合嵌入。不同模态的数据被很好地混合，同时不同的细胞类型形成了清晰、分离的簇。

3.2 实验二：跨模态插补

作者利用训练好的JAMIE模型，进行跨模态的特征预测，并与基线的KNN插补方法以及专门的插补工具（如BABEL）进行比较。

在所有测试的数据集上，JAMIE的跨模态插补性能（通过AUROC和Pearson相关性评估）都显著优于基线的KNN方法和SOTA的BABEL模型。例如，在人脑数据上，JAMIE在超过95%的基因上都比KNN插补得更准。

JAMIE的插补结果能够很好地保留原始数据中细胞类型特异性的特征表达模式。例如，在模拟数据中，插补后的特征分布与真实分布高度相似（JS距离小）。

3.3 实验三：可解释性

作者使用SHAP来解释JAMIE的插补模型，以探究哪些源模态特征对目标特征的预测贡献最大。

在从ATAC-seq插补基因表达的任务中，SHAP分析显示，对于目标基因（如DENND1B），其基因体附近（如10kb内）的ATAC peaks具有最高的贡献度。这与已知的顺式调控生物学原理完全吻合，证明了JAMIE模型确实学习到了有意义的调控逻辑，而不仅仅是黑箱式的预测。

在对齐任务中，SHAP分析识别出对区分细胞类型贡献最大的基因。例如，在小鼠大脑数据中，JAMIE识别出的SST基因是区分不同神经元亚型的关键，这与已知的生物学事实相符。