Inherently Faithful Attention Maps for Vision Transformers

2025-06-10

AI Summary

In vision Transformers, atypical backgrounds induce contextual bias, undermining model robustness and interpretability. To address this, we propose a two-stage differentiable binary attention framework: first, it localizes task-relevant image regions; second, it enforces prediction exclusively from these regions via an end-to-end learned binary mask, effectively suppressing background interference. Our core contributions are (i) intrinsic faithfulness of attention maps—attention weights directly and exclusively govern input visibility—and (ii) a joint optimization mechanism that synergistically enhances region discovery and focused prediction. Evaluated on object-centric tasks—including fine-grained recognition and joint localization-classification—the method significantly improves robustness against spurious correlations and out-of-distribution backgrounds, while boosting localization accuracy and prediction consistency.

Technology Category

• Computer Vision: Adversarial Attacks & Robustness
• Machine Learning: Calibration & Uncertainty Quantification
• Natural Language Processing: Safety and Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

We introduce an attention-based method that uses learned binary attention masks to ensure that only attended image regions influence the prediction. Context can strongly affect object perception, sometimes leading to biased representations, particularly when objects appear in out-of-distribution backgrounds. At the same time, many image-level object-centric tasks require identifying relevant regions, often requiring context. To address this conundrum, we propose a two-stage framework: stage 1 processes the full image to discover object parts and identify task-relevant regions, while stage 2 leverages input attention masking to restrict its receptive field to these regions, enabling a focused analysis while filtering out potentially spurious information. Both stages are trained jointly, allowing stage 2 to refine stage 1. Extensive experiments across diverse benchmarks demonstrate that our approach significantly improves robustness against spurious correlations and out-of-distribution backgrounds.

Problem

• Ensures only attended image regions influence predictions
• Addresses biased representations from out-of-distribution backgrounds
• Improves robustness against spurious correlations in object-centric tasks

Innovation

• Learned binary attention masks for faithful attention
• Two-stage framework for object-centric task processing
• Joint training to refine region identification