[論文レビュー] Toward Understanding BERT-Like Pre-Training for DNA Foundation Models

Chaoqi Liang, Lifeng Qiao+10 more|arXiv (Cornell University)|Oct 11, 2023

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固有の特性を考慮しないまま自然言語処理の適応に依存している点に焦点を当てる。
• 重複する対照的K-merトークン化と非重複のトークン化が、DNA基盤モデルの事前学習ダイナミクスおよび下流のファインチューニング性能に与える影響を理解すること。
• 重複トークン化に起因する高速収束と不十分な最適化の問題を解決し、より効果的な事前学習戦略を設計すること。
• DNA配列における局所的なヌクレオチドレベルのパターンと領域的な配列レベルのパターンを両方効果的に学習できる、シンプルだが効果的な手法を開発すること。
• エピゲノム的マーキング予測、転写因子バインディング、プロモーター検出を含む、多様な下流DNAタスクで最先端の性能を達成すること。

提案手法

• 事前学習中にマスクスパンのサイズを段階的に拡大する動的マスク戦略であるRandomMaskを導入する。初期段階では1つのトークンから始まり、徐々に長いスパンへと拡大する。
• 重複するK-merトークン化（サイズKのスライディングウインドウ、ストライド1）を適用して、より豊かな局所的文脈表現を可能にする。
• マスクド言語モデリングの目的関数を用い、段階的に拡大するマスクスパン内の元のヌクレオチドを予測させる。これにより、学習タスクの難易度を時間とともに増加させる。
• カリキュラム学習のアプローチを採用し、マスク境界をランダムに段階的に拡大することで、短距離依存性への過剰適合を防ぐ。
• 標準的なハイパーパramータと評価指標を用いて、7つの下流タスクにおける28のデータセットで、得られたモデルをファインチューニングする。
• 同じ訓練およびファインチューニング条件下で、DNABERT、Nucleotide Transformer、HyenaDNAなどのベースラインと、RandomMaskの性能を比較する。

実験結果

リサーチクエスチョン

• RQ1 重複するK-merトークン化は、DNA基盤モデルの事前学習ダイナミクスおよび下流のファインチューニング性能にどのように影響を与えるか？
• RQ2 なぜ重複トークン化は、DNA配列のBERT類似事前学習において、高速収束を引き起こすが、最適でないモデル最適化をもたらすのか？
• RQ3 難易度を時間とともに増加させる動的マスク戦略は、DNA基盤モデルの表現能力を向上させることができるか？
• RQ4 RandomMaskは、エピゲノム的マーキング予測や転写因子バインディングを含む多様な下流DNAタスクにおいて、既存の事前学習手法をどれほど上回るか？
• RQ5 RandomMaskは、DNA配列における局所的なヌクレオチドレベルの変化と、広範な領域的配列パターンを捉える能力をどのように向上させるか？

主な発見

• エピゲノム的マーキング予測において、RandomMaskは65.83%のMatthews相関係数（MCC）を達成した。これはベースラインのDNABERT比で14.02%向上、かつ以前のSOTA比で4.82%上回る。
• エンハンサー活性予測では、Pearson相関係数（PCC）が69.56%に達し、ベースライン（68.43%）およびSOTA手法を上回った。
• スプライスサイト予測では、MCCが87.20%に達し、ベースラインのDNABERT（85.44%）比で1.76%向上した。
• 転写因子予測では、マウスデータセットでベースライン比5.97%、ヒトデータセットで5.63%の性能向上を達成した。
• 7つの下流タスクにおける28のデータセットのうち26で、RandomMaskは最先端の性能を達成し、一貫性があり広範な向上を示した。
• 重複トークン化に起因する短距離依存性への過剰適合を効果的に緩和し、より深く、より強固な表現学習を可能にした。