本篇主要先熟悉Geneformer的方法设计
绘制基因网络图需要大量转录组数据来了解基因之间的联系,这阻碍了数据有限环境中的潜在发现,包括罕见疾病和临床上无法接触的组织的疾病。最近,迁移学习利用在大型通用数据集上预训练的深度学习模型,然后可以使用有限的任务特定数据对广泛下游任务进行微调,从而彻底改变了自然语言理解和计算机视觉等领域。在生信领域,Geneformer模型在约 3000 万个单细胞转录组的大规模语料库上进行预训练,以便在生物学数据有限的环境中进行上下文特定的预测。在预训练期间,Geneformer 对生物知识有了基本的了解,并以完全自监督的方式将生物知识编码到模型的注意力权重中。Geneformer 应用于有限患者数据的疾病建模,确定了心肌病的候选治疗靶点。总体而言,Geneformer 代表了一种预训练深度学习模型,可进行微调以适应广泛的下游应用,从而加速发现关键生物调控网络的regulators和候选治疗靶点。
来自:Transfer learning enables predictions in network biology,Nature,2023
目录
- 背景概述
- 网络架构和预训练
- Genecorpus-30M组成和rank value encoding
- 数据集收集
- rank value encoding
- 预训练和微调
- Geneformer架构
- 预训练
- 微调
- embedding和attention
- gene embedding
- cell embedding
- attention
- 硅基扰动
- 关于细胞注释
背景概述
通过绘制驱动疾病进程的基因调控网络图谱,可以筛选出通过使核心调控元件正常化来纠正网络的分子,而不是靶向那些不会改变疾病的外周下游效应因子。然而,绘制基因网络需要大量的转录组数据来了解基因之间的联系,数据有限的情况(包括罕见病和影响临床上难以获取组织的疾病)阻碍了用于校正网络的药物发现。测序技术的进步提供了大量不对齐的单细胞转录组数据,并且,单细胞分辨率有助于精确观察转录组状态,这为推断网络相互作用提供了高分辨率的数据,特别是在面对由多种细胞类型失调的疾病场景中。
最近,迁移学习的概念彻底改变了自然语言理解和计算机视觉,它利用在大型通用数据集上预训练的深度学习模型,针对大量下游任务进行微调(任务特定的数据有限,单独使用时不足以产生有意义的预测)。与需要为每个任务从头开始重新训练新模型的建模方法不同,这种方法将大规模预训练阶段学到的基础知识普及到与预训练学习目标不同的大量下游应用(图 1a)。自注意力机制的出现进一步改变了深度学习领域,它生成了能够关注大范围输入空间并学习在每个上下文中哪些元素最值得关注的上下文感知模型,从而提高了模型的预测能力。基因调控网络高度依赖于上下文,而基于注意力的模型(称为Transformer)正好特别适合特定上下文的调控网络动态建模。
- 图1a:迁移学习策略示意图,先自监督大规模预训练,将预训练权重复制到每个微调任务的模型中,添加微调层并使用有限的任务特定数据对每个下游任务进行微调。Geneformer在大规模转录组数据上进行了预训练,以便在数据有限的环境中进行预测。作者收集组装了一个大规模预训练语料库 Genecorpus-30M,包含来自广泛组织的 2990 万个人类单细胞转录组,这些数据来自公开数据集。然后,作者使用自监督mask学习目标在这个语料库上对 Geneformer 进行了预训练,以获得对调控网络知识的了解。
网络架构和预训练
Geneformer 是一种基于注意力机制的深度学习模型,该模型在大规模转录组数据上进行了预训练,可通过迁移学习在有限的数据下进行任务特定预测。Geneformer 利用自注意力机制来关注每个单细胞转录组中表达的大范围基因空间,并了解哪些基因是最重要的,需要重点关注,以在给定的学习目标内优化预测准确性。重要的是,生物调控网络可能因细胞类型、发育时间点或疾病状态而异。因此,上下文感知是 Geneformer 模型架构的独特优势,可实现针对每个细胞的预测。
首先,作者收集整理一个大规模预训练语料库 Genecorpus-30M,其中包含来自广泛组织的 2990 万个人类单细胞转录组,这些转录组来自公开可用的数据(图 1b)。作者排除了高突变的细胞(例如恶性细胞malignant cells和永生化细胞immortalized cell),这些细胞过度特异可能会导致网络学习不到通用知识,此外,作者建立了过滤指标,以排除可能的双联体(doublets)或受损细胞。
- 图1b:Genecorpus-30M 的摘要。NOS,未另行指定(not otherwise specified)。
然后,每个单细胞的转录组以rank编码(rank value encoding)的形式呈现给模型,在这种编码方式中,基因按照其在当前细胞中的表达情况进行排序,而该表达值是通过其在整个Genecorpus - 30M中的表达情况进行归一化得到的(图1c)。尽管这种基于rank的表示方法有局限性,比如没有充分利用转录计数中提供的精确基因表达测量值,但rank编码提供了每个单细胞转录组的非参数表示。具体来说,这种方法还会将普遍高表达的 housekeeping genes 的优先级降低,即将它们归一化为较低的rank。相反,像转录因子这类基因,它们在表达时可能处于低水平,但具有很强的区分细胞状态的能力,这些基因在编码中会移到更高的rank(扩展数据图1c)。此外,这种基于rank的方法可能对技术伪影(technical artefacts)更具鲁棒性,这些技术伪影可能会使绝对转录计数值出现系统性偏差,但每个细胞内基因的整体相对排名则是保持稳定的。
- 图1c:预训练的 Geneformer 架构。每个单细胞转录组被编码成一个rank编码(rank value encoding),然后通过六层Transformer编码器单元(transformer encoder units),其参数如下:输入大小为2048,嵌入维度为256,每层有四个注意力头(attention heads)。Geneformer 在2048的输入大小上使用完全密集自注意力(full dense self - attention)。可提取的输出包括基因和细胞嵌入、注意力权重和上下文预测。
然后,每个单细胞转录组的rank编码(rank value encoding)会通过六个Transformer编码器单元(transformer encoder units)进行处理,每个单元由一个自注意力层(self-attention layer)和前馈神经网络层(feed forward neural network layer)组成(图1c)。预训练是通过掩码学习目标(masked learning objective)来完成的,在其他领域已经表明,这种方法可以提高在预训练期间所学到的基础知识对各种下游微调目标的通用性。在预训练期间,每个转录组中15%的基因被掩码处理,然后训练模型,利用剩余未掩码基因的上下文信息来预测在特定细胞状态下每个掩码位置应该是哪个基因(扩展数据图1d - f)。这种方法的主要优势在于它是完全自监督的,并且可以在完全无标签的数据上完成,这使得可以纳入大量的训练数据,而不受限于带有相应标签的样本。作者应用了GPU以便在大规模数据集上进行高效的预训练。
- 扩展图1c:归一化后,与所有基因相比,转录因子(Transcription factors)在统计学上表现出的显著性更低。与所有基因相比,Housekeeping genes平均表现出更高归一化的趋势(Wilcoxon检验,所有基因数量n = 17903,管家基因数量n = 11,转录因子数量n = 1384)。
- 扩展图1d-f:d. 预训练是使用随机抽样的100,000个细胞的语料库进行的,保留10,000个细胞用于评估,并使用3种不同的随机种子。在这3次试验中,评估损失基本相同,这表明在预训练期间的鲁棒性。
e. 预训练是使用随机抽样的100,000个细胞的语料库进行的,保留10,000个细胞用于评估,并使用3种不同的mask百分比。与5%或30%的掩蔽相比,15%的掩蔽具有略低的评估损失。
f. 预训练是使用随机抽样的90,000个细胞的语料库进行的,然后对模型进行微调,以使用包含在(include)或排除在(exclude)90,000个细胞预训练语料库中的10,000个细胞来区分剂量敏感和剂量不敏感的转录因子。通过对这10,000个细胞进行五重交叉验证来测量下游微调任务的预测潜力,通过AUC、混淆矩阵和F1分数表明结果基本相同。由于微调应用是针对与mask学习目标完全分离的分类目标进行训练的,因此特定任务数据是否包含在预训练语料库中与下游分类预测无关。
Genecorpus-30M组成和rank value encoding
数据集收集
作者收集了一个大规模的预训练语料库,名为Genecorpus-30M,它包含了从公开数据中获取的、来自多种组织的2990万(29900531)个人类单细胞转录组(图1b)。作者排除了那些具有高突变的细胞(例如,恶性细胞和永生化细胞)。只纳入了基于droplet-based的测序平台,以确保表达值单位的可比性。总体而言,纳入了561个数据集,并将其以loom和HDF5格式作为统一文件存储,其中包括原始研究中的元数据,这些元数据以下面描述的行(特征)和列(细胞)属性形式存在。
用于将数据转换为统一的loom HDF5 文件的工具包括 loompy、scanpy、anndata、scipy、numpy、pandas、Cellranger 和 LoomExperiment。
行特征属性包括基因ID的Ensembl注释、名称和类型(例如,蛋白质编码、微小RNA、线粒体等)。注释数据是从Ensembl和MyGene中检索出来的。列细胞属性包括一个唯一的Genecorpus - 30M细胞ID,它由数据集名称、样本名称和该数据集中的细胞条形码组成。数据集和样本名称也作为单独的个体属性包含在内,如果需要的话,可以通过从唯一的Genecorpus - 30M细胞ID中检索这些名称来得到细胞条形码。列细胞属性还包括数据集中所包含的主要器官,作者将其注释为以下类别之一:脂肪、肾上腺、气道、膀胱、骨骼、骨髓、大脑、乳腺、脐带血、蜕膜、耳朵、胚胎、内皮、眼睛、心脏、免疫、未明确的肠道、肾脏、大肠、肝脏、肺、淋巴结、淋巴管、肌肉、鼻腔、食道、胰腺、胎盘、多能干细胞、前列腺、皮肤、小肠、脾脏、胃、睾丸、胸腺、扁桃体、其他、卵黄囊。列细胞属性还包括基于原始研究提供的元数据的数据集所包含的特定器官。如果原始研究包含细胞类型注释,作者也将这些作为每个细胞的细胞类型列属性包含在内。作者还包含了所使用的具体测序平台。
列细胞属性还包括对每个细胞的一些计算测量值:读取计数的总数、线粒体读取的百分比、被Ensembl注释为蛋白质编码或微小RNA(miRNA)基因的基因数量,以及细胞是否通过了作者为细胞的可扩展过滤而建立的质量控制指标,以排除可能的双峰(doublets)细胞和受损细胞。在这项工作中,只有通过这些过滤指标的细胞才用于下游分析。具体来说,对数据集进行过滤,以保留在该数据集内读取计数总数在均值和3个标准差范围内的细胞,以及在该数据集内线粒体读取在均值和3个标准差范围内的细胞。Ensembl注释的蛋白质编码和miRNA基因用于下游分析。将检测到的Ensembl注释的蛋白质编码或miRNA基因数少于7的细胞排除在外。最终,有2740万(27406217)个细胞通过了定义的质量控制。
rank value encoding
作者开发了一种rank编码(rank value encoding)方法,它为每个单细胞的转录组提供了一种非参数表示。这种方法利用了在Genecorpus - 30M中对每个基因表达的多次观测结果,对区分细胞状态的基因进行优先排序。具体来说,这种方法会将普遍高表达的housekeeping genes的优先级降低,即将它们归一化为较低的秩。相反,像转录因子这类基因,它们在表达时可能处于低水平,但具有很强的区分细胞状态的能力,这些基因在编码中会移到更高的秩。此外,这种基于秩的方法可能对技术伪影(technical artefacts)更具鲁棒性。
为此,首先计算在整个Genecorpus-30M中所有通过质量过滤的细胞中每个检测到的基因的非零表达中值。通过将每个细胞中的基因转录本计数除以该细胞的总转录本计数来校正不同的测序深度。然后,我们用每个基因在Genecorpus - 30M中的非零表达中值对每个单细胞转录组中的基因进行归一化,并根据这些基因在特定细胞中的归一化表达对它们进行排序。值得注意的是,作者选择使用非零表达中值,而不是将零包含在分布中。每个基因的归一化因子是根据预训练语料库计算的,并用于之后呈现给模型的所有数据集。所提供的分词器(tokenizer)代码包含此归一化程序,将其用于对呈现给Geneformer的新数据集进行分词,以确保每个基因所使用的归一化因子的一致性。
然后,基于在Genecorpus-30M中平均173152个细胞中检测到的总共25424个蛋白质编码或miRNA基因的词汇表,对每个单细胞转录组的秩值编码(rank value encodings)进行标记化(tokenized)处理。这个词汇表还包括另外两个用于填充(padding)和掩码(masking)的特殊标记。标记化的数据存储在Huggingface数据集(Datasets)结构中,该结构基于Apache Arrow格式,允许在没有内存限制的情况下处理大型数据集且实现零拷贝读取。值得注意的是,这种策略也节省空间,因为基因是作为排序的标记(ranked tokens)而不是精确的转录值存储的,而且用户只存储在每个细胞内检测到的基因(norm后表达值为0的基因不被考虑,预处理没有平移只有缩放,所以原始表达为0 的基因预处理后表达也是0),而不是包含所有未检测到基因的完整稀疏数据集。
预训练和微调
Geneformer架构
Geneformer由六个Transformer编码器单元(transformer encoder units)组成,每个单元由一个自注意力层(self - attention layer)和前馈神经网络层(feed forward neural network layer)构成,参数如下:输入大小为2048,嵌入维度为256,每层有四个注意力头(attention heads),前馈大小为512(图1c)。Geneformer在2048的输入大小上使用完全密集的自注意力。其他参数如下:非线性激活函数,修正线性单元(ReLU);所有全连接层的丢弃概率(dropout probability)为0.02;注意力概率的丢弃率(dropout ratio)为0.02;权重矩阵初始化器的标准差为0.02;层归一化层的epsilon为 1 × 1 0 − 12 1×10^{-12} 1×10−12。使用Huggingface Transformers库进行模型配置、数据加载和训练。
预训练
预训练是通过掩码学习目标(masked learning objective)完成的,这种方法可以提高在预训练期间所学基础知识对各种下游微调目标的通用性。在预训练期间,每个转录组中15%的基因被掩码,然后训练模型利用剩余未掩码基因的上下文来预测在特定细胞状态下每个掩码位置应该是哪个基因。这种方法的一个主要优点是它完全是自监督的,并且可以在完全无标签的数据上实现,这使得可以纳入大量的训练数据,而不受限于带有相应标签的样本。预训练的超参数被优化为以下最终值:最大学习率为 1 × 1 0 − 3 1×10^{-3} 1×10−3;学习调度器(learning scheduler),linear with warmup;优化器(optimizer),带权重衰减修正(weight decay fix)的Adam算法;warmup步数为10000;权重衰减为0.001;batch size为12。使用Tensorboard进行跟踪,模型进行了三个epoch的预训练。
由于对于一个密集的自注意力模型来说,2048的输入大小相当大(例如,BERT的输入大小为512),并且Transformers相对于输入大小具有二次方的内存和时间复杂度,作者实施了一些措施来优化大规模预训练的效率。
- 减小填充计算量:Huggingface Transformers库中的训练器(trainer)被用于预训练,并替换为自定义的标记器(tokenizer)来实现动态的、按长度分组的填充(padding),这将填充的计算量最小化,并使预训练速度提高了29.4倍。这个过程随机抽取一个大批次(megabatch),然后按其长度降序排列小批次(minibatches)(以确保尽早发现任何内存限制问题)。然后对小批次进行动态填充,由于按长度分组,从而将填充所浪费的计算量最小化。
- GPU训练的新技术:作者还应用了分布式GPU训练方面的最新进展,以使用Deepspeed在大规模数据集上进行高效的预训练,它将参数、梯度和优化器状态在可用的GPU之间进行划分,尽可能将处理/内存转移到中央处理器(CPUs)上,以便在GPU上容纳更多内容,并通过确保长短期内存分配不混合来减少内存碎片。总的来说,预训练在大约3天内完成,分布在三个节点上,每个节点有4个英伟达V100 32GB GPU(总共12个GPU)。
微调
Geneformer的微调(Fine - tuning)是通过用预训练的Geneformer权重初始化模型并添加一个最终的特定任务Transformer层来实现的。微调目标要么是基因分类,要么是细胞分类。如前所述,使用Huggingface Transformers库中的训练器进行预训练,并替换为自定义的标记器(tokenizer),以及用于动态标记基因或细胞类别(classes)的自定义数据整理器(data collator)。为了证明预训练的Geneformer在提高下游微调应用的预测潜力方面的功效,作者特意对所有应用使用相同的微调超参数。应该注意的是,深度学习应用的超参数调整通常会显著增强学习效果,因此Geneformer在这些下游应用中的最大预测潜力可能被大大低估了。用于微调的超参数如下:最大学习率为 5 × 1 0 − 5 5×10^{-5} 5×10−5;学习调度器(learning scheduler),线性预热;优化器(optimizer),带权重衰减修正(weight decay fix)的Adam算法;预热步数为500;权重衰减为0.001;批大小为12。所有微调均在单个epoch内进行,以避免过拟合。
一般来说,根据经验,与预训练目标更相关的应用程序受益于冻结更多的层,以防止对有限的特定任务数据过拟合;而与预训练目标关系较远的应用程序则受益于对更多层进行微调,以优化新任务的性能。基因分类应用的微调结果以AUCs±标准差和F1分数报告,这些结果是基于五重交叉验证策略计算得出的,该策略是在80%的基因训练标签上进行训练,并在20%预留的基因训练标签上测试性能,重复五次。值得注意的是,由于微调应用是针对与掩码学习目标完全分离的分类目标进行训练的,因此如扩展数据图1f所示,特定任务数据是否包含在预训练语料库中与分类预测无关。
然后,作者使用所有基因训练标签对剂量敏感性(dosage sensitivity)和二价性(bivalency)分类模型进行完全微调,以测试它们对样本外数据的泛化能力。测试了在没有进一步训练的情况下,经过微调以区分剂量敏感基因和不敏感基因的模型是否能够预测参考文献(A cross-disorder dosage sensitivity map of the human genome,Cell,2022)中最近报道的一组疾病基因的剂量敏感性。
embedding和attention
gene embedding
对于提供给Geneformer的每个单细胞转录组,该模型将每个基因嵌入到一个256维的空间中,这个空间编码了该基因在特定细胞环境下的特征。通过Geneformer模型的前向传递(forward pass)对给定单细胞转录组内的每个基因的256个嵌入维度的隐藏状态权重进行评估,从而提取出与上下文相关的Geneformer基因嵌入(Geneformer gene embeddings)。论文研究中分析的基因嵌入是从模型的倒数第二层提取的,因为已知最后一层包含的特征与学习目标预测更直接相关,而倒数第二层是一个更具通用性的表示。
cell embedding
Geneformer细胞嵌入(Geneformer cell embeddings)对单个细胞的状态特征进行编码,它是通过对该细胞中检测到的每个基因的嵌入进行平均而生成的,从而产生一个256维的嵌入。正如上面所讨论的,作者使用了倒数第二层的嵌入。
对于疾病建模,作者使用单核转录组数据对Geneformer进行微调,以区分心肌细胞与非衰竭心脏或受肥厚或扩张性心肌病影响的心脏,使用受影响患者的单核转录组数据。样本根据患者划分。训练数据包括93589个心肌细胞(非衰竭n=9,肥厚n=11,扩张n=9),样本外数据包括39,006个心肌细胞(非衰竭n=4,肥厚n=4,扩张n=2)。
用于微调的超参数为:
- 最大学习率:0.000804;学习调度:带预热的多项式;优化器:Adam与weight decay fix;warmup步:1812.678558;权重衰减:0.258828;seed:73.152431,批量:12,冻结层:2,epoch:0.9。
预测的准确度、精度和召回率在如上的样本外数据上进行计算,使用所有类别的预测的直接平均值,而不按出现频率加权。在建立了对样本外数据具有高预测准确度(90%)的模型后,作者将非衰竭、肥厚或扩张型心肌病状态定义为训练数据中每种情况下心肌细胞嵌入的平均值。然后,通过对来自训练数据中未衰竭心脏或受肥厚或扩张性心肌病影响的心脏的每个心肌细胞转录组中表达的每个基因进行硅基删除或激活,研究了定义每种状态的基因。
作者首先确定了在非衰竭心肌细胞中缺失或激活的基因,这些基因在256维嵌入空间中显著地将非衰竭嵌入转向肥厚或扩张性心肌病状态。对于硅基删除分析,首先删除每个非衰竭心肌细胞中的随机基因,并确定这种删除在嵌入空间中推动细胞嵌入的位置(图6a)。对于硅基激活分析,执行类似的程序,但是不删除基因,而是将基因移动到rank value编码的最前面,以模拟该基因的过表达。
- 图6a:对Geneformer进行微调以区分来自正常心脏(Non-failing)或受肥厚型或扩张型心肌病(HCM和DCM)影响的心脏的心肌细胞,从而定义了每种细胞状态的嵌入位置。然后,可以通过在硅基(in silico)删除或激活正常心肌细胞内的随机基因来进行疾病建模(左图),以定义随机分布(灰色云状),从而识别出在硅基删除或激活这些基因后可使嵌入位置明显向肥厚型或扩张型心肌病状态转变的基因。在硅基治疗分析(中间和右侧)中采用了相反的方法。
attention
Geneformer的六层中的每一层都有四个注意力头(attention heads),这些注意力头旨在以无监督的方式进行学习,从而自动关注不同类别的基因,提高预测能力。作者通过Geneformer模型的前向传递(forward pass)对给定单细胞转录组内的每个基因在每个自注意力层(self - attention layer)内的每个注意力头的上下文Geneformer注意力权重(Contextual Geneformer attention weights)进行提取。
对于实验,为了研究模型在预训练阶段是如何学习网络动态的,作者检查了经过预训练的Geneformer注意力权重。模型中每个基因经过训练的注意力权重反映了基因-基因的相互作用。
硅基扰动
作者设计了一种计算机模拟(in silico)扰动方法,通过扰动给定基因的rank来模拟它们的抑制或激活。在细胞和基因嵌入水平上测量计算机模拟扰动的影响,分别模拟扰动如何影响细胞状态和基因网络内下游基因的调控。模拟删除是通过:
- 从给定单细胞转录组的rank value编码中删除给定基因,并量化原始细胞嵌入和扰动后的细胞嵌入之间的余弦相似度,以确定在该细胞环境中删除该基因的影响
- 或者看单细胞转录组中剩余基因的基因嵌入,以确定哪些基因对给定基因的模拟删除最为敏感。
计算机模拟删除可以对单个基因或多个基因进行删除。计算机模拟激活是通过将给定基因移到rank value编码的前面来模拟的。理论上,可以通过在rank value编码内将基因向上或向下移动到更细微的程度来模拟更细微的下调和激活。
关于细胞注释
在原文中,作者说过,Geneformer最关注的是理解网络动力学,而不是细胞级的注释,但是作者依然研究了Geneformer在细胞类型注释方面的性能。
Geneformer随机划分训练和测试数据,比例为8:2,并且验证数据的所有细胞类型都在训练数据中存在。对于注释,预训练的Geneformer使用10个epoch的0冻结层进行微调,或者使用方法部分Geneformer微调中描述的学习超参数。