Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Unlocking Efficiency: Adaptive Masking for Gene Transformer Models

Soumyadeep Roy, Shamik Sural, Niloy Ganguly

TL;DR

This work focuses on using curriculum masking where it systematically increase the difficulty of masked token prediction task by using a Pointwise Mutual Information-based difficulty criterion, as gene sequences lack well-defined semantic units similar to words or sentences of NLP domain.

Abstract

Gene transformer models such as Nucleotide Transformer, DNABert, and LOGO are trained to learn optimal gene sequence representations by using the Masked Language Modeling (MLM) training objective over the complete Human Reference Genome. However, the typical tokenization methods employ a basic sliding window of tokens, such as k-mers, that fail to utilize gene-centric semantics. This could result in the (trivial) masking of easily predictable sequences, leading to inefficient MLM training. Time-variant training strategies are known to improve pretraining efficiency in both language and vision tasks. In this work, we focus on using curriculum masking where we systematically increase the difficulty of masked token prediction task by using a Pointwise Mutual Information-based difficulty criterion, as gene sequences lack well-defined semantic units similar to words or sentences of NLP domain. Our proposed Curriculum Masking-based Gene Masking Strategy (CM-GEMS) demonstrates superior representation learning capabilities compared to baseline masking approaches when evaluated on downstream gene sequence classification tasks. We perform extensive evaluation in both few-shot (five datasets) and full dataset settings (Genomic Understanding Evaluation benchmark consisting of 27 tasks). Our findings reveal that CM-GEMS outperforms state-of-the-art models (DNABert-2, Nucleotide transformer, DNABert) trained at 120K steps, achieving similar results in just 10K and 1K steps. We also demonstrate that Curriculum-Learned LOGO (a 2-layer DNABert-like model) can achieve nearly 90% of the state-of-the-art model performance of 120K steps. We will make the models and codes publicly available at https://github.com/roysoumya/curriculum-GeneMask.

Unlocking Efficiency: Adaptive Masking for Gene Transformer Models

TL;DR

This work focuses on using curriculum masking where it systematically increase the difficulty of masked token prediction task by using a Pointwise Mutual Information-based difficulty criterion, as gene sequences lack well-defined semantic units similar to words or sentences of NLP domain.

Abstract

Gene transformer models such as Nucleotide Transformer, DNABert, and LOGO are trained to learn optimal gene sequence representations by using the Masked Language Modeling (MLM) training objective over the complete Human Reference Genome. However, the typical tokenization methods employ a basic sliding window of tokens, such as k-mers, that fail to utilize gene-centric semantics. This could result in the (trivial) masking of easily predictable sequences, leading to inefficient MLM training. Time-variant training strategies are known to improve pretraining efficiency in both language and vision tasks. In this work, we focus on using curriculum masking where we systematically increase the difficulty of masked token prediction task by using a Pointwise Mutual Information-based difficulty criterion, as gene sequences lack well-defined semantic units similar to words or sentences of NLP domain. Our proposed Curriculum Masking-based Gene Masking Strategy (CM-GEMS) demonstrates superior representation learning capabilities compared to baseline masking approaches when evaluated on downstream gene sequence classification tasks. We perform extensive evaluation in both few-shot (five datasets) and full dataset settings (Genomic Understanding Evaluation benchmark consisting of 27 tasks). Our findings reveal that CM-GEMS outperforms state-of-the-art models (DNABert-2, Nucleotide transformer, DNABert) trained at 120K steps, achieving similar results in just 10K and 1K steps. We also demonstrate that Curriculum-Learned LOGO (a 2-layer DNABert-like model) can achieve nearly 90% of the state-of-the-art model performance of 120K steps. We will make the models and codes publicly available at https://github.com/roysoumya/curriculum-GeneMask.
Paper Structure (19 sections, 2 equations, 3 figures, 6 tables, 1 algorithm)

This paper contains 19 sections, 2 equations, 3 figures, 6 tables, 1 algorithm.

Table of Contents

  1. Related Works
  2. Research Background
  3. Preprocessing
  4. Pretraining
  5. GeneMask: State-of-the-art Masking Strategy
  6. Proposed Methodology
  7. Global PMI Masking Strategy
  8. CM-GEMS: Curriculum Masking-based Gene Masking Strategy
  9. Experimental Setup
  10. Datasets
  11. Evaluation Setup
  12. Training Details
  13. Baseline Models
  14. Experimental Results
  15. Conclusion
  16. ...and 4 more sections

Figures (3)

  • Figure 1: Comparison of CM-GEMS and time-invariant strategy of Global with existing masking strategies
  • Figure 2: Performance comparison on GUE benchmark at reduced compute 1K pretraining steps and reduced model size (LOGO). (left) Human species, (right) Non-human species. The SoTA models are pretrained for 120K steps on the Human Reference Genome
  • Figure 3: Impact of doubling the masking rate of DNABert 10K model. The performance stays almost the same at 1K steps, but then it slowly decreases at 2K and 10K steps