Analyzing Memory Capacity in Linear Layers and Transformers
Understanding how neural networks store information through linear maps, MLPs, and self-attention mechanisms.
1The Concept of Associative Memory
Large language models function as vast associative memory stores, recalling obscure facts upon query. A memory store is defined as a function holding key-value pairs where inputs retrieve specific outputs. The core challenge lies in balancing the number of stored memories against the interference generated during the retrieval process.
2Linear Layers as Memory Stores
Linear layers accumulate memories via outer products of keys and values. When querying with an input, the system returns a sum of values weighted by similarity. However, this basic mechanism suffers from significant interference once the number of stored memories exceeds the key dimension, leading to degraded recall accuracy.
3Interference-to-Signal Ratio Dynamics
The interference-to-signal ratio measures how much stored memories obstruct the retrieval of a target. In an outer-product store, interference grows as a random walk. Once the number of memories reaches the key dimension, interference matches the signal strength, effectively drowning out the desired output and limiting total storage capacity.
4Optimal Linear Memory Stores
By applying least-squares solutions, one can construct an optimal linear memory store. This method uses the pseudoinverse of the key matrix to whiten inputs and cancel interference. This approach doubles the capacity compared to standard Hebbian outer-product stores, allowing for exact recall until the memory count exceeds the key dimension.
5Limitations of Linear Scaling
Even optimal linear stores struggle when the number of memories significantly surpasses the key dimension. Because linear maps scale all similarities by the same factor, they cannot suppress interference from low-similarity keys. This necessitates the introduction of nonlinearities to selectively filter out irrelevant data during the memory recall process.
6MLPs as Nonlinear Memory Stores
Two-layer MLPs function as memory stores where the first weight matrix acts as keys and the second as values. By passing similarities through a nonlinear activation function, the network suppresses low-similarity interference. This architecture allows MLPs to store far more facts than their hidden unit count would suggest.
7The Role of Squared ReLUs
Squared ReLU activations improve recall by shrinking low-similarity weights more aggressively than large ones. This technique, utilized in SwiGLU architectures, ensures that only relevant memories contribute significantly to the output. By suppressing noise, these activations allow models to maintain high signal integrity even when storing a dense volume of information.
8Superposition and Sparse Activation
Modern MLPs store facts in superposition, blending multiple memories within the same weights. Research indicates that trained transformer MLPs often activate only 1% of hidden units for a given input. This sparse activation ensures that each recall operation draws on a small, focused fraction of the total stored memory bank.
9Historical Context of Transformers
The original transformer paper by Vaswani et al. briefly explored the connection between feedforward layers and attention mechanisms. Early versions experimented with replacing feedforward layers with attention over learned parameters. While translation quality remained comparable, training efficiency varied, leading to the eventual removal of these specific comparisons from later published versions.
10Self-Attention and Unbounded Capacity
Self-attention mechanisms represent the evolution of memory storage, offering superior recall and potentially unbounded capacity. By utilizing 64-dimensional keys, these systems jump from storing roughly 64 memories in linear layers to over a million. This shift fundamentally changes how models retrieve information compared to static, fixed-weight linear or MLP layers.
11Numerical Validation of Capacity
Numerical tests confirm that outer-product stores reach capacity at M equals 64 for 64-dimensional keys. Conversely, optimal stores maintain zero interference until M equals 64, with interference rising slowly thereafter. These empirical results validate the theoretical framework that increasing key dimensions or using selective nonlinearities directly enhances model memory.
12Future of Neural Memory
Understanding memory capacity through the lens of interference-to-signal ratios provides a rigorous foundation for architecture design. By treating layers as associative stores, researchers can better optimize how models learn and retrieve complex information. This perspective highlights why specific activation functions and weight structures are critical for scaling modern intelligence.