OmniGenome: Aligning RNA Sequences with Secondary Structures in Genomic Foundation Models

Recent studies [21, 22] have shown that data diversity can enhance FM performance without significantly increasing model capacity. For the OmniGenome pretraining, we collected transcriptome data from the OneKP initiative [23], which compiles large-scale RNA-Seq data from $1,124$ plant species. Because the raw biological data are sparse, noisy, and characterized by long sequences, they are far from ready for effective FM training.

To address this problem, we developed a four-step data curation protocol to improve data quality. $\blacktriangleright$ Raw RNA-Seq data are often excessively long with thousands of bases. We first sliced them into segments with a window size of $1,024$ bases to provide a sufficient context window for RNA sequence understanding. $\blacktriangleright$ To enhance training efficiency and reduce bias, we removed all duplicate sequences. $\blacktriangleright$ To tackle incomplete transcriptome data and other noises, we discard sequences shorter than $50$ bases. $\blacktriangleright$ To facilitate the sequence-structure alignment training, we adopt ViennaRNA to obtain the secondary structures for the sequences.

In the realm of genomic sequence modeling, modeling resolution, and performance highly depend on the tokenization strategy [3, 21]. Because the traditional k-mers [24, 13] and BPE [19, 22] methods combine multiple bases into single tokens, they may compromise modeling resolution and thus fail to capture granular features. To achieve a single nucleotide (SN) resolution necessary for genomic tasks like RNA design and structure prediction, we employed single nucleotide tokenization (SNT) method [3, 21] (see an illustrative example in Fig. 3). Additionally, to simplify the implementation, we adopted a vocabulary $\{$‘A’, ‘T’, ‘C’, ‘G’, ‘U’, ‘N’, ‘(’, ‘)’, ‘.’$\}$ to unify the tokenization of both nucleotide bases and RNA secondary structure information. Following BERT-style models [19], we incorporated special tokens, e.g., ‘<mask>’, to facilitate masked bases and structure prediction.

As discussed in Section 1, a key desideratum for complex genomic understanding and RNA sequence design is helping the model build an alignment between RNA sequences and their corresponding secondary structures. Bearing this in mind, we designed the following aggregated pretraining objective for OmniGenome:

$$\mathcal{L}_{\texttt{pretrain}}=\mathcal{L}_{\texttt{Str2Seq}}+\mathcal{L}_{\texttt{Seq2Str}}+\mathcal{L}_{\texttt{MRLM}}+\lambda||\theta||_{2},$$

(1)

where $\lambda$ is the $\ell_{2}$ regularization weight and $\theta$ represents the parameters of OmniGenome. The following paragraphs explain the design principles of each loss function used in equation (1).

• •

  $\mathcal{L}_{\texttt{Str2Seq}}$ is designed to enable OmniGenome to predict bases given partially masked sequences with the context of structures, so that OmniGenome can perceive RNA secondary structure by constructing Str2Seq prediction. To achieve this, we mask $15\%$ of the bases and structure tokens, encouraging the model to infer masked bases (i.e., $\{$‘A’, ‘T’, ‘C’, ‘G’, ‘U’, ‘N’$\}$) and structure tokens (i.e., $\{$‘(’, ‘)’, ‘.’$\}$). Specifically, $\mathcal{L}_{\texttt{Str2Seq}}$ is defined as the classic cross-entropy loss widely used in the masked language modeling:

  $$\mathcal{L}_{\texttt{Str2Seq}}=-\frac{1}{|m|}\sum_{i=1}^{m}\log p(x_{i}\mid x_{\setminus i}),$$

  (2)

  where $m$ is the number of masked nucleotide and structure tokens, and $p(x_{i}|x_{\setminus i})$ indicates the probability of predicting the masked nucleotide $x_{i}$ based on its context.

• •

  Complementary to $\mathcal{L}_{\texttt{Str2Seq}}$, $\mathcal{L}_{\texttt{Seq2Str}}$ is designed to enable OmniGenome for Seq2Str prediction. Instead of directly feeding the secondary structure into OmniGenome during training, this objective employs the RNA secondary structures as supervised-training labels. Unlike Str2Seq, the Seq2Str prediction is generally deterministic because the optimal structures are unique. This task is implemented as a token-level classification, where the $\mathcal{L}_{\texttt{Seq2Str}}$ loss is defined in the following cross-entropy loss:

  $$\mathcal{L}_{\texttt{Seq2Str}}=-\sum_{i=1}^{N}\sum_{c=1}^{C}s_{ic}\log(\hat{s}_{ic}),$$

  (3)

  where $s_{ic}$ denotes the label $c$ of secondary structure at the $i$-th position, and $\hat{s}_{ic}$ is the probability predicted by a linear classifier deployed on OmniGenome. $N$ is the length of an RNA sequence and $C=3$ denotes the number of the possible labels of structure, i.e., $\{$‘(’, ‘)’, ‘.’$\}$.

• •

  The last one $\mathcal{L}_{\texttt{MRLM}}$ is similar to the conventional masked language modeling loss in NLP. It aims to improve the model’s understanding of latent patterns in RNA sequences by masking or randomly replacing $15\%$ of nucleotide bases. The definition of $\mathcal{L}_{\texttt{MRLM}}$ is similar to that of $\mathcal{L}_{\texttt{Str2Seq}}$, but it only considers the prediction of masked nucleotide bases.

There will be a data leakage problem in $\mathcal{L}_{\texttt{Seq2Str}}$ objective while the structure is appended to the sequences. In practice, $70\%$ of RNA sequences are used for structure-contextualized pretraining for the $\mathcal{L}_{\texttt{Str2Seq}}$ objective, while the remaining $30\%$ are used for the other two objectives. This proportion is based on our experiment observation.

OmniGenome adopts the classic Transformer encoder architecture with bidirectional multi-head attention. We designed two variants, dubbed OmniGenome${}^{52\texttt{M}}$ and OmniGenome${}^{186\texttt{M}}$ with $52$ and $186$ million parameters respectively. Some key model specifications are summarized in Table 1.

The training specifications are summarized as follows. $\blacktriangleright$ The learning rate is set to $5\times 10^{-5}$ and the weight decay is set to $0.01$. $\blacktriangleright$ We use AdamW as the optimizer with hyperparameters $\beta_{1}=0.9$ and $\beta_{2}=0.999$. $\blacktriangleright$ We use a linear decay strategy with a warm-up period of $1,000$ steps in the learning rate scheduler. $\blacktriangleright$ The batch size is set to $2,048$. $\blacktriangleright$ No dropout is applied during pretraining, and we use the rotary position embeddings [25] to further enhance the model’s scalability to long RNA sequences. $\blacktriangleright$ We built a distributed training environment with $8$ Nvidia RTX $4090$ GPUs and pretrained for $4$ weeks, while its configuration is introduced in  Appendix 6.

There are no direct counterparts to OmniGenome in plant genomic modeling. However, we can compare it with the following recent genomic FMs as potential baselines to help evaluate the performance of OmniGenome as shown in Table 2. Please refer to Appendix 8 for brief introductions of these FMs. We are aware that some FMs are also developed for RNA, such as RNA-FM [26], RNA-MSM [27], Uni-RNA [28], 5UTR-LM [16], etc. However, we cannot compare OmniGenome with them because their source codes are very hard to work with in our efforts or the models are not publicly available. In the secondary structure prediction tasks, we input the secondary structures computed by ViennaRNA that are different from the true structures into OmniGenome$+$ to avoid data leakage problems.

In this experiment, we demonstrate the Str2Seq prediction capability of OmniGenome based on RNA sequence design. We employed the Eterna [17] V2 benchmark, which consists of $100$ specified secondary structures. This task aims to design plausible RNA sequences based on the reference structures. It is intuitive that higher performance in RNA design tasks indicates better structure-to-sequence prediction capability. We develop a simple GA algorithm based on OmniGenome models to solve RNA design puzzles. The implementation details can be found in Fig. 5 in Appendix 9. Within our GA implementation, the population size is set at $1000$, with $100$ iterations, and the mutation rate for each base is $0.5$. We set the maximum running time for each puzzle to $1$ hour because OmniGenome can easily solve most of the RNA design puzzles within $1$ hour, while each puzzle’s time budget is generally $24$ hours in existing RNA design methods. We utilize accuracy as our evaluation metric, meaning that only sequences with entirely correct structures are considered resolved. The experimental results are available in Table 3.

It is observed from Table 3 that the FM-free method RNAInverse solved 30 of the RNA design puzzles, indicating a competent but not optimal capability in sequence design. We select the FMs employing SNT to conduct comparison experiments, as tokenizers like k-mers and BPE cannot handle the SN-level resolution modeling, i.e., masked nucleotide prediction. We cannot compare with HyenaDNA because it does not support masked nucleotide prediction. The models SpliceBERT and CDSBERT demonstrated trivial proficiency in the RNA sequence design task, solving $3$ and $0$ puzzles, respectively. This suggests these FMs cannot precisely predict the bases without any Str2Seq mapping ability. OmniGenome variants without structure context input have performance similar to the baselines, i.e., OmniGenome${}^{52\texttt{M}}$ and OmniGenome${}^{186\texttt{M}}$ solved $4$ and $5$ puzzles, respectively. In contrast, OmniGenome${}^{52\texttt{M}}+$ and OmniGenome${}^{186\texttt{M}}+$ significantly outperformed other models with $71$ and $74$ puzzles solved, respectively, underscoring the impact of Str2Seq prediction in sequence-structure alignment, which evidently bolsters the models’ abilities to predict bases according to the structure contexts. Besides, we expect an increase in performance with sufficient computational budgets and the findings provide crucial evidence of the significance of Str2Seq for RNA sequence design.

This subsection assesses OmniGenome in both Seq2Str and Str2Seq prediction aspects of sequence-structure alignment. The evaluation of Seq2Str is based on zero-shot secondary structure prediction. We use OmniGenome and OmniGenome$+$ without fine-tuning to predict the secondary structures of sequences from the test datasets and measure the macro-F1 score, where better structure prediction performance indicates a stronger capability for Seq2Str prediction. We evaluate Str2Seq mapping ability based on masked nucleotide base prediction. This task leverages the FM’s capacity to predict sequences with the context of structure. We randomly mask one base in the sequences from the test datasets and leverage the Str2Seq ability of OmniGenome$+$ to predict the masked bases in a zero-shot scenario. We assess the model’s accuracy in predicting masked bases both with and without structure context. The experimental results are available in Table 4.

The results in Table 4 reveal OmniGenome’s effectiveness in both secondary structure prediction (i.e., Seq2Str) and masked nucleotide prediction (i.e., Str2Seq). More specifically, OmniGenome${}^{52\texttt{M}}$ and OmniGenome${}^{186\texttt{M}}$ show impressive performance for zero-shot sequence-structure alignment, while OmniGenome${}^{52\texttt{M}}+$ and OmniGenome${}^{186\texttt{M}}+$ outperform the base variants of OmniGenome. Even when ablating structure context, OmniGenome${}^{186\texttt{M}}$ achieves performance comparable with ViennaRNA. Interestingly, if the structures calculated by ViennaRNA are fed into the model, OmniGenome$+$ models can predict the structure more precisely and outperform ViennaRNA on most of the secondary structure prediction datasets, e.g., bpRNA, RNAStralign, and Rfam. From the results of masked nucleotide prediction, it is observed that sequence reconstruction performance can be enhanced by the structure context. OmniGenome${}^{52\texttt{M}}+$ and OmniGenome${}^{186\texttt{M}}+$ significantly outperform baseline FMs like SpliceBERT.

Overall, the results in Table 4 provide a comprehensive evaluation of the FMs’ capabilities in sequence-structure alignment. This underscores OmniGenome’s performance and robustness in generalizing among different species without prior fine-tuning, highlighting their potential utility in SN-level in-silico genomic modeling.

The results in Table 5 demonstrate the performance of OmniGenome and its generalizability across various fine-grained RNA genomic modeling tasks. It is observed that the base variants of OmniGenome models achieve better results than both RNA and DNA FM baselines, including Agro-NT and CDSBERT, which contain hundreds of millions of parameters. This is because the existing FMs usually adopt k-mers tokenization that cannot handle single nucleotide resolution tasks, e.g., single nucleotide mutation detection and repair. Because of the Seq2Str pretraining, OmniGenome and OmniGenome$+$ models exhibit strong results in secondary structure prediction, underscoring OmniGenome’s capabilities in SN-level RNA sequence understanding and manipulation.

The PGB is a plant-oriented genomic benchmark. Although the benchmark datasets in PGB are DNA-based tasks, we can still evaluate the performance of OmniGenome and its generalizability on multi-modal (i.e., DNA and RNA) genomic tasks. The results in Table 6 reveal substantial variability in the performance of different FMs, where OmniGenome${}^{52\texttt{M}}$ outperformed other baseline models across most tasks, particularly in tasks like Polyadenylation, Splice Site, and Enhancer Region classification, where they achieved the highest F1 scores. This suggests that OmniGenome’s architecture is particularly adept at handling complex genomic sequences. In comparison, existing FMs, e.g., CDSBERT and Agro-NT, showed lower performance with more parameters than OmniGenome. Besides, the performance of OmniGenome${}^{52\texttt{M}}+$ suggests that the structure context can further enhance the performance of genomic modeling.

Overall, OmniGenome models achieve state-of-the-art performance on both benchmarks, especially for OmniGenome$+$ variants. The results underscore the importance of sequence-structure alignment in achieving complex genomic modeling tasks.

The working mechanism of our designed genetic algorithm based on OmniGenome$+$ is implemented as the following five-step process:

1. Step $1$.

   Given the target RNA secondary structure, we use OmniGenome to generate a set of candidate sequences $\mathcal{P}=\{\mathbf{s}^{i}\}_{i=1}^{N}$.

2. Step $2$.

   If the termination criterion is not met, go to Step $3$; otherwise, output the current best sequence $\mathbf{s}^{\ast}=\operatorname*{argmax}_{\mathbf{s}\in\mathcal{P}}f(\mathbf{s})$.

3. Step $3$.

   Based on $\mathcal{P}$, use single-point crossover and mutation to generate a population of offspring sequences $\mathcal{O}=\{\tilde{\mathbf{s}}\}_{i=1}^{N}$.

4. Step $4$.

   Combine $\mathcal{P}$ and $\mathcal{O}$ to obtain $\mathcal{S}=\mathcal{P}\bigcup\mathcal{O}$, and use OmniGenome to predict the corresponding secondary structures of each sequence in $\mathcal{S}$. Evaluate the fitness values of sequences in $\mathcal{S}$.

5. Step $5$.

   Sort $\mathcal{S}$ according to the fitness values and preserve the best $N$ sequences to constitute a new $\mathcal{P}$. Return to Step $2$.

Note that the fitness value of a sequence $\mathbf{s}$, denoted as $f(\mathbf{s})$, is evaluated as the Hamming distance of the RNA secondary structure predicted by OmniGenome against the target structure. The above genetic algorithm is not terminated until the sequence for the target RNA secondary structure is identified or the allocated computational budget is exhausted.

As demonstrated in the zero-shot experiments in Table 3, OmniGenome$+$ models achieve state-of-the-art performance. We show several complex examples of puzzles from the EternaV2 benchmark. According to Fig. 6, puzzles #5 and #11 with approximately $200$+ bases are solved, while these puzzles are challenging to existing FMs. Even for puzzles that are not completely solved, e.g., puzzles #3 and #27, OmniGenome${}^{186\texttt{M}}+$ generates very similar structures, where the nucleotide base difference ratio between the designed structure and the target structure is only $\approx 3\%$. This finding indicates the proficiency of OmniGenome$+$ models in solving challenging single-nucleotide resolution genomic tasks.

Details of the RGB can be found in Table 7. RGB contains $7$ SN-level tasks that are curated or collected from published articles. The purpose of RGB is to benchmark genomic FMs in challenging SN-level modeling tasks such as detection and repair of SN mutations, mRNA sequence degradation rates, and RNA secondary structure prediction. Due to the lack of a plant RNA benchmark dataset, RGB includes the modeling of RNA sequences from a variety of species, e.g., plant and human. The sequence length in RGB ranges from $107$ to $512$, which is sufficient for most RNA understanding tasks. In summary, these multi-species and SN-level tasks in RGB serve as the first comprehensive benchmark utilized to assess the RNA sequence modeling capabilities of OmniGenome and its baseline models. The brief introduction of the datasets in RGB is as follows:

• •

  Single-Nucleotide Mutation Detection (SNMD): We developed a plant RNA dataset synthesizing the single-nucleotide mutations. Focused on identifying potential single nucleotide changes, this task is essential for detecting mutations linked to genetic disorders. The SNMD dataset introduces up to $10$ random mutations in the original sequences, regardless of variation ratios. Cross-entropy is utilized as the loss function for this binary token classification task.

• •

  Single-Nucleotide Mutation Repair (SNMR): This task challenges the model to suggest corrective actions at the single nucleotide level, aiding in gene therapy approaches. The SNMR dataset mirrors the SNMD dataset, with cross-entropy as the loss function, indicating a token 4-way (i.e., A, U, C, G) classification task.

• •

  mRNA Degrade Rate Prediction (mRNA): Estimating the decay rate of nucleotides in mRNA sequences, this task is vital for deciphering gene expression and regulation. The dataset originates from the Kaggle COVID-19 vaccine design competition⁶⁶6https://www.kaggle.com/competitions/stanford-covid-vaccine, focusing solely on sequence-based degradation rate prediction and excluding RNA structures. It’s a token regression task using MSE as the loss function, with the dataset resplit into training, validation, and testing sets for evaluation.

• •

  RNA Secondary Structure Prediction (bpRNA & Archive2 & RNAStralign & Rfam): Aiming to predict RNA folding into secondary structures, this task is fundamental to RNA functionality and interactions. We evaluated OmniGenome on four datasets, bpRNA [5] (TR0, VL0, TS0 sets), ArchiveII [6], RNAStralign [4] and Rfam [7]. Following existing works, we have excluded sequences over $512$ bases and complex structures, simplifying to three symbols: ‘(’, ‘.’, ‘)’Ṙesults may not directly compare with other studies due to these modifications. Cross-entropy serves as the loss function.

Please find the appendix for the input and output examples of each subtask in RGB. The detailed task descriptions for each nucleic acid and species, including the number of examples, classes, evaluation metric, and sequence length, are outlined in Table 7. Each task is carefully curated to reflect the complexity and variety inherent in genomic data, providing a robust framework for assessing the nuanced capabilities of state-of-the-art RNA FMs.

Table 8 show the virtual examples of different datasets in RGB. Please refer to our supplementary materials to find the datasets for more details.

The Plant Genomic Benchmark [14] (PGB) provides a comprehensive suite of datasets designed to evaluate and improve the predictive capabilities of genomic models in plant biology. This benchmark, as shown in Table 9, encompasses a range of critical genomic tasks⁷⁷7https://huggingface.co/InstaDeepAI/agro-nucleotide-transformer-1b, including binary classification, single and multi-variable regression, and multi-label classification, addressing various aspects of plant genomics such as RNA processing, gene expression, and chromatin accessibility. By integrating diverse genomic tasks, the PGB aims to facilitate advanced research and development in plant genomics, offering a robust platform for the assessment and enhancement of model performance across different plant species. To obtain a detailed description of PGB, please refer to Agro-NT [14].