RNA-MSM

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

Architecture

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

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)

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

Related Models

See the full RNA-MSM collection.

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

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

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

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

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

@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. Hugging Face port maintained by Taykhoom Dalal.

License

MIT, following the original repository.