HELM: Hierarchical Encoding for mRNA Language Modeling

## Abstract

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.

## 1 Introduction

RNA analysis is becoming increasingly important in molecular biology (Liu et al., ; Fu, ). Messenger RNA (mRNA) is of particular interest due to its unique role in protein synthesis  (Sahin et al., ). mRNA consists of triplets of nucleotides, called codons, which directly translate to protein amino acids through a surjective many-to-one mapping (Figure 1 left). This results in multiple synonymous mRNAs that encode identical amino acid sequences while exhibiting distinct physical and biological properties at the nucleotide level (Buhr et al., ). The role of mRNA, serving as an intermediary between DNA and proteins, underlies its therapeutic potential in applications such as vaccines and gene therapies (Caplen, ; Pardi et al., ), while also presenting challenges for analysis and optimization (Liu et al., ).

Language Models (LMs) have emerged as powerful tools for analyzing biological sequences, with notable successes in protein (Elnaggar et al., ; Ferruz et al., ; Lin et al., ; Hie et al., ) and DNA (Nguyen et al., ; Zhou et al., ) research. Despite the importance of mRNA, the field still lacks specialized LMs tailored for its analysis. Existing RNA LMs (Li et al., ; Chen et al., ) focus on non-coding sequences and do not account properly for codon hierarchy (Fig. 1 right) which, as we demonstrate, falls short when dealing with mRNA tasks. In this work, we aim to address this gap in mRNA language modeling by focusing specifically on the unique challenges presented by mRNA sequences.

Figure: 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.
Refer to caption: extracted/5931372/figures/teaser_new.png

To address the limitations of existing bio-language modeling methods, we introduce Hierarchical Encoding for mRNA Language Modeling (HELM), a novel pre-training strategy for mRNA sequences. Rather than directly adopting natural-language pre-training methods such as Masked Language Modeling (MLM) (Devlin et al., ; Liu et al., ; Cheng et al., ) or Causal Language Modeling (CLM) (Radford, ; Radford et al., ; Brown, ), HELM begins by recognizing and formalizing the hierarchical structure of mRNA sequences based on codon synonymity. HELM modulates the CLM or MLM loss based on a codon’s position in this hierarchy, treating errors between synonymous codons as less significant than those resulting in different amino acids. This approach requires minimal modifications to standard pipelines while aligning the model training process with the intrinsic hierarchical structure of mRNA data (Figure 1 right). In addition, to facilitate further development of mRNA LMs, we provide a consistent comparison of various tokenization methods, hierarchical and standard pre-training strategies, and different language model architectures, including recent Mamba and Hyena models (Gu & Dao, ; Poli et al., ).

To demonstrate the practical advantages of HELM, we focus on the important domain of antibody mRNAs where codon distribution can significantly impact physical and therapeutic properties  (Zhang et al., ). We conduct comprehensive experiments, evaluating our hierarchical model against its non-hierarchical counterparts pre-trained with natural-language CLM and MLM as well as against other state-of-the-art RNA foundation models. Our results show that hierarchical pre-training yields significantly better representations of antibody-encoding mRNA and even excels on non-antibody mRNA data. Our hierarchy-aware model outperforms non-hierarchical counterparts pre-trained with natural-language CLM and MLM as well as other state-of-the-art RNA foundation models by around 8% on average across six diverse downstream property prediction tasks and an antibody region annotation tasks while requiring up to 2 times fewer model parameters. Beyond its excellent discriminative performance, we demonstrate that HELM-based CLM pre-training yields strong generative models capable of producing diverse mRNA sequences while preserving closer alignment with the true data distribution. Ultimately, HELM addresses the limitations of current bio-LMs by effectively capturing the hierarchical structure of mRNA as a strong biological prior, leading to improved performance in both predictive and generative tasks while requiring minimal modification of the pre-training.

To sum up, we make the following contributions:

- •
We highlight the importance of biological hierarchy when modeling mRNA sequences and demonstrate that standard natural-language pre-training methods do not account for this inherent biological structure, leading to suboptimal performance in mRNA tasks.
- •
We propose Hierarchical Encoding for mRNA Language Modeling (HELM), a novel pre-training strategy that incorporates mRNA hierarchical structure into language model pre-training. HELM modulates the CLM or MLM loss based on codon hierarchy, aligning the model’s learning process with the biological structure of mRNA sequences.
- •
We provide a consistent comparison of various tokenization methods, hierarchical and standard pre-training strategies, and different LM architectures to facilitate further development of mRNA models.
- •
We demonstrate the advantages of representation learning with HELM where our method yields substantial downstream performance improvement on multiple diverse mRNA datasets spanning property prediction, automated sequence annotation and sequence generation tasks.

## 2 Related works

##### Pre-training bio-language models

Recent advancements in LMs have significantly improved the analysis of biological sequences such as DNA, RNA, and proteins (Laurent et al., ). The standard approach for bio-language modeling adopts pre-training strategies from natural language; state-of-the-art protein or DNA LMs rely on Masked Language Modeling (Liu et al., ; Lin et al., ; Dalla-Torre et al., ) or Causal Language Modeling (Brown, ; Lin et al., ) and pre-train representations on large bio-databases (Bairoch et al., ; Suzek et al., ; Consortium et al., ; Schoch et al., ). In addition, the modeling choices of pre-training bio-LMs involve network architecture and tokenization strategy. Regarding architectures, transformer-based models have been widely adopted (Elnaggar et al., ; Lin et al., ), while recent works explore structured state-space (Peng et al., ; Sgarbossa et al., ) and long-convolutional models (Nguyen et al., ; ; Moskalev et al., ) to better handle long-range dependencies in biological sequences. Tokenization strategies vary from character-level approaches (Elnaggar et al., ; Lin et al., ) to k-mer tokenization (Dalla-Torre et al., ) and Byte Pair Encoding  (Zhou et al., ). These existing pre-training and modeling choices primarily adapt natural language paradigms to biological sequences, overlooking the incorporation of inherent biological structures crucial for mRNA analysis. Moreover, the lack of consistent comparisons between these various approaches hinders the development of effective mRNA LMs. Our work addresses these limitations with HELM, a pre-training strategy incorporating mRNA’s hierarchy into the training, and we also provide analysis of various pre-training and modeling choices to facilitate further development of mRNA LMs.

##### RNA language models

Unlike the many protein and DNA LMs (Lin et al., ; Elnaggar et al., ; Dalla-Torre et al., ; Nguyen et al., ), there is a noticeable lack of models specifically designed for mRNA due to the smaller, specialized datasets available. Existing RNA models focus on specific RNA types or regions, such as RNA-FM (Chen et al., ) and RINALMO (Penić et al., ) for non-coding RNAs, UTR-LM (Chu et al., ) and UTRBERT (Yang et al., ) for untranslated regions, and SpliceBERT (Chen et al., ) for precursor mRNA. While these models are valuable, they are not tailored for protein-coding mRNAs. Proprietary models such as CodonBERT (Li et al., ) have been developed but they do not utilize the natural hierarchy of mRNA sequences. Our work addresses these limitations by introducing HELM to explicitly incorporate the biological knowledge of mRNA structure into the pre-training process of a LM.

##### Modeling hierarchical data

Hierarchical relationships are important across various domains, and several methods have been proposed to model them effectively. In computer vision, approaches range from embedding-based methods like DeViSE (Frome et al., ), which maximizes cosine similarity between image and label embeddings, to geometric approaches that rely on hyperbolic spaces (Liu et al., ; Atigh et al., ) or hyperspheres with structural constraints (Barz & Denzler, ; Bengio et al., ). Hierarchical loss functions, such as hierarchical cross-entropy (HXE) (Bertinetto et al., ; Garg et al., ), have been applied to tasks like image classification and astrophysical transient classification (Villar et al., ). In natural language processing, hierarchical structures have been modeled both implicitly through contrastive learning (Liu et al., ; Gosselin & Zouaq, ) and explicitly using hyperbolic losses (He et al., ) or hierarchical softmax (Schuurmans & Frasincar, ). While these methods show great promise in their respective domains, they are not tuned for biological data. Our HELM approach addresses this gap by extending the hierarchical learning to mRNA sequences, effectively capturing the codon hierarchy in the pre-training process. This demonstrates the potential of incorporating biological priors into language modeling pre-training, paving the way for more accurate and biologically relevant models in this domain.

## 3 Hierarchical Cross-Entropy Loss

Hierarchical cross-entropy (HXE) is a technique for training neural networks that incorporates hierarchical information into the learning process. HXE was originally proposed in computer vision for image classification (Bertinetto et al., ) tasks, where the data labels can be organized in a hierarchical structure, such as a tree. When the hierarchy $H$ is structured as a tree, HXE allows for a unique factorization of the categorical distribution $p(C)$ over the classes in terms of the conditional probabilities along the path from each class to the root of the tree. Specifically, for any leaf node/class $C$ with a path $C^{(0)}=C,\ldots,C^{(h)}=R$, the probability of $C$ can be written as
$p(C)=\prod_{l=0}^{h-1}p(C^{(l)}\mid C^{(l+1)}),$
where $h$ is the height of the node $C$.
The conditional probabilities can be expressed in terms of the class probabilities as:

$$ $p(C^{(l)}\mid C^{(l+1)})=\frac{\sum_{A\in\text{Leaves}(C^{(l)})}p(A)}{\sum_{B \in\text{Leaves}(C^{(l+1)})}p(B)},$ (1) $$

where $\text{Leaves}(C)$ denotes the set of leaf nodes within the subtree rooted at $C$.
To incorporate hierarchical information directly into the loss function, the classifier’s output can be factorized according to the hierarchical structure, and the total loss can be defined as a reweighted sum of the cross-entropies of these conditional probabilities along the hierarchical path. This leads to the hierarchical cross-entropy (HXE) loss, effectively incorporating the hierarchical information into the learning process:

$$ $L_{HXE}(p,C)=-\sum_{l=0}^{h-1}\lambda(C^{(l)})\log p(C^{(l)}\mid C^{(l+1)}),$ (2) $$

where $\lambda(C^{(l)})$ is the weight associated with the edge between nodes $C^{(l+1)}$ and $C^{(l)}$.

##### Formalizing the hierarchcy for mRNA sequences

We present a novel formalism for constructing a multi-level hierarchical tree $H=(V,E)$ that captures the structured relationships within mRNA sequences at the codon level (see Fig. 1 left, Appendix Fig. 6) and well-rooted in biology (Clancy & Brown, ). The root node $R$ of this tree represents a node corresponding to all codons. This root node has two immediate child nodes: coding codons $C_{\text{coding}}$ and non-coding codons $C_{\text{non-coding}}$.
The non-coding codons $C_{\text{non-coding}}$ further branch into two child nodes: the start codon node $C_{\text{start}}$ and the stop codon node $C_{\text{stop}}$. The start codon node $C_{\text{start}}$ has a single leaf node corresponding to the codon AUG. The stop codon node $C_{\text{stop}}$ has three leaf nodes corresponding to the codons UAA, UAG, and UGA.
On the other side of the hierarchy, the coding codons $C_{\text{coding}}$ are organized into multiple child nodes, each corresponding to a specific amino acid $A_{j}$. Each amino acid node $A_{j}$ has a number of child nodes, each corresponding to a synonymous codon that encodes for $A_{j}$. Let $\text{Codons}(A_{j})=\{C_{i_{1}},C_{i_{2}},\ldots,C_{i_{m}}\}$ denote the set of all synonymous codons for the amino acid $A_{j}$. These synonymous codons preserve the same amino acid, allowing for degeneracy in the genetic code.
The complete hierarchical structure is thus represented by the tree $H=(V,E)$, where $V$ includes the root node, all internal nodes (corresponding to coding and non-coding codons, start/stop codons, and amino acids), and all leaf nodes (corresponding to individual codons). The edges $E$ define the hierarchical relationships between these nodes.

##### Hierarchical MLM and CLM pre-training for mRNA

We adapt the Hierarchical Cross-Entropy to incorporate the biological structure of mRNA into Masked Language Modeling and Causal Language Modeling pre-training strategies. For the hierarchical MLM, we randomly mask a portion of codons in the input sequence and train a model to predict the masked codons supervising it by the HXE loss. In the hierarchical CLM, a model predicts the next codon given the previous ones with the HXE loss function.

The hierarchical structure of mRNA informs the HXE objective, allowing the model’s output token probabilities to be factorized according to the tree structure. Then, the HXE loss uses Eq.1 and Eq.2 weighting the MLM or CLM errors differently based on the codon’s position in the hierarchy. We employ a practical weighting function $\lambda(C)=\exp(-\alpha h(C))$ where $h(C)$ is the height of the node $C$ in the hierarchy and ($\alpha>0$). The value of $\alpha$ determines how much weight is given to information at different hierarchical levels. This weighting function is applied in both MLM and CLM pre-training, allowing HELM to capture both local codon-level information and the global hierarchical structure of the genetic code. By incorporating this hierarchical loss into pre-training, HELM encourages the model to learn biologically meaningful representations of mRNA sequences, which closely aligns the model’s learning with the biological structure of mRNA sequences. Note that our approach does not require modifying the architecture or tokenization of a model and permits operating on the standard codon vocabulary. mRNA hierarchy modeling with HXE will be referred as HELM from here on.

## 4 Experiments

In our experiments, we first compare various mRNA language modeling configurations by exploring different tokenization strategies and model architectures (Transformer, Mamba, Hyena) with MLM and CLM objectives. We demonstrate how HELM significantly improves mRNA property prediction over its non-hierarchical counterparts. Next, we investigate the effectiveness of hierarchical mRNA encoding, evaluating it through synonymous sequence clustering. Following this, we present generative evaluations include assessing generated sequence quality and diversity via Frechet Biological Distance (FBD), as well as testing the preservation of functional properties in generated mRNA sequences. Finally, we conclude with antibody region annotation, showcasing the model’s practical applicability in bioinformatics.

##### Datasets

Unlike protein and DNA modality which boast many large-scale datasets, mRNA does not have many curated high-quality pre-training datasets. For this reason, we curated the OAS database (Olsen et al., ) which contains antibody mRNA data from over 80 different studies with around 2 billion unpaired and 1.5 million paired sequences from various species. Although prior studies have curated this database on protein level (Ruffolo et al., ; Shuai et al., ; Kenlay et al., ) in the context of antibody-protein language modeling, a high-quality curated version of corresponding mRNA data does not exist. Due to its origin in high-throughput sequencing studies, the mRNA sequences in OAS exhibit varying levels of sequencing errors and noise. We conducted an extensive curation process to ensure high-quality mRNA data for pre-training. Dataset curation details can be found in Appendix A.9. Post-curation, we end up with 15.3 million mRNA sequences for pre-training mRNA LMs.

For downstream evaluation, we use the following six mRNA datasets, including antibody (Ab) encoding mRNA and general mRNA sequences:

- •
Ab1 (Anonymous, ) includes 1200 Ab-mRNA sequences with protein expression labels.
- •
Ab2 (Anonymous, ) contains 3442 Ab-mRNA sequences with protein expression labels.
- •
MLOS (Li et al., ) has 167 viral mRNA sequences with expression in HeLa cells.
- •
iCodon (Diez et al., ) includes 65357 mRNA sequences with thermostability profiles from humans, mice, frogs, and fish.
- •
Tc-Riboswitches (Groher et al., ) consists of 355 riboswitch mRNA sequences with switching factor measurements.
- •
mRFP (Nieuwkoop et al., ) has 1459 mRNA sequences with protein production levels for various gene variants in E. coli.

All downstream evaluation datasets provide regression labels for evaluating the quality of mRNA property prediction.

##### Evaluation

We ran all models for $3$ random data splits (train:val:test split of 70:15:15) and we report average performance across splits. When evaluating property prediction tasks, we use Spearman rank correlation as a standard metric under commonly used probing methodology (Marquet et al., ; Chen et al., ; Outeiral & Deane, ; Harmalkar et al., ) with a TextCNN (Kim, ) head to assess representation quality and transferability. For generative evaluation, we take the metrics used to evaluate generative models for DNA (Stark et al., ; Shuai et al., ) and we adopt them for mRNA data. For sequence region annotation, we report region annotation accuracy. We provide comprehensive experimental details in Appendix A.3.

### 4.1 Pre-training and modeling choices

We first aim to establish a strong non-hierarchical mRNA LM as a baseline. For this, we aim for consistent comparison between various standard pre-training strategies, different architectures, and various tokenization methods. We conduct the comparison on the property prediction mRNA tasks where we pre-train all our language language models on the same curated OAS data and test them on downstream property prediction tasks.

#### 4.1.1 Experimental details

##### Pre-training objective

To train an mRNA LM, we start with two standard pre-training objectives: Masked Language Modeling (MLM) and Causal Language Modeling (CLM). MLM is an inherently bidirectional pre-training method that relies on model learning to capture context from both directions of a sequence. CLM, on the other hand, is unidirectional, predicting each token based on previous tokens. We evaluate both objectives to determine their effectiveness in capturing the biological structure of mRNA sequences.

##### Model architectures

We evaluate various state-of-the-art architectures for sequence modeling: Transformer (Radford et al., ), Hyena (Poli et al., ), and Mamba (Gu & Dao, ). Transformers support both MLM and CLM objectives, while Hyena and Mamba are primary designed for CLM. All models use 50M parameters, balancing performance and efficiency. We found that models of this scale can outperforms larger existing models while maintaining reasonable run-times (see Appendix A.7 and Table 8). Detailed pre-training information is available in Appendix A.2.

##### Tokenization

Inspired by protein and DNA language modeling, we explored various tokenization strategies for mRNA sequences: nucleotide, codon-level, and 6-mer. We observed that codon-level tokenization consistently outperforms other strategies (results in Appendix Sec. A.4). We hypothesize that the superior performance of codon tokenization stems from codons’ role as the fundamental units of the genetic code, directly mapping to amino acids during protein synthesis. In addition, the codon level tokenization allows to naturally capture the variability in wobbling nucleotides, where changes in the third nucleotide do not always alter the amino acid produced (Barricelli, ; Outeiral & Deane, ). Here on, we adopt codon-level tokenization as the default strategy for our models. This way, our vocabulary includes $4^{3}=64$ codon tokens, representing all possible nucleotide combinations, plus special tokens for start, stop, padding, masking, and unknown characters.

##### Baselines

We additionally compare trained models against three state-of-the-art public foundation models: RNA-FM, SpliceBERT, and CodonBERT. While RNA-FM and CodonBERT are 100M parameter models, SpliceBERT is a 20M parameter model. We choose these baselines as CodonBERT is the only existing mRNA LM publicly available while SpliceBERT is pre-trained on pre-mRNA, and RNA-FM is known to be one of the strongest baselines in recent works because of its reported strong performance across multiple domains (Nguyen et al., ; Franke et al., ; Boyd et al., ). Additionally, we include a simple one-hot encoding-based baseline to evaluate how much LMs performances surpass basic sequence representation methods commonly used (Harmalkar et al., ; Boyd et al., ).

#### 4.1.2 Results

We first compare the performance of the transformer-based mRNA language models to state-space Mamba and long-convolutional Hyena-based architectures. In Table 1, we can see that transformer-based models outperform Hyena and Mamba on 4 out of 6 datasets, with the transformer CLM trailing closely behind Mamba on the remaining 2 datasets. The superior performance of transformers is likely due to the relatively short sequence lengths ($\sim$ 1000 tokens) of mRNA sequences in both pre-training and downstream datasets. Transformers, with their explicit global attention mechanism between each pair of tokens, excel at capturing fine-grained interactions in these relatively short sequences. In contrast, Hyena and Mamba architectures are optimized for handling much longer sequences via efficient convolutional or state-space mechanisms and may not fully leverage their strengths in this shorter sequence regime owing to their implicit handling of sequence dependencies, despite performing competitively. Given that transformer-based models generally show better performance than Hyena and Mamba models, we select transformer-based MLM and CLM models for all subsequent experiments.

Additionally, our results indicate that for the transformer models, both MLM and CLM perform comparably on 3 out of 6 datasets (Ab1, Ab2, and iCodon), while MLM outperforms CLM on MLOS and Tc-Riboswitches data, and CLM surpasses MLM on the mRFP dataset. The comparable performance of MLM and CLM across most cases suggests that the statistical and structural properties of mRNA are adequately captured by both unidirectional and bidirectional context modeling. This implies that the underlying biological information in mRNA sequences may be robust to the specific directional biases introduced by MLM or CLM pre-training, with neither approach consistently outperforming the other across varied bio-sequence contexts.

We then compare the performance of our transformer, Hyena, and Mamba models against various foundation model baselines. Table 1 illustrates that our transformer, Mamba, and Hyena-based LMs outperform all baselines by 5-17 percentage points across all downstream tasks which highlights the value of our curated mRNA pre-training dataset, tokenization, and architectural design choices.

**Table 1: Performance of mRNA LMs trained with standard cross-entropy loss on downstream tasks. Our models outperform general-purpose RNA baselines across various mRNA tasks. MLM and CLM show comparable results. Bold: best performance. Italics: second-best. Missing values: model unable to process due to sequence length limitations.**
| Model | Ab1 | Ab2 | MLOS | iCodon | Tc-Riboswitches | mRFP |
| --- | --- | --- | --- | --- | --- | --- |
| Baseline Models |  |  |  |  |  |  |
| One-hot | 0.431 | 0.421 | 0.462 | 0.152 | 0.378 | 0.511 |
| RNA-FM | 0.595 | 0.515 | - | - | 0.504 | 0.527 |
| SpliceBERT | 0.652 | 0.542 | - | - | 0.418 | 0.596 |
| CodonBERT | 0.686 | 0.557 | 0.543 | 0.350 | 0.502 | 0.770 |
| mRNA LMs (Ours) |  |  |  |  |  |  |
| Transformer XE (MLM) | 0.748 | 0.599 | 0.653 | 0.503 | 0.569 | 0.753 |
| Transformer XE (CLM) | 0.752 | 0.597 | 0.611 | 0.498 | 0.531 | 0.815 |
| Hyena XE | 0.743 | 0.594 | 0.623 | 0.526 | 0.517 | 0.844 |
| Mamba XE | 0.742 | 0.585 | 0.620 | 0.509 | 0.520 | 0.823 |

### 4.2 HELM improves mRNA property prediction

Next, we evaluate the effect of learning hierarchical mRNA encoding with HELM for the mRNA property prediction tasks. As can be observed from Table 2, training with biological hierarchical prior consistently improves model performance for both MLM and CLM objectives on all six datasets compared to baseline models trained with standard non-hierarchical cross-entropy. For causal language modeling, HELM-CLM outperforms the non-hierarchical XE-CLM model in five out of six datasets. For the masked language modeling, HELM-MLM outperforms the non-hierarchical XE-MLM model in all datasets. Within hierarchy-aware models, we again observe the same trend that MLM and CLM HELM models perform comparably for most tasks with each performing best on three of the six datasets.

**Table 2: Performance comparison of HELM vs. XE models. HELM models encoding mRNA hierarchy outperform non-hierarchical XE models across all downstream tasks and datasets, highlighting the importance of hierarchy as a strong biological prior. Bold indicates the best performing model.**
| Model | Ab1 | Ab2 | MLOS | iCodon | Tc-Riboswitches | mRFP |
| --- | --- | --- | --- | --- | --- | --- |
| Transformer XE (MLM) | 0.748 | 0.599 | 0.653 | 0.503 | 0.569 | 0.753 |
| Transformer HELM (MLM) | 0.767 | 0.609 | 0.701 | 0.525 | 0.626 | 0.822 |
| Transformer XE (CLM) | 0.752 | 0.597 | 0.611 | 0.498 | 0.531 | 0.815 |
| Transformer HELM (CLM) | 0.760 | 0.614 | 0.592 | 0.529 | 0.619 | 0.849 |

##### When and why learning hierarchical mRNA encoding is most effective?

Hierarchical HELM modeling improves the performance in all predictive tasks we considered. However, the extent of improvement varies across datasets. Greater gains are seen in MLOS, Tc-Riboswitches, and mRFP datasets, while Ab1, Ab2, and iCodon show smaller improvements (see Table 2). We hypothesize that this variation is driven by synonymous codon usage bias which indicates the preference for certain synonymous codons (coding for the same amino acid) to be used more frequently than others due to factors like gene expression, tRNA availability, or evolutionary pressures (Sharp & Li, ; Parvathy et al., ). We hypothesize that datasets with more pronounced synonymous codon usage bias benefit more from learning with hierarchy as HELM treats amino acids as parent nodes and their synonymous codons as child nodes, enabling the model to learn codon usage bias by capturing preferences at both the amino acid and codon levels.

To evaluate synonymous codon usage bias for any dataset, we first calculate the distribution of synonymous codons for each amino acid and then compute the entropy of codon usage for the five most frequent amino acids. Lower entropy indicates less uniform codon distribution hence the preference towards specific codons or codon bias. Fig. 3 and Appendix Table 7 demonstrate that MLOS, Tc-Riboswitches, and mRFP datasets have lower entropy, suggesting stronger codon bias, which aligns with the larger improvements seen in these datasets using HELM models. These findings suggest that HELM’s hierarchical approach is particularly beneficial for datasets with pronounced codon usage biases, highlighting its potential for improved mRNA modeling in biologically diverse contexts.

Figure: 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.
Refer to caption: extracted/5931372/figures/average_entropy.png

##### Evaluating hierarchical encoding through synonymous sequence clustering

To assess how effectively HELM captures the hierarchical relationships between amino acids and their synonymous codons in learned representations, we conduct a clustering analysis. We hypothesize that hierarchy-aware models should group biologically synonymous mRNAs (those coding for the same protein) more tightly than non-hierarchical models. To this end, we conduct experiment on a holdout set of 20000 curated OAS mRNA sequences not used in pre-training. For each sequence, we create 100 synonymous sequences by replacing codons with synonymous alternatives that code for the same amino acid, resulting in nucleotide-level variations but coding for the same protein sequences (see Fig. 3). We extract embeddings from both HELM and XE models and apply k-means clustering. The results, summarized in Table 4, support our hypothesis as HELM models produce significantly tighter clusters as measured by the Silhouette score (Rousseeuw, ), confirming their superior ability to learn and represent the hierarchical structure of mRNA sequences. This clustering behavior confirms HELM’s ability to effectively learn these hierarchical relationships.

**Table 3: Clustering performance metrics comparing hierarchical HELM models with vanilla XE models. Bold indicates better clustering.**
|  | MLM Silh. ($\uparrow$) | CLM Silh. ($\uparrow$) |
| --- | --- | --- |
| XE | 0.74 | 0.68 |
| HELM | 0.91 | 0.83 |

### 4.3 HELM improves generative mRNA sequence design

Our HELM-CLM models are inherently generative and can be used for mRNA sequence generation. We evaluate their ability to produce biologically relevant and diverse mRNA sequences through two experiments: Frechet Biological Distance and functional properties preservation.

##### Evaluating sample quality of generated mRNA

We adopt the Frechet Biological Distance (FBD) metric (Stark et al., ) to assess the quality of generated sequences. FBD is a standard generative evaluation metric which quantifies the similarity between generated and real sample distributions, and is used for biological sequences. Our experiment uses XE and HELM CLM models to generate 2000 mRNA sequences encoding antibody D and J regions, conditioned on signal peptide and V regions from the OAS hold-out dataset. Generated and real mRNAs are translated to protein sequences, from which we extract embeddings using the ESM-2 model (Lin et al., ). FBD is computed as the Wasserstein distance between Gaussian distributions fitted to these embeddings, with lower scores indicating better alignment with real data. We focus on amino acid (and hence protein) level FBD since the encoding hierarchy reflects the codon-amino acid relationship.

In the absence of any other existing generative models for mRNA, we benchmark the HELM model against the non-hierarchical baseline trained with vanilla cross-entropy (XE) loss. As an additional control, we establish a random baseline by generating sequences randomly, following the approach of Stark et al. (). The random baseline serves as a control to assess the significance of the sequence generation capabilities of both our HELM model and the XE baseline.

Results in Fig. 5 demonstrate that both the HELM model and the XE baseline significantly outperform the random baseline, indicating that both models generate mRNA sequences that are meaningfully aligned with real data distributions. Unsurprisingly, the generative performance varies with generation temperature, with higher temperatures leading to greater diversity but worse FBD scores. Notably, HELM consistently achieves better FBD scores than XE across all temperatures, suggesting it produces sequences more representative of real mRNA data while maintaining generation diversity. These results demonstrate HELM’s effectiveness in both predictive and generative tasks, highlighting its potential for improved mRNA sequence design.

##### Evaluating functional properties of generated mRNA

We assess the models’ ability to generate sequences that maintain functional properties across six datasets from our downstream predictive tasks (Sec. 4.1.2). For each dataset, we select up to 2000 sequences, condition the models on the first third of each sequence, and task them with generating the remaining two-thirds. To evaluate the retention of functional properties in generated sequences, we employ pre-trained property prediction models to estimate task-specific labels, following the protocol established in Stark et al. (); Avdeyev et al. (). We use Mean Squared Error (MSE) between predicted labels of generated sequences and true labels of corresponding real sequences as our performance metric, with lower MSE indicating better preservation of functional properties.

Fig.5 illustrates the relative improvement in MSE for HELM models compared to non-hierarchical XE models across all datasets. HELM consistently outperforms XE, with improvements ranging from 2% to 31%. This demonstrates that HELM’s hierarchical encoding better preserves functional properties across diverse tasks and datasets. Notably, we observe greater improvements for MLOS, Tc-Riboswitches, and mRFP datasets, which aligns with the stronger synonymous codon usage bias observed in these datasets (as discussed in Sec.4.2). These results further underscore HELM’s effectiveness in capturing and preserving the biological nuances of mRNA sequences during generation.

Figure: 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.
Refer to caption: x1.png

### 4.4 HELM improves Ab-mRNA sequence region annotation

We evaluate HELM’s performance in annotating specific regions within antibody-encoding mRNA sequences. This task is crucial for understanding immune diversity, guiding therapeutic antibody development, and studying immune-related diseases (Briney & Burton, ). Currently, no machine learning-based tool exists for this specific annotation task.

We use a hold-out set of 2000 antibody-encoding sequences from curated OAS data, previously employed in our clustering and generative experiments. The objective is to predict whether each nucleotide belongs to one of four regions: signal peptides, V, DJ, or constant regions, effectively annotating the entire antibody-encoding sequence. Region labels are curated from the OAS database.

Table 4 shows that HELM significantly improves annotation accuracy compared to the non-hierarchical XE model, particularly when using the Masked Language Modeling (MLM) objective. This improvement can be attributed to two factors. First, HELM’s hierarchical prior better captures the inherent structure of mRNA sequences, where certain regions are functionally and structurally dependent on others. Second, the bidirectional attention in MLM more effectively captures the context needed for precise sequence annotation, which is crucial for identifying distinct regions where both upstream and downstream nucleotide dependencies are important. For details see (Sec. A.8).

## 5 Conclusions

In this work, we highlight the importance of biological hierarchy in mRNA language modeling and show that standard natural language modeling approaches fall short in capturing the biological structure of mRNA sequences. We introduce HELM, a method to incorporate mRNA’s hierarchical structure into bio-language model pre-training. By modulating loss based on codon hierarchy, HELM aligns the learning process with mRNA’s inherent biological structure. In addition, we provide a consistent comparison of various tokenization methods, pre-training strategies, and model architectures to guide future mRNA LM development. HELM consistently outperforms non-hierarchical models by an average of 8% across six diverse downstream property prediction tasks with the biggest improvement observed on datasets with pronounced synonymous codon bias. Moreover, HELM models demonstrate significantly better generative capabilities than the non-hierarchical generative LMs enabling generation of more diverse and plausible mRNA sequences.

One limitation of our work is learning hierarchical relationships in Euclidean space, which may limit the model’s ability to fully leverage the hierarchical prior. In contrast, hyperbolic spaces are naturally better suited for capturing hierarchical relationships, as they can effectively model the exponential growth of tree-like structures. This can be explored in future work to improve hierarchical learning and further enhance mRNA sequence modeling.

## Appendix A Appendix

### A.1 Codon Hierarchy

Figure 6 shows the formalized codon hierarchy used in Helm pre-training.The root node $R$ of this tree represents a node corresponding to all codons. This root node has two immediate child nodes: coding codons $C_{\text{coding}}$ and non-coding codons $C_{\text{non-coding}}$.
The non-coding codons $C_{\text{non-coding}}$ further branch into two child nodes: the start codon node $C_{\text{start}}$ and the stop codon node $C_{\text{stop}}$. The start codon node $C_{\text{start}}$ has a single leaf node corresponding to the codon AUG. The stop codon node $C_{\text{stop}}$ has three leaf nodes corresponding to the codons UAA, UAG, and UGA.
On the other side of the hierarchy, the coding codons $C_{\text{coding}}$ are organized into multiple child nodes, each corresponding to a specific amino acid $A_{j}$. Each amino acid node $A_{j}$ has a number of child nodes, each corresponding to a synonymous codon that encodes for $A_{j}$. Let $\text{Codons}(A_{j})=\{C_{i_{1}},C_{i_{2}},\ldots,C_{i_{m}}\}$ denote the set of all synonymous codons for the amino acid $A_{j}$. These synonymous codons preserve the same amino acid, allowing for degeneracy in the genetic code.
The complete hierarchical structure is thus represented by the tree $H=(V,E)$, where $V$ includes the root node, all internal nodes (corresponding to coding and non-coding codons, start/stop codons, and amino acids), and all leaf nodes (corresponding to individual codons). The edges $E$ define the hierarchical relationships between these nodes.

Figure: Figure 6: Formalized tree of the codon hierarchy used for pre-training HELM models
Refer to caption: extracted/5931372/figures/hierarchy.png

### A.2 Pre-training details

We experimented with GPT-2, Mamba, and Hyena architectures pre-trained on a curated dataset. GPT-2 was trained with both MLM and CLM objective, while Mamba and Hyena were trained exclusively with CLM, given that Mamba is specifically designed for Causal Language Modeling, and recent findings show that CLM is superior to MLM for representation learning (Nguyen et al., ). The models are pre-trained with either vanilla cross-entropy loss (vanilla XE) or hierarchical cross-entropy loss (HXE), also referred to as the HELM model (See Sec. 4.1 in main text).

With codon tokenization, the maximum context length for pre-training is 444 tokens, equivalent to the longest sequence of 1332 nucleotides in the dataset. However, the models are designed to support sequences up to 2048 tokens, thanks to a positional embedding layer that scales to 2048 positional embeddings, offering flexibility for future use with longer sequences without architectural changes. All models were pre-trained for 40 epochs on 8 Nvidia A100 GPUs. For downstream tasks, a 1 million-parameter TextCNN(Kim, ) head was used for fine-tuning, following common probing techniques(Marquet et al., ; Chen et al., ; Outeiral & Deane, ; Harmalkar et al., ).

All models are designed to have approximately 50 million parameters. GPT-2 used 10 layers with a hidden size of 640 and an intermediate size of 2560. For position embedding, we follow the strategy used in GPT-2. We employ an embedding layer and add the embedded positional embeddings to the token embeddings. This approach is consistent across all models in our study.

Although GPT-2 supports CLM by design, we use attention without causal masking for creating MLM variant of it.
We employ the AdamW optimizer with a learning rate scheduler using linear warmup and cosine decay. Initially, we use a learning rate of 1e-3 for all models.
For hierarchical HELM models trained with HXE loss, we reduced the learning rate to 1e-4 for all models to accommodate the smaller scale of the loss compared to XE. We experimented with three different $\alpha$ values for HXE: 0.2, 0.4, and 0.6 (See Sec. 3 in main text for details) with $\alpha=0.2$ performing the best overall.

All models are trained using 8 NVIDIA A100 GPUs, each with 80GB of GPU memory. We utilized mixed-precision training (AMP) to improve computational efficiency. The distributed training setup allowed us to effectively use a batch size of 1024 (128 per GPU across 8 GPUs).

We initially experimented with longer training durations but found no significant improvement beyond 40 epochs. Therefore, all models are trained for 40 epochs to ensure consistent comparison and efficient use of computational resources. The hyperparameters used for training are summarized in the following table:

**Table 5: Hyperparameters for trained LM Models**
| Hyperparameter | GPT-2 | Mamba | Hyena |
| --- | --- | --- | --- |
| Number of layers | 10 | 40 | 7 |
| Hidden size | 640 | 256 | 768 |
| Intermediate size | 2560 | 1024 | 3072 |
| Batch size | 1024 | 1024 | 1024 |
| Learning rate (XE) | 1e-3 | 1e-3 | 1e-4 |
| Learning rate (HXE) | 1e-4 | 1e-4 | 1e-4 |
| Minimum learning rate | 1e-5 | 1e-5 | 1e-6 |
| Weight decay | 0.1 | 0.1 | 0.1 |
| Number of epochs | 40 | 40 | 40 |
| Vocabulary size | 70 | 70 | 70 |

### A.3 Downstream task finetuning details

To evaluate the effectiveness of our pre-trained LMs on downstream tasks, we employ probing approach (Marquet et al., ; Chen et al., ; Outeiral & Deane, ; Harmalkar et al., ) using TextCNN (Kim, ). This allows us to assess the quality and transferability of the learned representations across various downstream predictive tasks.

In our implementation of TextCNN, we set the embedding size to 640 to match the hidden size of our models, ensuring consistency across all probing experiments. The convolutional layer size is fixed at 100, which we find to provide a good balance between model capacity and computational efficiency. To focus solely on evaluating the quality of the learned representations, we freeze the weights of the pre-trained LMs, using them as fixed feature extractors to generate embeddings for the probing experiments.
We conduct a comprehensive hyperparameter search to optimize the performance of the TextCNN probe on each downstream task. Our search strategy employs a grid search methodology, exploring a range of learning rates and batch sizes. Specifically, we investigate learning rates of 3e-4, 1e-4, and 1e-5, which span a reasonable range for finetuning neural networks. For batch sizes, we examine values of 8, 16, 32, and 64, allowing us to balance between update frequency and the stability of gradient estimates.
The combination of three learning rates and four batch sizes results in twelve distinct configurations for each downstream task. We train the TextCNN probe using each of these configurations and evaluate their performance on a hold-out validation set to choose the best hyperparameters. For our final results, we report the performance as measured on a separate test set. This approach ensures that we are comparing the most optimized versions of each probing model, providing a fair assessment of the underlying representations learned by our pre-trained models.
For all datasets, whether they come with predefined splits or not, we maintained a consistent data partitioning strategy. We divide the data into training, validation, and test sets with ratios of 0.75, 0.15, and 0.15, respectively. This split ensures a substantial portion of data for training while allocating sufficient samples for validation and testing.
In cases where the datasets for downstream tasks does not come with predefined train-validation-test splits, we implement a robust splitting strategy. We randomly divide the data into training, validation, and test sets, maintaining consistent proportions across all tasks. To account for potential variability introduced by this random splitting, we repeat the entire process using three different random seeds. The reported results for these tasks represent the average performance across these three splits, providing a more reliable estimate of the model’s generalization capabilities.
This comprehensive finetuning and evaluation process allows us to systematically and fairly assess the transferability and quality of the representations learned by our pre-trained models.

### A.4 Tokenization Strategy Experiments

We conducted experiments to compare the performance of different tokenization strategies on various mRNA datasets, including nucleotide-level, 6-mer, and codon-level tokenization. The experiments were performed using a 50M parameter GPT-2 model with a Masked Language Modeling (MLM) objective. Other experimental setups, such as data processing, training pipelines, and evaluation metrics, were kept consistent across all models to ensure a fair comparison of tokenization strategies.

#### A.4.1 Tokenization Strategies

- •
Nucleotide-level Tokenization: This strategy treats each nucleotide (A, C, G, U) as a single token, resulting in a vocabulary size of 4.
- •
6-mer Tokenization: In this approach, sequences are split into overlapping or non-overlapping chunks of six nucleotides, generating a more complex vocabulary.
- •
Codon-level Tokenization: Codon-level tokenization represents sequences as triplets of nucleotides, resulting in a vocabulary of 64 codons, corresponding directly to amino acids.

#### A.4.2 Results

Table 6 presents a comparison of the performance of different tokenization strategies across several mRNA datasets in terms of Spearman rank correlation (higher the better). Codon-level tokenization generally performs best across most datasets, except for mRFP, where 6-mer tokenization outperformed codon-level tokenization.

**Table 6: Comparison of tokenization strategies across different mRNA datasets. Codon-level tokenization generally performs best, with 6-mer outperforming on mRFP.**
| Dataset | Nucleotide-level | Codon-level | 6-mer |
| --- | --- | --- | --- |
| Ab1 | 0.699 | 0.747 | 0.733 |
| Ab2 | 0.569 | 0.599 | 0.582 |
| MLOS | 0.528 | 0.654 | 0.625 |
| Tc-Riboswitches | 0.533 | 0.570 | 0.562 |
| mRFP | 0.820 | 0.754 | 0.857 |

The results show that codon-level tokenization captures the biological significance of mRNA sequences and is more suitable for most tasks. However, on MRFP and iCodon datasets, 6-mer tokenization yielded better performance.. The 50M parameter GPT-2 model with the MLM objective was used in all experiments, and all other experimental setups remained the same to facilitate fair comparison. For further experimental details and setup, please refer to this appendix.

### A.5 Entropy Details

Table 7 reveals that MLOS, Tc-Riboswitches, and mRFP datasets exhibit lower entropy values, indicating stronger codon usage bias compared to Ab1, Ab2, and iCodon datasets which have higher entropy and more uniform codon distributions. This observation aligns with our hypothesis, as HELM models achieve greater improvements in datasets with skewed codon usage, suggesting that hierarchical learning better captures and leverages these codon biases for enhanced accuracy.

**Table 7: Average entropy values of synonymous codon distributions for the top five most frequent amino acids in each dataset. Datasets with lower entropy values, such as MLOS, Tc-Riboswitches, mRFP, and iCodon, exhibit stronger synonymous codon usage bias, correlating with the enhanced performance of hierarchical HELM models on these datasets (see Fig. 3 and Table 2).**
| Dataset | Average Entropy (top-5) |
| --- | --- |
| Ab1 | 1.81 |
| Ab2 | 1.81 |
| MLOS | 1.62 |
| iCodon | 1.77 |
| Tc-Riboswitches | 1.63 |
| mRFP | 1.20 |

### A.6 Impact of the choice of α 𝛼 \alpha italic_α for HXE loss

In this section, we present a detailed analysis of the impact of the hyperparameter $\alpha$ on the performance of models trained using the Hierarchical Cross-Entropy (HXE) loss. The $\alpha$ parameter in the HXE loss plays a crucial role in weighting the importance of hierarchical relationships during model training. Specifically, $\alpha$ controls how much weight is given to mistakes that violate the hierarchical structure of mRNA sequences, where codons are organized based on their synonymous relationships.

As explained in Sec. 3, the HXE loss leverages the hierarchical nature of mRNA sequences by applying differential penalties to the prediction errors based on their position within the hierarchy. The weighting function $\lambda(C)$ is
$\lambda(C)=\exp(-\alpha h(C))$, where $h(C)$ is the height of the node $C$ and $\alpha>0$, determines the extent to which errors are penalized. Higher values of $\alpha$ increase the penalty for mistakes higher up in the hierarchy (e.g., confusing codons that code for different amino acids), while lower values of $\alpha$ result in a more uniform penalty across the hierarchy.

**Table 8: Ablation study on HXE loss $\alpha$ across downstream tasks**
| Hyperparameter | Ab1 | Ab2 | iCodon | mRFP | Tc-Riboswitches |
| --- | --- | --- | --- | --- | --- |
| Transformer HELM (MLM) |  |  |  |  |  |
| $\alpha=0.2$ | 0.767 | 0.609 | 0.525 | 0.822 | 0.626 |
| $\alpha=0.4$ | 0.767 | 0.608 | 0.511 | 0.831 | 0.653 |
| $\alpha=0.6$ | 0.762 | 0.599 | 0.508 | 0.796 | 0.582 |
| Transformer HELM (CLM) |  |  |  |  |  |
| $\alpha=0.2$ | 0.760 | 0.614 | 0.529 | 0.849 | 0.619 |
| $\alpha=0.4$ | 0.762 | 0.608 | 0.527 | 0.841 | 0.551 |
| $\alpha=0.6$ | 0.752 | 0.592 | 0.526 | 0.798 | 0.580 |

To assess the impact of $\alpha$ we evaluated the performance of the transformer models trained with the HXE loss using three different values of
$\alpha$: 0.2, 0.4, and 0.6. These values were chosen to explore a range of hierarchical weighting, from mild to strong penalties for hierarchical errors. Table 8 summarizes the performance of the models on each downstream task. We observe the following trends:

- •
$\alpha$ =0.2: At this lower value of $\alpha$, the model applies relatively mild penalties for hierarchical mistakes. The performance across most tasks is stable, indicating that even a modest incorporation of hierarchical weighting can enhance the model’s performance by aligning with the biological structure of mRNA sequences.
- •
$\alpha$ =0.4: Increasing $\alpha$ to 0.4 generally results in improved performance for MLM, particularly for the mRFP and Tc-Riboswitches tasks. This suggests that at this level, the model benefits from a stronger alignment with the hierarchical structure, which is critical for these tasks. However, the performance on some tasks (e.g., Ab1 and Ab2) remains stable, indicating that the benefit of increased $\alpha$ might be task-dependent.
- •
$\alpha$ =0.6: When $\alpha$ is further increased to 0.6, we observe a decline in performance for several tasks, including mRFP and Tc-Riboswitches. This suggests that an overly strong hierarchical penalty can restrict the model’s flexibility, leading to overfitting to the hierarchical relationships and a loss of generalization across diverse tasks.

Thus, our experiments show that a moderate value of
$\alpha$ (0.2 to 0.4) generally yields the best results across the diverse set of tasks. This value strikes a balance between enforcing hierarchical structure and maintaining model flexibility.

### A.7 Scalability Experiments

In this section, we provide the details of our experiments with three GPT-2-based models: 20 million (20M), 50 million (50M), and 100 million (100M) parameters. All models were trained using the HELM framework with a Causal Language Modeling (CLM) objective. The performance was evaluated on downstream mRNA property prediction tasks.

#### A.7.1 Model Configurations

- •
20M Model: This model consists of 12 transformer layers, each with hidden size of 384 and intermediate size of 1536.
- •
50M Model: This model consists of 10 transformer layers, each with a hidden size of 640 and intermediate size of 2560.
- •
100M Model: This model consists of 14 transformer layers, each with a hidden size of 768 and intermediate size of 3072.

All models were trained on the same mRNA dataset using the HELM-modified CLM objective, and we maintained a consistent training pipeline across all experiments.

#### A.7.2 Results

Figure: Figure 7: Performance comparison of GPT-2 models with 20M, 50M, and 100M parameters trained using the HELM framework.
Refer to caption: extracted/5931372/figures/scale.png

This analysis illustrates the scalability of HELM and highlights the balance between model size and performance efficiency (see Fig.  7. The diminishing returns beyond 50M parameters suggest that optimal model sizes for mRNA tasks may lie within the 50M-100M parameter range.

The training duration for each model configuration is illustrated in Figure 8.

Figure: Figure 8: Average Pre-training Time Across Models with Varying Parameter Counts
Refer to caption: extracted/5931372/figures/duration.png

### A.8 Ab-mRNA sequence region annotation experiment

In this experiment, we approached the problem of annotating specific regions within b-mRNA sequences as a multi-class classification task. Specifically, we aimed to predict the start and end positions of critical regions within the sequence: V (variable), SP (signal peptide), DJ (diversity and joining), and C (constant) regions. To achieve this, we designed a model that predicted seven key positions: the start and end of the V region, the start and end of the SP region, the start and end of the DJ region, and the start of the C region.
Notably, we did not predict the end of the C region since this boundary is trivially determined in the context of the sequences.

To ensure consistency across the dataset, we fixed the maximum number of nucleotides in the test dataset to 902. This allowed us to normalize the sequence length, ensuring that predictions were made for uniform positional slots across all samples. Each of the seven positions was treated as a separate prediction task within the classification framework.

As a comparison metric, we utilized the average accuracy across all predicted regions. By measuring the accuracy for each start and end position independently and averaging these results, we obtained a holistic view of the model’s performance across the entire Ab-mRNA sequence. This metric enabled us to evaluate how well the model generalized across various regions and how effectively it could predict the boundaries of biologically relevant areas within the sequence.

### A.9 Pre-training data curation and statistics

Figure: Figure 9: Comparison of gene distributions between the original OAS human data and the curated dataset. (a) shows the gene distribution for heavy chains, and (b) shows the gene distribution for light chains. (c) illustrates the balanced numbers of heavy and light chains, including the distribution of Kappa and Lambda isotypes, along with V-allele distribution.
Refer to caption: extracted/5931372/figures/heavy_genes_oas.png

As discussed in Sec. 4 in main text, we ensure that the mRNA data used for pre-training our models in this study is of the highest quality, with careful attention to preserving the biological diversity and representativeness of the sequences. We conduct statistical analyses to confirm that the curated dataset maintains a balanced representation of gene types, comparable to the original OAS human dataset and did not introduce any biases in the gene representation.

First, we filtered sequences based on the ANARCI status annotation in the dataset, excluding those with unusual residues, indels, truncations, or a lack of conserved cysteines, which are often problematic. We further refined our selection by focusing only on sequences with a V and J identity greater than 0.7, ensuring a high degree of similarity to known reference sequences. Additionally, we retained only those sequences labeled as productive and complete vdj, indicating they are fully functional and contain complete variable, diversity, and joining regions. For the purpose of this study, we restricted the dataset to human sequences only. Recognizing that the paired sequences in the database are significantly fewer in number compared to unpaired ones, we unpaired these sequences to increase the dataset’s size and diversity. We then performed sequence similarity analysis separately on paired and unpaired sequences to eliminate redundancy with a threshold of 0.5. To maintain a balanced representation, we manually adjusted the number of heavy and light chains. Finally, we compared the gene frequencies of heavy and light chains in our curated dataset against the original human OAS dataset to ensure that our curation process did not introduce artificial overrepresentation or underrepresentation of any gene types. This process yields $15.3$ million sequences with $7.7$ million heavy and $7.6$ million light chains each.

The following figures present the distribution of gene types for heavy and light chains in both the original and curated datasets, as well as a breakdown of heavy and light chain numbers, including Kappa (K) and Lambda (L) isotypes.

## Peer Reviews

confidence: 3
contribution: 4
initial_rating: 6
presentation: 4
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?
soundness: 4
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.
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.
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.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
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.
soundness: 2
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: See W3.
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: - Are you planning to open-source the curated dataset ?
soundness: 3
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.
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.
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.

Meta-review: 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.

Decision: Accept (Poster)