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

## Abstract

Abstract We introduce DNABERT-S, a tailored genome model that develops species-aware embeddings to naturally cluster and segregate DNA sequences of different species in the embedding space.
Differentiating species from genomic sequences (i.e., DNA and RNA) is vital yet challenging, since many real-world species remain uncharacterized, lacking known genomes for reference. Embedding-based methods are therefore used to differentiate species in an unsupervised manner.
DNABERT-S builds upon a pre-trained genome foundation model named DNABERT-2.
To encourage effective embeddings to error-prone long-read DNA sequences, we introduce Manifold Instance Mixup (MI-Mix), a contrastive objective that mixes the hidden representations of DNA sequences at randomly selected layers and trains the model to recognize and differentiate these mixed proportions at the output layer. We further enhance it with the proposed Curriculum Contrastive Learning (C 2 LR) strategy.
Empirical results on 23 diverse datasets show DNABERT-S’s effectiveness, especially in realistic label-scarce scenarios.
For example, it identifies twice more species from a mixture of unlabeled genomic sequences, doubles the Adjusted Rand Index (ARI) in species clustering, and outperforms the top baseline’s performance in 10-shot species classification with just a 2-shot training. Model, codes, and data is publicly available at https://github.com/MAGICS-LAB/DNABERT_S .

## 1 Introduction

Accurate differentiation of species from genomic sequences is a critical task in biology and ecology, supporting efforts in biodiversity conservation, epidemiology, understanding evolutionary processes, and exploring the roles of microbiomes in health and disease.
Traditional methods for species identification rely heavily on well-characterized reference genomes for comparative analysis. Thus, they are limited due to the vast and largely unexplored genetic diversity present in natural environments.

A prime example is metagenomics binning. Metagenomics binning [Kang et al., , ; Nissen et al., ; Meyer et al., ; Lamurias et al., ] is a pivotal process in microbiome research, aiming to group DNA sequences by species from complex mixtures containing DNA from potentially thousands of distinct, often uncharacterized species. In this context, effective DNA embeddings that can accurately segregate and cluster DNA sequences are more suitable than the methods that rely on known reference genomes for comparison and alignment.

Figure: Figure 1: TSNE visualization of the DNA embeddings generated by different methods on a CAMI2 [Meyer et al., ] dataset with $50$ different species. Each point represents an individual DNA sequence, with the color coding indicating the species affiliation. Notably, DNABERT-S demonstrates a pronounced ability to cluster and segregate different species within the embedding space.
Refer to caption: x1.png

Despite the critical role of DNA embeddings in various scenarios, there is a notable deficiency in the development of effective methods. Current approaches to achieving DNA embeddings include: 1) Descriptive textual features [Kang et al., , ; Nissen et al., ], 2) Pre-trained Kmer embeddings [Ng, ; Ren et al., ; Han et al., ], and 3) Genome foundation models [Ji et al., ; Nguyen et al., ; Zhou et al., ]. The first two methods, while straightforward, often fail to capture complex semantic relationships inherent in genomic data. Genome foundation models [Ji et al., ; Nguyen et al., ; Zhou et al., ], despite their successes in various genomic tasks through model fine-tuning, generally fail to develop embeddings that can discriminate certain properties like species. This is largely due to a mismatch between their pre-training objectives [Radford et al., ; Devlin et al., ] and the specific application scenarios. Our empirical analysis, as detailed in Table 1, reveals that in many scenarios, existing genome foundation models even underperform simple textual features.

In this work, we introduce DNABERT-S, a specialized genome model that harnesses the capabilities of genome foundation models to generate species-aware DNA embeddings. As depicted in Figure 1, DNABERT-S distinguishes itself from other methods by its ability to effectively cluster and separate different species within the embedding space. This enhanced performance stems from the proposed Manifold Instance Mixup (MI-Mix) loss and Curriculum Contrastive Learning (C^2LR) strategy.
Contrastive learning enables the model to discern between similar and dissimilar DNA sequences, and curriculum learning incrementally presents more challenging training samples, fostering better learning and generalization. The training of DNABERT-S includes two phases.
In the first phase, we adopt a Weighted SimCLR Chen et al. []; Zhou et al. [] training objective to encourage the model to group DNA sequences from the same species and separate DNA sequences from distinct species. In the second phase, we introduce Manifold Instance Mixup (MI-Mix), which mixes anchor instances at a randomly selected layer to create more challenging anchors for contrastive training.

To rigorously evaluate the models, we assemble a comprehensive benchmark that includes thousands of species, capturing the diversity of natural microbial communities. This benchmark incorporates complex samples from CAMI2 [Meyer et al., ], a leading metagenomics binning benchmark, and extensive reference genomes from Genbank [Benson et al., ]. Thus, it includes both natural yet error-prone long-read sequences and well-curated yet potentially biased reference genomes.

We evaluate the embedding quality from entirely unsupervised problems to classification tasks with abundant labels. Experimental results indicate the effective performance of DNABERT-S. Compared to the strongest existing method, DNABERT-S doubles its performance in the clustering task and achieves better performance with only $20\%$ of labeled data in the classification task (e.g., 2-shot v.s. 10-shot). Notably, in metagenomics binning, DNABERT-S is able to recover over $40\%$ and $80\%$ of species with an F1 score of over $0.5$ respectively from synthetics and more realistic datasets, which is also one time more than the strongest baseline. We also show that a simple K-Nearest-Neighbors species classifier, which is trained on DNABERT-S embeddings of a small portion of each target genome, can slightly outperform a well-established traditional method MMseqs2 [Steinegger and Söding, ] that relies on the entire reference genomes of target species for classification.

Our contribution can be summarized as follows:
1) We demonstrate the effectiveness of genome foundation models in learning DNA embeddings, opening new avenues for various genomics research problems;
2) We introduce DNABERT-S, a model that develops distinctly better embeddings for species differentiation;
3) We propose the Curriculum Contrastive Learning (C^2LR) strategy with the Manifold Instance Mixup (MI-Mix) loss;
4) We publish a large-scale evaluation benchmark.

## 2 Background

This study aims to build a species-aware DNA embedding model that maps each DNA sequence as a fixed-size numerical vector in an embedding space, where sequences from distinct species are naturally clustered and segregated. A DNA sequence is essentially a sentence composed of four unique characters: A, T, C, and G.

Existing works highly rely on descriptive textual features [Kang et al., , ; Nissen et al., ] and pre-trained K-mer embeddings [Ng, ; Ren et al., ; Han et al., ] to compute DNA embeddings. A representative descriptive textual feature is Tetra-Nucleotide Frequency (TNF), a $256$-dimensional vector where each position represents the frequency of each unique $4$-mer (e.g., TTCA, AACG) in the input DNA sequence. Despite its simplicity and effectiveness, this method is limited since it singly relies on the $4$-mer frequency and is not trainable to better fit downstream applications. Besides, our empirical analysis also suggests that a naive trainable model based on TNF, such as a Variational AutoEnoder [Kingma and Welling, ] with TNF as input, results in worse embeddings compared to TNF. With the success of Word2Vec [Mikolov et al., ], pre-trained Kmer embeddings have gained popularity in computing DNA embeddings for various applications [Ng, ; Ren et al., ; Han et al., ]. However, the emergence of deep learning advancements such as ELMo and BERT [Peters et al., ; Devlin et al., ] highlights the limitations of static word embeddings compared to contextual embeddings produced by foundation models.

Recently, genome foundation models such as DNABERT-2 and HyenaDNA have demonstrated their effectiveness in genome analysis [Ji et al., ; Dalla-Torre et al., ; Nguyen et al., ; Zhou et al., ].
However, these models do not naturally develop discriminative embeddings, largely due to the discrepancy between their language-modeling training objectives and the goal of segregating sequences in the embedding space [Li et al., ].
To leverage the power and potential of genome foundation models, we tailor a contrastive learning method [Reimers and Gurevych, ; Gao et al., ; Chen et al., ; Lee et al., ] for DNA embedding learning by introducing the curriculum contrastive learning (C^2LR) strategy with the Manifold Instance Mixup (MI-Mix) training objective.

## 3 Model

The proposed Curriculum Contrastive Learning ($\text{C}^{2}$LR) splits the training process into two phases, gradually creating more challenging anchors.
In phase I, we apply an effective contrastive learning method named Weighted SimCLR based on SimCLR and Hard-Negative sampling strategy (Sec. 3.1).
In phase II, we propose the Manifold Instance Mixup method which creates more challenging anchors by mixing intermediate hidden states of inputs in a randomly selected hidden layer of the model (Sec. 3.2).
Implementation details of DNABERT-S are presented in Sec. 3.3.

Notation: Let $x_{i}$ be an input sample. Given a batch $\{(x_{i},x_{i^{+}})\}_{i=1}^{B}$, where $B$ is the batch size and $(x_{i},x_{i^{+}})$ represents a pair of similar samples (a.k.a., positive pair).
In our setting, a positive pair $(x_{i},x_{i^{+}})$ represents two non-overlapping DNA sequences from the same genome.
Let $f(\cdot)$ define the embedding model, which takes $x_{i}$ as input and computes fixed-size embedding $f(x_{i})$.

### 3.1 Weighted SimCLR

SimCLR [Chen et al., ] is a simple and effective framework for contrastive learning. For an anchor $x_{i}$ in batch $\{(x_{i},x_{i^{+}})\}_{i=1}^{B}$, SimCLR treats all the other $2B-2$ samples in the same batch as negative samples. It encourages the model to increase the anchor’s similarity with its positive sample $x_{i^{+}}$ and reduces its similarity with the negative samples. It treats all negative samples equally. However, recent works [Zhang et al., ] have suggested that hard negatives that are closer to the anchor in the representation space offer more informative learning contrasts. Therefore, Weighted SimCLR [Zhang et al., ] gives higher weights to negative samples that are closer to the anchor.

Figure: Figure 2: Overview of DNABERT-S’s training process. We construct training data from massive reference genomes and train DNABERT-S with the proposed Curriculum Contrastive Learning (C^2LR) strategy that progressively provides more challenging contrastive anchors to the model in two different phases. We propose the Manifold Instance Mixup (MI-Mix) objective that mixes the intermediate hidden states of different inputs to construct more challenging contrastive anchor.
Refer to caption: x2.png

To align with subsequent sections, we introduce the virtual labels. The label for $(x_{i},x_{i^{+}})$ is $v_{i}\in\{0,1\}^{B}$, where $v_{i,i}=1$ indicates positive samples, and $v_{i,j\neq i}=0$ indicates negative samples. The Weighted SimCLR loss for $x_{i}$ is defined as:

$$ $\ell(f(x_{i}),v_{i})=-\sum_{n=1}^{B}v_{i,n}\log\frac{\exp\left(\operatorname{s }\left(f(x_{i}),f(x_{n^{+}})\right)/\tau\right)}{\sum_{j\neq i}\alpha_{ij}\exp \left(\operatorname{s}\left(f(x_{i}),f(x_{j})\right)/\tau\right)},$ (1) $$

where $\tau$ denotes the temperature and $\operatorname{s}(\cdot,\cdot)$ denotes the cosine similarity between two inputs. Weights $\alpha_{ij}$ denotes the relative importance of $x_{j}$ for optimizing the contrastive loss of anchor $x_{i}$ among all the $2B-2$ negative samples. A negative sample that is closer to the anchor receives a higher weight. We set $\alpha_{ii^{+}}=1$ and compute $\alpha_{ij}$ as:

$$ $\displaystyle\alpha_{ij}=\frac{\exp\left(\operatorname{s}\left(f(x_{i}),f(x_{j })\right)/\tau\right)}{\frac{1}{2B-2}\sum_{k\neq i,i^{+}}\exp\left( \operatorname{s}\left(f(x_{i}),f(x_{k})\right)/\tau\right)}.$ $$

For each positive pair $(x_{i},x_{i^{+}})$, Weighted SimCLR respectively takes $x_{i}$ and $x_{i^{+}}$ as the contrastive anchors to calculate the contrastive loss. It defines the loss $\ell(f(x_{i^{+}}),v_{i})$ for $x_{i^{+}}$ by exchanging the roles of instances $\{x_{i}\}_{i=1}^{B}$ and $\{x_{i^{+}}\}_{i=1}^{B}$ in Eq. (1) respectively. Therefore, the Weighted SimCLR loss on the entire batch is defined as:

$$ $\mathcal{L}=\frac{1}{2B}\sum_{i=1}^{B}(\ell(f(x_{i}),v_{i})+\ell(f(x_{i^{+}}), v_{i})).$ (2) $$

### 3.2 Curriculum Contrastive Learning ( C 2 superscript C 2 \text{C}^{2} C start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT LR)

In this part, we introduce our curriculum contrastive learning ($\text{C}^{2}$LR) method. Curriculum learning is an effective training method that first presents easy training batches and then progresses to more challenging ones [Hacohen and Weinshall, ]. Recent studies have successfully applied this technique to both positive pairs [Ye et al., ; Roy and Etemad, ] and negative pairs [Chu et al., ] in contrastive learning. We take this approach a step further by applying it to contrastive anchors, effectively using it for both types of pairs at the same time.

As shown in Figure 2, our $\text{C}^{2}$LR method includes two training phases, with anchors becoming progressively more challenging. In phase I, we use the Weighted SimCLR introduced in Sec. 3.1. In phase II, we propose the Manifold Instance Mixup (MI-Mix) method to mix up anchor instances in a random hidden layer, motivated by the instance mixup (i-Mix) method [Lee et al., ].

The i-Mix method mixes anchors at the input layer to create more challenging positive and negative pairs. It only uses the samples from $\{x_{i}\}_{i=1}^{B}$ as anchors and only considers the positive and negative samples from $\{x_{i^{+}}\}_{i=1}^{B}$. Otherwise, it nearly doubles the memory or training time compared to the Weighted SimCLR method in Sec. 3.1 (see Appendix D for details). To perform mixup within the anchor space, i-Mix first shuffles $\{(x_{i},v_{i})\}_{i=1}^{B}$ to generate $\{(\hat{x}_{i},\hat{v}_{i})\}_{i=1}^{B}$, i.e. a random permutation of $\{(x_{i},v_{i})\}_{i=1}^{B}$. Then for each anchor $(x_{i},v_{i})$, i-Mix mixes it with $(\hat{x}_{i},\hat{v}_{i})$ through weighted sum. The mixing weight $\lambda_{i}$ is drawn from $\text{Beta}(\alpha,\alpha)$, where $\alpha$ is a hyperparameter.

Despite i-Mix’s effectiveness on continuous data such as images and speeches, directly mixing DNA sequences may avoid biological plausibility. Thus, we proposed to instead mix hidden representations of DNA sequences at a deeper layer, which essentially combines more abstract, higher-level features of the sequences. We call it Manifold Instance Mixup, inspired by Verma et al. [].
Concretely, we denote the model $f(\cdot)$ as $f(x)=f_{m}(g_{m}(x))$. Here, $g_{m}(\cdot)$ maps input data to the intermediate hidden states at layer $m$, and $f_{m}(\cdot)$ maps these intermediate hidden states to the output $f(x)$.

The Manifold Instance Mixup includes four steps.
First, we uniformly select a random layer $m$ from a set of eligible layers $\mathcal{S}$ in the model, like one of the encoder layers in DNABERT-S.
Second, for a batch of anchors $\{(x_{i},v_{i})\}_{i=1}^{B}$, we process them up to layer $m$, resulting in a batch of intermediate hidden states $\{(g_{m}(x_{i}),v_{i})\}_{i=1}^{B}$.
Third, we shuffle $\{(g_{m}(x_{i}),v_{i})\}_{i=1}^{B}$ to get $\{(g_{m}(\hat{x}_{i}),\hat{v}_{i})\}_{i=1}^{B}$ and mix them up. This produces the mixed hidden states $h_{i}^{m}$ for each $(g_{m}(x_{i}),v_{i})$, where $h_{i}^{m}=\lambda_{i}g_{m}(x_{i})+(1-\lambda_{i})g_{m}(\hat{x}_{i})$, and $v_{i}^{mix}=\lambda_{i}v_{i}+(1-\lambda_{i})\hat{v}_{i}$.
Fourth, we feed $\{h_{i}^{m}\}_{i=1}^{B}$ through the remaining layers to get the last hidden states $\{f_{m}(h_{i}^{m})\}_{i=1}^{B}$.
Loss on $i$-th anchor $(x_{i},v_{i})$ is defined as:

$$ $\hat{\ell}(f_{m}(h_{i}^{m}),v_{i}^{mix})=-\sum_{n=1}^{B}v_{i,n}^{mix}\log\frac {\exp\left(\operatorname{s}\left(f_{m}(h_{i}^{m}),f(x_{n^{+}})\right)/\tau \right)}{\sum_{j=1}^{B}\alpha_{ij^{+}}\exp\left(\operatorname{s}\left(f_{m}(h_ {i}^{m}),f(x_{j^{+}})\right)/\tau\right)},$ $$

where weights $\alpha_{ii^{+}}=1$ and $\alpha_{ij^{+}}$ is computed as:

$$ $\displaystyle\alpha_{ij^{+}}=\frac{\exp\left(\operatorname{s}\left(f_{m}(h_{i} ^{m}),f(x_{j^{+}})\right)/\tau\right)}{\frac{1}{B-1}\sum_{k=1,k\neq i}^{B}\exp \left(\operatorname{s}\left(f_{m}(h_{i}^{m}),f(x_{k^{+}})\right)/\tau\right)}.$ $$

The Manifold Instance Mixup loss is defined as follows:

$$ $\hat{\mathcal{L}}=\frac{1}{B}\sum_{i=1}^{B}\hat{\ell}(f_{m}(h_{i}^{m}),v_{i}^{ mix}).$ (3) $$

### 3.3 Implementation

In the $\text{C}^{2}$LR method, we set temperature $\tau$ as $0.05$ and hyperparameter $\alpha$ as $1.0$. We train the model for one epoch in phase I using loss Eq. (2) and for two epochs in phase II using loss Eq. (3).
We use mean pooling of the last hidden states of all the tokens as the DNA embedding.
We employ the Adam optimizer [Kingma and Ba, ], with a learning rate of $3e-6$ and batch size of $48$. We save the model every 10000 training steps and select the best one based on the validation loss in the validation dataset.
We use the pre-trained DNABERT-2 [Zhou et al., ] as the starting point of contrastive training. We also conduct parallel experiments with HyenaDNA [Nguyen et al., ]. In Section 5.6, we show that DNABERT-2 outperforms HyenaDNA after the same contrastive training.
The training of DNABERT-S takes approximately 48 hours on 8 NVIDIA A100 80GB GPUs.

## 4 Data

In this section, we introduce the dataset we used for DNABERT-S training and evaluation.

##### Training

Each training sample of DNABERT-S is a pair of non-overlapping DNA sequences extracted from the same species. We focus on microbe species since the genetic diversity within this group provides a rich substrate for examining the nuances of species differentiation. The dataset is constructed with the reference genomes from GenBank [Benson et al., ]. We obtained $47923$ pairs from $17636$ viral genomes, $1$ million pairs from $5011$ fungi genomes, and $1$ million pairs from $6402$ bacteria genomes. We randomly selected $2$ million pairs from the entire $2047923$ pairs of DNA sequences to construct the training data. The rest pairs are treated as validation data. All the DNA sequences are 10000 bp in length.

##### Evaluation

Our evaluation spans $14$ long-read datasets from the Critical Assessment of Metagenome Interpretation (CAMI) II [Meyer et al., ] challenge benchmark and $9$ synthetic datasets from reference genomes. CAMI2 is one of the most comprehensive and rigorous benchmarks for metagenomics research. The datasets in CAMI2 are designed to mimic realistic microbiome environments and include a vast array of both new and known genomes, as well as plasmids and viruses. It aligns our study with real-world ecological and biological scenarios, providing a robust and contextually relevant evaluation for the DNA embedding models. We utilize $7$ datasets of long-read contigs respectively from the Marine and Plant-associated environments, where each dataset consists of $150$k-$200$k DNA sequences belonging to about $100-750$ different species sampled from $1680$ microbial genomes and $599$ circular elements. We also create $9$ Synthetic datasets by randomly extracting DNA sequences from fungi and viral reference genomes that do not overlap with our training data. Table 5 in Appendix 5 shows the statistics of the datasets we used for evaluation.

## 5 Experiments

We evaluate the model in a series of tasks, including: 1) metagenomics binning that identifies species from a mixture of sequences from an unknown number of species; 2) species clustering given the number of species; 3) species classification with a few labeled samples; and 4) long-read classification given reference genomes.
Since the CAMI2 datasets are highly imbalanced, for the clustering and classification tasks, we filtered the datasets to eliminate species with fewer than $100$ sequences and only kept $100$ sequences for each species, resulting in a set of perfectly balanced datasets. For the metagenomics binning problem,
to mimic real-world scenarios, we do not balance the data. Instead, following Kang et al. [], we only keep DNA sequences longer than $2500$bp and filter out species with fewer than $10$ sequences.

In this section, we present experimental design and empirical results. We introduce baselines in Sec. 5.1 and respectively present the results of clustering in Sec. 5.2, metagenomics binning in Sec. 5.3, and classification in Sec. 5.4. In Sec. 5.5, we present ablation studies on C^2LR and the proposed Manifold Instance Mixup training objective. Comparison with an alignment-based method for species classification is presented in Appendix 6.
We also provide empirical analysis on varying backbone models, different input lengths, reduced feature dimensions, scenarios with abundant training data, various other types of tasks, and results with error bars. Due to space limitations, we present them in Appendix C.
For all tasks involving randomness, we perform $5$ independent runs with different random seeds for each model and report the averaged results.

### 5.1 Baselines

We compare our model with four lines of work to examine its effectiveness.
TNF, TNF-K, and TNF-VAE are the most widely used DNA embedding methods in metagenomics binning tools [Kang et al., , ; Nissen et al., ]. TNF represents Tetra-Nucleotide Frequency, which uses the appearance frequency of each unique 4-mer ($4^{4}=256$ in total) in a DNA sequence as its embedding. TNF-K [Nissen et al., ] reduces TNF to $103$-dimension with a linear kernel, which utilizes DNA characteristics to reduce the correlations among different dimensions of the original TNF feature. TNF-VAE trains a Variational Autoencoder [Kingma and Welling, ] using TNF as input to extract features.
DNA2Vec [Ng, ] learns pre-trained K-mer embedding. We set $K=4$ to make it directly comparable with TNF and use the average of the $4$-mer embeddings as the DNA embedding.
DNABERT-2 [Zhou et al., ], HyenaDNA [Nguyen et al., ], and NT-v2 (Nucleotide Transformer-v2) [Dalla-Torre et al., ] are representative genome foundation models. We use the average of the last hidden states as the DNA embedding. For evaluations, we utilized the respective pre-trained models from Huggingface ModelHub, specifically zhihan1996/DNABERT-2-117M, LongSafari/hyenadna-medium-450k-seqlen-hf, and InstaDeepAI/nucleotide-transformer-v2-100m-multi-species.
DNA-Mutate, DNA-Dropout, and DNA-Double are variants of DNABERT-S, with the same hyperparameters but different positive pair construction strategies in contrastive training. DNA-Mutate views the same DNA sequence before and after random mutation (i.e., swap and delete $5\%$ of nucleotides) as a positive pair. DNA-Dropout passes the same DNA sequence through the embedding model (with a dropout rate $0.1$) twice and views the two distinct embeddings as a positive pair. DNA-Double views a DNA sequence and its reverse complementary (e.g., AATTC v.s. TTAAG) as a positive pair.
We also compare the proposed curriculum contrastive learning framework and MI-Mix training objective with well-established contrastive learning methods, including Weight SimCLR [Zhang et al., ], i-Mix [Lee et al., ] and Supervised Contrastive Learning (SupCon) [Khosla et al., ]. Results on different training objectives are presented in Section 5.5.

### 5.2 Clustering

**Table 1: Models’ performance on K-Means clustering measured by Adjusted Rand Index (ARI). DNABERT-S doubles the ARI of the strongest baseline on average.**
| Model | Synthetic | Marine | Plant | Ave. |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 0 | 1 | 0 | 1 | 2 | 3 | 4 | 0 | 1 | 2 | 3 | 4 |  |  |
| TNF | 38.75 | 37.76 | 25.65 | 25.31 | 26.05 | 20.67 | 23.47 | 25.80 | 24.23 | 24.81 | 22.72 | 22.39 | 26.47 |
| TNF-K | 36.26 | 35.66 | 25.99 | 25.00 | 26.27 | 21.15 | 23.27 | 25.60 | 25.58 | 26.45 | 22.59 | 21.76 | 26.30 |
| TNF-VAE | 25.94 | 24.60 | 16.28 | 16.52 | 16.27 | 12.92 | 15.02 | 18.40 | 16.51 | 17.53 | 14.08 | 14.38 | 17.37 |
| DNA2Vec | 24.68 | 23.34 | 16.07 | 15.99 | 16.18 | 12.62 | 14.51 | 20.13 | 19.77 | 20.25 | 17.24 | 16.37 | 18.10 |
| HyenaDNA | 20.04 | 18.99 | 16.54 | 16.64 | 16.47 | 13.35 | 14.85 | 24.06 | 25.33 | 26.18 | 21.01 | 21.16 | 19.55 |
| NT-v2 | 8.69 | 9.63 | 4.92 | 4.74 | 5.02 | 3.68 | 4.31 | 7.00 | 6.32 | 6.37 | 5.54 | 5.42 | 5.97 |
| DNABERT-2 | 15.73 | 16.74 | 13.24 | 13.53 | 12.99 | 10.41 | 11.87 | 15.70 | 16.28 | 16.32 | 13.99 | 13.66 | 14.21 |
| DNA-Dropout | 16.64 | 16.08 | 11.89 | 11.77 | 11.89 | 9.85 | 10.31 | 16.18 | 15.41 | 16.95 | 13.53 | 13.85 | 13.70 |
| DNA-Double | 35.11 | 34.14 | 27.05 | 27.23 | 26.56 | 21.47 | 24.39 | 22.35 | 21.35 | 23.03 | 19.44 | 19.06 | 25.10 |
| DNA-Mutate | 16.55 | 16.24 | 11.40 | 11.53 | 11.34 | 9.03 | 10.02 | 14.27 | 14.13 | 16.22 | 12.01 | 11.68 | 12.87 |
| DNABERT-S | 68.21 | 66.33 | 53.98 | 52.56 | 51.99 | 46.39 | 50.49 | 51.43 | 51.56 | 51.11 | 50.44 | 51.15 | 53.80 |

In this task, we evaluate the embedding quality by how well a standard clustering algorithm can distinguish and cluster different species based on the embedding.
To reduce the effects of other factors, we assume the number of species is known in this task.
For each dataset, we compute the embedding of each DNA sequence and perform K-means clustering by setting the num_clusters as the number of species that exist in this dataset. We employ the Adjusted Rand Index (ARI) as the evaluation metric. ARI is a measure of the similarity between two data clusterings, adjusted for chance, providing a normalized index that ranges from $-1$ to $1$; the higher, the better.

Table 1 shows the models’ performance on clustering. As shown in the table, DNABERT-S consistently achieves the best performance on all the datasets and doubles the performance of the strongest existing method on average. Among all the baselines, TNF and its variant TNF-K achieve the best performance, explaining their wide usage in metagenomics binning. Yet, TNF’s performance is heavily limited since it is not learnable. TNF-VAE represents a naive algorithm that enables learning with TNF, yet it leads to big performance degradation, potentially resulting from the large gap between its training objective and the specific downstream application. Similarly, pre-trained Kmer embeddings from DNA2Vec also fail to effectively cluster different species.

Existing genome foundation models training with language modeling objectives, such as HyenaDNA and DNABERT-2, despite their good performance on labeled datasets, also fail to generate representative embedding without fine-tuning. The phenomenon that pre-trained foundation models underperform descriptive textual features in generating embedding for clustering and retrieval is also observed in the field of natural language processing [Reimers and Gurevych, ].

Furthermore, by comparing the DNA-Dropout and DNA-Mutate with DNABERT-2, we found that those popular unsupervised positive pair methods used in contrastive learning in NLP, such as sentence swap/deletion and dropout, do not benefit DNA embedding learning. The DNA-Double, which utilizes the unique double-strain characteristics of DNA sequences, empowers DNABERT-2 to achieve a similar level of performance as TNF. Comparison between DNABERT-S and these variants indicates the importance of appropriate training data construction.

### 5.3 Metagenomics Binning

Metagenomics binning is a crucial process in microbial ecology, involving the categorization of DNA sequences into groups that represent individual species.
State-of-the-art metagenomics binning method [Kang et al., , ; Nissen et al., ] always formulate this problem as a clustering problem with an unknown number of clusters based on the feature of each DNA sequence. The DNA sequence feature is computed by combining sequence-based DNA embedding with various other features and the clustering algorithms are often complicated and strongly correlated with the features they utilize.
In our evaluation, to create a fair environment for DNA embedding benchmarking, instead of relying on any existing tool, we implement the modified K-medoid clustering algorithm proposed in Kang et al. [] for metagenomics binning due to its simplicity and effectiveness. Algorithm 1 describes the metagenomics binning algorithm.
Following Kang et al. [, ], we formulate this problem as identifying non-overlapping clusters of DNA sequences from the entire dataset, where each cluster of sequence is considered as an identified species. We iteratively identify the densest point in the embedding space and take all the sequences that are close (determined by a learned threshold) to it as the group of the sequences that belong to the same species. The embeddings of the taken sequences are removed from the embedding space. The iteration ends as there are no regions that contain enough number sequences within the threshold. This algorithm, evaluates how well the embedding method clusters and aggregates different species within the embedding space.
We then compare the predicted clusters with the true labels to count the number of species that have been successfully identified. A species is considered to be successfully identified if the F1 score of this species is over $0.5$. We compare different models by the number of species they identify with different levels of F1 scores (e.g., $0.5-0.6$, $0.8-0.9$).
We only use the DNA embeddings as the feature of each DNA sequence.

Figure: Figure 3: Metagenomics Binning Results. The bin size represents the number of unique species identified by each model and different colors represent the F1 score of the identified species. With high F1 scores, DNABERT-S identifies many more species than the baselines.
Refer to caption: x3.png

Figure 3 shows the models’ performance on $6$ metagenomics binning datasets. As shown in the figure, similar to our observation in clustering, DNABERT-S identifies twice the number of species with an F1 score of over $0.5$ compared to the strongest baseline, showing its great capability in tackling important real-world biology challenges. Notably, DNABERT-S identifies a large number of species with an F1 score over $0.9$. indicating its capability to accurately segregate different species in the embedding space, aligning with our observation in Figure 1. In the Synthetic datasets, where the sequences are error-less (extracted from reference genome) and the number of sequences in each species is more balanced, DNABERT-S recovers over $80\%$ of the species with an F1 score of over $0.5$ purely based on the DNA sequences themselves. In more realistic datasets such as Marine and Plant, where noise (e.g., error from sequences) exists in DNA sequence and species size is highly imbalanced, DNABERT-S is still able to recover $40\%$ of the species with an F1 score of over $0.5$.

### 5.4 Classification

In this task, we evaluate the embedding quality by how well a simple model can classify different species based on a few labeled embeddings. We conduct experiments with a linear regression model and a non-linear multi-layer perceptron (MLP) to examine the embeddings’ linear and non-linear descriptiveness. We present results on linear classification in this section. The results in non-linear settings are consistent with those in linear ones. Due to space limits, we present them in Table 6 and 7 in the Appendix. As shown in Table 5, all the datasets we use for classification consist of $100$ DNA sequence for each species. We first compute the embedding of each DNA sequence with each model. In each evaluation run, we independently select $80$ embeddings from each species to form the test set. For the rest DNA sequences, we respectively sample $1$, $2$, $5$, $10$, and $20$ embeddings from each species to form the training set. A Logistic Regression model is trained on the training set and evaluated on the test set. We use the macro F1 score as the evaluation metric.

Figure 4 shows the models’ performance on $6$ datasets. The remaining results are consistent and are presented in Appendix C.3. As shown in the figure, DNABERT-S consistently achieves the best performance. DNABERT-S achieves better performance than the strongest baseline with only $20\%$ of training data. For example, with only $2$ training samples per category, DNABERT-S achieves higher F1 scores than the strongest baseline with $10$ training samples. With the same amount of training samples, DNABERT-S outperforms the baselines by a large gap. Notably, in the Synthetic datasets, where none of the species are seen during the contrastive training, a linear model trained with DNABERT-S embeddings achieves an F1 score of over $0.8$ in 200 classes classification with only $5$ labeled samples in each species, showing DNABERT-S’s capability in generalizing well on unseen data. To further validate the generalizability of DNABERT-S on more distinct datasets, we compile three datasets, including genomes from invertebrate, protozoa, and mammalian species, that are largely different from the training species (microbial genomes) of DNABERT-S. As illustrated in Section C.5, DNABERT-S also achieves good performance in these datasets.

Figure: (a) Synthetic 0
Refer to caption: extracted/5944793/figures/cls_Reference_0.png

### 5.5 Ablation Study

In this section, we present our ablation studies on DNABERT-S. We perform the ablation study on CAMI2 datasets with both clustering and classification.
To validate the effectiveness of curriculum learning, we compare DNABERT-S with three of its variants, each of which is trained purely with the Weight SimCLR [Zhang et al., ], i-Mix [Lee et al., ], SupCon [Khosla et al., ], and our proposed Manifold Instance Mixup (MI-Mix) loss. To examine the effectiveness of MI-Mix, we also compare it with a variant trained with the curriculum contrastive method that replaces MI-Mix with i-Mix in the second phase. All the variants are trained with the same data and hyperparameters.

**Table 2: Ablation study on the Curriculum Contrastive Learning (C^2LR) and Manifold Instance Mixup (MI-Mix). -: Performance difference to W. SimCLR + MI-Mix.**
| Training Objective | Clustering | Classification |
| --- | --- | --- |
| W. SimCLR + MI-Mix | $51.11$ | $60.88$ |
| W. SimCLR + i-Mix | $-3.46$ | $-3.56$ |
| only W. SimCLR | $-1.13$ | $-1.17$ |
| only MI-Mix | $-0.66$ | $-0.42$ |
| only i-Mix | $-5.25$ | $-4.76$ |
| only SupCon | $-6.75$ | $-3.83$ |

As shown in Table 2, our curriculum learning strategy that combines Weighted SimCLR and MI-Mix achieves the best performance. Our method outperforms both variants that are trained purely with Weighted SimCLR and MI-Mix loss, showing the effectiveness of our proposed curriculum contrastive learning strategy. Moreover, the comparison among the four variants that are trained with a single loss function, including Weight SimCLR, i-Mix, and SupCon, indicates the effectiveness of MI-Mix in learning DNA embeddings.

### 5.6 Selection of Backbone Model

This subsection delineates a comparative analysis of existing genome foundation models in the context of DNA embedding generation. We evaluated four renowned models: DNABERT [Ji et al., ], DNABERT-2 [Zhou et al., ], Nucleotide Transformer [Dalla-Torre et al., ], and HyenaDNA [Nguyen et al., ].
Notably, DNABERT and the Nucleotide Transformer exhibit strict input sequence length limitations of 512 and 6144 (V1) or 12288 (V2) base pairs, respectively. Conversely, DNABERT-2 and HyenaDNA do not impose such constraints. Considering the potentially extensive length of genome sequences in metagenomics binning, our preliminary experiments focused solely on DNABERT-2 and HyenaDNA.

We train both models on our pre-training datasets with the same set of hyperparameters for $3$ epoch. We save checkpoints periodically and select the best checkpoint based on the models’ validation loss on the validation set. Since HyenaDNA, in general, requires larger learning rates than DNABERT-2, we train it with three different learning rates ($3e-4$, $3e-5$, and $3e-6$) and select the one that works best. For DNABERT-2, we only train it once with a learning rate of $3e-6$. To avoid the impact of other factors, such as the schedule of curriculum learning, we train both models with the Weighted SimCLR loss only in the entire training process. We evaluate the models before and after contrastive training on our evaluation benchmark.

Table 3 presents the performance of both pre-trained DNABERT-2 and HyenaDNA, with and without contrastive training, in K-means clustering evaluations. DNABERT-2, despite underperforming without contrastive training, demonstrated superior performance compared to HyenaDNA post-training, highlighting its efficacy in learning effective DNA embeddings. Similar trends are observed in few-shot species classification scenarios, as detailed in Table 4. While DNABERT-2 initially exhibited subpar performance, it largely improved post-contrastive training, achieving the best results. Thus, DNABERT-2 was chosen as the backbone model for DNABERT-S.

**Table 3: Performance of DNABERT-2 and HyenaDNA on K-Means clustering measured by Adjusted Rand Index (ARI). $\Delta$: the model’s performance improvement after contrastive training,**
|  | Synthetic | Marine | Plant | Ave. |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Dataset ID | 0 | 1 | 0 | 1 | 2 | 3 | 4 | 0 | 1 | 2 | 3 | 4 |  |
| DNABERT-2 w/o | 15.73 | 16.74 | 13.24 | 13.53 | 12.99 | 10.41 | 11.87 | 15.70 | 16.28 | 16.32 | 13.99 | 13.66 | 14.21 |
| DNABERT-2 w/ | 69.33 | 68.37 | 53.18 | 51.94 | 51.91 | 46.60 | 49.69 | 49.05 | 50.33 | 49.68 | 48.57 | 49.83 | 50.08 |
| $\Delta$ | 53.61 | 51.63 | 39.94 | 38.41 | 38.93 | 36.19 | 37.82 | 33.35 | 34.05 | 33.36 | 34.57 | 36.17 | 39.00 |
| HyenaDNA w/o | 20.04 | 18.99 | 16.54 | 16.64 | 16.47 | 13.35 | 14.85 | 24.06 | 25.33 | 26.18 | 21.01 | 21.16 | 19.55 |
| HyenaDNA w/ | 57.58 | 55.19 | 42.92 | 42.68 | 42.24 | 36.17 | 40.10 | 41.42 | 40.36 | 40.46 | 36.87 | 38.25 | 42.85 |
| $\Delta$ | 37.54 | 36.20 | 26.38 | 26.04 | 25.77 | 22.82 | 25.24 | 17.36 | 15.03 | 14.28 | 15.86 | 17.09 | 23.30 |

**Table 4: Performance of DNABERT-2 and HyenaDNA on few-shot classification measured by Macro F1 score. $\Delta$: the model’s performance improvement after contrastive training,**
|  | Synthetic | Marine | Plant | Ave. |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Num Shots | 1 | 2 | 5 | 10 | 1 | 2 | 5 | 10 | 1 | 2 | 5 | 10 |  |
| DNABERT-2 w/o | 24.43 | 34.81 | 48.93 | 58.58 | 19.50 | 28.45 | 40.64 | 48.98 | 21.04 | 28.16 | 38.50 | 45.46 | 36.46 |
| DNABERT-2 w/ | 72.02 | 78.63 | 84.55 | 86.99 | 48.23 | 57.90 | 64.80 | 67.94 | 44.97 | 52.35 | 60.35 | 64.63 | 65.28 |
| $\Delta$ | 47.60 | 43.83 | 35.62 | 28.41 | 28.73 | 29.45 | 24.16 | 18.96 | 23.92 | 24.20 | 21.85 | 19.17 | 28.82 |
| HyenaDNA w/o | 30.13 | 41.18 | 54.86 | 64.03 | 23.92 | 33.94 | 47.47 | 55.50 | 28.15 | 36.97 | 48.20 | 55.24 | 43.30 |
| HyenaDNA w/ | 59.58 | 67.79 | 74.62 | 78.53 | 43.60 | 53.70 | 62.00 | 65.55 | 43.46 | 52.12 | 59.40 | 62.80 | 60.26 |
| $\Delta$ | 29.45 | 26.62 | 19.76 | 14.50 | 19.68 | 19.76 | 14.52 | 10.05 | 15.31 | 15.15 | 11.20 | 7.55 | 16.96 |

## 6 Conclusion

We introduce DNABERT-S, a model tailored for species-aware DNA embeddings, which is empowered by the proposed Manifold Instance Mixup (MI-Mix) training objective and the Curriculum Contrastive Learning (C^2LR) strategy.
We perform extensive experiments on $23$ datasets across thousands of different species and a variety of challenging tasks, including species clustering, classification, and metagenomics binning, to demonstrate the DNABERT-S’s capability of species differentiation.
Furthermore, we conduct a series of experiments on the training objective, backbone model selection, impacts of sequence length, and feature dimension to provide empirical insights for DNA embedding learning.
We also compare DNABERT-S with traditional technique like MMSeq2 in the problem of species differentiation and demonstrate DNABERT-S can achieves slightly better performance with less amount of labeled samples.
We envision DNABERT-S to potentially change the way genomic problems are approached from an embedding perspective.
The primary limitation of DNABERT-S lies in its high computational demands, a common trait among deep learning models, when compared to more traditional, lightweight methods like Tetranucleotide Frequencies (TNF).

## Acknowledgments

HL is partially supported by NIH R01LM1372201, NSF CAREER1841569, DOE DE-AC02-07CH11359, DOE LAB 20-2261, and NSF TRIPODS1740735. HH and ZW, from the U.S. Department of Energy Joint Genome Institute ([https://ror.org/04xm1d337](https://ror.org/04xm1d337)), a DOE Office of Science User Facility, are supported by the Office of Science of the U.S. Department of Energy operated under Contract No. DE-AC02-05CH11231.

## Appendix A Algorithm for Metagenomics Binning

Algorithm 1 describes the unsupervised clustering algorithm we used for metagenomics binning, where $\texttt{s}({E_{i},E_{j}})$ represents the cosine similarity of two vectors $E_{i}$ and $E_{j}$.
Selection of threshold $\gamma$. As shown in Algorithm 1, the threshold $\gamma$ is the most important hyperparameter that greatly impacts the final binning results. A high threshold results in small and dense clusters while a low threshold results in large yet sparse clusters. Since different models generate embeddings with distinct distributions, a fixed threshold (e.g., $0.9$) could be too high for one model yet too low for another one. In practice, massive hyperparameter searches are needed to determine the best threshold for each model on different datasets. Due to the large size of our experiments and the various types of models we used, an automatic way is needed to fairly choose the threshold for each model on each dataset.
For each metagenomics binning dataset, we use the dataset from the same source (e.g., Marine) as it with ID 0 0 to compute a threshold for each model on it, Specifically, we generate embeddings for each DNA sequence in the dataset and compute the similarities between each DNA sequence and its species center (i.e., the average of all the DNA sequence belongs to this species). The $70$ percentile of all the similarities is used as the threshold.
Other hyperparameters. We set minimum bin size $m=10$, number of steps $Z=1000$, and number of iterations $T=3$. We also experimented with $T=3,4,5$ and $60,70,80,90$ percentile of all the similarities is used as the threshold $\gamma$, and found that the results are robust to these hyperparameters.

Figure: Algorithm 1 Modified $K$-Medoid Clustering

## Appendix B Data Statistics of Evaluation Benchmark

This section details the comprehensive statistics of the 18 datasets utilized for evaluating various DNA embedding models, as summarized in Table 5. For tasks involving clustering and classification, each dataset encompasses between 93 to 499 distinct species. From each species, 100 DNA sequences are sampled. These sequences vary in length, ranging from 2,000 to 20,000 base pairs. In the case of metagenomics binning, the datasets exhibit an unbalanced distribution of sequences across different species. Specifically, the number of sequences per species varies significantly, ranging from as few as 10 to as many as 4,599.

Our benchmark contains assets from GenBank [Benson et al., ] (license: Creative Commons Attribution Non-Commercial License [http://creativecommons.org/licenses/by-nc/2.0/uk/](http://creativecommons.org/licenses/by-nc/2.0/uk/)) and CAMI2 [Meyer et al., ] (license: Creative Commons Attribution 4.0 International License [http://creativecommons.org/licenses/by/4.0/](http://creativecommons.org/licenses/by/4.0/)). It is worth noting that both GenBank and CAMI2 preprocess the RNA sequences into DNA equivalents by replacing U with T. Thus, although we did not explicitly analyze RNA sequences, many RNA viruses are considered in both model training and evaluation.

**Table 5: Data statistics of the datasets for the DNA embedding evaluation. This table presents the sampling source, ID, number of sequences, number of sequences, and the minimum / maximum / medium values of the sequence lengths and number of sequences in each species. We use the same set of balanced datasets for clustering and classification and another set of datasets for metagenomics binning.**
| Tasks | Source | ID | Species | Sequences | Sequence Length | Num. Per Species |
| --- | --- | --- | --- | --- | --- | --- |
|  | Marine | 0 | 326 | 32600 | 2k / 20k / 7.6k | 100 / 100 / 100 |
| Unsupervised | Marine | 1 | 375 | 37500 | 2k / 20k / 8.2k | 100 / 100 / 100 |
| Clustering | Marine | 2 | 361 | 36100 | 2k / 20k / 8.5k | 100 / 100 / 100 |
| & | Marine | 3 | 499 | 49900 | 2k / 20k / 6.8k | 100 / 100 / 100 |
| Few-Shot | Marine | 4 | 360 | 36000 | 2k / 20k / 7.1k | 100 / 100 / 100 |
| Classification | Plant | 0 | 108 | 10800 | 2k / 20k / 6.6k | 100 / 100 / 100 |
|  | Plant | 1 | 100 | 10000 | 2k / 20k / 6.4k | 100 / 100 / 100 |
|  | Plant | 2 | 93 | 9300 | 2k / 20k / 6.2k | 100 / 100 / 100 |
|  | Plant | 3 | 129 | 12900 | 2k / 20k / 5.5k | 100 / 100 / 100 |
|  | Plant | 4 | 129 | 12900 | 2k / 20k / 5.7k | 100 / 100 / 100 |
| (Microbe) | Synthetic | 0 | 200 | 20000 | 10k / 10k / 10k | 100 / 100 / 100 |
| Synthetic | 1 | 200 | 20000 | 10k / 10k / 10k | 100 / 100 / 100 |  |
| (Mammalian, | Synthetic | 2 | 210 | 21000 | 10k / 10k / 10k | 100 / 100 / 100 |
| Invertebrate, | Synthetic | 3 | 210 | 21000 | 10k / 10k / 10k | 100 / 100 / 100 |
| Protozoa) | Synthetic | 4 | 210 | 21000 | 10k / 10k / 10k | 100 / 100 / 100 |
|  | Synthetic | 5 | 323 | 37278 | 10k / 10k / 10k | 31 / 200 / 111 |
|  | Synthetic | 6 | 249 | 28206 | 10k / 10k / 10k | 30 / 199 / 114 |
| Metagenomics | Marine | 5 | 515 | 119465 | 2.5k / 20k / 4.3k | 10 / 841 / 201 |
| Binning | Marine | 6 | 527 | 125194 | 2.5k / 20k / 4.4k | 10 / 915 / 223 |
| (Microbe) | Plant | 5 | 181 | 71642 | 2.5k / 20k / 3.7k | 10 / 4293 / 190 |
|  | Plant | 6 | 196 | 68426 | 2.5k / 20k / 3.7k | 10 / 4599 / 116 |
| Classification | Synthetic | 7 | 200 | 95309 | 10k / 10k / 10k | 452 / 600 / 600 |
| (Microbe) | Synthetic | 8 | 200 | 95475 | 10k / 10k / 10k | 442 / 600 / 600 |

## Appendix C More Experimental Results

### C.1 Results Comparison with Linear and Non-Linear Classifiers

Table: Table 6: DNABERT-S’s performance of species classification with varying numbers of training samples on $3$ datasets: Beyond using logistic regression (LR), we also train a Multi-Layer Perceptron (MLP) with non-linear activation function (ReLU). The term "Difference" denotes the performance gap between "DNABERT-S" and the "best-baseline". The results show that DNABERT-S embedding outperforms the best-existing baseline consistently in both linear and non-linear discriminativity.

Table: Table 7: DNABERT-S’s performance of species classification with varying numbers of training samples on $3$ datasets: Beyond using logistic regression (LR), we also train a Multi-Layer Perceptron (MLP) with non-linear activation function (ReLU). The term "Difference" denotes the performance gap between "DNABERT-S" and the "best-baseline". The results show that DNABERT-S embedding outperforms the best-existing baseline consistently in both linear and non-linear discriminativity.

### C.2 Results with Error Bars

In this section, we present detailed results on species clustering and few-shot classification on DNABERT-S and the most competitive baseline models. Table 8, 9, and 10 respectively show the models’ mean and std on each setting across 5 random seeds. As shown in the tables, DNABERT-S consistently outperforms the baselines with small variances.

**Table 8: Performance of models with error bars on Synthetic:0 dataset. We evaluate models using K-Means clustering (Sec. 5.2) and 1/2/5/10/20-shot classification (Sec. 5.4).**
|  | Synthetic 0 |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | ARI | 1 | 2 | 5 | 10 | 20 |
| TNF | 38.18 ± 1.27 | 44.30 ± 0.97 | 56.13 ± 1.10 | 68.68 ± 0.73 | 75.24 ± 0.37 | 79.48 ± 0.07 |
| TNF-K | 36.11 ± 0.09 | 39.51 ± 0.35 | 50.26 ± 0.89 | 62.43 ± 0.41 | 68.53 ± 0.50 | 72.95 ± 0.36 |
| HyenaDNA | 20.10 ± 0.50 | 30.03 ± 0.49 | 41.21 ± 0.92 | 54.42 ± 0.73 | 63.79 ± 0.59 | 70.53 ± 0.25 |
| DNA-Double | 34.91 ± 0.90 | 34.61 ± 0.91 | 46.79 ± 0.39 | 59.92 ± 0.53 | 67.45 ± 0.20 | 73.64 ± 0.25 |
| DNABERT-S | 66.94 ± 2.07 | 71.54 ± 0.51 | 77.77 ± 0.68 | 83.12 ± 0.18 | 85.63 ± 0.18 | 87.68 ± 0.15 |

**Table 9: Performance of models with error bars on Marine:0 dataset. We evaluate models using K-Means clustering (Sec. 5.2) and 1/2/5/10/20-shot classification (Sec. 5.4).**
|  | Marine 0 |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | ARI | 1 | 2 | 5 | 10 | 20 |
| TNF | 24.78 ± 0.23 | 27.89 ± 0.95 | 38.69 ± 0.04 | 52.36 ± 0.26 | 58.95 ± 0.05 | 62.69 ± 0.10 |
| TNF-K | 25.44 ± 0.50 | 22.82 ± 0.35 | 30.06 ± 0.61 | 40.50 ± 0.14 | 45.20 ± 0.25 | 49.35 ± 0.15 |
| HyenaDNA | 16.31 ± 0.07 | 23.89 ± 0.83 | 33.59 ± 0.01 | 47.50 ± 0.39 | 55.61 ± 0.18 | 61.75 ± 0.13 |
| DNA-Double | 26.82 ± 0.45 | 26.87 ± 1.16 | 36.98 ± 0.21 | 49.86 ± 0.12 | 57.93 ± 0.26 | 63.33 ± 0.06 |
| DNABERT-S | 53.91 ± 0.22 | 50.37 ± 0.74 | 59.71 ± 0.11 | 66.03 ± 0.11 | 69.00 ± 0.19 | 70.75 ± 0.13 |

**Table 10: Performance of models with error bars on Plant:0 dataset. We evaluate models using K-Means clustering (Sec. 5.2) and 1/2/5/10/20-shot classification (Sec. 5.4).**
|  | Plant 0 |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | ARI | 1 | 2 | 5 | 10 | 20 |
| TNF | 26.10 ± 0.70 | 24.07 ± 0.40 | 32.04 ± 1.00 | 43.35 ± 0.24 | 48.92 ± 0.27 | 53.29 ± 0.39 |
| TNF-K | 25.83 ± 0.67 | 22.82 ± 0.35 | 30.06 ± 0.61 | 40.50 ± 0.14 | 45.20 ± 0.25 | 49.35 ± 0.15 |
| HyenaDNA | 24.61 ± 0.77 | 28.45 ± 0.30 | 36.46 ± 0.61 | 48.56 ± 0.30 | 55.13 ± 0.53 | 59.80 ± 0.27 |
| DNA-Double | 22.10 ± 0.46 | 24.80 ± 0.56 | 32.91 ± 0.52 | 44.82 ± 0.18 | 52.77 ± 0.48 | 59.35 ± 0.37 |
| DNABERT-S | 51.15 ± 1.13 | 48.39 ± 1.33 | 55.92 ± 0.83 | 62.97 ± 0.36 | 67.14 ± 0.41 | 69.64 ± 0.15 |

### C.3 Remaining Results on Species Classification

In this section, we present the models’ performance on species classification in the other $6$ datasets that are not presented in Section 5.4 due to space limits. As shown in Figure 5, the results are consistent with those shown in Figure 4.

Figure: (a) Marine 2
Refer to caption: extracted/5944793/figures/cls_Marine_2.png

### C.4 Comparison with alignment-based method

While previous experiments have highlighted DNABERT-S’s exceptional performance in scenarios with limited or no labeled data, this section focuses on its effectiveness in situations where abundant labeled data is available. Specifically, we aim to understand the embedding-based species differentiation method in scenarios where reference genomes of the species to classify are available. We compare embedding-based methods with MMseqs2 [Steinegger and Söding, ], a leading alignment-based species classification tool.

For a fair comparison with MMseqs2, which relies on the reference genomes of each species when performing classification, we construct two datasets, each consisting of 200 distinct species. To mimic real-world setups, instead of classifying segment of reference genomes, we simulate 600 long-reads with PBSIM2 [Ono et al., ] from each selected species, 100 as the test set and 500 as the training set. For the embedding-based methods, such as DNABERT-S and TNF, we generate embedding for all these sequences and use K-Neareast-Neighbor (KNN) classifier with n equals to 5 for species classification. We respectively use 100, 200, 300, 400, and 500 sequences from each species to construct the training set. For the MMseqs2, we respectively use the reference genome and the simulated long-read sequences as the reference for alignment-based classification.

Figure: (a) Synthetic 7
Refer to caption: extracted/5944793/figures/mmseqs2_0.png

Figure 6 the results of the models on the datasets. As shown in the figure, DNABERT-S starts to outperforms MMseqs2 with 200-300 labeled sequence from each species with a simple KNN classier, while TNF achieves comparable performance as MMseqs2 with about 500 labels sequence per species. These results indicate the potential of embedding-based methods to replace traditional alignment-based method in species classification in data abundant scenarios. Yet a fact that cannot be ignored is that DNABERT-S is much more computational cost than MMseqs2. As a comparison, MMseqs2 requires about 30 seconds on a single CPU to make predictions while DNABERT-S requires about 1 hour on 2 NVIDIA A100 GPUs to do the same thing. Therefore, there is a long way to go to fully replace alignment-based methods in high-throughput genomics analysis.

### C.5 Results on Non-Microbe Species

Since DNABERT-S is trained on microbe species (e.g., viruses, fungi, and bacteria), a natural question is whether it can distinguish genomics sequences from species that are largely different from the ones in the training set. To answer this question, we construct 3 synthetic datasets (ID: 2, 3, and 4) that include genomes from invertebrate, protozoa, and mammalian species. We randomly select 70 species from each category to achieve 210 species in total. We perform the same clustering and few-shot classification as presented in Table 1 and Figure 4.

**Table 11: Performance of models on non-microbe species. We evaluate models on each dataset using K-Means clustering (Sec. 5.2) and 1/5/20-shot classification (Sec. 5.4).**
|  | Synthetic:2 | Synthetic:3 | Synthetic:4 |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Dataset ID | ARI | 1 | 5 | 20 | ARI | 1 | 5 | 20 | ARI | 1 | 5 | 20 |
| TNF | 20.70 | 26.28 | 45.71 | 56.24 | 19.11 | 25.05 | 42.69 | 52.27 | 21.49 | 25.94 | 44.96 | 54.60 |
| TNF-K | 18.63 | 22.42 | 40.07 | 50.75 | 16.91 | 20.38 | 36.48 | 46.74 | 19.30 | 21.38 | 39.23 | 48.94 |
| HyenaDNA | 11.20 | 17.29 | 34.77 | 48.76 | 10.86 | 16.50 | 32.69 | 45.26 | 11.58 | 17.30 | 35.38 | 47.86 |
| DNABERT-2 | 9.18 | 15.91 | 33.75 | 49.00 | 8.99 | 15.28 | 31.03 | 45.17 | 9.79 | 15.38 | 33.75 | 47.78 |
| DNA-Double | 13.01 | 16.04 | 30.56 | 42.11 | 12.80 | 15.91 | 28.82 | 39.28 | 14.15 | 16.20 | 31.22 | 42.25 |
| DNABERT-S | 32.70 | 33.21 | 49.78 | 59.01 | 29.44 | 29.78 | 45.65 | 54.63 | 33.60 | 32.58 | 49.67 | 57.97 |

As shown in Table 11, DNABERT-S consistently outperforms baselines across all the datasets and evaluation scenarios, indicating DNABERT-S’s transferability and robustness on species that are significantly different from its training set.
However, the improvements over the baselines are less significant, and the absolute scores, such as ARI in clustering and F1 in classification, are also lower than those in microbe datasets. On the one hand, there are higher genetic similarities among mammals compared to the often more significant genetic diversity found in microbes, making it more challenging to distinguish different mammalian species. On the other hand, due to the significant distinction between microbe and mammalian/protozoa species, some of the differentiation rules and markers learned from microbe genomes may not be applicable to genomes of species in other categories, which also suggests the needs of in-domain species-aware training.

### C.6 Impact of Sequence Length

This section delves into the influence of DNA sequence length on final model performance, examined from both training and evaluation standpoints.

#### C.6.1 Varying Sequence Length in Training

Training with longer DNA sequences increases the need for more memory and computing power. It also means we can only use smaller batches of data at a time. Therefore, the length of the sequences is an important factor in contrastive training as it affects how much it costs to train the model. To see how different sequence lengths affect training, we did three experiments using the same data. Our training data has sequences that are $10000$bp long. For experiments with shorter sequences $S$, we only used the first $S$ nucleotides of each DNA sequence. We tested sequence lengths of $500$bp, $2000$bp, and $10000$bp, training only with Weighted SimCLR loss and starting from the pre-trained DNABERT-2 model.

Figure 7 shows the results for the three models, along with the pre-trained DNABERT-2 without contrastive training and the strongest baseline, TNF. The findings reveal that sequence length significantly influences the model’s performance. Training even on short sequences, such as $500$bp, leads to substantial improvements. The model trained with $500$bp sequences performs nearly as well as TNF. When we increase the input sequence length from $500$bp to $2000$bp, there’s a marked improvement in performance. A similar trend is observed when increasing the sequence length from $2000$bp to $10000$bp. These results highlight the importance of sequence length in training an effective model. Therefore, we decided to train our model with $10000$bp sequences, despite the higher computational requirements.

Figure: Figure 7: Performance of DNABERT-S in clustering and classification with different sequence lengths used in contrastive training.
Refer to caption: x4.png

#### C.6.2 Varying Sequence Length in Evaluation

In this part, we assess how the length of DNA sequences in evaluation impacts performance. We use two synthetic datasets for clustering and classification tasks. Each sequence in these datasets is deliberately constructed to be $10000$bp long. This allows us to create a test set where all sequences have the same length. We test sequence lengths ranging from $32$bp ($2^{5}$) to $8192$bp ($2^{13}$). For each test with different sequence lengths, we keep everything else the same, like how we split the data into training and testing sets and the settings for logistic regression.

Figure 8 presents the performance of TNF and DNABERT-S with various sequence lengths. The results show that both models significantly benefit from longer sequences. When the sequence length is less than $256$bp ($2^{8}$), both models perform poorly in clustering and classifying samples. However, as the sequence length increases, starting from $512$bp ($2^{9}$), DNABERT-S begins to outperform TNF. The performance gap between the two models gets bigger as the sequence length increases. These findings highlight the crucial role of sequence length in effectively differentiating between species.

Figure: Figure 8: Performance of DNABERT-S and TNF on clustering (upper) and classification (lower) with different input sequence lengths during evaluation.
Refer to caption: x5.png

### C.7 Impact of Embedding Dimension Reduction

This section investigates how changes in embedding dimensions affect the performance of DNABERT-S, a key aspect influencing the scalability of DNA embeddings generated by the model. Initially, DNA embeddings for all clustering and classification datasets are computed using the pre-trained DNABERT-S. To reduce embedding dimensions, we use an average pooling layer with a consistent kernel size and stride $S$. This process effectively averages $S$ consecutive dimensions into one new dimension. We test with $S$ values of 96, 48, 24, 12, 6, 3, and 2, corresponding to reduced embedding dimensions of 8, 16, 32, 64, 128, 256, and 384, respectively.

Figure 9 illustrates DNABERT-S’s performance with these varying feature dimensions, in comparison to TNF. The results demonstrate that DNABERT-S’s embedding is quite resilient to dimension compression. It maintains nearly the same performance level even when reduced to 256 dimensions and only experiences a notable drop in performance when compressed to 32 dimensions. Notably, DNABERT-S still surpasses the 256-dimensional TNF feature even when its own dimensionality is reduced to just 16. This robustness to dimension reduction enhances its practical applicability in various genomic contexts.

Figure: Figure 9: DNABERT-S’s performance with varying embedding dimensions reduced by average pooling. DNABERT-S is robust to feature dimension reduction, and it even outperforms TNF with $16$-dimensional embedding.
Refer to caption: x6.png

### C.8 Species-Aware Embeddings on Genomics Function Predictions

This section evaluates the impact of species-aware embedding on other types of genomics analysis tasks. Specifically, we aim to understand how the species-aware embeddings perform compared to the embedding generated by the genome foundation model without contrastive training on various types of genome function prediction tasks.
We utilize the GUE benchmark [Zhou et al., ], which comprises a comprehensive collection of 28 datasets covering 7 diverse tasks, such as epigenetic marks prediction, promoter prediction, and transcription factor binding site prediction. Following our established methodology, each model is respectively used to generate embeddings for each DNA sequence, and a logistic regression model is trained for classification. The Matthews Correlation Coefficient (MCC) serves as the evaluation metric. We perform experiments on both DNABERT-2 and HyenaDNA.

**Table 12: Performance of DNABERT-2 and HyenaDNA before and after species-aware contrastive training on the GUE benchmark.**
|  | Epigenetic Marks Prediction |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | H3 | H3K14ac | H3K36me3 | H3K4me1 | H3K4me2 | H3K4me3 |
| DNABERT-2 | 66.87 | 38.92 | 43.39 | 31.79 | 30.16 | 23.68 |
| DNABERT-S | 69.02 | 37.45 | 41.91 | 32.76 | 27.86 | 22.16 |
| HyenaDNA w/o | 69.54 | 30.62 | 38.95 | 33.29 | 30.90 | 21.24 |
| HyenaDNA w/ | 69.67 | 32.06 | 38.61 | 32.29 | 30.32 | 23.21 |

The results, as detailed in Table 12, show that after species-aware conservative training, DNABERT-S underperforms DNABERT-2 on 19 of the 28 datasets, with an average loss of 2.12 in the MCC. Similarly, on HyenaDNA, the one went through species-aware training underperforms the original one on 18 out of 28 datasets.
As a model honed for species-aware tasks, DNABERT-S may exhibit reduced generalizability compared to broader genome foundation models (e.g., DNABERT-2). However, this specialized focus should not be viewed as a disadvantage. The intrinsic design of DNABERT-S—to prioritize species-specific features—may naturally limit its applicability to a broader range of tasks, a trade-off inherent to its specialized nature. For instance, when considering sequences from different species with varying functions, DNABERT-S’s training objective emphasizes species-specific similarities over functional commonalities, a choice that is deliberate for its targeted application.

## Appendix D i-Mix for Original Contrastive Learning

The Weighted SimCLR method treat each instance $\{x_{i}\}_{i=1}^{B}$ and $\{x_{i^{+}}\}_{i=1}^{B}$ as anchors, with each anchor associated with $2B-2$ negative samples. This section explains why integrating the i-Mix method into Weighted SimCLR (Sec. 3.1) doubles memory or training time.

For clarity, we define $x_{B+i}=x_{i^{+}}$. We also expand the virtual labels $v_{i}$, which are $2B$-dimensional, to identify the positive sample for each anchor. Here $v_{i,\tilde{i}}=1$ and $v_{i,j\neq\tilde{i}}=0$, $\tilde{i}=(B+i)\ \text{mod}\ 2B$.

When the i-Mix method treats every sample from $\{x_{i}\}_{i=1}^{2B}$ as anchor, it begins by shuffling $\{(x_{i},v_{i})\}_{i=1}^{2B}$ to generate $\{(\hat{x}_{i},\hat{v}_{i})\}_{i=1}^{2B}$. For each anchor $(x_{i},v_{i})$, it uses a simple mixup method to mix it up with $(\hat{x}_{i},\hat{v}_{i})$ before encoder layers of the model, using mixup coefficient $\lambda_{i}\sim\text{Beta}(\alpha,\alpha)$. This mixing results in:

$$ $(h_{i}^{0},v_{i}^{mix})=(\lambda_{i}x_{i}+(1-\lambda_{i})\hat{x}_{i},\lambda_{ i}v_{i}+(1-\lambda_{i})\hat{v}_{i}).$ $$

Moreover, for each anchor $x_{i}$, if i-Mix considers all samples from $\{x_{j}\}_{j=1}^{2B}\backslash\{x_{i}\}$ as either positive or negative samples, the model $f(\cdot)$ must process both the initial data instances $\{x_{i}\}_{i=1}^{2B}$ and mixed data instances $\{h_{i}^{0}\}_{i=1}^{2B}$ to generate their embeddings. Therefore, the i-Mix requires nearly twice more memory or training time compared to the method in Sec. 3.1 if using the same batch size.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1.	The paper lacks detailed information on hyperparameter selection and the tuning process, particularly regarding how these choices impact overall performance. Could the authors provide further details to clarify the influence of these selections on model stability and performance?
2.	What are the specific advantages of DNABERT-S over existing DNA classification models? A more detailed explanation would help elucidate the model’s unique contributions to the field.
3.	Could this method be extended to other biological sequences (e.g., RNA or protein sequences)? If so, what adjustments would be necessary to adapt to these cases?
4.	It would be helpful if the authors could further explain the measures taken to ensure fair comparisons in their experiments, including steps to prevent data leakage and whether model scales were controlled. While the downstream tasks performed well compared to multiple baselines, training/validation/testing based on sequence identity was not conducted, which could pose a risk of data leakage in this setup.
5.	One of the paper’s focuses is on exploring different embedding methods for DNA sequence classification, using a variety of pre-trained models to enhance classification performance. However, a detailed comparison of time and memory consumption across different embedding methods is missing, especially regarding:
(a)	Time and Memory Consumption of Embedding Methods: Could the authors clarify the computational time and memory usage differences between embedding methods during the training and inference stages?
(b)	Resource Analysis of Different Pre-Trained Models: How do time and memory consumption vary across different pre-trained models, and which models are most advantageous for specific tasks? A more detailed analysis could aid in model selection and optimization.
soundness: 3
strengths: 1.	Innovative enhancement of embedding representation for species differentiation: The paper improves the embedding representation of DNA sequences by introducing techniques such as DNA-Dropout and DNA-Double, which enable the model to better distinguish DNA sequences of different species. This improvement enhances the robustness of the embedding and the ability to capture the similarity of DNA structures, significantly improving the accuracy of species clustering and classification.
2.	Improving model generalization using contrastive learning: the paper’s Mixing of Streaming Instances (MI-Mix) and Course Contrastive Learning (C2LR) techniques gradually introduce training samples of increasing difficulty during fine-tuning, allowing the model to adapt more efficiently to species-rich macrogenomic data. This approach improves the model’s generalization ability in environments with scarce labels and high data diversity and is suitable for tasks such as macrogenomic binning and species classification.
3.	Practicality for macro-genomic data: The model is particularly suitable for macro-genomic data and biodiversity research. Through targeted fine-tuning and optimization, DNABERT-S significantly improves its performance on tasks such as species clustering and few-sample classification, providing a powerful tool for microbiomics and biodiversity research.
summary: The paper introduces DNABERT-S, a genome model focused on species-aware DNA embeddings to differentiate and cluster DNA sequences by species effectively. Building upon DNABERT-2, it incorporates two key innovations: Manifold Instance Mixup (MI-Mix) and Curriculum Contrastive Learning (C2LR). MI-Mix mixes hidden representations at random layers to enhance embedding robustness, whereas C2LR gradually presents increasingly challenging training samples to improve model generalization. Experiments across 23 datasets demonstrate DNABERT-S’s effectiveness, especially in species clustering, metagenomics binning, and few-shot classification tasks, showing that it significantly outperforms baselines, including by doubling clustering performance in Adjusted Rand Index (ARI). The model provides a robust and scalable solution to biodiversity studies and microbiome research in label-scarce environments, addressing limitations of previous genome foundation models in species differentiation.
weaknesses: 1.	Limited Novelty of Methodology: The paper employs Manifold Instance Mixup (MI-Mix) and Curriculum Contrastive Learning (C2LR), which are widely recognized in deep learning, limiting the originality of the methodology. The primary innovation lies in adapting these techniques specifically to metagenomic tasks rather than introducing novel technical advancements (see Sections 3.3 and 5.2). To strengthen this aspect, further evidence could clarify why these specific strategies are especially suited to metagenomics, particularly in addressing the shortcomings of traditional Mixup or contrastive learning for the specific challenges within this domain.
2.	Limited Scope of Comparative Experiments: The paper’s experimental validation is confined mainly to select metagenomic datasets, lacking a broader comparison with other current genomic models and widely used bioinformatics tools, such as database search techniques (refer to the experimental setup). Including these common baselines would provide a more comprehensive assessment of DNABERT-S’s effectiveness, highlighting the model’s practical applicability across diverse tasks and data types.
3.	Insufficient Visual Detail in Figures: In Figure 1, the current marker size obscures certain details, making it difficult to interpret the clustering and classification patterns. Adjusting the marker size could improve visibility, enhancing the visualization of data distribution across different methods.
4.	Figure Layout Issues: Figure 4’s layout partially overlaps with the text, which detracts from the paper’s readability and professionalism. Adjusting the figure’s placement could ensure proper spacing and clear separation between text and visuals.
5.	Ablation Study Lacks Detailed Discussion: While the ablation study indicates a substantial improvement when combining W. SimCLR and MI-Mix, the analysis does not sufficiently explore the mechanisms behind this synergy (see Section 5.3). A more detailed discussion, possibly with illustrative examples, would elucidate why the combined approach enhances data representation, providing stronger support for the method’s efficacy.
6.	Unclear Parameter Justification and Redundancy Reduction: The paper references data filtering criteria, such as selecting only species with at least 100 sequences for classification, but lacks a detailed rationale for these choices. Additionally, the redundancy reduction steps are not thoroughly explained, which could influence the results’ transparency and reliability. Supplementing the parameter selection with explicit reasoning in the data preprocessing steps would enhance the methodological rigor.
7.	Potential Data Leakage in Pre-Training: Given that GenBank serves as a substantial training source, there is a potential risk of overlap between the training and testing datasets. The study does not confirm whether this overlap was checked, which raises concerns about possible data leakage (see Section 5.1). Verifying this aspect and addressing any overlapping data would strengthen the reliability of the results.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Do any of the other baseline models see the same or similar labels to those used in the contrastive trianing?

How well does MI-Mix perform on other modalities - such as the original i-Mix task?
soundness: 2
strengths: This is a well written paper with a clear goal and solid choices in methods.

The main novel contribution is the Manifold Instance Mixup method(MI-Mix). They take previous work (i-Mix) which mix inputs in the batch such as images to create examples. They realize that there are no good ways to mix (blend) DNA sequence s at the input and instead apply the i-Mix methodology to the hidden states at a random layer of the network.

They perform extensive evaluations with baselines from VAE and transformer competitors and perform ablation studies.

The embeddings are clearly beneficial in downstream tasks and useful to the community.

They also create and share a benchmark dataset.
summary: This paper finetunes the DNABERT-2 model to generate species-aware embeddings from genomic sequences.
Current genome foundation models (such as DNABERT-2) are trained on language-modelling training tasks but do not develop discriminative embeddings.
The authors leverage genome species datasets and contrastive methods to learn embeddings that perform better on both unsupervised and supervised downstream tasks in species differentiation.

They develop a training scheme they name C^2LR for Curriculum Contrastive Learning.
In C^2LR the training of the model is in 2 phases:

Phase 1 : First, a weighted version of SimCLR is used to encourage embedding s from the same species to be near each other. Weighted SimCLR is SimCLR but with with higher weights for negative samples closer to the anchor.

Phase 2 :   Next, they introduce and use a contrastive loss called Manifold Instance Mixup. This is a more challenging task where they mix hidden states in a random layer and predict the proportion of the mix at the output.

They create and share an evaluation benchmark and perform extensive evaluations of the resulting embeddings.
weaknesses: In table 2, it is clear that MI-Mix performs very well on it's own. The paper would be much simpler and just as convincing re. performance if it focused on MI-Mix and dropped the curriculum and weighted SimCLR. Of course, getting the best result for a foundation model is also important.

In Table 1 there should also be at least one baseline for a model that has also gone through some kind of species differentiation training or finetuning. As noted in the text, baseline models are unlearnable or trained on generic language modelling objectives. As far as I can tell only DNABERT-S has had the luxury of using species labels in its training.

There could be more discussion about the original Manifold Mixup method (that i-Mix this method was partly inspired by), and how it relates to the new Manifold Instance Method.

Typos::

line 435 text overlaps with fig 4

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - While the authors claim to compare DNABERT-S to the strongest existing methods, there appear to be a number of comparable approaches that they overlook, especially in metagenomics binning. Some examples include COMEBin, a SOTA binning method based on contrastive multi-view representational learning (Wang et. al, 2024) and CLMB (Zhang et. al, 2022). How does DNABERT-S compare to these baselines?
- In the classification tasks, what is the performance of each baseline? While it is helpful to highlight the difference between DNABERT-S and the best-baseline, please include the full table of performance on each of the datasets for each of the models tested (in the appendix), as this information is still a helpful contribution.
- How does Manifold Instance Mixup differ from Manifold Mixup, introduced in Verma et a. (2019)? Please clarify this in the paper.
- Dataset complexity, meaning the number of genomes present, and the relative abundances of those genomes, can often influence the performance of a model in metagenomics binning. How does DNABERT-S perform at metagenomics binning when these two factors are changed? For example, what happens when the number of relative abundances of different species are highly imbalanced?
- Given the well established, recent CAMI II challenge for metagenomic binning (that the authors reference in the paper), how does DNABERT-2 compare to the tools benchmarked during this challenge (see Meyer et. al, 2022)? 


Meyer, F., Fritz, A., Deng, Z. L., Koslicki, D., Lesker, T. R., Gurevich, A., ... & McHardy, A. C. (2022). Critical assessment of metagenome interpretation: the second round of challenges. Nature methods, 19(4), 429-440.

Wang, Z., You, R., Han, H., Liu, W., Sun, F., & Zhu, S. (2024). Effective binning of metagenomic contigs using contrastive multi-view representation learning. Nature Communications, 15(1), 585.

Zhang, P., Jiang, Z., Wang, Y., & Li, Y. (2022). CLMB: Deep contrastive learning for robust metagenomic binning. In International Conference on Research in Computational Molecular Biology (pp. 326-348). Cham: Springer International Publishing.
soundness: 3
strengths: - The authors perform comprehensive testing across a variety of scenarios, as well as a thorough set of ablations. 
- While each of the separate components - Manifold Mixup, weighted SimCLR, curriculum learning -  have been introduced and utilized previously, this paper appears to be a new application of these methods to genome sequence data. 
- The model strongly outperforms baselines on metagenomics binning, species clustering, and species classification.
summary: This paper introduces DNABERT-S, a modified version of its precursor, DNABERT-2, that is applied to the task of differentiating DNA sequences between different species.
weaknesses: - In general, there is a lack of meta-level detail regarding the models that are being compared. It would be helpful to include tables that compare the number of parameters and embedding dimensions for each model/technique used.
- The use of Curriculum Contrastive Learning (C$^2$LR) strategy with the Manifold Instance Mixup (MI-Mix) loss could be an impactful contribution, however not enough work is done to show the utility of these approaches, and whether they are even helpful enough in light of computational tradeoffs. Indeed, removing C$^2$LR appears to change the performance by a small amount (only a -1.13 to -1.17 decrease).
- The paper is missing comparisons to similar models that employ contrastive learning for metagenomic binning tasks. 
- It is strange that the authors would not use the benchmarking set up in CAMI II to assess their model performance, given the built-in genome binning benchmark and comparison to SOTA tools.

Meta-review: The authors introduce techniques like DNA-Dropout and DNA-Double to improve embedding representations, enhancing the model's ability to distinguish between species. The paper utilizes contrastive learning techniques, such as MI-Mix and C2LR, to gradually introduce increasingly difficult training samples, improving the model's generalization ability, especially with limited labeled data. The paper was praised for its clear presentation and comprehensive testing across various scenarios.

However, the reviewers also felt that the techniques employed, like MI-Mix and C2LR, are not entirely novel and are adapted from existing deep learning methods, limiting the originality of the methodological contribution. The experimental validation focuses mainly on metagenomic datasets, lacking a broader comparison with other current genomic models and widely used bioinformatics tools.    
The ablation study, while showing improvement with combined methods, lacks a detailed exploration of the mechanisms behind the observed synergy The paper lacked detailed explanations for parameter choices and data preprocessing steps, raising concerns about transparency and potential data leakage from the training data.    

For these reasons, overall, the reviewers felt the paper is slightly below the acceptance threshold in its current state.

Decision: Reject