0. TL; DR

本文介绍了Nucleotide Transformer（NT），这是一个用于人类基因组学的强大基础模型。NT通过在大规模基因组序列上进行无监督预训练，能够学习基因组的结构和功能模式，并在多种下游基因组预测任务中表现出色。研究结果表明，NT不仅能够准确预测基因组学任务，还能通过零样本预测方法评估遗传变异的影响，为基因组学研究提供了一个高效且灵活的工具。

1. 背景介绍

基础模型（Foundation Models）在人工智能领域中扮演着重要角色，它们通常具有大规模的参数（数百万甚至数十亿），通过在大量数据集上进行训练，能够适应多种下游预测任务。在自然语言处理（NLP）领域，BERT和GPT等语言模型通过无监督学习在大规模未标记文本上进行预训练，已经取得了显著的成果。这些模型通过解决数十亿个填空测试（cloze tests）来全面理解语言，即预测句子中缺失的单词。受此启发，研究人员开始探索将类似的方法应用于生物序列（如蛋白质和DNA）的建模。

在基因组学中，DNA序列中的依赖模式对于理解基因组过程至关重要，从识别调控区域到评估单个变异的影响，都需要对基因组序列进行深入分析。近年来，深度学习（DL）模型被广泛用于从DNA序列中发现有意义的模式，例如预测基因表达、识别转录因子结合位点等。然而，现代基因组学研究产生了大量的数据，如何从这些未标记的数据中提取准确的信号是一个巨大的挑战。因此，训练大规模基础模型以处理核苷酸序列成为一个值得探索的方向。

本文提出了一种名为Nucleotide Transformer（NT）的基因组序列编码基础模型。研究人员构建了不同大小的NT模型（从5亿到25亿参数），并在人类参考基因组、3202个多样化人类基因组以及来自850个不同物种的基因组上进行预训练。通过在18个基因组预测任务上的系统性研究和基准测试，研究人员评估了这些模型的性能，并与其他DNA基础模型进行了比较。此外，研究人员还分析了模型的注意力图和困惑度，探索了预训练期间学习到的序列特征，并通过零样本评分评估了模型对功能重要遗传变异的影响。

2. Nucleotide Transformer 模型

Nucleotide Transformer（NT）是一个基于Transformer架构的DNA语言模型，能够从大规模未标记的基因组数据中学习通用的核苷酸序列表示。模型的核心是一个编码器，它将输入的DNA序列转换为高维嵌入向量，然后通过多层自注意力机制进行处理。

NT模型的架构如下：

输入表示：输入的DNA序列被分割成长度为6的核苷酸组合（6-mers），并被映射到一个固定大小的词汇表中。模型使用一个嵌入层将这些核苷酸组合转换为高维嵌入向量。
Transformer层：嵌入向量通过多层Transformer层进行处理。每一层包括一个多头自注意力机制和一个前馈网络（Feed-Forward Network, FFN）。自注意力机制允许模型捕捉序列中的长距离依赖关系，而FFN则对每个位置的嵌入向量进行非线性变换。
语言模型头：在预训练阶段，模型的输出通过一个语言模型头，预测被遮蔽位置的核苷酸。这个头是一个简单的线性层，将Transformer层的输出映射到词汇表的大小上，并使用softmax函数计算每个核苷酸的概率分布。

NT模型通过一个称为掩码语言建模（Masked Language Modeling, MLM）的任务进行预训练。具体来说，模型的目标是根据上下文信息预测被随机遮蔽的核苷酸。在训练过程中，模型的输入序列中随机选择15%的核苷酸位置，并将这些位置的核苷酸替换为一个特殊的掩码标记（MASK）。模型的任务是预测这些被遮蔽位置的核苷酸。预训练的目标是最小化预测分布与真实核苷酸之间的交叉熵损失：

\[\mathcal{L} = -\sum_{i \in \text{masked positions}} \log p(\text{true nucleotide}_i | \text{context})\]

在预训练完成后，NT模型可以通过微调（fine-tuning）适应多种下游基因组学任务。微调包括以下步骤：

替换语言模型头：将预训练模型的语言模型头替换为适合下游任务的分类或回归头。
冻结Transformer层：在微调过程中，Transformer层的权重被冻结，只更新新引入的分类或回归头的权重。
参数高效的微调技术：为了提高微调的效率，研究人员采用了参数高效的微调技术（如IA3），只引入少量的可训练权重。这些权重在Transformer层的自注意力机制和前馈网络中起到缩放作用，从而调整模型对下游任务的适应性。

除了微调，NT模型还可以通过零样本预测（zero-shot prediction）评估遗传变异的影响。零样本预测的核心思想是利用模型的嵌入空间来衡量变异前后序列之间的差异。具体来说，研究人员计算了变异序列和参考序列之间的嵌入向量距离（如余弦相似度、欧几里得距离等），并将这些距离作为变异效应的评分。这些评分可以用于预测变异的潜在致病性。

3. 实验分析

为了评估NT模型的性能，研究人员设计了一系列实验，包括：

下游任务的基准测试：在18个基因组预测任务上评估NT模型的性能，这些任务涵盖了基因表达预测、转录因子结合位点识别、增强子活性预测等多个方面。
与其他基础模型的比较：将NT模型与其他现有的DNA基础模型（如DNABERT、HyenaDNA和Enformer）进行比较。
零样本预测的评估：通过计算零样本评分，评估NT模型对遗传变异影响的预测能力。

下游任务的基准测试

NT模型在18个下游任务上的表现优于或至少与现有的监督学习模型相当。例如，在预测基因表达、转录因子结合位点和增强子活性等任务上，NT模型的性能与专门为此类任务优化的模型相当。具体来说，NT模型在多个任务上超过了BPNet等监督学习模型，尤其是在预测剪接位点和增强子活性方面表现出色。

与其他基础模型的比较

NT模型在大多数下游任务上优于或至少与DNABERT、HyenaDNA和Enformer等其他基础模型相当。特别是，NT的多物种模型（Multispecies 2.5B）在多个任务上表现最佳，这表明在多个物种的基因组上进行预训练可以提高模型对人类基因组任务的泛化能力。

零样本预测的评估

NT模型通过零样本评分能够有效预测遗传变异的影响。研究人员发现，基于余弦相似度的零样本评分与变异的严重性有较高的相关性。此外，NT模型在预测功能重要变异（如eQTLs、meQTLs、ClinVar和HGMD中的变异）方面表现出色，其性能与专门为此类任务训练的模型相当。

NT-v2

研究人员通过引入旋转位置编码（rotary embeddings）、改进的激活函数（swiGLU）、移除MLP bias 和 dropout以及扩展上下文长度至12 kb等技术，优化了NT模型的架构（NT-v2）。这些改进使得NT-v2在保持高性能的同时，显著减少了参数数量。例如，优化后的NT-v2模型在性能上与25亿参数的多物种模型相当，但参数数量减少了10倍。