0. TL; DR

CVQVAE (Cross Vector-Quantized Variational Autoencoder)是一个简洁而强大的多组学整合模型。它巧妙地结合了两种深度学习技术：交叉训练 (Cross-training) 和矢量量化 (Vector Quantization, VQ)。交叉训练是指用RNA编码器提取的特征，去驱动ATAC解码器重构ATAC数据，强制模型学习一个能被两种模态共同理解的、高度对齐的共享潜在空间。矢量量化将连续的潜在空间离散化为一个码本，解决了后验坍塌问题，使得训练过程更稳定，而且也更符合单细胞数据本质上是离散计数的生物学特性。

在NeurIPS竞赛的基准数据集上，CVQVAE在多个关键指标（如NMI、细胞周期保留、图连通性等）上取得了SOTA性能，总体得分超越了包括竞赛冠军CLUE在内的多种方法。相比于结构复杂的模型（如CLUE），CVQVAE在取得更高性能的同时，训练时间大幅缩短（在4万细胞的数据集上，从1个多小时缩短到10分钟），展现了极高的计算效率。

1. 背景介绍

单细胞多组学技术为我们打开了一扇前所未有的窗户，让我们能同时窥探一个细胞内多个分子层面的动态，如基因的表达（RNA）、染色质的开放状态（ATAC）和蛋白的丰度（ADT）。理想情况下，我们希望将这些来自不同“视角”的信息整合到一个统一的框架中，以获得对细胞状态的完整认知。一个理想的整合模型应该能生成一个共享的潜在空间，在这个空间里：

来自同一个细胞的不同模态的数据点应该彼此靠近。
原始数据中蕴含的生物学结构（如细胞类型聚类、细胞分化轨迹）应该被完好地保留下来。

然而，实现这一目标困难重重。现有的一些深度学习方法虽然强大，但依赖先验知识，预处理复杂，模型复杂且训练不稳定。一些基于VAE的模型，虽然不依赖强先验，但可能会面临训练不稳定的问题，例如后验坍塌 (posterior collapse)，解码器直接忽略了编码器传来的信息，导致潜在空间学不到任何有意义的表征。

因此亟需一个简单、高效、训练稳定且不依赖强先验知识的整合模型。

2. CVQVAE 模型

CVQVAE 由两个关键组件构成：交叉训练的自编码器和矢量量化层。

2.1 交叉训练 (Cross-Training)

假设有两种配对的模态数据，S1（如RNA）和S2（如ATAC）。传统的双模态VAE会为每个模态训练一个独立的VAE（Encoder1-Decoder1 和 Encoder2-Decoder2）。

CVQVAE则打破了这种常规，采用了交叉训练的策略：S1的数据经过 Encoder1 编码成潜在表征 $z_1$，S2的数据经过 Encoder2 编码成潜在表征 $z_2$。

此时 $z_1$ 不再送入Decoder1，而是送入 Decoder2 去尝试重构S2的数据（ $\hat{x}_{S2}$ ）；同理，$z_2$ 送入 Decoder1 去重构S1的数据（ $\hat{x}_{S1}$ ）。

这种训练迫使模型必须学习到一个通用的、跨模态的特征。Encoder1必须学习如何从S1数据中提取出对重构S2有用的信息，反之亦然。这自然而然地将两个模态的潜在空间推向对齐。

2.2 矢量量化 (Vector Quantization, VQ)

传统的VAE的潜在空间是连续的，编码器输出的是一个高斯分布的均值和方差。而CVQVAE则引入了VQ-VAE的思想，将潜在空间离散化。

模型维护一个可学习的码本（embedding space） $K$，它由一系列离散的编码向量（codewords）组成。当编码器输出一个连续的潜在向量 $z$ 后，它不会被直接送入解码器。取而代之的是，模型会在码本$K$中寻找与$z$最接近的那个编码向量 $e_z$ (nearest neighbor)。这个$e_z$才是解码器的真正输入。由于“寻找最近邻”这个操作是不可导的，使用直通估计器 (straight-through estimator) 的技巧，直接将 $ez$ 的梯度复制给 $z$，以实现端到端的训练。码本$K$本身通过指数移动平均（EMA）进行更新，不断适应训练数据的分布。

在传统VAE中，如果解码器能力过强，它可能会完全忽略编码器传来的连续、带噪声的$z$，导致KL散度项很快变为0，这就是后验坍塌。而在VQ-VAE中，解码器接收的是一个离散、确定的$e_z$，它无法“绕过”这个信息。理论上可以证明，VQ-VAE能够有效避免后验坍塌问题。此外单细胞测序数据（如UMI计数）本质上是离散的，使用离散的潜在空间来表征它们，可能更符合其生物学本质。

2.3 组合损失函数

CVQVAE的最终损失函数由三部分组成，权重由 $\lambda_1, \lambda_2, \lambda_3$ 控制：

\[L_{CVQVAE} = \lambda_1 \cdot L_1 + \lambda_2 \cdot L_2 + \lambda_3 \cdot L_3\]

其中$L_1$ 和 $L_2$ 是交叉的VQ-VAE损失：

$L_1 = L(\text{E}_1, \text{D}_2, \text{K}_1)$：用S1的数据，通过E1编码、D2解码重构S2的损失。
$L_2 = L(\text{E}_2, \text{D}_1, \text{K}_2)$：用S2的数据，通过E2编码、D1解码重构S1的损失。

每个VQ-VAE损失 $L(\cdot)$ 本身又包含三项：

\[L(E, D, K) = \underbrace{\log p(x|\mathbf{e_z}(x))}_{\text{重构损失}} + \underbrace{||\text{sg}[z(x)] - \mathbf{e_z}||_2^2}_{\text{VQ损失}} + \beta \underbrace{||z(x) - \text{sg}[\mathbf{e_z}]||_2^2}_{\text{承诺损失}}\]

重构损失衡量解码器输出与原始数据（目标模态的）的差距；VQ损失 (Codebook loss)促使码本中的编码向量$e_z$向编码器输出$z$靠拢；承诺损失 (Commitment loss)促使编码器输出$z$向其选择的码本向量$e_z$靠拢，防止$z$离$e_z$太远。$sg$是停止梯度操作。

$L_3$ 是潜在空间对齐损失，是一个简单的均方误差损失 (MSE)，用于直接拉近来自同一个细胞的两个潜在表征 $z_1$ 和 $z_2$ 的距离。

\[L_3 = \text{Loss}(z_1, z_2) = ||z_1 - z_2||^2\]

通过优化这个精心设计的组合损失函数，CVQVAE能够在没有额外监督信息的情况下，高效地学习到一个对齐良好且信息丰富的共享潜在空间。

3. 实验分析

作者在NeurIPS竞赛提供的10X Multiome（scRNA-seq + scATAC-seq）数据集上，对CVQVAE的性能进行了全面评估，并与Raw data（仅PCA+LSI）、scJoint、CLUE(GLUE)等方法进行了比较。

作者使用了一套由scIB工具包提供的标准评估指标，这些指标分为两大类：组学整合（评估批次/模态混合程度）和生物信息保留（评估细胞类型、细胞周期等信息的保留程度）。最终得分是这两类指标的加权平均（生物信息保留权重0.6，组学整合权重0.4）。

结果表明CVQVAE取得了最高的总分（0.840），超越了包括竞赛冠军CLUE（0.832）在内的所有基线方法。CVQVAE在多个关键单项指标上拔得头筹，包括NMI（评估聚类与真实标签的一致性）、CC score（细胞周期保留）和Graph connectivity（图连通性，评估同类细胞在潜在空间的连接紧密程度）。这表明CVQVAE在保留关键生物学信号和实现紧凑整合方面表现出色。

UMAP可视化清晰地显示，相比于原始数据中分离的两个模态，CVQVAE的整合结果中，scRNA-seq（圆形）和scATAC-seq（星形）的数据点被均匀地混合在一起，同时不同颜色的细胞类型（如B细胞、T细胞）又形成了清晰的聚类。

为了验证矢量量化（VQ）模块的必要性，作者进行了一个关键的消融实验：将CVQVAE中的VQ层去掉，变回一个普通的交叉训练VAE（CVAE）。CVAE的总分仅为0.690，远低于CVQVAE的0.840，性能下降非常明显。这有力地证明了VQ模块对于提升模型性能至关重要。

作者进一步分析了训练过程的损失曲线。CVAE的损失在训练初期迅速下降后就很快进入平台期，不再变化，这正是后验坍塌的典型表现。而CVQVAE的损失则在整个训练过程中持续稳定地下降，表明其训练过程更健康、更有效。

作者进一步检验了CVQVAE整合后的潜在空间是否保留了有意义的生物学信息。作者可视化了B细胞的marker基因BANK1和T细胞的marker基因INPP4B在整合后的UMAP图上的表达。结果显示，BANK1高表达的细胞清晰地聚集在被注释为B细胞的区域，INPP4B高表达的细胞也主要分布在T细胞区域。这表明细胞类型特异性的基因表达模式被完好地保留了下来。

作者分析了细胞的伪时间（pseudotime）或分化轨迹。在原始数据中，由于批次效应，分化轨迹是分散且不清晰的。而在CVQVAE整合后的空间中，细胞的分化轨迹（图中红色箭头所示）变得连续且清晰，批次效应被有效消除。这证明CVQVAE能够保留并凸显细胞动态发育过程的信息，为下游的轨迹分析提供了高质量的输入。