Dna2vec

🚀 DNA2Vec：基於Transformer的DNA序列嵌入模型

dna2vec 是一個基於Transformer的模型，用於對DNA序列進行嵌入表示。本倉庫提供了該模型的實現，包含Hugging Face模型（hf_model）和本地訓練模型（local_model）。該模型可用於DNA序列比對、分類及其他基因組學應用。

🚀 快速開始

要使用該模型，需安裝所需的依賴項：

pip install transformers torch

✨ 主要特性

• 基於Transformer架構：在基因組數據上進行訓練。
• 無參考嵌入：能夠實現高效的序列檢索。
• 基於對比損失的自監督訓練：確保對序列相似性的學習更加穩健。
• 支持Hugging Face模型和自定義本地訓練模型。
• 通過DNA向量庫進行高效搜索：將全基因組比對簡化為局部搜索。

📦 安裝指南

要使用該模型，安裝所需的依賴項：

pip install transformers torch

💻 使用示例

基礎用法

from transformers import AutoModel, AutoTokenizer
import torch
import torch.nn as nn

def load_hf_model():
    hf_model = AutoModel.from_pretrained("roychowdhuryresearch/dna2vec", trust_remote_code=True)
    hf_tokenizer = AutoTokenizer.from_pretrained("roychowdhuryresearch/dna2vec", trust_remote_code=True)
    
    class AveragePooler(nn.Module):
        def forward(self, last_hidden, attention_mask):
            return (last_hidden * attention_mask.unsqueeze(-1)).sum(1) / attention_mask.sum(-1).unsqueeze(-1)
    
    hf_model.pooler = AveragePooler()
    return hf_model, hf_tokenizer, hf_model.pooler

高級用法

def get_embedding(dna_sequence):
    model, tokenizer, pooler = load_hf_model()
    tokenized_input = tokenizer(dna_sequence, return_tensors="pt", padding=True)
    with torch.no_grad():
        output = model(**tokenized_input)
    embedding = pooler(output.last_hidden_state, tokenized_input.attention_mask)
    return embedding.numpy()

# 示例用法
dna_seq = "ATGCGTACGTAGCTAGCTAGC"
embedding = get_embedding(dna_seq)
print("Embedding shape:", embedding.shape)

📚 詳細文檔

模型概述

DNA序列比對是基因組學中的一項重要任務，涉及將短DNA讀取片段映射到參考基因組中最可能的位置。傳統方法依賴於基因組索引和高效的搜索算法，而最近的進展則利用基於Transformer的模型將DNA序列編碼為向量表示。

dna2vec 框架引入了一種無參考DNA嵌入（RDE）Transformer模型，該模型將DNA序列編碼到一個共享的向量空間中，從而實現高效的相似性搜索和序列比對。

模型細節

模型架構

該Transformer模型包含：

• 12個注意力頭
• 6個編碼器層
• 嵌入維度：1020
• 詞彙表大小：10,000
• 基於餘弦相似度的序列匹配
• Dropout：0.1
• 訓練：採用餘弦退火學習率調度

訓練細節

數據集

訓練數據由從不同物種的各種染色體中採樣的DNA序列組成。該數據集覆蓋了約2%的人類基因組，確保了模型在不同序列上的泛化能力。讀取片段使用ART MiSeq模擬生成，插入和刪除率存在變化。

訓練過程

• 自監督學習：基於對比損失的訓練。
• 動態長度序列：DNA片段長度為800 - 2000，讀取片段採樣範圍為 [150, 500]。
• 噪聲增強：40%的訓練讀取片段中存在1 - 5%的隨機鹼基替換。
• 批量大小：16，採用梯度累積。

評估

該模型與傳統比對工具（Bowtie - 2）和其他基於Transformer的基線模型（DNABERT - 2、HyenaDNA）進行了對比評估。評估指標包括：

• 比對召回率：高質量讀取片段的比對召回率 > 99%。
• 跨物種遷移：能夠成功比對不同物種的序列，包括嗜熱棲熱菌（Thermus Aquaticus）和褐家鼠（Rattus Norvegicus）。

🔧 技術細節

DNA序列比對是基因組學中的關鍵任務，傳統方法依賴於基因組索引和搜索算法，而本模型利用Transformer架構將DNA序列編碼為向量表示。dna2vec 框架的無參考DNA嵌入（RDE）Transformer模型，通過自監督學習和對比損失，將DNA序列映射到共享向量空間，實現高效的相似性搜索和序列比對。模型架構包含12個注意力頭、6個編碼器層，嵌入維度為1020，詞彙表大小為10,000。訓練過程中採用動態長度序列、噪聲增強等策略，以提高模型的泛化能力和魯棒性。

📄 許可證

本項目採用MIT許可證。

引用

如果您使用了該模型，請引用以下文獻：

@article{10.1093/bioinformatics/btaf041,
    author = {Holur, Pavan and Enevoldsen, K C and Rajesh, Shreyas and Mboning, Lajoyce and Georgiou, Thalia and Bouchard, Louis-S and Pellegrini, Matteo and Roychowdhury, Vwani},
    title = {Embed-Search-Align: DNA Sequence Alignment using Transformer models},
    journal = {Bioinformatics},
    pages = {btaf041},
    year = {2025},
    month = {02},
    issn = {1367-4811},
    doi = {10.1093/bioinformatics/btaf041},
    url = {https://doi.org/10.1093/bioinformatics/btaf041},
    eprint = {https://academic.oup.com/bioinformatics/advance-article-pdf/doi/10.1093/bioinformatics/btaf041/61778456/btaf041.pdf},
}

更多詳細信息，請查看完整論文。