HELM: Hierarchical Encoding for mRNA Language Modeling

TL;DR

HELM addresses the mismatch between natural-language pre-training and the hierarchical codon structure of mRNA by introducing a Hierarchical Cross-Entropy loss that enforces codon-level hierarchy in pre-training. The approach preserves standard architectures and tokenization (codon-level) while weighting errors according to hierarchical position, yielding consistent improvements across seven downstream properties and Ab-region annotation, plus enhanced sequence generation. Through extensive ablations on a curated 15.3M mRNA corpus (OAS) and diverse datasets, HELM demonstrates that incorporating biological priors—especially codon usage bias—improves both discriminative and generative performance by approximately 8% on average. The work suggests that hierarchical priors can be extended to other biological sequence tasks and motivates exploring hyperbolic representations to better capture hierarchical structure.

Abstract

Messenger RNA (mRNA) plays a crucial role in protein synthesis, with its codon structure directly impacting biological properties. While Language Models (LMs) have shown promise in analyzing biological sequences, existing approaches fail to account for the hierarchical nature of mRNA's codon structure. We introduce Hierarchical Encoding for mRNA Language Modeling (HELM), a novel pre-training strategy that incorporates codon-level hierarchical structure into language model training. HELM modulates the loss function based on codon synonymity, aligning the model's learning process with the biological reality of mRNA sequences. We evaluate HELM on diverse mRNA datasets and tasks, demonstrating that HELM outperforms standard language model pre-training as well as existing foundation model baselines on seven diverse downstream property prediction tasks and an antibody region annotation tasks on average by around 8%. Additionally, HELM enhances the generative capabilities of language model, producing diverse mRNA sequences that better align with the underlying true data distribution compared to non-hierarchical baselines.

Figures (7)

• Figure 1: Hierarchical encoding of mRNA sequences as a biological prior. Left: Hierarchical structure of codons for HELM and codon tokenization. The tree diagram illustrates the codon hierarchy used in the HELM approach, categorizing codons into Start, Coding (grouped by amino acids), and Stop. This hierarchy informs the loss calculation. The codon tokenizer demonstrates the process of converting an mRNA input sequence into codon tokens for modeling. Right: Codon prediction probabilities on a amino acid codon wheel. Segments represent amino acids, bars represent codons. Orange: HELM approach; Blue: cross-entropy (XE) loss. Bar height indicates probability. Non-hierarchical XE model assigns high probabilities to non-synonymous codons for masked tokens, while HELM assigns high probabilities to synonymous codons, even when errors occur.
• Figure 2: Average entropy of synonymous codon distributions correlates with HELM model performance improvement. Lower entropy, indicating stronger codon bias, is observed in datasets like MLOS, Tc-Riboswitches, and mRFP, where HELM shows greater improvements over XE models.
• Figure 3: Relationship between Codon Pair Bias (CPB) and % improvement over XE. Datasets with more negative CPB values (indicating stronger codon usage bias) tend to exhibit greater improvements with HELM.
• Figure 4: FBD comparison of generative HELM, XE, and random models across varying temperature. Higher temperatures increase diversity but worsen FBD scores. HELM consistently beats XE, suggesting better alignment with real mRNA data while maintaining diversity.
• Figure 5: Percentage MSE reduction: HELM vs XE models. MSE compares the predicted properties of generated sequences to the predicted properties of original sequences. HELM consistently outperforms non-hierarchical models, indicating better retention of key mRNA properties in generated sequences.