Uncertainty-Aware Subset Selection for Robust Visual Explainability under Distribution Shifts

This paper addresses the degradation of existing subset-based visual explanation methods under out-of-distribution conditions by introducing a training-free framework that integrates layer-wise uncertainty estimation with submodular optimization to generate robust, diverse, and informative attributions.

The Gist

The Big Picture: The "Overconfident Detective"

Imagine you have a super-smart AI detective (a Deep Vision Model) that is really good at identifying birds. If you show it a picture of a Cardinal, it confidently says, "That's a Cardinal!" and points to the red feathers and the beak. This works great when the bird looks exactly like the ones it studied in school (In-Distribution).

But what happens if you show it a Cardinal wearing a tiny hat, or a Cardinal in a foggy forest, or even a picture of a squirrel that looks a bit like a bird?

The old AI detective gets confused. It might still say "Cardinal," but when it tries to explain why, it points at the hat, the fog, or the squirrel's tail. It becomes brittle (breaks easily) and unreliable. It highlights the wrong parts of the image, making the explanation useless.

The Problem: Existing methods for making AI explain itself work perfectly in the classroom but fail miserably in the real world where conditions change (like weather, lighting, or new types of objects).

---

The Solution: The "Anxious but Smart" Detective

The authors of this paper built a new system to fix this. They call it Uncertainty-Aware Subset Selection.

Here is how it works, using a few analogies:

1. The "Stress Test" (Adaptive Weight Perturbations)

Imagine the AI detective is taking a test.

• Old Method: The detective just looks at the picture once and gives an answer. If the picture is weird, the detective guesses confidently but wrongly.
• New Method: Before giving an answer, the detective gives itself a "stress test." It slightly shakes its own brain (mathematically called perturbing weights) and asks, "If I change my mind just a tiny bit, does my answer change wildly?"
  • If the answer stays the same, the detective is confident.
  • If the answer jumps around like a nervous rabbit, the detective knows, "I am uncertain here."

This is like a pilot checking the instruments before takeoff. If the instruments are wobbling, the pilot knows something is wrong with the plane or the weather.

2. The "Smart Filter" (Submodular Subset Selection)

Usually, when an AI tries to explain an image, it highlights everything that looks important. This is like a student highlighting every single sentence in a textbook because they are scared of missing a test question. It's messy and redundant.

The new method uses a Submodular Filter. Think of this as a strict editor for a news article.

• The editor's job is to pick the top 5 sentences that tell the whole story.
• The editor doesn't just pick the loudest sentences; they pick the ones that are unique and essential.
• If two sentences say the same thing, the editor cuts one out to avoid redundancy.

3. Putting It Together: The "Trustworthy Highlighter"

The new system combines the Stress Test and the Smart Filter.

• Step 1: The AI looks at the image and runs the "stress test" on every tiny patch of the image.
• Step 2: It calculates a "Confidence Score."
  • High Confidence: "I know this is a bird's eye." (Keep it).
  • Low Confidence: "I'm not sure if this is a leaf or a wing because the image is blurry." (Discard it or lower its importance).
• Step 3: The Smart Filter picks the best patches based on this score. It ignores the blurry, confusing parts and focuses only on the clear, stable, and important features.

Why This Matters (The "Real World" Impact)

The paper tested this on two scenarios:

1. The "Related" Shift: Showing the AI a bird from a different continent (North American Birds vs. CUB dataset). The old AI got confused by the different background; the new AI still found the beak and eyes.
2. The "Weird" Shift: Showing the AI a picture of a car when it was trained on birds, or a picture with heavy static noise. The old AI pointed at the noise; the new AI realized, "I don't know what this is," and stopped pointing at random junk.

The Result:

• More Trustworthy: The AI stops lying to you by pointing at the background when it's actually unsure.
• More Efficient: It highlights fewer, but better, parts of the image.
• No Extra Training: The best part? They didn't have to retrain the AI from scratch. They just added this "stress test" and "filter" on top of the existing AI. It's like giving a new pair of glasses to an old detective rather than hiring a new one.

Summary in One Sentence

The authors created a system that makes AI detectives admit when they are unsure and ignore confusing parts of an image, ensuring that their explanations remain accurate and trustworthy even when the world gets messy or changes.

Technical Summary

1. Problem Statement

Deep vision models are increasingly deployed in safety-critical applications (e.g., autonomous driving, medical imaging), where interpretability is essential. Subset selection-based methods are a popular approach for visual explainability, aiming to identify the most influential image regions (patches) that drive a model's prediction.

However, the paper identifies a critical robustness gap:

• In-Distribution (ID) Performance: Existing subset selection methods (e.g., those based on submodular optimization) perform well on data similar to the training distribution.
• Out-of-Distribution (OOD) Failure: Under distribution shifts (e.g., different domains, corruptions, or complementary classes), these methods degrade significantly. They produce redundant, unstable, and uncertainty-sensitive explanations, often highlighting irrelevant background noise or fragmented patches rather than semantically meaningful features.
• Current Limitations: Existing approaches lack principled uncertainty modeling and rely heavily on confidence scores that fail under OOD conditions. They often require additional training or auxiliary models to estimate uncertainty, increasing computational overhead.

2. Methodology

The authors propose a lightweight, plug-and-play framework that integrates submodular subset selection with adaptive, gradient-based uncertainty estimation. The core innovation is a confidence score derived from weight perturbations that guides the selection of image patches without requiring retraining.

A. Uncertainty-Aware Submodular Objective Functions

The problem is framed as maximizing an objective function $F(S)$ over a subset of image regions $S$. The authors propose two specific formulations:

1. $F_{attr}(S)$ for General Visual Attribution: Builds upon the framework by Chen et al. [4], combining scores for Effectiveness, Consistency, and Collaboration. The key modification is replacing the standard confidence term with a novel gradient-based uncertainty score ($s_{conf}$).
2. $F_{obj}(S)$ for Object-Level Interpretation: Designed for foundation models (e.g., visual grounding). It integrates the uncertainty score into the Visual Precision Search (VPS) framework, combining a "Clue Score" (localization support) and "Collaboration Score" (collective importance).

B. Core Component: Gradient-Based Confidence Score ($s_{conf}$)

This is the primary technical contribution. Instead of relying on softmax probabilities, the method estimates uncertainty by simulating epistemic uncertainty via adaptive weight perturbations.

• Adaptive Weight Perturbation: During inference, Gaussian noise is injected into the network weights. Crucially, the noise scale is adaptive:
  • Layer-wise Scaling: Noise is scaled by the standard deviation of weights in each layer ($\sigma_\ell$) to account for different parameter magnitudes.
  • Input-Aware Modulation: The noise magnitude is modulated by the input's distance from the training data centroid (using Mahalanobis distance logic). If an input is OOD, the perturbation strength increases ($u(x) > 1$), amplifying sensitivity to atypical behavior.
• Gradient Sensitivity: The model performs $T$ stochastic forward passes. For each pass, it computes the gradient norms of the output with respect to layer activations. High gradient norms indicate high sensitivity (low confidence).
• Uncertainty Aggregation: Layer-wise gradient norms are aggregated into a descriptor vector. The final uncertainty score is computed as a regularized Mahalanobis distance between this descriptor and the distribution of descriptors from the training set.
  • $s_{conf}(S) = 1 - u_i$, where $u_i$ is the normalized uncertainty score.

C. Greedy Optimization

The framework uses a greedy maximization algorithm to select the subset $S$. Since the objective functions are monotonic and submodular, the greedy approach guarantees a solution within $(1 - 1/e)$ of the optimal solution, ensuring efficiency and high-quality selection.

3. Key Contributions

1. Empirical Evidence of Robustness Gap: The authors demonstrate that existing subset selection methods suffer severe performance degradation (up to 40% drop in metrics) under various OOD shifts (related, complementary, and transformed distributions).
2. Novel Framework: A lightweight method that combines submodular optimization with adaptive uncertainty estimation. It requires no additional training or auxiliary models, operating solely on a fine-tuned backbone.
3. Adaptive Uncertainty Estimation: A novel mechanism that uses input-aware, layer-wise weight perturbations to generate a robust confidence score, effectively distinguishing between ID and OOD inputs.
4. Dual-Task Applicability: The framework is validated on both fine-grained classification (identifying bird species) and object-level interpretation (visual grounding), showing improvements in both ID and OOD settings.

4. Experimental Results

The authors evaluated their method on two experimental settings using CUB-200-2011 (Birds) and COCO (Object Detection) as ID datasets, with corresponding OOD datasets (NABirds, CIFAR-100, iNaturalist, and transformed versions).

• Metrics: Evaluated using Insertion AUC (higher is better) and Deletion AUC (lower is better).
• Performance on Classification (CUB):
  • The proposed method consistently improved Insertion AUC over baselines (HSIC-Attribution + SMDL).
  • ID Improvement: +1.7% to +5.0% increase in Insertion AUC.
  • OOD Improvement: Significant gains under shifts. For example, on NABirds (Related OOD), Insertion AUC increased by +6.2% to +13.7%. On CIFAR-100 (Complementary OOD), gains reached +10.1% to +12.3%.
• Performance on Object Detection (COCO/GroundingDINO):
  • The method more than doubled the Insertion AUC on related OOD datasets (CIFAR-100) and transformed OOD datasets (COCO Transformed), with increases of +44.5% to +108.9% compared to the VPS baseline.
  • While Deletion AUC sometimes increased (indicating a trade-off), the method achieved significant stability improvements in specific transformed OOD scenarios.
• Qualitative Analysis: Visual comparisons showed that the proposed method selects compact, semantically coherent regions (e.g., focusing on the bird's beak/eyes) even under OOD conditions, whereas baselines selected fragmented or background regions.

5. Significance and Conclusion

This work addresses a critical limitation in current AI interpretability: the lack of reliability under distribution shifts. By integrating uncertainty-aware optimization directly into the subset selection process, the authors provide a method that:

• Enhances Trust: Produces explanations that remain faithful to the model's decision process even when inputs deviate from training data.
• Improves Generalization: Bridges the gap between ID and OOD performance without the computational cost of retraining or ensemble models.
• Enables Real-World Deployment: Offers a practical, lightweight solution for safety-critical applications where understanding model behavior under unexpected conditions is paramount.

The paper concludes that uncertainty-driven subset selection is essential for transparent and trustworthy AI, paving the way for more robust visual explainability in real-world scenarios.