Part 1: Reverse Engineering Transformers: Deconstructing Attention

💡

Other parts in the series:
Part 2: Reverse Engineering Transformers: Attention Heads & Circuits
Part 3 Reverse Engineering Transformers: Path Expansion

Part 4 Reverse Engineering Transformers: Induction Heads

This is a 4 part blog series on my notes and understanding of the foundational paper Elhage, N. et al. (2021). A Mathematical Framework for Transformer Circuits

The blogs will be discussing the following ideas:

• Reverse Engineering Transformers: The primary goal is to break down the complex computations within a transformer into understandable "circuits" or algorithmic patterns. This is analogous to reverse-engineering a compiled program to understand its source code.
• The Residual Stream as a Communication Channel: The paper conceptualizes the residual stream not just as an embedding, but as a central communication channel. Different components of the transformer (like attention heads and MLP layers) read from and write information to this shared space.
• Attention Heads as Independent, Additive Operations: The authors reframe attention layers as a collection of independent attention heads that operate in parallel and add their outputs to the residual stream. This is a more theoretically convenient way to think about them than the typical "concatenate and multiply" view.
• Separating Attention Head Components: The paper proposes that attention heads can be broken down into two largely independent circuits:
  • QK (Query-Key) Circuit: This determines the attention pattern, i.e., which tokens to move information from.
  • OV (Output-Value) Circuit: This determines what information is moved and how it's written to the destination.
• Path Expansion: A key analytical technique is to expand the transformer's computation into a sum of end-to-end paths from input tokens to output logits. This allows for a more granular analysis of how information flows and is transformed.
• Composition of Attention Heads: In multi-layer transformers, the paper explores how attention heads can compose with each other. This composition is what gives deeper transformers their power. The three types of composition are:
  • Q-Composition: A second-layer head's query is influenced by a first-layer head.
  • K-Composition: A second-layer head's key is influenced by a first-layer head.
  • V-Composition: A second-layer head's value is influenced by a first-layer head, creating "virtual attention heads".
• Induction Heads: A specific and powerful circuit that emerges in two-layer (or deeper) models. These heads are a mechanism for in-context learning, allowing the model to recognize repeated sequences and continue them. They work by searching for previous occurrences of the current token and then attending to the next token in that sequence.
• Circuits, Residual Streams, and Induction Heads

We start with the first two points in this blog

What is Reverse Engineering in Mechanistic Interpretability?

At its core, mechanistic interpretability (MI) is the scientific field of reverse engineering neural networks. The goal is to take a fully trained model, a complex and opaque mathematical function defined by billions of parameters, and understand the specific, human-understandable algorithms it has learned to perform its tasks.

Think of it this way:

• A traditional computer program starts with human-written source code (like Python). It's then compiled into a low-level binary file that a machine can execute. A reverse engineer would take that binary and try to reconstruct the original logic and algorithms.
• A neural network starts with a chosen architecture (the "virtual machine") and is trained on data. This process results in a set of optimized weights (the "program binary"). A mechanistic interpretability researcher takes those weights and tries to understand the algorithms they implement.

This analogy is quite deep and helps frame the entire endeavor.

Why is this Analogy So Powerful?

The Central Task: Decomposing Activations

The immediate serious challenge: activations are high-dimensional vectors. How can we possibly understand them?

This is where the concept of a privileged basis becomes critical. The hope is that a network's internal representations aren't just arbitrary directions in a vector space. Instead, MI researchers hypothesize and find evidence that networks often learn to align meaningful, independent "features" with the basis vectors of their activation space (i.e., with individual neurons or specific directions).

• If a privileged basis exists:

  We can decompose activations into independently understandable pieces. This is like figuring out that certain bytes in a program's memory correspond to a player's health, while other bytes correspond to their score. Each can be understood separately. This makes the task of reverse engineering tractable.

• If a privileged basis doesn't exist:

  And every "feature" is stored as some complex linear combination of all neurons (a phenomenon called superposition), the task becomes much harder. It would be like trying to understand a program where a player's health is determined by the 3rd bit of one byte, the 7th bit of another, and so on, all mixed up.

The papers you provided show that while networks don't necessarily converge to a single, unique basis across different training runs, the basis they do learn is not arbitrary. You can't just randomly rotate a layer's activations and expect the network to recover, which provides strong evidence that the learned basis directions are meaningful and privileged.

Summary

"Reverse engineering" is the guiding metaphor for mechanistic interpretability. It frames the goal as deciphering the learned algorithms from the model's weights and highlights that the central challenge is to find a way to break down the model's high-dimensional internal state (activations) into meaningful, understandable variables.

Transformer Circuits I: The Residual Stream as a Communication Channel

TL;DR

• The residual stream is the central data pathway in a Transformer, acting like a communication bus or an assembly line.
• It's a high-dimensional vector space (d_model) where information is stored and progressively refined.
• Transformer components (attention, FFNs) don't replace the stream's content; they read from it and write additive updates back to it.
• A "Read" is a linear projection from the stream (e.g., to create Q, K, V vectors).
• A "Write" is a linear projection that is added back into the stream (x_new = x_old + update).
• This additive structure is the key property that allows us to decompose the model's behavior, but non-linearities like LayerNorm complicate the pure linear view.

Context & Motivation

A Transformer is a deep stack of identical blocks. A natural question is: how does information from early layers persist and get refined by later layers? The answer lies in the residual connections that are present at every step.

To reverse-engineer a Transformer, we need a robust mental model of its architecture. The most powerful one, proposed by the "circuits" interpretability agenda, is to view the model not as a series of opaque transformations, but as a group of specialists collaborating on a shared workspace. That workspace is the residual stream.

Understanding the stream as a communication bus where components perform additive read/write operations is the first and most critical step. It provides the foundation for decomposing the model into understandable parts, allowing us to trace how specific features are computed from input to output.

Prereqs

• Basic familiarity with the Transformer architecture (embeddings, attention, FFN/MLP layers).
• Conceptual understanding of vectors and matrices (linear projections).
• Knowledge of PyTorch/NumPy tensor shapes.

Core Idea in One Picture

Picture the residual stream as an assembly line. For each token, a "package" of information (a vector) moves down the line. At each station, a component reads the package, adds something new, and puts it back on the line.

graph TD
    A[Input + Pos Embedding] --> B[Residual Stream]

    subgraph Layer0
        B --> Attn0[Attention Head]
        Attn0 --> B
        B --> FFN0[FFN Layer]
        FFN0 --> B
    end

    B --> MoreLayers[...]

    subgraph FinalLayer
        MoreLayers --> B
        B --> Unembed[Unembedding]
        Unembed --> Logits[Output Logits]
    end

    style B fill:#e6f3ff,stroke:#333,stroke-width:2px

Definitions & Setup

Let's formalize this. The residual stream is the state x that is passed from one layer to the next.

• Shape: The stream is a tensor of shape [batch_size, seq_len, d_model], where d_model is the model's hidden dimension (e.g., 768 for GPT-2 small). We will focus on the vector of a single token, which has shape [d_model].
• Initial State (x_0): The process starts with the sum of a token's embedding and its positional embedding. For a token t at position i:
  
  $\mathbf{x}0 = \mathbf{W_E}[t] + \mathbf{W{pos}}[i]$
  
• Read Operation: A component "reads" from the stream by applying a linear projection (i.e., multiplying by a weight matrix). For an attention head, this means creating its Query, Key, and Value vectors from the current stream state x_l:
  
  $\mathbf{q} = \mathbf{x}_l \mathbf{W_Q}, \quad \mathbf{k} = \mathbf{x}_l \mathbf{W_K}, \quad \mathbf{v} = \mathbf{x}_l \mathbf{W_V}$
  
• Write Operation: After a component computes its output (e.g., an attention head's output o or an FFN's output f), it "writes" this back to the stream as an additive update.
  
  $\mathbf{x}_{l+1} = \mathbf{x}_l + \mathbf{o}$
  
  This additive nature is the defining feature of residual connections and our interpretability model.

The Crucial Role of Layer Normalization

There is one key detail that complicates this clean, linear story: Layer Normalization (LayerNorm). Before most components read from the stream, the stream is normalized. The full update rule is closer to:

LayerNorm is a non-linear operation. This means that a component's output depends on the entire state of the input vector, not just on individual features within it. This breaks perfect linear decomposability, a limitation we must always keep in mind.

Walkthrough

Case Study: Tracing "Paris"

Let's trace the vector for the token "Paris" in the prompt "The capital of France is Paris."

This final, highly contextualized vector is then read by the unembedding matrix to predict the next token.

Implications & Limits

• Why this model is useful: Because the updates are (approximately) linear and additive, we can analyze the contribution of each component to the final output. This allows for attribution: we can measure how much attention head 5.2 contributed to the model predicting "Rome" instead of "Madrid."
• When it breaks (LayerNorm): As noted, LayerNorm's non-linearity means LN(a + b) ≠ LN(a) + LN(b). We cannot perfectly decompose a component's input into the sum of prior contributions. The linear model is a powerful and often accurate approximation, but it is not the ground truth.
• Competition for the bus: All components write to the same d_modeldimensional space. They must learn to write their updates in directions that don't destructively interfere with information written by other components. This pressure may be one reason why models learn to align features with specific basis directions (the "privileged basis" hypothesis).

Pitfalls

• Forgetting LayerNorm: It's easy to forget the non-linear LayerNorm step, which invalidates claims of perfect linearity. Always treat the additive model as a strong and useful linear approximation.
• Thinking of the stream as static: The residual stream is not just an "embedding." It is a dynamic workspace where the representation of a token is constantly being rewritten and augmented based on new context computed by the model.

Takeaways

• The residual stream is the central communication bus of a Transformer.
• Components "read" from the stream via linear projections and "write" additive updates back to it.
• This additive structure allows us to reason about and decompose the model's computation.
• The final representation of a token is the sum of its initial embedding and all updates from every attention head and FFN layer.
• Layer Normalization is a key non-linearity that complicates this pure linear picture, making it a powerful approximation rather than a perfect model.
• This "assembly line" mental model is the foundation for all further circuit-level analysis.