File size: 4,388 Bytes
6af52a6 c549bd6 6af52a6 fc7070b 6af52a6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | ---
language:
- rna
library_name: transformers
tags:
- RNA
- language-model
license: mit
---
# RNA-FM
A 12-layer BERT-style transformer pre-trained on 23.7 million non-coding RNA sequences via masked language modelling.
## Architecture
| Parameter | Value |
|---|---|
| Layers | 12 |
| Attention heads | 20 |
| Embedding dimension | 640 |
| FFN dimension | 5120 |
| Vocabulary size | 25 |
| Positional encoding | Learned |
| Architecture | ESM-1b-style pre-LN Transformer |
| Max sequence length | 1024 tokens |
Vocabulary: `<cls>`, `<pad>`, `<eos>`, `<unk>`, A, C, G, U, R, Y, K, M, S, W, B, D, H, V, N, `-`, and 4 null-padding tokens, `<mask>`.
## Pretraining
- **Objective:** Masked language modelling (BERT-style, 15% masking rate)
- **Data:** RNAcentral100 -- 23.7 million non-coding RNA sequences
- **Source checkpoint:** `RNA-FM_pretrained.pth` from [cuhkaih/rnafm](https://huggingface.co/cuhkaih/rnafm)
## Parity Verification
Hidden-state representations verified identical (max abs diff = 0.00) to the original
implementation at all 13 representation levels (embedding + 12 transformer layers).
Verified on GPU (CUDA) with PyTorch 2.7 / transformers 4.57.6. SDPA numerical
differences are expected (~1e-4 max diff over 12 layers) and are not a correctness issue.
## Related Models
See the full [RNA-FM collection](https://huggingface.co/collections/Taykhoom/rna-fm-6a22c8c778d29e6dd3d437af).
| Model | Training data | Embedding dim | Notes |
|---|---|---|---|
| **[RNA-FM](https://huggingface.co/Taykhoom/RNA-FM)** | 23.7 M ncRNA | 640 | This model |
| [mRNA-FM](https://huggingface.co/Taykhoom/mRNA-FM) | 45 M CDS | 1280 | Codon (3-mer) tokenisation |
## Usage
### Embedding generation
```python
import torch
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNA-FM", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/RNA-FM", trust_remote_code=True)
model.eval()
sequences = [
"GGGUGCGAUCAUACCAGCACUAAUGCCCUCCUGGGAAGUCCUCGUGUUGCACCCCU",
"AUCGGGCUUAGCAUAGCUU",
]
# RNA-FM was trained on RNA sequences (U not T). T is not in the vocabulary.
# If your sequences use DNA notation, convert first:
# sequences = [s.replace("T", "U") for s in sequences]
enc = tokenizer(sequences, return_tensors="pt", padding=True)
with torch.no_grad():
out = model(**enc)
cls_emb = out.last_hidden_state[:, 0, :] # (batch, 640) -- CLS token
token_emb = out.last_hidden_state # (batch, seq_len, 640) -- per-token
# Intermediate layers
out_all = model(**enc, output_hidden_states=True)
layer6_emb = out_all.hidden_states[6] # layer 0 = embedding, 1-12 = transformer layers
```
### MLM logits
```python
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNA-FM", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained("Taykhoom/RNA-FM", trust_remote_code=True)
model.eval()
enc = tokenizer(["GGG<mask>GCGAU"], return_tensors="pt")
with torch.no_grad():
logits = model(**enc).logits # (1, seq_len, 25)
```
### Fine-tuning
Standard HF conventions. Use the CLS token embedding (`out.last_hidden_state[:, 0, :]`) as
input to a classification or regression head for sequence-level tasks.
## Implementation Notes
The original implementation uses `F.multi_head_attention_forward` (eager). This HF port adds
`attn_implementation="sdpa"` and `attn_implementation="flash_attention_2"` support, which were
not part of the original codebase.
Input sequences are expected to use RNA notation (U not T).
## Citation
```bibtex
@article{chen2022_rnafm,
title = {Interpretable {RNA} Foundation Model from Unannotated Data for Highly Accurate {RNA} Structure and Function Predictions},
author = {Chen, Jiayang and Hu, Zhihang and Sun, Siqi and Tan, Qingxiong and Wang, Yixuan and Yu, Qinze and Zong, Licheng and Hong, Liang and Xiao, Jin and Shen, Tao and King, Irwin and Li, Yu},
journal = {arXiv preprint arXiv:2204.00300},
year = {2022},
doi = {10.48550/arXiv.2204.00300}
}
```
## Credits
Original model and code by Chen et al. Source: [GitHub](https://github.com/ml4bio/RNA-FM).
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.
|