# Two-stage Vision Transformers and Hard Masking offer Robust Object Representations

arXiv: https://arxiv.org/abs/2506.08915 · HTML: https://mlsanity.com/paper/2506.08915/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 56 of its 72 references that are in the map.
- First posted on arXiv 10 Jun 2025.

## Builds on (most cited first, 30 of 56)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [SGDR: Stochastic Gradient Descent with Warm Restarts](https://mlsanity.com/paper/1608.03983/) (2016)
- [Show, Attend and Tell: Neural Image Caption Generation with Visual Attention](https://mlsanity.com/paper/1502.03044/) (2015)
- [Energy-based Out-of-distribution Detection](https://mlsanity.com/paper/2010.03759/) (2020)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Vision Transformers Need Registers](https://mlsanity.com/paper/2309.16588/) (2023)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [Recognition in Terra Incognita](https://mlsanity.com/paper/1807.04975/) (2018)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [DeiT III: Revenge of the ViT](https://mlsanity.com/paper/2204.07118/) (2022)
- [ResNet strikes back: An improved training procedure in timm](https://mlsanity.com/paper/2110.00476/) (2021)
- [Recurrent Models of Visual Attention](https://mlsanity.com/paper/1406.6247/) (2014)
- [Noise or Signal: The Role of Image Backgrounds in Object Recognition](https://mlsanity.com/paper/2006.09994/) (2020)
- [Right for the Right Reasons: Training Differentiable Models by Constraining their Explanations](https://mlsanity.com/paper/1703.03717/) (2017)
- [One weird trick for parallelizing convolutional neural networks](https://mlsanity.com/paper/1404.5997/) (2014)
- [Learning to Explain: An Information-Theoretic Perspective on Model Interpretation](https://mlsanity.com/paper/1802.07814/) (2018)
- [Attention Mechanisms in Computer Vision: A Survey](https://mlsanity.com/paper/2111.07624/) (2021)
- [No Subclass Left Behind: Fine-Grained Robustness in Coarse-Grained Classification Problems](https://mlsanity.com/paper/2011.12945/) (2020)
- [Compositional Attention Networks for Machine Reasoning](https://mlsanity.com/paper/1803.03067/) (2018)
- [Chasing Sparsity in Vision Transformers: An End-to-End Exploration](https://mlsanity.com/paper/2106.04533/) (2021)
- [Exploring scalable medical image encoders beyond text supervision](https://mlsanity.com/paper/2401.10815/) (2024)
- [Revisiting Weakly Supervised Pre-Training of Visual Perception Models](https://mlsanity.com/paper/2201.08371/) (2022)
- [Scalable Pre-training of Large Autoregressive Image Models](https://mlsanity.com/paper/2401.08541/) (2024)
- [SCOPS: Self-Supervised Co-Part Segmentation](https://mlsanity.com/paper/1905.01298/) (2019)
- [A Whac-A-Mole Dilemma: Shortcuts Come in Multiples Where Mitigating One Amplifies Others](https://mlsanity.com/paper/2212.04825/) (2022)
- [MetaShift: A Dataset of Datasets for Evaluating Contextual Distribution Shifts and Training Conflicts](https://mlsanity.com/paper/2202.06523/) (2022)
- [Discover and Cure: Concept-aware Mitigation of Spurious Correlation](https://mlsanity.com/paper/2305.00650/) (2023)

## Built on by (most cited first, 1 of 1)

- [A2: Smaller Self-Supervised ViTs Localize Better than Larger Ones](https://mlsanity.com/paper/2606.03148/) (2026)

## Abstract (arXiv)

Context can strongly affect object representations, sometimes leading to undesired biases, particularly when objects appear in out-of-distribution backgrounds at inference. At the same time, many object-centric tasks require to leverage the context for identifying the relevant image regions. We posit that this conundrum, in which context is simultaneously needed and a potential nuisance, can be addressed by an attention-based approach that uses learned binary attention masks to ensure that only attended image regions influence the prediction. To test this hypothesis, we evaluate a two-stage framework: stage 1 processes the full image to discover object parts and identify task-relevant regions, for which context cues are likely to be needed, 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. The explicit nature of the semantic masks also makes the model's reasoning auditable, enabling powerful test-time interventions to further enhance robustness. Extensive experiments across diverse benchmarks demonstrate that this approach significantly improves robustness against spurious correlations and out-of-distribution backgrounds. Code: https://github.com/ananthu-aniraj/ifam