DNABERT: 一个基于 Transformer 双向编码器表征的预训练 DNA 语言模型

news2024/10/6 19:22:38

本文结合 DNABERT 的原文，主要介绍了：

Overview of DNABERT
开发 DNABERT 的背景
DNABERT 的 tokenization
DNABERT 的模型架构
DNABERT 的预训练
基于微调 DNABERT 的应用

1. Overview of DNABERT

我们之前介绍了 BERT，它是一个基于 Transformer 双向编码器表征的预训练语言模型。（更多关于 BERT 的知识请查看：BERT: 一个双向理解语言序列的模型 & 第一个用于微调的预训练模型）

DNABERT 是一个基于 Transformer 双向编码器表征的预训练 DNA 语言模型，它是专门针对基因组 DNA 语言设计的预训练模型，它受到了 BERT 的启发，并且在 BERT 的基础上做了一定的修改，它和 BERT 都是采用”预训练—微调“范式的模型。

作者使用人类基因组大型无标记数据集对 DNABERT 进行预训练，基于上下游核苷酸背景获取对基因组 DNA 语言的通用理解。我们可以使用特定任务的小型带标记数据集对 DNABERT 进行微调，应用到各种基因组 DNA 语言分析任务上面。DNABERT 通过微调，在基因组调控元件预测方面取得了性能领先的表现，包括启动子、剪接位点和转录因子结合位点等预测。

DNABERT 发表在论文 DNABERT: pre-trained Bidirectional Encoder Representations from Transformers model for DNA-language in genome 中，原文链接：https://academic.oup.com/bioinformatics/article/37/15/2112/6128680?login=false。

DNABERT 的源代码、预训练模型和微调可以在 GitHub 找到：https://github.com/jerryji1993/DNABERT

By the way，DNABERT 微调代码用到的包存在版本 bug，看一看即可，不要浪费时间尝试了。如果想用 DNABERT 微调模型，可以使用迭代的 DNABERT-2，原文地址：https://arxiv.org/pdf/2306.15006，GitHub地址：https://github.com/MAGICS-LAB/DNABERT_2，我后续会介绍 DNABERT-2。

2. 开发 DNABERT 的背景

解读非编码区 DNA 语言是基因组研究的基本问题之一，但是由于多义性和遥远的语义关系的存在，基因调控非常复杂，传统的信息学方法往往无法捕捉到这些内容。

多义性和遥远的语义关系正是自然语言的关键属性，这启发了研究者借用自然语言模型来研究基因组 DNA 语言。

此外，将 DNA 语言建模应该考虑以下两点：

全局考虑所有上下文信息以区分多义的顺式作用元件；
获取对 DNA 语言的通用理解，以“预训练-微调”范式应用到各种基因组 DNA 语言分析任务上。

鉴于上述内容，作者将（基于 Transformer 双向编码器表征的预训练语言模型）BERT 的思想改编到了基因组 DNA 语言模型的训练中，具体做法包括：

应用了 Transformer 的双向编码器，使用自注意力机制从整个输入序列全局捕获上下文信息；
在人类基因组大型无标记数据集对模型进行预训练，获取对基因组 DNA 语言的通用理解。

3. DNABERT 的 tokenization

在介绍 DNABERT 的预训练之前，我们先介绍一下它的 tokenization。

DNABERT 使用 k-mer 表示法对 DNA 序列进行 tokenization，k-mer 表示法通过将每个脱氧核苷酸碱基与其后续碱基连接起来，包含了更丰富的上下文信息。

作者在本研究中预训练了 4 个模型：DNABERT-3、DNABERT-4、DNABERT-5、DNABERT-6，分别对应 3-mer、4-mer、5-mer、6-mer。

对于 DNABERT-k，它的词汇表（vocabulary）由 k-mer 的所有排列以及 5 个特殊 token 组成，总计个 tokens。

5 个特殊 token：[CLS] 代表分类 token，[PAD] 表示填充 token，[UNK] 表示未知 token，[SEP] 表示分句 token，[MASK] 表示掩码 token。因此， DNABERT-k 词汇表中的标记。

更过关于 k-mer 和 tokenization 的知识请查看：从头实现一个处理 fasta 文件的 tokenizer

4. DNABERT 的模型架构

DNABERT 和 BERT 都是基于 Transformer 编码器设计的，两者的模型架构基本相同。

如上图所示：

DNABERT 使用 k-mer 作为输入（上图以 3-mer 为例），[CLS] token 是代表整个序列含义的标签，[SEP] token 是序列分隔符，表示序列结束，[MASK] token 表示预训练中的掩码 k-mer。

输入的 token 在嵌入层 embedding 为 Token Embedding，并添加 Positional Embedding 信息，两者组成 Input Embedding。

Input Embedding 被输入到 Transformer 编码器块中，总共有 12 个 Transformer 编码器块，每块有 768 个隐藏单元和 12 个注意力头。

对于最后隐藏状态的输出，开头第一个输出用于句子级别分类，单个掩码 token 的输出将用于 token 级分类。

5. DNABERT 的预训练

DNABERT 借鉴了 BERT 的思想，但对 BERT 的预训练过程做了一些修改：

删除了下一句预测；
调整了序列长度，通过截断和采样从人类基因组中提取 10~510 bp 的序列；
调整随机掩码处理：对输入的 token 序列进行随机掩码处理，单个掩码区域大小为 k 个连续 token 区域（k 为 k-mer 的 k），总掩码区域大小占序列全长的 15%。这样做是为了强制模型预测 k 个连续 token 区域，k 个连续 token 区域是和 DNA 语言场景相符合的，例如，一个功能元件一般是一小段序列。

注意：在微调中，不对输入的 token 序列进行随机掩码处理，直接将 token 序列输入到 Embedding 层。

损失函数：使用交叉熵损失函数对 DNABERT 进行预训练。