EAP-GP: Mitigating Saturation Effect in Gradient-based Automated Circuit Identification

2025-02-07

AI Summary

Gradient-based circuit identification methods suffer from zero-gradients and saturation effects, resulting in edge attributions that are insensitive to input perturbations, noisy, and causally unfaithful. To address this, we propose GradPath—a joint framework integrating path integration with Edge Attribution Patching (EAP). GradPath introduces an adaptive gradient-difference-based integration path that dynamically bypasses saturated regions; moreover, it is the first method to embed attribution patching directly into the path-integration paradigm, thereby enhancing attribution stability and causal fidelity. We evaluate GradPath across six benchmark datasets and the GPT-2 family of models. Results show up to a 17.7% improvement in circuit faithfulness over prior approaches, while precision and recall match or exceed human-annotated ground truth—achieving state-of-the-art performance.

Technology Category

• Machine Learning: Causal Learning
• Computer Vision: Adversarial Attacks & Robustness
• Reasoning under Uncertainty: Causality

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Understanding the internal mechanisms of transformer-based language models remains challenging. Mechanistic interpretability based on circuit discovery aims to reverse engineer neural networks by analyzing their internal processes at the level of computational subgraphs. In this paper, we revisit existing gradient-based circuit identification methods and find that their performance is either affected by the zero-gradient problem or saturation effects, where edge attribution scores become insensitive to input changes, resulting in noisy and unreliable attribution evaluations for circuit components. To address the saturation effect, we propose Edge Attribution Patching with GradPath (EAP-GP), EAP-GP introduces an integration path, starting from the input and adaptively following the direction of the difference between the gradients of corrupted and clean inputs to avoid the saturated region. This approach enhances attribution reliability and improves the faithfulness of circuit identification. We evaluate EAP-GP on 6 datasets using GPT-2 Small, GPT-2 Medium, and GPT-2 XL. Experimental results demonstrate that EAP-GP outperforms existing methods in circuit faithfulness, achieving improvements up to 17.7%. Comparisons with manually annotated ground-truth circuits demonstrate that EAP-GP achieves precision and recall comparable to or better than previous approaches, highlighting its effectiveness in identifying accurate circuits.

Problem

• Mitigate saturation effects in circuit identification.
• Enhance reliability of gradient-based attribution methods.
• Improve faithfulness in transformer model circuit discovery.

Innovation

• Edge Attribution Patching
• GradPath integration
• Mitigates saturation effects