Did Models Learn Sufficiently? Attribution-Guided Training via Subset-Selected Counterfactual Augmentation

Abstract

Current visual models often make predictions based on a limited set of discriminative visual cues. As a result, they may become unreliable when the distribution shifts or when these cues are missing. Faithful attribution methods can reveal such problematic reliance through localized explanations, but they are typically used post hoc and are not fed back into the model. To address this limitation, we propose Subset-Selected Counterfactual Augmentation (SS-CA), a training strategy that masks decision-relevant regions to construct counterfactual samples and guide the model toward more robust decision boundaries. Specifically, we extend LIMA, a subset-selection-based faithful attribution method, to Counterfactual LIMA to identify regions whose removal shifts the model toward a competing class. SS-CA then selects near-boundary masks that reduce the logit gap while preserving the original semantics, and applies an adaptive counterfactual filling strategy to replace the masked regions without introducing external semantics. Feeding these counterfactual samples back into training encourages the model to exploit the remaining informative evidence and shifts the decision boundary toward a more robust one. Extensive experiments across five ImageNet variants show that SS-CA effectively improves ID accuracy, OOD generalization, and perturbation robustness, achieving gains of 5.70%/18.04% on ImageNet-1k/ImageNet-R with CLIP ViT/32b, 9.52%/11.33% on ImageNet-R/ImageNet-S on TinyImageNet-200 with ResNet-101, and about 4% under Gaussian Noise corruption. The code will be released soon.