Concept-Guided Fine-Tuning: Steering ViTs away from Spurious Correlations to Improve Robustness

Yehonatan Elisha; Oren Barkan; Noam Koenigstein

cs.CVMar 9, 2026

Concept-Guided Fine-Tuning: Steering ViTs away from Spurious Correlations to Improve Robustness

Authors: Yehonatan Elisha, Oren Barkan, Noam Koenigstein

Organizations: Tel Aviv University · The Open University

Abstract

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
	ViT-B	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
	DINOv2	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
	DeiT	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
	ViT-B	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
	DINOv2	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
	DeiT	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” $\rightarrow$ “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
Model	Train Classes	Method	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	$\checkmark$	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
	$\checkmark$	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
	$\times$	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
	$\times$	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	$\checkmark$	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 $\mathcal{L}_{\text{cls}}$ (Eq. 6 )	79.82	17.96	34.29	34.84	69.41	39.65	37.27	43.09	58.46
w/o $\mathcal{L}_{\text{non-concept}}$ (Eq. 4 )	82.50	19.14	35.13	36.26	72.01	41.70	34.93	42.31	58.40
w/o $\mathcal{L}_{\text{concept}}$ (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	$\lambda_{\text{align}}$	$\lambda_{\text{cls}}$	$\lambda_{\text{non-concept}}$	$\lambda_{\text{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.

Model	Metric	IN-V	IN-V2	IN-A	ObjectNet	IN-R	IN-Sketch
ViT-B	R@1	81.35 $\pm$ 0.28	69.19 $\pm$ 0.51	27.76 $\pm$ 0.14	54.28 $\pm$ 0.82	48.47 $\pm$ 0.39	37.06 $\pm$ 0.66
ViT-B	R@5	95.51 $\pm$ 0.47	84.77 $\pm$ 0.19	62.75 $\pm$ 0.73	75.46 $\pm$ 0.32	70.50 $\pm$ 0.56	62.59 $\pm$ 0.21
DINOv2	R@1	81.44 $\pm$ 0.61	71.91 $\pm$ 0.34	27.71 $\pm$ 0.80	53.89 $\pm$ 0.17	48.53 $\pm$ 0.54	44.74 $\pm$ 0.42
DINOv2	R@5	95.65 $\pm$ 0.42	88.15 $\pm$ 0.77	62.36 $\pm$ 0.54	75.58 $\pm$ 0.59	70.73 $\pm$ 0.18	68.90 $\pm$ 0.69
DeiT	R@1	82.61 $\pm$ 0.44	73.11 $\pm$ 0.34	27.72 $\pm$ 0.57	54.24 $\pm$ 0.71	48.33 $\pm$ 0.23	44.83 $\pm$ 0.36
DeiT	R@5	95.77 $\pm$ 0.68	88.58 $\pm$ 0.49	62.20 $\pm$ 0.31	75.46 $\pm$ 0.15	70.65 $\pm$ 0.75	69.55 $\pm$ 0.53

Table 9: Evaluation over 5 different seeds.

Explore similar work

Sparse Autoencoders enable Robust and Interpretable Fine-tuning of CLIP models
May 15, 2026Fabian Morelli, Arnas Uselis, Ankit Sonthalia +1Distribution Shift Robustness VLM Robustness
University of Tübingen ∗ · University of Tübingen
A Mechanistic Analysis of Adversarial Fine-tuning of Vision Transformers
May 28, 2026Hannah Gao, Isha Agarwal, Dylan Hadfield-Menell +1Adversarial Training Fine-Tuning
Massachusetts Institute of Technology · MIT Cambridge, MA
Robustifying Vision-Language Models via Test-Time Prompt Adaptation
Jul 10, 2026Xingyu Zhu, Huanshen Wu, Shuo Wang +4VLM Robustness VLM Adaptation
National University of Singapore · University of Science and Technology of China · The Hong Kong University of Science and Technology

May 15, 2026J/K move · Enter open · S save

Sparse Autoencoders enable Robust and Interpretable Fine-tuning of CLIP models

Fabian Morelli, Arnas Uselis, Ankit Sonthalia, Seong Joon Oh

University of Tübingen ∗ · University of Tübingen

Distribution Shift Robustness VLM Robustness VLM Adaptation VLM Interpretability Sparse Autoencoders

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Sparse Autoencoders enable Robust and Interpretable Fine-tuning of CLIP models

A Mechanistic Analysis of Adversarial Fine-tuning of Vision Transformers

Robustifying Vision-Language Models via Test-Time Prompt Adaptation