Con-AAE：单细胞多组学对齐与整合的对比循环对抗自编码器.

0. TL; DR

Con-AAE (Contrastive cycle Adversarial Autoencoders)能够对齐和整合单细胞RNA-seq数据与单细胞ATAC-seq数据。Con-AAE的基础是两个耦合的自编码器，分别用于处理scRNA-seq和scATAC-seq数据。通过一个对抗性损失，模型被训练来学习两个既分离又无法被判别器区分的潜在表征，训练过程引入了循环一致性损失和自监督对比学习。

作者在多个模拟和真实的配对多组学数据集上，对Con-AAE进行了广泛的基准测试。结果表明，Con-AAE在细胞对齐和整合任务上的性能，显著优于Seurat, DCCA, CycleGAN, scJoint, MOFA+等一系列SOTA方法。

1. 背景介绍

单细胞多组学技术，特别是能同时测量scRNA-seq和scATAC-seq的配对（paired）技术，为我们全面理解细胞系统提供了前所未有的机会。然而，要真正释放这些数据的力量，我们必须解决两个核心的计算任务：

  1. 整合 (Integration): 如何将来自不同模态的数据融合在一起，得到一个统一的、能反映细胞真实状态的细胞类型划分。
  2. 对齐 (Alignment): 在不提供细胞配对信息的情况下，如何准确地找到scRNA-seq数据中的每个细胞，在scATAC-seq数据中对应的“另一半”。

这是一项极具挑战性的任务。单细胞数据固有的稀疏性、高维度和噪声，以及不同模态之间复杂的非线性关系，都为整合与对齐设置了重重障碍。现有的方法虽然众多，但各有其局限性：

为了在这些挑战中开辟出一条新路，作者提出了Con-AAE。它旨在通过一个集对抗自编码器、循环一致性和对比学习三者之长的强大框架，来解决单细胞多组学数据的对齐与整合问题。

2. Con-AAE 方法

Con-AAE的核心思想是利用两个耦合的对抗自编码器（Adversarial Autoencoders, AAE），并将它们置于一个包含循环一致性损失和对比损失的统一框架下进行训练。

Con-AAE它由两套独立的编码器-解码器（encoder-decoder）对和 一个共享的判别器（discriminator）组成。作者为scRNA-seq和scATAC-seq数据分别设计了一个自编码器。

相应的解码器 $D_j$ 和 $D_i$ 则负责从潜在空间重构回原始数据。模型的重构损失 $L_{recon}$ 确保了潜在表征能够保留原始数据的主要信息。

\[L_{recon} = \mathbb{E}_{x \sim p_{rna}} [d(x, D_j(E_j(x)))] + \mathbb{E}_{x \sim p_{atac}} [d(x, D_i(E_i(x)))]\]

为了使两个潜在空间 $Z_i$ 和 $Z_j$ 对齐，作者引入了一个判别器 $D$。判别器 $D$ 的任务是区分一个潜在向量是来自RNA编码器（$Z_j$）还是ATAC编码器（$Z_i$）。在训练过程中，判别器 $D$ 努力最大化分类准确率，而两个编码器 $E_i, E_j$ 则努力生成让判别器 $D$ 无法区分的嵌入。这通过一个对抗性损失 $L_{adv}$ 来实现。

\[L_{adv} = \mathbb{E}_{x \sim p_{rna}}[\log D(E_j(x))] + \mathbb{E}_{x \sim p_{atac}}[\log(1-D(E_i(x)))]\]

通过这个对抗过程，两个潜在空间被强迫对齐到一个协调的子空间（coordinated subspace）中。

单纯的AAE可能会导致模式坍塌，并且无法保证学习到有意义的跨模态映射。为了解决这个问题，作者引入了循环一致性损失，其思想借鉴于著名的CycleGAN。一个从RNA空间编码得到的向量，如果先经过ATAC解码器“翻译”成ATAC数据，再通过ATAC编码器翻译回潜在空间，它最终应该回到原来的位置。反之亦然。作者定义了循环一致性损失 $L_{cyc}$，来惩罚这种循环翻译前后的差异。

\[L_{cyc} = \mathbb{E}_{x \sim p_{rna}}[d(E_j(x), E_i(D_i(E_j(x))))] + \mathbb{E}_{x \sim p_{atac}}[d(E_i(x), E_j(D_j(E_i(x))))]\]

这个损失项强制编码器和解码器学习到结构上一致的、互为逆映射的跨模态转换，极大地增强了模型的稳定性和跨模态生成能力。

为了进一步提升模型的鲁棒性，特别是在处理带噪声的无监督对齐任务时，作者引入了自监督对比学习。

正样本对来自同一个细胞类型的任意两个细胞的嵌入。为了增加难度，作者特别定义了难正样本（hard positive），即在同一个簇中，与当前锚点距离最远的那个样本。负样本对来自不同细胞类型的任意两个细胞的嵌入。同样，作者定义了难负样本（hard negative），即在所有不同簇中，与当前锚点距离最近的那个样本。

对比学习的目标是，在潜在空间中，拉近正样本对，推开负样本对。这通过一个三元组损失（triplet loss）来实现：

\(L_{con} = \mathbb{E}_{z_a \sim Z}[d(z_a, z_p) - d(z_a, z_n) + \alpha]\) 其中，$z_a$ 是锚点，$z_p$ 是难正样本，$z_n$ 是难负样本，$\alpha$ 是一个预设的边界（margin）。这个损失项迫使模型学习一个类内距离小、类间距离大的嵌入空间，从而使得聚类边界更清晰，对噪声更鲁棒。

为了在有监督的整合任务中充分利用标签信息，作者还在协调子空间之上添加了一个简单的分类器 $C$，并使用交叉熵损失 $L_{classifier}$ 进行训练。

Con-AAE通过一个两阶段的训练方案，联合优化上述所有损失项。在每个迭代中，模型交替地更新判别器和自编码器-分类器部分，最终达到一个纳什均衡。

3. 实验分析

作者在24个不同大小和信噪比的模拟数据集，以及四个来自sci-CAR, SNARE-seq, 10X PBMC和SHARE-seq平台的真实配对多组学数据集上，对Con-AAE的性能进行了全面的评估。

3.1 实验一：模拟数据集评估

在整合任务（即细胞类型分类）中，作者测试了不同数据规模和信噪比（SNR）下的性能。结果显示，Con-AAE的性能在绝大多数情况下都超过了所有其他方法。特别是在低SNR和大规模数据这种最具挑战性的场景下，其他方法的性能急剧下降，而Con-AAE依然保持了优异的性能，展现了其强大的鲁棒性和可扩展性。

在对齐任务（通过recall@k评估）中，Con-AAE的性能同样全面领先。其箱线图的中位数、上下四分位数和边缘在大多数情况下都高于其他工具，表明其整体对齐精度更高且更稳定。

3.2 实验二：真实世界数据集评估

作者在四个真实的配对数据集上比较了Con-AAE与其他SOTA方法的性能。