Fix model correctness and Hugging Face compatibility

45a5e86 verified about 7 hours ago

6.54 kB

library_name: transformers
tags:
  - biology
  - RNA
  - language-model
  - 3-UTR
license: cc-by-4.0

UTRBERT-5mer

Minimal HuggingFace port of the 5-mer variant of 3UTRBERT -- a BERT-base language model pre-trained on aggregated human mRNA 3' UTR sequences.

Architecture

Parameter	Value
Layers	12
Attention heads	12
Embedding dimension	768
FFN hidden dimension	3072 (GELU)
Vocabulary size	1029 (5 special tokens + RNA 5-mers)
Positional encoding	Learned absolute (BERT-style)
Normalization	LayerNorm (post-LN, eps=1e-12)
Architecture	Post-LN BERT-base encoder
Max sequence length	512 tokens (up to 514 raw nucleotides)

Tokenization: raw RNA (or DNA) sequences are converted T->U, then split into overlapping 5-mers (stride 1). A sequence of length L produces L-4 tokens. A [CLS] and [SEP] token are prepended and appended by the tokenizer. The official preprocessing script capped raw sequences at 510 nucleotides.

Pretraining

Objective: Masked Language Modeling (MLM) on 5-mer tokens
Data: Human 3' UTR sequences
Source checkpoint: 5-new-12w-0/pytorch_model.bin from figshare software record 22851191 (direct download)

Checkpoint selection

The only publicly released pre-trained checkpoint for the 5-mer variant is 5-new-12w-0.

Parity Verification

All 13 representation levels (embedding + 12 transformer layers) and MLM logits were verified against the original 5-new-12w-0 weights. Maximum float32 absolute differences were 1.24e-5 / 6.72e-5 for eager hidden states / logits and 8.58e-6 / 7.34e-5 for SDPA. Verified on GPU with PyTorch 2.7.1 / CUDA 12.9 and transformers 4.57.6.

Related Models

See the full UTRBERT collection.

Model	k-mer	Vocab size
UTRBERT-3mer	3	69
UTRBERT-4mer	4	261
UTRBERT-5mer	5	1029
UTRBERT-6mer	6	4101

Usage

Embedding generation

import torch
from transformers import AutoTokenizer, AutoModel

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

sequences = ["AUGCAUGCAUGCAUGCAUGC", "GCGCGCGCGCGCGCGCGCGC"]
enc = tokenizer(
    sequences,
    return_tensors="pt",
    padding=True,
    truncation=True,
    max_length=512,
    return_special_tokens_mask=True,
)
model_inputs = {k: v for k, v in enc.items() if k != "special_tokens_mask"}

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

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

# Mean-pool only biological k-mer tokens (exclude padding, CLS, and SEP).
pool_mask = enc["attention_mask"].bool() & ~enc["special_tokens_mask"].bool()
mean_emb = (
    (token_emb * pool_mask.unsqueeze(-1)).sum(dim=1)
    / pool_mask.sum(dim=1, keepdim=True)
)

# Intermediate layers
out_all = model(**model_inputs, output_hidden_states=True)
layer6_emb = out_all.hidden_states[6]         # (batch, seq_len, 768)

MLM logits

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/UTRBERT-5mer", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained("Taykhoom/UTRBERT-5mer", trust_remote_code=True)
model.eval()

# Tokenize first, then replace one overlapping k-mer token with MASK.
enc = tokenizer(["AUGCAUGCAUG"], return_tensors="pt")
mask_position = 3  # position 0 is CLS
enc["input_ids"][0, mask_position] = tokenizer.mask_token_id
with torch.no_grad():
    logits = model(**enc).logits   # (1, seq_len, 1029)

Faster attention backends

# SDPA (PyTorch 2.0+)
model = AutoModel.from_pretrained(
    "Taykhoom/UTRBERT-5mer",
    trust_remote_code=True,
    attn_implementation="sdpa",
)

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

Fine-tuning

For sequence-level tasks, use the CLS embedding or the masked mean-pooled k-mer embedding above as input to a prediction head.

import torch.nn as nn
from transformers import AutoModel

model = AutoModel.from_pretrained("Taykhoom/UTRBERT-5mer", trust_remote_code=True)

class UTRClassifier(nn.Module):
    def __init__(self, base, num_labels):
        super().__init__()
        self.base = base
        self.head = nn.Linear(768, num_labels)

    def forward(self, input_ids, attention_mask):
        cls = self.base(input_ids, attention_mask=attention_mask).last_hidden_state[:, 0]
        return self.head(cls)

Implementation Notes

This checkpoint uses the shared BERT-updated code backend through its cross-repository auto_map, plus the custom k-mer tokenizer stored in this repository. trust_remote_code=True is required. Loading a local checkpoint directory also requires network access to BERT-updated, unless that code is already cached.

The original implementation uses eager scaled dot-product attention. This port adds selectable sdpa and flash_attention_2 inference backends.

Citation

@article{yang2024_3utrbert,
  title   = {Deciphering 3'{UTR} Mediated Gene Regulation Using Interpretable Deep Representation Learning},
  author  = {Yang, Yuning and Li, Gen and Pang, Kuan and Cao, Wuxinhao and Zhang, Zhaolei and Li, Xiangtao},
  journal = {Advanced Science},
  volume  = {11},
  number  = {39},
  pages   = {e2407013},
  year    = {2024},
  doi     = {10.1002/advs.202407013}
}

Credits

Original model and code by Yang et al. Source: GitHub. Hugging Face port maintained by Taykhoom Dalal.

License

The released checkpoint weights are CC BY 4.0, as specified by the source figshare record. The original repository's code is MIT licensed.