GENA-LM-bert-base

Minimal HuggingFace port of the bert-base variant of GENA-LM -- a transformer masked language model for long human / multi-species DNA sequences, using byte-pair (BPE) tokenization.

Architecture

Vocabulary: 32,000 BPE tokens trained on DNA, including [CLS], [SEP], [PAD], [UNK], and [MASK].

Pretraining

• Objective: Masked language modeling (15% masking, BigBird-style).
• Data: Human T2T genome assembly (GCA_009914755.3).
• Pretraining iterations: 500,000 (batch size 256, sequence length 512).
• Source checkpoint: AIRI-Institute/gena-lm-bert-base

Parity Verification

All 13 representation levels (embedding + 12 transformer blocks) and the masked-LM logits were verified to be bit-exact (max abs diff = 0.00) against the original AIRI-Institute/gena-lm-bert-base weights, for the eager backend. The added sdpa and flash_attention_2 backends agree with eager up to the expected fused-kernel floating-point tolerance. Verified on GPU with PyTorch 2.7 / CUDA 12.9.

Related Models

See the full GENA-LM collection.

Usage

Embedding generation

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True)
model.eval()

sequences = ["ACGTACGTACGTACGT", "TTACGGGCATACGACGT"]
enc = tokenizer(sequences, return_tensors="pt", padding=True)

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

cls_emb   = out.last_hidden_state[:, 0, :]   # (batch, dim) -- CLS token
token_emb = out.last_hidden_state             # (batch, seq_len, dim)

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

MLM logits

from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True)
model.eval()

enc = tokenizer(["ACGT[MASK]CGTACGT"], return_tensors="pt")
with torch.no_grad():
    logits = model(**enc).logits   # (1, seq_len, vocab_size)

Faster attention backends

# SDPA (PyTorch 2.0+) -- recommended for production
model = AutoModel.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True,
                                  attn_implementation="sdpa")

# Flash Attention 2 (requires flash-attn) -- fastest on long sequences
import torch
model = AutoModel.from_pretrained("Taykhoom/GENA-LM-bert-base", trust_remote_code=True,
                                  attn_implementation="flash_attention_2",
                                  dtype=torch.bfloat16)

Fine-tuning

Standard HuggingFace conventions. For sequence-level tasks, pool over non-padding positions or use the [CLS] token embedding as input to a prediction head.

Implementation Notes

This is a minimal, self-contained reimplementation of the GENA-LM pre-LayerNorm BERT backbone. The only behavioural additions over the original (eager-only) code are the sdpa and flash_attention_2 attention backends, selectable via attn_implementation; the eager backend reproduces the original outputs bit-for-bit. The original NSP head and pooler are not included, since this port targets embedding and masked-LM use. AutoModel returns the backbone without a pooler; use the [CLS] hidden state or masked mean pooling for sequence embeddings. The input embeddings and MLM decoder are tied.

Citation

@article{fishman2025_genalm,
  title   = {{GENA-LM}: a family of open-source foundational {DNA} language models for long sequences},
  author  = {Fishman, Veniamin and Kuratov, Yuri and Shmelev, Aleksei and Petrov, Maxim and Penzar, Dmitry and Shepelin, Denis and Chekanov, Nikolay and Kardymon, Olga and Burtsev, Mikhail},
  journal = {Nucleic Acids Research},
  volume  = {53},
  number  = {2},
  pages   = {gkae1310},
  year    = {2025},
  doi     = {10.1093/nar/gkae1310}
}

Credits

Original model and code by Fishman, Kuratov, et al. (AIRI Institute). Source: GitHub. Hugging Face port maintained by Taykhoom Dalal.

License

MIT, following the original repository.