通过最优传输进行单细胞数据整合的统一计算框架.

0. TL; DR

uniPort 是一个结合了耦合变分自编码器（coupled variational autoencoder, coupled-VAE）和小批量非均衡最优传输（Minibatch Unbalanced Optimal Transport, Minibatch-UOT）的统一单细胞数据整合框架。耦合VAE架构包含一个共享的解码器和多个数据集特异性的解码器，能够同时利用共享基因和数据集特异性基因的信息，Minibatch-UOT损失对齐不同数据集的细胞嵌入。

作者通过将其应用于整合包括单细胞转录组、染色质可及性和空间分辨转录组在内的多种数据集，展示了uniPort的强大能力。

1. 背景介绍

单细胞测序技术，如scRNA-seq和scATAC-seq，正在以前所未有的分辨率揭示细胞的复杂性。然而，单一的组学技术只能提供片面的信息。因此，将来自不同平台、不同模态的数据进行计算整合，已成为该领域的核心任务之一。

现有的整合方法虽然众多，但大多面临着各种挑战：

更重要的是，随着空间转录组学（spatial transcriptomics, ST）的兴起，如何将非空间的scRNA-seq数据与带空间坐标的ST数据进行整合，以实现空间细胞类型的反卷积（deconvolution），成为了一个新的、亟待解决的问题。

目前，还没有一个方法能够统一地解决这些不同的整合任务：从非配对的scRNA/scATAC整合，到scRNA/ST的整合。为了填补这一空白，作者开发了uniPort。它旨在通过一个结合了耦合VAE和最优传输（Optimal Transport, OT）的统一框架，来应对上述所有挑战。

2. uniPort 方法

uniPort的核心是一个创新的深度学习框架，它将一个耦合变分自编码器（coupled-VAE）和一个高效的小批量非均衡最优传输（Minibatch-UOT）算法结合在一起。

耦合变分自编码器是uniPort进行特征提取和数据生成的基础。与标准的VAE不同，uniPort的耦合VAE架构经过了精心设计，以同时处理共享和特异性信息。

uniPort使用一个共享的、数据集无关的概率编码器 $\psi(z|x)$。这个编码器负责将不同数据集（如scRNA-seq和scATAC-seq）中共同的高变基因的表达谱，都投影到一个共享的细胞嵌入潜在空间中。编码器输出一个多元高斯分布的均值 $\mu$ 和标准差 $\sigma$，潜变量 $z$ 通过重参数化技巧从中采样：$z = \mu + \sigma \times \nu$，其中 $\nu \sim \mathcal{N}(0, I)$。

从潜在空间 $z$ 出发，模型包含两组解码器：

这种双重解码器架构，使得uniPort能够同时学习和利用共享信息与特异性信息，从而更全面地处理数据集间的异质性。

uniPort的VAE部分通过最大化一个修改版的证据下界（ELBO）**来进行训练。这个**ELBO包含了对共享基因和特异性基因的联合重构损失。

\[L_{ELBO^*} = \mathbb{E}_{\psi(z|x)}[\log \phi(x|z) + \lambda_s (\log \phi_x(x_s|z) + \log \phi_y(y_s|z))] - \lambda D_{KL}(\psi(z|x) \| p(z))\]

为了在潜在空间中对齐来自不同数据集的细胞，uniPort引入了最优传输（Optimal Transport, OT）。传统的OT算法计算复杂度高，难以扩展到包含数十万细胞的大规模数据集。标准的OT假设两个分布的质量是相等的（即细胞数量相同且一一对应），这在处理细胞类型组成不同的异质性数据集时是不现实的。

uniPort采用了小批量非均衡最优传输（Minibatch-UOT）来解决这些问题。小批量（Minibatch） 是指在每次训练迭代中，作者不计算整个数据集间的OT距离，而是只在随机采样的两个小批量数据之间进行计算。这极大地降低了计算成本。非均衡（Unbalanced） 是指UOT通过在KL散度约束下进行优化，允许两个分布的质量不相等，从而能更灵活地处理细胞类型比例不匹配的异质性数据集。

uniPort的OT损失函数定义如下：

\[L_{UOT}^* = \sum_{i,j} C_{ij} \times T_{ij}^*\]

其中 $C_{ij}$ 是两个mini-batch中细胞i和j在潜在空间中的距离代价（基于它们编码后的高斯分布参数计算）。\(T^*\)是通过IPOT算法高效求解出的最优传输计划，它是一个概率矩阵，\(T_{ij}^*\)表示从细胞i运输多少质量到细胞j。

uniPort的最终训练目标是最小化ELBO*和Minibatch-UOT损失的加权和：

\[L_{uniPort} = -L_{ELBO^*} + \gamma L_{UOT}^*\]

通过这个统一的框架，uniPort在学习富有表现力的细胞嵌入的同时，也实现了跨数据集的精准对齐。

3. 实验分析

作者在一系列来自不同平台、不同物种的真实数据上，对uniPort的性能进行了全面的评估，并与Seurat, Harmony, GLUE, MultiMAP等多种SOTA方法进行了比较。

3.1 实验一：整合scATAC和scRNA数据

作者首先在三个配对或非配对的scRNA/scATAC数据集（PBMC和mouse spleen）上测试了uniPort。

在配对PBMC数据这个金标准数据集上，uniPort和GLUE取得了最好的、且性能相当的整合结果。uniPort在剪影系数（Silhouette coefficient）上排名第一（0.64），在细胞对齐准确性（FOSCTTM）上排名第二（0.0694），在聚类准确性上排名第二（总分2.321）。

在非配对小鼠脾脏数据这个更具挑战性的非配对任务中，uniPort, scMC, Harmony和Seurat表现最佳。uniPort取得了最高的剪影系数（0.709），而其他各项指标也名列前茅。

这些结果表明，无论数据是配对的还是非配对的，uniPort都能实现准确、鲁棒的整合。

3.2 实验二：处理非均衡数据集

为了测试uniPort处理细胞类型组成不匹配的异质性数据集的能力，作者在小鼠脾脏数据集上进行了非均衡匹配实验，即人为地从一个模态中移除部分细胞类型。

UMAP可视化显示，在两种非均衡场景下（UBM-ATAC和UBM-RNA），uniPort都能准确地将那些只存在于一个模态中的细胞类型（如DC, NK细胞）分离出来，同时很好地对齐了共享的细胞类型（图a, b）。

定量比较显示，相比于其他方法在非均衡场景下性能急剧下降，uniPort和Seurat表现出了最强的鲁棒性。特别是在UBM-ATAC场景下，uniPort的聚类总分和剪影系数均为最高。这证明了UOT在处理非均衡分布时的独特优势（图c, d）。

3.3 实验三：整合空间转录组学（MERFISH）与scRNA-Seq

作者进一步挑战了将高分辨率成像数据MERFISH与scRNA-seq进行整合的任务。

uniPort和scVI的整合效果最好，能够清晰地分离开OD Immature等细胞类型。特别地，uniPort还能准确地识别出MERFISH数据中特有的室管膜细胞（ependymal cells）（图a, b）。

uniPort在剪影系数（0.706）和聚类总分（2.404）上均排名第一，超越了包括gimVI在内的所有其他方法（图c, d）。

作者还测试了利用scRNA-seq数据为MERFISH数据插补缺失基因的能力。结果显示，uniPort的插补精度（通过Spearman和Pearson相关性评估）显著高于SOTA方法gimVI和Tangram。

3.4 实验四：空间细胞类型反卷积 (Deconvolution)

作者利用uniPort输出的最优传输计划 $T^*$，对基于条形码（barcoding-based）的空间转录组数据进行细胞类型反卷积。

在真实的10x Visium小鼠大脑数据上，uniPort成功地重建了大脑皮层的经典层级结构，其反卷积出的细胞类型比例和空间分布与已知的解剖学结构高度一致。

在HER2+乳腺癌数据中，uniPort的反卷积结果清晰地描绘了肿瘤微环境（TME）中的主要细胞成分，如T细胞、B细胞和癌细胞。更重要的是，通过计算T细胞和B细胞的共定位（colocalization），uniPort成功地定位了与良好预后相关的三级淋巴结构（Tertiary Lymphoid Structures, TLS）。

在一个分辨率更低的胰腺癌微阵列数据上，uniPort依然成功地揭示了肿瘤内部的异质性，区分出了两种具有不同功能（通过KEGG富集分析）和不同免疫浸润特征的癌症克隆（cancer clone）亚群。