Can Graph Learning Learn Circuits?

Circuit localization aims to identify a sparse subgraph (a 'circuit') within a transformer's computation graph that is sufficient to reproduce a specific behavior. Traditionally, this task is performed independently for each model-task pair, which is computationally expensive and fails to leverage regularities across different models. This paper asks whether graph machine learning can instead learn to identify these circuits in an amortized, cross-model fashion.

The authors frame circuit localization as a graph learning problem where components (e.g., attention heads, MLP blocks) are nodes and computational pathways are edges. They introduce Graph Circuit Learning (GCL), a supervised framework that trains a graph neural network (GNN) on a set of labeled model-task pairs to predict circuit masks for unseen cases. To facilitate this, they augment the existing InterpBench benchmark with additional cases derived from TracrBench. They also adapt PGExplainer, a GNN-specific explainability method, to serve as a baseline for this task.

The best-performing GCL configuration achieved a median edge AUROC of 0.902 on held-out test cases, which is competitive with established baselines like EAP-IG (0.910) and outperforms the adapted PGExplainer (0.858). Ablation studies show that removing message-passing edges significantly degrades performance, suggesting that the GNN successfully exploits the structural dependencies between computational pathways. These results indicate that graph learning provides a natural and effective inductive bias for mechanistic interpretability.

This work shifts the paradigm of circuit localization from isolated, per-model optimization to a transferable, amortized learning task. By demonstrating that GNNs can learn to identify circuits across different models, the authors open a new avenue for scaling mechanistic interpretability to larger, more complex systems where individual manual or per-case analysis is infeasible.