💬 NLP

Correcting Gradient-Based Circuit Localization via Interaction-Aware Backpropagation

This paper introduces Gradient Interaction Modifications (GIM), a novel technique that corrects systematic misestimations in circuit localization caused by ignoring component interactions like attention self-repair, thereby achieving state-of-the-art performance in identifying the model components responsible for specific behaviors in large language models.

Original authors: Joakim Edin, Casper L. Christensen, Róbert Csordás, Tuukka Ruotsalo, Zhengxuan Wu, Maria Maistro, Jing Huang, Lars Maaløe

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Joakim Edin, Casper L. Christensen, Róbert Csordás, Tuukka Ruotsalo, Zhengxuan Wu, Maria Maistro, Jing Huang, Lars Maaløe

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a Large Language Model (LLM) as a massive, bustling kitchen where thousands of chefs (neurons and attention heads) work together to cook a single dish (the model's answer). For a long time, scientists trying to understand how these models work have used a method called "circuit localization." Their goal is to figure out exactly which chefs are responsible for the flavor of the dish.

The Old Way: The "One-at-a-Time" Mistake
Traditionally, researchers tried to identify important chefs by firing them one by one. They would say, "Let's send Chef A home for the day and see if the soup tastes different." If the soup tastes the same, they assume Chef A wasn't important.

The problem is that these chefs don't work in isolation; they talk to each other and cover for one another. If Chef A is a master of chopping onions, but Chef B is also a master of chopping onions and is standing right next to them, sending Chef A home won't change the soup. Chef B just picks up the knife and keeps chopping. The researchers, seeing no change, wrongly conclude that Chef A was useless.

The New Discovery: "Attention Self-Repair"
The authors of this paper discovered a specific way this "covering for each other" happens in AI, which they call Attention Self-Repair.

Think of the AI's attention mechanism like a spotlight operator at a theater. The spotlight operator (the model) decides which actors (words) to shine a bright light on. Sometimes, two actors are standing in the same spot and saying the exact same thing. The operator splits the light 50/50 between them.

If you try to dim the light on Actor A to see if they are important, the operator instantly shifts that extra light to Actor B. Because Actor B is saying the same thing, the audience (the model's output) doesn't notice any difference. The "gradient" (the signal telling us who is important) disappears, making it look like neither actor mattered. But in reality, both were crucial; they just had a backup plan.

The Solution: GIM (Gradient Interaction Modifications)
To fix this, the authors created a new tool called GIM. Instead of firing one chef or dimming one spotlight at a time, GIM changes the rules of the game to account for the teamwork. It uses three clever tricks:

The "Warm Spotlight" (Temperature-adjusted Softmax):
Imagine the spotlight operator is usually very strict, giving 90% of the light to the top actor and 10% to everyone else. GIM tells the operator to be a bit more relaxed (increase the "temperature"). Now, the light is spread more evenly. When you dim one actor's light, the others don't just instantly take over because the light was already shared more broadly. This reveals who was actually holding the light, even if they had help.
Freezing the "Volume Knob" (Layernorm Freeze):
In the kitchen, there's a master volume knob that adjusts the overall loudness of the room. If one chef stops working, the volume knob automatically turns up the others to keep the noise level constant. This hides the fact that a chef left. GIM "freezes" this volume knob, so when a chef leaves, the others don't get artificially louder. This lets researchers see the true drop in volume caused by that chef leaving.
The "Shared Tool" Correction (Gradient Norm):
Sometimes, two chefs are using the same tool (like a shared knife) to chop. If you measure how much the knife was used, you might accidentally count the same chop twice—once for each chef. GIM fixes this math error by dividing the credit correctly, ensuring no one gets double-counted for the same action.

The Results
When the authors tested GIM on various AI models and tasks (like solving math problems or answering questions), it worked much better than the old methods.

It found the "real" important parts of the model that the old methods missed.
It proved that the old methods were systematically underestimating how important certain parts of the AI were because they didn't understand the "self-repair" teamwork.

In Summary
The paper argues that we can't understand complex AI models by looking at their parts in isolation. Because the parts help each other, removing one doesn't always show its true value. GIM is a new way of looking at the model that accounts for this teamwork, giving us a much more accurate map of how these digital brains actually think.

Technical Summary: Correcting Gradient-Based Circuit Localization via Interaction-Aware Backpropagation

Problem Statement

Circuit localization aims to identify the specific subset of model components (e.g., attention heads, neurons, connections) responsible for specific behaviors in Large Language Models (LLMs) to enable mechanistic interpretability. Most existing methods operate under the assumption that components act independently, estimating importance by perturbing each component in isolation.

However, neural network components interact, and ignoring these interactions leads to systematic misestimation of component importance. The authors identify a specific, problematic interaction termed attention self-repair. In this phenomenon, when multiple positions possess large attention weights and contain similar information (value vectors), ablating a single attention score has minimal effect on the output. This occurs because the softmax activation function redistributes weight to the remaining redundant positions to compensate. Consequently, standard perturbation-based and gradient-based explanation methods underestimate the importance of components contributing to these attention scores, often resulting in vanishing gradients for influential scores.

Methodology: Gradient Interaction Modifications (GIM)

The paper introduces Gradient Interaction Modifications (GIM), a gradient-based attribution method designed to estimate the impact of perturbing multiple interacting components simultaneously. GIM modifies standard backpropagation through three complementary techniques:

Temperature-adjusted Softmax Gradients (TSG):
- Motivation: Addressing attention self-repair, where standard gradients vanish because the gain from boosting one softmax weight is canceled by the loss from suppressing others when values are similar.
- Mechanism: During backpropagation, TSG recomputes the softmax with an elevated temperature ( $T > 1$ ) before calculating gradients. This flattens the probability distribution, shrinking large weights and growing small ones. By enlarging the set of "significant" positions, the assumption that contributions are identical across these positions is relaxed. This ensures that pairwise differences in gradients ( $g_j - g_k$ ) do not vanish, allowing the gradient signal to survive and approximate the effect of simultaneous perturbation.
Layernorm Freeze:
- Motivation: Layer normalization introduces a zero-sum dynamic where ablating one component reduces activation magnitude, causing the normalization factor to amplify remaining components (a form of self-repair).
- Mechanism: The normalization factor ( $\sigma$ ) is treated as a constant during backpropagation. This measures each component's direct contribution before the normalization rescaling occurs, preventing the underestimation of importance caused by compensatory amplification.
Gradient Norm:
- Motivation: Multiplicative operations (e.g., $A = QK^T$ where both $Q$ and $K$ depend on input $X$ ) lead to double-counting in standard first-order gradients. Perturbing $X$ affects both $Q$ and $K$ simultaneously, but standard chain rules sum their individual effects, overestimating the change.
- Mechanism: The gradients of multiplicative interactions are divided by the number of multiplicative inputs (e.g., dividing by 2 for $Q \times K$ ). This aligns with the "uniform rule" and corrects the overcounting of contributions in attention-value multiplication, query-key multiplication, and MLP gate-projection multiplication.

Key Contributions

GIM Framework: The introduction of a gradient-based attribution method that explicitly accounts for feature interactions, achieving state-of-the-art performance on the circuit localization track of the Mechanistic Interpretability Benchmark (MIB).
Identification of Attention Self-Repair: The formal identification and proof that attention self-repair (an OR-gate-like interaction) causes gradients to vanish for influential attention scores, leading to systematic underestimation of component importance.
Feature Attribution Superiority: Demonstration that GIM outperforms existing gradient-based methods in feature attribution across diverse tasks and LLM architectures.

Results

Circuit Localization: On the MIB, GIM significantly outperforms existing methods (including EAP, EAP-IG, Hybrid-ens, and IFR) across multiple models (GPT-2, Qwen-2.5, Gemma-2, Llama-3.1) and tasks (Indirect Object Identification, Arithmetic, MCQA, ARC). It achieves the highest average Circuit Performance Ratio (CPR), particularly on tasks requiring multi-step reasoning.
Feature Attribution: Across seven instruction-tuned LLMs and six datasets (covering QA, fact verification, sentiment, and hate speech detection), GIM generally achieves higher comprehensiveness and sufficiency scores compared to GradientXInput, Integrated Gradients, DeepLIFT, and TransformerLRP.
Validation of Self-Repair: Empirical analysis confirms that attention self-repair is prevalent in real LLMs, occurring on average 65–1200 times per input. Ablating two large attention scores jointly yields a significantly larger effect than the sum of their individual ablations, validating the theoretical analysis.
Ablation Study: Removing any single modification (TSG, Layernorm Freeze, or Gradient Norm) from GIM degrades performance across most model-dataset pairs, confirming that each component contributes meaningfully to the method's faithfulness.

Significance

The paper claims that GIM enables more faithful mechanistic analysis of LLMs by correcting the systematic underestimation of component importance caused by feature interactions. By explicitly modeling how components interact—specifically through softmax redistribution and normalization rescaling—GIM provides a more accurate picture of the internal mechanisms driving model behavior. The authors position GIM not just as an incremental improvement, but as a necessary correction to the "independence assumption" prevalent in current interpretability tools, allowing researchers to identify circuits that were previously obscured by vanishing gradients. The method is released as an open-source Python package to facilitate broader adoption.