---
language:
- rna
library_name: transformers
tags:
- RNA
- language-model
- 3-UTR
license: mit
---

# UTRBERT-6mer

A BERT-base language model pre-trained on human 3' UTR sequences using 6-mer tokenization.
Part of the 3UTRBERT model family introduced in Yang et al. (2024).

## Architecture

| Parameter | Value |
|---|---|
| Layers | 12 |
| Attention heads | 12 |
| Embedding dimension | 768 |
| Intermediate size | 3072 |
| Vocabulary size | 4101 (5 special tokens + RNA 6-mers) |
| Positional encoding | Learned absolute (BERT-style) |
| Architecture | BERT-base |
| Max sequence length | 512 tokens (~515 nucleotides for 6-mer) |

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

## 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 article 22847354](https://figshare.com/articles/dataset/22847354) ([direct download](https://ndownloader.figshare.com/files/40597961))

### Checkpoint selection

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

## Parity Verification

Hidden-state representations verified against the original BertForMaskedLM implementation
at all 13 representation levels (embedding + 12 transformer layers). Max abs diff < 2.5e-5
for both eager and SDPA backends (float32 accumulation across 12 layers). Verified on GPU
with PyTorch 2.7 / 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-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)

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

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

# Intermediate layers
out_all = model(**enc, 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-6mer", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained("Taykhoom/UTRBERT-6mer", trust_remote_code=True)
model.eval()

enc = tokenizer(["AUG[MASK]AUG"], return_tensors="pt")
with torch.no_grad():
    logits = model(**enc).logits   # (1, seq_len, 4101)
```

### Fine-tuning

For sequence-level tasks, use the CLS token embedding as input to a prediction head.

```python
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 port uses a standalone `UTRBertModel` (custom `PreTrainedModel` subclass, `model_type: "utrbert"`).
`trust_remote_code=True` is required for both the tokenizer and the model.

The original implementation uses standard scaled dot-product attention (post-LN BERT-base).
This HF port adds `attn_implementation="sdpa"` and `attn_implementation="flash_attention_2"`
support, which were not part of the original codebase.

```python
# Faster inference with SDPA (default on modern PyTorch)
model = AutoModel.from_pretrained("Taykhoom/UTRBERT-6mer", trust_remote_code=True,
                                   attn_implementation="sdpa")

# Flash Attention 2 (requires flash-attn installed)
model = AutoModel.from_pretrained("Taykhoom/UTRBERT-6mer", trust_remote_code=True,
                                   attn_implementation="flash_attention_2")
```

## Citation

```bibtex
@article{yang2024_utrbert,
  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).
The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
and reviewed manually by Taykhoom Dalal.

## License

MIT, following the original repository.