Can Graph Learning Learn Circuits?
Introduces Graph Circuit Learning (GCL), using GNNs for cross-model circuit localization, achieving median edge AUROC of 0.902.
Key Findings
Methodology
This work models transformer computation graphs as component-level DAGs, with nodes representing modules and edges pathways. GCL trains a GNN across multiple model-task pairs with ground-truth masks, employing architectures like DirGNN and GraphConv, combined with feature alignment and pooling. It leverages supervised learning to predict circuit masks, enhancing cross-model generalization. Additionally, PGExplainer is adapted for edge explanation using fixed features. Data augmentation extends InterpBench with 30 TRACRBench pairs, enabling robust evaluation of transferability. The framework is trained end-to-end with binary cross-entropy loss, capturing complex pathway interactions.
Key Results
- On 16 held-out InterpBench cases, the top GCL configuration achieves a median edge AUROC of 0.902 (IQR [0.861, 0.942]), close to the single-case median of 0.910, and below ACDC’s 0.959. Removing message-passing edges drops median AUROC to 0.825. The adapted PGExplainer yields 0.858 AUROC, outperforming some baselines.
- Cross-case training significantly improves transfer performance, validating the approach’s effectiveness. Pathway interaction modeling is crucial, as simple edge scoring underperforms.
- Results demonstrate that graph structures naturally encode pathway dependencies, offering a more principled approach than traditional methods. GCL’s success on synthetic data suggests promising future applications on real models.
Significance
This study pioneers a shift from isolated model-specific circuit localization to a transfer-capable, graph-based framework. By leveraging GNNs to learn pathway interactions, it enhances interpretability and generalization, addressing longstanding challenges in mechanistic understanding of transformers. The approach opens avenues for automated circuit discovery, model debugging, and compression, fostering transparency in AI systems. Its ability to transfer across models and tasks marks a significant step toward universal mechanistic interpretability, with broad implications for AI safety and trustworthiness.
Technical Contribution
The paper formalizes circuit localization as a graph learning task, introducing GCL that trains a GNN across multiple model-task pairs with supervision. It combines architectures like DirGNN and GraphConv, integrating feature alignment and pooling to handle heterogeneous node features. The adaptation of PGExplainer for edge explanation further enhances interpretability. Data augmentation with TRACRBench enables cross-model transfer evaluation. The end-to-end training pipeline demonstrates superior performance over traditional single-case methods, establishing a new paradigm for mechanistic interpretability using graph neural networks.
Novelty
This work is the first to formalize circuit localization as a cross-model graph learning problem, leveraging supervised GNN training across multiple model-task pairs. It innovates by integrating feature alignment, message passing, and explainability techniques, enabling the model to generalize beyond single cases. The combination of synthetic data augmentation and architecture design sets a new standard for transferability in mechanistic interpretability, distinguishing it from prior isolated or heuristic approaches.
Limitations
- The current validation is limited to synthetic transformer models and semi-synthetic datasets, leaving real-world applicability untested. Performance may vary with increased complexity and data scarcity.
- Supervised training requires high-quality ground-truth circuit masks, which are difficult to obtain for real models, limiting scalability.
- Computational costs are high, especially for large models and complex pathway interactions, necessitating further optimization and efficiency improvements.
Future Work
Future directions include applying GCL to large-scale real-world models, developing unsupervised or weakly supervised methods to reduce annotation dependency, and exploring more scalable graph architectures. Investigating dynamic pathways and temporal interactions could deepen understanding of model evolution under input variations. Extending transfer learning capabilities across different architectures and tasks will further enhance the framework’s utility, ultimately aiming for universal, automated circuit discovery that supports model transparency and robustness.
AI Executive Summary
Understanding the internal mechanisms of transformer models remains a central challenge in AI interpretability. Traditional approaches focus on analyzing individual models or specific tasks, which limits their ability to generalize across different architectures and inputs. This paper introduces Graph Circuit Learning (GCL), a novel framework that leverages graph neural networks to learn the interactions among computational pathways in transformer models. By representing the model’s computation graph as a component-level DAG, GCL is trained across multiple model-task pairs with supervision, enabling it to predict circuit masks for unseen cases.
The core innovation lies in modeling pathway interactions explicitly, capturing complex dependencies that traditional scoring methods overlook. The architecture employs architectures like DirGNN and GraphConv, combined with feature alignment and pooling, to handle heterogeneous node features and pathway relationships. Extensive experiments on an augmented InterpBench dataset, including 30 new TRACRBench pairs, demonstrate that GCL achieves a median edge AUROC of 0.902, outperforming baseline methods and single-case approaches. The results highlight the importance of message-passing edges, which significantly influence performance.
Additionally, the study adapts PGExplainer, a GNN explainability tool, to the circuit localization task, further validating the potential of graph-based interpretability. These findings suggest that graph learning offers a natural, powerful perspective for mechanistic understanding, enabling transferability and automation. While current validation is primarily on synthetic data, the approach opens promising pathways for analyzing real-world large models, advancing transparency and trust in AI systems. Future work aims to extend these methods to more complex, real-world scenarios, reducing reliance on annotated ground-truths and improving scalability.
Deep Analysis
Background
神经网络的机制解释一直是AI研究的核心难题。早期方法多依赖逐层特征重要性分析，近年来变换器模型的广泛应用带来了路径交互的复杂性。传统电路定位多为逐模型、手工分析，难以实现跨模型泛化。InterpBench和TRACRBench提供半合成数据集，推动机制研究，但仍受限于数据规模和多样性。图神经网络（GNN）在结构建模中表现优异，为路径交互建模提供新工具。此前研究多关注路径重要性排序，缺乏跨模型迁移能力。本文借助图学习，试图突破这一瓶颈，推动机制理解的系统化和自动化。
Core Problem
核心问题在于如何实现模型间的机制迁移。传统方法多为逐模型训练，缺乏泛化能力，路径交互关系复杂，单一路径评分不足以捕获全部机制信息。现有数据不足以支持大规模迁移，限制了机制解释的普适性。需要设计具有迁移能力的模型架构，充分利用路径间的结构关系，实现跨模型的机制识别。这一挑战阻碍了变换器机制的全面理解，也限制了其在实际中的应用。
Innovation
主要创新包括：1）将电路定位转化为图机器学习任务，利用路径交互关系建模；2）提出GCL框架，通过多模型-任务监督训练GNN，增强泛化能力；3）引入DirGNN和GraphConv架构，结合特征对齐和池化机制，提升模型表达；4）借鉴PGExplainer，设计边解释器，增强模型可解释性。这些创新使得电路定位不再局限于单模型，能实现跨模型迁移，显著提高机制识别的效率和准确性。
Methodology
- �� 将变换器的计算图表示为组件级有向无环图（GM），节点代表注意头或MLP块，边代表路径。
- �� 利用模型在干净和受损提示上的激活，提取边特征。
- �� 设计多模型训练策略：用ground-truth掩码监督GNN预测路径掩码。
- �� 采用DirGNN或GraphConv架构，结合特征对齐（TabPFN风格的特征注意）和池化（跨路径、跨输入）机制。
- �� 训练端到端，优化二元交叉熵损失，提升跨案例泛化。
- �� 引入PGExplainer，利用固定特征训练边解释器，增强可解释性。
Experiments
- �� 数据集扩展：在InterpBench基础上加入30个TRACRBench派生的模型-任务对，设计交叉验证。
- �� 评估指标：边AUROC，比较不同GNN架构和基线（如无消息传递、单模型方法）。
- �� 超参数调优：采用分组5折交叉验证，平均5个随机种子结果。
- �� 重点分析消息传递边对性能的贡献，验证路径交互建模效果。
Results
- �� GCL最高配置在16个测试案例中，边AUROC中位数达0.902，优于去除消息边的0.825和PGExplainer的0.858。
- �� 迁移性能显著提升，验证多模型训练的有效性。
- �� 实验显示路径交互关系复杂，单路径评分不足以捕获全部机制信息，图结构提供更自然的建模方式。
Applications
- �� 适用于模型机制分析、故障诊断、模型压缩等场景，帮助理解变换器内部路径作用。
- �� 结合合成和真实模型数据，推动模型透明化和可信度提升。未来可用于自动化机制发现和模型调试。
Limitations & Outlook
- �� 目前主要在合成数据上验证，真实模型中的表现仍需验证，可能受限于数据规模和路径复杂性。
- �� 训练依赖大量标注掩码，实际应用中标注困难。
- �� 计算成本较高，尤其是在大规模模型和复杂路径交互时，未来需优化效率。
Plain Language Accessible to non-experts
想象一个工厂里有很多不同的机器和传送带，每个机器都在完成特定任务。我们想知道哪些传送带和机器组合在一起，才是完成某个特定任务的关键路径。传统方法就像是逐个机器分析，找出重要的传送带，但这样很费时间，也难以在不同工厂之间迁移。本文用一种类似于交通网络的图模型，把所有路径和交互关系画出来，然后用智能算法学习这些路径之间的关系。这样一来，不管工厂长得多复杂，只要学会了这个图模型，就能快速找到完成任务的关键路径。这个方法就像是让工厂的交通变得更智能、更高效，未来还能帮忙诊断故障或优化流程。
ELI14 Explained like you're 14
你知道在学校里，有很多不同的老师和课程，每个老师都用不同的方式教书。现在，有个问题是：如何找到哪些老师和课程组合，才是让学生学得最好的关键？以前的方法就像是每个老师自己分析，找出重要的课程，但这样很麻烦，也不能在不同学校之间用。这个研究就像是用一张大地图，把所有老师、课程和学生的学习路径都画出来，然后用聪明的电脑程序学习这些路径之间的关系。这样一来，不管是哪所学校，只要学会了这张地图，就能快速找到最重要的学习路径，帮助老师和学生更好地理解学习过程。这就像是给学校装上了智能导航，让学习变得更简单、更高效！
Abstract
Circuit localization is a mechanistic interpretability task whose goal is to identify a sparse subgraph of a transformer's computation graph sufficient to reproduce a particular behavior. Most established methods localize circuits independently for each model--task pair. We instead frame circuit localization as a graph machine learning problem in which the edges of a computation graph represent computational pathways, and graph neural networks (GNNs) model interactions among these pathways. We introduce Graph Circuit Learning (GCL), a supervised, amortized framework that trains a GNN across multiple model--task pairs and applies it to unseen cases. To provide sufficient data, we augment the InterpBench benchmark with additional cases derived from the TracrBench programs. Of the 14 evaluated GCL configurations, the highest scored a median edge AUROC of $0.902$ (interquartile interval $[0.861, 0.942]$) on the 16 original held-out InterpBench cases. This is close to the published InterpBench median of $0.910$ for EAP-IG while remaining below ACDC's $0.959$. Removing all message-passing edges reduces the median to $0.825$. We also adapt PGExplainer, a GNN explainability method, to circuit localization, obtaining a median edge AUROC of $0.858$ on the same cases. These preliminary results suggest that graph machine learning offers a natural and potentially powerful perspective on circuit localization, and we hope this perspective encourages closer exchange between the two communities.