UTRBERT-6mer

UTRBERT-6mer

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

Architecture

Tokenization: raw RNA (or DNA) sequences are converted T->U, then split into overlapping 6-mers (stride 1). A sequence of length L produces L-5 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 6-mer tokens
• Data: Human 3' UTR sequences
• Source checkpoint: 6-new-12w-0/pytorch_model.bin from figshare software record 22851272 (direct download)

Checkpoint selection

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

Parity Verification

All 13 representation levels (embedding + 12 transformer layers) and MLM logits were verified against the original 6-new-12w-0 weights. Maximum float32 absolute differences were 3.31e-5 / 2.55e-5 for eager hidden states / logits and 2.24e-5 / 1.93e-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.

Usage

Embedding generation

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/UTRBERT-6mer", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/UTRBERT-6mer", 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-6mer", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained("Taykhoom/UTRBERT-6mer", 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, 4101)

Faster attention backends

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

# Flash Attention 2 (requires flash-attn)
model = AutoModel.from_pretrained(
    "Taykhoom/UTRBERT-6mer",
    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-6mer", 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.