Part 2: Reverse Engineering Transformers: Attention Heads & Circuits

💡

Other Parts in the series:
Part 1: Reverse Engineering Transformers: Deconstructing Attention
Part 3 Reverse Engineering Transformers: Path Expansion
Part 4 Reverse Engineering Transformers: Induction Heads

Transformer Circuits II: Attention Heads as Independent Operators

• The standard implementation of multi-head attention concatenates head outputs and uses a single large output projection (W_O).
• We can mathematically reframe this as summing the outputs of independent heads, each with its own smaller output projection matrix (W_O^h).
• This reframing is key for interpretability: it allows us to analyze the contribution of each attention head to the residual stream in isolation.
• An attention head's "write" to the residual stream is its value-weighted attention pattern, projected back into d_model space by its personal W_O^h.
• While heads in a layer compute in parallel, they are not fully independent; they read from the same source and their outputs may compete for capacity in the next layer's residual stream.

Context & Motivation

In our previous post, we established the residual stream as the central communication bus of a Transformer. We saw it as an assembly line where components "read" from the stream and "write" additive updates back to it. Our diagram represented "Attention" as a single block that performed one of these read-write operations.

Now, we need to look inside that block. A standard Transformer layer contains multi-head attention. The typical textbook explanation involves concatenating the outputs of all heads and projecting them back to the model dimension with a single output matrix, W_O. This implementation view is computationally efficient, but it's messy for interpretation. It mixes the outputs of all heads together, making it hard to ask: "What was the specific contribution of head 7?"

This post shows that we can adopt a more theoretically convenient, yet mathematically equivalent, perspective: the attention layer's total output is simply the sum of the outputs of each individual head. This allows us to treat each head as an independent operator writing to the residual stream, making our reverse engineering task dramatically simpler.

Prereqs

• Familiarity with the concepts of Query, Key, and Value in attention.
• Understanding of the residual stream as a communication bus (from Part I).

Core Idea in One Picture

The key is recognizing that the standard, monolithic attention output projection can be split into a sum of per-head contributions.

graph TD
    subgraph Standard View (Implementation)
        H1[Head 1 Out] --> C;
        H2[Head 2 Out] --> C;
        Hn[... Head N Out] --> C;
        C(Concat) --> WO1[W_O];
        WO1 --> LayerOut1[Layer Output];
    end

    subgraph Equivalent View (Interpretation)
        H1_i[Head 1 Out] --> WO_h1[W_O^1];
        H2_i[Head 2 Out] --> WO_h2[W_O^2];
        Hn_i[... Head N Out] --> WO_hn[W_O^n];
        WO_h1 --> S((+));
        WO_h2 --> S;
        WO_hn --> S;
        S --> LayerOut2[Layer Output];
    end

    LayerOut1 -- Mathematically Identical --> LayerOut2;

    style H1 fill:#f9f,stroke:#333
    style H2 fill:#f9f,stroke:#333
    style Hn fill:#f9f,stroke:#333
    style H1_i fill:#f9f,stroke:#333
    style H2_i fill:#f9f,stroke:#333
    style Hn_i fill:#f9f,stroke:#333

Figure 1: The standard implementation view (left) concatenates head outputs before a single projection. The interpretation view (right) shows this is equivalent to summing the outputs of heads projected by their own virtual output matrices.

Definitions & Setup

Let's formalize the two views. We assume a model with n_heads heads and a head dimension d_head. The model dimension is d_model, where typically d_model = n_heads * d_head. The input to the layer from the residual stream is x.

1. The Standard View (Implementation)
Each head h produces an output vector o_h of shape [seq_len, d_head]. These are concatenated along the last dimension to form a single large tensor O_cat of shape [seq_len, n_heads * d_head]. This tensor is then projected back to the residual stream's dimension by the output weight matrix W_O, which has shape [n_heads * d_head, d_model].

Output = Concat(o_1, o_2, ..., o_{n_heads}) @ W_O

2. The Theoretically Convenient View (Interpretation)
We can conceptually partition the large W_O matrix into n_heads smaller matrices, W_O^h, one for each head. Each W_O^h has shape [d_head, d_model].

The total output is then the sum of each head's output projected by its corresponding W_O^h:

$\text{Output} = \sum_{h=1}^{n_{\text{heads}}} o_h @ W_O^h$

The output of a single head, o_h @ W_O^h, is that head's direct, additive contribution to the residual stream. This is exactly the "write update" we discussed in Part I.

Walkthrough

Step 1: Proof of Equivalence

Why are these two views identical? Let's look at the structure of the matrices. The Concat(o_1, ..., o_n) operation creates a wide matrix. The W_O matrix is correspondingly tall.

Let O_cat = [o_1 | o_2 | ...] be the row-wise concatenation. Let W_O be the block matrix composed of the W_O^h matrices stacked vertically:

$\mathbf{W_O} = \begin{pmatrix} \mathbf{W_O^1} \\ \mathbf{W_O^2} \\ \vdots \\ \mathbf{W_O^{n_{heads}}} \end{pmatrix}$

The matrix multiplication O_cat @ W_O is then:

$[\mathbf{o_1} | \mathbf{o_2} | \dots] \begin{pmatrix} \mathbf{W_O^1} \\ \mathbf{W_O^2} \\ \vdots \end{pmatrix} = \mathbf{o_1 W_O^1} + \mathbf{o_2 W_O^2} + \dots = \sum_{h} \mathbf{o_h W_O^h}$

(Note: This is a slight abuse of notation for clarity; the multiplication is properly defined by partitioning columns and rows.)

Step 2: A Minimal Experiment

We can verify this with a small PyTorch snippet. We'll simulate the outputs of two heads and show the two methods produce the same result.

import torch

# Setup: 1 token, 2 heads, d_head=4, d_model=8
seq_len, n_heads, d_head = 1, 2, 4
d_model = n_heads * d_head

# Random head outputs
o_1 = torch.randn(seq_len, d_head) # Output of head 1
o_2 = torch.randn(seq_len, d_head) # Output of head 2

# Big W_O matrix for the standard view
W_O = torch.randn(d_model, d_model)

# 1. Standard View: Concatenate and multiply
o_cat = torch.cat([o_1, o_2], dim=-1) # Shape:
output_standard = o_cat @ W_O

# 2. Interpretation View: Split W_O and sum
W_O1 = W_O[:d_head, :]  # Shape:
W_O2 = W_O[d_head:, :]  # Shape:

output_interp = (o_1 @ W_O1) + (o_2 @ W_O2)

# Check if they are nearly identical
print("Standard output:", output_standard)
print("Interpretation output:", output_interp)
print("Outputs are close:", torch.allclose(output_standard, output_interp))
# Outputs are close: True

Step 3: Why This Reframing is Powerful

This equivalence is the foundation for analyzing individual heads. We can now cleanly state what a single head h does:

This allows us to ask targeted questions like, "What information does head 8.3 read from the 'Paris' token to write to the 'capital' token?" and answer it by analyzing only the matrices W_Q^{8.3}, W_K^{8.3}, W_V^{8.3}, and W_O^{8.3}. We have successfully decomposed a layer-level operation into a sum of head-level operations.

Implications & Limits

• Implication: The fundamental unit of analysis for attention inside a Transformer can be the individual head, not the entire layer. This makes reverse engineering tractable. We can study circuits composed of specific heads in different layers.
• Limits: The heads within a layer are not truly independent.
  1. Shared Input: They all read from the exact same normalized residual stream x_norm.
  2. Shared Output: Their outputs are all summed into the same residual stream. A powerful update from one head could change the vector's direction or magnitude in a way that affects how it is processed by the LayerNorm in the next block.

Pitfalls

• Confusing parallel computation with total independence. Just because heads are computed in parallel doesn't mean they don't influence each other's operating environment in subsequent layers. The additive nature of their outputs means they are in a "linear superposition" within the stream, but the subsequent LayerNorm is non-linear and will mix their signals.

Takeaways

• A multi-head attention layer's output is mathematically equivalent to the sum of individual head outputs.
• Each head can be thought of as having its own input (W_Q, W_K, W_V) and output (W_O) matrices.
• This decomposition allows us to isolate and analyze the function of a single attention head.
• A head's function is to read information from the residual stream, transform it via attention, and write an update back to the stream.
• This "independent, additive" view is a cornerstone of the circuits framework for understanding Transformers.

Transformer Circuits III: The QK and OV Circuits Inside an Attention Head

• An individual attention head's operation can be decomposed into two distinct sub-circuits: the QK-Circuit and the OV-Circuit.
• The QK-Circuit determines the attention pattern-where the head looks. It uses Query and Key vectors to decide which source tokens to get information from.
• The OV-Circuit determines the content being moved-what information the head moves. It uses Value vectors and the head's Output projection to specify the update written to the residual stream.
• This separation is a powerful analytical tool, allowing us to categorize heads by their QK behavior (e.g., "previous token heads") and their OV behavior (e.g., "copying heads") separately.
• The full effect of the OV-Circuit on the residual stream is captured by the composed matrix W_V @ W_O, which maps features from a source token's stream to the output update.

Context & Motivation

In our previous posts, we established that a Transformer layer's attention is best understood as a sum of contributions from independent heads, each writing an update to the residual stream. This was a crucial simplification, allowing us to isolate a single head for analysis.

Now, we zoom in further. What algorithm does a single head execute? An attention head performs a complex operation: for a given "query" token, it looks at all other "key" tokens, computes similarity scores, and then creates a weighted sum of their "value" vectors. This suggests two sub-tasks are happening at once:

The circuits framework makes this separation explicit, decomposing a head into a Query-Key (QK) circuit and an Output-Value (OV) circuit. By analyzing these two parts, we can achieve a much more precise understanding of a head's role in the model's overall algorithm.

Prereqs

• Understanding the residual stream as a communication bus.
• The view of attention layers as a sum of independent head outputs.
• Basic knowledge of the Query, Key, and Value formulation of attention.

Core Idea in One Picture

An attention head can be split into two parallel information paths: one calculates attention scores, the other prepares the information to be moved. They only combine at the very end.

graph TD
    subgraph Attention Head h
        direction LR
        x_norm[Normed Residual Stream]

        subgraph QK-Circuit (Where to look)
            x_norm -- W_Q^h --> Q[Query];
            x_norm -- W_K^h --> K[Key];
            Q --> AS(Attention Scores);
            K --> AS;
        end

        subgraph OV-Circuit (What to move)
           x_norm -- W_V^h --> V[Value];
           V -- W_O^h --> Projected_V[Projected Value];
        end

        AS -- Softmax --> A[Attention Pattern: A];
        A --> Combine((@));
        Projected_V --> Combine;
        Combine --> Output[Output Δ_h];
    end

    Output --> RS_Update{+ Add to Residual Stream};
    ```

Definitions & Setup

Let's formalize the two circuits for a single head h. The input from the previous layer is the normalized residual stream, x_norm.

1. The QK-Circuit (Attention Pattern)

The QK-circuit's job is to compute a scalar attention score between a query token t_q and a key token t_k.

• Inputs: Residual stream vectors x_q and x_k.
• Matrices: W_Q^h (Query) and W_K^h (Key), both of shape [d_model, d_head].
• Computation: The attention score s_{q,k} is computed via a dot product.
  
  $s_{q,k} = \frac{(\mathbf{x_q} \mathbf{W_Q^h}) \cdot (\mathbf{x_k} \mathbf{W_K^h})}{\sqrt{d_{head}}}$
  
• Output: A sequence of scores, which are passed through a softmax function to create the final attention pattern A. This pattern is a set of probabilities summing to 1, indicating where the head "looks".

The QK-circuit is fundamentally a search and matching mechanism.

2. The OV-Circuit (Content Movement)

The OV-circuit's job is to determine what information is moved from a source token t_v and how it is written to the destination token's residual stream.

• Input: The residual stream vector x_v from a source token.
• Matrices: W_V^h (Value) of shape [d_model, d_head] and the head's effective output matrix W_O^h of shape [d_head, d_model].
• Computation: For a single source token t_v, the "information payload" it offers is the projected value:
  
  $\text{Payload}_v = (\mathbf{x_v} \mathbf{W_V^h}) \mathbf{W_O^h}$
• Output: A d_modeldimensional vector. This is the update that will be added to the destination token's residual stream if this token receives 100% of the attention.

The OV-circuit is a feature extraction and transformation mechanism. The composed matrix W_V^h W_O^h can be seen as a single [d_model, d_model] operation that specifies "if you attend to this token, this is the update you will get."

Walkthrough

Step 1: Intuition in a Case Study

Consider the prompt "The Roman Empire fell in 476 AD. This date..." and a hypothetical head in a later layer processing the token "date".

Step 2: How to Inspect These Circuits

This decomposition isn't just a metaphor; it gives us concrete matrices to inspect.

• To analyze the QK-Circuit: We can study the matrix W_QK = W_Q^T W_K. This [d_model, d_model] matrix reveals the structure of attention patterns. A large positive value at (i, j) in this matrix means that if feature i is present at the query token and feature j is present at the key token, the attention score will be high. This allows us to find patterns like "attend to the previous token" or "attend to tokens that are nouns".
• To analyze the OV-Circuit: We study the composed matrix W_OV = W_V W_O. This [d_model, d_model] matrix tells us what kind of information processing the head performs. We can ask: if the input vector x_v has a feature in direction d_1, what features does the output vector (x_v W_{OV}) contain? For a "copying" head, this matrix would look similar to the identity matrix, meaning it moves information without much transformation. For other heads, it might map "proper noun" features to "capital city" features.

Implications & Limits

• Implication: A Taxonomy of Heads. This separation allows us to categorize heads. For example, many models learn "previous token heads" whose QK circuit reliably attends to the token at position - 1. However, these heads can have very different OV circuits: one might copy the previous token's embedding, while another might check if the previous token was a verb.
• Limit: Not Fully Independent. The circuits are optimized jointly during training. The gradients flowing back to W_Q and W_K depend on how useful the information moved by the OV-circuit was for the final task. A QK-circuit has no reason to create a sharp, meaningful attention pattern if the corresponding OV-circuit provides useless information. They co-evolve.

Pitfalls

• Analyzing W_V in isolation: A common error is to just look at W_V to understand what information is being moved. W_V projects into the d_headdimensional head-space, which is often uninterpretable. You must compose it with W_O to see the head's full effect on the d_modeldimensional residual stream.
• Over-interpreting Q/K vectors: It is tempting to think of q and k vectors as having rich semantic meaning. It's often more accurate to think of them as specialized pointers or search keys, designed only for the purpose of dot-product matching. The actual semantic content is carried by the v vectors.

Takeaways

• An attention head can be productively decomposed into a QK-circuit (where to look) and an OV-circuit (what to move).
• The QK-circuit's behavior is summarized by the composed matrix W_Q^T W_K.
• The OV-circuit's behavior is summarized by the composed matrix W_V W_O.
• This decomposition provides a powerful framework for categorizing and analyzing the function of individual heads.
• It simplifies the problem of understanding a head into two smaller, more focused problems.
• While analytically useful, the two circuits are not truly independent as they are co-optimized during training.