大规模基因组基础模型涌现出准确和通用的DNA表示.

paper：Accurate and General DNA Representations Emerge from Genome Foundation Models at Scale

0. TL; DR

AIDO.DNA 是一个用于AI驱动的数字生物体中的DNA表示的预训练模块。作为一个拥有70亿参数的、仅编码器的Transformer模型，它在包含796个物种的106亿个核苷酸数据集上进行了训练。通过在保持4k核苷酸的短上下文长度的同时扩展模型规模，AIDO.DNA在与功能基因组学、合成生物学和药物开发相关的广泛监督、生成和零样本任务中显示出显著的改进。

1. 背景介绍

基因组是数十亿年进化和选择的产物，它通过简单的四种化学词汇编码了生命的所有系统。跨越进化时间线的基因组元件的保守性与其功能角色高度相关。然而，仅从保守性很难确定确切的功能。

近年来，大规模无监督预训练的基础模型在蛋白质序列预测方面取得了巨大成功，解决了包括蛋白质折叠和从头设计在内的难题。这些模型能够准确预测蛋白质的结构和功能，其方式类似于传统的多序列比对，但又能推广到新的序列。尽管取得了这些成功，但用于蛋白质基础模型的架构在应用于更广泛的基因组时，并未表现出同样的成功。

最近的研究表明，基因组模型性能的差距可能与模型能处理的总基因组上下文长度有关。因此，研究人员提出了新的架构来学习长程相互作用。然而，将上下文长度增加到数百万个核苷酸，其性能提升却出现了递减效应。回顾基因组学中深度学习模型的发展，可以发现BERT风格的、仅编码器的Transformer架构已被广泛并成功地应用。

然而，现有的DNA编码器在参数规模上仍远小于蛋白质模型。考虑到一个通用的基因组语言模型应该隐含地包含一个蛋白质语言模型，并能表示所有其他由DNA编码的分子功能，作者认为这些模型需要比以往尝试的更大的表示能力。

为了学习准确和通用的基因组功能表示，作者开发了AIDO.DNA，这是迄今为止最大的无监督预训练DNA编码器之一，参数规模达到70亿。作者明确地将上下文长度设定为相对较短的4000个核苷酸，并在单核苷酸分辨率下进行训练。通过在一系列基准测试中取得新的SOTA性能，AIDO.DNA为开发AI驱动的数字生物体迈出了重要一步。

2. AIDO.DNA框架

作者开发了3亿和70亿两个参数规模的AIDO.DNA模型。该模型基于BERT架构，并采用了单核苷酸分词和掩码语言建模（MLM）的训练目标。

2.1 预训练

为了检验表示能力是否是限制DNA语言模型发展的瓶颈，作者使用了与Nucleotide Transformer相同的数据集和划分方式。该数据集包含来自796个物种的106亿个核苷酸。

DNA序列的“语法”是高度局部化和上下文相关的。将多个核苷酸压缩成“单词”的k-mer分词方法，会使学习局部语法（如密码子）变得困难，并稀释了单核苷酸变异等高影响效应的表示。

为了最大限度地减少偏见并学习高分辨率的单核苷酸依赖关系，作者选择了与真实数据紧密对齐的字符级分词。词汇表包含5个字符：A, T, C, G, N。序列以[CLS]开头，以[EOS]结尾。上下文长度被设定为4000个核苷酸。

作者采用了BERT风格的、仅编码器的密集Transformer架构，并进行了一些现代化更新：使用SwiGLU、使用LayerNorms、使用旋转位置嵌入RoPE。RoPE对于DNA建模尤为合适，因为它捕捉的是相对位置信息，而不是绝对位置，这与DNA语法不依赖于绝对位置的特性相符。

为了训练70亿参数的庞大模型，作者使用了Megatron-LM框架，在256个H100 GPU上进行模型并行训练。采用了bfloat16混合精度训练和FlashAttention-2，以支持在大规模下对长上下文进行训练。采用经典的掩码语言建模目标。在每个序列中随机选择15%的词元进行更改：其中80%被掩码，10%被随机替换，10%保持不变。模型通过交叉熵损失进行优化。

2.2 评估

作者通过一系列涵盖监督、无监督和生成目标的实验，全面评估了AIDO.DNA的性能。

序列属性分类：在与转录调控和转录本处理相关的标准基准上，通过对预训练模型进行全参数微调来评估性能。
零样本变异效应预测：使用ClinVar和GnomAD数据集，在不进行任何微调的情况下，通过计算参考序列和变异序列嵌入的L2距离来预测变异的致病性。
基因组挖掘的依赖关系映射：通过一种计算机模拟诱变（in silico mutagenesis）的方法，计算基因组上任意两个核苷酸之间的功能依赖关系，这个方法仅依赖于预训练模型，无需任何监督信息。依赖值$e_{i,j}$定义为：

\[e_{i,j} = \max_{k,q \in \{A,T,C,G\}} \left| \log_2 \left( \frac{\hat{\text{odds}}(n_j=k | n_1, \dots, n_i=q, \dots, n_L)}{\hat{\text{odds}}(n_j=k | n_1, \dots, n_L)} \right) \right|\]

从基因组上下文预测基因表达：在一个包含1000万个合成启动子序列及其配对表达测量值的数据集上，微调AIDO.DNA来预测基因表达。
定向启动子生成：将基因表达预测任务反过来，通过对AIDO.DNA进行一个简单的掩码扩散语言建模目标的微调，来实现根据期望的表达水平条件性地生成启动子序列。

3. 实验分析

3.1 规模化提升迁移学习性能

在一系列来自人类、小鼠和酵母基因组的标准基准任务上，评估AIDO.DNA在序列属性预测方面的性能。

AIDO.DNA在这些基准测试上取得了新的state-of-the-art性能。特别是在转录因子（TF）结合位点预测任务中，作者发现对测试集序列进行几轮MLM预训练可以进一步提升性能，这与密集Transformer需要良好对齐的归纳偏置的最新发现一致。

这些结果表明，通过大规模预训练学习到的DNA功能表示，可以实现数据高效的基因组挖掘和实验优先级排序。

3.2 精细的单核苷酸分词使能准确的变异效应预测

在ClinVar数据集上进行零样本变异效应预测。

AIDO.DNA的单核苷酸分词策略使其能够识别对人类健康有巨大影响的变异，为单序列预训练模型在该任务上树立了新的SOTA。有趣的是，3亿参数的模型（AUROC=0.807）在该任务上的表现优于70亿参数的模型（AUROC=0.785）。作者推测，这可能是因为较小的嵌入维度使得与致病性相关的信号更加集中，信噪比更高。

3.3 AIDO.DNA在无监督条件下学习共进化基因组元件

在不进行任何微调的情况下，利用计算机模拟诱变和依赖关系映射策略，探索人类基因组中功能未知的区域。

该方法揭示了DNA-DNA之间丰富的、已知的和未知的相互作用。以KIF26B基因为例，作者不仅发现了已知的启动子，还发现了一个未被注释的、与部分编码区镜像的微卫星（minisatellite）。这类微卫星被认为与表达调控有关，但其具体机制尚不清楚。这些依赖关系仅从无监督的预训练中涌现，并且只有在最大的模型规模下才变得清晰。这表明AIDO.DNA为通过预测共保守性和共进化来进行基因组挖掘提供了强大的基础。

3.4 预训练使能高效的基因表达预测

在一个包含1000万个合成启动子序列及其表达测量值的数据集上，评估AIDO.DNA在数据稀缺情况下的迁移学习能力。

AIDO.DNA能够准确高效地从启动子序列预测基因表达。在迁移学习中，AIDO.DNA利用了预训练阶段学习到的高级功能特征（如启动子、增强子等）。仅用0.01%的数据（1000个样本），AIDO.DNA就能保留其在完整数据集上73%的性能。仅用100个样本，就能保留50%的性能。这凸显了预训练在数据稀缺场景下的巨大价值。

3.5 调控DNA序列的定向生成

将基因表达预测任务反转，通过微调AIDO.DNA，根据期望的表达水平来条件性地生成启动子序列。

作者的方法在从头生成序列时，其碱基恢复的准确率达到了0.528，显著高于随机生成的理论上限（0.45-0.48）。这表明模型学会了与特定表达水平相对应的启动子“语法”。

可视化结果显示，为高表达生成的启动子和为低表达生成的启动子在序列上具有明显的、符合生物学意义的模式。例如，高表达启动子几乎完全使用G来连接启动子支架，而低表达启动子则在5’端富集G，并带有一个嘈杂的交替TATA基序。