GENA-LM-sparse-bigbird-base

GENA-LM-sparse-bigbird-base

Minimal HuggingFace port of the bigbird-base-sparse variant of GENA-LM -- a BigBird block-sparse masked language model for long human DNA sequences (up to 4096 BPE tokens), 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.
• Pretraining iterations: 810,000 (batch size 256).
• Source checkpoint: AIRI-Institute/gena-lm-bigbird-base-sparse

Block-sparse attention

The original model computes attention with DeepSpeed's block-sparse BigBird kernels. This port evaluates the same mathematical block-sparse equations in pure PyTorch: dense attention is restricted to allowed (query-block, key-block) pairs, with softmax over allowed keys only. The exact per-head checkpoint layout (master_layout) is consumed directly, using its top-left nb x nb region for nb sequence blocks. No DeepSpeed is required at inference.

The upstream wrapper forces Q/K/V to fp16 because its legacy Triton kernel accepts only fp16. This port instead evaluates the equations in the model's requested activation dtype for numerical stability. A standalone source-tensor reference is bit-exact with this mathematical implementation at all layers; exact legacy-kernel rounding/execution parity is not claimed because that kernel cannot run on Hopper (sm90).

For short sequences (<= 256 tokens / 4 blocks) the layout is fully dense, so the model behaves as an ordinary pre-LayerNorm BERT.

Parity Verification

For short sequences (fully dense layout), all 13 representation levels and the masked-LM logits are bit-exact (max abs diff = 0.00) against the original GENA-LM code run with dense attention (no DeepSpeed), confirming weight loading, the pre-LayerNorm structure, rotary embeddings, and the block-sparse reduction. For long sequences, every layer's probabilities match an independently rebuilt masked-softmax reference exactly, are zero outside the checkpoint's master_layout, and have rows that sum to 1. Verified on GPU with PyTorch 2.7 / CUDA 12.9. This is mathematical source parity, distinct from execution parity with the unavailable legacy DeepSpeed/Triton kernel.

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-sparse-bigbird-base", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/GENA-LM-sparse-bigbird-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)

MLM logits

from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/GENA-LM-sparse-bigbird-base", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained("Taykhoom/GENA-LM-sparse-bigbird-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+) evaluates the masked block-sparse attention with a fused kernel
model = AutoModel.from_pretrained("Taykhoom/GENA-LM-sparse-bigbird-base", trust_remote_code=True,
                                  attn_implementation="sdpa")

flash_attention_2 is not supported for the block-sparse checkpoints because Flash Attention cannot express the checkpoint-specific arbitrary block mask. Requesting it raises an explicit error rather than silently changing the attention pattern.

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 BigBird backbone. The block-sparse attention is computed as masked dense attention in pure PyTorch using the checkpoint's stored layout and requested activation dtype, removing the DeepSpeed dependency. The original NSP head and pooler are not included. 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. Rotary cache outputs are cloned in the caller's execution mode, so inference-mode logits can safely be followed by a grad-enabled embedding forward.

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.