Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models

TL;DR

The paper tackles the inefficiency of knowledge retrieval in Transformers by introducing conditional memory as a separate sparsity axis, instantiated with Engram—a scalable, constant-time $N$-gram memory with tokenizer compression, multi-head hashing, context-aware gating, and multi-branch integration. It formalizes the Sparsity Allocation problem, revealing a U-shaped scaling law that favors a hybrid allocation between MoE and Engram, and demonstrates this with Engram-27B and Engram-40B achieving superior performance under iso-parameters and iso-FLOPs. Mechanistic analyses (LogitLens and CKA) show Engram effectively deepens the network by offloading static pattern reconstruction from early layers, freeing attention for global context, which yields strong gains in long-context retrieval and reasoning tasks. Engram’s infrastructure-aware design enables deterministic memory addressing and host-memory offload with negligible overhead, illustrating practical scalability to hundreds of billions of parameters. Collectively, conditional memory emerges as a foundational primitive for next-generation sparse models, enabling robust knowledge lookup, improved reasoning, and efficient long-context processing.

Abstract

While Mixture-of-Experts (MoE) scales capacity via conditional computation, Transformers lack a native primitive for knowledge lookup, forcing them to inefficiently simulate retrieval through computation. To address this, we introduce conditional memory as a complementary sparsity axis, instantiated via Engram, a module that modernizes classic $N$-gram embedding for O(1) lookup. By formulating the Sparsity Allocation problem, we uncover a U-shaped scaling law that optimizes the trade-off between neural computation (MoE) and static memory (Engram). Guided by this law, we scale Engram to 27B parameters, achieving superior performance over a strictly iso-parameter and iso-FLOPs MoE baseline. Most notably, while the memory module is expected to aid knowledge retrieval (e.g., MMLU +3.4; CMMLU +4.0), we observe even larger gains in general reasoning (e.g., BBH +5.0; ARC-Challenge +3.7) and code/math domains~(HumanEval +3.0; MATH +2.4). Mechanistic analyses reveal that Engram relieves the backbone's early layers from static reconstruction, effectively deepening the network for complex reasoning. Furthermore, by delegating local dependencies to lookups, it frees up attention capacity for global context, substantially boosting long-context retrieval (e.g., Multi-Query NIAH: 84.2 to 97.0). Finally, Engram establishes infrastructure-aware efficiency: its deterministic addressing enables runtime prefetching from host memory, incurring negligible overhead. We envision conditional memory as an indispensable modeling primitive for next-generation sparse models.

Figures (8)

• Figure 1: The Engram Architecture. The module augments the backbone by retrieving static $N$-gram memory and fusing it with dynamic hidden states via context-aware gating. This module is applied only to specific layers to decouple memory from compute, leaving the standard input embedding and un-embedding module intact.
• Figure 2: System implementation of Engram. (a) Training Phase: The massive embedding tables are sharded across available GPUs. An All-to-All communication primitive is employed to retrieve active embedding rows across devices. (b) Inference Phase: Engram tables are offloaded to host memory. By exploiting the deterministic retrieval logic, the host asynchronously prefetches and transfers embeddings, overlapping communication with the on-device computation of preceding Transformer blocks.
• Figure 3: Sparsity allocation and Engram scaling. Left: Validation loss across allocation ratios $\rho$. Two compute budgets are shown ($2\text{e}20$ and $6\text{e}20$ FLOPs). Both regimes exhibit a U-shape, with hybrid allocation surpassing Pure MoE. Right: Scaling behavior in the infinite-memory regime. Validation loss exhibits a log-linear trend with respect to the number of embeddings.
• Figure 4: Analysis of representational alignment and convergence speed. (a) Layer-wise KL Divergence via LogitLens nostalgebraist2020logitlens. The consistently lower divergence in early layers indicates that Engram accelerates prediction convergence. (b-c) Similarity heatmap computed by CKA kornblith2019similarity. The distinct upward shift of the high-similarity diagonal demonstrates that =Engram's shallow layers are functionally equivalent to deeper layers of the MoE model, effectively increasing the model's depth.
• Figure 5: Architecture ablation results. We compare the 3B MoE baseline against Engram variations in two settings: (1) Layer Sensitivity (dark blue curve): Sweeping the insertion depth of a single Engram module confirms that early injection (Layer 2) is optimal, whereas efficacy degrades in deeper layers. (2) Component Ablation (Right Markers): Removing sub-modules from the reference configuration demonstrates the importance of multi-branch integration, tokenizer compression, and context-aware gating.