Table of Contents
Fetching ...

Dictionary Learning Improves Patch-Free Circuit Discovery in Mechanistic Interpretability: A Case Study on Othello-GPT

Zhengfu He, Xuyang Ge, Qiong Tang, Tianxiang Sun, Qinyuan Cheng, Xipeng Qiu

TL;DR

The paper introduces a patch-free circuit-discovery framework using sparse dictionary learning to decompose transformer activations into monosemantic features and trace information flow via OV, QK, and ADC mechanisms. Applied to Othello-GPT, it uncovers interpretable local circuits for board-state computation, attention patterns, and flipped-piece detection, avoiding activation patching and reducing computational overhead. The approach yields a set of actionable insights into how residual streams carry and transform information, and it provides a structured comparison to patch-based methods, including limitations related to QK nonlinearity and ADC approximations. The work advances mechanistic interpretability by offering a scalable, patch-free alternative with potential for extension to larger language models.

Abstract

Sparse dictionary learning has been a rapidly growing technique in mechanistic interpretability to attack superposition and extract more human-understandable features from model activations. We ask a further question based on the extracted more monosemantic features: How do we recognize circuits connecting the enormous amount of dictionary features? We propose a circuit discovery framework alternative to activation patching. Our framework suffers less from out-of-distribution and proves to be more efficient in terms of asymptotic complexity. The basic unit in our framework is dictionary features decomposed from all modules writing to the residual stream, including embedding, attention output and MLP output. Starting from any logit, dictionary feature or attention score, we manage to trace down to lower-level dictionary features of all tokens and compute their contribution to these more interpretable and local model behaviors. We dig in a small transformer trained on a synthetic task named Othello and find a number of human-understandable fine-grained circuits inside of it.

Dictionary Learning Improves Patch-Free Circuit Discovery in Mechanistic Interpretability: A Case Study on Othello-GPT

TL;DR

The paper introduces a patch-free circuit-discovery framework using sparse dictionary learning to decompose transformer activations into monosemantic features and trace information flow via OV, QK, and ADC mechanisms. Applied to Othello-GPT, it uncovers interpretable local circuits for board-state computation, attention patterns, and flipped-piece detection, avoiding activation patching and reducing computational overhead. The approach yields a set of actionable insights into how residual streams carry and transform information, and it provides a structured comparison to patch-based methods, including limitations related to QK nonlinearity and ADC approximations. The work advances mechanistic interpretability by offering a scalable, patch-free alternative with potential for extension to larger language models.

Abstract

Sparse dictionary learning has been a rapidly growing technique in mechanistic interpretability to attack superposition and extract more human-understandable features from model activations. We ask a further question based on the extracted more monosemantic features: How do we recognize circuits connecting the enormous amount of dictionary features? We propose a circuit discovery framework alternative to activation patching. Our framework suffers less from out-of-distribution and proves to be more efficient in terms of asymptotic complexity. The basic unit in our framework is dictionary features decomposed from all modules writing to the residual stream, including embedding, attention output and MLP output. Starting from any logit, dictionary feature or attention score, we manage to trace down to lower-level dictionary features of all tokens and compute their contribution to these more interpretable and local model behaviors. We dig in a small transformer trained on a synthetic task named Othello and find a number of human-understandable fine-grained circuits inside of it.
Paper Structure (37 sections, 9 equations, 13 figures, 1 table)

This paper contains 37 sections, 9 equations, 13 figures, 1 table.

Figures (13)

  • Figure 1: Activation of a given position and a given input decomposed into a sparse reconstruction of dictionary features.
  • Figure 2: Five choices of positions to decompose via dictionary learning. Prior work has studied decomposing $X_{P1}$, $X_{p2}$ and $X_{p5}$, namely word embedding, MLP hidden layer and the residual stream. We claim it preferable to decompose $X_{P1}$, $X_{P3}$ and $X_{P4}$ i.e. word representations, the output of each attention layer, and the output of each MLP layer.
  • Figure 3: Attention features decomposed into contributions of lower-level dictionary features of previous tokens via the OV circuit.
  • Figure 4: Attention scores (before softmax) decomposed into QK compositions of dictionary feature pairs between the residual stream of two tokens.
  • Figure 5: With Approximate Direct Contribution, the activation of an MLP feature is decomposed into Approximate Direct Contributions of lower-level dictionary features.
  • ...and 8 more figures

Theorems & Definitions (1)

  • Definition 2.1: Approximate Direct Contribution