GENA-LM:长DNA序列的开源基础模型家族.

0. TL; DR

本文介绍了一种名为GENA-LM的开源基础DNA语言模型家族,专门用于处理长DNA序列。该模型基于Transformer架构,能够处理长达36,000个碱基对的输入,并通过引入循环记忆机制进一步扩展处理能力。GENA-LM提供了多种预训练模型版本,包括多物种和特定分类群的模型,展示了在多种复杂生物学任务中的应用潜力。

1. 背景介绍

DNA序列编码了生物体的遗传信息,其功能元素(如启动子、增强子和转录因子结合位点等)的识别对于理解基因调控和非编码基因组的进化至关重要。然而,由于DNA序列的复杂性和表观遗传编码的多样性,准确解码基因组序列仍然是一个重大挑战。近年来,随着人工智能技术的发展,尤其是深度学习方法的应用,基因组学领域取得了显著进展。然而,现有的机器学习方法在识别长程依赖关系方面存在局限性,而这种依赖关系在人类和其他真核生物的基因组中非常常见。

为了解决这一问题,研究者们开发了基于Transformer架构的模型,如DNABERTBigBird,这些模型在处理长序列方面表现出色,但仍然受到输入长度的限制。为了进一步扩展模型的输入长度并提高其在基因组学中的应用能力,本文提出了GENA-LM模型家族。GENA-LM通过结合字节对编码(BPE)和稀疏注意力机制,能够处理长达36,000个碱基对的输入序列,并通过循环记忆机制进一步扩展其处理能力。

2. GENA-LM 模型

GENA-LM基于Transformer架构,采用字节对编码(BPE)对DNA序列进行分词,将序列分割为长度不等的子序列(tokens),从而能够处理更长的输入序列。模型架构包括BERTBigBird两种变体,其中BERT模型具有12层或24层的Transformer结构,而BigBird模型则结合了稀疏注意力机制,能够处理更长的输入序列(长达36,000个碱基对)。

稀疏注意力机制通过限制注意力连接的数量,扩展了模型的输入长度。GENA-LM中的BigBird模型采用了稀疏注意力机制,允许模型处理长达36,000个碱基对的输入序列。这种机制通过滑动窗口或块对角线等预定义或学习的注意力模式,将注意力计算的复杂度从输入长度的二次方降低到线性。

为了进一步扩展模型的输入长度,研究者们引入了循环记忆Transformer(RMT)架构。RMT通过将输入序列分割为多个段落,依次处理每个段落,并通过记忆tokens传递信息,从而实现对长序列的高效处理。这种方法允许模型处理长达数十万个碱基对的输入序列,同时保持计算资源的线性扩展。

GENA-LM模型通过掩码语言建模(MLM)任务进行预训练,其中模型需要根据上下文预测被掩盖的token。预训练过程中,模型的输入序列被随机掩盖15%的tokens,并通过优化交叉熵损失函数进行训练。训练使用了AdamW优化器,并在多个NVIDIA A100 GPU上并行进行。

在预训练完成后,GENA-LM模型可以通过微调(fine-tuning)应用于特定的生物学任务。微调过程中,模型在特定任务的数据集上进行训练,并通过添加一个任务特定的输出层来预测目标变量。例如,在启动子活性预测任务中,模型的输出层是一个二分类器,用于预测给定序列是否包含启动子;而在染色质分析任务中,模型的输出层是一个多标签分类器,用于预测多个表观遗传标记的存在与否。

3. 实验分析

研究者们使用了多种数据集来训练GENA-LM模型,包括人类T2T基因组组装、多物种基因组数据以及特定物种的数据集(如酵母、拟南芥和果蝇)。数据预处理包括去除线粒体序列、分割基因组序列为“句子”和“文档”,并应用数据增强技术(如反向互补序列和随机移位)以增加数据多样性。此外,还引入了1000基因组计划中的单核苷酸多态性(SNP)数据,以模拟基因组的变异情况。

染色质分析

研究者们使用DeepSEA数据集评估了GENA-LM模型在染色质分析任务中的表现。DeepSEA数据集包含了919种细胞类型特异性的染色质特征,包括DNA酶I超敏感位点(DHS)、组蛋白修饰(HM)和转录因子结合位点(TF)。

实验结果表明,GENA-LM模型在预测DHS、HM和TF结合位点方面均优于现有的深度学习模型(如DeepSEABigBird)。此外,当输入长度扩展到8 kb时,GENA-LM模型在HM预测任务中的表现进一步提升。

启动子活性预测

研究者们还评估了GENA-LM模型在启动子活性预测任务中的表现。实验使用了EPDnew数据库中的启动子序列,并将其与非启动子序列进行对比。

结果表明,随着输入序列长度的增加,模型的预测性能显著提升。此外,多物种预训练模型(如gena-lm-bert-base-t2t-multi)在启动子活性预测任务中的表现略优于单物种模型(如gena-lm-bert-base-t2t),表明多物种数据的引入可以提高模型的泛化能力。

剪接位点注释

在剪接位点注释任务中,GENA-LM模型需要预测给定序列中的剪接供体位点和剪接受体位点。实验使用了SpliceAI数据集,该数据集包含了15 kb的输入序列,其中心区域为5 kb的目标序列,两侧各5 kb的上下文序列。

结果表明,GENA-LM模型在剪接位点注释任务中的表现接近于专门为此任务设计的卷积神经网络SpliceAI。此外,能够处理更长输入序列的模型(如gena-lm-bigbird-base-t2t)在该任务中的表现优于参数更多的模型(如gena-lm-bert-large-t2t),表明上下文信息对于剪接位点预测的重要性。

多物种应用

GENA-LM模型不仅在人类基因组分析中表现出色,还能够应用于其他物种的基因组分析。研究者们评估了GENA-LM模型在多个物种上的表现,包括酵母、拟南芥、果蝇和几种哺乳动物。

结果表明,GENA-LM模型在预测启动子活性、CTCF结合位点和H3K27ac修饰位点等任务中,能够有效地迁移到其他物种的基因组数据上,尽管其在人类基因组上进行了微调。此外,针对特定物种预训练的模型(如酵母、果蝇和拟南芥的模型)在各自物种的任务中表现优于基于人类数据预训练的模型,表明物种特异性预训练可以进一步提高模型的性能。