GeneCompass：跨物种大模型用于破解基因调控机理

GeneCompass是第一个基于知识的跨物种基础模型，该模型预先训练了来自人类和小鼠的超过1.2亿个单细胞转录组。在预训练过程中，GeneCompass有效整合了四种生物先验知识，以自监督的方式增强了对基因调控机制的理解。对多个下游任务进行微调，GeneCompass在单个物种的多个任务中优于最先进的模型，并开启了跨物种生物研究的新领域。总的来说，GeneCompass标志着对通用基因调控机制的学习和加速发现关键调控因子和药物候选靶点的里程碑。

来自：GeneCompass: Deciphering Universal Gene Regulatory Mechanisms with Knowledge-Informed Cross-Species Foundation Model
工程地址：https://github.com/xCompass-AI/geneCompass

背景概述
概述：GeneCompass架构和预训练
方法：数据收集
方法：GeneCompass
- GeneCompass架构
- GeneCompass预训练和优化
- 消融实验与不同规模的数据
- 微调
- 细胞和基因embedding
- 先验知识embedding
下游任务方法
- 单物种的细胞类型注释
- GRN预测
- 细胞对药物及其剂量的反应预测

背景概述

脊椎动物是由数以万亿计的细胞组成的复杂系统，这些细胞被分为数百种不同的类型。细胞合作形成不同的组织和器官，每一个都有一套独特的生理功能。阐明这些组织和器官的基因调控机制对于破解个体发育模式和促进临床治疗至关重要。随着组学测序技术的快速发展，人们已经开始在单细胞分辨率上解剖各种器官中的细胞如何发挥其特定功能，并积累了大量的单细胞数据。然而，基因表达在多个水平上受到调控，从染色质可及性到转录后修饰。这意味着，仅通过湿生物学实验全面解读基因调控机制是劳动密集和耗时的。深度学习模型的出现，凭借其在大型数据集中表示复杂模式的能力，为剖析多层次和跨物种的调节机制提供了机会。

近年来，自然语言领域的BERT、GPT、PaLM和LLaMA等基础模型以及视觉领域的DALL-E在各种下游任务中取得了显著的性能。它们通常都采用自监督学习对大量数据进行预训练，然后通过微调对特定的下游任务进行适应。类似的，转录组是理解生物系统内基因调控活动的"语言"。一些研究利用单细胞转录组数据构建预训练基础模型，如scGPT、Geneformer和scFoundation。这些工作的共同点是利用数以千万计的人类单细胞转录组谱来预训练基础模型，并在广泛的下游任务中表现出卓越的性能，如细胞聚类、细胞类型注释、基因扰动模拟和药物靶标预测。然而，目前的模型有其局限性。

动机：
由于脊椎动物在物种间表现出巨大的多样性，而基因调控机制的进化守恒（尽管不同物种的基因组和表型可能有很大的差异，但它们在基因表达调控的基本原理和组件上存在共同之处），来自不同物种的数据集的整合为揭示基因调控提供了一个绝佳的机会，而目前的模型仅依赖于单个物种的数据。此外，过去几十年积累的丰富的生物学知识，包括核心调控区数据、实验证实的基因相互作用和基因家族注释，代表了我们迄今为止对生物过程的全面理解。将这些知识注入预训练过程中，可以指导模型以自监督的方式学习通用的基因调控机制。

贡献：
GeneCompass是一个知识丰富的跨物种基础模型，在scCompass-126M上预训练，scCompass-126M是目前最大的语料库，包含超过1.2亿个来自人类和小鼠的单细胞转录组。

重要的是，该模型结合了生物学知识，包括启动子序列、基因共表达网络、基因家族信息和转录因子-靶基因调控关系。通过在各种下游任务中微调预训练模型，模型取得了比SOTA模型更好的性能。总的来说，模型在基础模型的发展方面取得了重大突破。

概述：GeneCompass架构和预训练

GeneCompass是一个知识丰富的跨物种基础模型，预训练了来自人类和小鼠的超过1.2亿个细胞的转录组。通过对有限的数据进行微调，GeneCompass可以有效地应用于各种下游生物相关任务（图1a）。利用显性的自注意力机制，GeneCompass可以基于输入的转录组了解细胞和基因之间的复杂关系。

通过构建一个大规模的预训练语料库scCompass-126M推动了GeneCompass的开发。该语料库由来自人类和小鼠的1.26亿个单细胞转录组组成，这些转录组来自公开可用的数据集，包括广泛的器官和细胞类型（图1c, d）。为了确保数据质量，作者过滤掉了异常基因表达的细胞。此外，作者在数据集中保留了具有足够可变性或表达水平的基因，以捕获细胞类型特异性特征（即生物异质性）。由于人和小鼠之间存在不同的基因列表，这里通过共享相同的基因id来对齐它们的同源基因。在GeneCompass中，token字典包含了36,092个基因中的17,465个同源基因（图1a）。
fig1a

图1a：GeneCompass 的框架。该模型在人类和小鼠的大规模单细胞转录组上进行预训练。与从头开始训练不同，预训练的 GeneCompass 可用于多个下游任务，包括细胞类型注释、基因调控网络预测、剂量反应预测等。

fig1b

图1b：四种先验知识嵌入，包括基因调控网络、启动子序列、基因家族和共表达。

fig1c

图1c：scCompass-126M 中人类和小鼠的器官类型

fig1d

图1d：从 scCompass-126M 中采样的子集的不同细胞类型的 UMAP。

目前大规模转录组预训练模型主要利用相对基因排序或分组基因表达值作为输入，导致转录组的表达信息不足。为了克服这一限制，我们将基因的绝对表达值和相应的基因id连接起来，然后根据整个scCompass-126M的归一化表达值在细胞内对基因进行排序。绝对表达式值的引入为GeneCompass提供了更细粒度的信息。为了进一步增强预训练模型的能力，作者通过不同的方法将启动子序列、基因共表达网络、基因家族和转录因子-靶基因调控网络等四种不同类型的生物先验知识编码到统一的嵌入空间中（图1b）。

为了表示物种信息，在每个细胞前嵌入一个species token（图1a）。总的来说，GeneCompass将基因id，表达值，先验知识组成基因token，利用12层Transformer进行编码。

GeneCompass采用掩码语言建模（mask language modeling, MLM）策略，在预训练过程中随机掩码基因token。具体来说，每个细胞随机 mask 15%的基因，然后GeneCompass通过两个独立的Decoder heads同时恢复这些被mask的基因id和表达值。这种多任务自监督学习范式结合了相对rank和绝对expression的恢复，增强了捕捉基因表达中微妙关系的能力。

方法：数据收集

作者构建了一个大规模的预训练语料库scCompass-126M。该语料库由超过1.2亿个单细胞转录组组成，包括两种不同的物种，人类和小鼠。多物种单细胞数据为理解不同生物的细胞异质性提供了丰富的资源。然而，由于物种之间的生物过程和技术差异，收集和预处理这些数据具有挑战性。

在所有物种中，人类和鼠的细胞比例最高，各自有超过5000万个细胞数据。研究使用的数据来自不同来源的公开数据集平台，包括NCBI-GEO、NCBI-SRA、EMBL-EBI-ArrayExpress和CNCB-GSA。从这些数据库中下载FASTQ文件，并通过运行相同的pipeline获得基因原始计数。其他基因原始计数数据直接从CELLxGENE、Single Cell Portal、Curated Cancer Cell Atlas (3CA)、Cell BLAST、Human Cell Atlas、Temporal Expression during Development Database (TEDD)中下载。对于准备多物种单细胞数据进行下游分析，执行了几个预处理步骤。对于质量控制，排除了低质量和受损的细胞。然后进行归一化和log来统一数值范围。

方法：GeneCompass

GeneCompass架构

GeneCompass采用12层自注意力Transformer对输入嵌入进行编码，每层自注意力Transformer由12个heads组成。GeneCompass对2048基因token序列处理为768维，模型参数的数量达到1亿。在预训练和微调过程中，采用高斯误差线性单元(Gaussian Error Linear Units, GELUs)作为非线性激活，注意力层和密集层的dropout概率均为0.02。模型配置、数据加载和训练的代码由Pytorch和Huggingface Transformers库实现，用于模型配置、数据加载和训练。此外，作者扩展了库，使其能够输入可扩展的外部知识。

GeneCompass预训练和优化

GeneCompass将基因ID、表达值以及相应的先验知识(启动子、GRN、基因家族和共表达)整合在一起，编码细胞转录组。采用掩码语言建模策略在预训练过程中随机mask基因，包括基因的id、表达和先验知识。MLM已被证明可以帮助大型模型有效地学习更好的表示。详细地说，每个细胞随机选择15%的基因来mask。与已有工作相比，GeneCompass构建了一种多任务学习范式，同时基于embedding来预测被mask基因的ID和表达。

使用MSE优化表达预测：
eq1
其中， $n_{unk}$ 表示未知基因表达的数量， $\widetilde{x}_{j}^{(i)}$ 是预测的基因表达，真实值是 $x_{j}^{(i)}$ 。

交叉熵用于预测基因ID：
eq2
这里的基因ID是mask基因的rank。

预训练参数演示如下：使用AdamW优化器将学习率设置为线性衰减，10000个预热步骤，最大学习率为1e-3，使用基因填充策略，批大小为10。为了充分利用GPU的能力，Deepspeed被应用在代码框架中，因为它提供了优化，利用动态批处理、混合精度、模型并行性和内存优化等技术来加速大型神经网络的训练，同时减少它们的内存消耗，所有这些都在类似pytorch的API中。使用每节点带有8个Nvidia A800 GPU的4节点服务器，预训练过程在9天内完成。

消融实验与不同规模的数据

使用不同的细胞数进行预训练对GeneCompass进行消融研究。具体来说，对于人类单细胞转录组语料库预训练的GeneCompass，使用不同的细胞数 (即 $5\times 10^{4},5\times 10^{5},5\times 10^{6},5.5\times 10^{7}$ )。通过小鼠单细胞转录组语料库预训练的GeneCompass，使用不同的细胞数 (即 $5\times 10^{4},5\times 10^{5},5\times 10^{6},4.5\times 10^{7}$ )。通过人-小鼠单细胞转录组语料库预训练的GeneCompass，使用不同的细胞数，即前面两个列表相加。为了使GeneCompass完全收敛，对数据规模小的实验使用较大的epoch。

微调

GeneCompass是编码器解码器架构，编码器有12个Transformer层，这些层由预训练的权重初始化，解码器是任务特定层。微调模型用于各种下游任务，如细胞类型注释，剂量-反应预测和基因调控网络推断。由于任务范式的不同，不同下游任务的微调超参数可能会有所不同。为了公平比较，作者在不同数据集的相同任务中使用相同的超参数（可能不是最优的），以验证GeneCompass的鲁棒性。

细胞和基因embedding

768维的细胞嵌入是从<cls> token生成的，该token被添加到单细胞信息的开头以表示细胞个体。此外，GeneCompass可以将每个基因编码到768个维度，这些维度包含了单细胞转录组中基因的上下文信息。注意，细胞和基因嵌入都是从最后一层Transformer中获得的。

先验知识embedding

将启动子序列、基因共表达网络、基因家族信息、转录因子-靶基因调控关系等四种生物学知识纳入预训练模型。

启动子embedding：启动子是基因转录的激活信号。在GeneCompass中，每个基因的启动子由2500个碱基组成，包括转录起始位点(TSS)前的上游500个碱基和TSS后的下游2000个碱基。启动子序列在DNABert预训练模型上微调40 epochs，得到768维的启动子embedding。

启动子序列来自基因组注释文件。

共表达embedding：高共表达基因对是指在没有外界干扰的情况下，其转录表达谱在多种器官或组织中高度相关的基因。因此，高共表达基因在理论上是相关和相似的。通常我们希望编码后，共表达基因对之间的距离尽可能近，非共表达基因对之间的距离尽可能远。具体来说，作者计算了每个基因对的Pearson相关系数(PCC)，然后选择PCC大于0.8的基因对进行gene2vec输出embedding。

基因家族embedding：同一基因家族中的基因具有相同的祖基因，因此基因在功能上是相似的。因此，当基因嵌入到特征空间时，同一家族的基因应该更接近，形成基因簇。然而，一个基因可能属于多个基因家族，并且基因家族之间可能存在重叠。为了处理这种复杂的关系，在构建训练样本时，作者使用gene2vec嵌入方法列出所有具有家族共属关系的基因作为基因对。Gene2vec可以根据基因家族的大小和所属的家族数自适应调整基因的频率，并在属于多个家族时充分考虑单个基因与其他基因之间的关系。作者使用了1645个人类基因家族和1539个小鼠基因家族进行嵌入。每个基因的嵌入是768维的。

GRN embedding：GRN是指细胞内或特定基因组内基因间相互作用形成的网络。在众多相互作用关系中，以基因间转录表达为基础的调控关系尤为重要。一般来说，一些基因可以通过转录表达来控制其他基因的表达水平。通常，GRN是一个有向图，每个基因代表一个节点，GRN中基因之间的控制关系用加权有向图来描述。目标是在已经发现的调控关系的基础上，尽可能完善调控网络，并了解更多基因之间未知的相互作用。在不同的组织细胞中，基因之间的调控关系不是静态的。目标是使具有更多调控关系的基因在嵌入空间中更加紧密。因此，作者从GRN中统计具有调控关系的基因对。这些基因对通过gene2vec方法进行嵌入。基因对之间的调控关系越多，在训练过程中出现的次数越多，编码后的关系越接近。每个基因的嵌入是768维的。

下游任务方法

单物种的细胞类型注释

微调的目标是通过利用GeneCompass生成的细胞嵌入来识别每个细胞的类型。具体来说，在class token添加一个全连接层，用于预测细胞类型。预测的细胞类型概率和GT之间的交叉熵损失被优化得尽可能低。

对于人类特定的任务，作者在人类多发性硬化症(hMS)、人类肺(hLung)和人类肝脏(hLiver)数据集上对GeneCompass和基于Transformer的方法Geneformer进行了基准测试。对于小鼠特异性任务，在小鼠脑(mBrain)、小鼠肺(mLung)和小鼠胰腺(mp胰腺)数据集上比较了经过和未经过预训练的GeneCompass结果。同时，在小鼠胰腺(mpPancreas)数据集上对GeneCompass与TOSICA方法进行了基准测试。