基础模型在大量未标记的数据集上进行预训练,然后在特定任务上进行微调,最近基础模型在生物学应用中取得了成功。该文探讨了最近的scRNA-seq基础模型,scBERT和scGPT。专注于细胞类型注释的微调任务,主要探索了与简单基线,比如L1正则化的logistic regression相比,预训练模型的优势和缺陷,包括few shot设置。作者还进行了消融实验,以了解预训练是否真的能提高模型性能。最后,以scBERT为例,作者演示了微调对超参数设置和参数初始化的敏感性。综上所述,该文的结果强调了严格测试基础模型的重要性,并对模型学习的嵌入进行思考,以便更有效地利用这些模型来促进单细胞数据分析。
来自:A Deep Dive into Single-Cell RNA Sequencing Foundation Models
背景概述
scRNA-seq数据的高质量大规模基础模型可以大大提高单细胞RNA测序分析管道的性能,最近开发了几个基础模型,如scBERT、scGPT、Geneformer、scFoundation和scSimilarity。这些模型采用预训练+微调范式,在下游任务上对模型进行微调时,理论上需要收集比从头开始训练模型所需的小得多的特定于应用的数据集。这种样本效率在生物学领域具有特别的价值,在生物学领域,获取标签可能具有挑战性或昂贵,导致下游任务的可用数据集很小。
在该文中,作者探讨了两个模型,scBERT和scGPT。专注于细胞类型注释的微调任务,作者提供了原始论文研究以外的新基准,以检查这些基础模型在该任务的能力。然后,以scBERT为例,作者演示了微调对超参数设置和参数初始化的敏感性。实验结果强调了严格测试基础模型的重要性,以便更有效地利用这些模型促进单细胞数据分析。
方法
图1概述了的不同模型。对于研究的两个基础模型,scBERT和scGPT,都依赖于基于Transformer的架构来处理输入,但它们在如何表示输入数据、模型架构和训练过程方面有所不同。作者还分析了一个逻辑回归基线,该基线直接训练感兴趣的特定任务而不进行预训练。
scBERT
scBERT将每个基因编码为两个embedding的总和:一个代表基因的规范化,即binned log-scale表达水平,另一个通过先前开发的gene2vec嵌入指定基因身份,允许模型在无序数据中区分基因。由此产生的嵌入通过内存高效的Performer进行处理以产生输出。
与BERT中的预训练策略非常相似,scBERT通过对112万个人类细胞的imputation任务进行预训练,其中“被mask”的基因表达值被预测为细胞中所有其他基因嵌入的函数。在微调过程中,一个3层神经网络在较小的标记数据集上进行训练,以获取Transformer层输出的基因嵌入并预测细胞类型。在该文中,作者还评估了一个“无预训练”版本的scBERT,在微调时将Transformer设置为随机初始化的权重。
scGPT
scGPT在很大程度上遵循与scBERT相似的体系结构和预训练范例,只是有一些不同的设计选择。scGPT不是在log-scale上对输入数据进行分组,而是根据基因的表达对基因进行分组,从而使基因均匀分布在每个分组中。代替gene2vec,scGPT使用随机基因身份嵌入和附加的描述元信息的“条件嵌入”来区分每个基因。scGPT不是Performer架构,而是通过Flash-Attention处理嵌入。除了基因嵌入之外,scGPT还进一步训练一个CLS token来总结每个细胞。scGPT使用受OpenAI的GPT系列启发的因果掩蔽策略,实现了生成式预训练。
该模型在3300万个人类细胞上进行了预训练,并通过更广泛的微调任务进行了评估,包括细胞类型注释、遗传扰动预测、批次校正和多组学整合。本文主要研究细胞类型注释。
Logistic Regression
与之前描述的深度学习基础模型不同,逻辑回归基线不遵循预训练/微调范式。这种方法具有最少数量的可训练参数,并且简单地估计对数归一化基因表达数据的线性系数。作者使用L1正则化来提高稀疏性。
实验流程
在这项工作中,作者首先探索了预训练模型(如scGPT或scBERT)与简单基线(L1正则化逻辑回归)相比的相对性能。由于预训练表示对于训练数据有限的少样本预测特别有用,因此作者额外评估了每个模型在few shot设置中注释细胞类型的能力。作者发现,一个简单的逻辑回归基线与预训练的模型相比,即使在few shot设置中,也可以对细胞类型进行注释。
接下来,作者试图了解预训练对基于Transformer模型的性能有多大贡献。为此,跳过预训练过程,直接对模型进行细胞类型标注的微调。实验表明,对于scBERT,去除预训练不会对模型在细胞类型注释上的性能产生有意义的影响,而对于scGPT,它会。作者进一步证明,在不学习有意义的基因表示的情况下,scBERT可以在mask预训练和细胞类型注释上获得良好的准确性,注意到良好的准确性并不一定意味着丰富的表示学习。
最后,作者探讨了scBERT对超参数设置和随机种子的鲁棒性,发现这些设置显著影响模型性能。
- 图1:对不同方法的原理进行了比较。两个基础模型,scBERT和scGPT,构建输入数据的嵌入,并通过基于Transformer的体系结构处理它们。架构通过mask任务进行训练,并针对感兴趣的任务进行微调。作者将这些模型与逻辑回归和没有预训练的Transformer架构进行比较,以了解预训练/微调范式在几种不同实验环境中的潜在好处。
部分实验结果
对于scBERT,在不学习丰富表征的情况下,可以获得良好的预训练精度和微调精度
尽管在预训练目标上表现良好(准确率为78%),但该文的实验显示预训练的scBERT并不影响最终的微调性能。在本节中,作者将删除scBERT架构的另一个组件,并表明具有良好的预训练或微调accuracy不足以证明模型正在学习丰富的基因表示。
在scBERT中,给定基因的embedding是该基因的两个独立embedding的总和:“gene2vec”嵌入和“expression”嵌入。“gene2vec”嵌入将每个基因的身份编码,并且对于数据集中每个细胞中的每个基因都是恒定的。“expression”嵌入将给定基因的表达根据其分箱表达(分箱为[0-1)、[1-3)、[3-7)、[7-15)、[15-31]和31+library size-normalized UMIs)将其转换为六个随机初始化向量之一。expression嵌入(而不是gene2vec嵌入)驱动不同细胞下给定基因的可变性,并且是用于预训练模型的mask任务的重点。
在本实验中,作者从每个基因的表示中去除gene2vec嵌入。这严重限制了模型的表示能力,因为现在只有六种expression嵌入用于表示数据集中的每个基因。在预训练过程中,模型不知道基因的身份,每个被mask的基因都有相同的“上下文”。每个落在同一表达bin中的基因都是相同的。例如,在整个Panglao预训练数据集中,95%的基因被相同地表示为落在[0,1)中,2%的基因在[1-3)中,1.5%的基因在[3-7)中。
考虑到该模型不知道哪个基因是哪个基因,人们可能会认为,scBERT将无法高精度地预测基因的表达。然而,实验观察到,即使在这种具有挑战性的设置中,scBERT在Panglao验证集上预测隐藏基因表达的准确率达到72%(相比之下,当同时使用gene2vec +expression嵌入时,准确率为78%)。在对“no gene2vec”实验进行反思后,作者发现,在Panglao预训练验证集中的67,880个细胞中,有99.8%的细胞,scBERT预测了该细胞中所有被mask基因的每个细胞最常见的表达bin。这些结果表明,使用scBERT的bin表达嵌入方案,一个模型可以依靠简单的启发式方法,比如预测最常见的表达bin,在不学习上下文基因表示的情况下,在mask预训练中表现良好。
在微调方面,未经gene2vec嵌入的scBERT在细胞类型注释任务上表现出下降但仍合理的性能(“no gene2vec”准确率为0.701,F1为0.595)。注意到,作者没有专门为no gene2vec实验重新调优超参数,而是使用之前确定的超参数来微调完整的scBERT模型,因此,在这里报告的性能可能低估了专门为该任务调优超参数的潜在性能。虽然逻辑回归基线和无预训练消融实验已经证明,在PBMC数据中,准确的细胞类型注释不需要丰富的表示,但这一结果进一步强化了细胞类型注释的性能并不能代表基础模型的表示学习能力。