README.md

---
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.