从数据处理到模型训练,一文速览单细胞大语言模型进展

文摘   2024-08-10 12:03   中国香港  

引言


近年来,基于Transformer架构的预训练语言大模型层出不穷,其超大规模的训练数据与模型参数所带来的成果使得人工智能的发展进入新的阶段。大量生物数据的产生使得生物大模型应运而生,帮助人们更好的进行生物信息的整合和分析。在前面的推送:生物信息大模型卷起来了:应用与前景浅谈中,我们已经从总体的角度介绍了生物信息大模型。本文将作为我们关于生物大模型介绍的第一个具体场景-单细胞转录组大模型。在这篇文章中,我们将主要从方法的角度出发,介绍单细胞转录组大模型的数据处理和模型训练,之后我们还将另外专门介绍单细胞转录组大模型的具体应用~



什么是单细胞转录组数据?



单细胞转录组学是在生物学领域最重要的一种测序技术。单细胞转录组测序可以测量单个细胞中所有基因的表达。根据生物学中的中心法则,转录是遗传信息从DNA(基因)到RNA的过程。而RNA将会进一步翻译为蛋白质,组成细胞的结构并实现各种功能。因此,单细胞转录组是从分子层面研究细胞性质,分析疾病原因的重要参考,吸引了许多生物学家的注意,他们测得了大量来自不同物种,组织和状态的单细胞转录组数据,并希望通过这些数据分析细胞性质,找出导致疾病发生的关键通路,指导精准药物发现。


单细胞转录组数据通常以表达量矩阵的方式出现(行表示单个细胞,列表示基因名称,每个元素代表该基因在对应细胞的表达量)。由于我们是对mRNA逆转录后得到的cDNA(互补DNA)进行测序,每种基因都会通过转录过程得到相对应的mRNA参与此后的分子过程(翻译、调控等),不同基因有着不同的功能,因此,基因类型与基因表达量都很重要。由于目前有多种平台可以进行单细胞测序,加上每次测序所需进行的操作以及条件的差别,不同测序批次得到的数据会有较大的差异,即批次效应。此外,单细胞转录组数据一般包含上万个基因,且没有特定的序列性,也为大模型的引入带来一定困难,针对这些数据特点,我们将从数据处理和模型训练两方面进行讲述。



数据处理——如何进行单细胞转录组数据的Tokenization?



在大语言模型中,一个‘句子’会被分为若干个单独的‘词’,而整个模型会学习词汇表中所有词的Embedding,以及每个句子内部各位置的词语之间的关系。对于单细胞转录组,Tokenization在目前已发表的大模型中可分为两种方式


1.以基因为‘词’,细胞为‘句子’。2.以细胞为‘词’,细胞组成的组织为‘句子’,细胞内部视为基因为单位的‘词袋(bag of words)’,通过对基因Embedding的聚合得到细胞的Embedding。


这两种分词方式有着不同的侧重点,前者更关注细胞内部基因之间的相互关系,而后者则进一步加入了细胞间相互关系。接下来将通过几种大模型的比较,更详细的解释两者的实现方式和差异。




基因语言模型



对于表达量信息的处理,有的大模型选择将其转换为离散值(如排序、分组),有的通过线性转换转为向量,有的将其视为概率对基因进行采样,以下是一些具体例子。


Geneformer对于基因表达量信息使用了排序分词法,每个细胞所对应的基因顺序,是表达量从高到低的基因标识,减少了由绝对值带来的批次效应,对每个词语的位置和上下文赋予了相应的意义。


scGPT可以处理转录组,scATAC-seq和细胞表面蛋白数据(Chip-seq)。与Geneformer不同,scGPT设置了多个Tokenize通道,通过对各通道向量Embedding的整合运算,编码了多模态多层次的信息


  1. 基因名:按照基因名得到词汇表中相对应的数字标识。

  2. 表达量:通过分组方式将绝对表达量转换为离散的相对值,每个基因对应的组别即为该基因的表达量标识。

  3. 条件:表示了每个基因的条件,例如所受干扰,可用于干扰后预测。

  4. 批次:标记数据的批次信息,用于降低批次效应和批次整合。

  5. 模态:标记数据的种类,例如基因或者表面蛋白,可用于多模态整合。



与以上两种方法不同,scFoundation并没有将表达量信息转变为离散值(排序、分组),而是对其进行线性转换得到可训练的向量Embedding,与基因名所得到的Embedding共同作为模型的输入。此外,scFoundation在每个输入向量末尾添加了总表达量Token T和S,前者表示原总表达量,后者表示欠采样后的较低的表达量,这两种表达量标记用来模拟不同的测序深度,从而用于解决不同测序技术带来的结果差异和批次效应。


scMulan同样对表达量进行了线性转换得到向量Embedding。但是,它同时对Metadata(批次、器官等信息)、Tasks(要进行的生成和预测任务)也进行了Tokenization,从而可以利用更多的附加信息,根据任务指示词自动进行特定条件下的生成训练


以上几种大模型都只适用于人类单细胞数据,而UCE则使用来自八种不同物种的数据进行训练。在基因名Token Embeddings的初始化中,不同于其他模型的随机初始化,UCE对每个基因匹配其编码的蛋白质,使用来自于蛋白质大模型ESM2的Embeddings作为其Embedding的初始化。UCE没有对表达量进行直接的转换,而是根据其大小对每个细胞内部进行基因采样,表达量越大采样的概率越大,并将采样后的基因按照在染色体上的位置进行排序,这一排序的方式也和其他的模型有很大区别。然而,UCE仅使用了基因表达矩阵这一信息作为数据,而并没有利用其他的附加信息(批次等)。


为了加入更多信息,另一跨物种单细胞大模型GeneCompass引入了四种Prior Biological Knowledge(如下图),并用特定的方法将Prior Knowledge转换为与基因名和表达量同一空间下的Embedding,GeneCompass也对表达量进行线性转换,之后又将其按照表达量高低进行排序,作为其Embedding层的输入。




细胞语言模型



CellPLM将每个细胞看作一个词语单位,不同细胞在组织中构成‘句子’,每个细胞Embedding则由其中所包含的基因Embedding聚合计算得到。此外,CellPLM还加入了空间转录组数据作为训练数据的一部分,增加了细胞之间的空间信息作为额外的参考。这种方法与上文所讲述的基因语言模型的不同之处在于,CellPLM更倾向于捕捉不同细胞之间的作用关系,因此适合用来建模空间转录组学的数据。




训练数据规模



本节最后一部分列出了这些大模型进行预训练的数据情况供参考,包括数据量、参数量和物种组成。其中,由于scHyena与SCimilarity没有使用Transformer架构,这里没有列出其具体参数量。


模型数据量
(million)
参数量
(million)
是否多物种
Geneformer3010.4
scGPT3353
scFoundation50100
scMulan10368
scHyena0.7-
CellPLM1180
SCimilarity22.7-
UCE34650
GeneCompass120100



模型与预训练——如何设计适合于单细胞数据的模型和学习任务?



与传统的大模型训练类似,现有的单细胞生物大模型大部分使用基于自注意力机制的Transformer结构和掩码预测基因表达的方式进行预训练。然而,与文本数据不同的是,单细胞测序数据具有一些独特的性质:非序列性、高批次效应、长序列问题,因此不同的单细胞转录组大模型在模型结构和训练方法上也做了各自的针对性设计。


1
非序列性


现有的大模型主要采用掩码语言模型(MLM)或者自回归语言模型(AR)方式进行训练。和天然具备序列性的语言不同,一个细胞的基因表达并不具备序列性,因此如何将一个细胞的基因表达表示成一个“句子”进行预训练就是一个重要的研究问题。单细胞语言模型可以按照训练方式分为两类(MLM/AR)。


对于基于MLM训练的大模型,Geneformer的分词方式是按照正则化之后的表达量排序,直接赋予基因的排序和位置意义,因此在预训练中,通过位置掩码来预测被掩盖的各位置中应有的基因种类,从而学习基于周围基因的基因表达水平的条件概率分布。scFoundation使用预测随机掩码基因表达量的方法,能够学习到基因之间的共表达。scMulan则在训练中随机打乱基因序列来降低基因顺序带来的影响。总体来说,相比于文本MLM预训练,单细胞大模型的MLM预训练将传统的基于单词位置的位置编码改为基因编码,使得模型能够处理非序列数据。


对于基于AR训练的大模型,scGPT在模型中使用了一种注意力掩码机制,即对掩码的未知基因进行表达量预测。在训练时,对自注意力矩阵添加掩码,仅计算已知基因和未知基因间的自注意力矩阵,而不计算未知基因之间的注意力,在自回归的预测出未知基因的表达之后,再将未知基因的预测结果当作已知基因,并不断预测其他未知基因。这一设计模拟了语言模型训练中的‘下一词预测(Next Token Prediction)’任务,另外, scGPT还实现了分别从基因prompt(下图A)和细胞prompt(下图B)的基因表达矩阵的生成。



2
批次效应


批次效应是单细胞数据中常见的一种现象,由于测序平台和技术的不同,不同实验测试得到的单细胞数据在测序深度上会有较大差异,从而导致了表达量的不同。在数据处理方面,scGPT和GeneFormer对表达量的绝对值转换为离散的相对值信息,UCE仅将表达量作为采样概率,其模型的初始Embedding和基因序列都与表达量无关,这些方式在一定程度上减轻了测序深度带来的批次效应。


在模型与训练设计方面,scFoundation设计了RDA modeling的预训练任务(如下图),在输入数据中添加总表达量的同时,使用高斯欠采样方法模拟不同的测序深度样本,以欠采样矩阵学习重构原表达矩阵,从而为来自不同批次的数据学习统一的高质量表示



CellPLM设计了batch-aware解码器,在样本的latent embedding上添加对应的batch embedding来减少生成过程中批次效应的影响。此外,CellPLM还采用高斯混合模型作为隐空间的先验分布来解决训练数据的噪声和批次差异。


3
长序列问题


每个细胞通常都包含了成千上万个基因,在单细胞转录组大模型中,长序列的处理十分常见且重要。scGPT使用Flash-Attention方法,并且只考虑非零表达值的基因作为输入,其序列最长长度为1200,在训练时每一轮随机采样不同的基因组成句子。Geneformer设置句子最长为2048,同样只考虑非零基因。拥有更多参数的scFoundation则可以一次性输入大约两万个基因信息,针对这种长序列大规模的训练,scFoundation设计了非对称的自编码器结构,在得到每个词的Embedding后,将整个句子分成零(包含零表达基因和被掩码基因)和非零两部分,只有非零句子进入Transformer编码器,而编码器得到的非零Embedding会与原始零Embedding相接,共同进入以MLP为结构的简单解码器,重构出原始表达矩阵。UCE也使用对细胞内非零基因随机采样的方法减少输入序列的长度。


自注意力机制在长序列的处理上会消耗更多的时间(平方复杂度),同时模型的参数量也较大。不同于其他模型,scHyena使用更适合于长序列处理的Hyena模型为主要架构,其中的long convolution层在增加处理长度的同时不会导致参数量的剧烈增加。




其他预训练任务



除了对掩码基因表达水平的预测外,一些模型针对于不同的目标设计了不同的预训练任务。scMulan和GeneCompass不仅需要预测表达量,还需要预测基因名,也就是说,需要在已知的一些条件和基因信息下预测其他未知的可能表达的基因名,其中,scMulan采用了基于任务指示词和不同条件限定的训练方式,提高模型的条件生成能力。UCE更注重细胞Embedding的训练,因此在其训练时将得到的细胞潜在空间Embedding(从<cls>token中获得)与被掩码的基因(不作为模型的输入)的初始化Embedding合并作为Decoder的输入,来预测这些掩码基因是否会在该细胞中表达。


scTranslator使用成对的转录组和蛋白组数据,训练模型进行从基因表达矩阵到蛋白表达矩阵的推断。SCimilarity旨在构建一个单细胞转录组检索模型,因此,它根据细胞类型用三重采样法来生成参考样本和新样本用于后续的对比学习(如下图),从而实现了新样本在reference数据上的高效检索。对于给定的新样本,可以利用 SCimilarity 检索得到类似的细胞样本用来进行细胞注释等下游任务。




应用与讨论



总之,依据单细胞数据不同于自然语言的一些特点和特定的需求目标,不同大模型设计了更适合的数据处理方法、模型架构和训练方式。从目前已发表的单细胞转录组(以及单细胞多模态)生成式大模型来看,它们的应用方向很广泛。在基因水平上,细胞特异性基因Embedding可用于进行基因干扰后预测和细胞内基因调节网络推断,细胞水平上,可用细胞Embedding进行更准确的聚类和类型注释,多模态、多批次的数据整合,以及药物干扰后预测。这些应用任务可以帮助更好地理解亚细胞水平和细胞水平的动态网络,包括基因相互作用以及细胞间相互交流,辅助进行疾病模型的构建和药物治疗的计算机模拟。这里,关于单细胞转录组大模型的应用领域及其结果将在大模型系列的后续文章中介绍。


在数据方面,UCE和GeneCompass通过多物种的数据获得了跨物种的大模型;scTranslator尝试进行了转录组到蛋白组信息的转变,这引出了一些值得探索的问题:现存已经有多种物种和模态的数据,如何更好利用这些数据进行跨物种的大模型构建?如何使用大模型强大的学习推理和生成功能实现实现合理的多模态信息的链接,从各模态信息中帮助人类获取更全面的微观层面的生物网络?


基于大规模的数据和可学习参数,单细胞转录组大模型的编码和推理能力获得了很大提高,然而,其zero-shot生成能力则有待于进一步加强,也就是说,如何在无需微调的情况下应用于跨组织、跨平台的数据,有待于进一步的探究。一方面,局限于预测表达量的训练任务或许并不足够学习到分子世界的生物体运行逻辑,一些模型已经设计了多任务训练以提高其生成能力;另一方面,每个细胞和基因都有更多先验知识,UCE将基因编码蛋白质信息和染色体位置信息加入到模型学习中,带给我们了启发,合适的将这些有生物学意义的知识加入训练,进行特定条件的学习(基于Prompt生成)也具有很大意义。


总之,Transformers和大规模预训练在语言数据处理上取得了巨大的成功,这样的成功能否被迁移到生物数据,并解决传统方法所不能解决的生命科学相关问题,是值得进一步思考的事情,也是AI for Science所关注的主要问题之一。


参考文献:

Theodoris C V, Xiao L, Chopra A, et al. Transfer learning enables predictions in network biology[J]. Nature, 2023, 618(7965): 616-624.

Cui H, Wang C, Maan H, et al. scGPT: toward building a foundation model for single-cell multi-omics using generative AI[J]. Nature Methods, 2024: 1-11.

Hao M, Gong J, Zeng X, et al. Large-scale foundation model on single-cell transcriptomics[J]. Nature Methods, 2024: 1-11.

Wen H, Tang W, Dai X, et al. CellPLM: pre-training of cell language model beyond single cells[J]. bioRxiv, 2023: 2023.10. 03.560734.

Bian H, Chen Y, Dong X, et al. scMulan: a multitask generative pre-trained language model for single-cell analysis[C]//International Conference on Research in Computational Molecular Biology. Cham: Springer Nature Switzerland, 2024: 479-482.

Liu L, Li W, Wong K C, et al. A pre-trained large generative model for translating single-cell transcriptome to proteome[J]. bioRxiv, 2023: 2023.07. 04.547619.

Oh G, Choi B, Jung I, et al. scHyena: Foundation Model for Full-Length Single-Cell RNA-Seq Analysis in Brain[J]. arXiv preprint arXiv:2310.02713, 2023.

Yang X, Liu G, Feng G, et al. Genecompass: Deciphering universal gene regulatory mechanisms with knowledge-informed cross-species foundation model[J]. bioRxiv, 2023: 2023.09. 26.559542.

Heimberg G, Kuo T, DePianto D, et al. Scalable querying of human cell atlases via a foundational model reveals commonalities across fibrosis-associated macrophages[J]. bioRxiv, 2023: 2023.07. 18.549537.

Rosen Y, Roohani Y, Agrawal A, et al. Universal cell embeddings: A foundation model for cell biology[J]. bioRxiv, 2023: 2023.11. 28.568918.


如果您对我们的内容有任何疑问需要咨询或想要联系作者,可以备注“AIMS”联系课代表。也欢迎大家投稿!课代表微信:AimNina。转载请联系本公众号获得授权。



AIMShare
AlMShare(Al with Mission and Social responsibility)是以推动实现有社会使命的人工智能为目标的高质量知识分享平台。我们专注于AI交叉领域研究与人才培养,致力于打破该领域信息差。
 最新文章