SpliceBERT-1024nt

Minimal HuggingFace port of the vertebrate-1024nt variant of SpliceBERT -- a BERT-based RNA language model trained with masked language modeling on variable-length 64-1024 nt fragments from over 2 million primary RNA sequences across 72 vertebrates.

Architecture

Parameter	Value
Layers	6
Attention heads	16
Embedding dimension	512
FFN hidden dimension	2048 (GELU)
Vocabulary size	10
Positional encoding	Learned absolute
Normalization	LayerNorm (post-residual, eps=1e-12)
Architecture	Post-LN BERT encoder
Max sequence length	1024 nt (1026 tokens with [CLS]/[SEP])
Checkpoint size	~19.7M parameters

Vocabulary: [PAD]=0, [UNK]=1, [CLS]=2, [SEP]=3, [MASK]=4, N=5, A=6, C=7, G=8, T=9. Input U is normalized to T.

Pretraining

• Objective: Masked language modeling (MLM)
• Data: >2 million vertebrate primary RNA sequences from 72 species
• Sequence format: Single-nucleotide tokenization with spaces; U converted to T
• Source checkpoint: SpliceBERT.1024nt/pytorch_model.bin (from zenodo:7995778)

Checkpoint selection

The 1024nt variant is the primary SpliceBERT model trained on variable-length vertebrate sequences. Use this variant for general-purpose RNA embedding. The 510nt variants are trained on fixed-length fragments and require exact 510nt input.

Parity Verification

Hidden-state representations verified (max abs diff < 1e-5) against the original checkpoint at all 7 representation levels (embedding + 6 transformer layers), for both eager and sdpa attention backends. Verified on GPU with PyTorch 2.7.1 / CUDA 12.9.

Related Models

See the full SpliceBERT collection.

Model	Context	Training data	Notes
SpliceBERT-1024nt	1024 nt	72 vertebrates	This model
SpliceBERT-510nt	510 nt (fixed)	72 vertebrates	Fixed-length; requires exact 510 nt input
SpliceBERT-human-510nt	510 nt (fixed)	Human only	Human-specific; requires exact 510 nt input

Usage

Embedding generation

The tokenizer automatically handles U->T conversion and single-nucleotide spacing. Pass raw sequences directly.

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/SpliceBERT-1024nt", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/SpliceBERT-1024nt", trust_remote_code=True)
model.eval()

seq = "ACGUACGUACGUACGU"  # U->T handled automatically
enc = tokenizer(seq, return_tensors="pt")

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

# Mean pooling over non-special tokens
hidden = out.last_hidden_state[0]  # (seq_len+2, 512)
token_emb = hidden[1:-1]           # strip [CLS] and [SEP]
mean_emb = token_emb.mean(dim=0)   # (512,)

# Intermediate layers
layer3_emb = out.hidden_states[3]  # (1, seq_len+2, 512)

MLM logits

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/SpliceBERT-1024nt", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained("Taykhoom/SpliceBERT-1024nt", trust_remote_code=True)
model.eval()

seq = "A C G [MASK] A C G T"
enc = tokenizer(seq, return_tensors="pt")
with torch.no_grad():
    logits = model(**enc).logits  # (1, seq_len, 10)

Fine-tuning

Standard HF conventions. For sequence-level tasks, use mean pooling of non-special token positions (positions 1 to -1) as input to a prediction head.

Implementation Notes

The original checkpoint was saved as BertForMaskedLM with transformers==4.24.0. This port uses BERT-updated, which adds attn_implementation="sdpa" and attn_implementation="flash_attention_2" support not present in the original codebase.

The pooler weights (pooler.dense) are not present in the original checkpoint and are not included in the saved model.safetensors. add_pooling_layer=True (the default) allocates the pooler layer but its weights are randomly initialized -- do not use pooler_output without fine-tuning.

model = AutoModel.from_pretrained("Taykhoom/SpliceBERT-1024nt",
                                  trust_remote_code=True,
                                  attn_implementation="sdpa")

Citation

@article{chen2024_splicebert,
  title   = {Self-supervised learning on millions of primary {RNA} sequences from 72 vertebrates improves sequence-based {RNA} splicing prediction},
  author  = {Chen, Ken and Zhou, Yue and Ding, Maolin and Wang, Yu and Ren, Zhixiang and Yang, Yuedong},
  journal = {Briefings in Bioinformatics},
  volume  = {25},
  number  = {3},
  pages   = {bbae163},
  year    = {2024},
  doi     = {10.1093/bib/bbae163}
}

Credits

Original model and code by Chen et al. Source: GitHub. The HF conversion code was authored primarily by Claude Code and reviewed manually by Taykhoom Dalal.

License

The checkpoint weights are distributed under CC BY 4.0 by the upstream Zenodo record. The original SpliceBERT source code is BSD 3-Clause licensed.