[Paper Review] DNABERT-S: Pioneering Species Differentiation with Species-Aware DNA Embeddings

Zhihan Zhou, Weimin Wu+6 more|PubMed|Feb 13, 2024

Genomics and Phylogenetic Studies9 references11 citations

TL;DR

DNABERT-S learns species-aware DNA embeddings using curriculum contrastive learning and Manifold Instance Mixup to markedly improve clustering, few-shot classification, and metagenomics binning across diverse datasets.

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²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/MAGlCS-LAB/DNABERT_S.

Motivation & Objective

• Tackle the lack of effective DNA embeddings for species-aware tasks in unlabeled or limited-label settings.
• Develop a genome foundation model that yields discriminative embeddings for DNA sequences across many species.
• Introduce two training innovations—Curriculum Contrastive Learning (C2LR) and Manifold Instance Mixup (MI-Mix)—to improve embedding quality.
• Evaluate on large-scale metagenomics and clustering benchmarks to demonstrate improvements over baselines.

Proposed method

• Two-phase contrastive training: phase I uses Weighted SimCLR to encourage grouping of similar sequences and separation of dissimilar ones; phase II applies Manifold Instance Mixup (MI-Mix) to create harder contrastive anchors by mixing intermediate hidden representations.
• MI-Mix randomly selects a hidden layer, mixes intermediate representations of anchor sequences, and uses the mixed labels to compute a contrastive loss.
• Curriculum strategy progressively increases challenge of anchors and samples, enhancing generalization in low-label scenarios.
• DNABERT-S builds on pre-trained DNABERT-2 as a starting point and fine-tunes with contrastive objectives; uses mean-pooled last hidden state as embedding.

Figure 1: TSNE visualization of the DNA embeddings generated by different methods on a CAMI2 [Meyer et al., 2022 ] 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 pronou

Experimental results

Research questions

• RQ1 Can DNABERT-S produce embeddings that cluster by species across diverse genomes without extensive labeled data?
• RQ2 How do Curriculum Contrastive Learning and Manifold Instance Mixup affect clustering, classification (few-shot), and metagenomics binning performance?
• RQ3 Does DNABERT-S outperform existing genome foundation models and traditional embedding methods on clustering, classification with limited labels, and binning tasks?
• RQ4 What is the generalization of DNABERT-S to unseen species or synthetic datasets in classification?

Key findings

• DNABERT-S consistently achieves the best clustering performance (ARI) across evaluated datasets, doubling the ARI of the strongest baseline on average.
• In classification, DNABERT-S attains better performance with as little as 20% of labeled data compared to baselines with more labels, and in some synthetic settings achieves high F1 with few examples.
• For metagenomics binning, DNABERT-S identifies more species and achieves higher F1 scores, outperforming baselines by a substantial margin on multiple datasets.
• Ablation studies show that the combination of Weighted SimCLR and MI-Mix under the C2LR curriculum yields the best performance, with MI-Mix providing significant gains over single-loss baselines.
• The model demonstrates strong generalization to unseen species, achieving robust classification results on synthetic datasets with limited supervision.

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 2 LR) strategy that progressively provides more challenging contrastive anchors to the model in two different phases.