---
library_name: transformers
tags:
- biology
- DNA
- language-model
- genomics
license: apache-2.0
---

# DNABERT-4mer

Minimal HuggingFace port of the **4-mer** variant of
[DNABERT](https://huggingface.co/zhihan1996/DNA_bert_4) -- a BERT-base
masked language model pre-trained on the human reference genome using
overlapping 4-mer tokenization.

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

## Architecture

Standard BERT-base with a 4-mer DNA vocabulary.

| Parameter | Value |
|---|---|
| Layers | 12 |
| Attention heads | 12 |
| Embedding dimension | 768 |
| FFN hidden dimension | 3072 (GELU) |
| Vocabulary size | 261 (5 special + 256 DNA 4-mers) |
| Positional encoding | Learned absolute |
| Normalization | Post-LayerNorm (epsilon 1e-12) |
| Architecture | Bidirectional BERT encoder |
| Max sequence length | 512 tokens (510 k-mers; 513 nucleotides) |
| Runtime parameters | 87,034,629 |

### Tokenization

Input sequences must be pre-split into overlapping 4-mers (stride 1) with spaces
between tokens before calling the tokenizer. For example:

```
ATCGATG  ->  ATCG TCGA CGAT GATG
```

```python
def seq_to_kmers(seq, k=4):
    return " ".join(seq[i:i+k] for i in range(len(seq) - k + 1))
```

## Pretraining

- **Objective:** Masked Language Modeling
- **Data:** Human reference genome (GRCh38)
- **Source checkpoint:** `pytorch_model.bin` from [zhihan1996/DNA_bert_4](https://huggingface.co/zhihan1996/DNA_bert_4)

## Parity Verification

All 13 representation levels (embedding + 12 transformer layers) verified
against the source implementation (max abs diff = 9.79e-5); MLM logits match
with max abs diff = 2.52e-4. The source `dnabert_layer.BertModel` is a direct
subclass of `transformers.BertModel` with no modifications.
Verified on GPU with PyTorch 2.7.1 / CUDA 12.9.

## Related Models

See the full [DNABERT collection](https://huggingface.co/collections/Taykhoom/dnabert-6a20958f8ce004ea4e985e7b).

| Model | Architecture | Notes |
|---|---|---|
| [DNABERT-3mer](https://huggingface.co/Taykhoom/DNABERT-3mer) | BERT + k-mer | k=3 |
| **[DNABERT-4mer](https://huggingface.co/Taykhoom/DNABERT-4mer)** | **BERT + k-mer** | **k=4** |
| [DNABERT-5mer](https://huggingface.co/Taykhoom/DNABERT-5mer) | BERT + k-mer | k=5 |
| [DNABERT-6mer](https://huggingface.co/Taykhoom/DNABERT-6mer) | BERT + k-mer | k=6 |
| [DNABERT-2](https://huggingface.co/Taykhoom/DNABERT2) | MosaicBERT + BPE + ALiBi | Multi-species pre-trained |
| [DNABERT-S](https://huggingface.co/Taykhoom/DNABERT-S) | MosaicBERT + BPE + ALiBi | Species-aware |

## Usage

### Embedding generation

```python
import torch
from transformers import AutoTokenizer, AutoModel

def seq_to_kmers(seq, k=4):
    return " ".join(seq[i:i+k] for i in range(len(seq) - k + 1))

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

sequences = ["ATCGATCGATCG", "GCTAGCTAGCTA"]
kmer_seqs = [seq_to_kmers(s) for s in sequences]
enc = tokenizer(kmer_seqs, return_tensors="pt", padding=True)

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

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

# Mean-pool DNA k-mers only (exclude CLS, SEP, and padding)
content_mask = enc["attention_mask"].bool()
content_mask[:, 0] = False
sep_positions = enc["attention_mask"].sum(dim=1) - 1
batch_indices = torch.arange(len(sequences), device=content_mask.device)
content_mask[batch_indices, sep_positions] = False
mean_emb = (
    token_emb * content_mask.unsqueeze(-1)
).sum(dim=1) / content_mask.sum(dim=1, keepdim=True)

# Intermediate layers
out_all = model(**enc, output_hidden_states=True)
layer6_emb = out_all.hidden_states[6]
```

Sequences shorter than 4 nucleotides contain no k-mer tokens and therefore
cannot be mean-pooled or assigned a k-mer pseudo-likelihood.

### MLM logits

```python
from transformers import AutoModelForMaskedLM

model = AutoModelForMaskedLM.from_pretrained(
    "Taykhoom/DNABERT-4mer", trust_remote_code=True
)
tokens = seq_to_kmers("ATCGATCG", k=4).split()
tokens[2] = tokenizer.mask_token
enc = tokenizer(" ".join(tokens), return_tensors="pt")

with torch.no_grad():
    logits = model(**enc).logits   # (1, seq_len, 261)
```

### Faster attention backends

```python
# SDPA (PyTorch 2.0+)
model = AutoModel.from_pretrained("Taykhoom/DNABERT-4mer", trust_remote_code=True,
                                   attn_implementation="sdpa")

# Flash Attention 2 (requires flash-attn)
model = AutoModel.from_pretrained("Taykhoom/DNABERT-4mer", trust_remote_code=True,
                                   attn_implementation="flash_attention_2",
                                   dtype=torch.float16)
```

### Fine-tuning

For sequence-level tasks, mean-pool only k-mer positions as above or use the
CLS token embedding as input to a prediction head.

## Implementation Notes

The original DNABERT codebase has `BertModel` as a thin subclass of
`transformers.BertModel` with no modifications. This HF port uses
[Taykhoom/BERT-updated](https://huggingface.co/Taykhoom/BERT-updated) which adds
`attn_implementation="sdpa"` and `attn_implementation="flash_attention_2"`
support — these were not part of the original codebase.

## Citation

```bibtex
@article{ji2021_dnabert,
  title   = {{DNABERT}: pre-trained Bidirectional Encoder Representations from Transformers model for {DNA}-language in genome},
  author  = {Ji, Yanrong and Zhou, Zhihan and Liu, Han and Davuluri, Ramana V},
  journal = {Bioinformatics},
  volume  = {37},
  number  = {15},
  pages   = {2112--2120},
  year    = {2021},
  doi     = {10.1093/bioinformatics/btab083}
}
```

## Credits

Original DNABERT model and code by Ji et al. Source: [GitHub](https://github.com/jerryji1993/DNABERT).
Hugging Face port maintained by Taykhoom Dalal.

## License

Apache License 2.0, following the original repository.