ERNIE-RNA-SS

ERNIE-RNA fine-tuned on RNA secondary structure (SS) prediction. The task-specific structure prediction head has been discarded; the fine-tuned encoder and original pretrained MLM head are included.

Architecture

Parameter	Value
Layers	12
Attention heads	12
Embedding dimension	768
FFN hidden dimension	3072 (GELU)
Vocabulary size	25
Positional encoding	Sinusoidal (fairseq-style)
Normalization	LayerNorm (embedding and post-residual)
Architecture	Post-LN Transformer with recurrent 2D RNA pairing bias
Max sequence length	1024

See Taykhoom/ERNIE-RNA for the vocabulary table and full architecture description.

Pretraining + Fine-tuning

• Pretraining objective: Masked language modeling on RNAcentral
• Fine-tuning task: RNA secondary structure prediction (base-pair prediction)
• Fine-tuning data: RNA3DB
• Source checkpoint: RNA3DB.pt
• Included heads: Pretrained MLM head only; the SS prediction head is omitted

Checkpoint selection

Six SS fine-tuned checkpoints were available (bpRNA-1m, bpRNA-new, RIVAS, RNA3DB, RNAStralign, bpRNA-1m_RNAStralign). bpRNA-new was excluded because its backbone is identical to pretrained ERNIE-RNA. RNA3DB was the strongest general-purpose encoder among the remaining checkpoints and was selected for this release.

Parity Verification

Backbone and MLM head weights are extracted directly from the fine-tuned checkpoint. Every MLM tensor is verified bit-exact against the source checkpoint; the retained MLM transform, layer norm, and output bias are unchanged from pretrained ERNIE-RNA.

Only attn_implementation="eager" is supported (see Implementation Notes).

Related Models

See the full ERNIE-RNA collection.

Model	Notes
Taykhoom/ERNIE-RNA	Pretrained model
Taykhoom/ERNIE-RNA-SS	This model — SS fine-tuned
Taykhoom/ERNIE-RNA-MRL	UTR MRL fine-tuned

Usage

Embedding generation

import torch
from transformers import AutoTokenizer, AutoModel

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

sequences = ["AUGCAUGCAUGC", "GGGGCCCCGGGG"]
enc = tokenizer(sequences, return_tensors="pt", padding=True)

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

cls_emb   = out.last_hidden_state[:, 0, :]   # (batch, 768) -- CLS token
token_emb = out.last_hidden_state             # (batch, seq_len, 768)

# Intermediate layers
out_all = model(**enc, output_hidden_states=True)
layer6_emb = out_all.hidden_states[6]         # (batch, seq_len, 768)

MLM logits

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/ERNIE-RNA-SS", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained(
    "Taykhoom/ERNIE-RNA-SS", trust_remote_code=True
)
model.eval()

enc = tokenizer(["AUGC<mask>UGCA"], return_tensors="pt")
with torch.no_grad():
    logits = model(**enc).logits   # (1, seq_len, 25)

Fine-tuning

Use the CLS token embedding (last_hidden_state[:, 0, :]) as input to a prediction head for sequence-level tasks. For token-level tasks (e.g. base-pair prediction), use last_hidden_state directly. The original SS prediction head is not included.

Implementation Notes

ERNIE-RNA's recurrent 2D bias is updated from the pre-softmax attention scores at every layer (the raw QK logits become the bias input for the next layer). Fused attention kernels (SDPA, FlashAttention) do not expose pre-softmax scores, so they cannot maintain this recurrent pathway. Only attn_implementation="eager" is supported; requesting sdpa or flash_attention_2 raises a ValueError.

The twod_proj MLP is always run in float32 (matching the original) regardless of the model's compute dtype.

Citation

@article{yin2025_ernierna,
  title   = {{ERNIE-RNA}: an {RNA} language model with structure-enhanced representations},
  author  = {Yin, Weijie and Zhang, Zhaoyu and Zhang, Shuo and He, Liang and Zhang, Ruiyang and Jiang, Rui and Liu, Gan and Wang, Jingyi and Zhang, Xuegong and Qin, Tao and Xie, Zhen},
  journal = {Nature Communications},
  volume  = {16},
  number  = {1},
  pages   = {8407},
  year    = {2025},
  doi     = {10.1038/s41467-025-64972-0}
}

Credits

Original model and code by Yin et al. Source: GitHub. Hugging Face port maintained by Taykhoom Dalal.

License

MIT, following the original repository.