[논문 리뷰] DNABERT-S: Pioneering Species Differentiation with Species-Aware DNA Embeddings

Zhihan Zhou, Weimin Wu외 6명|PubMed|2024. 02. 13.

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 서열에 대해 구별 가능한 임베딩을 산출하는 게놈 파운데이션 모델을 개발합니다.
• 임베딩 품질 향상을 위한 두 가지 학습 혁신—Curriculum Contrastive Learning (C2LR) 및 Manifold Instance Mixup (MI-Mix)—을 도입합니다.
• 베이스라인 대비 향상을 입증하기 위해 대규모 메타게놈 및 클러스터링 벤치마크에서 평가합니다.

제안 방법

• 두 단계 대조 학습: 1단계는 Weighted SimCLR을 사용해 유사한 서열의 그룹화를 촉진하고 비유사한 서열의 분리를 촉진합니다; 2단계는 Manifold Instance Mixup (MI-Mix)을 적용해 중간 은닉 표현을 혼합하여 더 큰 난이도의 대조 중심(anchor)을 만듭니다.
• 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 커리큘럼 대조 학습과 MI-Mix가 클러스터링, 분류(소수 샷), 및 메타게놈 바이닝 성능에 어떤 영향을 미치는가?
• RQ3 DNABERT-S가 제한된 라벨로의 클러스터링, 분류 및 바이닝 작업에서 기존의 게놈 파운데이션 모델 및 전통적 임베딩 방법을 능가하는가?
• RQ4 분류에서 보지 못한 종이나 합성 데이터셋에 대한 DNABERT-S의 일반화는 어떠한가?

주요 결과

• DNABERT-S는 평가된 데이터세트 전반에서 지속적으로 최고의 클러스터링 성능(ARI)을 달성하며 평균적으로 가장 강력한 베이스라인의 ARI를 두 배로 향상시킵니다.
• 분류에서 DNABERT-S는 더 많은 라벨을 가진 베이스라인과 비교해도 라벨 데이터의 20%만으로도 더 나은 성능을 달성하며, 일부 합성 설정에서 소수 예제로도 높은 F1을 달성합니다.
• 메타게놈 바이닝에서 DNABERT-S는 더 많은 종을 식별하고 더 높은 F1 점수를 달성하여 여러 데이터세트에서 베이스라인보다 큰 차이로 우수한 성능을 보입니다.
• 없앰 연구는 Weighted SimCLR와 MI-Mix의 조합이 C2LR 커리큘럼 하에서 최상의 성능을 보이며, 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.