DNABERT-S

Weights and tokenizer for DNABERT-S (Zhou et al., Bioinformatics 2025), loaded with the shared MosaicBERT implementation from Taykhoom/MosaicBERT-updated.

DNABERT-S is a species-aware DNA embedding model fine-tuned from DNABERT-2 using curriculum contrastive learning. It generates embeddings that naturally cluster and segregate genomes from different species, enabling species identification, metagenomics binning, and evolutionary analysis.

This repo contains only weights and tokenizer files. The model code is loaded automatically from Taykhoom/MosaicBERT-updated via trust_remote_code=True.

Architecture

Parameter	Value
Layers	12
Attention heads	12
Embedding dimension	768
FFN hidden dimension	3,072 (GeGLU; bias-free 6,144-value gate projection)
Vocabulary size	4096 (BPE, identical to DNABERT-2)
Positional encoding	ALiBi (no hard length limit)
Normalization	LayerNorm (post-LN; eps=1e-12)
Architecture	Post-LN MosaicBERT encoder with unpadding and GeGLU
Max sequence length	~10,000 tokens (configured practical limit; ALiBi resizes dynamically)
Parameters	117,068,544 (including pooler; no MLM head)

Tokenization

Uses Byte Pair Encoding (BPE) tokenization via PreTrainedTokenizerFast, identical vocabulary to DNABERT-2. No k-mer pre-processing required.

Pretraining

• Objective: Curriculum contrastive learning (same-species pairs with i-Mix)
• Initialization: Fine-tuned from zhihan1996/DNABERT-2-117M
• Source checkpoint: pytorch_model.bin from zhihan1996/DNABERT-S

Parity Verification

Hidden-state representations verified identical (max abs diff = 0.00) to the original implementation at all 13 representation levels (embedding + 12 transformer layers). SDPA verified (max abs diff < 1e-4). Verified on GPU with PyTorch 2.7 / CUDA 12.9.

Related Models

See the full DNABERT collection.

Model	Architecture	Notes
DNABERT-3mer	BERT + k-mer	k=3
DNABERT-4mer	BERT + k-mer	k=4
DNABERT-5mer	BERT + k-mer	k=5
DNABERT-6mer	BERT + k-mer	k=6
DNABERT-2	MosaicBERT + BPE + ALiBi	Pre-trained
DNABERT-S	MosaicBERT + BPE + ALiBi	This model

Usage

Embedding generation

The current mean-pooling policy excludes padding, [CLS], and the terminal [SEP] token.

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/DNABERT-S", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/DNABERT-S", trust_remote_code=True)
model.eval()

sequences = ["ACGTAGCATCGGATCTATCTATCGACACTTGG", "ATCGATCGATCGATCG"]
enc = tokenizer(sequences, return_tensors="pt", padding=True)

with torch.no_grad():
    out = model(**enc)

cls_emb = out.last_hidden_state[:, 0, :]   # (batch, 768)

pool_mask = enc["attention_mask"].clone()
pool_mask[:, 0] = 0  # exclude [CLS]
sep_positions = enc["attention_mask"].sum(dim=1) - 1
pool_mask[torch.arange(pool_mask.size(0)), sep_positions] = 0  # exclude [SEP]

pool_mask = pool_mask.unsqueeze(-1).to(out.last_hidden_state.dtype)
mean_emb = (out.last_hidden_state * pool_mask).sum(dim=1)
mean_emb = mean_emb / pool_mask.sum(dim=1).clamp_min(1)  # (batch, 768)

Attention implementation

# SDPA (default on PyTorch >= 2.0)
model = AutoModel.from_pretrained("Taykhoom/DNABERT-S", trust_remote_code=True,
                                   attn_implementation="sdpa")

# Flash Attention 2
model = AutoModel.from_pretrained("Taykhoom/DNABERT-S", trust_remote_code=True,
                                   attn_implementation="flash_attention_2",
                                   torch_dtype=torch.bfloat16)

Implementation Notes

DNABERT-S is an embedding checkpoint and does not contain a trained masked-language-model head. Load it with AutoModel; AutoModelForMaskedLM raises a clear error rather than initializing random prediction weights.

The original DNABERT-S codebase uses a Triton-based flash attention implementation (flash_attn_triton.py). This HF port uses Taykhoom/MosaicBERT-updated which replaces it with the standard flash-attn package, and also adds attn_implementation="sdpa" support. These were not part of the original codebase.

Citation

@article{zhou2025_dnaberts,
  title   = {{DNABERT}-S: Pioneering Species Differentiation with Species-Aware {DNA} Embeddings},
  author  = {Zhou, Zhihan and Wu, Weimin and Ho, Harrison and Wang, Jiayi and Shi, Lizhen and Davuluri, Ramana V. and Wang, Zhong and Liu, Han},
  journal = {Bioinformatics},
  volume  = {41},
  number  = {Supplement_1},
  pages   = {i255--i264},
  year    = {2025},
  doi     = {10.1093/bioinformatics/btaf188}
}

Credits

Original DNABERT-S model and code by Zhou et al. Source: GitHub. Hugging Face port maintained by Taykhoom Dalal.

License

Apache 2.0, following the original repository.