PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ajay Jaiswal, Lauren Hannah, Han-Byul Kim, Duc Hoang, Arnav Kundu, Mehrdad Farajtabar, Minsik Cho

arXiv:2602.00398v1cs.LG

TL;DR

Existing FFNs are difficult to interpret because their inputs are contextualized residual-stream representations, leaving their memory behavior and task importance unclear. MemoryLLM decouples FFNs from self-attention and trains them on token embeddings, while Flex-MemoryLLM provides an intermediate design; the resulting FFNs support token-level retrieval analysis and pre-computed lookups with reduced memory and compute cost.

  • Problem

    FFNs are difficult to interpret because contextualized residual-stream inputs obscure their memory locations, while their influence on downstream tasks remains underexplored.

  • Method

    MemoryLLM decouples FFNs from self-attention, trains them independently on context-free token embeddings, and analyzes them with a token-key-value retrieval-memory framework.

  • Results

    MemoryLLM FFNs store token-level parametric knowledge, support retrieval-based tasks more strongly than inferential or logical tasks, and can be pre-computed as token-wise lookups.

  • Takeaways & Limitations

    The architecture makes FFNs analyzable over a finite vocabulary and enables reduced memory footprint and compute cost through pre-computed lookups.

  • Takeaways & Limitations

    The approach assumes that token distributions in modern LLM-generated content follow Zipf’s law irrespective of tokenizer.

Abstract

from arXiv · show

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.

1 Introduction

MemoryLLM addresses FFN interpretability by decoupling FFNs from self-attention and training them on token-indexed embeddings, creating a context-free retrieval-memory view. The design also supports FFN interpretability, storage offloading, and Flex-MemoryLLM’s intermediate architecture.

  • FFNs are difficult to study because they consume non-interpretable residual-stream mixtures that include prior self-attention outputs.
  • MemoryLLM trains self-attention and FFNs independently and in parallel, with FFNs trained directly on context-free token-indexed embeddings.
  • The TKV framework treats FFNs as token-key-value retrieval memory over a finite vocabulary and examines how similar tokens access similar memory locations.
  • FFNs act as reservoirs of token-level parametric knowledge and are more important for retrieval-based tasks than for inferential or logical-thinking tasks.
  • Token-wise lookups enable pre-computation and storage transfer of FFNs, addressing inference-time memory and computational bottlenecks.
  • Flex-MemoryLLM splits FFN parameters between context-aware and context-free modules to bridge the performance gap with conventional LLMs.

2 MemoryLLM: LLMs with Interpretable Token-Indexed Feed-Forward Memory

MemoryLLM disentangles FFNs from self-attention by training them on context-free token embeddings, making FFNs interpretable as token-indexed retrieval memory. This design also enables pre-computed lookups that can be transferred between storage and VRAM according to layer importance.

  • Motivation: Conventional FFNs are difficult to interpret because they process a context-dependent mixture of self-attention output and residual information, despite comprising about two-thirds of LLM parameters.Prior analyses also face difficulty mapping intermediate FFN keys back to initial input tokens because contextual representations shift across layers.
  • MemoryLLM Architecture: MemoryLLM trains self-attention conventionally while training every FFN independently on token-indexed, context-free embedding vectors.All FFNs across the transformer receive embeddings generated directly from tokenizer token IDs rather than intermediate contextual features.
  • TKV Framework: The TKV framework treats FFN up-projection and down-projection matrices as keys and values, while the gate projection reweights memory-cell coefficients.For each token query, key coefficients are estimated and then used to form a weighted combination of value vectors.
  • TKV Framework: Token-indexed queries define a finite, human-interpretable memory space and remove the need to reverse-engineer contextual input prefixes from calibration data.This framework supports studying how vocabulary tokens access FFN memory locations and how those locations matter for downstream tasks.
  • Efficiency: Because FFN inputs are static and vocabulary-bounded, outputs for every token and transformer layer can be pre-computed as token-wise lookups and offloaded to storage.The lookups support asynchronous prefetching and reduce computational and VRAM overhead during inference.
  • Efficiency: Later FFN layers can be permanently offloaded with minimal residual-flow disruption because their contribution to MemoryLLM performance drops after the first few layers.The corresponding layer-importance pattern differs from the non-uniform U-shaped behavior reported for a conventional base LLM.

3 Empirical Study of FFN Neural Memory

MemoryLLM exposes FFN memory as token-indexed key-value retrieval and shows that semantically related tokens access related locations. FFN contribution is especially important for retrieval tasks, while decoupling supports controlled analysis of model behavior.

  • 3.1 Spatial Distribution of Neural Memory in FFNs: The TKV framework models FFNs as token-wise key-value retrieval memory, where query-token coefficients determine contributions from memory keys.Each memory output is formed from key-activation coefficients and corresponding value vectors.
  • 3.1 Spatial Distribution of Neural Memory in FFNs: Semantically similar tokens tend to activate similar FFN memory keys, enabling targeted knowledge editing, injection, and toxicity suppression.The observed clustering links token semantics to specific memory locations.
  • 3.1 Spatial Distribution of Neural Memory in FFNs: High clustering coefficients appear across MemoryLLM-1B layers, while terminal layers contain more outlier keys associated predominantly with output formatting.The pattern suggests token-level information converges within limited keys.
  • 3.2 Probing FFN Memory Across Downstream Tasks: MemoryLLM’s decoupled architecture permits direct study of FFN effects, with decreasing FFN contribution causing comparatively less performance degradation than in conventional checkpoints.The regulated-contribution experiment compares MemoryLLM and base checkpoints.
  • 3.2 Probing FFN Memory Across Downstream Tasks: Reducing FFN contribution hurts recall and retrieval tasks more than reasoning or logical-thinking tasks.This comparison uses controlled reductions of the FFN contribution coefficient α across two task categories.

4 MemoryLLM: Performance and Efficiency

MemoryLLM trades dense computational capacity for token-wise lookup memory, reducing active parameters but falling short of dense performance at equal total size. Flex-MemoryLLM restores capacity by splitting FFN parameters between context-aware computation and context-free memory, narrowing or eliminating this gap.

  • 4.1 MemoryLLM Comparison with Conventional LLM: Token-wise lookups can be precomputed and offloaded, reducing active VRAM parameters because FFNs typically contain about two-thirds of model parameters.The lookup strategy makes effective active parameters approximately one-third of total model parameters.
  • 4.1 MemoryLLM Comparison with Conventional LLM: At equal total parameter counts, MemoryLLM underperforms conventional dense LLMs but notably outperforms its dense counterpart when comparing effective active parameters.ToLs are excluded from the active-parameter count in this comparison.
  • 4.2 Flex-MemoryLLM: Bridging Conventional LLM and MemoryLLM: Flex-MemoryLLM splits FFN parameters between FFN Compute, which operates on residual flow, and FFN Memory, which remains context-free and token-embedding-trained.Moving parameters from FFN Memory to FFN Compute increases active computational capacity.
  • 4.2 Flex-MemoryLLM: Bridging Conventional LLM and MemoryLLM: Flex-MemoryLLM balances efficiency and performance, achieving close to base performance with approximately 5h^2 fewer active parameters.The result follows experiments using progressively different parameter divisions between computation and memory.
  • 4.2 Flex-MemoryLLM: Bridging Conventional LLM and MemoryLLM: With 25B-to-150B token scaling, the gap to dense models diminishes, and Flex-MemoryLLM-3h^2 with 704M active parameters closely matches a dense model with 1.2B active parameters.The 1B Flex-MemoryLLM-3h^2 model also outperforms the base-737M model at the stated active-parameter count.
  • 4.3 Comparison with Pruning Methods: The paper compares MemoryLLM and Flex-MemoryLLM against Magnitude, SparseGPT, and Wanda pruning methods at 1B and 750M total parameter counts.Figure 11 uses dense base models as reference lines for active-parameter comparisons.

5 Conclusion

MemoryLLM decouples FFNs from residual flow and self-attention, making them interpretable as token-indexed neural memories. The resulting memory structure links similar tokens to similar locations, supports retrieval-task performance, and enables precomputed token-wise lookups.

  • 5 Conclusion: MemoryLLM trains FFNs with context-free token embeddings, making them interpretable as neural key-value memory over the vocabulary.The vocabulary provides a finite, human-interpretable query space.
  • 5 Conclusion: Lexically and semantically similar tokens are indexed across similar FFN memory locations, and this knowledge is crucial for retrieval-based tasks.The conclusion connects spatial organization of memory with downstream retrieval behavior.
  • 5 Conclusion: FFN parameters can be precomputed as token-wise lookups, reducing memory footprint and compute cost.The lookup representation supports the efficiency benefits reported for MemoryLLM.

A Implementation Details

The implementation details report common training configurations and inference measurements for the model variants. Decoding-speed and memory experiments use a single A100 GPU with BF16 token-wise lookups and sequence length 2048.

  • A Implementation Details: Base, MemoryLLM, and Flex-MemoryLLM checkpoints use the same configuration for fair comparison.The shared training setup is summarized in Table 3.
  • A Implementation Details: The implementation varies how FFN parameters are divided between context-dependent computation and context-free memory.Table 4 reports the parameter division used by the model variants.
  • A Implementation Details: Inference experiments measure decoding speed and memory requirements on one A100 GPU with BF16 ToLs and sequence length 2048.The reported decoding speed accounts for loading ToLs from storage into VRAM.

C.1 Memory Augmented Architectures

Memory-augmented architectures expand effective model capacity with trainable or external memory, while this paper situates its designs within that broader lineage.

  • Memory Augmented Architectures: Memory-augmented models expand a model’s effective parameter space without large computational overhead.The passage traces this goal from early memory networks through differentiable external-memory systems.
  • Memory Augmented Architectures: MemoryLLM and Flex-MemoryLLM are evaluated alongside Base transformer models at the 1B total-parameter scale.The empirical comparisons include memory requirements and token decoding estimates.
  • Memory Augmented Architectures: The paper connects persistent memory systems with earlier memory-augmented architectures and rank-one matrix memory with mixture-of-experts models.These connections extend the related-work context beyond conventional vector-valued memory.
  • Memory Augmented Architectures: Accurate factual generation is commonly assessed with open-domain question answering and other knowledge-intensive tasks.The passage motivates memory mechanisms through the importance of encoding factual knowledge for correct responses.

C.2 Understanding Feed-Forward Networks in Transformers.

Prior FFN analyses interpret these modules as knowledge memories, but contextualized residual activations make their query space indirect. The paper instead trains dense FFNs on context-free token embeddings for deterministic token-level interpretability.

  • Understanding Feed-Forward Networks in Transformers: Prior studies interpret FFNs as key–value memories associated with linguistic features and factual recall.They map FFN keys to lexical or semantic patterns and examine vocabulary-level concepts during prediction.
  • Understanding Feed-Forward Networks in Transformers: Contextualized residual activations and calibration-data mining make inferred FFN query spaces indirect and difficult to interpret.MoLE also retains contextual dependence through routers trained with attention output.
  • Understanding Feed-Forward Networks in Transformers: MemoryLLM trains dense transformer FFNs directly on context-free token embeddings, eliminating contextual ambiguity for deterministic token-level interpretability.This design addresses whether dense FFN computation can be disentangled from intermediate activations without substantially harming trainability.

D Understanding Storage Challenges of ToLs

Token-wise lookup tables make FFN memory transferable between VRAM and storage, but their storage cost motivates quantization, low-rank, and layer-wise compression studies.

  • Understanding Storage Challenges of ToLs: Token-wise lookup tables can be pre-computed and transferred between VRAM and storage to address resource-constrained inference.The paper frames this as a trade-off between VRAM and storage-device capacity.
  • Understanding Storage Challenges of ToLs: MemoryLLM-1B requires approximately 12.6 GB of F16 storage for ToLs across 24 layers and a 128,256-token vocabulary.The model uses hidden dimension 2048 and a LLaMa-3.1 tokenizer.
  • Understanding Storage Challenges of ToLs: The storage investigation considers quantization, low-rank compression, and layer-wise ToL compression.These are presented as three distinct perspectives on reducing lookup-table storage cost.
  • Understanding Storage Challenges of ToLs: Table 7 compares MemoryLLM-1B performance across low-precision token-wise lookup tables.The table targets the quantization perspective of ToL storage reduction.

D.2 Low Rank Compression of Token-wise ToLs

The paper studies low-rank representations of token-wise lookup tables and finds that uniform SVD can substantially reduce storage with only marginal performance change, while layer-specific compression may improve the trade-off further.

  • D.2 Low Rank Compression of Token-wise ToLs: The paper places ToL compression within broader efforts to exploit low-rank characteristics in large-model weights and gradients.Those prior approaches target storage demands and computational complexity.
  • D.2 Low Rank Compression of Token-wise ToLs: Each layer’s ToL has dimension vocab_size × hidden_dim and can be replaced by rank-r SVD factors U and V.The factors have dimensions |V|×r and r×d, respectively, instead of storing the full table.
  • D.2 Low Rank Compression of Token-wise ToLs: For MemoryLLM-1B, storage-saving low-rank compression requires r ≤ 2015, corresponding to at least 2% rank reduction.The estimate uses 24 layers, hidden dimension 2048, and vocabulary size 128,256.
  • D.2 Low Rank Compression of Token-wise ToLs: Uniform SVD compression reduces ToL storage cost by approximately 2× with a marginal change in model performance.The experiment applies the same compression strategy across all layers.
  • D.2 Low Rank Compression of Token-wise ToLs: Non-uniform rank reduction is proposed as a future direction because ToLs exhibit different low-rank properties across layers.The paper reports normalized singular values across 24 transformer blocks and expects layer-specific compression to offer a better performance trade-off.

D.3 Layer-wise ToLs Compression

Layer-wise ToL removal shows that early layers are most important for performance, while many middle-layer ToLs can be omitted with marginal impact. The resulting redundancy suggests middle-layer ToL dropping as a storage-compression strategy.

  • Dropping ToLs from the first few layers causes the major performance degradation across all eight tasks.Figure 12 reports task performance as a function of the dropped layer.
  • Middle-layer ToLs are highly redundant and have marginal impact on performance when removed.
  • Dropping middle-layer ToLs is a promising compression direction under limited storage availability.
  • The singular values of ToLs are normalized and sorted across 24 layers in MemoryLLM and Flex-MemoryLLM at 1B scale.Figure 13 compares the 2048 singular values associated with the layer-wise ToLs.
Loading 2602.00398v1…