GENERator:长上下文生成式基因基础模型.
0. TL; DR
本文介绍了一种名为 GENERator 的生成式基因组基础模型,它采用了基于 Transformer 解码器架构,训练数据集包含 3860 亿个碱基对(bp)的真核 DNA 序列,具有 98k bp 的上下文长度和 12 亿参数。GENERator 在多个基准测试中展现出了最先进的性能,能够准确生成编码蛋白质的 DNA 序列,并在启动子设计等序列优化任务中展现出巨大潜力。该模型的出现为基因组研究和生物技术进步提供了新的工具,有望在精准医学和生物技术干预中发挥重要作用。
1. 背景介绍
基因组序列包含了驱动复杂生物过程(如基因调控和蛋白质合成)的大量生物信息,对表型特征和疾病状态产生影响。尽管 DNA 测序技术的快速进步极大提高了我们解码各种生物基因组序列的能力,但由于遗传物质的复杂性以及缺乏高质量的任务特定数据集,对基因组序列的预测和解释仍然是一个巨大挑战。
近年来,大型语言模型(LLMs)为理解生物序列提供了新的途径。在自然语言处理(NLP)领域,像 BERT 和 GPT 这样的 LLMs 通过大规模预训练数据展现了强大的泛化能力。然而,在基因组学领域,早期的努力主要集中在掩码语言模型(如 DNABERT、GROVER、Caduceus 和 Nucleotide Transformer),这些模型在理解 DNA 语义方面表现出色,但缺乏生成能力,并且通常受限于相对较短的序列长度。
为了解决现有模型的局限性,作者提出了 GENERator,一个基于 Transformer 解码器架构的生成式基因组基础模型。该模型在包含 3860 亿个碱基对的真核 DNA 数据集上进行训练,具有 98k bp 的上下文长度和 12 亿参数。GENERator 在多个基准测试中展现出了最先进的性能,能够准确生成编码蛋白质的 DNA 序列,并在启动子设计等序列优化任务中展现出巨大潜力。
2. GENERator 模型
GENERator 的训练数据来源于 RefSeq 数据库中的所有真核生物的原始 DNA 序列。作者探讨了两种数据处理策略:基因序列训练和全序列训练。基因序列训练策略从基因组序列中分离出基因区域,这些区域涵盖了多种功能,包括转录成各种 RNA 分子、翻译成复杂蛋白质以及控制基因表达的调控功能。全序列训练策略则直接将 RefSeq 中所有真核生物的基因和非基因 DNA 序列的混合数据输入语言模型进行训练。尽管全序列训练策略的预训练损失较低,但基因序列训练策略在多个下游任务中表现更好。
作者探讨了将 DNA 序列分词为适合输入语言模型的 token 的过程。常用的分词器包括单核苷酸分词器、K-mer 分词器和字节对编码(BPE)分词器。经过实验,作者发现 K-mer 分词器在下一个 token 预测(NTP)预训练中显著优于其他分词器,特别是 6-mer 分词器在实验中表现最佳。
GENERator 的预训练采用了与 Llama 类似的架构,使用 AdamW 优化器和余弦学习率调度器。预训练过程使用了 200 万 token 的批大小,总共进行了 6 个 epoch,处理了 3860 亿个 token。为了提高长上下文预训练的效率,作者使用了 Flash Attention 和 Zero Redundancy Optimizer 等优化技术。
GENERator的下游应用包括:
- 序列理解:GENERator 在多个序列分类任务中进行了评估,包括基因分类和分类任务。基因分类任务评估模型对短至中等长度序列(100 到 5000 bp)的理解能力,而分类任务则评估模型对长序列(10000 到 100000 bp)的理解能力。
- 中心法则:为了测试 GENERator 是否能够生成编码蛋白质的 DNA 序列,作者从 UniProt 数据库中选择了两个目标蛋白家族(组蛋白和细胞色素 P450 家族),并使用 GENERator 生成类似的蛋白质编码 DNA 序列。通过与自然家族的分布进行比较,验证了生成的 DNA 序列的稳定性。
- 启动子设计:作者使用 DeepSTARR 数据集中的启动子活性数据,通过监督微调(SFT)的方式训练 GENERator 生成具有特定活性轮廓的启动子序列。实验结果表明,GENERator 生成的启动子序列在预测活性和自然样本之间存在显著差异。
3. 实验分析
3.1 下一个 K-mer 预测
为了评估模型的生成能力,作者设计了下一个 K-mer 预测任务。在这个任务中,模型被输入一个序列片段,并被要求预测接下来的 K 个碱基对。预测的序列与实际序列进行比较以评估准确性。
作者通过训练多个使用不同分词器的模型来探索训练因果 DNA 语言模型的最佳分词器。实验结果表明,K-mer 分词器在下一个 K-mer 预测任务中表现优于 BPE 分词器,特别是 6-mer 分词器在有限的输入 token 下表现稳健,并且随着输入 token 数量的增加,其性能保持顶级。
3.2 基准测试评估
GENERator 在多个基准测试中与最先进的基因组基础模型进行了比较,包括 Nucleotide Transformer 任务、Genomic Benchmarks 和 Gener 任务。在 Nucleotide Transformer 任务中,GENERator 在多个数据集上表现出色。在 Genomic Benchmarks 中,GENERator 同样在多个任务中取得了最佳性能。在 Gener 任务中,GENERator 在基因分类和分类任务中均取得了最佳性能。
3.3 中心法则
作者评估了 GENERator 生成的组蛋白和细胞色素 P450 家族的 DNA 序列的质量。生成的 DNA 序列和翻译后的蛋白质序列的长度分布与目标家族的分布非常接近。通过 Progen2 计算生成序列的困惑度(PPL),发现其分布与自然家族的分布非常接近,并且与随机打乱的序列显著不同。此外,使用 AlphaFold3 预测生成的蛋白质序列的折叠结构,并使用 Foldseek 在蛋白质数据银行(PDB)中寻找类似结构,发现生成序列的构象与 PDB 中的已知结构高度相似。
3.4 启动子设计
作者使用 DeepSTARR 数据集中的启动子活性数据,通过监督微调(SFT)的方式训练 GENERator 生成具有特定活性轮廓的启动子序列。生成的启动子序列在预测活性方面与自然样本存在显著差异,表明 GENERator 在启动子设计任务中的潜力。