[論文レビュー] 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は k-merトークン化をByte Pair Encoding (BPE)に置き換え、効率的な多種属ゲノム基盤モデルを構築します。パラメータ数を大幅に減らし、計算量を削減しつつ、最先端に近い性能を達成し、Genome Understanding Evaluation benchmark (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.

研究の動機と目的

• Identify bottlenecks of k-mer tokenization in genome language models and propose a more efficient tokenization approach.
• Develop DNABERT-2 with an architecture and training strategy that scales efficiently to multi-species genomes.
• Introduce a standardized, comprehensive benchmark (GUE) for fair cross-model comparison across tasks and species.
• Demonstrate that multi-species pre-training with BPE-based tokenization yields strong performance with reduced compute.

提案手法

• Replace overlapping/non-overlapping k-mer tokenization with Byte Pair Encoding (BPE) via SentencePiece for DNA sequences.
• Incorporate ALiBi (Attention with Linear Biases) to remove input-length limits and enable long-sequence processing.
• Use Flash Attention and Low Precision LayerNorm to improve speed and memory usage.
• Apply LoRA (Low-Rank Adaptation) for parameter-efficient fine-tuning when needed.
• Train with MLM objective and a 15% token masking, using a 128-length max sequence and 8 GPUs, ~14 days.
• Evaluate with GUE benchmark across 7 tasks and 28 datasets to assess cross-species generalization.

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 Can BPE-based tokenization provide both compression and better sample efficiency compared to k-mer tokenization for genome sequences?
• RQ2 Does multi-species pre-training with a modern genome foundation model improve performance across diverse species and tasks?
• RQ3 How do efficiency techniques (ALiBi, Flash Attention, LoRA) affect training/inference cost and scalability for genome models?
• RQ4 Is there a standardized benchmark to reliably compare genome foundation models across species and tasks?

主な発見

• DNABERT-2 achieves comparable performance to the state-of-the-art with 21× fewer parameters and ~56× less GPU time than NT on GUE.
• On average, DNABERT-2 outperforms DNABERT on 23 of 28 datasets in GUE, with an average improvement of 6 absolute score points.
• BPE-based tokenization reduces sequence length by ~5× and eliminates information leakage issues of k-mer tokenization, improving efficiency and robustness.
• Multi-species pre-training (4 species across 135 species) enhances performance on non-human tasks and transfers well to human genome analysis.
• Additional in-domain pre-training on GUE data provides top-two consistency across tasks, though benefits vary by task.
• DNABERT-2 remains efficient while scaling to longer inputs thanks to ALiBi and Flash Attention.]
• table_headers: ["Model","Num. Params. ↓","FLOPs ↓","Trn. Tokens","Num. Top-2 ↑","Ave. Scores ↑"]
• table_rows: [["DNABERT (3-mer)","86M","3.27","122B","2 | 0","61.62"],["DNABERT (4-mer)","86M","3.26","122B","0 | 1","61.14"],["DNABERT (5-mer)","87M","3.26","122B","0 | 1","60.05"],["DNABERT (6-mer)","89M","3.25","122B","0 | 1","60.51"],["NT-500M-human","480M","3.19","50B","0 | 0","55.43"],["NT-500M-1000g","480M","3.19","50B","0 | 1","58.23"],["NT-2500M-1000g","2537M","19.44","300B","0 | 1","61.41"],["NT-2500M-multi","2537M","19.44","300B","7 | 9","66.93"],["DNABERT-2","117M","1.00","262B","8 | 4","66.80"],["DNABERT-2 ◊","117M","1.00","263B","11 | 10","67.77"]]} is malformed and cannot be parsed as valid JSON. Please provide a well-formed JSON object.

Figure 2: Illustration of the BPE vocabulary constructions.