--- 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](https://github.com/yangyn533/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](https://doi.org/10.6084/m9.figshare.22851191.v1) ([direct download](https://ndownloader.figshare.com/files/40597919)) ### 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](https://huggingface.co/collections/Taykhoom/utrbert-6a2059e7d24778aee83af7bc). | Model | k-mer | Vocab size | Notes | |---|---|---|---| | [UTRBERT-3mer](https://huggingface.co/Taykhoom/UTRBERT-3mer) | 3 | 69 | | | [UTRBERT-4mer](https://huggingface.co/Taykhoom/UTRBERT-4mer) | 4 | 261 | | | **[UTRBERT-5mer](https://huggingface.co/Taykhoom/UTRBERT-5mer)** | 5 | 1029 | | | [UTRBERT-6mer](https://huggingface.co/Taykhoom/UTRBERT-6mer) | 6 | 4101 | | ## Usage ### Embedding generation ```python 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 ```python 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 ```python # 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. ```python 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`](https://huggingface.co/Taykhoom/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 ```bibtex @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](https://github.com/yangyn533/3UTRBERT). Hugging Face port maintained by Taykhoom Dalal. ## License The released checkpoint weights are [CC BY 4.0](https://creativecommons.org/licenses/by/4.0/), as specified by the source figshare record. The original repository's code is MIT licensed.