DNABERT-S: Pioneering Species Differentiation with Species-Aware DNA Embeddings

Q4: Performance in unbalanced data

We agree that the relative abundances of those genomes is very important to model's performance. So we choose the raw samples from CAMI2 as our evaluation data, which already contains largely unbalanced data. We use the 0/25/50/75/100 percentile of number of sequences in each species as the data balance statistics and present them in the table below.

As the data is already very unbalanced, and the robustness of clustering results largely depends on the clustering algorithm, we validate the model's robustness to data balancing with K-means clustering and consider 10 datasets used in our clustering & classification evaluation.

For each dataset, we kept 100 species and evaluated species clustering with 3 cases.

• Case 1: Balanced. We keep 100 sequences in each species.
• Case 2: Less balanced. We keep 100 sequences in the first 10 species, 90 sequences in the next 10 species, 80 sequences in the next 10, ...., 10 sequences in the last 10.
• Case 3: Very unbalanced. We keep 100 in the first species, 99 in the second species, ... , 1 in the last species

We then set K=100 for K-means and use Adjusted Random Index (ARI) as the clustering metrics. We ran each experiment with 5 random seeds and report the mean and std of the runs. As shown in the Table, DNABERT-S demonstrates relatively robust performance as the data goes from purely balanced to largely unbalanced.

We really appreciate your reviews and suggestions. We hope our reply can solve your concerns. Please don't hesitate to share any other thoughts.

Q2: Performance of other baselines in classification.

Thanks for asking this. We agree it is also beneficial to include the results of other models in the appendix. We have organized the data and integrated the statistics in Table 14 (page 23) in the revised version. Please see below for the statistics on the first datasets of synthetic, marine, and plant-associated datasets.

Q3: How does Manifold Instance Mixup differ from Manifold Mixup

Thanks for your comment. Here are some clarifications.

1. Manifold Mixup [Verma19] is a regularization method that linearly interpolates between hidden states and labels of different data samples at a randomly selected layer. It improves deep neural network representations by training on these interpolations. This technique smoothens decision boundaries, flattens class-specific representations, and promotes less confident predictions on unseen data.
2. Manifold I-Mix adapts Manifold Mixup for contrastive learning by applying it to the anchor set. It assigns continuous values between 0 and 1 to indicate the similarity between sequences (for both positive and negative pairs). This approach is particularly effective for species differentiation, teaching the model nuanced similarities rather than binary classifications. For example, humans are more similar to monkeys than to viruses. We aim for the DNABERT-S embeddings not only to segregate different species but also to reflect the relative similarities among them (placing humans closer to monkeys and further from viruses). Therefore, the regression-like nature of Manifold I-Mix method makes it ideal for our problem.

[Verma19] Manifold Mixup: Better Representations by Interpolating Hidden States, ICML 2019

W3 & Q1: Missing comparison with similar models like COMEBin

We appreciate this concern and would like to clarify that DNABERT-S serves a distinct purpose from complete metagenomics binning methods, making direct comparisons inappropriate. Modern metagenomics binning methods (e.g., MetaBat2 [Kang19] and COMEBin [Wang24]) typically follow a five-step pipeline utilizing three data types:

1. Generate DNA embeddings from sequences (contigs)
2. Extract abundance information (alignment files)
3. Combine DNA embeddings and abundance information to generate final embeddings
4. Perform clustering
5. Refine results using external features (e.g., length, single-copy gene markers)

DNABERT-S specifically targets step 1, rather than end-to-end metagenomics binning. To evaluate its effectiveness, we implemented step 4 using the straightforward approach detailed in Algorithm 1, comparing DNABERT-S against other methods for step 1 only. We deliberately omitted steps 2, 3, and 5 to isolate and assess the impact of DNA embeddings alone.

Consequently, compared to complete binning methods like COMEBin, our approach uses only one-third of the available information (sequences only, without alignments or external features) and employs a simplified clustering algorithm. DNABERT-S is designed to enhance existing binning methods by replacing their step 1 (TNF), which explains our focus on TNF as the primary baseline. As discussed in our response to W4, despite using significantly less information and a simpler clustering approach, we achieve comparable binning performance to SOTA methods on CAMI2, demonstrating DNABERT-S's effectiveness.

[Kang19] MetaBAT 2: an adaptive binning algorithm for robust and efficient genome reconstruction from metagenome assemblies, PeerJ 7, 2019

[Wang24] Effective binning of metagenomic contigs using contrastive multi-view representation learning, Nature Communications, 2024

W4 & Q5: Compare with tools directly on CAMI2 benchmark

We appreciate this suggestion. As explained above, our implementation utilizes only a subset of available information and omits three steps from the standard pipeline, making direct comparisons with existing binners not entirely appropriate.

Nevertheless, our results are competitive with existing methods. We have compared our approach against official results from SOTA models on the CAMI2 benchmark, including a baseline implementation using TNF (the default DNA embedding method in most metagenomics binners).

Results for the 'Plant 5' and 'Plant 6' datasets.

Results for the 'Marine 5' and 'Marine 6' datasets.

Our comparisons yield three key insights:

1. DNABERT-S largely outperforms TNF, demonstrating its superiority as a DNA embedding method
2. The performance gap between TNF and complete binning solutions highlights the importance of integrating additional features (e.g., abundance information)
3. Despite our simplified implementation using limited information, we achieve comparable performance to comprehensive binning methods

While integrating DNABERT-S embeddings into existing metagenomics binners (replacing TNF) would require substantial engineering effort, our results suggest this would be a promising direction for future work.

Thank you very much for your insightful review and suggestions!

W1: Number of parameters and embedding dimensions for each models/techniques used in comparision

Thanks for pointing this out. We acknowledge that the number of parameters and embedding dimensions are important parts in the comparsion.

We have included a table in the latest revised version that compares the number of parameters (million), embedding dimension, inference time (seconds), and inference memory (MB) in Appendix E. The symbol "-" denotes that the inference time or memory is not comparable with the LLM-based models.

W2: Utility of $\text{C}^{2}$LR and MI-Mix & Computational tradeoffs

Thanks for the suggestions.

(i) There is no extra computational cost when incorporating C$^2$LR since the computation costs of MI-Mix and Weighted SimCLR are nearly the same. For DNABERT-S, we train it with SimCLR for 1 epoch and MI-Mix for 2 epochs. For variants without curriculum learning, we train them for 3 epochs with the same loss function.

(ii) Intuition for our proposed curriculum contrastive learning ($\text{C}^{2}$LR) and MI-Mix.

Weighted SimCLR teaches the model whether two sequences are similar or not (a binary outcome of 0 or 1), whereas MI-Mix teaches the model how similar two sequences are, assigning a continuous value between 0 and 1.

MI-Mix method is especially suitable for species differentiation because species similarities are more nuanced than a binary classification. For example, humans are more similar to monkeys than to viruses. We aim for the DNABERT-S embeddings not only to segregate different species but also to reflect the relative similarities among them (placing humans closer to monkeys and further from viruses). Therefore, the regression-like nature of MI-Mix method makes it ideal for our problem.

However, predicting these finer-grained similarities is challenging, which is why we introduce Curriculum Contrastive Learning. We view Weighted SimCLR as a warm-up phase in model training, where the model first learns to segregate different species, and then we fine-tune it to adjust the distances in a more fine-grained manner.

(iii) Further cases study on the benefit of MI-Mix in genomics data.

Based on our motivation behind the method design, we conduct a case study to empirically validate our intuition. Specifically, we collect 50 5000-bp genome sequences from 3 species: human, monkey, and a randomly selected bacteria named Salmonella enterica. We compute the embedding of each genome sequence, and achieve the species embedding by averaging the embedding of all its 50 sequences. We then compute the cosine distance between human-monkey (H-M), human-bacteria (H-B), and monkey-bacteria (M-B). We then compute the relative distance between humans and bacteria (H-B/H-M) and monkeys and bacteria (M-B/H-M). As shown in the table below, models trained with MI-Mix loss naturally segregate very dissimilar species like humans and bacteria further while keeping similar species like humans and monkeys closer. We observe the same pattern in several different bacteria. These case studies can be an illustration of why MI-Mix is more suitable than SimCLR for metagenomics data, besides the scores in the ablation study. We will conduct more comprehensive case studies and discuss them in the revised version. Thanks for this suggestion.

Q2: Performance of MI-Mix on the original I-Mix tasks

Thank you for your question. We conduct experiments on the CIFAR-10 dataset in the I-Mix tasks with two separate models: ResNet-18 and ResNet-50.

Following a similar setting as outlined in the I-Mix tasks (Sec. 4 in [Lee23]), we undertake contrastive representation learning on a pretext dataset and assess the quality of representations through supervised classification on a downstream dataset. In all experiments, we train the model for varying epochs for contrastive learning, specifically 50 or 100. For supervised learning, we train the model for 50 epochs across all settings. We employ the N-pair [Sohn16] as the base contrastive learning method. We integrate I-Mix or MI-Mix with N-pair to compare their performances. The results are presented as follows.

Table 1: Results for ResNet-18 on the CIFAR-10 Dataset

Table 2: Results for ResNet-50 on the CIFAR-10 Dataset

As shown in the table, I-Mix and MI-Mix achieve very similar performance on images. This is expected as MI-Mix is specifically motivated by the genomics context. I-Mix is less suitable for genomic sequences since sequences from different species may share common segments. If the embedding mixup happens at the beginning of the model, where no contextual information is involved, it becomes very challenging for the model to distinguish the source of a sequence (whether the common segment comes from species A or B). Consequently, the model can become confused when species share common segments. By mixing at an intermediate layer, the common segments incorporate contextual information, allowing for better differentiation of closely related species during model training.

Yet in visions tasks, it is very unlikely that two images will share common segments. Thus, mixing at intermediate layers (MI-Mix) does not have much benefit over mixing at the begining (i-Mix).

[Lee23] i-Mix: A Domain-Agnostic Strategy for Contrastive Representation Learning, ICLR 2023. [Sohn16] Improved Deep Metric Learning with Multi-class N-pair Loss Objective, NeurIPS 2016.

Thanks a lot for your comments and suggestions. We hope our response can help solve your concerns with this work. Please don't hesitate to share any other thoughts!

Thank you very much for your detailed and insightful review!

W1: Intuition for Curriculum Contrastive Learning

Thank you for pointing this out. MI-Mix indeed performs very well on its own. The reason we choose to use curriculum contrastive learning in this work it that: $\text{C}^{2}$LR does not involve any extra computational costs or engineering efforts, while lead to slightly better performance. As a foundation model for species differentiation, we aim to get the best possible results. Nevertheless, we agree that using MI-Mix alone is good choice for conceptual simplicity in genome representation learning. We will further discuss this in the revised version.

The intuition behind the proposed curriculum contrastive learning ($\text{C}^{2}$LR) and MI-Mix are:

1. Weighted SimCLR teaches the model whether two sequences are similar or not (a binary outcome of 0 or 1), whereas MI-Mix teaches the model how similar two sequences are, assigning a continuous value between 0 and 1.
2. MI-Mix method is especially suitable for species differentiation because species similarities are more nuanced than a binary classification. For example, humans are more similar to monkeys than to viruses. We aim for the DNABERT-S embeddings not only to segregate different species but also to reflect the relative similarities among them (placing humans closer to monkeys and further from viruses). Therefore, the regression-like nature of MI-Mix method makes it ideal for our problem.
3. However, predicting these finer-grained similarities is challenging, which is why we introduce Curriculum Contrastive Learning. We view Weighted SimCLR as a warm-up phase in model training, where the model first learns to segregate different species, and then we fine-tune it to adjust the distances in a more fine-grained manner.

W2: Baselines with species differentiation training

Thanks for indicating this. Yes, only DNABERT-S has gone through species differentiation training among all the models in previous Table 1. We have results of other models with species differentiation training, such as HyeneDNA trained with our proposed data construction and Weighted SimCLR. We have included it in the Table 1 in the revised version to better reflect it.

W3: Discussion on original mixup

Thanks for your suggestions.

1. Manifold Mixup [Verma19] is a regularization method that linearly interpolates between hidden states and labels of different data samples at a randomly selected layer. It improves deep neural network representations by training on these interpolations. This technique smoothens decision boundaries, flattens class-specific representations, and promotes less confident predictions on unseen data.
2. Manifold I-Mix adapts Manifold Mixup for contrastive learning by applying it to the anchor set. It assigns continuous values between 0 and 1 to indicate the similarity between sequences (for both positive and negative pairs). This approach is particularly effective for species differentiation, teaching the model nuanced similarities rather than binary classifications. For example, humans are more similar to monkeys than to viruses. We aim for the DNABERT-S embeddings not only to segregate different species but also to reflect the relative similarities among them (placing humans closer to monkeys and further from viruses). Therefore, the regression-like nature of Manifold I-Mix method makes it ideal for our problem.

[Verma19] Manifold Mixup: Better Representations by Interpolating Hidden States, ICML 2019

W4: Typo in Figure 4's layout

Thank you for bringing this to our attention. We apologize for the oversight. We have modified it in the revised version to optimize the layout.

Q1: Any baseline model see the same labels

Thanks for your question.

1. In previous Table 1, none of the listed baselines utilize species differentiation labels, as we consider only existing DNA embedding methods for baselines.
2. We have undergone species differentiation training, including HyenaDNA with Weighted SimCLR (Sec. 5.6) and various DNABERT-S variants trained with different loss functions (Sec. 5.5). We primarily view these models as ablation studies focusing on the base model and training loss functions. As discussed in W2, we have also included HyenaDNA with Weighted SimCLR as a baseline in Table 1 in the revised version.

Thank you very much for your detailed and insightful review! Your comments and suggestions are very helpful in improving the quality of our manuscript.

W1: Novelty

Thank you for sharing your perspective on this matter. We agree that providing further illustration behind our method design is important.

1. SimCLR or I-Mix? In summary, SimCLR teaches the model whether two sequences are similar or not (a binary outcome of 0 or 1), whereas I-Mix teaches the model how similar two sequences are, assigning a continuous value between 0 and 1. I-Mix methods are especially suitable for species differentiation because species similarities are more nuanced than a binary classification. For example, humans are more similar to monkeys than to bacterias. We aim for the DNABERT-S embeddings not only to segregate different species but also to reflect the relative similarities among them (placing humans closer to monkeys and further from bacterias). Therefore, the regression-like nature of I-Mix method makes it ideal for our problem. However, predicting these finer-grained similarities is challenging, which is why we introduce Curriculum Contrastive Learning. We view SimCLR as a warm-up phase in model training, where the model first learns to segregate different species, and then we fine-tune it to adjust the distances in a more fine-grained manner. In our reply to your W5, we provide a case study to empirically validate this.
2. MI-Mix vs. I-Mix. I-Mix is less suitable for genomic sequences since sequences from different species may share common segments. If the embedding mixup happens at the beginning of the model, where no contextual information is involved, it becomes very challenging for the model to distinguish the source of a sequence (whether the common segment comes from species A or B). Consequently, the model can become confused when species share common segments. By mixing at an intermediate layer, the common segments incorporate contextual information, allowing for better differentiation of closely related species during model training.

Besides the above methodology, our novelty also lies at data construction and benchmark design.

• Data Construction. Unlike well-explored areas like NLP and CV, data construction for DNA representation is non-standard with respect to data source, data augmentation, sequence length, and preprocessing. With inappropriate data construction, the trained model is likely to underperform textual features like TNF, as illustrated in Table 1 and Figures 7–9. We demonstrate the viability of learning species-awareness by treating non-overlapping segments of the same species as positive pairs and analyze the effectiveness of different positive pair construction strategies for DNA sequences.
• Benchmark Availability. There is a lack of standard datasets and evaluation strategies for this problem. DNA datasets are diverse in many aspects, such as being balanced or raw, containing seen or unknown species, being data-scarce or abundant, and consisting of reference or long-read sequences. We have therefore compiled and published a benchmark and evaluation pipeline after iterative refinement to address these challenges.

W2: Scope of experiments

Thank you for highlighting this issue.

We have conducted the experiments, but due to space limitations, we did not include all of them in the main text. We have highlighted them in the revised version (line 273-282). Specifically:

1. We have compared our model with most of the widely-used and state-of-the-art genomics models, including Nucleotide Transformer v2, DNABERT-2, and HyenaDNA (as shown in Table 1, Figures 3 and 4).
2. In Appendix C.4, we compare DNABERT-S with MMSeqs2, one of the state-of-the-art database search methods. We show that DNABERT-S achieves slightly better performance than MMSeqs2 with fewer labeled data, indicating the potential of embedding-based methods in taxonomy classification.
3. In Appendix C.8, we also compare DNABERT-S with DNABERT-2 on several genome function prediction tasks and show that species-aware training does not significantly improve genome function predictions.

W3: Insufficient visual detail

Thank you for your suggestion. Our aim with this figure is to provide a global view showing that DNABERT-S is able to segregate species into separate clusters. We agree that the marker size was not optimal, and we have adjusted the figure accordingly to improve visual clarity and enhance the visualization of data distribution in the revised version.

W4: Figure layout

Thank you for bringing this to our attention. We apologize for the oversight. We have adjusted the figure's placement in the revised version.

Q1: Hyperparameter selection

Thank you for your question regarding hyperparameter selection.

1. In our problem, we consider the most important hyperparameters to be sequence length and hidden dimension size, as they are directly related to real-world applications and have a significant impact on performance. In Appendices C.6 and C.7, we present the impact of these two parameters.
2. For batch size, we set it to the maximum value possible with BF16 precision on 80GB GPUs, as larger batch sizes generally benefit contrastive learning. Due to the high memory cost of handling long input sequences, we are limited in batch size. Preliminary experiments suggest that our base model, DNABERT-2 [Zhou23], is robust to most hyperparameters, including learning rate, weight decay, and dropout. Therefore, we use the same values suggested in the DNABERT-2 paper.

[Zhou23] DNABERT-2: Efficient Foundation Model and Benchmark For Multi-Species Genomes, ICLR 2023

Q2: Advantage over classification model

Thanks for your question. DNABERT-S's most significant advantage over existing DNA classification models is its generalizability to unseen species, which is one of the main motivations behind this work. Traditional classification models are only applicable to the species they are trained on, limiting their utility in real metagenomics research where a large portion of observed sequences belong to unknown species. DNABERT-S, on the other hand, generates embeddings that naturally cluster and segregate sequences from different species, making it applicable to any given DNA sequence regardless of prior knowledge about the species.

Q3: Extend to other biological sequences

Thanks for your question. Yes, our method can be readily extended to other biological sequences, such as RNA or protein sequences. The only requirements are:

1. A deep learning model that generate an embedding for the input biological sequence
2. A dataset contains pairs of similar sequences. The similarity can be defined in any desired ways (e.g., species, function, and structure).

Q4: Fair comparison

Thanks for your question. Please refer to our response to W7 for the measures we have taken to prevent data leakage and ensure fair comparisons in our experiments. Regarding model scales, we have tried to reduce their effect by selecting the appropriate version of baseline models with a similar number of parameters. DNABERT-S contains 117 million parameters. Among the deep learning baselines, variants of DNABERT-S (including DNA-Dropout, DNA-Double, DNA-Mutate, and DNABERT-2) have the same number of parameters. For the Nucleotide Transformer, we chose the 97.9 million parameter version to maintain consistency in model size across comparisons.

Q5: Time / Memory / Resources

Thank you for pointing out the importance of analyzing time and memory consumption. We agree that including this information is helpful, and we have added it to the revised version (Appendix E).

1. Time and memory usage of inference.

Since time and memory usage are primarily impacted by the base model, we compare DNABERT-2, HyenaDNA, and Nucleotide Transformer using sequences of 10,000 base pairs and BF16 precision. The memory usage is measured with a batch size of 1, and time is computed using the largest possible batch size when encoding 512 sequences. As shown below, the models demonstrate similar inference speeds, while DNABERT-2 uses more memory than the Nucleotide Transformer and HyenaDNA.

2. Resources for training different models

Comparing training costs directly is challenging because DNABERT-S is trained upon the pre-trained DNABERT-2, and different genomics models are trained with different datasets. However, for reference, pre-training DNABERT-2 takes approximately 3 days on 8 A100 80GB GPUs, while training DNABERT-S takes around 2 days using the same resources.

3. Model selection

Based on our empirical study, DNABERT-S is most advantageous for tasks where species differentiation is important. For single-species tasks, such as promoter prediction on the human genome, genome foundation models like DNABERT-2 and Nucleotide Transformer are preferred choices. For tasks involving extra-long sequences (e.g., 1 million base pairs), HyenaDNA offers advantages due to its ability to handle longer sequences efficiently.

Thanks again for all the suggestions. We hope our reply can solve your concerns. Please don't hesitate to share any other thoughts!

W5: Ablation study discussion

Thanks for highlighting this! Based on our motivation behind the method design (discussed in reply to W1), we conduct a case study to empirically validate our intuition. Specifically, we collect 50 5000-bp genome sequences from 3 species: human, monkey, and a randomly selected bacteria named Salmonella enterica. We compute the embedding of each genome sequence, and achieve the species embedding by averaging the embedding of all its 50 sequences. We then compute the cosine distance between human-monkey (H-M), human-bacteria (H-B), and monkey-bacteria (M-B). We then compute the relative distance between humans and bacteria (H-B/H-M) and monkeys and bacteria (M-B/H-M). As shown in the table below, models trained with MI-Mix loss naturally segregate very dissimilar species like humans and bacteria further while keeping similar species like humans and monkeys closer. We observe the same pattern in several different bacteria. These case studies can illustrate why MI-Mix is more suitable than SimCLR for metagenomics data, besides the scores in the ablation study.

W6: Unclear parameter justification and redundancy reduction

Thanks for indicating this! We agree that supplementing the parameter selection is helpful and have included them in the revised version (line 263-272).

1. Datasets for metagenomics binning. To mimic real-world applications, we use the raw datasets without any data filtering for the metagenomics binning experiments.
2. Datasets for species classification and clustering. We apply data filtering only to the datasets used for species classification and clustering. Clustering algorithms and few-shot classification are often sensitive to unbalanced data. Therefore, filtering allows us to rule out the data balance factor and fairly compare embedding quality. We chose the threshold of 100 sequences per species based on dataset statistics, which allows us to maintain a sufficient number of species while ensuring enough sequences in each species for reliable analysis. The selection of 100 sequences from each species is purely random, and we will share the code used for this process to enhance transparency.

W7: Potential data leakage

Thank you for expressing this concern. In species differentiation tasks, we consider data leakage to occur when the same species are present in both training and evaluation datasets. We have performed careful validation to prevent this and provided the details in the revised version (Appendix F).

Our experiments use two categories of data: CAMI2 and synthetic datasets. We constructed the synthetic data to ensure they do not include any species present in the training data. For the CAMI2 datasets, due to discrepancies in species annotations between CAMI2 and GenBank, direct validation was challenging. Therefore, we performed an alignment-based estimation using minimap2. We aligned each evaluation dataset to the training data and considered sequences with over $90\%$ alignment to the training sequences as present in the training data.

We computed the presence rate of each evaluation dataset as number of presented sequences / total number of sequences. It's important to note that different species can share common or highly similar genome sequences, so a non-zero presence rate is expected in real-world scenarios. As a reference, the two synthetic datasets with non-overlapping species have presence rates of $6.88\%$ and $8.92\%$. For the CAMI2 datasets, the plant-associated ones have presence rates between $3.51\%$ and $4.98\%$, which are even lower than the synthetic datasets. The marine datasets have presence rates between $7.99\%$ and $9.45\%$, comparable to the synthetic ones. Based on these statistics, there is negligible species leakage between our training and evaluation data.