HELM: Hierarchical Encoding for mRNA Language Modeling

Authors: Mehdi Yazdani-Jahromi, Mangal Prakash, Tommaso Mansi, Artem Moskalev, Rui Liao
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: Messenger RNA (mRNA), Codon structure, Hierarchical modeling, Bio-language model, Property prediction, mRNA sequence generation

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

Decision: accept (poster)
Reviewer scores: 6, 5, 6, 6

Reviews and discussion:
<thread>
creation date:
1730641713883

thread_type:
official_review

<thread_object>
creator: reviewer_sjyv

<content>
summary:
The paper introduces a pre-training strategy, **Hierarchical Encoding for mRNA Language Modelling (HELM)**, that incorporates the hierarchical structure of mRNA codons into language model training. By modulating the model's loss based on codon synonymity, HELM aligns learning with the biological structure of mRNA. Results show HELM outperforms non-hierarchical models by around 8% on various mRNA tasks, including antibody annotation and property prediction, while also enhancing generative capabilities for mRNA sequences.

strengths:
- HELM integrates codon hierarchy into model training, capturing the inherent structure of mRNA data.
- Demonstrates a good performance improvement over existing models on multiple downstream tasks.
- The model can generate more biologically plausible and diverse mRNA sequences than existing approaches.

weaknesses:
- The authors curate a dataset from OAS specifically to train their models. This may explain much of the performance gains observed in Table-1, as the model is trained and tested on data with similar statistical characteristics. An ablation study isolating the effect of the curated dataset would clarify the extent to which performance gains are due to the data source versus tokenization and architecture design.
- As the authors noted it, representing hierarchical relationships in Euclidean space might limit HELM’s ability to capture the full complexity of these structures, potentially affecting performance.
- Although the results are encouraging, they are primarily based on antibody and property prediction tasks, which could restrict the model's applicability to other mRNA-related contexts.
- Additionally, the strategy for splitting the data into training, validation, and test sets is crucial and should be detailed in the main paper. In sequence-based datasets, it is common practice to split data based on clusters after clustering sequences, which minimizes data leakage. The reported performance gains may be overstated if there is any leakage between the splits.


I am open to revisiting my score; however, I will wait for: (i) the authors' response to my comments, and (ii) feedback from other reviewers.

questions:
- Are you planning to open-source the curated dataset ?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730644470907

thread_type:
official_review

<thread_object>
creator: reviewer_erze

<content>
summary:
The paper introduces HELM, a hierarchical encoding approach tailored for mRNA language modeling that leverages the codon-level structure of mRNA sequences. Unlike traditional models, HELM incorporates a hierarchical cross-entropy loss function to align with the biological structure of mRNA, particularly focusing on synonymous codon usage and its functional implications. The model is evaluated on multiple mRNA datasets, demonstrating improvements in downstream prediction tasks, generative diversity, and sequence annotation accuracy over non-hierarchical baselines. Overall, HELM represents a biologically-informed advancement in mRNA modeling, showing enhanced performance on tasks relevant to protein synthesis and gene expression.

strengths:
1. Biological Prior Integration: The hierarchical encoding strategy effectively incorporates biological knowledge of mRNA structure, particularly the synonymous codon usage, which enhances the model’s performance on property prediction and generative tasks.
2. Diverse Evaluations: HELM's performance is thoroughly evaluated across multiple tasks, including property prediction, generative sequence design, and antibody sequence region annotation, showcasing the model’s versatility and relevance to bioinformatics applications.
3. Improved Representational Quality: By aligning model training with the codon hierarchy, HELM captures the underlying biological structure more effectively, achieving better clustering of synonymous sequences and improved predictive accuracy in key bioinformatics tasks.

weaknesses:
1. Lack of Scaling Experiments: The paper primarily employs models with 50M parameters, which is relatively small compared to large language models (LLMs) in NLP. This limitation raises concerns about the necessity of HELM’s hierarchical encoding, as larger models might naturally learn these hierarchical relationships without explicit design. Experiments with larger models could help clarify if the hierarchical loss function offers unique advantages or if it becomes redundant with increased scale.
2. Limited technical contribution:  The primary contribution of this work lies in adapting hierarchical cross-entropy (HXE) to the mRNA structure, which, while valuable for bioinformatics, represents a limited advancement from a machine learning perspective. The adaptation of HXE to this domain may not provide significant innovation to the broader ML community.
3. Baseline suggestion: A simpler way to capture the hierarchical structure could be to introduce separate embeddings for different types of codons and combine these with the token embeddings as input. For instance, embeddings such as [start], [stop], [S], and [Q] (as shown in Figure 6) could be used to represent codon types without modifying the loss function. This approach would provide a straightforward baseline for evaluating the benefits of the proposed HXE solution.

questions:
See W3.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730686930813

thread_type:
official_review

<thread_object>
creator: reviewer_9tth

<content>
summary:
The paper introduces Hierarchical Encoding for mRNA Language Modeling (HELM), a novel pretraining strategy that incorporates the hierarchical codon structure of mRNA into language model training. It addresses the limitations of existing models that overlook codon synonymity, which can lead to suboptimal performance in mRNA tasks. HELM modifies loss functions to prioritize errors between different amino acids over synonymous codons, enhancing the model's alignment with biological realities. Evaluations show that HELM outperforms standard language model pretraining and other state-of-the-art RNA models by approximately 8% across six diverse downstream tasks, including property prediction and antibody region annotation, while using fewer model parameters. Additionally, HELM demonstrates improved generative capabilities, producing mRNA sequences that better align with true data distributions. Overall, HELM effectively captures the hierarchical nature of mRNA, leading to enhanced performance in both predictive and generative tasks.

strengths:
1. This study proposed a novel approach for embedding biological hierarchical structures into language models to enhance interpretability and accuracy in biological sequence analysis.
2. A comprehensive technical evaluation spanning multiple model architectures was presented, benchmarked across diverse datasets and use cases.
3. Clear and precise presentation of technical methods and experimental setup were provided to facilitate reproducibility and further research.

weaknesses:
1. The performance improvement seems more attributable to data selection than methodology: - HELM uses antibody mRNA while baselines use different data types (ncRNA, pre-mRNA, diverse organisms mRNA).
2.  The paper fails to justify why antibody mRNA pre-training would generalize to: - Viral RNA sequences - Riboswitch sequences.
3.  The evaluation methodology is insufficient: a. Over-reliance on single metric (Spearman correlation) b. No analysis across sequence lengths c. No biological interpretation of prediction errors d. The evaluation lacks comprehensive stress testing on edge cases (e.g., extreme sequence lengths, unusual GC contents, rare codon clusters), raising concerns about the model's reliability in challenging real-world scenarios.
4. Methodological issues: (1) α parameter (0.2-0.6) lacks biological basis (2) Oversimplified codon bias analysis (3) Insufficient validation of hierarchical structure's biological relevance.
5. Limited practical applications: (1) Too narrow focus on antibody sequences and expression (2) Critical applications remain untested: a. RNA vaccine design applications b. Abnormal or mutated sequences c. Real-world biological scenarios.

questions:
1. What is the biological justification for the α parameter range selection?
2. How does the model handle sequences with unusual codon usage patterns? Such evaluation is crucial because unusual codon usage patterns could significantly impact translation efficiency and protein expression levels, yet these scenarios are not addressed in the current validation.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730837454146

thread_type:
official_review

<thread_object>
creator: reviewer_tlr1

<content>
summary:
The paper proposes an elegant and novel solution for incorporating codon hierarchical structure into language model training, which is hierarchical cross-entropy loss. Although the method is simple, the biological insight is deep and the performance is convincing. The authors also curate the hierarchical information and a pre-train dataset, which may benefit future research. Overall, I feel this would be an impactful work for the community.

strengths:
1. The authors combine biological insight with language modeling.
2. The proposed HELM is solid and performs well for mRNA-related tasks. It achieves SOTA performance on several practical and impactful tasks, e.g., mRNA sequence design, and sequence region annotation.
3. The authors provide a curated dataset, and curated domain knowledge, which may benefit future computation-oriented research.
4. The authors provide solid benchmark experiments of various backbone architectures and tokenization methods.

weaknesses:
1. The methodology itself is relatively simple.
2. It is not clear to what extent the model, the codes, and the data will be made public.

questions:
I do have a strong concern for the reproducibility of this paper. Despite the fact that I think highly of this paper, I don't think the paper should be accepted without further declaration on reproducibility. Specifically, will the model checkpoint, source codes, and datasets be released?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732530372346

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Comment on Revisions and Additions

comment:
We thank the reviewers for their valuable feedback and constructive suggestions, which have enhanced the quality and clarity of our submission. We are glad that the reviewers appreciate our elegant and novel solution  (***TLr1, 9ttH***) for incorporating codon hierarchy into mRNA language model, thus representing biologically-informed advancement in mRNA modeling (***eRzE***) while outperforming baselines by around 8% on various mRNA tasks (***TLr1, 9ttH, eRzE, SJYV***). 

In terms of main contributions,
- We introduce HELM, a novel yet simple framework for mRNA language model pre-training, aligning learning with mRNA hierarchical structure (***TLr1, 9ttH***).
- We introduce a high-quality mRNA dataset (***9ttH, SJYV***) and evaluate diverse tokenization, pre-training, and LM architectures.
- HELM improves downstream performance by ~8% on diverse mRNA tasks, including property prediction, sequence generation, and region annotation (***TLr1, 9ttH, eRzE, SJYV***).

**Revisions during rebuttal**

Below, we summarize key changes made to the paper and additional experiments conducted in response to the reviewers’ comments:

## Declaration of Reproducibility
We have added a formal declaration of reproducibility in main text (page 11), explicitly committing to open-sourcing all code, model weights, and datasets after acceptance. Further, we provided implementation details for training and inference in Appendix Sec. A.2 and A.3 to ensure clarity and transparency.

## Inclusion of Additional Metrics
In addition to Spearman correlation, we now report Pearson correlation and R² values in Appendix Table 9. These metrics confirm the robustness of HELM’s performance across tasks and provide a more nuanced evaluation of the results.

## Performance Analysis on Sequence Lengths and GC Content
We conducted detailed experiments to evaluate HELM’s performance across varying sequence lengths and GC content ranges, utilizing the iCodon mRNA thermostability prediction dataset. The results, presented in Appendix A.10, demonstrate that HELM consistently outperforms baselines across these diverse tasks.

## Scaling Experiments
To explore scalability, we conducted additional experiments with 100M parameter models. The results, included in Appendix A.7, reveal that HELM’s hierarchical pre-training consistently outperforms non-hierarchical baselines even at larger scales, reinforcing utility of explicitly incorporating hierarchical priors.

## Codon Pair Bias (CPB) Analysis
We have now included an established Codon Pair Bias (CPB) analysis, which measures frequency of specific codon pairs in an organism's genome compared to their expected frequencies. This analysis is detailed in Appendix A.5.2 (lines 991-1004 on page 19) and visually presented in Appendix Fig. 7 (page 20, highlighted in blue). This analysis reveals that HELM shows greater improvements on datasets with stronger codon bias. We have also referred to this analysis in the main text (lines 419-421, highlighted in blue).

## Alternative Baseline Experiments
We implemented a reviewer-suggested baseline by introducing separate embeddings for codon types combined with token embeddings. The results, included in Appendix A.11, show that HELM still outperforms this alternative baseline across all datasets, underscoring the advantages of hierarchical pretraining.

## New Dataset and Expanded Vaccine Relevance
We incorporated a new public COVID-19 mRNA Vaccine Degradation Prediction dataset to further evaluate HELM’s relevance for RNA vaccine design. Results on this dataset, presented in Tables 1 and 2, demonstrate HELM’s effectiveness in this critical application area.

## Improved Biological Justifications
We have expanded the discussion on HELM’s biological relevance, including:
- Generalizability to diverse mRNA types due to shared structural and functional motifs.
- Alignment of prediction errors with synonymous codons for better biological interpretability.
- The role of the α parameter in the hierarchical loss, allowing flexible weighting of hierarchical levels.

## Clarification on Data Splitting
To prevent data leakage, we detailed clustering-based partitioning methodologies (e.g., LinClust with a similarity threshold of 0.9) for datasets without predefined splits. These details are now clarified in the main text.

## Quantitative Analysis of mRNA Structural Motif Conservation Between Antibody and Non-Antibody mRNA Sequences to Demonstrate Generalizability of Pre-Trained Models
We conducted a quantitative analysis of RNA secondary structure conservation between antibody and non-antibody mRNA datasets using motif-based similarity metrics, demonstrating meaningful structural overlap that supports the generalizability of our model’s learned features across diverse mRNA types.

These revisions and additions address the reviewers’ comments comprehensively, strengthening the contributions of our work. We thank the reviewers for their insights, which have improved our submission.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents a pretraining strategy that incorporates codon-level hierarchical structure into language model training. The reviewers found the paper's methodology to be a useful combination of biological insights with language modeling. While the technical novelty from the ML point of view is limited, the method's significance to bioinformatics and the solid experimental results make up for this concern. There were some concerns about reproducibility and the evaluation, but these were addressed during the rebuttal period. Given all this, I am happy to recommend acceptance. Please incorporate all the reviewer comments in the final version, and make sure to release the model, code, and data.