[논문 리뷰] Toward Understanding BERT-Like Pre-Training for DNA Foundation Models

Chaoqi Liang, Lifeng Qiao외 10명|arXiv (Cornell University)|2023. 10. 11.

Genomics and Phylogenetic StudiesBiochemistry, Genetics and Molecular Biology인용 수 3

한 줄 요약

이 논문은 DNA 서열에서 BERT 유사 모델을 위한 새로운 사전 훈련 전략인 RandomMask를 제안한다. 이 전략은 마스크 경계를 확장함으로써 점차 마스크의 난이도를 높여, 겹치는 K-mer 토큰화의 한계를 극복한다. 이는 빠른 수렴과 최적화가 부족한 문제를 야기한다. RandomMask는 7개의 하류 작업에서 28개 데이터셋 전반에 걸쳐 최신 기준 성능을 달성하였으며, 에피제네틱 마크 예측에서 65.83%의 Matthews 상관계수를 기록했다. 이는 기준 모델 대비 14.02% 향상되었고, 이전 최고 성능(SOTA) 대비 4.82% 높은 성능이다.

ABSTRACT

With the success of large-scale pre-training in language tasks, there is an increasing trend of applying it to the domain of life sciences. In particular, pre-training methods based on DNA sequences have received increasing attention because of their potential to capture general information about genes. However, existing pre-training methods for DNA sequences largely rely on direct adoptions of BERT pre-training from NLP, lacking a comprehensive understanding and a specifically tailored approach. To address this research gap, we provide the first empirical study with three insightful observations. Based on the empirical study, we notice that overlapping tokenizer can benefit the fine-tuning of downstream tasks but leads to inadequate pre-training with fast convergence. To unleash the pre-training potential, we introduce a novel approach called RandomMask, which gradually increases the task difficulty of BERT-like pre-training by continuously expanding its mask boundary, forcing the model to learn more knowledge. RandomMask is simple but effective, achieving state-of-the-art performance across 6 downstream tasks. RandomMask achieves a staggering 68.16\% in Matthew's correlation coefficient for Epigenetic Mark Prediction, a groundbreaking increase of 19.85\% over the baseline and a remarkable 3.69\% improvement over the previous state-of-the-art result.

연구 동기 및 목표

• 기존 BERT 유사 사전 훈련 방법의 한계를 DNA 서열 모델링에서 조사한다. 특히, DNA 고유의 특성을 고려하지 않은 채 직접 NLP 방법을 적용한 데서 기인하는 문제를 다룬다.
• 겹치는 vs. 겹치지 않는 K-mer 토큰화가 DNA 기초 모델의 사전 훈련 동역학과 하류 미세조정 성능에 미치는 영향을 이해한다.
• 겹치는 토큰화로 인해 발생하는 빠른 수렴과 부적절한 최적화 문제를 해결하기 위해 더 효과적인 사전 훈련 전략을 설계한다.
• 단순하면서도 효과적인 방법을 개발하여, DNA 서열에서 국소 뉴클레오티드 수준과 지역적 서열 수준의 패턴을 모두 학습할 수 있는 모델 능력을 향상시킨다.
• 에피제네틱 마크 예측, 전사 인자 결합, 프로모터 탐지 등 다양한 하류 DNA 작업에서 최신 기준 성능을 달성한다.

제안 방법

• 사전 훈련 중 마스크 구간의 크기를 점차 증가시키는 동적 마스킹 전략인 RandomMask를 도입한다. 처음에는 단일 토큰에서 시작하여 점차 더 긴 구간으로 확장한다.
• 더 풍부한 국소적 맥락 표현을 가능하게 하기 위해 겹치는 K-mer 토큰화(크기 K의 슬라이딩 윈도우, 스트라이드 1)를 적용하여 토큰 시퀀스를 생성한다.
• 모델이 점차 커지는 마스크 구간 내 원래의 뉴클레오티드를 예측하는 마스크 언어 모델링 목적을 사용하며, 이로 인해 과제 난이도가 시간이 지남에 따라 증가한다.
• 커리큘럼 학습 방식을 사용하여 마스크 경계를 무작위로 단계적으로 확장함으로써, 짧은 거리의 의존성에 대한 과적합을 방지한다.
• 표준 하이퍼파라미터와 평가 지표를 사용하여 7개의 하류 작업에서 28개 데이터셋에 대해 최종 모델을 미세조정한다.
• 동일한 훈련 및 미세조정 조건에서 RandomMask의 성능을 DNABERT, Nucleotide Transformer, HyenaDNA 등의 기준 모델과 비교한다.

실험 결과

연구 질문

• RQ1 겹치는 K-mer 토큰화는 DNA 기초 모델의 사전 훈련 동역학과 하류 미세조정 성능에 어떤 영향을 미치는가?
• RQ2 왜 겹치는 토큰화는 BERT 유사 사전 훈련에서 DNA 서열에 대해 빠른 수렴을 유도하지만 최적화가 부적절한가?
• RQ3 시간이 지남에 따라 과제 난이도를 점차 높이는 동적 마스킹 전략이 DNA 기초 모델의 표현 능력을 향상시킬 수 있는가?
• RQ4 RandomMask는 에피제네틱 마크 예측 및 전사 인자 결합과 같은 다양한 하류 DNA 작업에서 기존 사전 훈련 방법보다 얼마나 뛰어나게 성능을 높이는가?
• RQ5 RandomMask는 DNA 서열에서 국소 뉴클레오티드 수준의 변화와 더 넓은 지역적 서열 패턴을 포착하는 데 모델 능력을 어떻게 향상시키는가?

주요 결과

• RandomMask는 에피제네틱 마크 예측에서 Matthews 상관계수(MCC) 65.83%를 기록했으며, 이는 기준 모델인 DNABERT 대비 14.02% 향상되고 이전 최고 성능(SOTA) 대비 4.82% 높은 성능이다.
• 강화자 활성도 예측에서는 Pearson 상관계수(PCC) 69.56%를 기록했으며, 이는 기준 모델(68.43%)과 최고 성능 방법을 초월했다.
• 스プライ스 사이트 예측에서는 MCC 87.20%를 달성했으며, 기준 모델인 DNABERT(85.44%) 대비 1.76% 향상되었다.
• 전사 인자 예측에서는 쥐 데이터셋 기준 5.97% 향상되었고, 인간 데이터셋 기준 5.63% 향상되었다.
• 7개의 하류 작업에서 28개 데이터셋 중 26개에서 최신 기준 성능을 달성하여 일관되고 광범위한 성능 향상을 입증했다.
• 겹치는 토큰화로 인한 짧은 거리 의존성에 대한 과적합 문제를 효과적으로 완화하여, 더 깊고 더 강력한 표현 학습을 가능하게 하였다.