iFAM (waterbirds-k4) Model Checkpoint

This is the official pre-trained checkpoint of the iFAM (Inherently Faithful Attention Maps for Vision Transformers) framework, proposed in the paper "Two-stage Vision Transformers and Hard Masking offer Robust Object Representations" (accepted as an oral presentation at ICPR 2026).

• Paper: Two-stage Vision Transformers and Hard Masking offer Robust Object Representations
• Repository: GitHub - ananthu-aniraj/ifam

Model Description

iFAM model trained on the Waterbirds dataset with 4 parts (K=4).

The iFAM framework is a two-stage approach:

1. Stage 1 (Selector): Processes the full image to discover object parts and identify task-relevant regions.
2. Stage 2 (Predictor): Restricts its receptive field to the selected regions using input attention masking, preventing spurious background details from affecting the classification.

Usage

To use this model, ensure you have the ifam repository cloned and the dependencies installed. You can load this model directly from the Hugging Face Hub:

import torch
import sys
# Append the path to the ifam repository
sys.path.append("/path/to/ifam")

from models import FullTwoStageModelDoubleClassify

# Load the model directly from Hugging Face Hub
model = FullTwoStageModelDoubleClassify.from_pretrained("ananthu-aniraj/ifam-waterbirds-k4")
model.eval()

# Example: Run inference
# x = torch.rand(1, 3, 518, 518)
# with torch.no_grad():
#     outputs = model(x)

Citation

If you use this model in your research, please cite:

@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}
}