Beyond Conditional Computation: Retrieval-Augmented Genomic Foundation Models with Gengram

TL;DR

Gengram introduces a hash-based motif memory module that explicitly stores multi-base motifs and integrates via a local-window, gated memory pathway into Transformer-based genomic foundation models. By combining multi-scale N-gram memory with deduplicated local-context retrieval and a gated residual, Gengram achieves up to 14% gains on motif-dominated tasks while preserving training efficiency and enabling stable load balancing in sparse MoE settings. The approach demonstrates robust generalizability across architectures and attention mechanisms and reveals mechanistic, biologically meaningful structures in memory representations, including reverse-complement symmetry and motif-aligned residual writes. The work provides a scalable architectural advancement for GFMs that improves both empirical performance and interpretability, with potential to extend long-range sequence modeling and multi-omics applications.

Abstract

Current genomic foundation models (GFMs) rely on extensive neural computation to implicitly approximate conserved biological motifs from single-nucleotide inputs. We propose Gengram, a conditional memory module that introduces an explicit and highly efficient lookup primitive for multi-base motifs via a genomic-specific hashing scheme, establishing genomic "syntax". Integrated into the backbone of state-of-the-art GFMs, Gengram achieves substantial gains (up to 14%) across several functional genomics tasks. The module demonstrates robust architectural generalization, while further inspection of Gengram's latent space reveals the emergence of meaningful representations that align closely with fundamental biological knowledge. By establishing structured motif memory as a modeling primitive, Gengram simultaneously boosts empirical performance and mechanistic interpretability, providing a scalable and biology-aligned pathway for the next generation of GFMs. The code is available at https://github.com/zhejianglab/Genos, and the model checkpoint is available at https://huggingface.co/ZhejiangLab/Gengram.

Figures (9)

• Figure 1: Overview of the Gengram architecture and evaluation. (Left) Gengram is integrated into a Transformer block via a residual connection and applied before the attention module. It maintains a motif memory implemented as a hash table with static keys derived from all possible k-mers ( k=1,2,3,4,5,6) and learnable embedding values. For a given local window, k-mers are mapped to memory entries, aggregated within each k-mer level, and concatenated across levels. The resulting representation is then passed through a gate-controlled module to modulate the hidden states. (Right) Evaluation and analysis of Gengram. We conduct extensive benchmarks across diverse genomic tasks and Transformer architectures, demonstrating strong compatibility with different attention mechanisms (e.g., MHA, GQA, MLA). Biological inspection further shows that Gengram’s activations align with functional genomic regions.
• Figure 2: Parameter Selection Experiments.Top: Validation loss when inserting Gengram into individual Transformer layers in a 1.2B-parameter (0.3B activated) model trained on 50B tokens. The red curve denotes the baseline without Gengram. Bottom: Validation loss under different Gengram window sizes.
• Figure 3: Performance Comparison of Baseline and Gengram Models at 8k and 32k Context Lengths
• Figure 4: Validation-Loss Comparison:Validation-loss bar chart under MoE and Dense architectures across GQA/MLA/MHA.
• Figure 5: Cross-Sparsity Load-Balancing:Load-balancing loss curves with and without the Gengram module under Top-2 / 128, 64, and 32 experts, showing consistent stabilization across sparsity settings.