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

ICPR 2026 Oral Presentation

Ananthu Aniraj¹, Cassio F. Dantas^1,2, Dino Ienco^1,2, Diego Marcos¹

¹ Inria, EVERGREEN, University of Montpellier, 34090 Montpellier, France

² INRAE, UMR TETIS, University of Montpellier, 34090 Montpellier, France

Springer (Paper) arXiv Code 🤗 HF Checkpoints HF Paper

Early Hard Masking vs. Late Feature Masking

Comparison of Late Feature Masking versus Early Hard Masking in ViTs — **Figure 1: Receptive fields and information leakage.** (Top) Previous attention-based approaches apply masks to high-level feature tensors (Late Masking). Because deep layers possess large, global receptive fields, background information has already leaked into all tokens before masking. (Bottom) **iFAM (Early Hard Masking)** constrains self-attention at the input level via attention masking (\(M_{ij} = -\infty\)), mathematically guaranteeing that only attended foreground regions can influence the representation and downstream prediction.

Abstract

Context can strongly affect object representations, sometimes leading to undesired biases, particularly when objects appear in out-of-distribution (OOD) backgrounds at inference. At the same time, many object-centric tasks require to leverage 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.

Methodology & Architecture

iFAM Architecture and Test-Time Interventions — **Figure 2: Overview of the iFAM Framework.** (Left) **Two-Stage Architecture:** Stage 1 processes the whole image to identify task-relevant regions via part prototypes. The parts are discretized and merged into a binary input mask \(\mathbf{s}\). Stage 2 classifies the image using input-level attention masking, ensuring zero background information enters the prediction. (Right) **Auditable Test-Time Interventions:** Leveraging semantically consistent part attribution, users can perform training-free interventions: dropping parts that capture spurious correlations (Part-Removal) or pruning low-confidence tokens distant from learned prototypes (Token-Removal).

1. Early vs. Late Masking

Conventional attention-based methods (Late Masking) learn a selector \(h_{\theta_s}(\cdot)\) and a feature extractor \(h_{\theta_p}(\cdot)\), computing the output as:

\[ \mathbf{y} = g_{\phi}\left(h_{\theta_p}(\mathbf{x}) \odot h_{\theta_s}(\mathbf{x})\right) \]

Because \(h_{\theta_p}(\mathbf{x})\) is a deep feature tensor whose tokens have already exchanged information globally, spurious background cues leak into the foreground features.

In contrast, iFAM (Early Masking) computes a binary token selection mask \(\mathbf{s} \in \{0, 1\}^N\) from Stage 1, and enforces hard input masking inside Stage 2 by directly modulating the Vision Transformer's multi-head self-attention:

\[ \text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax} \left( \frac{\mathbf{QK}^\top}{\sqrt{D}} + \mathbf{M} \right) \mathbf{V} \]

where the elements of the attention mask \(\mathbf{M} \in \mathbb{R}^{N \times N}\) are defined as:

\[ M_{ij} = \begin{cases} -\infty, & \text{if } s_i = 0 \text{ or } s_j = 0 \\ 0, & \text{otherwise} \end{cases} \]

Setting \(M_{ij} = -\infty\) forces attention weights to and from masked-out tokens to be exactly zero after the softmax, guaranteeing that background tokens cannot contaminate the final image representation.

2. Stage 1: Discovering Relevant Regions

Stage 1 discovers \(K\) distinct foreground object parts plus 1 background channel using prototype-driven part discovery guided by classification and Total Variation priors. The resulting foreground part maps are discretized using a Gumbel-Softmax straight-through gradient estimator: hard binary assignments are used in the forward pass while continuous soft probabilities backpropagate gradients, allowing Stage 2 to end-to-end refine Stage 1.

3. Auditable Robustness via Test-Time Interventions

Because iFAM’s part representations are semantically consistent and auditable, users can perform training-free interventions at inference:

Part-Removal Intervention: If a learned part inadvertently latches onto a spurious object cue (e.g., hospital tubes in medical scans, or water in bird images), that entire part channel can be excluded from the Stage 2 input mask—either via manual visual inspection or automated Leave-One-Out (LOO) validation.
Token-Removal Intervention: When encountering out-of-distribution inputs that produce false-positive activations, tokens whose feature distance to the prototype exceeds a calibrated training percentile threshold \(\tau_k^q\) are filtered out and assigned to background.

Jump to Full Interventions Section, Part Maps & Empirical Results ↓

Quantitative Benchmarks

iFAM is extensively benchmarked across four standard out-of-distribution and debiasing datasets: MetaShift (Cat vs. Dog with spurious context), Waterbirds (waterbirds/landbirds with conflicting backgrounds), ImageNet-9 Backgrounds Challenge, and SIIM-ACR Pneumothorax chest radiographs.

Table 1(a): Results on MetaShift and Waterbird

Category	Method	Arch.	MetaShift (%)			Waterbird (%)
Category	Method	Arch.	\(K\)	Avg. Acc (AA)	Worst-Group (WGA) \(\uparrow\)	\(K\)	Avg. Acc (AA)	Worst-Group (WGA) \(\uparrow\)
Upper Bounds	Early mask (GT)^†	ViT-B	--	--	--	1	99.2	97.2
Supervised	ERM	ViT-B	--	75.8	62.5	--	95.0	80.7
Self-Supervised	DinoV2 (Frozen)	ViT-B	--	83.2	72.6	--	95.9	88.5
Self-Supervised	DinoV2 (Fine-tuned)	ViT-B	--	84.7	76.8	--	98.6	95.8
Specialized De-biasing	GroupDRO	R-50	--	73.6	66.0	--	91.8	90.6
Specialized De-biasing	DISC	R-50	--	75.5	73.5	--	93.8	88.7
Late Masking	DinoV2 + AIM	ViT-B	--	88.9	83.0	--	97.6	93.9
Late Masking	PDiscoFormer	ViT-B	4	83.2	75.5	8	94.2	84.3
Early Masking	iFAM (Ours)	ViT-B	4	88.7	88.6	8	99.0	97.0

^† Uses ground truth bounding boxes/segmentations as upper bound. Shaded columns denote out-of-distribution robustness metrics (WGA).

Table 1(b): Results on ImageNet-9 (IN-9) Backgrounds Challenge

Category	Method	Arch.	IN-1K	IN-9O	Mixed-Same (MS)	Mixed-Rand (MR)	BG-GAP \(\downarrow\)
Supervised	ERM	R-50	81.2	96.4	90.0	84.6	5.4
Supervised	ERM	ViT-B	83.8	97.9	92.4	87.9	4.6
Self-Supervised	DinoV2	ViT-B	84.6	98.1	93.1	87.1	6.0
Self-Supervised	DinoV2	ViT-L	86.7	98.3	95.5	90.2	5.3
Specialized De-biasing	LLE	R-50	76.3	95.5	88.3	83.4	4.9
Late Masking	PDiscoFormer (\(K=1\))	ViT-B	83.3	98.4	93.9	88.6	5.3
Early Masking	iFAM (Ours, \(K=1\))	ViT-B	84.3	97.5	93.5	91.1	2.4

\(\text{BG-GAP} = \text{MS} - \text{MR}\) measures reliance on spurious background cues (lower is better). iFAM slashes BG-GAP to just 2.4%.

Table 2: Generalization from CUB to Waterbird200 (OOD)

Category	Method	\(K\)	CUB (In-distribution)	Waterbird200 (OOD Backgrounds) \(\uparrow\)
Upper Bounds	Early mask (GT Seg)^†	1	91.4	88.8
Upper Bounds	Late mask (GT Seg)^†	1	90.7	74.8
Self-Supervised	ViT-B DinoV2 (Frozen)	--	89.2	76.6
Self-Supervised	ViT-B DinoV2 (Fine-tuned)	--	91.6	68.4
Late Masking	PDiscoFormer	4	89.1	76.0
	PDiscoFormer	8	88.8	76.8
	PDiscoFormer	16	88.7	75.8
Early Masking	iFAM (Ours)	4	90.1	86.1
	iFAM (Ours)	8	90.4	86.2
	iFAM (Ours)	16	90.6	86.2

Training-Free Human-AI Collaboration

Auditable Robustness via Test-Time Interventions

Unlike black-box models, iFAM's explicit part decomposition and Early Hard Masking allow practitioners to audit and intervene at inference—mitigating spurious shortcuts with zero model retraining.

When a conventional deep neural network or late-masking model learns a spurious shortcut (e.g., predicting cat from an indoor sofa, or detecting pneumothorax based on chest drainage tubes), remediating the bias typically demands collecting expensive counterfactual data or retraining from scratch.

In contrast, iFAM provides inherent auditability: Stage 1 discovers semantically consistent, discrete part masks, while Stage 2 guarantees that unmasked regions are physically excluded from the Vision Transformer's receptive field. This unlocks two powerful, training-free intervention strategies that can be executed manually by a human expert or programmatically via automated validation:

1. Part-Removal Intervention (Leave-One-Out / Human-in-the-Loop)

The Mechanism: When the number of parts \(K\) is set high, weakly supervised part discovery may allocate an entire part to a strongly correlated background feature or medical artifact. Because learned parts are shared across classes and semantically consistent across all images, a user can visually inspect a handful of images to spot the spurious part, or an automated Leave-One-Out (LOO) validation search can programmatically exclude parts whose removal improves worst-group performance:

\[ \mathbf{s}_{\text{intervened}} = \bigvee_{j \in \{1,\dots,K\} \setminus \{k\}} \mathbf{s}_j \]

Because Stage 2 enforces early hard masking at the attention layer, the excluded part is strictly eliminated from the model's receptive field, preventing the classifier from exploiting the confounder.

MetaShift (\(K=8\)) LOO Part-Removal

Part assignment maps (top/bottom rows: pure, middle row: spurious).

Active Parts	WGA (%)	Impact
All Parts	78.8	Baseline
-- Part 1	64.7	-14.1
-- Part 2	75.8	-3.0
-- Part 3	75.8	-3.0
-- Part 4	78.8	0.0
-- Part 5	75.5	-3.3
-- Part 6 (Brown)	81.7	+2.9 (Spurious Indoor Cue)
-- Part 7	77.1	-1.7
-- Part 8	69.9	-8.9

Finding: Part 6 consistently tracked indoor furniture correlated with cats. Dropping this single part boosts Worst-Group Accuracy by +2.9%.

SIIM-ACR Chest X-Ray (\(K=8\)) LOO Part-Removal

Part assignment maps + ground truth pneumothorax occurrence heatmap.

Active Parts	WG AUC (%)	Impact
All Parts	65.9	Baseline
-- Part 1	65.0	-0.9
-- Part 2	59.9	-6.0
-- Part 3	63.7	-2.2
-- Part 4 (Red)	67.3	+1.4 (Drainage Tube Artifact)
-- Part 5	65.2	-0.7
-- Part 6	65.6	-0.3
-- Part 7	66.7	+0.8
-- Part 8	65.5	-0.4

Finding: Part 4 covers the central chest where chest drainage tubes (a confounding artifact) are placed, far from peripheral pneumothorax lesions. Removing it boosts Worst-Group AUC by +1.4%.

2. Token-Removal Intervention (Prototype Confidence Thresholding)

The Mechanism: On out-of-distribution (OOD) test inputs with unfamiliar backgrounds or noisy artifacts, Stage 1 can produce false-positive token activations. We prune unconfident tokens by comparing their feature distance to the learned part prototypes:

\[ s_i = \begin{cases} 0, & \text{if } d(\mathbf{h}_i, \mathbf{p}_k) > \tau_k^q \\ s_i, & \text{otherwise} \end{cases} \]

The threshold \(\tau_k^q\) is calibrated on the training set using a high percentile \(q\) (e.g., \(q = 97\%\) or \(99\%\)), representing the proportion of training tokens assigned to part \(k\) closer than \(\tau_k^q\) to prototype \(\mathbf{p}_k\). Tokens exceeding this distance at test time are reassigned to background.

Table 3: Results of Token-Removal Intervention Across Benchmarks

Method	MetaShift (\(K=8\))		Waterbird (\(K=16\))		SIIM-ACR (\(K=8\))		Waterbird200 (OOD)
Method	AA (%)	WGA (%)	AA (%)	WGA (%)	Avg AUC	WG AUC	\(K=4\)	\(K=8\)	\(K=16\)
iFAM (Baseline)	84.5	78.8	98.8	97.0	92.1	65.9	86.1	86.2	86.2
Token-Removal (\(q=97\%\))	84.7 (+0.2)	79.1 (+0.3)	98.7 (-0.1)	96.6 (-0.4)	92.0 (-0.1)	66.0 (+0.1)	86.8 (+0.7)	86.7 (+0.5)	87.3 (+1.1)
Token-Removal (\(q=99\%\))	84.7 (+0.2)	80.1 (+1.3)	98.8 (0.0)	97.4 (+0.4)	92.2 (+0.1)	66.4 (+0.5)	86.6 (+0.5)	86.9 (+0.7)	86.9 (+0.7)

Token-removal filtering consistently improves OOD generalization across all benchmarks while preserving in-distribution accuracy.

3. Combining Interventions: Synergy & The Early Masking Guarantee

Because Part-Removal operates at the macro semantic level (pruning entire confounded part channels) and Token-Removal operates at the micro confidence level (pruning noisy boundary tokens across all remaining parts), their effects are highly complementary.

Crucially, Table 4 demonstrates that Early Masking is an architectural prerequisite for interventions to succeed:

Table 4: Combined Interventions on MetaShift and SIIM-ACR (\(K=8\))

Model	Intervention Strategy	MetaShift		SIIM-ACR
Model	Intervention Strategy	AA (%)	WGA (%)	Avg AUC (%)	Worst-Group AUC (%)
PDiscoFormer (Late Masking)	No Intervention	83.2	75.5	92.6	48.1
	+ LOO Part-Removal	85.2	76.8 (+1.3)	92.6	48.1 (0.0)
	+ LOO + Token-Removal (\(q=99\%\))	85.4	76.8 (+1.3)	92.6	48.2 (+0.1)
iFAM (Ours) (Early Hard Masking)	No Intervention	84.5	78.8	92.1	65.9
	+ LOO Part-Removal	84.7	81.7 (+2.9)	90.6	67.3 (+1.4)
	+ LOO + Token-Removal (\(q=99\%\))	84.8	83.0 (+4.2)	91.1	69.0 (+3.1)

Why Late-Masking Models Cannot Intervene: In late-masking architectures (like PDiscoFormer), attention layers operate on the entire image before feature pooling occurs. By the time a part or token is masked out at the end, spurious background information has already leaked across self-attention layers into the foreground representations. Only iFAM's Early Hard Masking strictly restricts the receptive field from the very first layer, ensuring that intervened regions are physically undetectable by the classifier.

Qualitative Part Maps & Robustness

Visualizations of discovered semantic regions across diverse benchmarks from the main paper and supplementary material. iFAM selectively segments the task-relevant object foreground across various part granularities (\(K\)), preventing spurious background correlations from entering the Vision Transformer's receptive field.

OOD Transfer & Thresholding (CUB → WB200)
CUB (\(K=1, 4, 8, 16\))
Waterbirds (\(K=1, 4, 8, 16\))
MetaShift (\(K=1, 2, 4, 8\))
SIIM-ACR Medical X-Ray

CUB to Waterbird-200 OOD Transfer & Test-Time Token Removal Intervention

Qualitative results on CUB (\(K=8\)), corresponding out-of-distribution (OOD) images from Waterbird-200 with conflicting synthetic backgrounds, and the effect of test-time thresholding intervention (\(q=99\%\) and \(q=97\%\)) from Supplementary Fig. 1.

CUB Input	iFAM (CUB)	Waterbird-200 (OOD)	iFAM (WB200)	Intervention (\(q=99\%\))	Intervention (\(q=97\%\))

Observations: On Waterbird-200, synthetic background replacement causes small background false positives. Token-removal intervention with \(\tau_k^q\) successfully reclassifies low-confidence boundary tokens as background, recovering crisp part segmentations without model retraining.

CUB-200: Part Discovery Across Granularities (\(K=1, 4, 8, 16\))

Qualitative results of iFAM trained on CUB for different numbers of foreground parts \(K\) (from Supplementary Fig. 4).

Image
\(K=1\)
\(K=4\)
\(K=8\)
\(K=16\)

At \(K=1\), the model extracts holistic foreground. At higher \(K\) (\(4, 8, 16\)), it automatically discovers consistent fine-grained anatomical components (head, wings, belly, tail, legs) with zero background leakage.

Waterbirds: Part Discovery Across Granularities (\(K=1, 4, 8, 16\))

Qualitative results of iFAM trained on Waterbirds across different numbers of foreground parts \(K\) (from Supplementary Fig. 5).

Image
\(K=1\)
\(K=4\)
\(K=8\)
\(K=16\)

Water ripples, forest canopies, and marshland backgrounds are discarded across all part settings. Notice the high part consistency across varied bird species and orientations.

MetaShift: Part Discovery Across Granularities (\(K=1, 2, 4, 8\))

Qualitative results of iFAM trained on MetaShift across different numbers of foreground parts \(K\) (from Supplementary Fig. 6).

Image
\(K=1\)
\(K=2\)
\(K=4\)
\(K=8\)

Even with extreme background diversity (indoor furniture, beds, outdoor streets), iFAM reliably discovers torso, head, and facial regions of cats and dogs across granularities.

SIIM-ACR: Pneumothorax Lesion Localization in Chest Radiographs

Discovered torso part maps on chest radiographs alongside the dataset-wide ground truth pneumothorax occurrence heatmap.

On chest radiographs, standard classifiers latch onto chest drainage tubes (a confounding artifact). iFAM localizes true lung pathologies, and the central tube region can be pruned using our test-time intervention.

Pre-trained Checkpoints & Usage

All pre-trained iFAM models across Waterbirds, CUB, MetaShift, and SIIM-ACR are hosted on the Hugging Face Hub and can be loaded with minimal lines of code via PyTorchModelHubMixin:

from models import FullTwoStageModelDoubleClassify, FullTwoStageModelDoubleClassifyHF

# Load pre-trained Waterbirds model (8 parts, timm ViT-B backbone)
model = FullTwoStageModelDoubleClassify.from_pretrained("ananthu-aniraj/ifam-waterbirds-k8")
model.eval()

# Load pre-trained SIIM-ACR Pneumothorax medical model (transformers-based backbone)
model_med = FullTwoStageModelDoubleClassifyHF.from_pretrained("ananthu-aniraj/ifam-siim-acr-k8")
model_med.eval()

# Run forward pass:
# output contains stage 1 discovered masks and stage 2 robust classification logits
output = model(images)

🤗 Hugging Face Model Collection GitHub Repository

BibTeX

@inproceedings{aniraj2026ifam,
  title={Two-stage Vision Transformers and Hard Masking offer Robust Object Representations},
  author={Aniraj, Ananthu and Dantas, Cassio F. and Ienco, Dino and Marcos, Diego},
  booktitle={International Conference on Pattern Recognition (ICPR)},
  year={2026},
  organization={Springer},
  doi={10.1007/978-3-032-31673-8_20},
  url={https://link.springer.com/chapter/10.1007/978-3-032-31673-8_20}
}