Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Efficient Automated Circuit Discovery in Transformers using Contextual Decomposition

Aliyah R. Hsu, Georgia Zhou, Yeshwanth Cherapanamjeri, Yaxuan Huang, Anobel Y. Odisho, Peter R. Carroll, Bin Yu

TL;DR

This work tackles the problem of interpreting large transformers by proposing CD-T, a contextual decomposition framework that yields interpretable circuits at arbitrary granularity, including specific sequence positions. CD-T uses exact decomposition rules for linear modules and a specialized treatment for self attention to quantify how internal components contribute to model outputs, enabling automated circuit discovery that prunes away irrelevant parts. On IOI, Greater-than, and Docstring tasks CD-T outperforms prior baselines in recovering manual circuits with high ROC AUC and demonstrates faithfulness to the full model that surpasses random circuits, often achieving faithfulness equal to 1 with relatively small circuits. The method significantly accelerates circuit discovery (seconds rather than hours) and is compatible with common transformer architectures, promising scalable mechanistic interpretability in large language models.

Abstract

Automated mechanistic interpretation research has attracted great interest due to its potential to scale explanations of neural network internals to large models. Existing automated circuit discovery work relies on activation patching or its approximations to identify subgraphs in models for specific tasks (circuits). They often suffer from slow runtime, approximation errors, and specific requirements of metrics, such as non-zero gradients. In this work, we introduce contextual decomposition for transformers (CD-T) to build interpretable circuits in large language models. CD-T can produce circuits of arbitrary level of abstraction, and is the first able to produce circuits as fine-grained as attention heads at specific sequence positions efficiently. CD-T consists of a set of mathematical equations to isolate contribution of model features. Through recursively computing contribution of all nodes in a computational graph of a model using CD-T followed by pruning, we are able to reduce circuit discovery runtime from hours to seconds compared to state-of-the-art baselines. On three standard circuit evaluation datasets (indirect object identification, greater-than comparisons, and docstring completion), we demonstrate that CD-T outperforms ACDC and EAP by better recovering the manual circuits with an average of 97% ROC AUC under low runtimes. In addition, we provide evidence that faithfulness of CD-T circuits is not due to random chance by showing our circuits are 80% more faithful than random circuits of up to 60% of the original model size. Finally, we show CD-T circuits are able to perfectly replicate original models' behavior (faithfulness $ = 1$) using fewer nodes than the baselines for all tasks. Our results underscore the great promise of CD-T for efficient automated mechanistic interpretability, paving the way for new insights into the workings of large language models.

Efficient Automated Circuit Discovery in Transformers using Contextual Decomposition

TL;DR

This work tackles the problem of interpreting large transformers by proposing CD-T, a contextual decomposition framework that yields interpretable circuits at arbitrary granularity, including specific sequence positions. CD-T uses exact decomposition rules for linear modules and a specialized treatment for self attention to quantify how internal components contribute to model outputs, enabling automated circuit discovery that prunes away irrelevant parts. On IOI, Greater-than, and Docstring tasks CD-T outperforms prior baselines in recovering manual circuits with high ROC AUC and demonstrates faithfulness to the full model that surpasses random circuits, often achieving faithfulness equal to 1 with relatively small circuits. The method significantly accelerates circuit discovery (seconds rather than hours) and is compatible with common transformer architectures, promising scalable mechanistic interpretability in large language models.

Abstract

Automated mechanistic interpretation research has attracted great interest due to its potential to scale explanations of neural network internals to large models. Existing automated circuit discovery work relies on activation patching or its approximations to identify subgraphs in models for specific tasks (circuits). They often suffer from slow runtime, approximation errors, and specific requirements of metrics, such as non-zero gradients. In this work, we introduce contextual decomposition for transformers (CD-T) to build interpretable circuits in large language models. CD-T can produce circuits of arbitrary level of abstraction, and is the first able to produce circuits as fine-grained as attention heads at specific sequence positions efficiently. CD-T consists of a set of mathematical equations to isolate contribution of model features. Through recursively computing contribution of all nodes in a computational graph of a model using CD-T followed by pruning, we are able to reduce circuit discovery runtime from hours to seconds compared to state-of-the-art baselines. On three standard circuit evaluation datasets (indirect object identification, greater-than comparisons, and docstring completion), we demonstrate that CD-T outperforms ACDC and EAP by better recovering the manual circuits with an average of 97% ROC AUC under low runtimes. In addition, we provide evidence that faithfulness of CD-T circuits is not due to random chance by showing our circuits are 80% more faithful than random circuits of up to 60% of the original model size. Finally, we show CD-T circuits are able to perfectly replicate original models' behavior (faithfulness ) using fewer nodes than the baselines for all tasks. Our results underscore the great promise of CD-T for efficient automated mechanistic interpretability, paving the way for new insights into the workings of large language models.
Paper Structure (30 sections, 4 equations, 7 figures, 2 tables, 1 algorithm)

This paper contains 30 sections, 4 equations, 7 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related work
  3. Causal Interpretation
  4. Circuit discovery methods
  5. Contextual Decomposition in Neural Networks
  6. Our method: CD-T and mechanistic interpretability
  7. Transformers
  8. Contextual decomposition for transformers (CD-T)
  9. Choosing a decomposition of source nodes for mechanistic interpretability
  10. Automating circuit discovery
  11. Definitions of circuit discovery
  12. Method
  13. Evaluating circuit discovery algorithms
  14. Experimental setup
  15. CD-T excels in identifying the circuit responsible for the underlying algorithm efficiently
  16. ...and 15 more sections

Figures (7)

  • Figure 1: Left: Log algorithm runtime and ROC AUC comparison. Each dot represents an average measurement on one task with specific methods differentiated by colors.; Right: The relative faithfulness of CD-T circuits compared to random circuits from the reference distribution of varying sizes (x-axis). Dotted vertical lines indicate the actual size of the circuits. $C$ denotes a CD-T circuit. $C^r$ denotes a random circuit. $M$ denotes the full model.
  • Figure 2: Faithfulness of CD-T circuits, EAP circuits, and randomly selected circuits of equivalent size for IOI, Greater-than and Docstring tasks. CD-T circuits obtains the full model’s performance (faithfulness of 1) faster than EAP as attention heads are added in order of importance.
  • Figure :
  • Figure :
  • Figure :
  • ...and 2 more figures