README.md

---
language:
- rna
library_name: transformers
tags:
- RNA
- language-model
- MSA
license: mit
---

# RNA-MSM

Multiple sequence alignment-based RNA language model trained on homologous RNA
sequence alignments from the RNAcmap pipeline.

## Architecture

| Parameter | Value |
|---|---|
| Layers | 10 |
| Attention heads | 12 |
| Embedding dimension | 768 |
| FFN dimension | 3072 |
| Vocabulary size | 12 |
| Positional encoding | Learned (sequence) + learned scalar (alignment row) |
| Architecture | Axial MSA Transformer (row + column self-attention) |
| Max sequence length | 1024 |
| Max alignment depth | 1024 |

**Input format:** RNA-MSM takes 3D input `(batch, num_alignments, seqlen)`. Each
alignment is a set of homologous RNA sequences of equal length (an MSA). The model
applies row self-attention (across sequence positions) and column self-attention
(across alignment rows) at each of the 10 transformer layers.

### Vocabulary

| Token | ID | Token | ID |
|---|---|---|---|
| `<cls>` | 0 | `U` | 7 |
| `<pad>` | 1 | `X` | 8 |
| `<eos>` | 2 | `N` | 9 |
| `<unk>` | 3 | `-` | 10 |
| `A` | 4 | `<mask>` | 11 |
| `G` | 5 | | |
| `C` | 6 | | |

Each sequence is prepended with `<cls>` (id 0). No `<eos>` token is appended.

## Pretraining

- **Objective:** Masked language modeling on RNA MSAs (masking ~15% of tokens)
- **Data:** RNA homologous sequences searched by RNAcmap from non-redundant RNA
  databases
- **Source checkpoint:** `RNA_MSM_pretrained.ckpt`
  ([original Google Drive link](https://drive.google.com/file/d/11A-S13qAb5wiBi1YLs3EOrnixSDq7Q0q/view))

### Checkpoint selection

There is one publicly released RNA-MSM pretrained checkpoint. This is that checkpoint,
converted from the original PyTorch Lightning `.ckpt` format.

## Parity Verification

Hidden-state representations verified identical (max abs diff = 0.00, exact match) to
the reference implementation at all 11 representation levels (embedding + 10 transformer
layers), both on padded and unpadded batches. Verified on GPU with PyTorch 2.7 /
CUDA 12.6.

## Related Models

See the full [RNA-MSM collection](https://huggingface.co/collections/Taykhoom/rna-msm-6a18b5c2b0181ebbc71ff777).

## Usage

RNA-MSM is an **MSA model** -- it performs best when given multiple homologous
sequences as input. For single-sequence embedding, each sequence is treated as a
1-row MSA.

### Single-sequence embedding

```python
import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)
model.eval()

sequences = ["AGCUAGCUAGCU", "GCUAGCUA"]
enc = tokenizer(sequences, return_tensors="pt", padding=True)
# enc["input_ids"]: (2, 1, seqlen)  -- each sequence treated as 1-row MSA

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

# last_hidden_state: (batch, num_alignments, seqlen, 768)
lhs = out.last_hidden_state   # (2, 1, seqlen, 768)

# Per-token embeddings for the query sequence (row 0), excluding CLS
token_emb = lhs[:, 0, 1:, :]  # (2, seqlen-1, 768)

# Mean-pool over non-padding positions for sequence-level embedding
mask = enc["attention_mask"][:, 0, 1:].unsqueeze(-1).float()  # (2, seqlen-1, 1)
seq_emb = (token_emb * mask).sum(1) / mask.sum(1).clamp(min=1)  # (2, 768)
```

### MSA embedding

```python
import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)
model.eval()

# One MSA: 3 aligned homologous sequences of equal length
msa = [
    "AGCUAGCUAGCU",
    "AGCUAGCUAGC-",
    "AGCU--CUAGCU",
]
enc = tokenizer.encode_msa([msa], return_tensors="pt", padding=True)
# enc["input_ids"]: (1, 3, seqlen)

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

# last_hidden_state: (1, 3, seqlen, 768)
# Use row 0 (query sequence) for downstream tasks
query_emb = out.last_hidden_state[:, 0, 1:, :]  # (1, seqlen-1, 768)
```

### Intermediate layers

```python
with torch.no_grad():
    out = model(**enc, output_hidden_states=True)

# hidden_states: tuple of 11 tensors, each (batch, num_alignments, seqlen, 768)
# Index 0 = embedding, 1..10 = transformer layer outputs
layer5_emb = out.hidden_states[5][:, 0, :, :]  # (batch, seqlen, 768)
```

### MLM logits

```python
from transformers import AutoModelForMaskedLM

mlm = AutoModelForMaskedLM.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)
mlm.eval()

enc = tokenizer(["AGCU<mask>AGCU"], return_tensors="pt", padding=True)
with torch.no_grad():
    logits = mlm(**enc).logits  # (1, 1, seqlen, 12)
```

### Fine-tuning

For sequence-level downstream tasks (e.g., solvent accessibility), extract the
embedding from the query row (row 0) of the last hidden state, then apply a
prediction head. The model's attention maps (row attention) are also useful for
2D structural tasks (e.g., secondary structure prediction).

## Implementation Notes

RNA-MSM uses **axial attention**: each transformer layer applies row self-attention
(attending across sequence positions, summed over alignment rows) followed by column
self-attention (attending across alignment rows per position). This custom attention
pattern is not compatible with `attn_implementation="sdpa"` or
`attn_implementation="flash_attention_2"` -- only `"eager"` is supported.

`last_hidden_state` has shape `(batch, num_alignments, seqlen, embed_dim)` -- note
the 4D output, reflecting the MSA structure. For single-sequence use (1-row MSA),
this is `(batch, 1, seqlen, embed_dim)`.

## Citation

```bibtex
@article{zhang2024_rnamsm,
  title   = {Multiple sequence alignment-based {RNA} language model and its application to structural inference},
  author  = {Zhang, Yikun and Lang, Mei and Jiang, Jiuhong and Gao, Zhiqiang and Xu, Fan and Litfin, Thomas and Chen, Ke and Singh, Jaswinder and Huang, Xiansong and Song, Guoli and Tian, Yonghong and Zhan, Jian and Chen, Jie and Zhou, Yaoqi},
  journal = {Nucleic Acids Research},
  volume  = {52},
  number  = {1},
  pages   = {e3},
  year    = {2024},
  doi     = {10.1093/nar/gkad1031}
}
```

## Credits

Original model and code by Zhang et al. Source: [GitHub](https://github.com/yikunpku/RNA-MSM).
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.