Does Engram Do Memory Retrieval in Autoregressive Image Generation?

Technical Summary: Does Engram Do Memory Retrieval in Autoregressive Image Generation?

Problem Statement
The Engram module, a hash-keyed, $O(1)$ associative memory injected into Transformer layers, was recently proposed for Large Language Models (LLMs) to provide a content-addressed shortcut for recurring local token patterns. The prevailing hypothesis is that it retrieves specific content corrections based on the surface form of recent context. This paper investigates whether this mechanism transfers effectively to autoregressive (AR) image generation. Specifically, it asks: (Q1) Does injecting Engram into a class-conditional AR generator improve sample quality (FID) at matched parameter counts compared to a pure AR baseline? (Q2) If performance changes, is the driver genuine content-addressed retrieval from the memory table, or is it an architectural side-pathway effect (e.g., gating, residual streams, or regularization)?

Methodology
The authors adapt the Engram module from 1D language to 2D vision and integrate it into a strong class-conditional AR generator trained on ImageNet 256×256.

• Architecture Adaptation: The module utilizes 2D spatial n-gram hashing over a causal grid, gated fusion, and KV-cache-compatible incremental inference. Unlike the NLP version, the vision adaptation employs spatial neighbors (e.g., left, top, diagonal) to define n-gram contexts.
• Tokenization Prerequisite: Before adaptation, the authors verified that visual tokens possess recurrent local n-gram structures. Using DINO-stratified Jaccard analysis, they found that the AliTok tokenizer exhibits consistent local patch patterns correlated with semantic similarity, whereas MaskGIT-VQ does not. Consequently, AliTok was selected for all experiments.
• Experimental Design:
  • Parameter Sweep: A controlled sweep of the backbone-to-total parameter ratio ($\rho \in [0.17, 0.90]$) was conducted, trading backbone width for memory table capacity.
  • Gate-Clamp Analysis: The learned gating mechanism was replaced with constant values ($g \in \{0.00, 0.10, \dots, 1.00\}$) to determine the optimal usage of the memory pathway.
  • Donor Probe: The hash inputs were swapped with tokens from matched (same-class), adversarial (different-class), or random exemplars to test if the model retrieves semantically relevant content.
  • Frozen-Noise Control: A model was trained from scratch with the entire memory table frozen to $N(0, 1)$ noise, isolating the architectural contribution from the learned content.

Key Results
The empirical findings challenge the "content-addressed retrieval" hypothesis for AR image generation:

1. No Quality Gain at Matched Parameters: Across all backbone-to-memory ratios, every Engram-augmented variant achieved a worse Fréchet Inception Distance (FID) than the pure AR baseline (2.02). While the module saves backbone FLOPs by offloading local pattern reconstruction, it does not inherently improve sample quality. Adding 2D n-gram banks (1×3, 3×1, 2×2) yielded negligible improvements (<0.01 FID) over the 1D variant.
2. Pathway Necessity vs. Content Over-mixing: Disabling the Engram pathway entirely ($g=0$) caused catastrophic generation failure (FID-10K jumped from 4.89 to 115). However, the learned gate was found to be "over-mixing"; a tiny constant gate ($g=0.10$ or $0.25$) matched or outperformed the learned gate, while higher gates degraded quality. This suggests the model benefits from the pathway being barely used, inconsistent with a mechanism that heavily relies on retrieving specific content.
3. Hash Address Does Not Select Content: In the donor probe, swapping hash inputs with matched, adversarial, or random same-class exemplars produced statistically indistinguishable next-token distributions. Conversely, collapsing the hash buckets or randomizing the table degraded performance by two to three orders of magnitude. This indicates the model relies on the diversity and learned structure of the table, not on the semantic match between the query and the retrieved entry.
4. Architectural Dominance: Training a model from scratch with the memory table frozen to $N(0, 1)$ noise resulted in only a marginal FID increase ($\Delta FID = 0.10$) compared to a fully learned table, while actually raising the Inception Score (259 to 285). This confirms that the bulk of the benefit comes from the gated, hashed side-pathway (additional parameters and a parallel residual stream) rather than the specific contents of the memory table.

Significance and Claims
The paper concludes that in the context of autoregressive image generation, the Engram module functions primarily as a gated architectural side-pathway rather than a content-addressed retriever.

• Mechanism: The module acts as a hash-keyed residual stream where the learned table contributes only a minor distributional refinement. The primary utility is the architectural addition of a parallel stream and gating mechanism, which serves as a regularizer or prior on local image-token statistics.
• Implication for Vision: The "content-addressed" interpretation that motivated the original Engram design does not hold for visual token streams produced by current tokenizers. The observed gains (or lack thereof) are dominated by the pathway itself, not the retrieval of recurring semantic patterns.
• Scope: The authors note their study is limited to dense AR generators at moderate scales and does not explore extreme low-data regimes, long-tail classes, or higher resolutions where genuine memory retrieval might play a larger role.

In summary, the Engram in AR image generation is best described as a necessary but weakly content-addressed architectural prior, where the learned memory table offers only a small refinement over the structural benefits of the side-pathway.