Circuits are all you need (or how to look inside transformers)

A walkthrough of "A Mathematical Framework for Transformer Circuits"

Having an LLM is good. Really. But researchers also want to know how it works inside: what makes it predict “subscribe” after “hit the like button and ___” That’s when mechanistic interpretability comes into play.

The paper we’re discussing here is one of the must-reads in this domain: “A Mathematical Framework for Transformers Circuits” by Anthropic. The authors try to reverse-engineer transformers mathematically (that’s why it looks so scary when skimming and reading it)

CORE IDEA: convert the transformer’s operation into mathematically equivalent operation and trace how information goes through it. That’s how we can start read the circuits hidden in the weights - not just look at final result

Important

The paper focuses on attention only transforerms. This means no MLP layers. Authors did this because MLP layers are hard to interpret:

• They’re non-linear

• Individual neurons are polysemantic - a single neuron can fire for many unrelated things - because of superposition (model packs more features than it has dimensions, so one neuron ends up storing information about salsa ingredients and Newton’s second law at the same time)

Transformer’s Overview (but make it different)

The core trick of the paper is to look at transformers differently

The Residual Stream

Residual stream sums output of all previous layers and original embeddings. Each layer reads from it and writes information back to it

Image is taken from A Mathematical Framework for Transformer Circuits

The residual stream doesn’t have a privileged basis. That means you can’t just take a single dimension and say “this dimension makes the model friendly”. - features are spread across many dimensions as direction (something like “0.8 dimension 1 + 0.15 dimension 2”) This is why rotating the coordinate system doesn’t change anything: the model will still work the same. What matters is geometry - direction, angles between directions, subspaces

Virtual Weights

The residual stream is lineaer. That means you can multiply the read and write matrices of any two layers (no matter if they are layer 115 and layer 999) together to get a single matrix describing how one directly influences the other - skipping all intermediate layers.
This matrix is called “virtual weights”

\(\text{Virtual Connection}=W_{I(layer999)} \times W_{O(layer115)}\)

The image was taken from A Mathematical Framework for Transformer Circuits

Attention Heads

The transformers discussed here are attention-only, so of course we’ll discuss attention heads.

Each attention head works in two places:

• QK (query-key) circuit

• OV (output-value) circuit

So what is circuit? To put simply,

circuit = traceable paths of how information flows through the weights.

Generally, it’s just another matrix multiplication (because the whole point of AI is beefing with matrix multiplication)

WARNING! A bit of math will be there!

---

• QK-circuit - uses the query and key matrices to decide which tokens attend to which. It’s the one that build the attention heatmap. They’re always operate together.

• OV-circuit - uses the output and value matrices to decide what information is moved once the connection is made. These ones too always operate together.

Q, K, V, and O are all learned matrices.
Q, K, V are famous to the community, but where did O come from? It’s simple: it is always there, but it just always there, just rarely shown on diagrams

During forward pass:

• There are activations, which are the product of input tokens against weight matrices, They’re pass-specific (discarded after each forward pass)
  You’re getting:

  \(Q=x \cdot W_Q \text{, } K=x \cdot W_K \text{, } V=x \cdot W_V\)

• Then you compute the attention pattern. In attention-only transformers, softmax is the only non-linear component in the attention-only transformer, which makes analysis traceable

  \(A=\text{softmax}(Q \cdot K^T)\)

• Then you compute output, which will be written into residual stream

  \(O=A \cdot V \cdot W_O\)

---

Things that models do

Researches ran experiment on zero-, one- and two-layer transformer and analyzed their behavior. Here’s what they do:

Zero-Layer Transformers: Pure Statistics

Zero-layer transformers have no attention. They’re simply token→embed→unembed
All its behavior is matrix W_U x W_E (unembedding and embedding matrices), which is trained to approximate bigram likelihoods: how often token B follows token A.
You can read it from the weights

One-Layer Transformers: Skip-Trigrams

You need one layer to predict .__init__() after self

When having one attention layer, attention heads can look back at previous token and influence prediction of what comes next. They create skip-trigrams: pattern in the form of [source] … [destination] → [output]

This creates interesting things:

• a big part of their capacity is spent on copying: OV-circuit looks at token and if attention head attended it, they increase this token’s probability. QK-circuit ensures that copying happens only where bigram statistics make the token likely continuation

  • from this happens interesting quirk, related to tokenization: for tokens “Ralph” and “ Ralph” (with space) predictions would be different.
    Imagine you’re predicting what goes after “R”.
    If you had “ Ralph” (there’s a space) token, copying will boost probability of “alph”
    If you had “Ralph”, your prediction will depend on what was in text before

• fundamental limitation of one-layer models: because they represent skip-trigrams, they can’t capture three-way interactions. This happens because QK and OV circuits operate independently - they can’t coordinate on which source-destination-output combinations are valid.
  Take the phrase “keep in mind”:

  • ”keep” - source token (the one being attended to)

  • ”in” - destination token (current position, triggers the lookup)

  • “mind” - output token (what gets predicted)

  The QK circuit handles the source→destination relationship: “in” attends back to ”keep”

  The OV circuit handles the source→output relationship: attending to keep “keep” boosts “mind”
  

  Suppose the model also learned “keep at bay”. Here’s where trouble arises:

  • The QK circuit learns both “in” and “at” attend back to “keep”.

  • The OV circuit learns that attending to “keep” boosts both “mind” and “bay”.

    But they can’t communicate with each other. That’s how you get four combinations, instead of two: “keep in mind” (correct), “keep at bay” (correct), “keep in bay” (bug), “keep at mind” (bug)

Two-Layer Transformers: Compositions

Now you’re not just getting more space for skip-trigrams, you’re composing (this is taking output of one function into input of another, f(g(x)))

Here, you have three types of composition:

• Q-composition: query matrix reads from early head modifications and changes what the later head is looking for

• K-composition: key matrix reads from early head modifications and changes what later head finds relevant

• V-composition: value matrix reads from a modified subspace and creates “virtual attention heads”. It allows a layer 2 head’s OV circuit to read from a subspace written by layer 1 head, creating a combined head whose behavior isn’t visible when inspecting either head alone

Two layer model create a special mechanism, which is called “induction heads”.

Induction heads

• Induction heads increase the likelihood of [B] given [A][B]…[A] (because there was [A][B] pattern before)

• They search over the context for previous examples of how this token was used. If they don’t find it, they do nothing. But if they find it, they look at what next token was and copy it.
  This allows:

  • more confident predictions

  • robustness to distributional shift, as long as there are repeating patterns in the sequence patterns to capture

The Mechanism:

Central trick: at every position, the key encodes information about token that came just before

Image taken from In-context Learning and Induction Heads (actually, this is a spoiler for Transformer Circuits article)

Take a sequence where "node" → "struction" appears twice:

1. Layer1 “previous-token head” does its work during the forward pass: when model processes the first “struction”, previous-token head copies information from “node” into residual stream at “struction” position

2. Layer2 induction head builds keys of “what came before me” (K-composition) “struction” has key that says “node” is before it

3. When going through the next occurrence of “node”, query-key matching will find “node”→ “struction” pair. Attention fires and OV circuit copies “struction” as prediction of what is next.

Key Takeaways

1. Attention-only transformers are largely linear - excluding softmax - which makes them traceable end-to-end

2. Model behavior can be decomposed into circuits: interpretable paths showing how information flows through the weights

3. The residual stream acts as a shared memory that each layer reads from and writes to independently, which lets us see contributions to specific components

4. Composition of attention heads across layers gives models their expressivity and enables things like induction heads and in-context learning