GPT-2 Mechanistic Interpretability on SCAN

Training a GPT-2 style model on SCAN to mechanistically investigate how transformers fail at compositional generalization using TransformerLens.

Research Question

Which internal mechanisms cause transformer models to fail at compositional generalization?

Specifically: why do transformer models systematically fail to generate the correct number of repeated actions — a failure pattern observed consistently across all SCAN splits in T5?

Motivation

T5 exhibits a systematic failure across all SCAN splits: it consistently generates the wrong number of repeated actions, regardless of split difficulty (see T5 Compositional Analysis Project). This suggests the failure is not about generalization to novel compositions, but about a deeper inability to track and reproduce action repetitions correctly.

Rather than interpreting T5 directly (intractable at scale), we use it as a diagnostic tool — its failure pattern motivates a precise mechanistic question. We train a minimal GPT-2 style model and use TransformerLens to investigate: which internal components are responsible for tracking action repetitions, and why do they fail?

This follows the model organism methodology: use a tractable, transparent model to study a phenomenon mechanistically, then reason about what this implies for larger systems.

Approach

Model

Small GPT-2 style transformer trained from scratch:

Results

Evaluation (Seq2Seq Generation)

Mechanistic Interpretability Findings

Notebook 01: Logit Lens Analysis

Question: Where in the model does the failure happen?

Key Finding: Two distinct failure modes identified:

• Layer 1 Dissolution (~60%): Failure originates in early layers —
  foundational features fail to preserve structural context

• Layer 2 Overwrite (~40%): Layer 1 computes correctly but Layer 2
  overwrites with wrong prediction

• Length split shows delayed failure onset (step 18.50) — model
  holds on longer before collapsing on longer sequences

Notebook 02: Attention Pattern Analysis

Question: Which heads track modifier tokens ('twice', 'thrice')?

Modifier Attention Scores (Simple Split):

Cross-Split L1H1 Modifier Attention:

Key Findings:

• H1 Confirmed: L1H1 is the primary modifier-tracking head —
  consistently highest modifier attention across all splits

• H2 Rejected: Modifier attention is preserved in failure cases —
  the model attends to 'twice'/'thrice' equally in success and failure

• H3 Partially confirmed: L1H1 shows diagonal attention pattern
  suggesting previous-token tracking for repetition counting

• New finding: L0H2 shows largest attention drop in AddPrim failures
  (+0.0226), suggesting primitive-specific structural encoding

• New finding: Length split shows L1H1 attention HIGHER in failures —
  model compensates on longer sequences but still fails

Core insight: The model correctly SEES modifiers but fails to USE
them — failure is downstream of attention.

Notebook 03: Activation Patching & MLP Ablation

Question: Which downstream MLP circuits fail to translate modifier
attention into correct repetition counts?

Section 3: L1H1 Ablation

Section 4: MLP Activation Patching

Section 5: Recovery Rates

Key Findings:

• H4 Partially confirmed: Ablating L1H1 changes behavior
  across all splits — negative delta (slight improvement) suggests
  L1H1 contributes to structural tracking but is not the sole
  repetition-counting circuit. L0H2 shows larger ablation effect.

• H5 Rejected: Patching average success MLP activations
  DECREASES token accuracy (~20% drop for L1, ~40% for L0) —
  MLP activations are highly context-specific and non-transferable

• H6 Partially supported: L1 MLP consistently more
  transferable than L0 MLP across all splits, but neither recovers
  failures meaningfully

• Core finding: Failures are localized (70-75% tokens correct)
  but MLP activations are too context-specific to patch — the
  failure mechanism is distributed rather than localized to one circuit

Summary of Hypotheses

Core Research Finding

The model correctly attends to modifier tokens (L1H1, ~70-75%
token accuracy on failure cases) but fails to translate this
attention into correct repetition counts at specific generation
steps. MLP activations are highly context-specific — averaged
success activations corrupt rather than fix failure computations.
The failure mechanism appears distributed across circuits rather
than localized to a single component.

Setup

git clone https://github.com/suehuynh/gpt2-mech-interp
cd gpt2-mech-interp
pip install -r

git clone https://github.com/suehuynh/gpt2-mech-interp
cd gpt2-mech-interp
pip install -r

git clone https://github.com/suehuynh/gpt2-mech-interp
cd gpt2-mech-interp
pip install -r

Key dependencies: torch, transformer_lens, datasets,
wandb, matplotlib, circuitsvis, pandas, seaborn

Repository Structure

Related Work

• Project 2: T5 failure analysis on SCAN — [https://github.com/suehuynh/scan-compositional-generalization]

• Keysers et al. (2019) — Measuring Compositional Generalization

• Nanda et al. (2023) — Progress measures for grokking via mechanistic interpretability

• Conmy et al. (2023) — Towards Automated Circuit Discovery for Mechanistic Interpretability