PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2402.12201v1cs.LG

TL;DR

The paper asks how to connect the interpretable features extracted by sparse dictionary learning into circuits. It develops a patch-free tracing framework and uses it to identify interpretable circuits in Othello-GPT.

  • Problem

    Sparse dictionary learning extracts more interpretable features from superposed activations, but circuit analysis seeks functional connections between internal features.

    Source: §1 · p. 1

  • Method

    The framework decomposes contributions into dictionary features from residual-stream-writing modules and traces from logits or features to lower-level features without activation patching.

    Source: §6 · p. 15

  • Results

    Dictionary learning finds feature types previously identified by probing, and the framework locates interpretable circuits in the Othello model.

    Source: §1 · p. 2, §3.2 · p. 8

  • Takeaways & Limitations

    The authors argue that dictionaries should cover every module writing to the residual stream, including the input embedding.

    Source: §6 · p. 15

    The circuit analysis is based on particular inputs and does not claim a universal understanding of Othello-GPT; discovering general Transformer motifs is left for future work.

    Source: §3.3 · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Sparse dictionary learning extracts more interpretable features from superposed activations, but connecting internal features into circuits remains a challenge. This work proposes a patch-free framework and applies it to discover interpretable circuits in Othello.

1 Introduction

  • Activation-patching and ablation approaches face out-of-distribution and hydra-effect problems and have quadratic asymptotic complexity.The paper presents its circuit-discovery framework as an alternative to these approaches.Source: §1 · p. 2
  • Dictionary features from attention and MLP outputs can be traced iteratively from logits or features to lower-level feature contributions.The method trains dictionaries on attention and MLP layer outputs in an Othello model.Source: §1 · p. 2
  • The paper argues for dictionaries on every module writing to the residual stream, including embeddings, attention outputs, and MLP outputs.It compares choices of positions on which to train dictionaries.Source: §1 · p. 2
  • In Othello, dictionary learning automatically finds, without supervision, all feature types previously discovered by probing.The paper presents this as a minor contribution.Source: §1 · p. 2
  • The framework identifies meaningful Othello-model subgraphs and reports a patch-free route to interpretable circuits.The authors describe the subgraphs as revealing a large partition of the model’s inner information flow.Source: §1 · p. 2

2 Patch-Free Circuit Discovery with Dictionary Learning

The section builds a patch-free circuit method by decomposing model activations into dictionary features and tracing their contributions through attention and MLP computations. It argues for training dictionaries on module outputs and uses distinct attribution strategies for attention and nonlinear MLPs.

2.1 Interpretable Sparse Coding with Dictionary Learning

  • Sparse dictionary learning represents an activation as a sparse weighted sum of overcomplete feature directions, with feature magnitudes indicating each direction’s contribution.Source: §2.1 · p. 3

2.2 Where to Train Dictionaries on?

  • The authors favor dictionaries on word representations and each attention/MLP output, avoiding residual-stream depth shifts or incompatible bases and MLP hidden-layer parameter and nullspace costs.Source: §2.2 · p. 3, Figure 2 · p. 4

2.3 Identifying Information Flow in OV Circuits

  • Attention outputs can be traced to lower-level features across tokens by decomposing module inputs and linearizing LayerNorm, yielding feature-level attribution through the OV circuit.Source: §2.3 · pp. 4–5

2.4 Attributing Attention Patterns to Dictionary Feature Pairs

  • The bilinear QK circuit decomposes pre-softmax attention scores into contributions from dictionary-feature pairs across two tokens’ residual streams.Source: §2.4 · pp. 5–6, Figure 4 · p. 6

2.5 Attributing MLP Features to Lower-level Features

  • Because MLP nonlinearities block direct linear attribution, Approximate Direct Contribution treats the activation-function input as constant to assign feature contributions.Source: §2.5 · p. 7

2.6 Circuit Discovery

  • Top-down attribution supports end-to-end or local circuits from a logit or dictionary feature to another feature or the input embedding, with attention-head and neuron-level contributions.Source: §2.6 · pp. 7–8

3 Dissecting Othello-GPT

This section examines dictionary features and selected circuits in a 1.2M-parameter, six-layer Transformer trained to predict legal Othello moves; its circuit analyses are input-specific.

3.1 Background and Notations

  • The study uses a 1.2M-parameter, six-layer decoder-only Transformer trained to predict legal moves in Othello, not tactics.The model has hidden dimension 128, and inputs encode board tiles numbered 1 to 64.Source: §3.1 · p. 8

3.2 Dictionary Features in Othello-GPT

  • Dictionary features include current-move and board-state features, empty-cell features, and legal-move features, which directly contribute to their corresponding logits.The paper reports that dictionary learning finds feature types previously identified by probing.Source: §3.2 · p. 8

3.3 A Case Study on Circuits in Othello-GPT

  • The circuit case studies examine selected dictionary features and explicitly make no claim to a universal understanding of Othello-GPT; discovering general Transformer motifs is left for future work.Source: §3.3 · p. 9

3.3.1 A Local OV Circuit Computing the Board State

  • An OV circuit copies board-state information from earlier moves into L2A474, with head 5 retrieving own-move information and head 2 retrieving opponent-move information.The top three contributors explain six tiles, but account for only a small proportion of L2A474's activation magnitude; many other contributors have similar semantics.Source: Figure 6 · p. 9, §3.3.1 · p. 10

3.3.2 How Do QK Circuits Implement ’Attending to Theirs’ Heads?

  • Position embeddings and low-level features encoding the current move's color contribute to L1A head 0's pattern of attending to opponent moves.The example decomposes a salient attention score into feature-pair contributions, including white cues for token 13 and black cues for token 14.Source: §3.3.2 · pp. 10–11, Figure 7 · p. 10

3.3.3 How Do Early MLPs Identify Filpped Pieces?

  • L1M combines three consecutive-move cues with an AND operation to identify a flipped piece, while attention heads 6 and 7 both contribute to one lower-level feature.The three largest Approximate Direct Contribution features correspond to consecutive moves; head 6 and head 7 both contribute to L0A629, while two other features are fired only by head 7.Source: §3.3.3 · p. 12

4 Related Work

Mechanistic interpretability seeks to reverse engineer neural networks, but superposition and sparse features complicate the recovery of interpretable components. Prior work developed dictionary learning and patch-based circuit discovery; the paper situates its approach as a patch-free alternative and reviews Othello-GPT findings that motivate its case study.

4.1 Mechanistic Interpretability

  • Mechanistic interpretability aims to reverse engineer neural networks, often by defining a behavior and using activation patching or direct logit attribution to locate relevant components.Prior work applied these tools to behaviors including factual retrieval and in-context learning.Source: §4.1 · p. 12

4.2 Attacking Superposition and Extracting More Monosemantic Features

  • Superposition packs more linear features into activation space than there are hidden dimensions, while sparse language-model features make monosemantic neurons harder to find.Sparse dictionary learning has decomposed embeddings, MLP hidden layers, and residual streams into more interpretable linear features.Source: §4.2 · p. 12

4.3 Circuit Discovery

  • Circuit-discovery workflows patch components, quantify their effects on a target logit, and prune the computational graph, but patching can face distribution shift and hydra effects.The authors identify little prior work on patch-free circuit discovery despite these problems.Source: §4.3 · pp. 12–13

4.4 Understanding Othello-GPT

  • Othello-GPT probing work supported world-model formation, while changing probe targets to my color, opponent’s color, and empty also yielded strong linear-probe performance.That probe-target change provided evidence for feature linearity, complicating the earlier claim that nonlinear probes were needed.Source: §4.4 · p. 13

5 Discussions

The discussion compares the proposed patch-free analysis with patch-based methods, examines approximation limits, and identifies future work. It reports substantial contributor overlap with direct patching, while noting that the framework is an efficient alternative rather than an overwhelmingly superior replacement.

5.1 Comparison to Patch-Based Methods

  • The authors present their framework as an efficient alternative to feature ablation, not as overwhelmingly superior to patching dictionary features for circuit discovery.This qualifies the scope of the comparison rather than claiming general superiority.Source: §5.1 · p. 13

5.1.1 Dimensions of Choice for Feature Patching

  • Feature patching choices include zero versus mean resets and causal versus direct ablation; sparse features make zero and mean patching similar in most circumstances.Less-than-rank-one ablation changes other features too, and is excluded from later discussion because its dictionary-training configuration differs.Source: §5.1.1 · p. 13

5.1.2 Computational Complexity

  • Causal patching requires inference count proportional to the number of lower-level features, whereas direct patching and this framework need only rerun the interested module.The authors characterize their framework as optimal in complexity and free from out-of-distribution effects.Source: §5.1.2 · pp. 13–14, Table 1 · p. 14

5.1.3 Correspondence between Approximate Direct Contribution and Direct Patching

  • 0.68 average IOU between the top five contributors identified by Approximate Direct Contribution and direct patching across 180 sampled MLP features indicates relatively high correspondence.The features came from the top three activated features in each of six MLP layers across 10 sampled input sequences.Source: §5.1.3 · p. 14

5.2.1 QK Decomposition

  • QK decomposition explains individual attention scores but not their relative strength, because softmax makes score changes affect the sequence-wide attention pattern.The authors say more work is needed for a comprehensive account of attention patterns.Source: §5.2.1 · p. 14

5.2.2 Approximate Direct Contribution

  • Approximate Direct Contribution captures positive effects on MLP features, but poorly captures suppression by neurons with negative effects when their activation is inhibited to zero.The discussion considers self-gated activations and fixes the activation function value when estimating contribution.Source: §5.2.2 · pp. 14–15
  • Most MLP features are interpretable under Approximate Direct Contribution with little decoder-bias contribution, but an OR-gate case attributes activation to decoder biases.The authors say both their method and patch-based methods theoretically cannot identify both features in an OR gate.Source: §5.2.2 · p. 15

5.3 Future Work and Open Problems

  • Future work includes cleaner Othello-GPT structure for Transformer Pathology and scalable circuit analysis, including feature-family taxonomies.Even the small Othello model has about 1,000 active dictionary features in each residual stream.Source: §5.3 · p. 15

6 Conclusions

The paper proposes a computationally efficient, patch-free framework for discovering Transformer circuits with dictionary learning. It argues that dictionaries should cover every module writing to the residual stream and reports interpretable circuits in Othello-GPT, while noting that only representative clean cases are presented.

6 Conclusions

  • The proposed framework discovers Transformer circuits with dictionary learning without activation patching and is described as computationally efficient.The proposed dictionary-training positions include every module writing to the residual stream, including input embeddings.Source: §6 · p. 15
  • The authors test the framework on Othello-GPT and report discovering a number of interpretable circuits.They present representative cases with clean structures rather than all findings.Source: §6 · p. 15
  • The authors identify better dictionary-training techniques and presentation interfaces as possible ways to share more universal Othello-GPT findings.They also express interest in applying the framework to real-world language models.Source: §6 · pp. 15–16

A.1 Dealing with Non-linearity of LayerNorm

LayerNorm’s standard-deviation calculation is nonlinear, complicating attribution of consequences to individual linear components. The appendix treats that standard deviation as constant to linearize LayerNorm and estimate each component’s impact separately.

A.1 Dealing with Non-linearity of LayerNorm

  • LayerNorm’s standard-deviation calculation is nonlinear, preventing attribution of a consequence to each linear component of its input.This is the specific nonlinearity the appendix addresses.Source: §A.1 · p. 20
  • Treating the input standard deviation as constant transforms LayerNorm into a linear function without changing the computational result.The resulting modified LayerNorm can be applied separately to linear components of the input activation.Source: §A.1 · p. 20

A.2.1 Our Re-trained Version of Othello-GPT

The authors retrained a smaller Othello-GPT, removing dropout and reducing its depth while retaining 99.5% performance. The model has about 1.2 million parameters and was trained on about 1e9 sequences.

A.2.1 Our Re-trained Version of Othello-GPT

  • The retrained Othello model has about 1.2 million parameters, a hidden dimension of 128, and was trained on about 1e9 sequences.Source: §A.2.1 · p. 20
  • Dropout was removed because the authors say it aggravates redundancy and makes circuit analysis more difficult.They note that this setting is also widely applied in mainstream language-model training pipelines.Source: §A.2.1 · p. 20
  • Reducing depth from 8 to 6 layers lowered model size without sacrificing interpretability conclusions; performance changed from 99.9% to 99.5%.The authors say the change reduced repetitive experimental sections.Source: §A.2.1 · p. 20
  • Figure 10 describes a dictionary-training limitation: near-perfect reconstruction error alongside too many dead neurons.Source: Figure 10 · p. 21

A.2.2 Dictionary Training

The authors train dictionaries with an untied encoder-decoder autoencoder and an L1 sparsity constraint. They report that using too-small sparsity coefficients yielded near-perfect reconstruction but about 75% dead neurons per dictionary.

A.2.2 Dictionary Training

  • The dictionary pipeline uses an untied encoder-decoder autoencoder with 128 input and output neurons, 1024 hidden neurons, and an L1 sparsity constraint.Source: §A.2.2 · p. 21
  • Because transformer activations grow larger in later layers, the authors empirically use smaller L1 sparsity coefficients α there.They report the same activation trend as Dettmers et al. (2022).Source: §A.2.2 · p. 21
  • Using too-small α values produced near-perfect reconstruction performance but left about 75% dead neurons in each dictionary.The authors say this may reflect an overly small choice of α.Source: §A.2.2 · p. 22
  • The authors expect deeper understanding of Othello-GPT if the dictionary-training problem is alleviated.Source: §A.2.2 · p. 22

A.3 Interface for Interpretability

The interpretability interface summarizes the top-activating board states for each dictionary feature, and its example identifies a feature associated with the current move at f-1. The Othello dictionaries are relatively dense, and the authors note both the task’s value as a conceptual language-model substitute and a likely non-optimal training implementation.

A.3 Interface for Interpretability

  • The interface replaces difficult-to-read images, which the authors say invite interpretability illusions, with plots of statistics over each feature’s top-k activating inputs.Figure 12 shows seven plots for the 2048 most activating inputs out of 1.2 million board states.Source: §A.3 · p. 22, Figure 12 · p. 23
  • The interface visualizes current-player piece-color proportions among k inputs, tile play counts across k moves, legal-move counts across k board states, summed board-state values, and per-cell flips across k moves.For each tile, board-state values are 1 for the current player’s color, -1 for a different color, and 0 if empty, summed over the k boards.Source: §A.3 · pp. 22–23
  • L0A622 is interpreted as a “current move = f-1” feature because its top 2048 activating inputs play moves at position f-1.The example examines the 2048 most activating inputs among 1.2 million board states.Source: §A.3 · pp. 23–24

A.4 Othello-GPT and Language Models

  • Othello is used as a conceptual verification task for language-model interpretability because it has abstract concepts and moderate complexity without an excessively large vocabulary.The task has 60 tokens, and the authors say its lack of duplicate inputs to some extent enhances explainability.Source: §A.4 · p. 24
  • Othello features are denser than language-model features in the authors’ account, and become more frequently active as board states grow more complex during play.They say this is consistent with existing observations of language models.Source: §A.4 · p. 24
  • Late MLP dictionary decompositions sometimes approach an L0 norm of 128, versus typically 10–20 sparse activations for a 512-d layer in Bricken et al. (2023).Figure 13 plots average L0 norm across 10,000 tokens for each dictionary.Source: §A.4 · p. 24, Figure 13 · p. 24
  • The authors attribute Othello’s dense features largely to its simpler, denser-feature world and say their dictionary-training implementation is likely not optimal.Source: §A.4 · p. 24
Loading 2402.12201v1…