[论文解读] DNABERT-2: Efficient Foundation Model and Benchmark For Multi-Species Genome

Zhihan Zhou, Yanrong Ji+4 more|arXiv (Cornell University)|Jun 26, 2023

Genomics and Phylogenetic Studies被引用 154

一句话总结

DNABERT-2 使用 Byte Pair Encoding (BPE) 取代 k-mer 标记化，构建高效的多物种基因组基础模型。它在显著减少参数和计算量的情况下实现接近最先进水平的性能，并引入 Genome Understanding Evaluation 基准（GUE）。

ABSTRACT

Decoding the linguistic intricacies of the genome is a crucial problem in biology, and pre-trained foundational models such as DNABERT and Nucleotide Transformer have made significant strides in this area. Existing works have largely hinged on k-mer, fixed-length permutations of A, T, C, and G, as the token of the genome language due to its simplicity. However, we argue that the computation and sample inefficiencies introduced by k-mer tokenization are primary obstacles in developing large genome foundational models. We provide conceptual and empirical insights into genome tokenization, building on which we propose to replace k-mer tokenization with Byte Pair Encoding (BPE), a statistics-based data compression algorithm that constructs tokens by iteratively merging the most frequent co-occurring genome segment in the corpus. We demonstrate that BPE not only overcomes the limitations of k-mer tokenization but also benefits from the computational efficiency of non-overlapping tokenization. Based on these insights, we introduce DNABERT-2, a refined genome foundation model that adapts an efficient tokenizer and employs multiple strategies to overcome input length constraints, reduce time and memory expenditure, and enhance model capability. Furthermore, we identify the absence of a comprehensive and standardized benchmark for genome understanding as another significant impediment to fair comparative analysis. In response, we propose the Genome Understanding Evaluation (GUE), a comprehensive multi-species genome classification dataset that amalgamates $36$ distinct datasets across $9$ tasks, with input lengths ranging from $70$ to $10000$. Through comprehensive experiments on the GUE benchmark, we demonstrate that DNABERT-2 achieves comparable performance to the state-of-the-art model with $21 imes$ fewer parameters and approximately $92 imes$ less GPU time in pre-training.

研究动机与目标

• 确定基因组语言模型中 k-mer 标记化的瓶颈并提出更高效的标记化方法。
• 开发 DNABERT-2，具备可扩展到多物种基因组的架构和训练策略。
• 引入一个标准化、全面的基准（GUE），用于跨任务和跨物种的模型公平比较。
• 证明基于 BPE 标记化的多物种预训练在降低计算的情况下仍能提供出色性能。

提出的方法

• 用 Byte Pair Encoding (BPE) 通过 SentencePiece 对 DNA 序列进行标记化，取代重叠/非重叠的 k-mer 标记化。
• 结合 ALiBi（带线性偏置的注意力）以消除输入长度限制并实现长序列处理。
• 使用 Flash Attention 和低精度 LayerNorm 提高速度与内存使用效率。
• 在需要时应用 LoRA（低秩自适应）实现参数高效微调。
• 以 MLM 目标进行训练，掩码比例为 15%，最大序列长度为 128，使用 8 张 GPU，约 14 天。
• 在 GUE 基准上跨 7 个任务、28 个数据集进行评估，以评估跨物种的泛化能力。

Figure 1: Illustration of the drawbacks of k-mer tokenization. In the overlapping setting, information about a masked token is leaked by its adjacent tokens, while in the non-overlapping setting, adding/deleting one nucleotide base leads to a dramatic change in the tokenized sequence.

实验结果

研究问题

• RQ1 基于 BPE 的标记化是否能够在对基因组序列进行压缩的同时提高样本效率，相较于 k-mer 标记化？
• RQ2 包含多物种的预训练是否能在多样物种和任务上提升性能？
• RQ3 效率技术（ALiBi、Flash Attention、LoRA）如何影响基因组模型的训练/推理成本与可扩展性？
• RQ4 是否存在一个标准化的基准，用于在不同物种和任务之间可靠地比较基因组基础模型？

主要发现

• DNABERT-2 在 GUE 上的表现与现有最先进水平相当，但参数量少 21 倍、GPU 时间少约 56 倍。
• 在 GUE 的 28 个数据集中，DNABERT-2 平均在 23 个数据集上优于 DNABERT，平均提升 6 个绝对分数点。
• 基于 BPE 的标记化将序列长度缩短约 5×，并解决了 k-mer 标记化的信息泄露问题，提高了效率与鲁棒性。
• 多物种预训练（4 个物种，覆盖 135 个物种）提升了对非人类任务的性能，并能较好迁移到人类基因组分析。
• 在 GUE 数据上的额外领域内预训练在任务间保持了前两名的稳定性，尽管收益因任务而异。
• 由于 ALiBi 与 Flash Attention 的作用，DNABERT-2 在扩展到更长输入时仍保持高效。

Figure 2: Illustration of the BPE vocabulary constructions.