Mechanistic Interpretability & Mathematics: A Whirlwind Tour

Slides: https://neelnanda.io/maths-mech-interp-slides

​

Sequence: https://neelnanda.io/concrete-open-problems

Getting Started: https://neelnanda.io/getting-started

Tooling: https://github.com/neelnanda-io/TransformerLens

What is Mechanistic Interpretability?

• Hypothesis: Models learn human-comprehensible algorithms and can be understood, if we learn how to make it legible
• Achieve rigorous understanding without tricking ourselves
• Goal: Reverse engineer train neural networks, like program binary -> source code

Key Claim: Mechanistic Understanding of Neural Networks is Possible

Setup

The Fourier Multiplication Algorithm for Modular Addition

Mystery: Why do models grok?

(Power et al, 2022)

Background

Inspiration: Mechanistic Interpretability

• Hypothesis: Models learn human-comprehensible algorithms and can be understood, if we learn how to make it legible
• Achieve rigorous understanding without tricking ourselves
• Goal: Reverse engineer train neural networks, like program binary -> source code

A Growing Area of Research

Does Localization Inform Editing? (Hase et al, 2023)

A Mathematical Framework for Transformer Circuits (Elhage et al, Anthropic 2021)

Transformer Feed-Forward Layers Are Key-Value Memories (Geva et al, EMNLP 2021)

Investigating Gender Bias in Language Models Using Causal Mediation Analysis (Vig et al, NeurIPS 2020)

Toy Models of Superposition (Elhage, Anthropic 2022)

Locating and Editing Factual Associations in GPT (Meng et al, NeurIPS 2022)

A Growing Area of Research

Multimodal Neurons in Artificial Neural Networks (Goh et al, Distill 2021)

Compositional Explanations of Neurons (Mu and Andreas, NeurIPS 2020)

Causal Abstractions of Neural Networks (Geiger et al, NeurIPS 2021)

SGD Learns Parities Near the Computational Limit (Barak et al, NeurIPS 2022)

Curve Circuits (Cammarata et al, Distill 2020)

The Quantization Model of Neural Scaling (Michaud et al, 2023)

Proposal: Apply mechanistic interpretability techniques to reverse-engineer modular addition, to understand grokking

Getting Traction

Understanding the Embedding

t-SNE of embedding, Power et al

Principal components of embedding, Liu et al

Insight: Apply the Fourier Transform

At initialisation

Grokked model

Logits are Sums of Cosines

​

Explains 95% of the variance!

What are the logits?

Inputs a,b -> output c

Logits are Sums of Cosines

​

Explains 95% of the variance!

Why cos?

What are the logits?

Inputs a,b -> output c

The Learned Algorithm

The Fourier Multiplication Algorithm for Modular Addition

Lines of Evidence

1. Suggestive Evidence: Surprising Periodicity
2. Mechanistic Evidence: Reading off the algorithm from model weights

Lines of Evidence

• Suggestive Evidence: Surprising Periodicity
• Mechanistic Evidence: Composing Model Weights
• Zooming In: Approximating Neurons with Sines and Cosines

Lines of Evidence

• Suggestive Evidence: Surprising Periodicity
• Mechanistic Evidence: Composing Model Weights
• Zooming In: Approximating Neurons with Sines and Cosines
• Correctness checks: Ablations

Understanding Grokking

Mystery: Why do models grok?

Question: Is mechanistic understanding useful?

1. It’s a random walk (Millidge, 2022)
2. It takes time to learn representations (Liu et al, 2022)
3. Double descent (Davies et al, 2022)
4. The slingshot mechanism (Thilak et al, 2022)
5. The weight norm is too high (Liu et al, 2023)

The Three Phases of Grokking

Two learned algorithms: Memorising circuit and Fourier circuit

1. Memorisation
2. Circuit Formation
3. Cleanup

Understanding Grokking via Progress Measures

• Progress Measure: A smooth metric that identifies hidden progress
• Key: Using our mechanistic understanding to derive the measures
• Circuit formation with Excluded loss: Removes the Fourier circuit.

Understanding Grokking via Progress Measures

• Progress Measure: A smooth metric that identifies hidden progress
• Key: Using our mechanistic understanding to derive the measures
• Circuit formation with Excluded loss: Removes the Fourier circuit.
• Cleanup with Restricted loss: Removes the memorisation circuit.

Generalising the Algorithm with Representation Theory�Follow Up: A Toy Model of Universality (Chughtai et al)

The Group Composition via Representations Algorithm

A Toy Model of Universality (Chughtai et al)

Analysing Learned Representations in S5 Composition

A Toy Model of Universality (Chughtai et al)

Testing Universality: Which Representations Are Learned?

A Toy Model of Universality (Chughtai et al)

Key Takeaway 1: Deep Learning Can Learn Rich Mathematical Structure

Key Takeaway 2: Mechanistic Understanding Is Possible

Progress Measures for Grokking via Mechanistic Interpretability (Nanda et al)

Open Problems: Interpreting Algorithmic Models

A Growing Area of Research

Does Localization Inform Editing? (Hase et al, 2023)

A Mathematical Framework for Transformer Circuits (Elhage et al, Anthropic 2021)

Transformer Feed-Forward Layers Are Key-Value Memories (Geva et al, EMNLP 2021)

Investigating Gender Bias in Language Models Using Causal Mediation Analysis (Vig et al, NeurIPS 2020)

Toy Models of Superposition (Elhage, Anthropic 2022)

Locating and Editing Factual Associations in GPT (Meng et al, NeurIPS 2022)

A Growing Area of Research

Multimodal Neurons in Artificial Neural Networks (Goh et al, Distill 2021)

Compositional Explanations of Neurons (Mu and Andreas, NeurIPS 2020)

Causal Abstractions of Neural Networks (Geiger et al, NeurIPS 2021)

SGD Learns Parities Near the Computational Limit (Barak et al, NeurIPS 2022)

Curve Circuits (Cammarata et al, Distill 2020)

The Quantization Model of Neural Scaling (Michaud et al, 2023)

What is Mechanistic Interpretability?

• Goal: Reverse engineer neural networks
• Like reverse-engineering a compiled program binary to source code
• Hypothesis: Models learn human-comprehensible algorithms and can be understood, if we learn how to make it legible
• Understanding features - the variables inside the model
• Understanding circuits - the algorithms learned to compute features
• Forming an epistemic foundation for a rigorous science of interpretability

Interpretability Request for Proposals (Chris Olah)

Motivating Mechanistic Interpretability

Goal: Understand Model Cognition

​

Is it aligned, or telling us what we want to hear?

Motivation

• Key Q: What should interpretability look like in a post GPT-4 world?
• Large, generative language models are a big deal
• Models will keep scaling. What work done now will matter in the future?
• Emergent capabilities keep arising
• Many mundane problems go away
• A single massive foundation model

Interpretability in a Post GPT-4 World

Inputs and Outputs Are Not Enough

​

How do models represent their thoughts?

Understanding Superposition

How do models represent thoughts?

Understanding Superposition

Decomposing Language Models With Dictionary Learning (Bricken et al)

​

• Goal: Decompose models to independently meaningful + composable units/features
• Curse of dimensionality �=> this is crucial
• Hope: Neurons = features
• Problem: Polysemanticity

How do models represent thoughts?

Understanding Superposition

• Goal: Decompose models to composable + meaningful units
• Curse of dimensionality -> this is crucial
• Hope: Neurons = features
• Problem: Polysemanticity
• Hypothesis: Superposition

Toy Models of Superposition (Nelson Elhage et al)

Approach: Forming conceptual frameworks

Q: Can we do it for neuron superposition?

Eg n sparse binary inputs, nC2 outputs for x_i AND x_j

Toy Models of Superposition (Nelson Elhage et al)

Conceptual Frameworks:

Geometry of Superposition

Toy Models of Superposition (Elhage et al)

Open Problems: Exploring Polysemanticity & Superposition + Analysing Toy Language Models

Approach: Extracting features from superposition

Finding an interpretable, overcomplete basis

Q: What to learn from neuroscience, compressed sensing, etc?

[Interim research report] Taking features out of superposition with sparse autoencoders (Lee Sharkey et al)

Data

Reconstruction

Dictionary

Example: Compound Word Detection

How big a deal is this?

Finding Neurons In A Haystack (Wes Gurnee et al)

Extracting Features in Superposition

Decomposing Language Models With Dictionary Learning (Bricken et al)

Idea: Train a Sparse Autoencoder

Decomposing Language Models With Dictionary Learning (Bricken et al)

Key Idea: True Features Are Sparse

Decomposing Language Models With Dictionary Learning (Bricken et al)�Sparse Autoencoders Find Highly Interpretable Features in Language Models (Cunningham et al)

L1 Regularisation

Finding Monosemantic Features: The Arabic Feature

Decomposing Language Models With Dictionary Learning (Bricken et al)

Feature Splitting: Surprising Geometric Structure

Decomposing Language Models With Dictionary Learning (Bricken et al)

Example: A Split Feature

Fact Localisation: Study Computation in Superposition�Nanda & Rajamanoharan, forthcoming

How Do Language Models Store & Recall Facts?

• Eg “Fact: Michael Jordan plays the sport of” -> “basketball”

Dissecting Recall of Factual Associations (Geva et al)

How Do Language Models Store & Recall Facts?

• Eg “Fact: Michael Jordan plays the sport of” -> “basketball”
• Sport is linearly separable on the Jordan token!
• Mystery: How do early MLP layers combine the name tokens to be linearly separable?

Fact Localisation (Nanda & Rajamanoharan, forthcoming)

Distilling to a Toy Model

• Memorise a lookup table from {0, …, n-1} x {0, …, n-1} -> {0, 1}
• Mystery: Can we reverse-engineer it?

Integer 0 “embedding”

(Frozen, Random)

Integer 1 “embedding”

(Frozen, Random)

MLP 0

Binary Linear Classifier

+

+

MLP 1

Fact Localisation (Nanda & Rajamanoharan, forthcoming)

MLP(x) = W_out GELU(W_in x + b_in)

Case Study: Emergent World Representations in Othello

Case Study: Emergent World Representations in Othello-GPT

Networks have real underlying principles with predictive power

Seemingly Non-Linear Representations?!

Emergent World Representations (Li et al)

Linear Representation Hypothesis:

Models represent features as directions in space

Models have underlying principles with predictive power

• My colour vs their’s
• Linear representation hypothesis
• Generalises
• Survived falsification
• Has predictive power

Actually, Othello-GPT Has A Linear Emergent Representation (Neel Nanda)

Open Problems: Future Work on Othello-GPT

Case Study: Emergent World Representations in Othello

Key takeaway: Models have universal principles with predictive power

Learning More

• 200 Concrete Open Problems in Mechanistic Interpretability
• https://neelnanda.io/concrete-open-problems
• Getting Started in Mechanistic Interpretability
• https://neelnanda.io/getting-started
• A Comprehensive Mechanistic Interpretability Explainer
• https://neelnanda.io/glossary
• TransformerLens
• https://github.com/neelnanda-io/TransformerLens
• Slides: https://neelnanda.io/whirlwind-slides

Model Debugging

Fixing behaviour without breaking everything else

Challenge: Memory editing = Deletion + Insertion

ROME Tries Memory Editing

Locating and Editing Factual Associations in GPT (Kevin Meng et al)

Challenge: Memory editing = Deletion + Insertion

Memory insertion

Does Localization Inform Editing? (Peter Hase et al) & Detecting Edit Failures In Large Language Models (Jason Hoelscher-Obermaier et al)

Challenge: Memory editing = Deletion + Insertion

True memory editing?

Challenge:

• Zoom in: How are facts represented in neurons?
• Probing + interventions + how it’s used
• Superposition?
• A minimal deletion
• A minimal insertion

Challenge: What Does Finetuning Do To Circuits?

Are circuits rearranged, or formed anew?

Challenge: Fixing Bad In-Context Learning

Finding the minimal edit

​

Examples: Sycophancy, less educated, buggy code

Challenge:

• Identify key heads
• Find the minimal edit
• Superposition? Context-specific?

Overthinking the Truth (Danny Halawi et al)