0. TL; DR

DNABERT-2 是一种基于 Transformer 的多物种基因组基础模型，它通过引入字节对编码（Byte Pair Encoding, BPE）来替代传统的 k-mer 分词方法，显著提高了模型的计算效率和样本效率。DNABERT-2 在多物种基因组理解任务上表现出色，与现有的最先进模型相比，在参数数量减少 21 倍、预训练 GPU 时间减少 92 倍的情况下，仍然能够实现相当的性能。此外，DNABERT-2 还引入了 Genome Understanding Evaluation（GUE）基准测试，涵盖了 9 个任务和 36 个数据集，用于全面评估基因组基础模型的性能。

1. 背景介绍

基因组学是现代生物学的核心领域之一，其目标是解析 DNA 序列中的语言规则，以理解基因调控机制。近年来，基于 Transformer 的基础模型在自然语言处理（NLP）领域取得了巨大成功，类似的模型也被应用于基因组学，例如 DNABERT 和 Nucleotide Transformer。然而，现有的基因组基础模型大多依赖于 k-mer 分词方法，这种方法虽然简单，但在计算效率和样本效率方面存在显著问题，限制了模型的扩展性。

k-mer 分词方法将 DNA 序列分割成固定长度的子序列（k-mer），但这种方法会导致信息泄露和计算冗余。为了克服这些限制，DNABERT-2 引入了 BPE 分词方法，它通过统计方法动态生成变长的分词，从而提高模型的效率和性能。此外，DNABERT-2 还引入了多种技术来克服输入长度限制、减少计算和内存开销，并增强模型能力。

2. DNABERT-2 模型

DNABERT-2 采用 SentencePiece 框架实现 BPE 分词。BPE 是一种基于统计的数据压缩算法，通过迭代合并最频繁共现的基因片段来构建词汇表。具体来说，BPE 的词汇表构建过程如下：

初始化词汇表，包含所有唯一的字符（A、T、C、G）。
在每次迭代中，找到最频繁的字符对（例如 “TA“），将其作为一个新的词加入词汇表，并将所有相同的片段替换为这个新词。
重复上述过程，直到词汇表达到目标大小。

BPE 分词方法不仅避免了 k-mer 分词中的信息泄露问题，还显著减少了序列长度，从而提高了计算效率。例如，对于长度为 $L$ 的 DNA 序列，BPE 分词后的序列长度约为 $L/5$。

DNABERT-2 基于 Transformer 编码器架构，并引入了以下改进：

Attention with Linear Biases (ALiBi)：ALiBi 通过在注意力计算中引入线性偏置来替代传统的学习位置嵌入，从而克服输入长度限制。具体来说，ALiBi 在每个注意力头中添加一个固定的偏置项，以惩罚距离较远的键-查询对。这种偏置项的计算公式如下。其中 $m$ 是一个固定的头特定常数，用于控制惩罚率。

\[\text{softmax}(q_i K + m \cdot [-(i-1), \dots, -2, -1, 0, -1, -2, \dots, -(L-1-i)])\]

FlashAttention：FlashAttention 是一种高效的注意力实现方法，通过优化 GPU 的内存访问模式，显著减少了计算时间和内存占用。FlashAttention 将键/查询/值矩阵分割成块，并逐步计算整个输入的 softmax，从而减少了对 GPU 高带宽内存（HBM）的读写操作。
Low-Rank Adaptation (LoRA)：LoRA 是一种参数高效的微调方法，通过低秩分解来减少可训练参数的数量。具体来说，LoRA 将权重矩阵的变化表示为两个低秩矩阵的乘积：$\Delta W = BA$，其中，$B \in \mathbb{R}^{m \times r},A \in \mathbb{R}^{r \times n}$，且$r \ll m, r \ll n$。这种方法显著减少了训练时间和内存使用。
GEGLU 激活函数：GEGLU 是 GLU 的一个变体，通过结合 GELU 激活函数和线性变换来提高模型性能。GEGLU 的计算公式如下，其中$\text{GELU}(x) = x \Phi(x)$，$\Phi(x)$是标准正态分布的累积分布函数。