Dna2vec

🚀 DNA2Vec: TransformerベースのDNA配列埋め込みモデル

このリポジトリは、dna2vecの実装を提供します。dna2vecは、DNA配列の埋め込みに特化したTransformerベースのモデルです。Hugging Face (hf_model)とローカルで学習したモデル (local_model)の両方を含んでいます。このモデルは、DNA配列のアライメント、分類、その他のゲノミクスアプリケーションに使用できます。

🚀 クイックスタート

このモデルを使用するには、必要な依存関係をインストールする必要があります。

pip install transformers torch

✨ 主な機能

• ゲノムデータで学習されたTransformerベースのアーキテクチャ。
• 効率的な配列検索を可能にする参照フリーの埋め込み。
• 堅牢な配列類似性学習を保証する自己教師付き学習のためのコントラスト損失。
• Hugging Faceとカスタム学習したローカルモデルの両方をサポート。
• DNAベクトルストアを通じた効率的な検索で、全ゲノムアライメントを局所検索に削減。

📚 ドキュメント

モデルの概要

DNA配列アライメントは、短いDNAリードを参照ゲノム内の最も可能性の高い位置にマッピングする重要なゲノミクスタスクです。従来の方法はゲノムインデックスと効率的な検索アルゴリズムに依存していましたが、最近の進歩ではTransformerベースのモデルを利用してDNA配列をベクトル表現にエンコードしています。

dna2vecフレームワークは、参照フリーDNA埋め込み (RDE) Transformerモデルを導入しています。このモデルはDNA配列を共有ベクトル空間にエンコードし、効率的な類似性検索と配列アライメントを可能にします。

モデルの詳細

モデルアーキテクチャ

Transformerモデルは以下の要素で構成されています。

• 12個のアテンションヘッド
• 6つのエンコーダ層
• 埋め込み次元: 1020
• 語彙サイズ: 10,000
• コサイン類似度に基づく配列マッチング
• ドロップアウト: 0.1
• 学習: コサインアニーリング学習率スケジューリング

モデルの使用方法

Hugging Faceモデルの読み込み

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

モデルの使用

モデルを読み込んだら、DNA配列の埋め込みを取得するために使用できます。

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配列で構成されています。このデータセットはヒトゲノムの約2%をカバーしており、異なる配列に対する汎化能力を保証します。リードはART MiSeqシミュレーションを使用して生成され、挿入と欠失率に変動があります。

学習手順

• 自己教師付き学習: コントラスト損失ベースの学習。
• 動的な長さの配列: 長さ800 - 2000のDNA断片で、リードは[150, 500]でサンプリング。
• ノイズ増幅: 学習リードの40%に1 - 5%のランダムな塩基置換。
• バッチサイズ: 勾配累積を伴う16。

評価

このモデルは、従来のアライナー (Bowtie-2) と他のTransformerベースのベースライン (DNABERT-2, HyenaDNA) と比較して評価されました。評価指標は以下の通りです。

• アライメント再現率: 高品質のリードで99%以上。
• 種をまたぐ転移: Thermus Aquaticus や Rattus Norvegicus などの異なる種の配列を正常にアライメント。

引用

このモデルを使用する場合は、以下を引用してください。

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

詳細については、全文を確認してください。

📄 ライセンス

このプロジェクトはMITライセンスの下で公開されています。