Part 3 Reverse Engineering Transformers: Path Expansion

💡

Other parts in the series:
Part 1: Reverse Engineering Transformers: Deconstructing Attention
Part 2: Reverse Engineering Transformers: Attention Heads & Circuits
Part 4 Reverse Engineering Transformers: Induction Heads

Transformer Circuits IV: Path Expansion

TL;DR

• Path expansion is a technique for decomposing a Transformer's full forward pass into a sum of individual end-to-end paths.
• A "path" is a sequence of matrix multiplications representing information flow from an input token, through a series of attention heads and FFN layers, to a final output logit.
• The total logit for a specific token is the sum of the contributions from all possible paths leading to it.
• For this linear decomposition to work, we must ignore the non-linear LayerNorm operations, treating path expansion as a useful "linear approximation" of the true computation.
• This technique allows us to move from analyzing what a single component does to analyzing how components compose into circuits that solve a task.

Context & Motivation

So far in our series, we have established a powerful mental model:

We have successfully broken the model down into its smallest functional parts. The next logical question is: how do these parts combine across layers to implement complex algorithms? How does a "previous token head" in Layer 0 enable an "induction head" in Layer 5?

The answer lies in path expansion. By recursively expanding the definition of the residual stream, we can re-express the entire model's computation for a single output logit as a massive sum of individual end-to-end information paths. This allows us to attribute the model's final output to specific combinations of components, which we call "circuits."

Prereqs

• Understanding of the residual stream, independent heads, and QK/OV circuits (Parts I-III).
• Comfort with the idea of matrix composition (multiplying matrices in a sequence).

Core Idea in One Picture

A Transformer's output for a token is the sum of all the ways information can travel from the input embeddings to the final unembedding layer.

graph TD
    subgraph Input Embeddings
        T_A[Token A]
        T_B[Token B]
    end

    subgraph Layer 0
        H0_1[Head 0.1]
        F0[FFN 0]
    end

    subgraph Layer 1
        H1_1[Head 1.1]
        F1[FFN 1]
    end

    subgraph Output Logits
        U[Unembedding] --> Logit_C[Logit for 'C']
    end

    T_A -- Path 1 --> U
    T_A -- Path 2 --> H0_1 -- Path 2 --> U
    T_B -- Path 3 --> F0 -- Path 3 --> H1_1 -- Path 3 --> U

    style Path 1 stroke:#ff0000,stroke-width:2px,stroke-dasharray: 5 5
    style Path 2 stroke:#0000ff,stroke-width:2px,stroke-dasharray: 5 5
    style Path 3 stroke:#008000,stroke-width:2px,stroke-dasharray: 5 5```
*Figure 1: A conceptual illustration of path expansion. The final logit for token 'C' is a sum of contributions from many paths. Path 1 is a direct connection. Path 2 goes through one head. Path 3 goes from a different token through an FFN and another head. Each path corresponds to a product of weight matrices.*

Definitions & Setup

Let's trace the computation algebraically. The final logit for a specific token T is computed by taking the dot product of the final residual stream state x_final with the token's corresponding row in the unembedding matrix, W_U[T].

Logit(T) = x_final @ W_U[T]

We know that x_final is the sum of the initial embedding x_0 and all component updates:
x_final = x_0 + ∑_{l, h} Δ_{l,h}^{attn} + ∑_{l} Δ_{l}^{ffn}

If we ignore LayerNorm, we can substitute this in:
Logit(T) = (x_0 + ∑ Δ_{l,h}^{attn} + ∑ Δ_{l}^{ffn}) @ W_U[T]

Since the dot product distributes over addition, this becomes:
Logit(T) = (x_0 @ W_U[T]) + ∑ (Δ_{l,h}^{attn} @ W_U[T]) + ∑ (Δ_{l}^{ffn} @ W_U[T])

This is the first level of expansion. We've expressed the logit as a sum of contributions from each component's direct write to the final residual stream. But now we can expand each Δ term, because each component's output depends on the state of the stream before it, which is itself a sum of previous terms. This recursive expansion is path expansion.

A path is a term in this fully expanded sum. Each path corresponds to a unique sequence of components that information flows through.

Walkthrough: Expanding Two Simple Paths

Let's analyze the paths for predicting the token 'C' that follows the prompt "A B".

Path 1: The Direct "Bigram" Path

The simplest path is the information flowing directly from an input token's embedding to the output.

• Path: Embedding('B') → Unembedding('C')
• Description: This path represents the model's learned bigram statistics. How likely is 'C' to follow 'B' without any deeper context?
• Algebraic Contribution: The contribution to the logit of 'C' from this path is:
  
  \text{Contribution}1 = (\mathbf{W_E}[\text{'B'}] + \mathbf{W{pos}}) \cdot \mathbf{W_U}[\text{'C'}]
  
  This is a single scalar value, computed from a product of two embedding vectors. This path exists for every input token.

Path 2: The "Skip-Trigram" Attention Path

Now consider a path that involves one attention head. Let's trace the information from token 'A' to the output at position 'B', which then predicts 'C'.

• Path: Embedding('A') → L0 Head OV Circuit → Unembedding('C')
• Description: This path models skip-trigrams (e.g., in "The dog ran", it helps predict "ran" based on "The" after seeing "dog"). The head at position 'B' attends to 'A' and moves information to 'B's residual stream.
• Algebraic Contribution: Let's assume head h in layer 0 at position q=B is attending to source position k=A.
  • The information moved by the OV circuit is (x_A @ W_V^h) @ W_O^h.
  • This is projected by the unembedding: ((x_A @ W_V^h) @ W_O^h) @ W_U[C].
  • The full contribution is this value multiplied by the attention probability A_{B ← A} paid from B to A.
    
    \text{Contribution}2 = A{B \leftarrow A} \cdot \left( (\mathbf{W_E}[\text{'A'}] + \mathbf{W_{pos}}) \mathbf{W_V^h} \mathbf{W_O^h} \mathbf{W_U}[\text{'C'}] \right)
    
    This shows how the matrices along a path compose. The term W_V^h W_O^h W_U forms a single composed operator that describes the end-to-end effect of this path.

The full logit for 'C' is Contribution_1 + Contribution_2 + contributions from all other paths (from 'B' through heads, from 'A' through FFNs, through multiple layers, etc.).

Implications & Limits

• Implication: Finding Circuits. Path expansion is the theoretical tool that allows us to find circuits. The "Indirect Object Identification" circuit, for instance, is a set of paths involving specific heads that compose to perform that task. By summing the contributions of just the paths in a hypothesized circuit, we can see how much of the model's behavior on a specific task that circuit explains.
• Implication: Explaining Logits, Not Probabilities. This is an analysis of the pre-softmax logits. Logits are additive, so this decomposition is mathematically sound (ignoring LayerNorm). You cannot do this with probabilities, which are the output of the non-linear softmax function. softmax(a+b) != softmax(a) + softmax(b).
• Limit #1: LayerNorm Breaks Everything. As stated before, this entire framework relies on the linear additivity of the residual stream. LayerNorm is non-linear and breaks this assumption. LN(x_0 + Δ_1) is not LN(x_0) + LN(Δ_1). Therefore, path expansion is a linear approximation. For many circuits and models, this approximation holds up surprisingly well, but it is not the ground truth.
• Limit #2: Combinatorial Explosion. The number of paths is exponential in the depth of the model. Fully expanding all paths for a large model is computationally intractable. In practice, researchers analyze the top-K most important paths or analyze contributions on a component-by-component basis, which is a coarser-grained version of path expansion.

Pitfalls

• Forgetting the Linear Approximation. It is the most common mistake to treat path expansion results as the exact ground truth of the model's computation. Always remember the LayerNorm caveat.
• Summing Probabilities. Never attempt to sum the "probability contributions" of different paths. The analysis is valid only for the additive logits.

Takeaways

• Path expansion reframes a Transformer's computation as a sum over end-to-end information flow paths.
• Each path has a value (its contribution to a final logit) and is composed of a product of weight matrices.
• This technique is the key to moving from analyzing components to analyzing multi-component circuits.
• It provides a granular way to attribute a model's output to specific model parameters and interactions.
• The validity of simple path expansion rests on a linear approximation that ignores LayerNorm.
• This analysis applies to logits, not probabilities, due to the additivity requirement.

Transformer Circuits V: The Power of Depth is Composition

TL;DR

• The power of deep Transformers comes from composition, where components in later layers operate on the outputs of components in earlier layers.
• An earlier head (H0) can influence a later head (H1) in three ways: Q-Composition, K-Composition, and V-Composition.
• Q-Composition: The output of H0 affects H1's Query vector, changing what H1 is looking for.
• K-Composition: The output of H0 affects H1's Key vector, changing how H1's source tokens appear to other queries.
• V-Composition: The output of H0 affects H1's Value vector, changing what information H1 moves. This can create powerful "virtual attention heads."
• These composition mechanisms are the building blocks of complex circuits, like the famous "induction heads."

Context & Motivation

In our last post on Path Expansion, we established that a Transformer's output can be viewed as a massive sum of information paths. However, this view can be misleading if it suggests paths are just independent, parallel computations. The most interesting behaviors in Transformers arise when these paths interact-when the output of one component directly becomes the input to another.

A single attention head in Layer 0 is limited. It can only implement simple algorithms like "attend to the previous token" or "attend to nouns." To build a program that does something complex, like "find the last time the current token appeared, and copy the token that came after it," you need multiple steps. You need depth.

Composition is the formal term for how these steps are chained together. Specifically, we'll examine how an attention head in an earlier layer (H0) can pass information through the residual stream to be read by a head in a later layer (H1), fundamentally altering H1's behavior.

Prereqs

• Solid understanding of the residual stream, independent heads, QK/OV circuits, and the concept of path expansion (Parts I-IV).

Core Idea in One Picture

A head in Layer 0 (H0) writes an update to the residual stream. A head in Layer 1 (H1) then reads from this updated stream. Depending on which part of H1 reads the update (its Query, Key, or Value projection), we get a different type of composition.

graph TD
    subgraph Layer 0
        x0[Residual Stream @ L0] --> H0[Head H0]
        H0 -- Δ_0 --> x0_updated((+))
    end

    subgraph Layer 1
        x1[Residual Stream @ L1] --> H1[Head H1]
    end

    x0 --> x0_updated
    x0_updated --> x1

    subgraph H1
        direction LR
        x1 -- W_Q^1 --> Q1[Query]
        x1 -- W_K^1 --> K1[Key]
        x1 -- W_V^1 --> V1[Value]
    end

    linkStyle 0 stroke-width:0px
    linkStyle 4 stroke-width:0px

    subgraph Legend
        QComp[Q-Composition]
        KComp[K-Composition]
        VComp[V-Composition]
    end

    H0 -- "affects" --> Q1;
    H0 -- "affects" --> K1;
    H0 -- "affects" --> V1;

    style H0 fill:#ccf
    style H1 fill:#fcc
    style QComp fill:#f9d,stroke:#333
    style KComp fill:#d9f,stroke:#333
    style VComp fill:#9df,stroke:#333

Figure 1: The three types of head composition. An earlier head, H0, writes an update (Δ_0) to the stream. A later head, H1, reads this update. If the update influences H1's Query, it's Q-Composition; if its Key, K-Composition; if its Value, V-Composition.

Definitions & Setup

For simplicity, let's consider a two-layer, attention-only model. The residual stream entering Layer 1 is the sum of the initial embeddings (x_0) and the output of Layer 0 (Δ_{L0}), which is the sum of all Layer 0 head outputs. Let's focus on the interaction between one head H0 and one head H1.

The stream after H0 is x_1 = x_0 + Δ_0 (again, ignoring LayerNorm for our linear analysis). H1 now computes its Q, K, and V vectors from this x_1.

Walkthrough: What Each Composition Type Does

Step 1: Q-Composition ("Changing What to Look For")

Imagine H0 is a "proper noun detector." When it sees a name like "Mary," it writes a +is_a_proper_noun feature into the residual stream at that position. H1 might be a head that wants to find the main verb of a sentence. Through Q-composition, H1 can learn to modify its query based on H0's output. Its query at the "Mary" token, influenced by Δ_0, might become something like "I am a proper noun, now I am looking for the verb I am the subject of." This allows H1's search to be context-dependent.

Step 2: K-Composition ("Changing How You Are Seen")

This is subtler but crucial. H0's output modifies the key of a token, changing how it "appears" to queries from other positions. The canonical example is in forming induction heads.

• Let H0 be a previous token head. At every position t, it attends to position t-1 and copies its embedding. So, Δ_0 at position t contains information about token t-1.
• Now consider H1, the induction head. It wants to find a previous occurrence of the current token. Let's say we're at token[t] = 'B' and we've seen "... A B ..." earlier in the sequence. H1 needs to attend to token[t-1] = 'A'.
• How? At position t-1, H0 wrote information about token[t-2] = 'A'. H1's key at t-1 is thus modified by H0 to effectively broadcast "the token after me is 'A'". H1's query at t is looking for the token 'A'. The QK-dot-product will match, directing H1's attention to the correct previous token.

Step 3: V-Composition ("Changing What You Move")

This is perhaps the most powerful type, creating what the paper calls virtual attention heads. Here, H0 determines what information H1 will move. H1's attention pattern might direct it to attend to token A, but the information it moves can come from token B, as mediated by H0.

Let's expand the algebra for a V-composition path, where H1 at position q attends to position k_1, and H0 at k_1 attends to k_0.

• H1's output: Output_1 = A_{q ← k_1} · (x_{k_1} @ W_V^1 @ W_O^1)
• But x_{k_1} = x_0 + Δ_0 = x_0 + (A_{k_1 ← k_0} · (x_{k_0} @ W_V^0 @ W_O^0))
• Plugging this in, the contribution to H1's output from this two-head path is:
  
  $\text{Path Output} = A_{q \leftarrow k_1} \cdot A_{k_1 \leftarrow k_0} \cdot ( \mathbf{x_{k_0}} @ \mathbf{W_V^0 W_O^0 W_V^1 W_O^1} )$
  
  This path moves information from x_{k_0} to position q via a two-step hop. It attends from q to k_1 and k_1 to k_0, but the matrices W_V^0 W_O^0 W_V^1 W_O^1 form a a new, composed OV circuit that operates directly on the information at x_{k_0}. We've created a "virtual head" that has the attention pattern of H1 but an OV circuit defined by both heads.

Implications & Limits

• Implication: Head composition is the primary mechanism for building complex, multi-step algorithms in Transformers. It explains how models can perform tasks that are impossible for a single layer.
• Implication: It allows us to reason about models as computational graphs where heads are nodes and composition links are edges, leading to a much richer understanding of model behavior.
• Limit: This clean decomposition is, as always, complicated by LayerNorm, which mixes all parts of the residual stream together non-linearly. The real interactions are not quite as simple as adding up path contributions.
• Limit: FFN/MLP layers also compose with heads and with each other. A full circuit analysis must account for their role, which is often less clear than that of attention heads.

Pitfalls

• Viewing heads as independent agents. H1 and H0 are not collaborating consciously. They are a single system optimized end-to-end. H1 learns to use H0's output because the joint behavior is useful for minimizing the loss function. There is no guarantee H0's "purpose" is cleanly separable from H1's.
• Assuming one type of composition per head pair. A single head pair (H0, H1) can and often does exhibit all three types of composition simultaneously. H0's output can influence H1's Q, K, and V paths all at once.

Takeaways

• Composition is how Transformers build complex algorithms from simple components, and it is the key reason for the power of model depth.
• The three types of composition-Q, K, and V-provide a vocabulary for describing how heads in different layers interact.
• Q-Composition allows a head's search query to be context-dependent.
• K-Composition allows a token's representation for search purposes to be context-dependent, enabling induction heads.
• V-Composition allows a head to move information from sources it isn't directly attending to, creating powerful "virtual heads."
• Analyzing these compositional circuits is a core activity in mechanistic interpretability.