RAGenome: Scaling Retrieval-Based Genomic Language Models to Long Contexts

AI summary

Overview

Research area: Machine learning for genomics — specifically self-supervised pretraining of genomic language models (gLMs) using multiple sequence alignments and retrieval.

Technical level: Intermediate. The paper assumes familiarity with transformer architectures, masked language modeling, and attention mechanisms, but the core idea (retrieve homologous sequences, drop gap tokens, attend across species) is explained clearly enough for a reader with general deep learning background.

Scope: The paper introduces RAGenome, a 168M-parameter retrieval-based genomic language model pretrained on whole-genome alignments of 100 vertebrates, and evaluates it on gene finding and variant effect prediction.

What This Paper Is About

Genomic language models trained on raw DNA are enormously expensive, and models that instead learn from evolutionary information via multiple sequence alignments (MSAs) are stuck at very short input windows of around 128–256 nucleotides. That short window makes them unable to handle tasks like gene finding, where the signals that define exon and intron boundaries can be thousands of base pairs apart. RAGenome's goal is to keep the explicit evolutionary signal of MSA-based models while extending the usable context by roughly 100×, so that one model can handle both short-range evolutionary tasks and long-range structural tasks.

Key Contributions

1. A retrieval-based gLM over whole-genome alignments. RAGenome is described as the first retrieval-based gLM trained on whole-genome alignments from 100 vertebrates, explicitly modeling evolutionary relationships while scaling to 100× longer context windows than MSA-based methods, at far lower compute than large-scale gLMs.
2. Substantially improved long-range capability on gene finding. Extending context raises gene finding performance from an MCC of 0.45 (GPN-Star, the strongest MSA-based gLM) to 0.60 for RAGenome, narrowing the gap to much larger long-context gLMs.
3. Preserved performance on variant effect prediction. RAGenome reaches an AUROC of 0.87 on ClinVar pathogenic-vs-benign classification, indicating that the retrieval mechanism still captures evolutionary conservation signals even with long input contexts.
4. A single efficient architecture that unifies both capabilities. Aside from Evo2, RAGenome is the only model in the comparison table that performs competitively on both tasks, at 168M parameters and 16,032 GPU hours of training.

Main Findings

• Gene finding improves sharply at each context extension. MCC rises from 0.47 to 0.52 after extending context to 4,096 nucleotides, and from 0.54 to 0.57 after extending to 13,312 nucleotides. The largest gains occur immediately after each transition, suggesting context length — not training compute alone — drives the improvement.
• Variant effect prediction stays flat across pretraining stages. The authors interpret this as evidence that the retrieval mechanism preserves the short-range evolutionary signal rather than trading it away for long-range ability.
• Context length helps even at a fixed checkpoint. An ablation holding the model fixed and varying only inference context shows monotonic gains: for the final checkpoint, MCC is 0.53 at 1,024 nucleotides, 0.57 at 4,096, and 0.60 at 13,312.
• Retrieval is essential. With a retrieval budget of 0, performance collapses (expected, since the model was never trained without retrieved context). Performance improves substantially with a small budget; variant effect prediction saturates by a budget of 40,000, while gene finding continues to improve slightly up to the full budget.
• Gap tokens dominate alignments. Analyzing the 100-vertebrate WGA used by GPN-Star, the authors found 73% of tokens are gaps. In the top 5% most conserved regions the gap fraction is lower at 41%. For evolutionarily distant species, the fraction of the human genome that aligns can fall below 2%.
• Comparison against baselines on gene finding (MCC). GPN-MSA 0.43, GPN-Star 0.45, HyenaDNA 0.33, DNABERT-2 0.41, GENA-LM 0.50, NT-MS 0.66, Evo2 0.81, RAGenome 0.60. RAGenome beats all MSA-based and mid-scale long-context models but trails the far larger NT-MS and Evo2.
• Comparison against baselines on variant effect prediction (AUROC). GPN-MSA 0.90, GPN-Star 0.92, HyenaDNA 0.51, NT-MS 0.57, Evo2 0.80, RAGenome 0.87. DNABERT-2 and GENA-LM results are not applicable due to tokenization strategy.
• Per-class differences on gene finding. RAGenome substantially outperforms NT-MS at identifying donor and acceptor splice sites on both strands, while NT-MS is better only at introns — and because introns span many more positions in the test set, NT-MS's aggregate MCC is higher. Evo2 outperforms both on most classes.
• Taxonomy embeddings recover phylogenetic structure. A UMAP projection of the learned embeddings for the 100 vertebrate species shows species of the same taxonomic class form well-separated clusters, and within Mammalia they further separate by order.

Methodology in Plain English

The core insight is that MSAs waste almost all of their input on gap tokens — positions where a species simply has nothing aligned. Standard MSA-based gLMs process every token, so their cost grows with both sequence length and species count, which forces them into tiny windows.

RAGenome instead retrieves only the aligned nucleotides and discards gaps entirely. Each retained token keeps its original alignment column as its positional index in rotary position embeddings (RoPE), so cross-species alignment information is preserved without any extra machinery. All retrieved sequences are concatenated into a single array and tagged with taxonomy embeddings so the model knows which species each token came from.

Two design elements keep the cost bounded:

• A fixed retrieval budget. The number of retrieved tokens is capped at a budget B, selected by walking the phylogenetic tree to sample from as many distinct lineages as possible. In the largest context setting, B = 80,000 tokens, chosen to fit within an NVIDIA H100 GPU.
• Two attention types. Intra-species local (sliding-window) attention with window size w = 512 handles within-species dependencies at linear cost; inter-species cross-attention layers let the query sequence attend over all retrieved species at once. The 15-layer model uses 12 local attention layers and 3 cross-attention layers.

Pretraining uses Homo sapiens as the query, retrieving 89 additional species (excluding species in the same clade as the query). Masked language modeling masks contiguous spans of 3 to 6 nucleotides, covering 15% of non-gap positions in each query sequence, with matching positions selected across species in the same clade.

Training proceeds in four stages that gradually increase both context length and genome diversity, to stabilize training:

1. L = 1,024 nucleotides restricted to the top 5% most conserved regions.
2. Same context length, extended to the top 40% most conserved regions.
3. Context extended to L = 4,096 nucleotides.
4. Full target context of L = 13,312 nucleotides, covering the entire genome.

The model is trained with AdamW at a learning rate of 10⁻⁴ with a cosine schedule of 4,000 warmup steps per stage. Evaluation uses two tasks: gene finding from BEND (framed as nucleotide-level 9-class classification with frozen embeddings fed to a two-layer CNN, scored by MCC) and zero-shot variant effect prediction on ClinVar (log-likelihood ratio between reference and alternate alleles, scored by AUROC).

Why This Matters

The paper argues implicitly that scaling genomic language models does not have to mean billion-parameter models trained for months on thousands of GPUs. By retrieving only informative tokens and explicitly conditioning on evolutionary information, RAGenome gets competitive results on two very different tasks at 168M parameters and 16,032 GPU hours — compared to Evo2's 7B parameters and estimated >1M GPU hours, or NT-MS's 2.5B parameters and 215,000 GPU hours. That reframes efficiency as a first-class design axis in genomic modeling.

Real-world applications (as framed by the tasks the paper evaluates):

• Clinical variant interpretation. Prioritizing pathogenic versus benign variants in ClinVar is directly relevant to diagnosing genetic disease; RAGenome's 0.87 AUROC puts it between the specialized MSA models and the large-scale gLMs.
• Genome annotation. Gene finding and splice-site identification are foundational for annotating newly sequenced genomes, and RAGenome's strength at donor and acceptor splice sites is directly applicable here.
• Comparative and evolutionary genomics. The model's taxonomy embeddings and cross-species attention make it a tool for reasoning about conservation across the 100-vertebrate alignment.
• Lower-resource genomics research. A 168M-parameter model that trains in a fraction of the compute broadens who can build and run genomic foundation models.

Industry relevance: The affiliations include Novo Nordisk A/S alongside University of Copenhagen and the Technical University of Denmark, and the work was supported by the Novo Nordisk Foundation and run on the Gefion supercomputer. Pharmaceutical and biotech groups interested in variant interpretation, target discovery, and efficient foundation-model training are the natural industrial audience.

Future Directions

1. Removing the alignment requirement. RAGenome currently requires a whole-genome alignment at inference and cannot be applied to sequences or species with no available alignment. The authors propose training a single model that works both with and without an alignment.
2. Generalizing beyond the human query. The model is trained with Homo sapiens as the query species and its behavior on other species remains untested. The architecture supports any query species and alignment, but applying it to a new one would likely require additional pretraining.
3. Scaling model size and data. Only a single 168M-parameter RAGenome model was trained; the authors leave scaling in model size and data to future work, and expect it to close the remaining gap to large-scale gLMs — particularly since gene finding MCC was still improving with longer context.
4. Combining with supervised post-training. The authors note that supervised post-training on functional tracks, as done in NTv3, is orthogonal to their approach and could in principle be applied to RAGenome.

Target Audience

This paper is most useful to machine learning researchers working on foundation models for biological sequences, computational biologists interested in genome annotation and variant interpretation, and practitioners who need strong genomic representations without access to massive compute budgets. It is also relevant to groups evaluating MSA-based versus scale-based pretraining strategies, and to readers tracking the broader question of whether explicit evolutionary modeling or raw scale is the better route to genomic understanding.

Authors’ abstract

The genome holds the blueprint that governs the biological properties of the cell. Consequently, advancing our knowledge of genomic function is crucial both for a broader understanding of biology and for continued biomedical advances. The success of large language models on natural language and protein sequences has motivated similar efforts on genomic data. However, standard genomic language models (gLMs) often require extremely large computational resources and still fall behind traditional methods on some downstream tasks. Recently, MSA-based pretraining has been proposed as an efficient alternative, but existing models are limited to short input contexts, restricting their use to short-range tasks, such as variant effect prediction. In this work, we present RAGenome, the first retrieval-based gLM that scales pretraining to longer contexts (100$\times$ longer than existing MSA-based gLMs), allowing it to capture both across-species evolutionary relationships and within-species longer-range interactions. Trained on whole-genome alignments from 100 vertebrates, RAGenome substantially improves the long-range capabilities of MSA-based gLMs, raising gene finding performance from 0.45 to 0.60, while remaining competitive on purely evolutionary-based tasks like prioritizing pathogenic variants. RAGenome provides competitive gLM performance at a fraction of the training cost, unifying evolutionary modeling and long-range capabilities within a single, flexible, scalable framework. Code is available at https://github.com/PanosAntoniadis/RAGenome.