使用对比学习处理大规模多模态单细胞数据

摘要

单细胞数据集的规模不断扩大，对处理规模正在扩大、模态正在扩展和不可避免的批次效应引发了计算挑战。最近出现了很多基于深度学习的方法，通过将细胞表示为非线性embedding来解决这些问题。这里将介绍细胞表示的对比学习，Concerto，它利用自监督蒸馏框架来建模多模态单细胞图谱（multimodal single-cell atlases）。通过简单地学习区分每个细胞和其他细胞，Concerto就可以适应各种下游任务，例如自动细胞类型分类、data integration，特别是reference mapping。与当前的主流软件包不同，Concerto的对比设置很好地支持对所有基因进行操作，以保留生物特异性。Concerto可以灵活地推广到多组学数据，并获得统一的细胞表示。在模拟和真实数据集的基准测试中，Concerto大大优于其他方法。通过映射到综合的reference，Concerto概括了新冠肺炎患者的差异免疫反应（differential immune responses），并发现了疾病特异性（disease-specific）的细胞状态。Concerto易于并行，可高效扩展，可在1.5小时内构建1000万个细胞reference，并在8秒内查询10000个细胞。总的来说，Concerto将通过迭代构建单细胞reference atlases和快速映射新数据集来转移相关细胞标注，从而促进生物医学研究。

引言

最近的单细胞多组学工具正在以前所未有的高分辨率彻底改变组织的特征。大规模的人类细胞图谱（Human Cell Atlas）和小鼠脑皮层图谱（Tabula Muris Atlas）正在接近数百万规模。目前已经开发了很多pipeline，如Seurat、Scanpy和Pegasus，并对其进行了基准测试。注意到，单细胞分析具有一些独特的特点。首先，过多的零计数可能是由于技术中断或生物信号干扰。关于单细胞数据潜在分布的争论仍在继续。主流软件包依赖特征选择（高变基因，HVG，highly variable genes）和线性降维（主成分分析PCA），但这可能导致信息丢失。深度学习方法为所有基因之间的非线性关系建模提供了一个有前景的解决方案。变分自编码器（VAEs）用具有重构函数的编码器-解码器结构来学习细胞的低维embedding（比如scVI、scETM，GLUE）。然而，强迫模型重构稀疏的计数矩阵得进一步讨论。第二，批次效应广泛存在于技术、条件之间。从混杂的测序数据中提取出生物信号对于数据整合非常重要。Seurat v3 使用 mutual nearest neighbours（MNN）识别跨批次的anchor cell pairs，MNN一次仅允许两个批次的集成，并在处理更多细胞时导致了指数级的内存消耗。Harmony迭代使用模糊聚类和线性校正，trVAE利用conditional VAE校正批次效应。理想的方法应该将样本数量扩展到百万规模，同时整合多个批次，避免混合本身不重叠的簇。最后，query-to-reference mapping已变得流行，这能够快速集成并分类新的数据集，而无需费力的从头聚类或手动注释。与硬性的监督分类不同，我们将query-to-reference mapping视为无监督迁移学习问题，以基于学习的query embedding导出基于voting的细胞注释。Seurat v4在MNN上使用有监督PCA传输reference annotation。Symphony使用混合建模框架将query定位到稳定的reference上。ScArches使用conditional AE通过微调以映射query cell。

对比学习最近在计算机视觉领域取得了巨大成功，如SimCLR和MoCo。这种类型的方法为未标记图像定义了一个pretext task（代理任务，不是目标任务），并通过最小化超球空间中增强视图之间的对比损失来进行自监督学习。学习的embedding可用于图像分类。受对比学习在建模未标记数据方面的优势的启发，我们猜想，只要以自监督的方式区分每个细胞，就可以获得高质量的表示。蒸馏方案被用于在非对称神经网络（即教师-学生网络）之间传递知识，从有监督环境中的模型压缩和在线协同蒸馏（online co-distillation），到半监督环境中（例如，noisy student）的自训练，再到自监督蒸馏以获得更好的表示（例如，SEED、DINO）。利用类似蒸馏的方案在增强视图之间的共享知识上提供了一种简洁的解决方案，可以在典型的对比学习框架中生成类内一致但类间独特的embedding。

这里，我们提出一个用于单细胞分析的自蒸馏对比学习框架Concerto。通过对真实数据集和模拟数据集的综合基准测试，学习的embedding可以针对各种下游需求进行微调，包括自动细胞类型分类、聚类、批次效应校正以及query-to-reference mapping。Concerto 可以灵活地处理多组学数据集，并在每个任务中实现优于竞争方法的性能。实验表明，Concerto的注意力权重提供了模型的可解释性。此外，Concerto根据包含健康和感染样本的integrated reference atlas查询新冠肺炎免疫细胞数据集，概括了不同疾病状态患者的几种不同免疫特征。Concerto是一个强大、准确、可扩展的表示学习框架，可以用于1000万细胞规模的单细胞多模态分析。

Results

Overview of Concerto

Concerto利用了一个自蒸馏对比学习框架，该框架被设置为非对称的teacher-student架构（图1a）。不对称设计注入了不平衡的模型复杂性，其中较大的教师网络通过注意力机制将基因embedding聚合到细胞embedding中，而较小的学生网络使用dense操作简单地将离散输入转换为细胞embedding。表征知识通过自蒸馏在两者之间传递。通过为每个未标记细胞定义实例区分pretext task，Concerto通过最大化每个细胞的教师和学生视图之间的一致性来学习语义不变embedding。在输出层之前添加了一个random dropout mask，以在模型层生成最小的数据增强，这是由SimCSE（回顾SimCSE：用于句子嵌入的对比学习）的语句处理方案驱动的（SimCSE用dropout mask生成句子的正样本）。添加了一个domain-specific的批标准化层，以校正批次效应。处理多组学数据集时，每个模态的element-wise求和可以生成统一的细胞embedding（图1b）。通过投影到单位超球空间上，Concerto通过将同一个细胞的教师和学生的视图整合为正对，同时将一批细胞中的其他细胞分开，从而区分细胞。学习的embedding可以针对各种下游任务进行微调，包括自动细胞类型分类、聚类、批次效应校正以及query-to-reference mapping（图1c）。

• 图1a：每个细胞通过不对称的teacher-student网络结构获得两个视图（蓝色框，形成自蒸馏方案，其中实线和虚线分别表示教师和学生视图的输出。每个细胞的两个视图是正对，将被拉到一起，而其他细胞的视图被推离，一起形成优化目标，作为对比学习。
• 图1b：当使用多组学数据集作为输入时，Concerto只需在教师或学生模块中对不同模态的输出进行element-wise求和，即可获得统一的细胞embedding。
• 图1c：学习的细胞embedding被用到各种下游任务（有监督注释、聚类、数据整合和query-to-reference mapping）。

对比学习的embedding通过微调显著提高了自动细胞分类的性能，并支持跨组织的新细胞类型发现。

为了证明对比学习的embedding满足刚性细胞分类（rigid cell classification），我们使用现有注释作为训练标签，在Concerto上实现有监督的微调。首先，我们使用经典的人类外周血单核细胞数据集（PBMC45k，n=31021，七个protocols）来比较不同的分类方法，包括基于似然的SciBet、基于神经网络的Cell BLAST、基于相关性的SingleR、基于支持向量机的Moana和元学习方法MARS。Concerto是一种分两步的方法（预训练和微调），而其他的则是端到端的训练。我们还实现了Concerto（Concerto-E2E）的端到端版本，方法是丢弃对比损失并以完全监督的方式进行训练。对于数据集内评估，我们在每个批次（n=9）中应用fivefold交叉验证，并评估所有细胞类型的中值F1分数。Concerto获得最高分数（0.926），在每个fold中表现最稳定（图2a），而Concerto-E2E获得较低分数（0.867），证明了对比学习预训练带来的效用。对于数据集间评估，我们使用一个protocol作为测试集，其他protocol作为训练集（bootstrapping五次）。Concerto在几乎所有train-test splits上都大大优于其他方法（图2b）。

如果测试集包含训练样本中不存在的细胞类型，则好的分类器应将NOTA（none-of-the-above）细胞中的任何一个都标记为拒绝选项（rejection option）。我们下载PBMC CITE-seq数据集（PBMC160k，n=161764个细胞，仅在NOTA研究中使用的RNA数据集），在三个水平上进行注释，并从训练集中删除不同的T细胞，以评估NOTA设置。图2d显示，Concerto可以清楚地分离1级和2级mask的验证和测试集的置信曲线。即使是最具挑战性的3级场景，Concerto也获得了与验证曲线部分重叠的双峰曲线；然而，SciBet将CD4 Mem T细胞误分配到其他类型。

为了对更细粒度的分类进行基准测试，我们将胸腺scRNA-seq atlas（n=107969个细胞）与PBMC45k结合，构建了多层次免疫细胞数据集。合并高分辨率胸腺数据集在区分沿T细胞发育轨迹的细微状态差异方面带来了更大的挑战。Concerto仍然达到最高F1中位数0.830（fivefold交叉验证的平均值），大大优于SingleR（0.705）和SciBet（0.667）（图2c）。Concerto可以很好地区分不同的发育阶段，包括双阴性T细胞、双阳性T细胞和单阳性T细胞。我们还使用了异源板鼠Senis atlas（TMS，heterogenous Tabula Muris Senis，n=101045个细胞，23个小鼠组织）来训练组织分类器（组织内预测）。Concerto在所有组织上都大大优于SciBet（图2f），实现了膀胱（0.999）、脑髓细胞（0.99）和乳腺（0.996）的最高平均ACC。

对于跨组织注释，我们采用与MARS类似的实验设计，将一个组织作为未标注的测试集，并在所有其他组织（TMS数据集）上训练Concerto。通过添加域自适应模块，Concerto在22个保留组织（hold-out tissue）上实现了优于MARS的 adjusted Rand指数（ARI），范围从脾脏的最大ARI（+89.4%）到膀胱的最小ARI（+0.613%）（图2g）。保留组织通常包含训练组织中不存在的几种细胞类型。与MARS类似，Concerto有效地传递知识，发现跨组织的新细胞类型。特别是，当肢体肌肉伸出时，Concerto将来自其他组织的功能相似的细胞类型放置在更靠近肢体肌肉的六个主要注释处（图2h）。Sankey图（图2i）显示，来自其他组织的general B cells、T cells、endothelial cells 和 macrophages cells被正确转移到肢体肌肉。来自肢体肌肉的skeletal muscle satellite cells 和 mesenchymal stem cells正确地映射到其他肌肉和脂肪组织中的对应细胞，而MARS错误地使用T细胞来注释肢体肌肉中的一些satellite cells。

• 图2a和b：比较数据集内和数据集间的性能；
• 图2c：在包括胸腺scRNA-seq（n=107969个细胞）和PBMC45k数据的具有挑战性的数据集上，通过中值F1得分测量基准性能。
• 图2d：比较Concerto和SciBet对训练集中不存在的细胞类型分配低置信度评分的能力；
• 图2e：PBMC160k数据集（RNA、蛋白质和RNA+蛋白质）上多模态注释的性能比较。
• 图2f：TMS数据集上组织内预测精度的比较；
• 图2g：跨组织预测的性能比较；
• 图2h：UMAP显示真实标签与Concerto预测保留肢体肌肉组织；
• 图2i：显示跨组织相关细胞类型标记转移的Sankey图（以肢体肌肉为例）；

其他结果

其他结果包括：

• Concerto能够在多模态数据集上实现有效的无监督聚类，并能够以单细胞分辨率自动从注意力权重中提取分子特征；
• Concerto通过消除不需要的批次效应实现从头数据集成，并支持集成部分重叠的数据集。
• Concerto实现了最先进的query-to-reference mapping精度，并支持在reference中投影看不见的细胞类型。
• Concerto可以有效地扩展到1000万个细胞图谱的构建和reference mapping。
• 绘制新冠肺炎免疫细胞与疾病对照图谱，揭示不同感染状态下的不同免疫反应。

• 图3：Concerto支持在多模态数据集上进行有效的无监督聚类和批次校正数据集成。
• 图3a：PBMC45k数据集上的平均ARI测量，评估Concerto的embedding与PCA、Seurat和scDeepCluster聚类性能的影响。Louvain和Leiden聚类方法用于Concerto、PCA和Seurat。
• 图3b：真实细胞类型标签与簇分配的UMAP可视化。
• 图3c：在PBMC160k数据集（RNA、蛋白质和RNA+蛋白质）上的Concerto学习embedding的UMAP可视化，由Azimuth标记（单模态的1级类别和双模态的2级类别）。蓝色箭头显示了CD8 T细胞、CD4 T细胞和B细胞内亚型的方向分布。黑色箭头用于指示某些细胞类型，如CD4 TCM、CD8 Naive，以避免图中的重叠和模糊。
• 图3d：heatmap显示了注意力权重与主要免疫细胞类型中某些典型模态特异marked gene的关系。热图的第i行、第j列表示第i个细胞的第j个基因（或蛋白质）的注意力权重。
• 图3e：整合人胰岛scRNA-seq数据集后，通过细胞类型标签和批次标签着色的UMAP可视化。基准方法包括Seurat v3、Harmony、trVAE和未校正baseline（仅PCA）。
• 图3f：通过去除除了CEL-Seq2以外的所有β细胞（红色并用红色虚线椭圆表示）进行过校正分析，如UMAP可视化所示。
• 图3g：利用kBET和ASW对六种HVG场景（前2000、5000、10000、15000和20000个HVG以及所有基因）测量的批次校正的比较。

• 图4：Concerto实现了最先进的query-to-reference mapping精度，并支持以数百万个细胞的规模在reference中投影未知的细胞类型。
• 图4a：对Symphony、scArches和Seurat v4的query-to-reference mapping的性能比较。左边：HP->inDrop，HP数据集是reference，但inDrop用作query时除外。右边：HP->MP，HP是reference，inDrop和MP用作query。
• 图4b：ACC测量的Concerto预测混淆矩阵。
• 图4c：在HP数据集上Concerto、scArches、Seurat和Symphony的对齐一致性plot。
• 图4d：非对称师生网络结构的消融研究，2名教师或2名学生分别表示使用注意力操作或密集操作的两个网络。教师+学生表示非对称设计，最终细胞embedding从网络中提取并用括号表示。方框图显示了HP数据集的query-to-reference mapping的准确性。
• 图4e：顶：热图显示Concerto可以成功识别在reference中的CD8 T细胞，并表达典型的CD8 protein、GZMA和CCR7 RNA标记基因。底部：Concerto保留了生物信号，显示标记基因GMLY和NKG7呈负相关，而标记物CCR7呈正相关（负表示CD8 T细胞在距离NK细胞较近的地方表达增加；正表示在距离NK细胞较远的地方表达增多）。
• 图4f：通过UMAP可视化的CLEC12A、CD55、CD11c和CD14 protein的Ground Truth蛋白表达与Concerto预测的热图。
• 图4g：Concerto通过对多模态PBMC160k数据集（RNA+蛋白质）上评估的所有基因进行操作，将未知的细胞类型投影到reference上的能力说明。删除query集中的所有CD8 T细胞。
• 图4h：Concerto的规模通过reference构建和query相同大小的细胞所用的时间来衡量。
• 图4i：Concerto与其他三个主流软件包的示意图比较。

• 图5：层级query-to-reference mapping保留了COVID-19患者的差异免疫响应。
• 图5a：Schulte Schrepping及PBMC-scRNA数据集（n=99049个细胞）映射到整合的COVID-19 reference的图示。
• 图5b：在健康对照（HC）、中度（M）和重度（S）疾病状态中，注释CD8 T细胞分为五种亚型的UMAP可视化。
• 图5c：Expression heatmap显示了CD8初始（naive）、细胞毒性（cytotoxic）、增殖（proliferating）和衰竭（exhaustion）状态的典型标记。
• 图5d：方框图显示了CD8 T细胞亚型在不同疾病状态下的相对百分比。
• 图5e：UMAP显示增殖-衰竭的CD8 T细胞和其他衰竭的CD8T细胞（顶部）。热图显示了UMAP可视化中的功能特定归一化标记（底部）。
• 图5f：方框图显示了不同疾病状态下CD4 T细胞中CD4初始T细胞的相对百分比（左）和不同疾病状态（右）下活化CD4 T淋巴细胞中T细胞活化相关基因的表达水平。
• 图5g：方框图显示了不同疾病状态下注释NK细胞中NK CD56dimCD16bright和NK CD56bright细胞的相对百分比。
• 图5h：UMAP显示抗原呈递标记物（HLA-DR）缺乏和炎症标记物（S100A）富集，其共同定位于注释单核细胞的左上区域。