MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers

Abstract: Understanding how transformer components operate in LLMs is important, as it is at the core of recent technological advances in artificial intelligence. In this work, we revisit the challenges associated with interpretability of feed-forward modules (FFNs) and propose MemoryLLM, which aims to decouple FFNs from self-attention and enables us to study the decoupled FFNs as context-free token-wise neural retrieval memory. In detail, we investigate how input tokens access memory locations within FFN parameters and the importance of FFN memory across different downstream tasks. MemoryLLM achieves context-free FFNs by training them in isolation from self-attention directly using the token embeddings. This approach allows FFNs to be pre-computed as token-wise lookups (ToLs), enabling on-demand transfer between VRAM and storage, additionally enhancing inference efficiency. We also introduce Flex-MemoryLLM, positioning it between a conventional transformer design and MemoryLLM. This architecture bridges the performance gap caused by training FFNs with context-free token-wise embeddings.

• The paper decouples feed-forward networks from residual streams to form token-indexed memories, enabling deterministic key-value retrieval.
• It introduces pre-computed FFN lookups leveraging Zipfian token distributions and a plug-and-play design to significantly reduce VRAM usage.
• Using Flex-MemoryLLM, it balances efficiency and performance while revealing semantic clustering in token activations for enhanced interpretability.

MemoryLLM: Interpretable and Resource-Efficient Token-Indexed Memory for Transformers

Introduction and Motivation

Understanding the operational role of feed-forward networks (FFNs) in transformer-based LLMs remains substantially underexplored, especially compared to the rich literature on self-attention. The work "MemoryLLM: Plug-n-Play Interpretable Feed-Forward Memory for Transformers" (2602.00398) presents a principled architectural departure from conventional transformers by explicitly decoupling FFNs from the residual stream and self-attention modules. The central claim is that FFNs, when supplied only with static, context-free token embeddings as input, act as deterministic key-value retrieval memories over the model's vocabulary. This context-free design enables efficient pre-computation and offloading, thus mitigating VRAM and inference costs. To further balance performance and interpretability, the paper introduces Flex-MemoryLLM, allocating FFN capacity across both token-indexed memory and conventional computation.

Figure 1: Architecture comparison between a conventional Transformer and MemoryLLM, highlighting the decoupled FFN input and its direct dependence on static token embeddings.

Architecture: Decoupling FFNs via Token-Key-Value Memory

Conventional transformers mix residuals and attention in FFN inputs, which renders the interpretation of FFN internal mechanics intractable. MemoryLLM, however, feeds only static token embeddings to all FFNs, removing dependencies on changing residual states or attention outputs. The TKV (Token-Key-Value) framework formally casts each FFN as a collection of context-free memory cells: tokens index into the key (up projection) and value (down projection) matrices, while the gate projection modulates contribution strength per memory cell.

Figure 2: The TKV framework, illustrating mapping from discrete token IDs to static FFN memory cells via key, value, and gating projections.

Because token activation is now explicit and finite, interpretability and ablation become straightforward: each memory cell is directly and solely associated with a vocabulary item. Crucially, this enables the FFN outputs to be pre-calculated for all tokens and layers in advance, stored as static lookup tables (ToLs).

Pre-Computed FFN Lookups and Plug-and-Play Design

A salient technical contribution is the formulation of FFN outputs over all tokens—across all layers—as pre-computed token-indexed lookups. These ToLs are offloadable from VRAM, loaded on-demand, or even pruned entirely at certain layers or vocabulary slices with minimal computational disruption. This yields a substantial efficiency benefit over dense FFNs that must be computed per input sequence and cannot be decoupled from the transformer’s forward path.

Figure 3: FFN outputs per token can be pre-computed for all vocabulary items and layers, enabling lightweight token-indexed lookups during inference.

The ToL structure leverages Zipfian token distributions—storing frequent ToLs in cache, while rarely needed ones are asynchronously loaded. Experiments show that pruning late-layer FFNs in MemoryLLM induces only mild performance degradation, in contrast to significant U-shaped loss in conventional transformers when FFNs are dropped mid-residual flow.

Figure 4: Percentage perplexity increase when FFN at layer $L$ is dropped: MemoryLLM exhibits robust ablation tolerance compared to conventional LLMs.

Interpretable Memory: Token Semantics and Spatial Structure

Applying t-SNE and K-means clustering on the “importance score” vectors ($c_k$) across all tokens reveals that semantically and lexically similar tokens naturally cluster and activate similar FFN memory cells, confirming that token-level knowledge is locally organized in the learned memory.

Figure 5: Visualization of $c_k$ clustering: tokens with similar semantics access adjacent memory cells, validating spatial semantic structure in FFN memory.

Moving across layers, the clustering coefficient of these importance vectors remains consistently high, especially in early and late layers, while outlier token activation increases towards the output, hinting at convergence to sparse memory access.

Figure 6: (a) Clustering coefficient of token-key vectors per layer. (b) Average number of outlier coefficients, showing sparse memory usage in later layers.

Functional Analysis Across Task Categories

The paper highlights a differential contribution of token-indexed FFN memory across task types. When FFN contribution is continuously reduced in MemoryLLM, performance on retrieval-based and fact-centric tasks decays more sharply than on logical or inferential tasks. Conversely, reasoning tasks maintain relative robustness, indicating that FFN memory predominantly encodes factual knowledge while attention retains compositional capabilities.

Figure 7: Model performance as FFN contribution is interpolated: MemoryLLM is tolerant while conventional LLMs degrade rapidly, supporting architectural decoupling.

Flex-MemoryLLM: Bridging Performance and Interpretability

MemoryLLM’s decoupling, while yielding interpretability and resource savings, introduces a performance gap compared to conventional LLMs of equivalent parameter count due to reduced representational flexibility. Flex-MemoryLLM mitigates this by splitting FFN capacity between "compute" FFN (operating on residuals) and "memory" FFN (token-indexed), tuning the allocation with a parameter $\beta$.

Figure 8: Comparative block diagrams of MemoryLLM, Flex-MemoryLLM, and conventional transformers with identical total parameter count.

Scaling $\beta$ recovers dense-LLM performance, with Flex-MemoryLLM closely matching conventional models in perplexity and downstream evaluation even with substantially fewer active parameters. This points to over-parameterization of FFNs in standard architectures and suggests that explicit token memories can function as drop-in parameter-efficient surrogates for much of the factual knowledge in LLMs.

Figure 9: Model performance at 1B scale—Flex-MemoryLLM bridges the gap between MemoryLLM efficiency and base LLM performance.

Figure 10: Scaling analysis at multiple model sizes and token budgets—Flex-MemoryLLM with reduced active parameters can outperform pruned dense LLMs.

Further, active parameter counts in MemoryLLM and Flex-MemoryLLM substantially undercut those of dense LLMs at comparable performance, and they robustly outperform magnitude, SparseGPT, and Wanda pruning methods at equivalent parameter budgets.

Figure 11: Perplexity vs. active parameters—MemoryLLM and Flex-MemoryLLM dominate conventional pruning approaches at equivalent compression rates.

Implications for Architecture, Deployment, and Model Compression

By proving that a large portion of an LLM’s factual knowledge can be decoupled into static token-indexed memory, MemoryLLM enables modular plug-and-play deployment: only the most relevant memory slices need be resident in RAM/VRAM, with the remainder stored or pruned as dictated by task requirements or hardware constraints. The architecture is also inherently interpretable: memory editability, targeted knowledge injection, and toxicity rectification can be performed by manipulating explicit token memory cells.

Moreover, the plug-and-play approach facilitates new storage-compression strategies, such as quantizing or low-rank compressing per-layer ToLs, which the experiments show can halve memory requirements with negligible accuracy loss. Performance is robust to substantial ablation: dropping middle-layer ToLs has little effect on a wide range of language modeling and reasoning tasks.

Figure 12: Effect of dropping layer $L$ ToLs: task performance is insensitive to ablations in most layers except earliest ones.

Figure 13: Singular value spectra of ToLs across MemoryLLM and Flex-MemoryLLM layers indicate broad compressibility and redundancy.

Conclusion

MemoryLLM rigorously disentangles FFN computation from residual and attention streams, recasting FFNs as explicit, interpretable, context-free key-value memories indexed by tokens. This yields models that are both more transparent and more amenable to resource-constrained inference, with the additional advantage of modular, plug-and-play memory capacity. Empirical results show differential impact of FFN memory on retrieval versus reasoning tasks, the semantic structure of token-level memory, and competitive performance under heavy parameter compression compared to pruning methods. Flex-MemoryLLM architectures, which reintroduce partial FFN computation, can approach dense-LM performance at a fraction of the active memory cost. This framework challenges the necessity of monolithic FFN architectures and opens avenues for interpretable, efficient transformers—in both research and deployment contexts.