Learn to Rank: Visual Attribution by Learning Importance Ranking

1. Problem Statement

The paper addresses the challenge of Visual Attribution (generating heatmaps to explain model decisions) for complex computer vision models, particularly Vision Transformers (ViTs). Existing methods face a "three-way trade-off":

• Propagation-based methods (e.g., Grad-CAM, Attention Rollout): Efficient and model-consistent but often biased toward low-level features (edges) and architecture-specific. They may fail to capture non-attention pathways in Transformers.
• Perturbation-based methods (e.g., RISE, SHAP, MDA): Causally grounded and architecture-agnostic but computationally expensive (requiring thousands of forward passes). For ViTs, they often operate on patch-level tokens, resulting in coarse, blocky attributions that lack pixel-level precision.
• Learning-based (Amortized) methods: Fast at inference but typically rely on surrogate objectives or heuristic teachers (like other attribution methods), inheriting their biases or making restrictive assumptions about data distributions.

The Core Gap: There is a lack of a method that is simultaneously causally grounded (like perturbation), efficient (like learning-based), and capable of producing dense, pixel-level explanations for Transformer models without relying on heuristics.

2. Methodology: Amortized Hybrid Attribution (AHA)

The authors propose AHA, a learning-based framework that directly optimizes standard attribution quality metrics (Deletion and Insertion) via a differentiable relaxation of the sorting process.

A. Objective: Direct Metric Optimization

Instead of using heuristics, AHA trains an explainer model $\Phi_{f,\theta}$ to minimize the Deletion AUC and maximize the Insertion AUC.

• Deletion: Measures how quickly the model's confidence drops when high-attribution pixels are removed.
• Insertion: Measures how quickly confidence rises when high-attribution pixels are added to a blank reference.
• Challenge: These metrics require sorting pixels by importance and creating binary masks, which involves non-differentiable operations (sorting/permutation).

B. Differentiable Relaxation via Gumbel-Sinkhorn

To enable end-to-end training, the authors reformulate the sorting problem as a permutation learning task:

1. Soft Permutation: Instead of hard sorting, they use the Gumbel-Sinkhorn algorithm to generate a soft, doubly-stochastic permutation matrix $P^{soft}$. This approximates the sorting of attribution scores $a$ against target ranks $p$.
2. Soft Top-k Masks: By taking the cumulative sum of the rows of the soft permutation matrix, they generate soft top-k masks ($m^{soft}$) with continuous values in $[0, 1]$.
3. Differentiable Perturbation: These soft masks are used to blend the input image with a reference image (e.g., black or blurred), creating differentiable perturbed images. This allows gradients to flow from the target model's output back to the explainer's parameters.

C. Training Efficiency and Robustness

To make training feasible for high-resolution images, AHA employs several strategies:

• Region-Based Permutation: Instead of sorting $N$ individual pixels, the image is divided into a $G \times G$ grid. Sorting happens at the region level, and masks are upsampled. This reduces computational complexity and prevents overfitting to high-frequency pixel artifacts.
• Sampling Perturbation Steps: Instead of evaluating all $K$ steps for the AUC curve, the method samples a subset $S$ of steps during training.
• Grid Augmentation: Randomly varying the grid size and position prevents the model from learning biases specific to a fixed grid alignment.
• Regularization: An $L_2$ loss penalizes isolated high-amplitude peaks to encourage spatially coherent attributions.

D. Test-Time Refinement

While AHA is designed for single-pass efficiency, it supports an optional Test-Time Refinement step. For a specific input, the explainer head can be fine-tuned for a few gradient steps ($T$) using the same differentiable AUC objectives. This adapts the explanation to the specific sample, further sharpening the attribution map.

3. Key Contributions

1. Metric-Driven Learning: A novel framework that directly optimizes Deletion and Insertion metrics via differentiable surrogates, bridging the gap between perturbation-based causality and learning-based efficiency.
2. Gumbel-Sinkhorn Relaxation: The application of Gumbel-Sinkhorn to relax hard sorting operations, enabling end-to-end training of an explainer model guided by ranking metrics.
3. Dense Pixel-Level Attributions for ViTs: Unlike previous ViT-specific methods that produce coarse patch-level maps, AHA generates dense, pixel-aligned heatmaps that respect object boundaries.
4. Amortized Efficiency: The method achieves inference speeds comparable to propagation-based methods (single forward pass) while offering the fidelity of metric-optimized approaches.
5. Flexible Trade-off: The optional test-time refinement allows users to balance computational cost against explanation quality.

4. Experimental Results

The method was evaluated on ImageNet using ViT-B/16, ViT-B/32, and ConvNeXt Small classifiers.

• Quantitative Performance:
  • AHA (without refinement, $T=0$) consistently outperforms state-of-the-art baselines (including MDA, TIS, ViT-CX, and Grad-CAM) on Deletion AUC and Positive AUC.
  • With test-time refinement ($T=3$), AHA achieves the best overall Faithfulness scores (Ins-Del and Neg-Pos differences), surpassing even per-sample optimization methods like MDA.
  • It significantly reduces the Average Drop Percentage (ADP), indicating better foreground isolation.
• Qualitative Performance:
  • Visualizations show AHA produces sharper, boundary-aligned heatmaps.
  • It correctly highlights specific object parts (e.g., the head of an impala or bottle caps on a can opener) whereas patch-based methods (TIS, ViT-CX) produce blocky, diffuse maps that often include background noise.
• Runtime:
  • At $T=0$, AHA runs in ~0.016s per image, comparable to Grad-CAM and significantly faster than MDA (13s) or TIS (1.1s).
  • Even with refinement ($T=3$), it remains faster than most perturbation-based baselines.

5. Significance

This work represents a significant advancement in Explainable AI (XAI) for modern vision architectures:

• Solves the Resolution Problem: It overcomes the "patch-level" limitation of current ViT explainers, providing human-interpretable, pixel-precise explanations.
• Eliminates Heuristics: By directly optimizing the ground-truth evaluation metrics (Deletion/Insertion) rather than distilling from other methods, it avoids inheriting the biases of surrogate teachers.
• Practical Deployment: The amortized nature of the approach makes it viable for real-world applications where both speed and high-fidelity explanations are required (e.g., autonomous driving or medical imaging).
• Robustness: The optional refinement step allows the system to adapt to out-of-distribution data or specific model quirks (e.g., detecting shortcut learning) without retraining the entire explainer.

In summary, AHA successfully unifies the causal grounding of perturbation methods with the efficiency of learning-based methods, setting a new state-of-the-art for faithful and dense visual attribution in Vision Transformers.