Mila唐建团队新作：可迁移、可解释的单细胞RNA测序模型

【栏目：前沿进展】近日，McGill大学的李岳老师和魁北克人工智能研究所Mila唐建老师团队共同提出了一种高效、易用、可拓展、可迁移、可解释的模型——scETM，用于单细胞RNA测序工作，并于Nature Communications发表。该论文提出了一个嵌入主题模型（Embedded Topic Model）可以帮助更好地理解细胞的功能和区分不同细胞的类型，未来具有广泛的生物意义。（本文为智源社区独家约稿）

作者：赵伊凡、蔡辉宇（Mila唐建团队）

编校：李梦佳

背景和摘要

单细胞转录组测序（scRNA-seq）技术能以单细胞精度测量样本的表达谱，是加深人类对生物学的认知、推动疾病诊断和治疗的进步的有力工具。近年来已经出现了多个图谱（atlas）级别的scRNA-seq数据集。

然而，对scRNA-seq数据的整合分析（integrated analysis）仍存在很大的挑战，不仅数据规模越来越大，测序技术、实验批次、dropout 等因素的干扰也屡见不鲜，如何从大型参考数据集（reference dataset）上迁移知识、如何解释模型学到的参数等问题尚未得到充分解决。

McGill大学的李岳老师和魁北克人工智能研究所Mila唐建老师团队共同提出了一种高效、易用、可拓展、可迁移、可解释的模型——scETM，近日于Nature Communications发表。利用GPU，该方法可轻松在百万级数据集上训练scRNA-seq数据模型，且训练后的模型参数可以迁移到其他（如规模更小的、其他组织甚至其他物种的）数据集上。scETM能够为数据集中的细胞、基因和主题计算嵌入向量，利用这些嵌入可以进行聚类、差异表达分析、富集分析等多种下游任务。这些嵌入向量为模型提供了高度的可解释性（见3.3节）。

scETM的代码和教程已经在GitHub上开源（https://github.com/hui2000ji/scETM），其Python包可直接用pip安装（pip install scETM）。scETM与scanpy和tensorboard兼容，研究者可以利用scanpy对scETM得到的低维细胞嵌入进行后续分析，也可以利用tensorboard探索细胞、基因和主题嵌入，获得新的生物发现。

方法介绍

我们知道，要建模scRNA-seq数据，核心的目标有两个：因为形状为细胞数×基因数的表达矩阵太过庞大，所以需要降维；因为测序过程中存在批次效应和技术瑕疵，所以需要降噪。

主题模型（topic models）可以很好地实现以上目标，它将细胞看作一个个文档，细胞内的基因看作词语，并用“主题”将两者连接起来。以LDA（Latent Dirichlet Allocation）为例，一个主题对应于一个基因组上的分布，分布上概率较高的词语即为主题词。文档被视为各主题的混合，文档的主题分布可以用于对文档进行分类，或者比较文档之间的相似性。

应用到scRNA-seq数据建模上，我们便可以用主题模型，将表达矩阵Y分解为形状为细胞数×主题数的矩阵θ和形状为主题数×基因数的矩阵β，分别代表数据集内细胞的主题混合（topic mixture）和主题的基因分布（gene distribution）。

为了更好地建模长尾的基因分布，获得可解释性更强的主题，scETM采用嵌入主题模型（Embedded Topic Model, ETM），它借鉴了词嵌入模型的思想，将基因分布矩阵β进一步分解为主题嵌入矩阵α和基因嵌入矩阵ρ。这样我们就可以在嵌入空间中观察所有基因和主题之间的关系。

针对scRNA-seq数据的特点，scETM为每个批次s引入了长度为基因数的批次效应校正向量λ。对于批次s内的细胞d，基因表达量的预测值为softmax(θs,dαρ + λs)。参数λ的引入，使得主题分布θs,d中不必包含批次信息，有利于得到批次无关的主题分布。实验证明这是scETM去除批次效应的关键。

scETM的完整模型如下图所示。

其输入为经过L1-正则化使每个细胞总表达量为1的表达矩阵
编码器（encoder）完成降维和降噪，计算每个细胞的主题分布，用于对细胞进行聚类。训练完成后，编码器可以直接迁移到新的数据集上完成聚类等任务。
解码器（decoder）用一种可解释的方式试图将主题分布还原为表达矩阵，其参数包含主题与基因嵌入和批次效应校正矩阵。
模型的损失函数为训练数据对数似然概率的证据下界（evidence lower bound, ELBO）。

实验结果

又好又快的scRNA-seq数据融合

作者在Mouse Pancreas (MP), Human Pancreas (HP), Tabula Muris (TM), Alzheimer’s Disease (AD), Major Depressive Disorder (MDD), Mouse Retina (MR)共六个大小和批次效应各不相同的数据集上，评测了scETM与Harmony, Seurat, LIGER, scVI, scVI-LD, scVAE-GM七个强大的scRNA-seq模型。评测指标为ARI（越高越好，表1）和kBET（越高越好，表2），分别衡量模型的聚类能力和批次效应校正能力。

这一系统的评测结果表明scETM在各种数据集上均表现良好，位于第一梯队。进一步的实验（见原文Batch overcorrection analysis一节）表明scETM能够在校正批次效应和识别细胞类型之间取得良好的平衡。对比scETM和scETM－λ可以发现，批次效应校正向量λ的引入对于scETM的批次效应校正能力影响极为显著。

如果需要进一步加强scETM的批次效应能力，可以在主题分布上面加入批次判别器，引入对抗损失训练模型，使得编码器生成的主题分布包含的批次信息尽可能减少。由此得到的scETM + adv模型有着一流的批次效应校正能力，同时也基本保持了scETM优秀的聚类能力。

scETM不仅建模效果好，而且消耗的资源少。得益于简单优雅的模型设计和高效的实现，scETM比其他基于深度学习的模型训练速度高出3-10倍，在处理大规模数据时，其内存占用也低于基于相互最近邻（mutual nearest neighbors, MNN）的Seurat和基于融合型非负矩阵分解（integrative non-negative matrix factorization, iNMF）的LIGER等传统方法。Harmony和Scanorama由于把表达矩阵降维到了50-100维，其时间和内存消耗更低，但这些方法不具备可解释性和可迁移性。

跨批次、跨组织、跨物种迁移学习

作者在三对数据集上对比了scETM和scVI、scVI-LD的迁移学习能力：

TM (FACS) ↔ MP：两者都是小鼠的数据，但测序技术和样本来源不同，前者是基于FACS得到的多组织全景数据，后者是使用inDrop技术得到的胰岛数据。
HP (inDrop) ↔ MP：两者都是用inDrop在胰岛上测得的数据，但前者来自人类，后者来自小鼠。
MusMOp ↔ HumM1C：两者都是用10x Genomics在初级运动皮质上测得的数据，但前者来自人类，后者来自小鼠。

其中A → B表示模型在数据集A上训练，随后用于聚类数据集B中的细胞。由下图可知，scETM的ARI在所有实验中均高于scVI和scVI-LD，kBET也在大多数情况中更优。

实验还得到了几条有趣的结论：

在TM (FACS)上训练的scETM模型，相比在MP上训练的scETM模型，区分MP中的B细胞、T细胞和巨噬细胞的能力显著增强。这体现了迁移学习的优势：在大规模、高质量的参考数据集上训练的模型，可以帮助建模小规模的新数据集。
仅在1886个MP细胞上训练的scETM模型，在TM (FACS)中展现了不俗的聚类能力。这提示scETM具有较强的跨组织泛化能力，可能能够从小规模、单一器官的数据中提炼出更加普适的可迁移的信息。
MusMOp ↔ HumM1C上scVI和scVI-LD的kBET大幅领先，但这是由于它们完全没能区分细胞类型，ARI极低，将所有细胞——包括不同批次的细胞都混在了一起。

从功能和主题表达差异的角度探索scETM主题的生物意义

为了探索scETM主题的生物意义，我们应用了基因集富集分析（Gene Set Enrichment Analysis, GSEA），即根据我们所关注的基因集（如GO gene set）在某一scETM主题的基因分布中的位置，从左到右计算累计的统计量（running-sum statistic)，在此过程中得到的最大值也就是富集得分（Enrichment Score, ES）。

在HP数据集中，我们发现许多scETM主题与胰腺功能相关，包括胰岛素加工、胰岛素受体再循环、胰腺癌等。相应的，在AD、MDD数据集中，我们找到了多个与淀粉样蛋白纤维的形成以及抑郁症呈现显著富集的scETM主题。而在scVI-LD学习到的gene loading (与scETM主题对应）则只能找到较弱的相关性。scETM同时学习基因嵌入和主题嵌入的特性也使得我们可以在同一嵌入空间可视化基因和主题（见图5c）。

那么scETM是否也能学到细胞类型特异性、疫病状况特异性主题呢？通过对主题进行表达差异分析，作者在HP、AD、MDD均发现scETM学习到的部分主题对细胞类型标记基因（marker genes) 具有高度选择性，并且对细胞类型具有高度辨别力。此外，一些主题在病例组和对照组的表达也呈现显著差异，例如原文提到的AD主题58，除了在AD组和对照组的表达有所不同外，作者发现主题58显著富集线粒体基因，而线粒体恰好与阿尔兹海默症密切相关。

结语

综上，scETM为大规模scRNA-seq数据集的整合分析提供了一个统一、可拓展的框架。与现有方法相比，它在数据融合、迁移学习、可拓展性、可解释性等方面表现突出。

当然，scETM也不是尽善尽美的，作者在文末给出了一些未来的改进方向：

1）建模更复杂的批次效应，例如，用树状结构建模层次性的批次效应；用分段嵌入或MLP建模连续的批次效应变量（如年龄）

2）将scETM推广至多组学数据融合

3）融合scETM与预训练的基因嵌入