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

TL;DR

This work tackles the interpretability of feed-forward networks (FFNs) in transformers by decoupling FFNs from self-attention, reframing FFNs as context-free, token-indexed memory. MemoryLLM trains FFNs in isolation on token embeddings to form a deterministic key-value memory over the vocabulary, enabling pre-computation as token-wise lookups (ToLs) and plug-n-play memory transfer to storage, thereby reducing active parameters in memory and increasing efficiency. The authors introduce the token-key-value (TKV) framework to analyze memory locations and show that semantically similar tokens access related memory keys, with memory contributing more to retrieval-based tasks. To bridge the performance gap with conventional dense models, Flex-MemoryLLM splits FFN capacity into memory and compute components, achieving near-base performance with substantial memory offloading. Across comparisons with pruning approaches, MemoryLLM and Flex-MemoryLLM demonstrate favorable trade-offs between interpretability, efficiency, and task performance, highlighting a viable path for memory-augmented, interpretable LLM architectures.

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.

Figures (13)

• Figure 1: Architecture comparison of Conventional Transformer ($\mathrm{base}$) v/s MemoryLLM with Residual Stream Perspective: (a) FFN input in conventional transformers is a sequential and non-interpretable latent snapshot of a residual stream, including prior self-attention module output; (b) MemoryLLM decouples FFNs across all transformer blocks completely from self-attention modules and trains them in isolation of the residual stream, directly on token-indexed input embeddings.
• Figure 2: TKV Framework: Input text is tokenized into discrete token IDs as context-free query vectors for FFN memory cells. $W_{Up}$ and $W_{Down}$ projection matrices emulate the behavior of Keys and Values while the $W_{Gate}$ matrix can be interpreted as a reweighting function of token memory coefficients.
• Figure 3: FFNs as Pre-computed Token-wise Lookups: Outputs corresponding to each vocabulary tokens for all FFN modules across $N$ transformer blocks can be pre-computed offline and stored as static token-indexed lookups (ToLs) in storage devices.
• Figure 4: Percentage increase in perplexity when FFN computation for layer $L$ is dropped in $\mathrm{Base}$ and MemoryLLM.
• Figure 5: Semantically Similar Tokens Build Memory Outputs with Similar Keys: t-SNE plot with K-Means clustering of $c_k$ vectors, which represent each key's contribution to memory outputs, yields clusters of tokens with semantically similar properties.