[KongchangAI]
· 3 min read· 1,923 words

How Transformers Think and Store Facts: A Detailed Guide to Mechanistic Interpretability

How Transformers Think and Store Facts: A Detailed Guide to Mechanistic Interpretability

How Transformers store facts through MLP key-value lookups and high-dimensional geometry.

This article explains Transformer internals through mechanistic interpretability: MLP layers function as key-value lookup systems storing long-term factual memory in static weights, while high-dimensional spaces enable storing 100,000+ near-orthogonal concept vectors in just 1024 dimensions. Attention layers handle dynamic context routing (short-term memory), and the residual stream serves as an additive information bus connecting all components.

Introduction: Opening the Transformer Black Box

Large Language Models (LLMs) are everywhere, but how they actually "think" and store vast amounts of factual knowledge has long remained a black box. Recently, a Reddit user compiled clear notes on Transformer internals based on modern Mechanistic Interpretability research, focusing on the collaboration between Attention layers and Feed-Forward (MLP) layers, how MLPs serve as "key-value lookup tables," and how high-dimensional geometric spaces enable models to store thousands of facts.

Mechanistic Interpretability is an important branch of AI safety and alignment research, championed by organizations like Anthropic and DeepMind. Unlike traditional "behavioral interpretability" (which only observes input-output relationships), mechanistic interpretability attempts to reverse-engineer the internal computations of neural networks, understanding what each component (neurons, attention heads, layers) specifically computes. Representative works in this field include Chris Olah's team's "Circuits" research framework, Neel Nanda's TransformerLens tool, and Anthropic's recently released large-scale Sparse Autoencoder (SAE) feature dictionaries. The core hypothesis of this field is that neural networks contain understandable algorithmic structures, rather than uninterpretable statistical noise.

The value of these notes lies in reorganizing obscure concepts from academic papers using intuitive mathematical and geometric perspectives, helping us understand the essence of how Transformers "store knowledge." This article provides further elaboration and interpretation.

reddit source: How Transformers Think And Store Facts

MLP Layers: The Model's Long-Term Memory

MLPs Are Actually Key-Value Lookup Systems

A standard Transformer MLP layer (ignoring bias terms) takes an input vector $\mathbf{x}$ from the residual stream and performs two linear transformations with a nonlinear activation function $\sigma$ (such as GELU) in between:

$$\text{MLP}(\mathbf{x}) = \sigma(\mathbf{x} W_1) W_2$$

These notes reference the classic study by Geva et al. (2020), Transformer Feed-Forward Layers Are Key-Value Memories, which shows that $W_1$ can be viewed as a set of Key Vectors $[k_1, ..., k_d]$. When the input $\mathbf{x}$ takes a dot product with a key vector $\mathbf{k}_i$, it's actually measuring the directional alignment between the current token state and that key (unnormalized cosine similarity, scaled by vector magnitude). The closer the direction and the larger the magnitude, the larger the dot product result.

In other words, each key vector $\mathbf{k}_i$ is responsible for detecting a specific semantic pattern or condition—for example, "the current token is some kind of edible thing."

Activation Functions Serve as "Switches"

The nonlinear activation function (GELU or ReLU) after the dot product serves as a gate/threshold: it returns a high positive value when $\mathbf{x}$ matches a certain pattern, and approaches zero otherwise, suppressing irrelevant signals.

GELU (Gaussian Error Linear Unit) was proposed by Hendrycks and Gimpel in 2016, with the mathematical form GELU(x) = x · Φ(x), where Φ(x) is the cumulative distribution function of the standard normal distribution. Unlike ReLU's hard threshold cutoff, GELU provides a probabilistic "soft gate": the larger the input value, the higher the probability of being preserved; the smaller the input value (negative), the higher the probability of being suppressed. The intuition behind this design is that neuron activation shouldn't be a deterministic 0/1 switch, but should decide whether to "pass through" based on the input's relative position in the data distribution. Major Transformer models including the GPT series and BERT all use GELU as the activation function in MLP layers, and experiments show it outperforms ReLU and ELU on language modeling tasks.

Correspondingly, $W_2$ can be viewed as a set of Value Vectors $[v_1, ..., v_d]$. Once a key is successfully matched, the activation function "lights up" the corresponding neuron $i$, multiplying $\sigma(\mathbf{x} W_1)$ by the value vector $\mathbf{v}_i$. This value vector carries factual or linguistic updates associated with that key—for example, boosting the output probability of "Apple" or "Banana" in the vocabulary. Finally, the retrieved value is written directly back into the token's residual stream.

This mechanism clearly explains: MLP layers are the model's long-term factual memory store, with all key-value pairs being static parameters solidified in the weights during pretraining.

High-Dimensional Geometry: Why Millions of Concepts Can Be Stored

The Remarkable Property of Approximate Orthogonality

The most illuminating part of the notes is the geometric interpretation of "facts." We can think of mutually exclusive facts—such as "x is a fruit" and "x is a vehicle"—as value vectors stored in weight matrices that are nearly perpendicular (orthogonal) to each other. Because $\cos(90°)=0$, the dot product between orthogonal vectors is approximately zero, meaning they don't interfere with each other.

The key insight is: high-dimensional spaces allow for an extremely large number of approximately orthogonal vectors. This is why Transformers can accommodate thousands of non-interfering concept directions in a relatively low-dimensional latent space.

This phenomenon has rigorous mathematical backing. The Johnson-Lindenstrauss Lemma (1984) proves that any n high-dimensional points can be projected into O(log n / ε²)-dimensional space while preserving all pairwise distances within a factor of (1±ε). This means high-dimensional spaces have extremely strong "capacity"—randomly sampled vectors are almost always approximately orthogonal. Specifically, when randomly sampling two unit vectors in d-dimensional space, their inner product has an expectation of 0 and a standard deviation of approximately 1/√d. When d=1024, the standard deviation is only about 0.031, which explains why the vast majority of random vector pairs have angles concentrated near 90°. This property has wide applications in compressed sensing, random projections, and neural network theory.

An Intuitive Numerical Comparison

The notes provide an impressive example:

  • In 1024-dimensional space, if vectors must be perfectly orthogonal, you can fit at most 1024 vectors.
  • But if you allow approximate orthogonality (angles between 81.4° and 98.6°), the same space can accommodate over 100,000 vectors.

This is because the capacity of high-dimensional space expands exponentially. It is precisely this property that allows LLMs to store millions of concepts within finite parameter spaces. This also mathematically addresses a common puzzle: why models with seemingly limited parameters can remember such vast world knowledge.

Attention and MLP: Two Dimensions of the Thinking Engine

Clear Division of Labor in Collaboration

Transformer's power comes from the elegant cooperation of two modules. The notes summarize their roles as "information router" and "processing and memory":

  • Attention Layers (Information Router — across sequence dimension): Aggregate context across different sequence positions, deciding "where information flows." Their key-values are dynamic representations derived from other tokens in the prompt, constituting short-term memory.
  • MLP Layers (Processing and Memory — across channel dimension): Serve as key-value lookup memory and nonlinear processing units, transforming features into factual and logical outputs. Their key-values are static weight vectors learned during pretraining, constituting long-term factual memory.
Memory TypeKeys and ValuesInformation Source
Attention Key-ValuesDynamic representations derived from other tokensContext/Sequence (Short-term memory)
MLP Key-ValuesStatic weights $\mathbf{k}_i$, $\mathbf{v}_i$ learned during pretrainingModel parameters (Long-term memory)

By repeatedly stacking these two types of layers, the network can learn increasingly abstract and higher-level semantic associations and facts layer by layer.

Residual Stream: The "Memory Bus" Running Through the Entire Network

The residual stream $\mathbf{x}_l$ is a shared ultra-high-dimensional "memory bus" running through the entire depth of the network. A single encoder layer is defined as:

$$\mathbf{x}_{l+1} = \mathbf{x}_l + \text{Attention}(\mathbf{x}_l) + \text{MLP}(\mathbf{x}_l)$$

Attention layers and MLP layers read information from the residual stream through projections (geometrically equivalent to transforming into another high-dimensional space), complete their respective tasks, and then write updates back to the residual stream. They never erase past states but continuously accumulate information. This additive communication design is one of the key reasons Transformers can stably train very deep networks.

Residual connections were originally proposed by Kaiming He et al. in the 2015 ResNet paper to address the vanishing gradient problem in deep convolutional networks. The core idea is that each layer only needs to learn the "residual" (difference) between input and output, rather than the complete mapping. The Transformer architecture inherited and extended this design—from the perspective of mechanistic interpretability, the residual stream is reinterpreted as an information bus: various layer components read from and write to it in parallel, rather than the traditional understanding of sequential layer-by-layer processing. This perspective was systematically articulated by Elhage et al. (2021) in A Mathematical Framework for Transformer Circuits, providing a theoretical framework for understanding cross-layer feature composition and phenomena like "virtual attention heads."

Transformers as Universal Function Approximators

The synergy between attention's dynamic context aggregation and MLP's information processing and storage is at the core of Transformer's remarkable scalability—from sequence modeling and generative modeling (flow matching, diffusion models) to computer vision, it's virtually everywhere.

This trend began with the success of Vision Transformer (ViT) in 2020, where images are split into patch sequences and fed into a Transformer. In generative modeling, the denoising network of diffusion models (such as the DiT architecture) replaced the traditional U-Net with Transformers; Flow Matching is a continuous normalizing flow training method recently adopted by Meta's MovieGen and Stability AI's SD3, with its core network also based on Transformers. Additionally, Transformers are used in protein structure prediction (AlphaFold2), weather forecasting (Pangu-Weather), robot control (RT-2), and other domains, demonstrating their powerful generalization capability as a universal sequence-to-sequence mapping architecture.

The notes also mention the Universal Approximation Theorem: a feedforward network with a single hidden layer and a non-polynomial activation function can theoretically approximate any continuous function to arbitrary precision. This theorem was first proved by Cybenko (1989) and Hornik (1991) for sigmoid and general activation functions respectively. However, the author specifically highlights a noteworthy "caveat":

The theorem only guarantees that such a network exists, but does not provide a method for finding the network's parameters, nor does it say how large the network needs to be.

This limitation deserves deeper understanding: First, the theorem doesn't guarantee that the required hidden layer width is polynomial—some functions may require exponential width; Second, it doesn't address learning algorithms—gradient descent may not find the approximation solution; Third, it doesn't consider sample efficiency—even if the network exists, training it may require astronomical amounts of data. Deep learning practice shows that deep networks are usually more efficient than shallow networks with equivalent parameters, hinting at the "compositional representation" advantage that depth provides.

This reminds us that there is still a huge gap between theoretical "universality" and engineering "trainability and scalability." The real contribution of the Transformer architecture lies precisely in finding a feasible path to efficiently approximate these functions in practice—through attention mechanisms enabling dynamic computation graphs, through residual connections ensuring gradient flow, and through layer normalization stabilizing training dynamics.

Conclusion

These community-sourced mechanistic interpretability notes sketch out the internal landscape of how Transformers "think and remember" using clear mathematical and geometric intuition:

  • MLP layers are key-value lookup-style long-term memory stores;
  • The approximate orthogonality property of high-dimensional spaces allows massive facts to coexist;
  • Attention and MLP collaborate across sequence and channel dimensions respectively;
  • The residual stream is an additive information bus running through the entire network.

Understanding these mechanisms not only helps us view the capabilities and limitations of LLMs more rationally, but also provides a solid intuitive foundation for future research in model editing, knowledge injection, and interpretability. For example, based on the understanding of MLP key-value memory, researchers have already developed techniques like ROME (Rank-One Model Editing) that can precisely locate and modify specific facts stored in a model without affecting other knowledge—a vivid example of mechanistic interpretability moving from theory to application. References include Vaswani et al. (2017) Attention Is All You Need and Geva et al. (2020) on key-value memories; interested readers are encouraged to explore the original papers.

#mechanistic interpretability#Transformer architecture#MLP key-value memory#high-dimensional geometry#attention mechanism#residual stream#universal approximation theorem
Share:

Related articles