Efficient Automated Circuit Discovery in Transformers using Contextual Decomposition

Authors: Aliyah R. Hsu, Georgia Zhou, Yeshwanth Cherapanamjeri, Yaxuan Huang, Anobel Odisho, Peter R. Carroll, Bin Yu
Primary area: interpretability and explainable AI
Keywords: Automated Circuit Discovery, Explainable AI, Interpretation, Machine Learning, Language Models, Transformers

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 at any level of abstraction and is the first to efficiently produce circuits as fine-grained as attention heads at specific sequence positions.
CD-T is compatible to all transformer types, and requires no training or manually-crafted examples.
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.

Decision: accept (poster)
Reviewer scores: 5, 6, 8

Reviews and discussion:
<thread>
creation date:
1730486399087

thread_type:
official_review

<thread_object>
creator: reviewer_lhid

<content>
summary:
The paper proposes an automatic, mechanistic interpretability approach for finding circuits in transformer models. Their approach propagates backwards from the model output, finding a circuit of nodes most relevant to the given output, and from there the nodes most relevant to those nodes, and so on. These circuits are task-specific subgraphs in the larger model compute graph that are most stimulated (and most relevant) to the input type or task under consideration. 

The authors consider circuit discovery for the attention heads of transformer models, choosing to leave the MLP layers in an attention block for later work. Their linear decomposition model is mechanistic in that it solely depends on the weights and biases of each layer of the network and the relevance metric calculated for the previous layer. It does appear to be data-dependent in that it still relies on a dataset for calculating these relevance metrics. It also relies on precomputed mean activations to do so.

strengths:
1) The paper's originality is in applying a novel formulation of an existing method - Context Decomposition - to transformer models, albeit to the attention heads of the transformer models. While novel, I am not in the space of circuit discovery in neural networks and unable to judge how novel this is. 
2) The presentation is clear and the limitations in data dependence and precomputation are made clear to the reader in algorithm pseudocode, which is also concise and clear. Experimentation is presented to clearly establish the work's superior performance over the baselines. Notation remains consistent through the paper to my eyes.
3) Significance: The work presents the an automatic circuit discovery approach using CD (extending it as CD-T) in transformer models, and clearly outperforms its chosen baselines. The approach is moreover mechanistic and iterative rather than probabilistic and therefore also comes with a guarantee of circuit discovery given certain assumptions on available data and time.

weaknesses:
1) I am unsure of the significance of the work due to the iterative nature of the process and the need for data to adequately stimulate different 'tasks' (or behavior modes, perhaps?) and therefore discover circuits in the model. This is to my eyes not clearly explained. [Addressed]
2) Prior work is noted to have conducted circuit discovery on large language models, and I can't easily find the models used for circuit discovery here. That lack of clarity is concerning, especially when using the manually discovered circuits as a target and prior art as baselines. [Addressed]
3) I am also uncertain of scalability of this approach, considering that the network architectures under evaluation are not clear and that it has been restricted to attention heads rather than MLP layers. I am also uncertain of the effects of sequence length, autoregressive or non-autoregressive decoders. Section 5 does not make these things clear. [Addressed in part; prior art remains deficient as well here - future work?]

questions:
1) I would suggest clearly presenting model architectures under evaluation and assumptions being made on input data for the approach - the mechanistic relevance metrics are ultimately calculated from the outputs derived from input data.
2) How would this scale to large training datasets in LLMs? Is there a minimum dataset required? Can the authors comment on this?
3) What is the effect of autoregressive models and sequence to sequence decoding on this approach? Is there a difference? I don't see a clear analysis in Section 5, and while that is not necessary it would be nice to have the authors comment on this in Section 6.
Questions addressed in comments.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730656326591

thread_type:
official_review

<thread_object>
creator: reviewer_tnfo

<content>
summary:
The paper introduces Contextual Decomposition for Transformers (CD-T) for building interpretable circuits in llms for better mechanistic interpretability in transformers efficiently. In contrast to other methods, CD-T employs a mathematical decomposition that isolates the contributions of specific features, allowing CD-T to discover circuits at various levels of granularity. Results show significant improvements in runtime and interpretability over existing methods such as ACDC and EAP on several tasks.

strengths:
1.	The paper is relatively well presented. It provides a clear description of the method, experiment details, and contribution. 
2.	Using the decomposition approach for designing/scaling LLM models is relatively novel and effective. Results show promising improvements in computational efficiency while maintaining circuit quality, and the method seems to be agnostic to the specific types of transformer models.

weaknesses:
1.	The description of the algorithm can be improved. In particular, the algorithm description is relatively informal and could benefit from more details and formalism. For “Prune nodes from S for which doing so increases the task metric”, how is increasing the task metric evaluated?

questions:
1. See W1; In addition, what’s the computational complexity of the algorithm in terms of the input parameter?
2. How does pruning affect the mechanistic interpretability of the model?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730756064826

thread_type:
official_review

<thread_object>
creator: reviewer_zeek

<content>
summary:
This paper introduces a novel method CD-T that leverages contextual decomposition for mechanistic interpretability in transformers. CD-T's granularity is fine-grained up to attention heads (does not include MLP layers). They report the results of their method (classification metrics and faithfulness of the circuit as compared to randomly sampled circuits) for three tasks, namely indirect object identification, greater-than comparisons, and docstring completion.

strengths:
1. Motivation, relevant scientific terms and prior works are well-written, making the paper accessible to researchers who are non-experts in this area. 

2. The proposed method can be easily used for all transformer architectures (encoder-decoder and decoder-only, as per line 72). Further, the proposed method significantly reduces the runtime required for inference.

weaknesses:
1. Manual circuits are not fully explained (definition, how they are computed, cost of computation, etc.) - manual circuits are used as reference during evaluation (line 405), and it is necessary to provide these details regarding them. 

2. It is unclear how CD-T works for inference - is there a distinct circuit discovered for each inference datapoint, or is the circuit found on a broader scale, so as to tradeoff size vs. performance for a larger test set?

3. The authors are encouraged to discuss the practical utility of CD-T. Once their method yields a circuit, could additional analyses, such as on attention heads, be conducted? For instance, are there heads that consistently produce positive or negative outputs across all data points, or do the importance of attention heads vary based on input data characteristics?

4. What transformer architecture is used for the experimental results?

questions:
Line 20: "CD-T is the first ---?" there seems to be a missing word here
</content>
</thread_object>

</thread>


Metareview:
metareview:
Thank you for your submission to ICLR. This paper presents CD-T (contextual decomposition for transformers), a method which aims to build interpretable circuits in large language models. It computes these circuits efficiently, with possible reductions in runtime from hours to seconds, in comparison with existing baselines. In addition, CD-T performs well empirically, compared to baselines, on multiple circuit evaluation datasets.

Reviewers appreciate the clarity of presentation, the method novelty, and the performance of the method. While there was some concern about certain relevant background descriptions and definitions, as well as the method’s scalability, the authors eventually addressed all of the reviewers’ concerns.