A Mathematical Framework for Transformer Circuits (2021)

Authors: Nelson Elhage$^*{^\dagger}$, Neel Nanda$^*$, Catherine Olsson$^*$, Tom Henighan$^\dagger$, Nicholas Joseph$^\dagger$, Ben Mann$^\dagger$, 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, Jared Kaplan, Sam McCandlish, Chris Olah$^\ddagger$

Affiliation: Anthropic Published: December 22, 2021 $^*$ Core Research Contributor; $^\dagger$ Core Infrastructure Contributor; $^\ddagger$ Correspondence: colah@anthropic.com

---

Introduction

The landscape of technology is currently being reshaped by transformer-based language models. We see this in the deployment of high-capacity systems such as GPT-3, LaMDA, Codex, Meena, and Gopher. While powerful, the sheer scale and open-ended nature of these models introduce a significant risk: unexpected and potentially hazardous behaviors.

A recurring issue is that even long after training is complete, users and developers continue to stumble upon "emergent" capabilities—some of which are problematic—that were not anticipated during the initial build.

To combat this, the authors propose a method called mechanistic interpretability.

This approach is analogous to a software engineer attempting to reverse engineer a complex binary file back into human-readable source code to understand exactly how the program functions.

If we can successfully reverse engineer transformers, we can move toward a systematic method for:

• Explaining existing safety failures.
• Detecting new vulnerabilities.
• Predicting the risks associated with future, more powerful architectures.

While the Distill Circuits project previously applied this logic to vision models, no such comprehensive effort had been applied to language models until now.

Research Strategy

Given the staggering complexity of modern LLMs, the researchers adopted a "bottom-up" strategy. Rather than attacking a 96-layer model like GPT-3 immediately, they focused on the simplest possible iterations to identify fundamental algorithmic "motifs."

The Scope of this Study:

• Study transformers with $\le 2$ layers.
• Focus exclusively on attention-only blocks (removing the MLP blocks found in standard transformers).
• Develop a mathematically equivalent way to conceptualize model operations.

The authors discovered that "induction heads" are critical for in-context learning and only emerge in models with at least two layers. While this paper focuses on toy models, subsequent work will demonstrate that these findings scale to larger, more realistic architectures.

---

Summary of Results

Reverse Engineering Findings

By analyzing attention-only toy models, the team uncovered a hierarchy of complexity:

Key Discoveries:

1. Weight Accessibility: Bigram and skip-trigram tables can be extracted directly from the weights $\text{W}$ without needing to run an inference pass.
2. In-Context Learning (ICL):
   • 1-layer models perform a primitive version of ICL.
   • 2-layer models utilize Induction Heads, a sophisticated algorithm for general ICL.
3. The Transition: There is a qualitative leap in inference-time algorithms between the first and second layers.

---

Conceptual Take-Aways

The researchers noted that reverse engineering a transformer requires a different mental model than that used for InceptionV1.

The Framework

The core of their mathematical shift involves viewing the transformer not as a sequence of monolithic layers, but as a collection of independent operations.

• Residual Stream: Attention heads act as independent agents that output a result, which is then added to the residual stream.
• Computational Efficiency vs. Math: While the $\text{concatenate-and-multiply}$ method is used for speed in code, it is $\text{mathematically equivalent}$ to the additive view.
• Path-Based Analysis: Attention-only models can be decomposed into a sum of end-to-end functions.

$\text{Logits} = \sum \text{Paths}(\text{Tokens})$

If the attention patterns are frozen, these paths become linear. This allows researchers to treat the model as a black box break the model apart into interpretable sums.

Implementation Logic

To visualize how a "path" might be conceptualized in a simplified code-like structure:

def transformer_path(token_input):
    # Simplified view of a single path through the circuit
    residual_stream = embed(token_input)
    for head in attention_heads:
        # Each head adds its contribution to the stream
        residual_stream += head.compute(residual_stream)
    return project_to_logits(residual_stream)

Final Note: The authors acknowledge that these concepts may exist implicitly in other literature, but they provide them here as a formal foundation for the "circuits" approach to interpretability.