Paper: DNABERT-2: Efficient Foundation Model and Benchmark For Multi-Species Genomes (ICLR-2024)

Reviewer comment (Novelty):
Novelty is questionable since BPE is a well-known technique. The use of FlashAttention, LoRA, and AliBi are also not new. So methodologically, it is hard to gauge its novelty.

Author response:
Thank you for your comments regarding the novelty of our work. We understand your concerns about the application of well-known techniques in our study. However, we would like to emphasize that the novelty of our work lies not in the invention of new techniques, but in their innovative application and combination to address specific challenges in genome language modeling, as well as provide solid and well-tailored foundation model and benchmark to this area.

1. While BPE is established in NLP, our work is among the first to explore the use of BPE for genome sequences. Instead of trivially applying it, we provide a comprehensive theoretical foundation (Section 2) and empirical evidence (Section 1 of our common response) to demonstrate the advantages of BPE in this context. We also conduct a comprehensive empirical analysis of the model’s performance and efficiency with different vocab sizes. We believe our work provides key insights and guidance for future research in this domain.
2. The novelty also lies in our unique combination of advanced techniques (FlashAttention, LoRA, and ALiBi) to overcome specific barriers in large genome foundational models. This combination, tailored to the unique challenges, demonstrates how existing tools can be synergistically applied to create a more efficient and effective model for genome analysis. And the pre-trained model itself can serve as an efficient yet powerful tool for countless downstream tasks.
3. Additionally, our work introduces the Genome Understanding Evaluation (GUE) benchmark, the largest and most comprehensive benchmark to date for genome foundational models that address the lack of standardized and comprehensive benchmarks in this field. 

In summary, while the individual techniques we employed are known, their application in the context of genome language modeling, the way they are combined to address specific challenges, the pre-trained DNABERT-2 model, and the introduction of a comprehensive new benchmark constitute significant novel contributions to our work. We believe these aspects collectively advance the field of genome foundational models and provide valuable insights and resources for future research.

Reviewer comment (Experiments):
Is there potential for cross-species information leakage? For instance, given the substantial overlap in genomes between humans and primates, the model might easily predict the masked token.

Author response:
Thank you for indicating this! We agree with the concern of cross-species information leakage. In an extreme case, if all the species have the same genome, than training on the multi-species genome is identical to training the model on a single species genome for lots of epochs, which leads to severe overfitting problems. This is one of the main reasons we chose to train DNABERT-2 on multi-species genomes. 

We can also find some evidence in existing works, for example, if we compare the Nucleotide Transformers variants trained on 1000G and multi-species genome, the one trained on the multi-species genome actually performs much better despite the size of the 1000G dataset being larger than the multi-species one. We think the diversity and comprehensiveness of pre-training data is essential to the model’s performance. There may exist more strategical ways of selecting different species to achieve better diversity and comprehensiveness and we will investigate more in this direction in the future.

Reviewer comment (Experiments):
How does this compare to HyenaDNA?

Author response:
Thanks for asking. HyenaDNA is a concurrent work of us. We have added it as a baseline and reported its results on the common results. In short, we think DNABERT-2 and HyenaDNA have different focuses. Based on our current evaluation, DNABERT-2 is more effective on different types of tasks, while HyenaDNA, as a convolution-based model, has the potential to handle much longer sequences given the same GPU memory. We will compare with it rigorously and add our comparison results and conclusion in the camera-ready version.

Reviewer comment (Experiments):
On page 7, the authors note that they utilize LoRA for NT but opt for full fine-tuning for DNABERT/DNABERT-2. However, in the methods section, LoRA is described as an integral part of the approach. This is somewhat perplexing.

Author response:
Sorry for causing the confusion. We did make LoRA available in our DNABERT-2 implementation but we did not use it in the final results since it leads to little efficiency improvements given the small size and efficiency of DNABERT-2 and does not affect the performances. We will re-structure the method section in the camera-ready version to avoid confusion. Thanks for pointing out this!

Reviewer comment (Experiments):
While the authors suggest further pre-training on GUE sequences, this might raise concerns regarding its ability to generalize to datasets with novel sequences. For a balanced comparison, it might be best if the authors refrain from additional pre-training on GUE sequences.

Author response:
Sorry for causing the confusion! We will explain it better in the camera-ready version. We evaluate two variants of the DNABERT-2 model, **DNABERT-2** and **DNABERT-2♢**, as shown in Tables 2, 3, and 4 in the paper. The first one is the model purely trained on the pre-training corpus (never seen GUE before), while the second one is achieved by further training the first one with MLM loss on GUE's training set. As shown in Table 2, without further pre-training on GUE, **DNABERT-2** already performs on par with **NT-2.5B-multi**. The fact that **DNABERT-2♢** outperforms **DNABERT-2** by 1 on the average score on GUE shows the benefit of further pre-training.

Reviewer comment (Evaluation):
Did the authors evaluate the sequence statistics of the GUE sequences in relation to the sequences from the pre-training corpus?

Author response:
Thanks for asking! The average sequence length during pre-training is about 700 bp. After adding the 8 datasets (as described in the common response), the sequence length in GUE ranges from 70bp to 10000bp. Based on our evaluations, DNABERT-2 performs consistently well across tasks with different sequence lengths, either the input is 10 times shorter or 15 times longer than the sequences it saw during pre-training. We will add this discussion in the camera-ready version, too.

Reviewer comment (Experiments):
The authors claim the method requires significantly less computational power and memory. Did they test the performance with a larger model size? If there wasn't a notable performance enhancement, it would be noteworthy to highlight this.

Author response:
Unfortunately, we do not have enough resources yet to train a much larger model. But we agree with your point on this. The phenomenon you described can imply the existence of some sort of "ceiling" in DNA language models. This is something we are interested in diving into in the future. However, given the fact that a small-scale in-domain pre-training can lead to big performance improvement, we think we are still far from the ceiling of DNA language models.

Reviewer comment (Experiments):
Have the authors assessed how the model's performance varies with different dataset sizes?

Author response:
Thanks for pointing out this! We agree it is important, so we design experiments to assess the model on this. We show the results in section 4 of our common response and will add it to the camera-ready version. Thanks!

Reviewer comment (Experiments):
Have the authors conducted ablation on FlashAttention, AliBi, and LoRA?

Author response:
Thanks for asking! As shown in Section 1 of the common response, we have done ablation studies on the effect of BPE.

The FlashAttention module, in theory, does not affect the final result at all, since it only changes the process of attention calculation. So we did not perform an ablation study on it. We are performing the ablation study on ALiBi and it may take a few days. We will share the results here is we manage to get it before the end of the discussion period, and we will add it to the camera-ready version, too. As for LoRA, we only perform an ablation study on its effects on Nucleotide Transformers, since NT are the only models we use LoRA. In Appendix A.4 we compare the results of our implementation of NT with LoRA with the results they reported in the paper to show our comparison with them is fair. We do not compare NT on LoRA vs full finetuning, since we do not have enough resources to fully finetune NTs.

Reviewer comment (Evaluation):
Deep learning models applied to one-hot encoded genomic sequences appear to have a much higher level of interpretability than those that utilize k-mer tokenization, and I expect this to be even worse for the BPE encoding used in this work. Unlike other areas of application, in computational biology applications, interpretability is a key factor in choosing a model.

Author response:
Thank you for highlighting the critical aspect of interpretability in computational biology applications. We acknowledge that models trained on one-hot encoded sequences offer a high level of interpretability, especially when it comes to understanding the interactions at the level of individual nucleotide bases. In fact, there does exist an inevitable trade-off between fine-grant interpretability and the model’s computational efficiency. 

We think BPE tokenization also offers a unique perspective on genomic sequence analysis. BPE tokenization groups frequently co-occurring nucleotide sequences, providing insights into how these 'groups'  play a role in the genome's function. For instance, consider a scenario where a mutation occurs. We may be able to understand how impactful a mutation is by looking at both the difference between the tokenized sequences before and after the mutation, and the model’s predictions/attentions among different tokens.

Reviewer comment (Experiments):
The statement "This underscores the importance of providing the model with adequate data, particularly when the model size is scaled up, and further highlights the inefficiency of overlapping k-mer tokenization. The comparison between DNABERT and NT-2500M-1000g exposes the sample inefficiency of non-overlapping k-mer tokenization. Despite being trained on 2.5 times more tokens, NT-2500M-1000g achieves a performance similar to that of DNABERT." could also be questioned — there are other differences between the models, so ascribing this to the difference in tokenization method is a stretch. To demonstrate the advantage of BPE tokenization, an experiment should be performed on two different versions of DNABERT2 — one with k-mer tokenization, and one with BPE tokenization. **The authors have addressed this point with a thorough ablation study**.

Author response:
Thank you for highlighting the importance of directly comparing the effectiveness of different tokenization methods in DNABERT2. We concur that an ablation study focusing on tokenization techniques is crucial for substantiating the superiority of BPE (Byte Pair Encoding) over k-mer tokenization. In response to your suggestion, we have conducted additional experiments with two distinct versions of DNABERT2 - one utilizing k-mer tokenization and the other employing BPE tokenization. As detailed in Section 1 of our common response, these experiments have yielded compelling results. The version of DNABERT2 with BPE tokenization demonstrates a marked improvement in performance over its k-mer counterpart. This outcome not only confirms the efficacy of BPE tokenization in the context of DNA language modeling but also provides robust empirical evidence to support our claim. We believe that these findings significantly bolster the validity of our approach and contribute valuable insights into the optimization of language models for genomic data analysis

Reviewer comment (Novelty):
Please compare your benchmark datasets with the recently published "Genomic benchmarks": Grešová, K., Martinek, V., Čechák, D. et al. Genomic benchmarks: a collection of datasets for genomic sequence classification. BMC Genom Data 24, 25 (2023). https://doi.org/10.1186/s12863-023-01123-8

Author response:
Thanks for indicating this! In short, GenomicsBenchmark is a concurrent work with an aim similar to us, to create a comprehensive and discriminative space to compare different genome analysis models. It contains 9 datasets from 3 species with input sizes ranging from 200 to ~4000, while GUE contains 36 datasets from multiple species with input sizes ranging from 70 to 1000. Compared to GUE, GenomicsBenchmark also provides a Python package for data loading and processing, which is very helpful and we plan to do this in the future. We will include and discuss this paper in the camera-ready version.

Reviewer comment (Experiments):
Perform an ablation study to quantify the contribution of BPE and ALiBi independently.

Author response:
Thanks for indicating this! We totally agree that the ablation study is important for the audience to understand the model. As shown in the common response, we have done an ablation study to compare BPE vs k-mer with the same DNABERT-2 architecture and shown that BPE is indeed more data-efficient. Thank you for this suggestion! We are performing the ablation study on ALiBi and it may take a few days. We will share the results here if we manage to get it before the end of the discussion period, and we will add it to the camera-ready version, too.

Reviewer comment (Reproducibility):
Explain why the code, data and pre-trained model could not be made public now.

Author response:
Sorry for the confusion. To avoid violating the double-blind reviewing process, we follow the common practice to describe our open-sourcing plan in this way. In fact, our code has been publicly available in the supplementary file since the paper submission time. The link to the model and data is also shared in the README file. We will add links to the non-anonymous DNABERT-2 resources in the camera-ready version since we are not allowed to add them now. Thanks for your support to the open-source community.

Reviewer comment (Evaluation):
Explain why mcc and f1 are used as the comparison metric for different tasks.

Author response:
As suggested by [1], MCC is more reliable than F1 in binary classification and in scenarios when the test set is unbalanced. Since most of our datasets are binary classification and are unbalanced, we use MCC in most cases. The only exception is the virus classification dataset, which is multi-task and perfectly balanced. Thus, we used F1, which we considered to be more common in this scenario. Thanks for indicating the confusion! Since MCC also works well in the multi-class classification case, we plan to always use it as the metric to avoid confusion and will explain it in the camera-ready version.

Reviewer comment (Writing):
Explain the benefit of further pre-training. The reviewer gets lost in understanding the sentence "This results in 0.41B training tokens..." right above section 5.3.

Author response:
Sorry for causing the confusion! We will explain it better in the camera-ready version. In general, if we consider there exists a distributional difference between the pre-training corpus and downstream task, further pre-training on the downstream task datasets helps the model to adapt tot the downstream domain. Here we also investigate the effect of further pre-training in the DNA domain.

We evaluate two variants of the DNABERT-2 model, **DNABERT-2** and **DNABERT-2♢**, as shown in Tables 2, 3, and 4 in the paper. The first one is the model purely trained on the pre-training corpus, while the second one is achieved by further training the first one with MLM loss on GUE's training set. The sentence "This results in 0.41B training tokens..." refers to the process of training DNABERT-2 on GUE's training set. Please kindly let me know if this is clearly explained.

As shown in Table 2, **DNABERT-2♢** outperforms **DNABERT-2** by 1 on the average score on GUE, showing the benefit of further pre-training on genome sequence modeling.

Reviewer comment (Experiments):
The input size for the proposed benchmark seems to be on the shorter side for genomics, potentially limiting its applicability to broader genomics problems.

Author response:
We appreciate your observation regarding the input size limitation in the original version of our Genome Understanding Evaluation (GUE) benchmark. You are correct in noting that the initial input sizes, ranging up to 1000 base pairs (bp), might not fully encompass the broader spectrum of genomics problems that involve longer genomic sequences.

The decision to initially focus on shorter input sizes was driven by practical considerations related to the capabilities of existing models and the computational resources available to us. Specifically, DNABERT's architecture inherently limits it to handling sequences no longer than 512 bp. Similarly, the Nucleotide Transformer is limited for inputs up to 6000 bp, and processing sequences of this length requires substantial memory (45-50GB) on each single GPU, which exceeded our resource capabilities.

However, recognizing the importance of including longer sequences to enhance the comprehensiveness and applicability of GUE, we have now incorporated 8 additional datasets with input sizes ranging from 5000 to 10000 bp. This expansion allows us to better represent the diversity of genomic sequence lengths encountered in real-world genomics problems.

Our updated results with these additional datasets demonstrate DNABERT-2's remarkable capability in handling longer sequences. Despite being trained on sequences up to 700 bp, DNABERT-2 shows impressive performance on sequences between 5k and 10k bp, outperforming the strongest baseline by a significant margin on all 8 of these datasets. This underscores DNABERT-2's robustness and versatility in modeling genomic sequences of varying lengths.

We agree that extra-long input sequences, such as those around 1 million bp, are also very meaningful to study. However, they present unique challenges that might require specifically designed architectures or different approaches. Our current benchmark focuses on sequence lengths that can be effectively handled by existing pre-trained models without necessitating specialized architectural modifications.

Reviewer comment (Experiments):
The benchmark's design appears constrained, lacking baseline models like CNNs and omits language model training from scratch, which could provide comparative insights.

Author response:
Thank you for your insightful feedback regarding the design of our benchmark and the selection of baseline models. We understand your concern that the initial version of our benchmark may have appeared constrained due to the focus on state-of-the-art (SOTA) models, potentially overlooking the insights that could be gained from including a broader range of baselines.

Recognizing the importance of a more comprehensive comparative analysis, we have expanded our set of baseline models in the revised manuscript. Specifically, we have added two recent CNN models from referenced papers [1] and [2]. These models represent a different approach to genomic sequence analysis and provide a valuable point of comparison to highlight the strengths and potential limitations of our approach.

Furthermore, to address the aspect of language model training from scratch, we have included an additional baseline: DNABERT-2 without pre-training. This allows us to directly assess the impact and effectiveness of pre-training in our model. By comparing DNABERT-2 without pre-training to the pre-trained version, especially when finetuned on the entire dataset versus finetuned on only 5% of the training data, we offer a clearer picture of how pre-training contributes to the model's performance. This comparison is detailed in Section 4 of our common response.

[1] Nguyen, Eric, et al. "Hyenadna: Long-range genomic sequence modeling at single nucleotide resolution." *arXiv preprint arXiv:2306.15794* (2023).
[2] Grešová, Katarína, et al. "Genomic benchmarks: a collection of datasets for genomic sequence classification." *BMC Genomic Data* 24.1 (2023): 25

Reviewer comment (Experiments):
While the paper is apt for an ML conference, there is a discernible deficiency in the depth of biological insights. Better downstream tasks, such as CAGE-seq prediction and so on.... (longer sequence context)

Author response:
We sincerely appreciate your constructive feedback regarding the depth of biological insights in our paper. We acknowledge the importance of incorporating biologically relevant downstream tasks, such as CAGE-seq prediction for transcriptional regulation and activity, to enhance the biological applicability and insights of our model. 

However, it is crucial to highlight that the primary goal of this paper was to tackle specific computational challenges in genome foundation models. Our focus was on enhancing the efficiency and effectiveness of these models in handling large and diverse genomic datasets. This approach was driven by a need for more computationally practical yet powerful tools that can manage the increasing scale and complexity of genomic data, a challenge that is pivotal for advancing genome research.

In developing the Genome Understanding Evaluation (GUE) benchmark, our emphasis was slightly different from providing in-depth biological insights. Our aim was to construct a benchmark that, for evaluation purposes, could reveal a clear distinction between the computational capabilities of various methods in genome modeling. To achieve this, we considered several metrics, including the difficulty of tasks, comprehensiveness of datasets, standardization of benchmarks, and applicability to existing genome foundational models. With this focus in mind, tasks and datasets that could best examine and showcase such computational strengths of these models were prioritized.

We recognize the importance of bridging the gap between computational methodology and biological application. In future works, we plan to incorporate more biologically oriented downstream tasks, such as CAGE-seq prediction and other tasks requiring longer sequence contexts and are actively working on it. We believe that incorporating these downstream tasks will not only enrich the biological insights provided by our model but also enhance its practical relevance and utility in real-world genomic research.

Reviewer comment (Writing):
In the introduction, can you clarify what you specifically mean by "genome language modeling"?

Author response:
In our study, the terms 'DNA language modeling' and 'genome language modeling' are used synonymously. By 'genome language modeling', we specifically mean the development of a language model tailored to genomic sequences. This involves leveraging language model techniques to enhance our understanding of the genome, such as predicting genome functionality and interpreting genomic data.

Reviewer comment (Novelty):
Following up on the theme, why was there no citation or reference to models like deepbind/deepSEA? For instance, the TF-DNA binding prediction from Wang et al., 2022, seems not a great citation? Not a genomics language modeling.

Author response:
We apologize for the oversight in not citing seminal works such as DeepBind and DeepSEA, which indeed represent significant advancements in the field. To rectify this, we will revise the statement 'Recent advancements in genome language modeling have demonstrated their superiority …' to 'Recent advancements in deep learning techniques for genomic data analysis have demonstrated their superiority ...' in the camera-ready version. This change more accurately reflects the scope of the referenced literature. We will also ensure to include and discuss DeepBind and DeepSEA appropriately, recognizing their critical contributions to the field.

Reviewer comment (Theory):
Given the unique structure and function of DNA, why is there a continued emphasis on tokenization in DNA language modeling?

Author response:
The concept of tokenization is pivotal in DNA language modeling, mirroring its importance in natural language processing (NLP). In NLP, tokenization involves converting text into smaller units (tokens) for language models to process. Similarly, in DNA language modeling, we tokenize long, continuous sequences of nucleotides into small chunks. This step is crucial as it transforms genomic sequences into a form amenable to language model analysis. The choice of tokenization method significantly influences the model's perception and interpretation of the genomic data. As such, our work delves into various DNA tokenization strategies, assessing their impact and providing valuable insights for future research in this domain. 

We have now also demonstrated BPE leads to better model performance than k-mer tokenizations (Ablation study 1 in common response), which may be attributed to the fact that BPE iteratively merges the most frequent co-occurring genome segments that may be biologically relevant. Such observation emphasizes the critical role of proper tokenization on more effectively applying language model techniques to genomic sequences.

Reviewer comment (Presentation):
I recommend adding the count of sequences for each dataset in Table 1 to provide a clearer understanding of the dataset sizes.

Author response:
We sincerely appreciate the reviewer’s suggestion. We completely agree and will add the counts to the camera-ready version.

Reviewer comment (Experiments):
Why weren't tasks involving longer sequences incorporated after introducing DNABERT?

Author response:
We appreciate the reviewer for raising this important question. We have addressed this in our responses to the weaknesses and have now added 8 additional datasets involving longer sequences.