使用预训练DNA语言模型进行植物基因组跨物种单核苷酸分辨率建模.

0. TL; DR

PlantCaduceus 是一个多功能的DNA语言模型,在16种多样化的被子植物基因组上进行了预训练。在基因注释任务中,PlantCaduceus的性能显著优于现有的最佳DNA语言模型。在变异效应预测中,它的表现与state-of-the-art的蛋白质语言模型相当。PlantCaduceus预测为有害的突变,其平均次要等位基因频率(MAF)比基于多序列比对的方法识别出的突变低三倍,表明其预测更准确。此外,PlantCaduceus成功识别了拟南芥和玉米中已知的因果变异。

1. 背景介绍

在过去的20年里,已有超过1000种植物的基因组被公布,这一数字在未来几十年还将显著增加。理解这些基因组中的功能元件及其对适应性的影响,对于推进植物基因组学和作物育种至关重要。然而,植物基因组在大小、组成和复杂性上表现出惊人的多样性,即使在近缘物种之间也存在巨大差异。

虽然为拟南芥、水稻和玉米等模式植物已经生成了丰富的基因组资源,但为所有植物基因组进行类似的实验性资源生成是耗时、昂贵且不切实际的。这凸显了开发能够捕捉跨物种进化保守性的、可迁移的计算模型的迫切需求。

有监督的深度学习(DL)序列模型在理解DNA序列功能方面取得了成功,但它们通常需要大规模的标记数据,而这在植物基因组学中往往是稀缺的。此外,在模式物种上训练的模型,向其他物种迁移时常常面临挑战。

语言模型(LMs)通过自监督预训练,可以从大量未标记的序列数据中学习生物学序列模式。这些预训练好的模型可以通过微调(fine-tuning)或零样本预测(zero-shot prediction)的方式,应用于下游任务。蛋白质语言模型已在预测蛋白质结构、识别致病变异等方面显示出巨大成功。然而,它们局限于编码区,而DNA语言模型则能提供全基因组的视角,这对于理解非编码区的功能至关重要。

然而,构建DNA语言模型也面临挑战,如基因组中大量的重复序列会干扰预训练,以及需要考虑DNA双链的反向互补(RC)对称性。为应对这些挑战,作者引入了PlantCaduceus,一个专门为植物基因组设计的DNA语言模型。

2. PlantCaduceus框架

PlantCaduceus是一个基于Caduceus架构的DNA语言模型,后者建立在Mamba(一种选择性状态空间序列模型)之上,并专为DNA序列进行了优化。

2.1 预训练数据与策略

作者在16种被子植物的基因组上对PlantCaduceus进行了预训练,这些物种涵盖了1.6亿年的进化历史。

受到GPN模型的启发,作者采用了一种特殊的策略来处理基因组中大量的重复序列和非编码区。通过对非编码区进行降采样和对重复序列进行降权,模型可以更专注于学习功能更相关的编码区和非重复区域的模式。

模型采用掩码语言建模(Masked Language Modeling, MLM)作为预训练目标。即在输入的512 bp序列中随机掩盖15%的核苷酸,然后训练模型来预测这些被掩盖的核苷酸。

2.2 Caduceus模型架构

PlantCaduceus的核心架构Caduceus有几个关键设计,使其特别适合DNA建模:

作者训练了一系列不同规模的PlantCaduceus模型,参数量从2000万到2.25亿不等。

2.3 下游任务的应用方式

在预训练之后,PlantCaduceus可以通过两种主要方式应用于下游任务:

  1. 特征提取 + 分类器:将预训练模型的权重冻结,将其最后一层的隐藏状态作为序列的嵌入表示。然后,在这个嵌入表示之上训练一个简单的分类器(如XGBoost或线性层)来完成特定任务。
  2. 零样本预测:不进行任何额外的训练,直接利用预训练模型预测参考等位基因与变异等位基因的对数似然比,从而评估突变的功能影响。

3. 实验分析

作者通过三项互补的分析,系统地评估了PlantCaduceus的性能。

3.1 基因注释任务的跨物种迁移能力

作者设计了四个基因注释任务:翻译起始位点(TIS)、翻译终止位点(TTS)、剪接供体位点和剪接受体位点的预测。模型在拟南芥的标记数据上进行训练,然后在包括水稻、高粱和玉米在内的其他五种植物上进行测试,评估其跨物种的泛化能力。

在跨物种预测场景中,几乎所有基线模型(包括GPN, AgroNT和有监督的CNN+LSTM模型)的性能都出现了显著下降。相比之下,PlantCaduceus在所有物种和所有任务中都保持了极高的预测性能(以AUPRC衡量),平均AUPRC达到0.764,远超其他模型。例如,在与拟南芥相距1.6亿年的玉米上,PlantCaduceus在剪接供体预测任务上的性能是GPN的1.45倍,在TIS预测任务上更是达到了7.23倍。这表明PlantCaduceus学习到了跨越广阔进化距离的、高度保守的基因组语法。

作者通过一系列消融实验证明,PlantCaduceus的优越性能归功于其独特的Caduceus架构(优于GPN的卷积架构)以及对RC对称性的内建处理。

3.2 进化约束预测

作者利用高粱及其近缘物种的全基因组比对数据,定义了基因组中的“保守”和“中性”位点。然后,在拟南芥数据上微调后的PlantCaduceus被用于预测高粱基因组中每个位点的保守性。

3.3 零样本变异效应预测

不经过任何微调,直接使用预训练的PlantCaduceus模型计算突变的“零样本分数”(zero-shot score),以评估其有害性。作者通过三种方式验证了这些分数的有效性:1)计算机饱和诱变;2)在自然群体中与等位基因频率的关联;3)识别已知的因果变异。

作者以玉米中一个著名的控制甜度的基因sugary1 (Su1)为例。尽管GWAS信号由于强连锁不平衡(LD)而模糊不清(图A),但PlantCaduceus的零样本分数能够准确地从众多候选变异中 pinpoint 出唯一的、已知的因果突变W578R(图B, C, D)。