💬 NLP

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

This paper introduces Engram, a scalable conditional memory module that modernizes $N$ -gram embeddings for O(1) lookup, revealing a U-shaped sparsity allocation law that optimizes the trade-off between neural computation and static memory to significantly boost reasoning, code, and long-context retrieval performance in large language models.

Original authors: Xin Cheng, Rui Tian, Wangding Zeng, Damai Dai, Qinyu Chen, Bingxuan Wang, Zhenda Xie, Kezhao Huang, Xingkai Yu, Chengqi Deng, Shangyan Zhou, Chenggang Zhao, Zhewen Hao, Yukun Li, Han Zhang, Zhengyan Z

Published 2026-07-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Xin Cheng, Rui Tian, Wangding Zeng, Damai Dai, Qinyu Chen, Bingxuan Wang, Zhenda Xie, Kezhao Huang, Xingkai Yu, Chengqi Deng, Shangyan Zhou, Chenggang Zhao, Zhewen Hao, Yukun Li, Han Zhang, Zhengyan Zhang, Yixu Wei, M. Y Xu, Huishuai Zhang, Dongyan Zhao, Wenfeng Liang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you're trying to solve a giant, complex puzzle. For a long time, the smartest puzzle-solvers (AI models) have been built like massive factories. They have thousands of specialized workers (called "Mixture-of-Experts" or MoE) who jump in only when needed to figure out tricky logic or new ideas. This is great for thinking, but it's a bit clumsy when it comes to remembering simple, boring facts.

Think of it like this: If you asked a brilliant mathematician, "What is the capital of France?" they wouldn't just say "Paris." They would have to run a full mental simulation, re-deriving the answer from scratch every single time, burning through a lot of brainpower just to recall something they've known forever. It's like using a supercomputer to look up a phone number in your own head.

The Big Discovery: A "Cheat Sheet" for the Brain
The researchers at DeepSeek-AI and Peking University asked: What if we gave these models a dedicated "cheat sheet" for the stuff they just need to remember?

They introduced a new tool called Engram. Instead of forcing the AI to "think" its way through every single word, Engram acts like a super-fast, O(1) lookup table (that means it takes the same tiny amount of time whether you're looking up one word or a million). It's based on an old-school idea called N-grams (looking at groups of words together), but they upgraded it with modern tech to make it work perfectly inside a giant AI.

The "U-Shaped" Secret
The team discovered a fascinating rule about how to build these models, which they call a U-shaped scaling law.

Imagine you have a fixed budget of "brain power" (computing steps). You can spend it all on the "thinking workers" (MoE), or you can spend it all on the "memory cheat sheet" (Engram).

If you spend 100% on thinking workers, the model is good at logic but bad at remembering facts.
If you spend 100% on the cheat sheet, it remembers facts but can't reason well.
The Sweet Spot: The paper found that the best performance happens when you split the budget. You need both. Specifically, they found that taking about 20–25% of the "spare" memory budget away from the thinking workers and giving it to the Engram memory module makes the whole system smarter. It's like realizing that a genius needs a good library, not just a fast brain.

What Actually Happened? (The Proof)
They built a model called Engram-27B and compared it to a standard "thinking-only" model of the exact same size and speed. The results were surprisingly strong:

Knowledge: It got better at trivia and facts (like scoring +3.4 higher on MMLU and +4.0 on CMMLU).
Reasoning: Even more surprisingly, it got much better at general reasoning and logic puzzles (scoring +5.0 higher on BBH and +3.7 on ARC-Challenge).
Math & Code: It also improved at coding and math (like +3.0 on HumanEval).

The paper suggests this happens because the Engram module handles the "boring" stuff (like remembering that "Alexander the Great" is a specific name) so the main brain doesn't have to waste time reconstructing it. This frees up the deep layers of the AI to focus on complex, deep thinking. It's like having a secretary who handles all the scheduling so the CEO can focus on strategy.

The "Long Memory" Trick
One of the coolest side effects is how well these models handle long stories. Because the Engram module grabs local details instantly, the main AI has more "attention" left over to remember the whole story from start to finish. In tests, the Engram model could find a specific needle in a haystack of text 97.0% of the time, while the standard model only managed 84.2%.

The Hardware Magic
Finally, the paper argues that this isn't just a software trick; it's a hardware-friendly one. Because the "cheat sheet" uses fixed addresses (deterministic IDs), the computer can start fetching the next piece of memory while it's still doing the math for the current piece. This means they can store a massive 100 billion parameter table on a regular computer's hard drive (host memory) and still get results almost as fast as if it were on the super-fast GPU. The speed penalty was negligible, less than 3%.

What They Explicitly Say It's NOT
The paper is very clear about what this is not:

It is not just a bigger vocabulary. They tried just making the vocabulary bigger (OverEncoding), and it didn't work as well as Engram.
It is not a replacement for the "thinking" part. If you remove the thinking workers (MoE) entirely, the model fails at reasoning. The memory is a helper, not a replacement.
It is not a magic fix for everything. The paper notes that while it helps with facts and reasoning, the "backbone" (the main AI) is still the one doing the heavy lifting for things like reading comprehension, which relies more on context than static memory.

How Sure Are They?
The authors are very confident in their measurements. They didn't just simulate this; they trained real models on 262 billion tokens of data and tested them on real-world benchmarks. They showed that the "U-shaped" rule holds true across different sizes and that the performance gains are real, measurable, and repeatable. They even visualized the "gates" inside the model to show that it really is turning on the memory specifically for things like names and phrases, just as they designed.

In short, the paper suggests that the future of AI isn't just about making bigger brains; it's about giving those brains a better, faster, and smarter way to look things up.

Technical Summary: Conditional Memory via Scalable Lookup (Engram)

1. Problem Statement

Current Large Language Models (LLMs), particularly those utilizing Mixture-of-Experts (MoE) architectures, rely on conditional computation to scale capacity. However, standard Transformers lack a native primitive for knowledge lookup. Consequently, models must inefficiently simulate retrieval through dynamic computation.

The authors argue that linguistic signals exhibit intrinsic heterogeneity:

Compositional Reasoning: Requires deep, dynamic computation.
Knowledge Retrieval: A substantial portion of text (e.g., named entities, formulaic patterns) is local, static, and highly stereotyped.

Standard Transformers force the model to consume multiple early layers of attention and feed-forward networks (FFNs) to reconstruct static lookups (e.g., resolving "Diana, Princess of Wales") that could theoretically be retrieved via a simple table lookup. This wastes valuable sequential depth on trivial operations, preventing the model from allocating resources to higher-level reasoning.

2. Methodology: The Engram Module

To address this, the paper introduces Conditional Memory as a complementary sparsity axis to conditional computation. This is instantiated via Engram, a module that modernizes classic $N$ -gram embeddings to enable $O(1)$ constant-time lookups.

2.1. Architecture

Engram augments the Transformer backbone by structurally separating static pattern storage from dynamic computation. It operates in two phases:

Sparse Retrieval via Hashed $N$ -grams:
- Tokenizer Compression: A surjective function maps raw token IDs to canonical identifiers (normalizing case, punctuation, etc.) to maximize semantic density and reduce vocabulary size.
- Multi-Head Hashing: To handle the combinatorial space of $N$ -grams, the module uses $K$ distinct hash heads per $N$ -gram order. These map compressed contexts to indices in embedding tables via deterministic functions, retrieving static vectors.
- Concatenation: Retrieved embeddings from various $N$ -gram orders and hash heads are concatenated to form the memory vector.
Context-Aware Gating & Fusion:
- Since static memories lack context, a gating mechanism modulates the retrieved memory using the current hidden state (aggregated global context) as a Query.
- A scalar gate $\alpha_t$ is computed via a dot product between the normalized hidden state and the memory Key. If the memory contradicts the context, the gate suppresses the memory (noise reduction).
- A lightweight depthwise causal convolution expands the receptive field before the gated memory is fused into the backbone via a residual connection.

2.2. System Efficiency & Infrastructure Co-Design

Unlike MoE, which requires dynamic routing and hidden state communication, Engram uses deterministic addressing based solely on input token IDs.

Prefetching: Indices are known before the forward pass, allowing asynchronous prefetching of embeddings from host memory (CPU/DRAM) over PCIe.
Overlap: This enables communication to overlap with the computation of preceding Transformer blocks, incurring negligible overhead.
Offloading: Massive embedding tables (e.g., 100B parameters) can be offloaded to host memory, bypassing GPU memory constraints.

2.3. Sparsity Allocation

The authors formulate the Sparsity Allocation problem: given a fixed parameter budget, how should capacity be split between MoE experts (dynamic computation) and Engram memory (static lookup)?

They define an allocation ratio $\rho$ representing the fraction of inactive parameters assigned to MoE.
Experiments reveal a U-shaped scaling law: Pure MoE ( $\rho=100\%$ ) and pure Engram ( $\rho=0\%$ ) are suboptimal. The best performance is achieved with a hybrid allocation (e.g., $\rho \approx 75\%–80\%$ ), where roughly 20–25% of the sparse budget is reallocated to Engram.

3. Key Results

3.1. Pre-training Performance

The authors trained four models on 262B tokens, strictly matching activated parameters (3.8B) and training FLOPs:

Dense-4B: Baseline.
MoE-27B: Standard MoE baseline.
Engram-27B: Iso-parameter to MoE-27B, with 5.7B parameters reallocated to Engram (reducing experts from 72 to 55).
Engram-40B: Scales Engram memory to 18.5B parameters (Total 39.5B).

Findings:

Superiority over MoE: Engram-27B outperformed the iso-parameter MoE-27B baseline across diverse domains.
Reasoning Gains: Contrary to the intuition that memory only helps knowledge retrieval, Engram showed significant gains in general reasoning (BBH +5.0, ARC-Challenge +3.7) and code/math (HumanEval +3.0, MATH +2.4).
Knowledge Gains: Improved performance on knowledge-intensive tasks (MMLU +3.0, CMMLU +4.0).
Scaling: Engram-40B further reduced loss, suggesting the memory axis scales effectively without additional compute.

3.2. Long-Context Capabilities

By offloading local dependencies to lookups, Engram frees attention capacity for global context.

RULER Benchmark: Engram-27B achieved 97.0% accuracy on Multi-Query Needle-in-a-Haystack (NIAH) compared to 84.2% for the MoE baseline.
Efficiency: Engram-27B matched the baseline's LongPPL performance with only 82% of the pre-training FLOPs.

3.3. Mechanistic Analysis

Effective Depth: Using LogitLens and CKA (Centered Kernel Alignment), the authors found that Engram's early layers produce representations semantically equivalent to much deeper layers in the MoE baseline. This suggests Engram "deepens" the network by bypassing the need for early layers to reconstruct static features.
Functional Dichotomy: Ablation studies showed that factual knowledge benchmarks collapsed (retaining only 29–44% performance) when Engram was disabled, while reading comprehension remained robust (81–93%), confirming Engram acts as the primary repository for parametric knowledge.

3.4. System Efficiency

Throughput: Offloading a 100B-parameter Engram table to host memory resulted in a negligible throughput penalty of <3% (2.8% on an 8B backbone), validating the prefetch-and-overlap strategy.

4. Significance and Claims

The paper claims that Conditional Memory is an indispensable modeling primitive for next-generation sparse models. Its significance lies in three main areas:

Structural Complementarity: It establishes that static lookup and dynamic computation are not mutually exclusive but structurally complementary. Optimizing the allocation between them yields a U-shaped scaling law that outperforms pure MoE.
Efficiency via Decoupling: By decoupling parameter storage from computation through deterministic addressing, Engram allows models to scale to massive parameter counts (e.g., 40B+) without being constrained by GPU memory or compute costs, effectively bypassing hardware bottlenecks.
Deepening the Network: The mechanism analysis suggests that providing explicit knowledge lookup relieves the backbone's early layers from static reconstruction tasks. This effectively increases the network's depth available for complex reasoning, explaining the unexpected gains in reasoning and code tasks.

The authors position Engram not merely as a retrieval augmentation but as a fundamental architectural shift that modernizes $N$ -gram concepts for the era of large-scale sparse models, offering a new axis of sparsity that complements the prevailing MoE paradigm.