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

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.

• The paper's main contribution is the introduction of conditional memory via a scalable, hash-based Engram module that offloads static pattern retrieval from standard computation.
• It demonstrates a hybrid MoE-Engram architecture achieving superior performance on knowledge tasks, reasoning benchmarks, and long-context retrieval through optimal sparsity allocation.
• The study highlights that decoupling memory from computation enables efficient scaling, improved hardware utilization, and enhanced interpretability in large language models.

Conditional Memory via Scalable Lookup: A New Axis of Sparsity for LLMs

Introduction and Motivation

This paper introduces conditional memory as an axis of architectural sparsity, orthogonal to the established paradigm of conditional computation via Mixture-of-Experts (MoE). While MoE mechanisms allocate parameters to dynamic computational paths on a per-token basis, standard Transformers lack primitives for efficient knowledge retrieval. In existing architectures, the resolution of static, local linguistic patterns (e.g., named entities, multi-token formulae) is achieved via repeated application of attention and FFN layers—transforming a trivial $O(1)$ lookup task into a sequence of expensive computational operations.

The authors propose Engram, a module for scalable, hash-based $N$-gram memory, decoupling static pattern retrieval from neural computation. By formulating the Sparsity Allocation problem—how to optimally allocate a fixed parameter budget between MoE capacity and memory—the work demonstrates both theoretically and empirically that a hybrid design outperforms pure compute-centric models.

Engram Architecture

The Engram module injects a scalable, embedding-based, hash-lookup memory into selected layers of the Transformer backbone. Each token position retrieves and fuses $N$-gram static memory (computed via hashed and compressed suffixes of length $N$) into the dynamic hidden state by a context-sensitive gating mechanism.

Figure 1: The Engram Architecture augments the Transformer by fusing static $N$-gram memory with dynamic hidden states via context-aware gating in selected layers.

Retrieval is performed via multiple-hash heads over compressed and normalized token ids, using deterministic functions as memory indices. Fusion with the neural state is handled with a single-gate scalar $\alpha_t$ per position—computed by projecting both the retrieved embedding and the dynamic hidden state, then applying RMS normalization and a sigmoid activation to produce the gating value. This result is processed by a lightweight causal convolution, with strong residual integration into the backbone.

Engram is agnostic to the backbone topology but optimized for multi-branch architectures (e.g., manifold-constrained hyper-connections), enabling efficient parameter sharing and branch-specific gating.

System Implementation and Decoupling Memory from Compute

A critical design is the decoupling of memory storage from computational resources. Engram's deterministic addressing allows memory tables to be sharded across devices for training (via All-to-All collectives) and offloaded to host DRAM for inference; with embedding indices pre-computable, embedding fetch can be overlapped with neural computation, masking communication latency.

Figure 2: System implementation of Engram, leveraging GPU sharding for training and asynchronous host prefetch for inference, exploiting static retrieval logic to overlap communication and compute.

The storage requirements of massive embedding tables are managed by co-designing memory hierarchy—frequent patterns are cached at fast tiers (HBM), while rare patterns are stored in slower host or disk memory, exploiting natural Zipfian frequency distributions.

Scaling Laws and Sparsity Allocation

The Sparsity Allocation problem formalizes the tradeoff between MoE expert capacity and Engram memory under fixed parameter and FLOPs budgets. The empirical result is a robust U-shaped validation loss curve: allocating 20–25% of sparse parameters to Engram (i.e., using only $\sim$75–80% for MoE) yields the best performance, with pure MoE being strictly suboptimal.

Figure 3: Sparsity allocation and Engram scaling: Validation loss is minimized at an intermediate allocation ratio between MoE and Engram memory – not at the extremes.

When the Engram memory is aggressively scaled (beyond what is possible with computation alone), further gains are observed with loss decreasing log-linearly in the number of embedding slots, indicating that Engram provides a unique, scalable axis of parameter expansion.

Empirical Results and Performance Analysis

The hybrid Engram-MoE architecture is scaled to 27B total parameters with a fixed computation budget. Compared to a parameter- and FLOPs-matched MoE-27B baseline, Engram-27B achieves strictly improved results in every tested domain:

• Knowledge tasks: MMLU $+3.4$, CMMLU $+4.0$;
• General reasoning: BBH $+5.0$, ARC-Challenge $+3.7$;
• Code/math: HumanEval $+3.0$, MATH $+2.4$;

Gains are not limited to knowledge-centric tasks; improvements are equally evident in domains such as reasoning and code, contradicting initial expectations that static memory would only benefit retrieval-heavy benchmarks.

Mechanistic and Internal Analysis

Detailed representational analysis elucidates why Engram improves generalization:

• Prediction convergence: Layer-wise KL divergence (LogitLens) indicates that early layers in Engram models reach prediction-readiness significantly faster, relieving the backbone from slow static reconstruction.

  Figure 4: Analysis of representational alignment: Engram leads to faster convergence (lower divergence) in early layers and deeper functional equivalence via CKA.

• Effective depth increase: Centered Kernel Alignment (CKA) confirms an upward shift in representational similarity between Engram and MoE models, showing that early Engram layers functionally correspond to deeper MoE layers.

Module ablations reveal that multi-branch integration, tokenizer compression, and context-aware gating are critical—removing any leads to substantial degradation. Placement sensitivity experiments demonstrate that early-layer insertion (Layer 2) is optimal, maximizing the offloading of static pattern modeling while retaining strong global context.

Figure 5: Architecture ablations show that early Engram injection and specific module components critically determine performance.

Task Sensitivity and Long-Context Benefits

Inference-time ablations establish that Engram is the primary carrier of factual knowledge: suppressing its output causes catastrophic drops in knowledge recall, but reading comprehension tasks remain largely preserved (attesting to the backbone's retained reasoning ability).

Figure 6: Factual knowledge tasks are highly dependent on Engram, while reading comprehension remains robust to its ablation.

Engram dramatically improves long-context retrieval benchmarks, particularly in multi-key or variable tracking scenarios—multi-query NIAH accuracy increases from 84.2 to 97.0, demonstrating superior scaling of context management.

Gating Mechanism and Interpretability

Visualizations of the Engram gating mechanism confirm hypothesized behavior: strong activation occurs precisely at the positions concluding static, multi-token patterns (e.g., named entities, idioms) across languages.

Figure 7: The gating heatmap reveals Engram selectively activates on detected $N$-gram patterns, consistent for structured phrases and entities.

Implications and Future Outlook

The evidence positions conditional memory—instantiated by scalable $N$-gram lookup and Engram-style modules—as a fundamental modeling primitive for future LLM sparsity. By leveraging deterministic retrieval, hybrid allocation, and memory-compute decoupling, Engram architectures unlock new parameter scaling regimes, more efficient hardware utilization, and qualitative improvements in reasoning and long-range context management.

Theoretically, these results strongly support a modular interface split between static knowledge retrieval and dynamic reasoning in large-scale models. Practically, such an approach may enable models to expand well beyond current GPU memory constraints with negligible inference overhead—memory can be made arbitrarily large, and communication bandwidth is the only binding constraint.

Future development should investigate dynamic memory updates, non-parametric retrieval integration, and optimal allocation as model and hardware scales increase. Memory modules capable of learning, editing, or augmenting their key spaces will broaden applicability to ever-changing knowledge domains.

Conclusion

This work establishes conditional memory via scalable lookup as a distinct and essential axis of sparse modeling for large-scale LLMs. The proposed Engram module achieves strict improvements across knowledge-intensive, reasoning, and code/math benchmarks under equal computation budgets. Mechanistic analysis demonstrates Engram's role in offloading static pattern reconstruction, functionally deepening the neural backbone and freeing up capacity for higher-level reasoning and long-context management. The hybrid memory-compute approach validated here is poised to shape next-generation LLMs, enabling efficient, interpretable, and hardware-friendly scaling far beyond current architectures (2601.07372).