Learn to Rank: Visual Attribution by Learning Importance Ranking

Conference: ECCV 2026
Paper: ECCV Official
Code: https://github.com/dschinagl/AHA
Area: Interpretability
Keywords: Visual Attribution, Amortized Explanation, Permutation Learning, Gumbel-Sinkhorn, Vision Transformer

TL;DR

To overcome the three-way trade-off among computational efficiency, causal faithfulness, and fine-grained spatial resolution in visual attribution, this paper proposes Amortized Hybrid Attribution (AHA), which relaxes discrete sorting into differentiable permutation learning via Gumbel-Sinkhorn to directly optimize Deletion and Insertion metrics end-to-end, producing dense pixel-level heatmaps in a single forward pass.

Background & Motivation

The growing complexity of deep vision models makes their internal decision-making opaque, which hinders trust and accountability in safety-critical domains such as healthcare and autonomous driving. Visual attribution methods aim to produce saliency maps that highlight the input features most influential to a model's prediction. However, existing post-hoc attribution paradigms suffer from an inherent three-way trade-off: propagation-based approaches (such as Grad-CAM, Integrated Gradients, and LeGrad) are fast and faithfully mirror internal computation graphs, but they are architecture-specific and heavily biased toward low-level edges and textures; perturbation-based methods (such as RISE, LIME, SHAP, and TIS) offer rigorous causal grounding by evaluating output shifts under input occlusion, yet they require hundreds of model evaluations per image and, for Vision Transformers, are constrained to coarse patch-level tokens that yield blocky maps; meanwhile, existing learning-based amortized explainers largely depend on distilling heuristic teacher methods or making restrictive binary signal-to-noise assumptions.

The core tension stems from the fact that Deletion and Insertion curves—the widely accepted gold standards for evaluating causal attribution quality—inherently require descending sorting of feature importances and discrete Top-k thresholding. Because discrete rank permutations have zero or undefined gradients, prior methods could not optimize an explainer network directly against these metric objectives via standard backpropagation; earlier metric-aligned efforts like MDA had to settle for prohibitive per-sample iterative search.

This paper's angle of attack is that both Deletion and Insertion metrics depend exclusively on the relative ordering of attribution scores rather than their exact magnitudes. Hence, metric optimization can be reformulated as a continuous permutation learning task. Core idea: by relaxing discrete sorting and Top-k selection into smooth doubly-stochastic permutation matrices via Gumbel-Sinkhorn, the Deletion and Insertion integrals become end-to-end differentiable, enabling an amortized explainer to generate dense pixel-level attributions in a single forward pass with optional test-time gradient refinement.

Method

Overall Architecture

AHA trains an amortized attribution model \(\Phi_{f,\theta}\) for a frozen target classifier \(f\). Given an input image \(I\) and a target class \(t\), the explainer predicts a dense attribution map \(A\) matching the input's spatial dimensions. During training, AHA circumvents discrete sorting barriers by casting the ranking of attribution values as an optimal transport assignment problem, using the Gumbel-Sinkhorn algorithm to compute continuous soft permutation matrices. These matrices yield soft Top-k selection masks that generate differentiable sequences of perturbed images, allowing the classifier's output probabilities across the perturbation curve (AUC) to serve directly as an end-to-end training loss. During inference, the explainer produces high-fidelity, boundary-aligned attribution maps in a single feedforward pass, with the option to run a few gradient steps (e.g., \(T=3\)) of test-time refinement for extreme fidelity.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Input Image I & Target Class t"] --> B["Explainer Network Predicts Attribution A"]
    B --> C["Region Partitioning & Grid Augmentation<br/>G×G Mean Pooling"]
    C --> D["Gumbel-Sinkhorn Permutation Sorting<br/>Compute Soft Permutation Matrix"]
    D --> E["Soft Top-k Mask Construction<br/>Cumulative Sum for Smooth Perturbations"]
    E --> F["Differentiable Deletion & Insertion Loss<br/>End-to-End Backprop to Explainer"]
    F --> G["Inference Output<br/>Single Forward Pass or Test-Time Refinement"]

Key Designs

1. Gumbel-Sinkhorn Permutation Sorting: Formulating Discrete Sorting as Differentiable Optimal Transport

Traditional Deletion and Insertion metrics sort vectorized attribution values \(a = \text{vec}(A) \in \mathbb{R}^N\) in descending order, introducing step-function discontinuities where gradients vanish. AHA recognizes that the metrics depend solely on relative rank order. It defines an ideal sorted reference vector \(p = \frac{1}{N}[N, N-1, \dots, 1]^\top\) and sets up an optimal transport cost matrix based on negative squared distances \(L_{i,j} = -(a_i - p_j)^2\). To facilitate stochastic exploration during training, Gumbel noise is injected to obtain \(\tilde{L} = (L + G)/\tau\), where \(\tau > 0\) denotes temperature. Applying the Sinkhorn-Knopp algorithm performs alternating row and column normalizations to output a doubly-stochastic matrix \(P^{\text{soft}}_{\pi_A} \in [0, 1]^{N \times N}\). This provides a smooth, continuous relaxation of the permutation matrix, allowing perturbation-driven downstream gradients to flow smoothly back into the explainer parameters.

2. Soft Top-k Mask Construction: Generating Continuous Perturbation Trajectories

Once the soft permutation matrix is obtained, the discrete Top-k binary selection mask is generalized by summing the first \(k\) rows of \(P^{\text{soft}}_{\pi_A}\):

\[m^{(k)}_{\text{soft}, A} = \sum_{i=1}^k P^{\text{soft}}_{\pi_A}[i, :]^\top \in [0, 1]^N\]

These monotonically expanding continuous masks smoothly blend the input image \(I\) with a reference image \(I_0\) (e.g., black, mean gray, or blurred) to generate soft perturbation sequences \(I_{\text{del}}^{\text{soft}, (k)}\) and \(I_{\text{ins}}^{\text{soft}, (k)}\). Crucially, Sinkhorn assignment links score magnitude separation to assignment certainty: diffuse attribution maps create ambiguous, low-contrast masks that blur foreground and background, yielding flat class probability trajectories; only when the explainer sharply segregates true causal features from the background do the soft masks produce steep probability drops under Deletion and rapid gains under Insertion, directly maximizing the AUC loss gradient.

3. Region-Based Permutation & Grid Augmentation: Overcoming Computational Bottlenecks and Shortcut Artifacts

Performing an \(N \times N\) Sinkhorn operation over raw high-resolution pixels (\(224 \times 224 = 50{,}176\)) is computationally prohibitive and prone to overfitting on high-frequency adversarial shortcuts of the target model. AHA addresses this by average-pooling attribution maps into \(K = G \times G\) disjoint spatial regions before running soft sorting, then bilinearly upsampling the soft masks back to full resolution. To prevent the network from memorizing static spatial partitions, the grid size \(G\) is randomly sampled from \(\{7, \dots, 28\}\) per iteration with random spatial offsets. Furthermore, the continuous AUC integrals are efficiently approximated by uniformly subsampling \(S \ll K\) perturbation steps (e.g., \(S=16\)), restricting forward and backward passes to a fixed compute budget.

4. Test-Time Refinement: Hybrid Inference Balancing High Throughput and Extreme Fidelity

To serve both latency-critical deployment and deep diagnostic inspection, AHA provides a flexible test-time adaptation mechanism. By default (\(T=0\)), the frozen explainer runs in a single forward pass taking only 16ms. In settings prioritizing explanation fidelity over throughput, the explainer head can be briefly fine-tuned on the individual test sample for \(T\) gradient steps (typically \(T=3\) to \(5\)) using the exact same differentiable metric loss. To ensure numerical stability and determinism during test-time adaptation, Gumbel noise is disabled and the temperature is locked to its final training value (\(\tau = \tau_{\text{final}}\)), rapidly suppressing spurious background noise and sharpening object boundaries.

Loss & Training

The overall training loss balances the differentiable Deletion AUC, differentiable Insertion AUC, and a spatial smoothness regularization term:

\[\mathcal{L}(\theta) = \lambda_{\text{del}} \cdot \text{del}_{\text{AUC}}^{\text{soft}}(A, I, t) - \lambda_{\text{ins}} \cdot \text{ins}_{\text{AUC}}^{\text{soft}}(A, I, t) + \lambda_{\text{reg}} \cdot \mathcal{L}_{\text{reg}}(A)\]

where \(\mathcal{L}_{\text{reg}}\) penalizes \(L_2\) deviations between attribution map \(A\) and its average-filtered version to suppress isolated high-frequency spikes. Experiments set \(\lambda_{\text{del}} = 1.0\), \(\lambda_{\text{ins}} = 1.0\), and \(\lambda_{\text{reg}} = 2.5 \times 10^{-3}\). The explainer employs a frozen DINOv3 ViT-L/16 backbone combined with a trainable class-conditioned DPT-style decoder. Training requires only 1 epoch on the ImageNet-1K training set using AdamW (learning rate \(3 \times 10^{-4}\), weight decay \(10^{-3}\)) with a one-cycle cosine learning rate schedule.

Key Experimental Results

Main Results

Evaluations are conducted on the complete ImageNet-1K validation set (50,000 images) with a frozen ViT-B/16 classifier. All perturbation metrics are averaged across three baseline reference modes \(I_0\) (black, mean color, Gaussian blur) for the predicted class.

Method	Deletion↓	Insertion↑	Ins. - Del.↑	Positive↓	Negative↑	Neg. - Pos.↑	ADP↓	PIC↑	Runtime
Grad-CAM	0.2858	0.4982	0.2124	0.3346	0.5788	0.2442	84.33	4.36	0.014s
Integrated Gradients	0.1857	0.5544	0.3686	0.2248	0.6347	0.4099	22.93	33.44	0.251s
Trans-Att.	0.1654	0.5833	0.4178	0.1961	0.6654	0.4633	28.46	30.85	0.254s
Bi-Att.	0.1600	0.5958	0.4358	0.2022	0.6789	0.4828	25.70	28.26	0.254s
TIS	0.1420	0.6381	0.4961	0.1731	0.7222	0.5492	11.25	45.57	1.167s
ViT-CX	0.1750	0.5714	0.3964	0.2149	0.6478	0.4328	12.55	43.50	0.860s
MDA	0.1556	0.6112	0.4556	0.1933	0.6983	0.5050	48.67	13.27	13.645s
LeGrad	0.1666	0.5666	0.4000	0.2036	0.6461	0.4425	20.39	34.34	0.006s
AHA (\(T=0\), Ours)	0.1381	0.6112	0.4731	0.1714	0.6947	0.5233	15.87	36.82	0.016s
AHA (\(T=3\), Ours)	0.1215	0.6485	0.5271	0.1491	0.7378	0.5887	12.47	42.27	0.547s

Ablation Study

Ablation experiments evaluate the sampling sensitivity of training perturbation steps \(S\) and the progressive impact of test-time refinement steps \(T\) on Insertion AUC and Deletion AUC (ViT-B/16, ImageNet validation set):

Config / Variable	Value	Insertion AUC (Pred)↑	Deletion AUC (Pred)↓	Note
Perturbation steps \(S\)	\(S = 4\)	~0.535	-	Under-sampled Riemann sum hinders gradient alignment
Perturbation steps \(S\)	\(S = 8\)	~0.580	-	Rapid performance surge
Perturbation steps \(S\) (Baseline)	\(S = 16\)	0.6112	0.1381	Optimal accuracy-efficiency saturation point
Perturbation steps \(S\)	\(S = 32\)	~0.613	-	Double the compute with marginal gain
Refinement steps \(T\)	\(T = 0\) (single forward)	0.6112	0.1381	Pure amortized inference at 0.016s
Refinement steps \(T\)	\(T = 1\)	~0.628	~0.130	Initial gradient step yields largest relative leap
Refinement steps \(T\)	\(T = 2\)	~0.640	~0.125	Further clears ambiguous background activations
Refinement steps \(T\)	\(T = 3\)	0.6485	0.1215	Sweet spot between latency (0.547s) and fidelity
Refinement steps \(T\)	\(T = 5\)	~0.655	~0.118	Diminishing returns with razor-sharp boundaries

Key Findings

• Zero-Shot Amortization Outperforms Iterative Baselines: In the feedforward regime (\(T=0\)), AHA achieves a Deletion score of 0.1381, substantially outperforming MDA (0.1556)—despite MDA requiring 13.6 seconds of per-sample iterative search. With \(T=3\) refinement steps, AHA establishes new state-of-the-art faithfulness (Ins. - Del. reaches 0.5271).
• Pixel-Level Boundary Alignment vs. Patch Grid Artifacts: While native transformer explainers (TIS, ViT-CX) produce coarse patch blocks that spill onto the background, AHA generates dense, object-aligned heatmaps that cleanly trace morphological contours, reflected in lower background drop percentages (ADP).
• Clear Saturation Thresholds: Training curves plateau around \(S=16\) perturbation steps, and test-time refinement captures over 80% of its total potential gain within \(T=3\) steps, maintaining a massive speed advantage over iterative perturbation schemes.

Highlights & Insights

• Bridging Non-Differentiable Metrics via Permutation Learning: By recognizing that ranking-based metrics depend only on relative orders, the authors leverage Gumbel-Sinkhorn to make Deletion and Insertion directly optimizable as neural network losses.
• Robustness Against Adversarial Cheating: Combining region-level pooling, multi-scale grid randomization, and spatial regularizers effectively prevents the explainer from exploiting high-frequency classification quirks of the target model.
• Unified Amortized and Optimization Paradigm: AHA successfully unifies single-pass feedforward efficiency with optional sample-level test-time backpropagation, providing a practical model-agnostic blueprint for explainable computer vision.

Limitations & Future Work

• Target-Model Retraining Overhead: The explainer is trained for one specific target model; explaining a different classifier architecture requires retraining the explainer head.
• Reference Image Sensitivity: Deletion and Insertion curves remain sensitive to baseline reference choices (black, mean, or blurred frames), which can introduce out-of-distribution artifacts into the evaluation.
• Reliance on In-Distribution Training Data: Training the explainer requires access to representative image data matching the target model's training distribution, limiting its direct out-of-the-box use on closed-source, proprietary models without test-time adaptation.

Related Work & Insights

• vs MDA (Metric-Driven Attribution): MDA also seeks to optimize Deletion/Insertion metrics, but relies on per-sample patch mask optimization that takes ~13.6 seconds per image; AHA shifts the optimization offline, cutting inference latency to 0.016 seconds (>800x faster) while offering smooth pixel-level heatmaps.
• vs TIS / ViT-CX: TIS and ViT-CX operate exclusively in the token domain, outputting coarse \(14 \times 14\) patch grids; AHA utilizes a high-capacity pixel decoder that resolves fine anatomical details without grid blockiness.
• vs LeGrad / Transformer-Attribution: Attention propagation techniques are millisecond-fast but often contaminated by superficial edge gradients and gradient saturation; AHA preserves the strict causal fidelity of perturbation methods while matching propagation speed.

Rating

• Novelty: ⭐⭐⭐⭐⭐ Elegant use of Gumbel-Sinkhorn permutation learning to resolve the long-standing non-differentiability of visual attribution metrics.
• Experimental Thoroughness: ⭐⭐⭐⭐⭐ Comprehensive evaluation across 50,000 ImageNet validation images, two ViT architectures (ViT-B/16, ViT-B/32), three reference image baselines, and runtime benchmarks.
• Writing Quality: ⭐⭐⭐⭐⭐ Clear mathematical derivations, well-structured exposition, and concise problem formulation.
• Value: ⭐⭐⭐⭐⭐ A highly practical, scalable framework for faithful, high-resolution visual explanation in modern vision models.

Related Papers

• [ICLR 2026] Escaping Low-Rank Traps: Interpretable Visual Concept Learning via Implicit Vector Quantization
• [ICLR 2026] Flow-Disentangled Feature Importance
• [ICLR 2026] Concept-TRAK: Understanding how diffusion models learn concepts through concept attribution
• [ICLR 2026] Learning to Weight Parameters for Training Data Attribution
• [ICLR 2026] TreeGrad-Ranker: Feature Ranking via O(L)-Time Gradients for Decision Trees