Source-linked AI summary
Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models
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
TL;DR
Transformers lack a native primitive for knowledge lookup and therefore simulate retrieval through computation. The paper introduces Engram as conditional memory built from scalable N-gram lookup, studies its allocation with MoE, and reports superior iso-parameter, iso-FLOPs performance across diverse tasks. The analyses associate these gains with deeper effective computation, greater attention capacity for global context, and low-overhead host-memory expansion.
Problem
Transformers lack a native knowledge-lookup primitive for static, local linguistic patterns and instead rely on computation to simulate memory retrieval.
Method
The paper introduces Engram, a conditional-memory module that modernizes N-gram embeddings with constant-time O(1) lookup and combines retrieved static embeddings with dynamic hidden states.
Results
Engram-27B outperforms a strictly iso-parameter and iso-FLOPs MoE-27B baseline across knowledge, reasoning, code, mathematics, and long-context evaluations.
Takeaways & Limitations
Conditional memory provides a complementary sparsity axis whose hybrid allocation with MoE improves diverse capabilities while enabling deterministic prefetching and host-memory parameter offloading.
Abstract
from arXiv · showhide
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.
1. Introduction
The paper proposes conditional memory as a complementary sparsity axis to MoE, using Engram for efficient lookup of static linguistic patterns. A U-shaped allocation law and 27B-scale experiments show gains across knowledge, reasoning, code, mathematics, and long-context tasks, with deterministic addressing supporting low-overhead host-memory offloading.
- Motivation and approach: Conditional memory complements MoE by retrieving static, local linguistic patterns through sparse lookups rather than dynamic computation.Engram modernizes N-gram embeddings with constant-time O(1) lookup and combines retrieved memory with dynamic hidden states through context-aware gating.
- Sparsity allocation: A U-shaped Sparsity Allocation law identifies how to divide a fixed parameter budget between MoE experts and Engram memory.Guided by this law, the authors scale Engram to 27B parameters and compare it with a strictly iso-parameter and iso-FLOPs MoE-27B baseline.
- Experimental results: MMLU +3.4, CMMLU +4.0, BBH +5.0, ARC-Challenge +3.7, HumanEval +3.0, and MATH +2.4 demonstrate improvements across knowledge, reasoning, code, and mathematics.The gains are larger in general reasoning and code/math domains than in the knowledge-intensive examples reported.
- Mechanistic analysis: Engram relieves early layers from reconstructing static knowledge and delegates local dependencies to lookups, increasing effective reasoning depth and freeing attention for global context.These mechanisms are associated with stronger long-context performance, including Multi-Query NIAH: 97.0 vs. 84.2.
- Infrastructure efficiency: Deterministic IDs enable runtime prefetching that overlaps communication with computation, while offloading a 100B-parameter table to host memory incurs overhead below 3%.This design addresses GPU memory constraints while allowing aggressive parameter expansion.
2. Architecture
Engram augments Transformers with conditional memory that deterministically retrieves static local-pattern embeddings, then adapts and integrates them with contextual computation. Its design separates storage from computation and supports distributed training, host-memory prefetching, and multi-branch backbones.
- Sparse Retrieval via Hashed N-grams: Engram extracts suffix N-grams and retrieves static embeddings through tokenizer compression and deterministic hashing.The retrieval phase maps local contexts to memory entries before fusion with the Transformer backbone.
- Context-aware Gating: Context-aware gating uses the current hidden state to modulate retrieved memory and suppress contextually inconsistent signals.Retrieved embeddings serve as keys and values, while the hidden state provides the query; contradictory memory is downweighted through the gate.
- Context-aware Gating: A short depthwise causal convolution expands the receptive field and adds nonlinearity before Engram contributes through a residual connection.The convolution uses kernel size 4, dilation equal to the maximum N-gram order, and SiLU activation.
- Integration with Multi-branch Architecture: In multi-branch backbones, Engram shares its embedding table and Value projection while using branch-specific Key projections for independent gating.The projections can be fused into a single dense FP8 matrix multiplication; experiments use four branches unless otherwise stated.
- System Implementation: Deterministic token-ID addressing decouples parameter storage from computation, enabling sharding during training and asynchronous host-memory prefetching during inference.Inference places Engram at selected layers so preceding computation can overlap communication; frequent patterns can also be cached across storage tiers.
3. Scaling Laws and Sparsity Allocation
The paper studies how sparse capacity should be divided between MoE computation and Engram memory under fixed budgets and during aggressive memory scaling. It finds that hybrid allocation is superior to pure MoE and that larger Engram tables provide predictable improvements without added per-token FLOPs.
- 3.1. Optimal Allocation Ratio Between MoE and Engram: The Sparsity Allocation problem distributes a fixed inactive-parameter budget between MoE experts and Engram embeddings.The allocation ratio ρ denotes the fraction assigned to MoE expert capacity, with the remainder assigned to Engram memory.
- Results and Analysis: Validation loss follows a U-shaped relationship with allocation ratio ρ across both tested compute regimes.The two regimes use 2 × 10^20 and 6 × 10^20 FLOPs, with constant sparsity ratio Ptot/Pact ≈ 10.
- Compute-matched formulation: Models are compared with matched total parameters and activated parameters, while Engram slot scaling increases storage without increasing per-token FLOPs.The protocol holds Ptot and Pact fixed within each compute regime and uses a constant number of retrieved slots per token.
- Results and Analysis: Reallocating roughly 20%–25% of sparse parameters from MoE to Engram yields the best performance, while pure MoE is suboptimal.Engram remains comparable to pure MoE even when MoE allocation falls to approximately 40%.
- Infinite Memory Regime: Validation loss improves consistently as memory slots increase, following a strict power law that is linear in log-space.Engram offers a predictable scaling knob, and it extracts more scaling potential from a memory budget than OverEncoding.
4. Large Scale Pre-training
Large-scale pre-training compares dense, MoE, and Engram models under controlled token, activation, parameter, and compute settings. Engram-27B improves over the matched MoE-27B baseline across knowledge, reasoning, code, and mathematics, while Engram-40B generally improves further but remains affected by under-training.
- Experimental setup: Table 1 compares models trained for 262B tokens with matched activated parameters of 3.8B.Engram-27B matches MoE-27B in total parameters by reallocating routed-expert capacity to a 5.7B-parameter memory module.
- Overall results: All three sparse variants significantly outperform the iso-FLOPs Dense-4B baseline across all benchmarks.The sparse variants are MoE-27B, Engram-27B, and Engram-40B.
- Overall results: Engram-27B consistently improves over the iso-parameter and iso-FLOPs MoE-27B baseline.The gains extend beyond knowledge-intensive tasks to general reasoning, code, and mathematical reasoning.
- Benchmark results: BBH: +5.0, ARC-Challenge: +3.7, and DROP: +3.3 over the MoE-27B baseline.Knowledge results include MMLU: +3.0, MMLU-Pro: +1.8, and CMMLU: +4.0.
- Benchmark results: HumanEval: +3.0, MBPP: +1.6, GSM8K: +2.2, and MATH: +2.4 over the MoE-27B baseline.These results support improved representation efficiency from adding dedicated knowledge lookup to conditional computation.
- Scaling results: Engram-40B further reduces pre-training loss and improves performance across most benchmarks, but does not strictly dominate Engram-27B on every task.The authors attribute this pattern to under-training and report that the loss gap continues widening near the end of training.
5. Long Context Training
Long-context experiments evaluate Engram under matched training procedures, loss, and compute against MoE-27B. Engram improves retrieval performance under controlled comparisons and can remain competitive even with substantially less pre-training compute.
- Training and evaluation: Long-context evaluation uses LongPPL and RULER after a 32768-token extension stage trained for 5,000 steps.The extension uses 30B tokens of high-quality long-context data with YaRN scaling.
- Training and evaluation: The study compares final MoE-27B and Engram-27B checkpoints with intermediate Engram checkpoints at 41k and 46k steps.All variants undergo the same context-extension training protocol.
- Controlled analysis: Engram long-context performance improves monotonically with pre-training progression from 41k to 50k steps.The analysis controls for base-model capability when assessing this progression.
- Architectural comparison: Under both iso-pretraining-loss and iso-pretraining-FLOPs settings, Engram-27B substantially outperforms the baseline across all metrics.The controlled analysis separates architectural efficiency from differences in base-model capability.
- Architectural comparison: With approximately 82% of pre-training compute, Engram-27B matches the baseline on LongPPL and surpasses it on RULER.The 41k-step Engram model is compared with the 50k-step MoE-27B baseline.
6. Analysis
The analysis finds that Engram changes how computation is allocated across layers: explicit lookups accelerate static feature formation, deepen effective representations, and improve long-context and ablation performance.
- Effective Depth: Engram reaches prediction-ready representations faster than the MoE baseline, with the largest KL-divergence gap in early blocks.The steeper Engram curves indicate faster feature composition after explicit knowledge access.
- Representational Alignment: CKA shows that Engram representations at earlier layers align with deeper MoE representations, supporting an effective-depth increase.Engram-27B layer 5 aligns most closely with approximately layer 12 of the MoE baseline.
- Layer Sensitivity: Layer 2 achieves the best single-layer insertion result, with Val Loss = 1.770; performance degrades when memory is inserted deeper.The result indicates that one attention round provides contextualized gating while preserving early local aggregation replacement.
- Layer Sensitivity: Splitting the fixed 1.6B memory across Layers 2 and 6 improves validation loss to 1.768, combining early intervention with later contextual gating.The layered design also improves utilization of the memory hierarchy.
- Structural Ablation: Multi-branch fusion, context-aware gating, and tokenizer compression produce the largest ablation regressions, while depthwise convolution has only a marginal effect.Four-grams are slightly suboptimal at the fixed 1.6B budget, although higher-order n-grams may help at larger scales.
- Infrastructure Efficiency: Offloading a 100B-parameter embedding table raises throughput cost by at most 2.8% on the 8B backbone, with communication scaling by activated slots.The experiment retrieves from host memory over PCIe, making the overhead a conservative baseline for locality-aware designs.
7. Related Work
Related work spans classical and neural N-gram embeddings, sparse expert computation, memory-augmented networks, categorical embeddings, and studies of Transformer knowledge storage. Engram is positioned as a fair-comparison, layer-integrated approach that combines scalable lookup with conditional computation.
- N-gram Modeling and Embedding Scaling: Classical N-gram models use local history and efficient lookups, while modern embedding-scaling methods expand capacity through massive tables or compositional N-gram structures.The related approaches include Per-Layer Embeddings, STEM, L3, DeepEmbed, and N-Grammer.
- Positioning Engram: Engram differs from prior N-gram augmentation by injecting memory beyond the input layer, overlapping retrieval with computation and exploiting Zipfian locality.The design is presented as enabling massive parameter scaling with negligible overhead under fairer efficiency considerations.
- High-Cardinality Categorical Embeddings: Large-scale recommender systems address embedding trade-offs using multi-hash, compositional, and frequency-aware representations for high-cardinality categorical features.These methods target vocabularies ranging from millions to billions of IDs.
- Mixture-of-Experts: MoE scales model capacity by conditionally activating sparse expert subsets while maintaining relatively stable inference computation.The section situates Engram as complementary to this conditional-computation paradigm.
- Memory Networks: Memory-augmented networks expand capacity through parametric or non-parametric memory, including sparse key-value stores integrated into model layers.Related parametric methods include PKM, PEER, Selfmem, Memory+, and Ultra-Mem.
- Mechanisms of Knowledge Storage: Transformer knowledge-storage studies commonly interpret FFNs as key-value memories, with early layers detecting patterns and later layers projecting information into the residual stream.This provides a conceptual background for separating static memory from dynamic computation.
8. Conclusion
The paper introduces conditional memory as a complement to MoE conditional computation and instantiates it with Engram’s scalable N-gram lookup. A hybrid allocation of computation and memory outperforms pure MoE across domains, while mechanistic and systems analyses explain its broader benefits.
- Contribution: Conditional memory addresses the inefficiency of simulating knowledge retrieval through dynamic computation by providing scalable O(1) lookups for static patterns.Engram modernizes classic N-gram embeddings as the paper’s conditional-memory module.
- Scaling and Results: The Sparsity Allocation problem yields a U-shaped scaling law in which hybrid MoE–Engram capacity allocation strictly outperforms pure MoE baselines.The law guides scaling Engram to 27B parameters and supports superior performance across diverse domains.
- Mechanism and Efficiency: Engram relieves early layers from static reconstruction, frees attention for global context and complex reasoning, and improves long-context capabilities.The paper reports gains on LongPPL and RULER as evidence for this mechanistic account.
- Mechanism and Efficiency: Deterministic addressing decouples storage from computation, enabling massive embedding tables to be offloaded to host memory with negligible inference overhead.This motivates infrastructure-aware efficiency as a design principle for sparse models.
A. Detailed Model Architecture and Hyper Parameters
The detailed architecture section provides a table of model architecture information and training hyperparameters.
- Table 5 contains the detailed model architecture information and training hyperparameters.
B. Full Benchmark Curves
Figure 8 presents the benchmark curve for the last 10k pre-training steps.
- Figure 8 shows the benchmark curve over the last 10k pre-training steps.
C. Case Study of Tokenizer Compression
Table 6 presents the top-5 merged tokens by tokenizer compression and reports a 23.43% overall compression ratio for the 128k tokenizer.
- 23.43% is the overall compression ratio for the 128k tokenizer.
- Table 6 lists the top-5 merged tokens by tokenizer compression.