对比学习能够快速映射数百万尺度的多模态单细胞图谱.

0. TL; DR

Concerto（contrastive learning of cell representations）是一种基于自监督蒸馏（self-supervised distillation）和对比学习（contrastive learning）的深度学习方法，用于单细胞多组学图谱的建模。Concerto 采用非对称的教师-学生网络架构，通过最大化同一细胞在不同增强视图下的一致性来学习高质量的细胞嵌入。该方法仅需通过区分每个细胞与其他细胞，即可适应多种下游任务，包括自动细胞类型分类、数据整合、批次校正以及参考映射。

在模拟和真实数据集上的基准测试表明，Concerto 在各项任务中均显著优于竞争方法。作者利用 Concerto 将 COVID-19 患者免疫细胞映射到综合参考图谱，成功重现了不同疾病状态下的差异免疫响应，并发现了疾病特异性细胞状态。此外，Concerto 具有出色的可扩展性，可在 1.5 小时内构建包含 1,000 万个细胞的参考图谱，并在 8 秒内完成 10,000 个细胞的查询映射。

1. 背景介绍

近年来，单细胞多组学工具（single-cell multiomics tools）以前所未有的分辨率彻底改变了对组织的表征方式。大规模项目如 Human Cell Atlas 和 Tabula Muris Atlas 正接近百万级（multimillion）规模。现有的计算流程如 Seurat、SCANPY 和 Pegasus 已被广泛开发和基准测试。

单细胞数据分析具有一些独特特征：

对比学习（Contrastive learning）最近在计算机视觉领域（如 SimCLR 和 MoCo）取得了巨大成功。这类方法通过最小化对比损失（contrastive loss）来学习增强视图之间的一致性。受其自然语言处理中 SimCSE 的启发，作者提出 Concerto，一种用于单细胞分析的自蒸馏对比学习框架。通过将每个未标记细胞定义为实例判别（instance discrimination）的预训练任务，Concerto 以自监督方式学习语义不变的细胞表示。

2. 方法介绍

Concerto 利用自蒸馏（self-distillation）对比学习框架，配置为非对称的教师-学生架构（asymmetric teacher-student architecture）：

对于单细胞 RNA-seq 数据，归一化的基因 count 矩阵被转换为索引-值格式（index-value format）：索引是基因身份，由包含某物种所有基因的字典定义。值是细胞中对应的 count 值。这种编码方案支持稀疏高维输入并提高计算效率。

⚪ 教师网络

教师网络接受 $X_{\text{indices}} \in \mathbb{R}^G$ 和 $X_{\text{counts}} \in \mathbb{R}^G$，其中 $G$ 表示基因数量。

首先，通过嵌入层将基因索引映射到 $d$ 维向量空间（默认 $d=128$）：

\[\text{emb} = \text{Embedding}(x_{\text{indices}})\]

然后，计算加权的隐藏向量：

\[\text{hidden}_i = \text{emb} \times x_{\text{counts}}\]

接下来，使用注意力机制聚合基因嵌入。隐藏向量首先经过带有 tanh 激活的多层感知机（MLP）：

\[\text{hidden}_i' = \tanh(\text{hidden}_i)\]

然后，通过细胞上下文向量（cellular context vector）$u \in \mathbb{R}^d$ 计算注意力权重：

\[\text{attention}_i = \text{softmax}(\text{hidden}_i' \cdot u)\]

通过加权求和获得聚合向量：

\[\text{hidden} = \sum_i (\text{attention}_i \times \text{hidden}_i')\]

最后，经过批次归一化、dropout 层和 ReLU 激活的稠密层，得到教师网络的输出 $z_{\text{teacher}} \in \mathbb{R}^d$：

\[z_{\text{teacher}} = \text{ReLU}(\text{Dropout}(\text{BatchNormalization}(\text{hidden})))\]

为校正批次效应，添加了域特定的批次归一化层（domain-specific batch normalization layer）。

⚪ 学生网络

学生网络仅接受 $X_{\text{counts}} \in \mathbb{R}^G$ 作为输入，依次通过稠密层、批次归一化、dropout 层和 ReLU 激活，得到输出 $z_{\text{student}} \in \mathbb{R}^d$：

\[z_{\text{student}} = \text{ReLU}(\text{Dropout}(\text{BatchNormalization}(\text{Dense}(x_{\text{counts}}))))\]

⚪ 自蒸馏

知识通过自蒸馏在两者之间传递。通过为每个未标记细胞定义实例判别（instance discrimination）的预训练任务，Concerto 通过最大化每个细胞的教师视图和学生视图之间的一致性来学习语义不变的嵌入。

数据增强通过在输出层前添加随机 dropout 掩码（mask）在模型级别生成，这种方法受到 SimCSE 的启发，可在不改变原始分子输入的情况下实现最小化数据增强。

对比学习在单位超球面空间上进行，显式比较细胞嵌入对。对于批次中的 $N$ 个细胞，通过教师和学生网络生成 $2N$ 个数据点（每个细胞两个视图）。

给定正样本对 $(j, j^+)$ ，其中 $j$ 来自教师视图，$j^+$ 来自学生视图，对比损失（normalized temperature-scaled cross-entropy loss, NT-Xent）定义为：

\[\ell(j, j^+) = -\log \frac{\exp(s_{j,j^+}/\tau)}{\sum_{k=1}^{2N} \mathbb{I}_{[k \neq j]} [\exp(s_{k,j}) + \exp(s_{k,j^+})]}\]

其中相似度 $s_{\alpha,\beta}$ 通过余弦相似度计算：

\[s_{\alpha,\beta} = \frac{z_{\text{teacher},\alpha}^T \cdot z_{\text{student},\beta}}{\tau \cdot \|z_{\text{teacher},\alpha}\| \cdot \|z_{\text{student},\beta}\|}\]

总损失为：

\[\mathcal{L} = \frac{1}{2N} \sum_{k=1}^N [\ell(j, j^+) + \ell(j^+, j)]\]

其中 $\tau$ 是可调的温度系数（temperature coefficient），默认设为 0.1。

⚪ 多组学整合

对于多组学数据（如 RNA 和蛋白质），Concerto 通过对各模态的教师或学生输出进行元素级求和（element-wise summation）来生成统一的细胞视图：

\[z_{\text{student}}^{\text{multi}} = \text{Add}(z_{\text{student}}^{\text{RNA}}, z_{\text{student}}^{\text{protein}}) \\ z_{\text{teacher}}^{\text{multi}} = \text{Add}(z_{\text{teacher}}^{\text{RNA}}, z_{\text{teacher}}^{\text{protein}})\]

⚪ 下游任务适配

3. 实验分析

3.1 自动细胞类型分类性能

作者在 PBMC45k 数据集（$n=31,021$，9 个批次，7 种实验方案）上对比了 Concerto 与 SciBet、Cell BLAST、SingleR、Moana、MARS 以及端到端版本的 Concerto-E2E。

3.2 无监督聚类与多组学整合

3.3 批次校正与过度校正分析

作者在包含 8 个批次、5 种技术（$n=14,890$）的多供体人胰岛（HP）数据集上评估了批次校正性能。

3.4 查询到参考映射（Query-to-reference Mapping）

3.5 COVID-19 免疫分析

作者将 Schulte-Schrepping 等人的 COVID-19 PBMC 数据集（$n=99,049$）映射到由 Ren 等人（$n=451,096$）和 Zhu 等人（$n=42,752$）构建的综合参考上，无需微调。