Taykhoom/RNA-MSM

1---2library_name: transformers3tags:4- RNA5- language-model6- MSA7license: mit8---9 10# RNA-MSM11 12Multiple sequence alignment-based RNA language model trained on homologous RNA13sequence alignments from the RNAcmap pipeline.14 15## Architecture16 17| Parameter | Value |18|---|---|19| Layers | 10 |20| Attention heads | 12 |21| Embedding dimension | 768 |22| FFN hidden dimension | 3072 (GELU) |23| Vocabulary size | 12 |24| Positional encoding | Learned (sequence) + learned scalar (alignment row) |25| Normalization | Pre-LayerNorm + final LayerNorm |26| Architecture | Axial MSA Transformer (row + column self-attention) |27| Max sequence length | 1024 tokens (including the prepended `<cls>`) |28| Max alignment depth | 1024 rows |29 30**Input format:** RNA-MSM takes 3D input `(batch, num_alignments, seqlen)`. Each31alignment is a set of homologous RNA sequences of equal length (an MSA). The model32applies row self-attention (across sequence positions) and column self-attention33(across alignment rows) at each of the 10 transformer layers.34 35### Vocabulary36 37| Token | ID | Token | ID |38|---|---|---|---|39| `<cls>` | 0 | `U` | 7 |40| `<pad>` | 1 | `X` | 8 |41| `<eos>` | 2 | `N` | 9 |42| `<unk>` | 3 | `-` | 10 |43| `A` | 4 | `<mask>` | 11 |44| `G` | 5 | | |45| `C` | 6 | | |46 47Each sequence is prepended with `<cls>` (id 0). No `<eos>` token is appended.48 49## Pretraining50 51- **Objective:** Masked language modeling on RNA MSAs (masking ~15% of tokens)52- **Data:** RNA homologous sequences searched by RNAcmap from non-redundant RNA53  databases54- **Source checkpoint:** `RNA_MSM_pretrained.ckpt`55  ([original Google Drive link](https://drive.google.com/file/d/11A-S13qAb5wiBi1YLs3EOrnixSDq7Q0q/view))56 57### Checkpoint selection58 59There is one publicly released RNA-MSM pretrained checkpoint. This is that checkpoint,60converted from the original PyTorch Lightning `.ckpt` format.61 62## Parity Verification63 64Hidden-state representations verified identical (max abs diff = 0.00, exact match) to65the reference implementation at all 11 representation levels (embedding + 10 transformer66layers), both on padded and unpadded batches. Verified on GPU with PyTorch 2.7 /67CUDA 12.9.68 69## Related Models70 71See the full [RNA-MSM collection](https://huggingface.co/collections/Taykhoom/rna-msm-6a18b5c2b0181ebbc71ff777).72 73## Usage74 75RNA-MSM is an **MSA model** -- it performs best when given multiple homologous76sequences as input. For single-sequence embedding, each sequence is treated as a771-row MSA.78 79### Single-sequence embedding80 81```python82import torch83from transformers import AutoTokenizer, AutoModel84 85tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)86model = AutoModel.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)87model.eval()88 89sequences = ["AGCUAGCUAGCU", "GCUAGCUA"]90enc = tokenizer(sequences, return_tensors="pt", padding=True)91# enc["input_ids"]: (2, 1, seqlen)  -- each sequence treated as 1-row MSA92 93with torch.no_grad():94    out = model(**enc)95 96# last_hidden_state: (batch, num_alignments, seqlen, 768)97lhs = out.last_hidden_state   # (2, 1, seqlen, 768)98 99# Per-token embeddings for the query sequence (row 0), excluding CLS100token_emb = lhs[:, 0, 1:, :]  # (2, seqlen-1, 768)101 102# Mean-pool over non-padding positions for sequence-level embedding103mask = enc["attention_mask"][:, 0, 1:].unsqueeze(-1).float()  # (2, seqlen-1, 1)104seq_emb = (token_emb * mask).sum(1) / mask.sum(1).clamp(min=1)  # (2, 768)105```106 107### MSA embedding108 109```python110import torch111from transformers import AutoTokenizer, AutoModel112 113tokenizer = AutoTokenizer.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)114model = AutoModel.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)115model.eval()116 117# One MSA: 3 aligned homologous sequences of equal length118msa = [119    "AGCUAGCUAGCU",120    "AGCUAGCUAGC-",121    "AGCU--CUAGCU",122]123enc = tokenizer.encode_msa([msa], return_tensors="pt", padding=True)124# enc["input_ids"]: (1, 3, seqlen)125 126with torch.no_grad():127    out = model(**enc)128 129# last_hidden_state: (1, 3, seqlen, 768)130# Use row 0 (query sequence) for downstream tasks131query_emb = out.last_hidden_state[:, 0, 1:, :]  # (1, seqlen-1, 768)132```133 134### Intermediate layers135 136```python137with torch.no_grad():138    out = model(**enc, output_hidden_states=True)139 140# hidden_states: tuple of 11 tensors, each (batch, num_alignments, seqlen, 768)141# Index 0 = embedding, 1..10 = transformer layer outputs142layer5_emb = out.hidden_states[5][:, 0, :, :]  # (batch, seqlen, 768)143```144 145### MLM logits146 147```python148from transformers import AutoModelForMaskedLM149 150mlm = AutoModelForMaskedLM.from_pretrained("Taykhoom/RNA-MSM", trust_remote_code=True)151mlm.eval()152 153enc = tokenizer(["AGCU<mask>AGCU"], return_tensors="pt", padding=True)154with torch.no_grad():155    logits = mlm(**enc).logits  # (1, 1, seqlen, 12)156```157 158### Fine-tuning159 160For sequence-level downstream tasks (e.g., solvent accessibility), extract the161embedding from the query row (row 0) of the last hidden state, then apply a162prediction head. The model's attention maps (row attention) are also useful for1632D structural tasks (e.g., secondary structure prediction).164 165## Implementation Notes166 167RNA-MSM uses **axial attention**: each transformer layer applies row self-attention168(attending across sequence positions, summed over alignment rows) followed by column169self-attention (attending across alignment rows per position). This custom attention170pattern is not compatible with `attn_implementation="sdpa"` or171`attn_implementation="flash_attention_2"` -- only `"eager"` is supported.172 173`last_hidden_state` has shape `(batch, num_alignments, seqlen, embed_dim)` -- note174the 4D output, reflecting the MSA structure. For single-sequence use (1-row MSA),175this is `(batch, 1, seqlen, embed_dim)`.176 177## Citation178 179```bibtex180@article{zhang2024_rnamsm,181  title   = {Multiple sequence alignment-based {RNA} language model and its application to structural inference},182  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},183  journal = {Nucleic Acids Research},184  volume  = {52},185  number  = {1},186  pages   = {e3},187  year    = {2024},188  doi     = {10.1093/nar/gkad1031}189}190```191 192## Credits193 194Original model and code by Zhang et al. Source: [GitHub](https://github.com/yikunpku/RNA-MSM).195Hugging Face port maintained by Taykhoom Dalal.196 197## License198 199MIT, following the original repository.200