GENA-LM-t2t-bigbird-base

GENA-LM-t2t-bigbird-base

Minimal HuggingFace port of the bigbird-base-t2t variant of GENA-LM -- a BigBird masked language model for long human DNA sequences, using byte-pair (BPE) tokenization.

This checkpoint uses the native transformers BigBird architecture (model_type=big_bird), so it loads with stock transformers and does not require trust_remote_code.

Architecture

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

BigBird automatically uses full ("original_full") attention for short sequences and block-sparse attention for long ones, following the standard transformers behaviour.

Pretraining

• Objective: Masked language modeling (15% masking, BigBird-style).
• Data: Human T2T genome assembly, augmented with 1000-genomes / gnomAD SNPs.
• Pretraining iterations: 1,070,000 (batch size 256).
• Source checkpoint: AIRI-Institute/gena-lm-bigbird-base-t2t

Parity Verification

This port re-serializes the original AIRI-Institute/gena-lm-bigbird-base-t2t weights into model.safetensors using the stock transformers BigBird implementation. All 13 representation levels and masked-LM logits are bit-exact (max abs diff = 0.00) against the original checkpoint for both full-attention and block-sparse inputs. Verified with PyTorch 2.7 / transformers 4.57.

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-t2t-bigbird-base")
model = AutoModel.from_pretrained("Taykhoom/GENA-LM-t2t-bigbird-base")
model.eval()

sequences = ["ACGTACGT" * 8]
enc = tokenizer(sequences, return_tensors="pt", padding=True)

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

cls_emb   = out.last_hidden_state[:, 0, :]
token_emb = out.last_hidden_state

MLM logits

from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/GENA-LM-t2t-bigbird-base")
model = AutoModelForMaskedLM.from_pretrained("Taykhoom/GENA-LM-t2t-bigbird-base")
model.eval()

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

Fine-tuning

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

Implementation Notes

This repository contains no custom modeling code -- it is the stock transformers BigBird model with the original GENA-LM weights, repackaged as model.safetensors. Stock BigBird exposes its native original_full / block_sparse modes via attention_type; it does not support the sdpa or flash_attention_2 dispatch backends. The input embeddings and MLM decoder are tied.

Stock transformers 4.57.6 has an invalid block-sparse output_attentions diagnostic. It returns an fp32 tensor shaped (batch, heads, query, key) with values in [0, 1]; padded keys receive zero mass and absolute global-query blocks have normalized rows. However, its upstream reconstruction is marked TODO and loses probability mass for valid non-global queries, so those rows do not sum to 1. Padded non-global query rows are zero, while padded queries in the final global block still sum to 1. Requesting the diagnostic in bf16 also raises a dtype mismatch. This does not change attention context, hidden states, or logits, and short/full-attention diagnostics are valid. Do not interpret the stock block-sparse diagnostic as standard post-softmax attention probabilities.

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.