AI · Tech · Science · Crypto · Linux · Gaming · DIY · Guides
🤖 AI · AI

A Mathematical Framework for Transformer Circuits (2021)

4567 words · 22 min read

A Mathematical Framework for Transformer Circuits (2021)

In December 2021, a team at Anthropic published a paper that gave researchers something they had been missing: a way to look inside a transformer and describe what they found using math instead of metaphor. "A Mathematical Framework for Transformer Circuits" did not explain how GPT-3 works. It did something more useful. It built a vocabulary—residual streams, QK circuits, OV circuits, induction heads—that turned transformer interpretability from a vague aspiration into an engineering discipline.

This deep-dive walks through the paper's core ideas, the math behind them, and the code you can write to reproduce its central findings. We'll work through how attention heads decompose into circuits you can read, how two attention layers compose into an induction head, and why superposition makes all of this harder than it first appears.


Introduction: The Quest to Reverse-Engineer Transformers

The Black-Box Problem in Large Language Models

By 2021, transformer language models were producing fluent text, answering questions, and completing code. Nobody could say precisely why. We had architecture diagrams and loss curves, but no account of the internal computation that produced a given output. Interpretability research at the time leaned on probing classifiers, attention visualizations, and input-output correlations. These tools described behavior. They did not describe mechanism.

The gap mattered. If you can't explain why a model produces an output, you can't reliably fix it when it's wrong, predict when it will fail, or verify that it's not using a shortcut you'd object to. The field needed methods that treated the model as a program to be reverse-engineered rather than a function to be approximated.

Anthropic's Transformer Circuits Thread and the December 2021 Breakthrough

Anthropic launched the Transformer Circuits Thread as a venue for mechanistic interpretability work—research aimed at understanding the internal algorithms of neural networks. The first paper on the thread, "A Mathematical Framework for Transformer Circuits," appeared in December 2021 and is also available as arXiv:2112.04510.

The authors—Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, and Jared Kaplan—brought a specific stance: decompose the model into components small enough to understand, then describe how those components compose.

Why a Mathematical Framework Is Essential for Mechanistic Interpretability

Interpretability without formalism tends to produce stories. A researcher looks at an attention pattern, sees something that looks like syntax, and writes it up. The next researcher sees something else. Nothing accumulates.

A mathematical framework changes that. When you can write an attention head as a product of matrices and show that one part determines where it attends while another determines what it moves, you get claims that can be checked, refuted, and built upon. The paper's contribution is not a single discovery about a single model. It's a set of definitions and derivations that make claims about transformer internals precise.

Overview of the Paper's Core Contributions and Impact

The paper makes five contributions that have shaped subsequent work:

  1. The residual stream as a communication channel. Every layer reads from and writes to a shared residual stream via addition. This reframes the transformer as a set of independent modules communicating through a common bus.
  2. QK and OV circuit decomposition. Each attention head factors into a circuit that determines attention patterns and a circuit that determines what information moves.
  3. Circuits as computational subgraphs. Specific behaviors can be attributed to compositions of heads and MLPs.
  4. Induction heads. A two-layer attention-only transformer can implement induction—copying with a shift—through composition of two heads.
  5. Superposition. Models can represent more features than they have dimensions, which complicates naive interpretation.

The paper has been cited over 500 times as of 2024 and spawned follow-ups including "In-context Learning and Induction Heads" and "Toy Models of Superposition," both also published on the Transformer Circuits Thread.

Key Takeaway: The framework's value is not that it explains large models end-to-end. It's that it gives researchers a shared, checkable language for describing internal computation—one that subsequent work could extend rather than reinvent.


The Transformer as a Residual Stream Computer

The Residual Stream as the Central Communication Channel

The paper's first move is conceptual. Instead of thinking of a transformer as a stack of layers that transform a hidden state, think of it as a central memory—the residual stream—that every sublayer reads from and writes to.

At each position, the residual stream is a vector of dimension $d_{model}$. It starts as the sum of token and positional embeddings. Every attention head and every MLP layer takes the current stream as input, computes something, and adds its output back.

How Attention Heads and MLP Layers Read From and Write to the Stream

An attention head reads from the residual stream at every position by projecting it through its query, key, and value matrices. It writes to the stream at each position by projecting its output through the output matrix and adding the result.

An MLP layer reads from the stream at each position independently, applies a nonlinearity, and writes back. The crucial structural fact: all writes are additions. Nothing overwrites. Nothing replaces. The stream accumulates.

Mathematical Formulation: Updates as Vector Additions

If we let $x_i^{(l)}$ denote the residual stream at position $i$ after layer $l$, then:

$$x_i^{(l+1)} = x_i^{(l)} + \sum_{h} \text{Attn}_h^{(l)}(x^{(l)})_i + \text{MLP}^{(l)}(x_i^{(l)})$$

Each attention head's contribution is a vector added to the stream. Each MLP's contribution is a vector added to the stream. The stream is the sum of all prior contributions.

Implications for Information Flow and Gradient Paths

This structure has two important consequences.

First, information flows additively. A head in layer 5 can read anything written by any head in layers 1–4, plus the original embeddings. There's no bottleneck where information gets squashed into a fixed-size hidden state that then gets overwritten.

Second, gradients flow cleanly. The residual connection means gradients can propagate directly from the loss to early layers without passing through every intermediate nonlinearity. This is part of why deep transformers train.

Common Misconception: The Residual Stream Is Not a Simple Sum—It's a Rich, Structured Representation

It's tempting to read "sum of contributions" as "unstructured pile." That's wrong. The stream is a high-dimensional vector space in which different directions carry different information. A head that writes a specific direction is placing a specific fact or feature into the stream where other heads can read it. The additive structure is what makes this possible: a head can write a small perturbation to a specific subspace without disturbing everything else.

Key Takeaway: Treating the residual stream as a shared communication bus—rather than a per-layer hidden state—makes the transformer's internals far more tractable. Reads and writes are linear projections; composition is addition.


Decomposing Attention Heads: QK and OV Circuits

The Anatomy of an Attention Head: Query, Key, Value, Output Matrices

An attention head has four learned weight matrices: $W_Q$, $W_K$, $W_V$, and $W_O$. For a head with head dimension $d_{head}$ operating on a residual stream of dimension $d_{model}$, each of these is a linear map.

Given a residual stream $X \in \mathbb{R}^{n \times d_{model}}$ for a sequence of length $n$:

$$Q = X W_Q, \quad K = X W_K, \quad V = X W_V$$

Attention scores are computed as $QK^T / \sqrt{d_{head}}$, passed through a softmax, and used to weight $V$. The result is then projected by $W_O$ and added back to the stream.

QK Circuit: Computing Attention Patterns (Where to Attend)

The attention pattern depends only on $Q$ and $K$:

$$A = \text{softmax}\left(\frac{X W_Q W_K^T X^T}{\sqrt{d_{head}}}\right)$$

The product $W_Q W_K^T$ is what the paper calls the QK circuit. It's a matrix of shape $d_{model} \times d_{model}$ that determines which pairs of positions attend to each other. To understand where a head looks, you analyze the QK circuit.

OV Circuit: Determining Information Movement (What to Copy)

The information that flows from an attended position to the current position depends on $W_V$ and $W_O$:

$$\text{Output}i = \sum_j A (x_j W_V W_O)$$

The product $W_V W_O$ is the OV circuit. It's a $d_{model} \times d_{model}$ matrix that determines what transformation is applied to the value vector before it's added to the stream. To understand what a head moves, you analyze the OV circuit.

Mathematical Derivation: W_QK = W_Q^T W_K and W_OV = W_O W_V

The paper writes the QK circuit as $W_{QK} = W_Q^T W_K$ and the OV circuit as $W_{OV} = W_O W_V$ (using the convention that weight matrices act on column vectors). The exact index conventions vary, but the substance is fixed: both circuits are products of two learned matrices, and both operate on the residual stream directly.

This factorization is the paper's central technical move. It says: don't analyze a head as a black box. Analyze two separate circuits. One tells you where the head attends; the other tells you what it moves.

Interpreting Head Behavior Through Circuit Analysis

If the QK circuit has high similarity between position $i$ and position $j$ when token $j$ is the previous token, the head attends to previous tokens. If the OV circuit maps a token's embedding to a direction representing its part of speech, the head moves part-of-speech information.

You can read these properties off the matrices. No training, no probing, no input-output experiments required.

Example: Analyzing a Copying Head via QK and OV Circuits

Consider a head that copies the token at the previous position. Its QK circuit should give high scores for the pair (current, previous). Its OV circuit should approximately map the value of a token back to its own embedding direction—an identity-like operation on the relevant subspace.

A copying head is the simplest non-trivial circuit. It's also a building block for more complex behaviors, including induction.

Key Takeaway: Factoring an attention head into QK and OV circuits separates "where" from "what." Each circuit is a product of two learned matrices operating on the residual stream, and both can be analyzed directly.


Circuits: Computational Subgraphs in Transformers

Definition: Circuits as Subgraphs Implementing Specific Behaviors

A circuit is a subgraph of the model's computational graph—a subset of heads, MLPs, and connections between them—that implements a specific behavior. Circuits are the paper's unit of explanation. Instead of saying "the model does X," you say "this circuit does X."

The residual stream makes circuits possible. Because every head can read from the stream and every head writes to it, heads in later layers can use information written by heads in earlier layers. Composition is the mechanism by which simple circuits become complex behaviors.

Composition of Attention Heads Across Layers

There are several ways heads compose. The most important for the paper's examples is Q-composition and K-composition: a head in layer 2 uses information written by a head in layer 1 to construct its queries or keys. This lets layer 2 heads attend based on properties that layer 1 heads computed.

The Role of the Residual Stream in Enabling Composition

Composition works because of the residual stream. If layer 1 writes a feature into the stream, layer 2 can read it via its own Q, K, or V projections. The stream is the medium through which circuits span layers.

Case Study: The Induction Head Circuit in a Two-Layer Attention-Only Transformer

The paper's flagship example is the induction head. In a two-layer attention-only transformer, an induction head can be formed by composing a previous token head in layer 1 with a matching head in layer 2.

The behavior: given a sequence like ... A B ... A, the model should predict B at the final A. This is copying with a shift—attend to the token that followed the previous occurrence of the current token.

Step-by-Step Breakdown: Previous Token Head + Matching Head = Induction

Layer 1: previous token head. This head attends to the position immediately before the current one. For each position $i$, it writes into the residual stream information about the token at position $i-1$.

Layer 2: matching head. This head uses the information written by layer 1 to construct its queries and keys. Its query at position $i$ encodes the current token. Its key at position $j$ encodes the token that followed position $j$—because layer 1 wrote the previous-token information into the stream, and layer 2's key projection can use it. The head attends to the position $j$ where the token after $j$ matches the current token.

Result. At the final A, the model attends to the B that followed the previous A, and copies it. The prediction is B.

Two heads, one circuit, one behavior. Neither head alone implements induction. The composition does.

Why Induction Heads Are Crucial for In-Context Learning

Induction is a general pattern: find a previous occurrence of the current context and copy what followed. It's the simplest mechanism that produces in-context learning—the ability to use examples in the prompt to guide predictions. The follow-up paper "In-context Learning and Induction Heads" (also on the Transformer Circuits Thread) developed this connection further.

Key Takeaway: Circuits are the paper's unit of explanation. The induction head is the canonical example: two heads in different layers compose via the residual stream to implement copying-with-a-shift, a behavior that neither head implements alone.


Superposition: When Features Overlap

The Challenge: Representing More Features Than Dimensions

A residual stream of dimension $d_{model}$ can represent at most $d_{model}$ orthogonal directions. But models need to track far more than $d_{model}$ features—syntax, semantics, position, task context, and more. The paper confronts this directly: models can represent more features than dimensions by using non-orthogonal directions.

Mathematical Definition: Non-Orthogonal Feature Directions

If a model represents feature $i$ as a direction $v_i$ in the residual stream, and two features have $v_i \cdot v_j \neq 0$, they interfere. In superposition, the model accepts this interference in exchange for representing more features.

How Superposition Complicates Interpretability

If features are not orthogonal, you can't read them off by projecting onto individual dimensions. A single dimension may carry a mix of several features. A single feature may be spread across several dimensions. Naive interpretation—"this dimension means X"—breaks down.

Trade-Offs: Efficiency vs. Interpretability

Superposition is a trade-off. The model gets more representational capacity per dimension. Interpretability gets harder. The paper doesn't resolve this trade-off; it names it and provides tools for analyzing it.

Techniques to Disentangle Superposed Features (e.g., Sparse Autoencoders)

The paper points toward techniques for recovering features from superposed representations. Subsequent work developed sparse autoencoders as a practical method: train an overcomplete sparse dictionary to reconstruct activations, and the dictionary atoms tend to correspond to interpretable features.

Misconception: Superposition Is Not Always Detrimental—It's a Feature, Not a Bug

Superposition is often framed as a problem for interpretability. It's also an efficient use of capacity. A model that represents 10,000 features in 1,000 dimensions is doing something clever. The goal is not to eliminate superposition but to understand when it occurs and how to recover features from it.

Key Takeaway: Superposition means features can overlap in the residual stream. This makes naive "one dimension, one feature" interpretation unreliable, but it also explains how models achieve high capacity. Sparse autoencoders are one practical route to disentangling superposed features.


MLP Layers as Key-Value Memories

The MLP as a Memory Mechanism Reading From and Writing to the Residual Stream

An MLP layer reads the residual stream at each position, applies a nonlinearity, and writes back. The paper frames this as a key-value memory: the input weights $W_{in}$ act as keys that detect patterns in the stream, and the output weights $W_{out}$ act as values that write associated information back.

Mathematical Formulation: W_in and W_out as Keys and Values

For an MLP with hidden dimension $d_{ff}$:

$$\text{MLP}(x) = \text{ReLU}(x W_{in}) W_{out}$$

Each row of $W_{in}$ is a key vector. When the input $x$ has high dot product with a key, the corresponding hidden unit activates. The corresponding row of $W_{out}$ is written back to the stream. This is a pattern-matching memory: keys detect, values write.

How MLPs Interact With Attention Heads in Circuits

MLPs participate in circuits the same way attention heads do. They read from the stream and write to it. A head in layer 3 can read information an MLP in layer 2 wrote. Composition across the stream is uniform across head types.

Example: MLP Layers in Factual Recall and Pattern Completion

A common pattern: an MLP detects a subject entity in the stream and writes a fact associated with it—the entity's capital city, for example. A later head can then read that fact and use it in a prediction. This is a plausible mechanism for factual recall, though the paper treats it as an illustration rather than a proven account of any specific model.

Extending the Framework Beyond Attention-Only Models

The paper's detailed analysis focuses on attention-only transformers. MLPs are included in the framework—they're just another module reading from and writing to the stream—but the paper's concrete examples (induction heads, copying heads) don't require them. Extending the analysis to models with MLPs is a natural next step, and subsequent work has taken it.

Key Takeaway: MLP layers fit the same framework as attention heads. $W_{in}$ acts as keys that detect patterns; $W_{out}$ acts as values that write information back to the stream. Circuits can span heads and MLPs alike.


The Framework in Action: Analyzing a Two-Layer Transformer

Experimental Setup: A Small Attention-Only Transformer

The paper's concrete experiments use a two-layer attention-only transformer trained on a simple task. No MLPs, no layer norm complications beyond what's needed, small enough to analyze by hand. The goal is not performance but clarity: can we identify the circuit that implements a behavior?

Identifying Induction Heads Through Circuit Analysis

Train the model on sequences with repeated subsequences. The model learns to predict the token that followed a previous occurrence of the current context. Inspect the attention patterns: layer 1 heads attend to previous tokens; layer 2 heads attend to earlier occurrences of the current token's context. That's the induction circuit.

Visualizing Attention Patterns and OV Projections

Attention patterns are matrices of shape $n \times n$ showing where each position attends. Plotting them reveals the previous-token and matching behaviors directly. OV projections—the products $W_V W_O$—show what information moves. For an induction head, the OV projection should approximately copy the value of the matched token.

Code Walkthrough: Implementing QK/OV Decomposition in PyTorch

Here's a minimal implementation of the decomposition:

import torch
import torch.nn.functional as F

def qk_circuit(W_Q, W_K):
    """Compute the QK circuit: W_Q^T W_K (shape d_model x d_model)."""
    return W_Q.T @ W_K

def ov_circuit(W_V, W_O):
    """Compute the OV circuit: W_O W_V (shape d_model x d_model)."""
    return W_O @ W_V

def attention_pattern(X, W_Q, W_K, d_head):
    """Compute attention pattern for a sequence X (n x d_model)."""
    Q = X @ W_Q          # n x d_head
    K = X @ W_K          # n x d_head
    scores = Q @ K.T / (d_head ** 0.5)
    return F.softmax(scores, dim=-1)

def head_output(X, W_Q, W_K, W_V, W_O, d_head):
    """Full attention head output, added to residual stream."""
    A = attention_pattern(X, W_Q, W_K, d_head)
    V = X @ W_V          # n x d_head
    return A @ V @ W_O   # n x d_model

# Inspect the QK circuit to see where the head attends.
# Inspect the OV circuit to see what it moves.

The QK circuit is a $d_{model} \times d_{model}$ matrix. Its eigenvectors and dominant directions tell you what token properties the head is sensitive to. The OV circuit tells you what transformation the head applies to the attended value before writing back.

Results: How Induction Emerges From Composition

In the trained model, layer 2 heads show attention patterns that depend on information written by layer 1. Ablating the layer 1 previous-token head breaks the layer 2 induction behavior. The circuit is necessary and sufficient for the observed behavior in this small setting.

Limitations: Scaling to Deeper Models and the Challenge of Superposition

The paper is explicit about limits. Two-layer attention-only models are tractable. Real language models are deeper, wider, and full of superposition. Circuits in larger models may be harder to isolate because features overlap and because many heads may contribute to the same behavior. The framework is a starting point, not a finished tool.

Key Takeaway: The two-layer analysis demonstrates that the framework can identify a complete circuit—previous token head plus matching head—for a specific behavior. Scaling this to production models remains an open problem.


Impact and Legacy: Shaping Mechanistic Interpretability

Immediate Influence: Follow-Up Papers

Two follow-ups on the Transformer Circuits Thread extended the work directly. "In-context Learning and Induction Heads" connected the induction circuit to in-context learning at scale. "Toy Models of Superposition" turned the superposition discussion into a concrete research program, training small models on synthetic tasks to study when and how features overlap.

Adoption by the Broader Research Community

The paper has been cited over 500 times as of 2024. The vocabulary it introduced—residual stream, QK circuit, OV circuit, induction head, superposition—is now standard in mechanistic interpretability. When researchers describe a head as "attending to previous tokens," they're using the paper's framing whether they cite it or not.

Foundation for Sparse Autoencoders and Circuit-Level Analysis

Sparse autoencoders, now a central tool in interpretability, grew out of the superposition problem the paper named. Circuit-level analysis—identifying subgraphs responsible for specific behaviors—is the paper's core methodological proposal, and it has become a standard approach.

Open Questions: Scaling Interpretability to Large Language Models

The paper's framework was developed on small models. Scaling it to production-scale language models remains the field's central challenge. Superposition is more severe, circuits are more tangled, and the number of components is larger by orders of magnitude. Progress since 2021 has been real but partial.

The Transformer Circuits Thread as a Collaborative Research Hub

The Transformer Circuits Thread functions as an ongoing publication venue for mechanistic interpretability. It hosts the original paper, the follow-ups, and a growing body of related work. Its existence is part of the paper's legacy: it created a place for this research to accumulate.

Key Takeaway: The paper's influence is visible in the vocabulary the field now uses, the follow-up research it enabled, and the research venue it helped establish. Its limitations—small models, superposition, scaling—define much of the work that followed.


Conclusion: A Roadmap for Reverse-Engineering Intelligence

Summary of Key Contributions

The paper delivers five things: a reframing of the transformer as a residual stream computer, a decomposition of attention heads into QK and OV circuits, a definition of circuits as computational subgraphs, a concrete demonstration that two heads can compose into an induction circuit, and a clear statement of the superposition problem.

Why This Framework Is a Starting Point, Not a Complete Solution

The paper does not explain how GPT-3 works. It explains how to think about transformer internals in a way that makes explanation possible. The two-layer analysis is a proof of concept, not a general method. Superposition, depth, and scale remain unsolved.

Future Directions

Three directions stand out. First, scaling circuit analysis to deeper models—finding circuits in models with dozens of layers rather than two. Second, disentangling superposition—recovering features from overlapping representations, which sparse autoencoders attempt. Third, automating circuit discovery—replacing hand analysis with methods that find circuits programmatically.

Practical Takeaways for Interpretability Researchers

If you're working in this area, three practices follow from the paper. Decompose before you interpret: analyze QK and OV circuits separately rather than treating a head as a monolith. Look for composition: behaviors often emerge from heads in different layers working together. Expect superposition: don't assume a dimension means a feature, and don't assume a feature lives in one dimension.

Final Thoughts

The path toward transparent AI systems runs through work like this. Not because the framework is complete, but because it makes progress cumulative. Claims about transformer internals can now be stated precisely, tested, and built upon. That's what a mathematical framework buys you.

Key Takeaway: The paper's lasting contribution is a shared language for describing transformer internals. It doesn't solve interpretability, but it makes the problem tractable—and that's the precondition for solving it.


FAQ

What is the main contribution of "A Mathematical Framework for Transformer Circuits"?

It introduces a mathematical framework for analyzing transformer language models by decomposing them into interpretable components: the residual stream, QK and OV circuits within attention heads, and circuits as computational subgraphs. It also names and analyzes superposition.

Who authored the paper?

Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Sam McCandlish, Chris Olah, and Jared Kaplan.

When was the paper published?

December 2021, on the Transformer Circuits Thread. It is also available as arXiv:2112.04510.

What is a residual stream?

The central communication channel in a transformer. Every attention head and MLP layer reads from it and writes to it via addition. It starts as the sum of token and positional embeddings and accumulates contributions from every sublayer.

What is an induction head?

An attention head that implements copying-with-a-shift: given a sequence where the current token appeared before, it attends to the token that followed the previous occurrence and copies it. Induction heads can be formed by composing a previous token head with a matching head across two layers.

What is superposition in the context of transformers?

A phenomenon where a model represents more features than it has dimensions by using non-orthogonal directions in activation space. Features overlap, which increases representational capacity but complicates interpretation.

How does the framework help with interpretability?

It provides a precise vocabulary and mathematical tools for describing internal computation. Attention heads decompose into QK and OV circuits, behaviors can be attributed to circuits, and superposition can be analyzed rather than ignored.

What are QK and OV circuits?

The QK circuit ($W_Q^T W_K$) determines where an attention head attends. The OV circuit ($W_O W_V$) determines what information the head moves. Analyzing them separately separates "where" from "what."

Is the framework applicable to large language models?

The framework is general, but the paper's detailed analysis focuses on small attention-only transformers. Scaling circuit analysis to production-scale models remains an open research problem, complicated by depth and superposition.

What is the Transformer Circuits Thread?

A research publication venue maintained by Anthropic for mechanistic interpretability work. It hosts the original paper, follow-ups like "In-context Learning and Induction Heads" and "Toy Models of Superposition," and related research.


Citations

  • Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., ... & Kaplan, J. (2021). A Mathematical Framework for Transformer Circuits. Transformer Circuits Thread.
  • Elhage, N., et al. (2021). A Mathematical Framework for Transformer Circuits. arXiv:2112.04510.
  • Anthropic. (2021). Transformer Circuits Thread. https://transformer-circuits.pub/

Dive deeper into mechanistic interpretability: read the original paper on the Transformer Circuits Thread, experiment with the provided code examples, and explore follow-up works like "Toy Models of Superposition" to continue your journey into reverse-engineering transformers.