--- 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 | |---|---|---|---| | `` | 0 | `U` | 7 | | `` | 1 | `X` | 8 | | `` | 2 | `N` | 9 | | `` | 3 | `-` | 10 | | `A` | 4 | `` | 11 | | `G` | 5 | | | | `C` | 6 | | | Each sequence is prepended with `` (id 0). No `` 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(["AGCUAGCU"], 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.