Vision Transformers (ViTs) often degrade under distribution shifts because they rely on spurious correlations, such as background cues, rather than semantically meaningful features. Existing regularization methods, typically relying on simple foreground-background masks, which fail to capture the fine-grained semantic concepts that define an object (e.g., long beak'' and wings'' for a ``bird''). As a result, these methods provide limited robustness to distribution shifts. To address this limitation, we introduce a novel finetuning framework that steers model reasoning toward concept-level semantics. Our approach optimizes the model's internal relevance maps to align with spatially grounded concept masks. These masks are generated automatically, without manual annotation: class-relevant concepts are first proposed using an LLM-based, label-free method, and then segmented using a VLM. The finetuning objective aligns relevance with these concept regions while simultaneously suppressing focus on spurious background areas. Notably, this process requires only a minimal set of images and uses half of the dataset classes. Extensive experiments on five out-of-distribution benchmarks demonstrate that our method improves robustness across multiple ViT-based models. Furthermore, we show that the resulting relevance maps exhibit stronger alignment with semantic object parts, offering a scalable path toward more robust and interpretable vision models. Finally, we confirm that concept-guided masks provide more effective supervision for model robustness than conventional segmentation maps, supporting our central hypothesis.
Figures & tables
Figure 1: Motivation for CFT: Relevance maps produced by ViTs often concentrate on spurious background cues rather than semantically meaningful concepts. The figure illustrates this issue using ViT-B on ImageNet-A and ImageNet-R, showing relevance maps before and after applying CFT. By encouraging the model to focus on class-relevant, discriminative concepts, CFT substantially improves the semantic alignment of relevance maps. Notably, after CFT, the model highlights meaningful object parts, such as the beak and wings of the bird (top row) or the fins and mouth of the fish (bottom row), despite never being fine-tuned on these datasets.
Dataset
Model
Metric
Original
GradMask
RRR
RRDA
RobustViT
CFT
IN-V
ViT-B
R@1
80.41
80.65
80.89
81.12
80.24/80.11
81.35
R@5
94.88
95.03
95.17
95.39
94.16/94.25
95.51
DINOv2
R@1
81.07
80.95
80.48
81.61
80.38/80.72
81.44
R@5
95.20
95.08
94.83
95.53
95.05/95.12
95.65
DeiT
R@1
82.20
82.04
81.85
82.78
81.93/81.78
82.61
R@5
95.04
95.23
95.37
95.64
94.82/94.89
95.77
Table 1: Out-of-distribution (OOD) robustness and in-distribution accuracy: metrics are Top-1 and Top-5 accuracy (%).
Dataset
Model
Metric
Original
GradMask
RRR
RRDA
RobustViT
CFT
SI-location
ViT-B
R@1
34.26
35.18
36.42
38.54
39.02 /38.51
39.15
R@5
50.54
52.33
54.21
60.95
61.74 /60.88
62.48
DINOv2
R@1
35.11
36.05
37.22
38.95
38.79/38.21
39.07
R@5
52.85
54.67
56.40
61.33
61.18/60.75
62.15
DeiT
R@1
34.74
35.41
36.93
38.79
38.34/38.13
39.08
R@5
51.24
53.06
55.38
60.65
60.37/59.81
62.04
Table 2: Robustness to geometric transformations on the SI-Score benchmark: metrics are Top-1 and Top-5 accuracy (%).
Figure 2: Qualitative examples of CFT correcting prediction failures on OOD datasets using the ViT-B model: the baseline model (Original) misclassifies the images, with relevance maps often highlighting misleading context. Our CFT-finetuned model successfully corrects the prediction (e.g., “scorpion” → “common newt”) by focusing its relevance on the object’s core semantic concepts, demonstrating improved reasoning.
Metric
ViT-B
DINOv2
DeiT
CNv2
Original
mIoU
62.91
60.35
60.45
62.64
mAP
78.67
80.14
82.37
84.46
PA
72.23
74.85
77.16
79.32
CFT
mIoU
68.23
70.84
72.91
74.21
mAP
84.26
86.45
88.16
89.32
PA
80.34
81.92
83.74
84.58
Table 3: Alignment of relevance maps with ground-truth object masks: pixel-level agreement between model relevance maps and human-annotated masks. Additional details are provided in Sec. 4 .
Model
Train Classes
Method
IN-V
IN-A
IN-R
IN-Sketch
IN-V2
ObjectNet
R@1
R@5
R@1
R@5
R@1
R@5
R@1
R@5
R@1
R@5
R@1
R@5
ViT-B
✓
Original
82.02
96.26
14.86
35.26
32.96
47.45
36.85
58.42
71.39
89.86
36.03
55.82
CFT
82.39
96.42
22.95
46.10
35.53
50.59
38.16
59.73
71.46
89.92
43.51
64.67
×
Original
81.04
95.95
17.20
38.74
34.73
49.62
34.02
56.41
71.33
89.85
34.33
53.88
CFT
81.33
96.12
25.43
48.74
37.04
52.17
35.73
57.94
71.21
89.63
41.01
62.33
DINOv2
✓
Original
82.89
96.04
24.10
50.39
38.65
55.75
43.52
65.29
73.22
91.09
38.58
63.68
Table 4: Generalization to unseen classes: robustness evaluation was performed separately on classes included in the fine-tuning set and those excluded from it. Last row reports average change for both training and non-training classes across all models and datasets.
Model
Method
IN-V
IN-A
IN-R
IN-Sketch
IN-V2
ObjectNet
SI-location
SI-rotation
SI-size
ViT-B
Original
81.53
16.09
33.81
35.47
71.13
35.12
33.36
39.18
55.65
CFT
82.57
26.05
36.69
36.28
71.65
42.29
38.60
46.82
61.94
Segmentation
80.25
24.12
35.95
36.01
70.84
41.70
38.13
46.31
61.40
DINOv2
Original
83.15
17.52
35.20
36.83
72.81
36.78
34.63
40.25
57.09
CFT
83.62
27.88
38.02
37.61
73.04
43.91
39.82
48.19
63.11
Segmentation
82.04
25.91
37.14
37.12
72.25
43.15
39.07
47.73
62.85
Table 5: Ablation study: Concept-level guidance vs. object-level guidance.
Method
IN-V
IN-A
IN-R
IN-Sketch
IN-V2
ObjectNet
SI-location
SI-rotation
SI-size
Original
81.53
16.09
33.81
35.47
71.15
35.12
33.36
39.18
55.65
CFT
82.57
24.11
36.34
36.28
72.03
42.29
38.60
46.25
62.18
w/o Lcls (Eq. 6 )
79.82
17.96
34.29
34.84
69.41
39.65
37.27
43.09
58.46
w/o Lnon-concept (Eq. 4 )
82.50
19.14
35.13
36.26
72.01
41.70
34.93
42.31
58.40
w/o Lconcept (Eq. 3 )
81.26
24.10
34.37
35.22
71.69
42.14
38.55
45.61
61.79
w/ ground-truth
82.61
21.35
34.56
35.72
72.08
41.55
37.02
44.68
61.36
Table 6: Ablation study on loss components: we evaluate the impact of removing each of our three main loss terms using the ViT-B model.
Method
GradCAM
Gradient-Rollout
IIA
AttnLRP
IN-A
25.88
26.43
26.94
27.82
IN-R
43.32
46.92
47.58
48.54
Table 7: Ablation study on relevance methods: Top-1 accuracy using different relevance methods.
Model
λalign
λcls
λnon-concept
λconcept
Learning rate
CFT
ViT-B
0.8
0.2
1.2
0.5
5e − 7
DINOv2
0.8
0.2
1.2
0.5
6e − 7
DeiT
0.8
0.2
1.2
0.5
8e − 7
CNv2
0.8
0.2
1.2
0.5
3e − 6
RRR
ViT-B
-
2e − 6
1e − 10
-
2e − 6
DINOv2
-
2e − 8
1e − 10
-
1e − 5
Table 8: Hyperparameter selection for all methods.
Large-scale pre-trained vision-language models like CLIP demonstrate remarkable zero-shot performance across diverse tasks. However, fine-tuning these models to improve downstream performance often degrades robustness against distribution shifts. Recent approaches have attempted to mitigate this trade-off, but often rely on computationally expensive text-guidance. We propose a novel method for robust fine-tuning, SAE-FT, which operates only on the model's visual representations. SAE-FT regularizes changes to these representations by penalizing the addition and removal of semantically meaningful features identified by a Sparse Autoencoder trained on the pre-trained model. This constraint prevents catastrophic forgetting and makes the fine-tuning process interpretable, enabling direct analysis of semantic changes. SAE-FT is both mechanistically transparent and computationally efficient, matching or exceeding state-of-the-art performance on ImageNet and its associated distribution shift benchmarks. Code is publicly available at: https://github.com/Fabian-Mor/sae-ft.
The widespread use of image classification models in high-risk, real-world situations necessitates making these models robust to slight disturbances or perturbations, such as blurring or sharpening, in the input images. While vision transformers (ViTs) play an integral role in many modern-day multi-modal models like Vision-Language-Models (VLMs) and Vision-Language-Action (VLA) models, they have received a lack of attention in the setting of robustness. In this work, we analyze the effects of adversarial fine-tuning, a popular method for improving model robustness to image perturbations, on a ViT's performance on perturbed and regular images through a mechanistic lens. We adversarially train a ViT on low-frequency and high-frequency image corruptions, and attempt to explain changes in downstream model performance through an examination of the model's attention mechanisms, internal representations, and knowledge evolution. Overall, our results suggest that, while fine-tuning on inputs with common corruptions improves model performance and certainty on new instances of corrupted data, these improvements do not transfer to other classes of corruptions not seen in the training. Additionally, despite observing changes in visual attention and knowledge evolution across layers, we found that adversarial training did not lead to fundamental changes in the sparse representations learned by ViTs.
Pre-trained Vision-Language Models (VLMs) such as CLIP achieve strong zero-shot generalization, but their performance degrades sharply under adversarial perturbations. Existing test-time adaptation methods typically rely on sample-level confidence heuristics, overlooking the intrinsic distributional structure of the data. This sample-centric approach limits robustness, as it fails to distinguish confident adversarial mispredictions from true semantic consistency. In this work, we observe that adversarial distortion is structurally brittle: while holistic representations are corrupted, semantic integrity is often preserved in the distribution of augmented views. Motivated by this insight, we propose RITA, a Robust test-tIme prompt-TAdaptation framework that shifts from sample-level estimates to distribution-level alignment. Specifically, RITA employs optimal transport to align the distribution of augmented visual features with textual prototypes, mitigating adversarial outliers and rectifying cross-modal semantic misalignment. Furthermore, we introduce a dynamic cache to progressively accumulate reliable cues from the test stream for online refinement. Extensive experiments demonstrate that RITA significantly improves adversarial robustness without compromising clean accuracy.
Xingyu Zhu, Huanshen Wu, Shuo Wang +4
National University of Singapore · University of Science and Technology of China · The Hong Kong University of Science and Technology