[논문 리뷰] DNABERT-2: Efficient Foundation Model and Benchmark For Multi-Species Genome

Zhihan Zhou, Yanrong Ji외 4명|arXiv (Cornell University)|2023. 06. 26.

Genomics and Phylogenetic Studies인용 수 154

한 줄 요약

DNABERT-2는 k-mer 토큰화를 바이트 페어 인코딩(BPE)로 대체하여 효율적인 다종 종 게놈 기초 모델을 구축합니다. 더 적은 매개변수와 감소된 계산으로 최첨단에 필적하는 성능을 달성하고, 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 기반 토큰화로 다종 종 사전 학습이 계산량을 줄이면서도 강력한 성능을 낳는다는 것을 입증합니다.

제안 방법

• DNA 서열에 대해 SentencePiece를 통해 Byte Pair Encoding(BPE)로 중첩/비중첩 k-mer 토큰화를 대체합니다.
• 입력 길이 한계를 제거하고 긴 시퀀스 처리를 가능하게 하기 위해 ALiBi(Attention with Linear Biases)를 도입합니다.
• 속도 및 메모리 사용을 개선하기 위해 Flash Attention과 저정밀도 LayerNorm을 사용합니다.
• 필요 시 매개변수 효율적 미세조정을 위한 LoRA(Low-Rank Adaptation)를 적용합니다.
• MLM 목표와 15% 토큰 마스킹으로 학습하며 최대 시퀀스 길이 128, 8 GPUs를 사용하여 약 14일 간 학습합니다.
• 7개 태스크와 28개 데이터셋에 걸친 GUE 벤치마크로 교차 종 일반화 가능성을 평가합니다.

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에서 NT보다 21배 적은 매개변수와 약 56배 적은 GPU 시간으로 최신 수준의 성능에 상응합니다.
• 평균적으로 DNABERT-2는 GUE의 28개 데이터셋 중 23개에서 DNABERT를 능가하며 평균 점수 차이 6점 포인트 증가를 기록합니다.
• BPE 기반 토큰화는 시퀀스 길이를 약 5배 감소시키고 k-mer 토큰화의 정보 누출 문제를 제거하여 효율성과 로버스트성을 높입니다.
• 다종 종 사전 학습(135종에 걸친 4종)은 비인간 태스크의 성능을 향상시키고 인간 게놈 분석으로의 전이도 잘 됩니다.
• GUE 데이터에 대한 추가 도메인 내 사전 학습은 태스크 간 상위 2위 일관성을 제공하지만 태스크에 따라 이점이 다릅니다.
• ALiBi와 Flash Attention 덕분에 더 긴 입력으로 확장하더라도 DNABERT-2는 여전히 효율적입니다.

Figure 2: Illustration of the BPE vocabulary constructions.