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},
}

更多详细信息，请查看完整论文。