ERNIE-RNA-MRL

ERNIE-RNA fine-tuned on UTR mean ribosome load (MRL) prediction. The task-specific CNN regression 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: UTR mean ribosome load (MRL) prediction
• Source checkpoint: ERNIE-RNA-UTR_ML_CNN.pt
• Included heads: Pretrained MLM head only; the MRL regression head is omitted

Checkpoint selection

Single MRL fine-tuned checkpoint from the original repository. The original model uses a CNN head on top of the ERNIE-RNA encoder. This port keeps the fine-tuned encoder and the MLM head retained in the source checkpoint, while omitting the CNN head.

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	SS fine-tuned
Taykhoom/ERNIE-RNA-MRL	This model -- UTR MRL fine-tuned

Usage

Embedding generation

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("Taykhoom/ERNIE-RNA-MRL", trust_remote_code=True)
model = AutoModel.from_pretrained("Taykhoom/ERNIE-RNA-MRL", 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-MRL", trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained(
    "Taykhoom/ERNIE-RNA-MRL", 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. The original MRL regression 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.