[论文解读] DNABERT-S: Pioneering Species Differentiation with Species-Aware DNA Embeddings

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

Genomics and Phylogenetic Studies参考文献 9被引用 11

一句话总结

DNABERT-S 使用课程对比学习和流形实例混合（Manifold Instance Mixup）来学习物种感知的 DNA 表达，使聚类、少样本分类和宏基因组分箱在跨多样本数据集上显著提升。

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.

研究动机与目标

• 解决在无标签或标注有限的情况下，针对物种感知任务缺乏有效DNA嵌入的问题。
• 开发一个基因组基础模型，能够为跨多种物种的DNA序列提供判别性嵌入。
• 引入两项训练创新——课程对比学习（C2LR）与流形实例混合（MI-Mix）——以提升嵌入质量。
• 在大规模宏基因组学和聚类基准上进行评估，以证明相较基线的改进。

提出的方法

• 两阶段对比学习：阶段I 使用加权的 SimCLR 以促进相似序列的聚集、不同序列的分离；阶段 II 应用流形实例混合（MI-Mix）通过混合中间隐藏表示来创建更难的对比锚点。
• MI-Mix 随机选择一个隐藏层，混合锚序列的中间表示，并使用混合标签来计算对比损失。
• 课程策略逐步增加锚点和样本的难度，在低标签场景下提升泛化能力。
• DNABERT-S 以预训练的 DNABERT-2 为起点，并通过对比目标进行微调；使用最后一层的均值池化隐藏状态作为嵌入。

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

实验结果

研究问题

• RQ1 在不大量标注数据的情况下，DNABERT-S 能否产生按物种聚类的嵌入，在多样的基因组之间？
• RQ2 课程对比学习和流形实例混合如何影响聚类、分类（少样本）和宏基因组分箱的性能？
• RQ3 在聚类、少量标注的分类以及分箱任务上，DNABERT-S 是否优于现有的基因组基础模型和传统嵌入方法？
• RQ4 DNABERT-S 在未见物种或合成数据集上的分类泛化能力如何？

主要发现

• DNABERT-S 在所评估的数据集上一直实现最佳聚类性能（ARI），平均ARI是最强基线的两倍。
• 在分类方面，若仅有约20%的带标签数据，DNABERT-S 的表现就优于使用更多标签的基线，在某些合成场景下，少量样本也能达到较高的F1。
• 在宏基因组分箱方面，DNABERT-S 识别出更多物种并获得更高的 F1 分数，在多个数据集上对基线有显著超越。
• 消融研究表明，在 C2LR 课程下，结合加权 SimCLR 与 MI-Mix 能获得最佳性能，且 MI-Mix 相比单损失基线带来显著提升。
• 该模型在未见物种上的泛化能力强，在有限监督的合成数据集上实现稳健的分类结果。

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.