🤖 machine learning

AGOP as Explanation: From Feature Learning to Per-Sample Attribution in Image Classifiers

This paper introduces AGOP-Weighted, a novel post-hoc attribution method that leverages the Average Gradient Outer Product (AGOP) to suppress gradient noise and amplify consistently important pixels, demonstrating superior performance over existing techniques like Integrated Gradients and GradCAM on both synthetic and photorealistic benchmarks with minimal inference cost.

Original authors: Raj Kiran Gupta Katakam

Published 2026-05-14

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Raj Kiran Gupta Katakam

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 you have a neural network (a type of AI) that looks at pictures and guesses what they are. You want to know: "What exactly is the AI looking at when it makes that guess?"

For a long time, scientists have used tools to highlight the important pixels in an image, like a heat map showing where the AI is "focusing." But the paper you provided argues that these old tools are like trying to understand a person by only watching them once in a specific moment. They miss the bigger picture of how the AI learned to see things.

Here is a simple breakdown of the paper's new ideas, using everyday analogies.

1. The Old Way: The "Snapshot" Approach

Most current methods (like Integrated Gradients or GradCAM) act like a photographer taking a single snapshot of the AI's brain while it's solving a specific problem.

The Problem: If the AI is confused or if the math gets messy (like when pixels multiply each other), these snapshots get noisy or blurry. They might highlight the wrong spots or miss the signal entirely.
The Analogy: Imagine trying to figure out what a chef is good at by watching them chop one onion. You might see them struggle with that specific onion, but you miss the fact that they are actually a master at chopping carrots.

2. The New Idea: The "Training Diary" (AGOP)

The authors introduce a concept called AGOP (Average Gradient Outer Product). Instead of just looking at one snapshot, they look at the AI's entire "training diary."

What is it? It's a record of every time the AI made a mistake or a correction while learning. It tracks which pixels the AI consistently relied on across thousands of examples.
The Analogy: Instead of watching the chef chop one onion, you watch a video of them chopping onions, carrots, and potatoes for a whole month. You notice a pattern: "Ah, every time they see a round shape, they use a specific knife motion." This pattern is the AGOP.

3. The Three New Tools

The paper turns this "training diary" into three specific tools to explain the AI:

A. AGOP-Local (The "Standard Snapshot")

This is just the old way of looking at a single image, but the authors realized it's actually just a tiny piece of the bigger AGOP puzzle. It's like looking at one frame of the chef's video. It works okay, but it's not special.

B. AGOP-Weighted (The "Smart Highlighter")

This is the paper's star invention. It combines the "snapshot" of the current image with the "training diary."

How it works: It looks at the current image, but it uses the diary to say, "Hey, in the past, the AI always ignored this noisy pixel, but it always paid attention to that one." It boosts the important pixels and silences the noise.
The Analogy: Imagine a detective solving a crime. The old method looks at the suspect's current alibi. AGOP-Weighted looks at the alibi and cross-references it with the suspect's entire criminal history. If the suspect usually lies about Tuesdays, the detective knows to be extra skeptical of a Tuesday alibi.
The Result: On simple tasks, this method was 44% better at finding the right pixels than the previous best method.

C. AGOP-Global (The "Zero-Cost Cheat Sheet")

This tool doesn't look at the specific image at all. It just looks at the "training diary" to create a generic map of what the AI usually cares about.

How it works: If the AI always looks at the center of the image to find a "Tetris block," this tool draws a permanent spotlight on the center.
The Result: It costs nothing to use (you just look up a saved file). On tricky tasks where pixels multiply each other (making math very noisy), this tool was 7 times better than the old methods. It's like having a cheat sheet that says, "The answer is always in the middle," which works perfectly if the game is predictable.

4. What They Discovered (The "Gotchas")

The authors tested these tools on two different types of puzzles:

The "Linear" Puzzle (Simple): The AI just adds up pixel brightness.
- Winner: AGOP-Weighted. It combined the current view with the history perfectly.
The "Multiplicative" Puzzle (Tricky): The AI has to multiply pixel values. This creates a lot of mathematical noise that confuses standard tools.
- Winner: AGOP-Global. Because the noise was random, the "average" from the training diary smoothed it all out, revealing the true signal.
The "Moving Target" Puzzle: The object moves around the image.
- Winner: Integrated Gradients (The old method). Since the object moves, a "global" map (AGOP-Global) fails because it tries to highlight everywhere at once. The old method, which looks at the specific image, wins here.

5. A Warning About "GradCAM"

The paper also found that a popular tool called GradCAM completely fails on small, low-resolution images.

The Analogy: Imagine trying to read a tiny 8x8 pixel image, but the tool you are using only has a 3x3 grid to work with. It's like trying to paint a detailed portrait using only three giant brushstrokes. The detail is lost, and the tool points to the wrong place.

Summary

The paper argues that to truly understand an AI, you shouldn't just look at what it's doing right now. You should look at how it learned.

AGOP-Weighted is like a smart assistant that knows the AI's habits and helps you interpret its current thoughts.
AGOP-Global is a free, instant cheat sheet that works amazingly well when the AI's job is consistent.
The Big Takeaway: By using the AI's own training history (the AGOP), we can explain its decisions much more accurately and cheaply than before, especially when the math gets messy.

Technical Summary: AGOP as Explanation

Problem Statement

The field of Explainable AI (XAI) faces a central challenge in image classification: explaining which input pixels drive a neural network's prediction. Current attribution methods (e.g., Integrated Gradients, GradCAM, SmoothGrad) largely treat the trained model as a black box, applying geometric constructs independently of the network's learning dynamics. Furthermore, these methods often struggle with specific task characteristics, such as multiplicative interactions, spatial resolution constraints, or tasks requiring second-order reasoning. There is a lack of rigorous comparison between established baselines and methods grounded in the theoretical mechanics of feature learning.

Methodology

This work introduces AGOP-Weighted, a novel attribution method grounded in the Neural Feature Ansatz (NFA). The NFA posits that the weight Gram matrices of a neural network align with the Average Gradient Outer Product (AGOP) matrices computed over the training distribution. The authors leverage this theoretical insight to propose that the diagonal of the AGOP matrix, $diag(M)$, serves as a natural, training-distribution prior for attribution.

The methodology involves three core components:

AGOP Variants:
- AGOP-Local: Equivalent to VanillaGrad, representing the per-sample gradient.
- AGOP-Weighted: A hybrid method that multiplies the per-sample gradient by a normalized training-distribution prior, $v = \sqrt{diag(M)} / \max(diag(M))$ . This suppresses gradient noise and amplifies pixels consistently important across the training set.
- AGOP-Global: Uses $diag(M)$ directly as a zero-cost saliency map for the entire dataset, ignoring per-instance variations.
AGOPTrainingHook: An implementation mechanism that accumulates $diag(M)$ during standard training via a separate backward pass (using autograd.grad) without altering weight updates. This allows for the efficient computation of the global prior.
Evaluation Framework: The authors conduct a rigorous comparison against Integrated Gradients (IG), SmoothGrad, GradCAM, and VanillaGrad across two benchmarks with pixel-level ground truth:
- XAI-TRIS: A synthetic benchmark with 8×8 images and four classification scenarios (Linear, Multiplicative, Translations+Rotations, XOR).
- CLEVR-XAI: A photorealistic benchmark using ResNet-18 fine-tuned on 3D scenes.

Key Contributions

Novel Attribution Formulation: The introduction of AGOP-Weighted, which bridges per-sample and global attribution by modulating gradients with a training-distribution prior.
Efficient Implementation: The development of AGOPTrainingHook, enabling AGOP-Global to operate at zero inference cost (disk lookup) and AGOP-Weighted at a single gradient pass cost.
Systematic Evaluation: The first rigorous comparison of AGOP-based methods against standard baselines across diverse task types (linear, multiplicative, spatially varying) and resolutions.
Empirical Insights:
- Demonstration that AGOP-Weighted achieves significantly higher mean Intersection over Union (mIoU) than Integrated Gradients on linear tasks.
- Demonstration that AGOP-Global drastically outperforms Integrated Gradients on multiplicative tasks (where IG fails below random chance) at zero inference cost.
- Identification of GradCAM's failure on small-resolution images due to spatial resolution collapse.
- Observation that the quality of $diag(M)$ improves monotonically throughout training, even after classification accuracy has plateaued.

Results

The experiments yield distinct performance profiles based on task characteristics:

Linear Tasks (XAI-TRIS): AGOP-Weighted achieves a 44% higher mIoU than Integrated Gradients (0.314 vs. 0.218) with half the inference cost (1.0 ms vs. 2.4 ms).
Multiplicative Tasks (XAI-TRIS): AGOP-Global dominates, achieving 7× the mIoU of Integrated Gradients (0.337 vs. 0.048). In this scenario, IG falls below the random baseline, while AGOP-Global's averaging of squared gradients cancels out per-sample noise to recover the stable signal.
Spatially Varying Tasks (Translations+Rotations): Integrated Gradients outperforms all methods (mIoU 0.540) because it tracks specific per-instance shapes, whereas AGOP-Global fails (mIoU 0.080) due to the diffusion of global statistics.
CLEVR-XAI (ResNet-18): Both AGOP-Weighted (+18% mIoU over IG) and AGOP-Global (+37% mIoU over IG) outperform Integrated Gradients. AGOP-Global achieves competitive results at 6,600× lower inference cost (0.02 ms/sample vs. 131.98 ms/sample).
GradCAM Limitations: GradCAM and GradCAM++ fail completely on small-resolution images (8×8 and 64×64 inputs) due to the collapse of spatial feature maps (e.g., 3×3 or 2×2 maps), resulting in a Pointing Game score of 0.000.

Significance and Claims

The paper claims to establish AGOP not merely as a theoretical construct but as a practical, post-hoc attribution tool. The significance lies in:

Bridging Theory and Practice: Validating the Neural Feature Ansatz by demonstrating that gradient statistics accumulated during training ($diag(M)$) provide superior explanatory power compared to methods that ignore the learning process.
Task-Dependent Superiority: Showing that no single method dominates all scenarios; however, AGOP variants offer specific, high-value advantages: AGOP-Weighted for linear tasks and AGOP-Global for multiplicative or positionally consistent tasks at zero inference cost.
Efficiency: Providing a method (AGOP-Global) that requires no inference-time computation, making it ideal for batch pipelines and dataset-level diagnostics.
Limitations of Current Baselines: Highlighting that standard methods like GradCAM are fundamentally unsuited for low-resolution inputs and that Integrated Gradients can fail catastrophically on multiplicative interactions.

The authors conclude that while AGOP-Weighted and AGOP-Global offer significant improvements in specific contexts, the choice of attribution method must be guided by the nature of the task (linear vs. multiplicative, fixed vs. varying spatial signal). Future work is suggested to extend these findings to larger networks and explore off-diagonal interaction terms.