Model Editing for Vision Transformers
Authors: Xinyi Huang
Kangfei Zhao
Long-Kai Huang
Venue: NeurIPS 2025 poster ()
Keywords: model editing
vision transformer
Primary area: general_machine_learning

Abstract: Model editing offers a promising paradigm for efficiently and precisely updating knowledge in pre-trained transformers without costly retraining. While extensively studied in language models (LMs), model editing for vision transformers (ViTs) remains underexplored. Existing methods typically adapt LM-based techniques by modifying the multi-layer perceptron (MLP) modules, overlooking the unique characteristics of ViTs. In this work, we show that ViT predictions are more strongly influenced by the multi-head self-attention (MSA) modules than by the MLPs. Building on this observation, we propose a two-stage framework for editing ViTs. First, we identify which attention heads are most responsible for incorrect predictions. Next, we selectively remove the corresponding features to correct the model’s prediction. To further balance error correction with predictive stability on unrelated data, we learn a projection matrix that refines the image representations. Extensive experiments across multiple real-world datasets and model editing benchmarks demonstrate that our method consistently outperforms existing model editing methods for ViTs, achieving superior generalization and locality. Our code is available at https://github.com/shanghxy/Model-editing-for-vision-transformers.

Model Editing for Vision Transformers
Xinyi Huang Independent Researcher xinyi0060e. ntu. edu. 5g
Kangfei Zhao
Long-Kai Huang
Beijing Institute of Technology zkf 1 11e1111@@maiil. com
Hong Kong Baptist University longkai @comp. hkbu. edu hkk
Abstract
Model editing offers promising paradigm for efficiently and precisely updating knowledge in pre-trained transformers without costly retraining. While exten- sively studied in language models (LMs), model editing for vision transform- ers (ViTs) remains underexplored. Existing methods typically adapt LM-based techniques by modifying the multi-layer perceptron (MLP) modules, overlook- ing the unique characteristics of ViTss In this work, we show that ViT pre- dictions more strongly influenced by the multi-head self-attention (MSA) modules than by the MLPs. Building on this observation, we propose two- stage framework for editing ViTs. First, we identify which attention heads are most responsible for incorrect predictions. Next, we selectively remove the corresponding features to correct the model's prediction. To further balance error correction with predictive stability on unrelated data, we learn a projec- tion matrix that refines the image representations. Extensive experiments across multiple real-world datasets and model editing benchmarks demonstrate that our method consistently outperforms existing model editing methods for ViTs, achieving superior generalization and locality. Our code is available at https: //github. com/ beaagghry/MMdell-diiiinn for--isioonninnoo
vision- transformers
Introduction
Model editing has emerged promising paradigm for efficiently and precisely updating the knowledge encoded in pre-trained transformers, without the need for expensive retraining. While numerous methods have been proposed for editing language models (LMs), there is equally pressing need to extend these capabilities to vision transformers (ViTs). In real-world deployments, vision models frequently exhibit unexpected prediction failures, highlighted in recent studies [23 2, 31. These failures are particularly common when the downstream data distribution differs from the pre-traine pre-training distribution. A key reason that pre-training datasets cannot fully capture the diversity of real-world subpopulations, causing vision models to rely spurious cues for predictions. For example, models may use background contextual attributes in images for recognition 34 20), leading to incorrect predictions when these contexts change. Therefore, there is a strong need for cost-effective model editing approaches rectify errors pre-trained vision models.
Editing ViTs, however, presents unique challenges that require specialized solutions. Although LMs and ViTs both use transformer architectures, their training objectives and internal dependencies differ. LMs, trained with auto-regressive loss to predict the next token, primarily encode factual associations in their MI.P modules 10 21). In contrast, ViTs are trained with classification loss, and their predictions depend more heavily multi- head self-attention (MSA) modules, which extract relationships between tokens form high-level concepts. Our empirical studies (see See indicate that MSA modules contribute significantly more to predictions than MLP modules in ViTs, a finding also supported by prior work on CLIP-VIT [ This reveals a key limitation of existing
Coorepsonding to Long-Kat Kai Huang
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(e)
MLP
MLP
Gan
Attettion head
Idertifeed diated lead
Condty
represorationn
MEP
VII
Trainable projection mEEi
Figure I: Overview of ViT: Stage One identifies ceanntionnhhd-leeee features linked to prediction errors. Stage Two learns projection matrix by optimizing editing success loss (Cunoss). ) locality loss (Cincaity) ) and cross-entropy loss (Ccs).
VIT editing methods, such as LiT-ViTs [32], which adapt ideas from LMs by focusing on editing MLP modules. Since ViTs' predictions rely mainly attention mechanisms rather than factual associations, editing the MLP modules may force the model to overfit the editing data, introducing unexpected information and leading to suboptimal performance. Therefore, specialized solutions that target the unique architectural and functional characteristics ViTs are necessary for effective model editing.
In this paper, we propose a novel model editing method tailored for ViTs, aimed at rectifying predictive errors using a limited number of samples. Our method is built on the observation that individual attention heads within MSA modules specialize in extracting features related to specific semantic concepts [5]. These modules can detect various concepts from input images and integrate these features into the residual stream. This suggests that classification errors in ViTs often arise not from failure to capture essential features, but from over-reliance spurious
non-causal ones.
To correct the prediction errors, we decompose the final representation of ViTs into features both the MLP-layer level and attention-heat level. Leveraging this decomposition, we develop two-stage framework called RefineVTT. In the first stage, we identify iitention-head-lles to prediction errors by analyzing small set of samples. In the second stage, we rectify the final representation by selectively ablating these problematic features. The framework enables effective model editing with only few examples and is illustrated in Fig.
features that contribute
The second stage, representation rectification, can be implemented in different ways depending on the task. For binary classification, we find that simply zeroing out the identified attention-head features yields substantial performance gains. However, in multi-label classification, directly removing these features may degrade performance on unrelated classes, as the ablated heads may still be important for other predictions. To address this issue, we learn projection matrix that aligns the final representation of erroneous samples with their ablated counterparts, while preserving the original representations of correctly predicted data. This approach balances error correction with predictive stability, allowing precise model edits while minimizing side effects. The projection can be integrated into the classifier, requiring no changes to the model architecture. Notably, our framework can also be seamlessly applied to ViT-based models like CLIP-VIT, which use cosine similarity between image and text embeddings for classification.
Summary of contribution:
We empirically reveal a key difference between LMs and ViTs for model editing: while editing MLP layers is effective in LMs, ViT predictions depend mainly on MSA modules, and errors are often driven by spurious features extracted by specific attention heads. This suggests that effective VIT editing should target MSA modules rather than MLPs.
We propose a two-stage method that identifies and removes problematic attention-head features from the final representation. To ensure both editing success and locality in multi-label classifica- tion, we learn a projection matrix that aligns erroneous samples with their corrected counterparts while preserving correct predictions.
We validate the effectiveness of our framework on multiple real-world datasets and model editing benchmarks. Our method achieves strong error correction with minimal impact unrelated data,

---

outperforming existing ViT editing baselines in both generalization and locality. Moreover, it generalizes well to ViT-based architectures such as CLIP-ViT.
Preliminaries
2.1 Problem Formulation:
Model Editing for ViTs
In this paper, tackle the challenge of error correction in Vision Transformers (ViTs). When ViT misclassifies an image a as ý instead of its true label y, this error often reflects a broader pattern rather than isolated mistake. Such errors typically recur across similar images with the same true label y. We define this subset of misclassified data as Iy 9 = {=') fal = D}. where fo is the ViT model with parameters 0.
Model editing aims to fix these misclassification by refining the model to correct not only the specific image æ but also similar instances in Iy.9. This process focuses two key goals:
Edit success measures how well the modified model correctly classifies previously misclassified data. We evaluate this by measuring the edited model's accuracy on Iy.0-
Edit locality measures how well the modified model preserves the prediction on unrelated data. This measured by the edited model's accuracy on data outside the target subset: Oy {ar ar ]. High edit locality indicates that the editing has minimal impact on unrelated data.
In practice, failures in object recognition are typically identified by testing the model's performance on small set of samples from each class. These samples, along with their true labels and model predictions, serve as the foundational data for model editing.
2.2 VIT Image Representation Decomposition
Vision Transformer (ViT) consists of residual blocks. Each block contains multi-head self- attention (MSA) layer, followed by multi-layer perceptron (MLP) layer. Both layers are preceded by layer normalization. The ViT processes an input image by dividing into N patches. Each patch is embedded into a d-dimensional token, resulting in N tokens:
{2() Ji=1 Additionally,
a special class token z added. Together, these form the initial state the residual stream
z = I 2f. >f..... .... 2x] E Ix(((+1), which updated by the residual blocks. The residual blocks
sequentially update this stream, and the layer-normalized class token from the final layer serves as the image representation (z)):
ž' MSA' Z (Z Z -'-1, z' = -MLP (2)) + 2', LE (1,2...., L}.
(1)
L
(x) = 1 dls = 12 [ds +
Σ 1=1
(M + [MMP((Z))
(2)
1=1
where (24 refers to the first column of z corresponding the class token. Here, we omit the layer normalization for notational simplicity.
Following [4) the output of each MSA layer can be written as the sum of the outputs from all attention heads, each multiplied by its own output matrix W Specifically, the class token output from the /-th MSA layer can be decomposed as:
H
H
N
MSA' (Zloe =
Σ h=1
He ih (z'
Σ h=1
Σ
W w !
(3)
Here, И is the number of attention heads in each layer; Head 1. is the h-th head in the /-th MSA 1.k layer; d0 is the attention weights from the class token to the i-th token; W and w are the output and value projection matrix; 2! z!-1 the i-th token from the (1 1)-th MLP layer.
By substituting Eq. 3 into Eqs. and defining h. = Headd.A (1 for simplicity, we can express the final image representation sum of contributions from the initial class token, all

---

VIT-L/16
VIT-B/16
VIISS/16
100
100
100
90
80
80
Accuracy (%)
60
70
40
60
MLP Ablation
20
MSA Ablation
0
2
4
6 8
10
12
14
16
18
20
22
24
6
8
10
12
2
4
6
8
10
12
Accumulated Aelan-e Layers
Figure 2: Accuracy changes after mean ablation of MLP and MSA modules.
MLP layers, and all attention heads:
4
((x) [2 Jlts EMMLP' (Zi)ete
(MMLP
ΣΣΣN [=1 h=1
h Ah
(4)
The decomposition in Eq. illustrates the direct contribution of the initial class token, each MLP
layer, and each attention head in the MSA layers to the final image representation []
Differences in Model Editing between LMs and ViTs
suggest that feed-forward layers in Transformers play crucial role in shaping predictions, making them effective targets for model editing. However, these findings largely limited Language Models (LMs) and may not be generalized to Visual Transformers (ViTs). In this section, we analyze the roles MSA and MLP modules in ViTs, and highlight key differences in model editing between LMs and ViTs.
Previous studies [2 10
3.1 Direct Effects of MSA and MLP Modules in ViTs
To quantify the direct effect of each component, we measure the drop in classification accuracy after mean ablation. Specifically, replace the output features of target modules with their dataset-wide mean, recompute the representations (see Eq. 4). and evaluate the resulting accuracy. A larger drop indicates stronger direct effect on predictions.
We conduct this analysis on ViT-S/16, ViT-B/16, and ViT-L/16 using 10,000 randomly sampled images from the ImageNet test set 3 These samples are used to compute mean values for each component and to assess classification accuracy.
Fig. 2 presents classification accuracies when replacing all direct contributions of MSAs or MLPs up given layer with their dataset-wide mean. Two key observations emerge. First, ablating MLPs causes only marginal drop in accuracy, indicating that MLPs have a minimal direct impact on predictions. Second, the effect of ablating MSAs varies across layers: replacing early MSAs has little effect, while modifying the last two leads to a notable accuracy decline, highlighting their critical role final predictions.
These findings suggest that, unlike in LMs, editing MLP layers in ViTs is less likely to yield significant changes predictions. Therefore, alternative strategies should be explored for effective model editing in ViTs.
3.2 The Devil is in the Spurious Correlations
Based these observations, we hypothesize that many errors in ViTs caused by the model relying on misleading features, a phenomenon known as spurious correlation 7 24 [1]. Simply fine-tuning MSA modules using the edit sample can distort the features extracted by attention heads, affecting predictions on unrelated data and resulting low edit locality. A more effective strategy may to remove the influence of misleading features only for the target samples.
To test this hypothesis, we conduct an empirical study to examine whether removing the direct contribution of feature from the image representation can facilitate error correction. We use the

---

100%
Attn Head Min
MLP
2006
100%
Relatte
2806
3000
400%
Layar
Figure 3: Relative change in logit gap between ground-trune and incorrect classes after ablating attention heads or MLP module at different layers. For attention heads, we show the maximal and minimal relative changes after ablating four heads. Results are based on ViT-B/16.
same editing dataset as in 32 For each editing case, we remove the direct effect of the feature MLP module
attention head in MSA modules, and report the relative change in the logit gap between the ground-truth class and the incorrectly predicted class. Specifically, let Go and g denote the logit gaps between the ground-truth class and the incorrectly predicted class before and after feature removal, respectively. The relative change is calculated 9-9-/19ol- A relative change greater than 100% indicates that the prediction is likely to shift from the original incorrect class to the ground-true truth class.
extracted by specific component, such
The results in Fig. 3 show that ablating certain attention heads achieve positive relative change, even greater than T00%. This supports our hypothesis that errors are often caused by misleading features extracted by specific attention heads, and that removing their direct effect can correct these erros.. Additionally, we observe that ablating attention heads has a much greater impact than ablating MLP modules, further confirming the central role of MSA modules in ViT predictions.
4 Refine vIT: Model Editing for ViT
In the previous section, we show that misclassifications in ViTs often stem from misleading features extracted by attention heads. Building this observation, we propose R tiot-aitsstsssaage model editing framework designed to mitigate such biases. In the first stage, we identify the attention heads most responsible for misclassification. In the second stage, temporarily disable (ablate) them to create a modified model that corrects the error. To improve edit locality. we use the knowledge gained from this ablation to update the original model, enabling it to fix the error while preserving predictions on unrelated inputs.
4.1 Identify the Cause of Errors
Our analysis shows that attention heads in the later MSA layers have a strong direct impact on the model's final predictions. Accordingly. we focus on these layers to find the specific heads most responsible for misclassification.
To this end, we define ablation matrix A, where each element Am {0, 1} indicates whether the h-th head in (-th layer is ablated (Ai, 1) or not (Au, = 0). When a head is ablated, its feature is replaced by zero in the representation. Therefore, the image representation after ablation ((x A) is calculated
1.
H
1. ΣΣ 1=1 R=1
(xx A) 12 +
MLPP' (2MLF'(2))a. Σ (
+(x))
Auhht
(5)
1=1
1=1 k=1
where + (x)) is the original image representation.
To measure how effective an ablation is at correcting an error, we define the following utility function:
U( v. ú.
($(z, A)) (a A)
(6)
where is the input image, y the ground-truth label, ý the incorrect prediction, and f(),, represents the output logit for class y from the pretrained classifier.
5

---

From the perspective of ablation study, (xx y. B. A) - Uz y, ú. 0) simultaneously quantifies the of direct effect of the ablated attention heads to the incorrectly predicted label ÿ as (0(z))
f (0(2, A)) and the negative direct effect to the ground label as ($(x)), -f (z) A))
Since U (xx y, y, 0) is constant, we can ignore it and identify the attention heads within these layers that contribute most to prediction errors by maximizing this utility
A = arg max U y. D. A
However, solving the problem regarding the binary matrix A is NP-hard and generally requires searching over all the possible solutions, which infeasible. To make this tractable, we consider abla- tions where only one head is ablated time. For each head, we compute the utility Ua, v, D. In), where Lin is a matrix with a single at position (1, h) and O elsewhere. We then rank all L x H heads by their utility scores and select the top Tr, heads with the highest scores. Empirically, observe that the selected heads are predominantly located in the later layers, consistent with our earlier findings. We denote the set of these heads Sablate and construct the approximate optimal ablation matrix A by setting Aln 1 for (L, h) E Sablate and 0 otherwise.
In most cases, only single error sample available for editing, so we use zero-ablation to estimate the contribution of each attention head. When more error samples correctly classified samples available, more advanced techniques, such as mean ablation, can be used for a finer analysis. Details of these extended methods provided in Appendix A
4.2 Refine ViT Through Representation Rectification
After identifying the attention heads that cause errors, straightforward way to fix mistakes is to directly remove their direct contributions from the image representation. For binary classification tasks, this simple ablation can already lead to large improvements. Specifically, we compute the ablated image representation as (xx A).
However, in more complex settings such as multi-label classification, simply removing these features can hurt performance on other classes. This because the ablated heads may still be important for unrelated predictions. Moreover, modifying the forward pass to remove features can make deployment more difficult.
To address these issues, we introduce a feature projection matrix P(O) E Rddd, where is the trainable parameters and d is the dimension of the image representation. This matrix is applied after the Transformer blocks and rectifies the image representation as follows:
(x)) (())((())= P(0)
Jel=
P(O)MMPP' (Z')) )es. 1Σ P(O)
(8)
To achieve edit success, we propose transferring knowledge from the ablated representation, ((x A). to the projected representation by minimizing the MSE loss for samples where the ViT fails make accurate predictions (æ E Draited) as follows:
Luncess (0) = Ee Diated P(O)(()) 6(x, A)
To achieve edit locality, we want the projection keep the representation unchanged, preserving the model's original behavior. We use another MSE loss to align the representations before and after projection for samples where the model is already correct (xr D as follows:
Ciocalty (0) =
E Downe
(())
+())]]
(10)
If D Duncess is not available, we constrain the projection matrix to be close to the identity matrix 7 as
Ciocality (0) PP()) - 1
(11)
We combine these two objectives with the standard cross-entropy loss (Cx(#) for classification. The total loss for learning the projection is:
0" arg min javCuncers (0) + 3 Stucaaltt (0) + (0))
(12)

---

Natural image subset
Al oil painting subset
100
100
95
95
90
90
85
GR (%)
85
GR (%)
80
75
80
70
75
65
70
60
55
65
50
60
45
55
60
65
70
75
80
85
90 95
100
55
60
65
70
75
80
85
90
95
100
LR (%)
LR (%)
FT-12
-.... T-Patcher
MEND
-K--
SPT
LWE-VITs
LoRA
KE
KN
ROME
RefineViT (ours)
Figure 4: Editing results for ViT/B-16.
where a and /3 are hyperparameters that balance error correction and locality preservation.
This projection-based approach allows to correct specific errors while minimizing side effects on other predictions. Importantly, it can be easily integrated into ViT-based models, including those like CLIP-ViT that use cosine similarity for classification, since the projection only modifies the image representation before the classifier.
Related Work
Model editing techniques [22 33
aim to refine the behavior LLMs for specific input-output pairs, while preserving their performance on other data. These methods can be grouped into three main categories: bussificr-based geetaleearning-bee based, and bocate-then-edie methods.
Classifier- model editing methods retain the pre-trained parameters and use an auxiliary classifier determine behavioral modifications. This method ensures that the modifications are applied only to targeted samples, leaving the original model predictions unchanged for unrelated samples outside the edited scope. locate-then-edee methods first identify model parameters associated with specific knowledge, often through causal tracing, and then directly update these parameters to achieve the desired edits. Meta-learning-basee
methods utilize a hyper-network, known as an editor, to update parameters. This editor is meta-trained across multiple editing tasks to learn how to generate the necessary updates based on the provided edit samples. For a comprehensive review of these methods for language models, please refer to (33].
Despite its strides language models, adapting similar techniques to visual models like Vision Transformers (ViTs) and CLIP remains largely untapped. [26] adapted classifiers in convolutional neural networks to mitigate concept-level spurious features by mapping misleading visual concepts to correct targets. However, this requires prior knowledge of the erroneous visual concept, its location, and the target concept, which may not always be available. Another line of work [ proposes to interpret model components and ablate spurious components to rectify errors. Yet, these methods also rely on knowing which visual concept triggers the error. Recently, (32) introduced method inspired by LMs editing techniques, which fine-tunes MLP modules in ViT's and employs hyper-network to determine where to apply edits. However, as discussed in Section 3 this approach may yield suboptimal performance due to fundamental differences between LMs and ViTs.
Our method shares the goal of correcting errors by removing misleading features, as 16 26
However, unlike these approaches, we identify spurious features using data-driven strategy, without requiring prior knowledge of biases or their locations. This makes our method more practical for real-world scenarios where such information is not readily available.
7

---

Table 1: Average-gras gropp accuracy (%) and gorst accuracy (%) on Waterbirds Dataset. Methods marked with an asterisk (*) use additional data for training or validation. The best results are highlighted in bold, while the second-bot best results are underlined.
VIT-B/16
VIT-1114
VIT-H/14
Method
Avg. (t)
Wst. (1)
Avg. (1)
Wst. (1)
Avg. (*)
Wst. (f)
Base
72.8
45.6
75.5
47.7
68.6
37.2
Tip-Adapter (training- free)
74.4
46.9
77.4
52.6
70.3
38.0
Tip-Adapter (training- based)
76.3
49.9
78.0
52.2
74.8
59.3
Refine ViI (ours)
81.1
61.4
85.5
72.1
75.9
51.3
TextSpan *
78.5
57.5
84.4
72.9
72.9
43.3
ViT (ours)
80.4
65.9
85.6
75.6
75.9
51.3
6 Experiments
We evaluate the proposed method, RiT, the VIT editing benchmark from [32], the Binary Waterbirds dataset [25]. CelebA [9], ImageNet-R [33, and ImageNet-A [1]. Our experiments are designed to answer the following research questions: Q1: Does Refine which edits the MSA modules, outperform
tftte-of-thearrn methods that primarily focus on editing the MLP modules, in
terms of both editing success and locality in ViTs? (Sections 6.1) Q2: Can Refine ViT generalize
to other ViT-based models such as CLIP-VIT? (Section 6.2) Q3: Why does Refine ViT work?
Specifically, is it effective in identifying the attention-head-levll features responsible for prediction Ablation studies and sensitivity analyses are provided in Appendix D
failures? (Appendix C
6.1 Edit success and edit locality
To address Q1 and assess RefineViT's editing success and locality, we evaluate on the benchmark proposed by (32) This benchmark collects misclassified samples by ViTs from naturally underrep- resented images and Al-generated images. Further details about this benchmark are provided in Appendix [...
Evaluation Metrics. Following the experimental setup in 32 we evaluate all model editing methods on the single-sample editing task and compare their performance using three evaluation metrics: 1) Success Rate (SR): the prediction success rate of the edited model on the single sample used for correction; 2) Generalization Rate (GR): the accuracy of the edited model on neighboring samples within the editing scope; 3) Loculity Rate (LR): the accuracy of the edited model on unrelated samples outside the editing scope.
We rigorously evaluate Refine under the same conditions as [32]. and compare its performance with recent model editing methods as follows:1) Learning- Learning-Where- To-Edir (LWE-Vo ViTs) [32]: A meta-learning-
Competing Methods.
approach that selects edit locations within the 8th 10th
MLP layers in ViT-B/16; 2) Standard Fine- Tine-Tuning (FT): standard fine-tuning targeting the 8th to 10th MLP layers; FT-Ez: extends standard fine-tuning by incorporating f--norm regularization; 4) KE 22 and 5) MEND (22]: leverage neper-networks to guide parameter updates in the last three MLP layers; T-Patcher 17 introduces and trains a small set of additional neurons in the final MLP
layer; 7) SPT [12]: sensitivity-awar parameter-efficient fine-tuning: 8) ROME (2I]: modifies the
last MLP layer to update specific factual associations; 9) LoRA 15 low-rank updates to all MSA layers in the Transformer.
Main Results. Since most methods achieve nearly 100% SR. demonstrating their effectiveness in correcting single predictive errors, we focus on the GR-LR curve in Fig As shown, Refine ViT outperforms LWE-ViTs 32, which the previous state-of-the-art on this benchmark, and achieves substantially better GR-LR performance than all baselines. This improvement likely stems from the fact that most existing methods (LWE- ViT',, FT, FT-&,, KE, MEND, T-Patcher ROME) focus on editing MLP layers, strategy inspired by findings in language models that highlight the efficacy of modifying MLP modules NLP tasks. In contrast, RifineVTT targets the last few MSA layers, which proves more effective for computer vision tasks. Although LoRA also targets MSA layers, it updates all of them directly and simultaneously, resulting a large number of trainable parameters and potential distortion of attention head features.

---

Avg. Acc. (langet classes)
Avg. A.c. classes)
Avg. Avg.Acc. (target classes)
Augg (ochee classes)
85
80
80
Accuncy (%6
50
80
50
Std. Flae tume RefneVi
70
Su.. Fine oune RefieVI
25
Sod. Pine-tune
Sed. Fine-tune
20
ReineViT (oun)
RefineVIT (ouns)
70
0
50 spoch
100
90 opoch
10n
50 opoch
100
spoch
Figure 5: Comparison of ViT and Standard Fine-Tuning with ViT-L/14. From left to right, the first two figures show average accuracy the target classes and other classes on the Waterbirds ImageNet-R combined dataset, while the next two figures display the corresponding results for the ImageNet-A dataset. We repeat the experiment with three random seeds and report the average results.
Moreover, RiT is highly efficient, as it avoids modifying the Transformer and instead trains a lightweight projection matrix. This allows updates in under 0.3 seconds for 50 epochs on a single NVIDIA A100 (40GB) as the output of the backbone can be cached and reused. In comparison, WiTs takes around 12 seconds for fine-tuning under the same conditions, without considering the extra cost of training hyper-network before editing.
6.2 Generalization to ViT-Based Models
Although ViT is designed for ViTs, it generalizes effectively to ViT-based models such as CLIP- ViT, which uses cosine similarity between textual and visual embeddings for zero-shot classification. We evaluate Refine ViT within CLIP-ViT for both performance enhancement and model editing tasks.
Performance Enhancement and Debiasing, We evaluate ReT on the Binary Waterbirds [25 and CelebA [99 datasets, both widely used benchmarks for CLIP debiasing. Our goal is to enhance CLIP-ViTs' zero-shot ability to distinguish between waterbirds and landbirds despite background interference in the Binary Waterbirds dataset, and to improve its accuracy in predicting age (young vs. old) in the CelebA dataset. More implementation details can be found in Appendix (B2
Competing Methods. On the Waterbirds dataset, we compare Refine VIT with: (1) TextSpan [5]. which treats background 5,000 test images and identify spurious attention heads; and (2) Tip-Adapter [35], a robust method that improves CLIP-ViTT s accuracy while preserving its zero-shot capabilities. We evaluate both its training-free and training-
known spurious feature and uses human expertise to analyze over
variants. For the CelebA dataset, where spurious features are unknown, we compare only against Tip-Adapter. Since these baselines focus binary classification rather than model editing, we use a simplified version Refine applying zero ablation in the second stage instead of learning a projection matrix.
Results. Table D presents the performance of various methods and CLIP-VIT variants on the Binary Waterbirds dataset. We report both average and worst-case accuracy across the four bird groups (landbirds on land, landbirds on water, waterbirds on water, and waterbirds on land). For CelebA, Table shows the results of Refine and Tip-Adapter in predicting the "Young" vs. 'Old' attribute using
Key observations include: (i) Given the same number of samples per class, Refine ViT outperforms Tip-Adapter both datasets. (ii) Refine also generally surpasses TextSpan on Waterbirds, despite using fewer samples, incurring lower computational cost, and requiring no human intervention. Its performance can be further improved with a validation set.
Model Editing. We evaluate ViT for CLIP-ViT in model editing scenarios using ImageNet-A and combined dataset of Waterbirds and ImageNet-R, Additional details are available in Appendix
Results. We compare Rif with standard fine-tuning, which also updates the projection matrix using the same samples. As shown in Fig. [5] Refine ViT achieves higher accuracy target classes while substantially mitigating performance drops on non-target classes, avoiding the common issue caused by overfitting in standard fine-tuning-
Ablation study and Sensitive Analysis. We conduct experiments to analyze the impact of the hy perparaaaetss and 3 in the loss function Eq. [2), the number of candidate attention heads

---

selected (I") in stage one, the update strategies employed, and other factors to fully assess our method. The results and analyses are presented in Appendix D
Conclusion and Future Directions
In this work, we study model editing for vision transformers and uncover key difference from language models. While MLP modules are often the main editing targets language models, we find that the predictions of vision transformers self-attention layers. Based on this finding, we introduce Refine a two-stage framework that first identifies the attention heads responsible for prediction errors and then refines their representations to correct these errors. Experiments on standard benchmarks show that Refine achieves SOTA performance in both edit accuracy and generalization. The framework also generalizes well to ViT-based models such as CLIP-VitE demonstrating its broader applicability. Additional ablation and sensitivity studies confirm the robustness and effectiveness of the proposed approach.
more sensitive to changes in the multi-head
Our framework still faces several limitations. The current analysis considers only the direct effect (first-order effect) of individual attention heads and does not account for the indirect effect that may accumulate across layers. Incorporating these higher-order effects may further improve reliability. In addition, our experiments focus image classification tasks. Extending Refine to dense- prediction tasks will require new attribution strategies and evaluation protocols. Finally, model editing involves trade-off between success, generalization, and locality. Although our projection- based refinement helps reduce unintended side effects, achieving complete locality remains an open challenge for future work.
Acknowledgments and Disclosure of Funding
Kangfei Zhao is supported by National Key Research and Development Plan, No. 2023
References
[1] Sriram Balasubramanian, Samyadeep Basu, and Soheil Feizi Decomposing and interpreting image representations via text in vits beyond clip. Advances in Neural Information Processing Systems, 046-810076 2024.
[2] Nicola De Cao, Wilker Aziz, and Ivan Titov. Editing factual knowledge in language models. In Marie-Francine Moens, Xuanjing Huang, Lucia Specia, and Scott Wen-tau Yih, editors, Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 6491-6506 2021.
[3] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248-255, 2009.
[4] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Sattield-Dodds Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https:/transffrmer--cice www.htsssppppo2t0le/ wttpuiissppo//222//maaaeooooo//o wwtp///ww//////ee...iii html.
[5] Yossi Gandelsman, Alexei A. Efros, and Jacob Steinhardt. Interpreting clip's image representa- tion via text- decomposition. In International Conference on Learning Representations. OpenReview.net, 2024.
[6] Irena Gao, Gabriel Ilharco, Scott M. Lundberg, and Marco Túlio Ribeiro. Adaptive testing of computer vision models. In IEEE/CVF International Conference on Computer Vision, pages 3980 3991. IEEE, 2023.
[7] Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A. Wichmann, and Wieland Brendel. Imagenet-traine cnns are biased towards texture; increasing shape bias
10

---

improves accuracy and robustness. In International Conference on Learning Representations. OpenReview.net, 2019.
[8] Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in auto-regressive language models. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 1221 2216-1 2023.
[9] Mor Geva, Avi Caciularu, Kevin Ro Wang, and Yoav Goldberg. Transformer reww-fowward layers build predictions by promoting concepts in the vocabulary space. In Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pages 30-45, 2022.
[10] Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer
layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages
2021.
[11] Soumya Suvra Ghosal and Yixuan Li. Are vision transformers robust to spurious correlations? International Journal of Computer Vision, 132(3):689-7777 1009999799290 2024.
[12] Haoyu He, Jianfei Cai, Jing Zhang, Dacheng Tao, and Bohan Zhuang. Sensitivity-aware visual parameter-efficient fine-tuning. In International Conference Computer Vision, pages 11791-11801. IEEE, 2023.
[13] Dan Hendrycks, Steven Basart, Norman Mu, Saurav Kadavath, Frank Wang, Evan Dorundo, Rahul Desai, Tyler Zhu, Samyak Parajuli, Mike Guo, Dawn Song, Jacob Steinhardt, and Justin Gilmer. The many faces of robustmess: A critical analysis of ICCV, 2021.
generalization.
[14] Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. Natural adversarial examples. CVPR, 2021.
[15] Edward Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations. OpenReview.net, 2022.
[16] Long-Kai Huang, Peilin Zhao, Junzhou Huang, and Sinno Pan. Retaining beneficial information from detrimental data for neural network repair. Advances in Neural Information Processing Systems, 36:48051-48006 2023.
[17] Zeyu Huang, Yikang Shen, Xiaofeng Zhang, Jie Zhou, Wenge Rong, and Zhang Xiong. Transformer-patcher: One mistake worth one neuron, 2023.
[18] Gabriel Ilharco, Mitchell Wortsman, Ross Wightman, Cade Gordon. Nicholas Carlini, Rohan Taori, Achal Dave, Vaishaal Shankar, Hongseok Namkoong, John Miller, Hannaneh Hajishirzi, Ali Farhadi, and Ludwig Schmidt. Openclip, July 2021.
[19] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision, December 2015.
[20] Huan Ma, Yan Zhu, Changqing Zhang, Peilin Zhao, Baoyuan Wu, Long-Kai Huang, Qinghua Hu, and Bingzhe Wu. Spurious feature eraser: Stabilizing test-time adaptation for vision- language foundation model. In Toby Walsh, Julie Shah, and Zico Kolter, editors, AAAI, pages 19296-19304. AAAI Press, 2025.
[21] Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in GPT. Advances in Neural Information Processing Systems, 2022.
[22] Eric Mitchell, Charles Lin, Antoine Bosselut, Chelsea Finn, and Christopher D. Manning. Fast model editing at scale. In International Conference Learning Representations. OpenRe- view.net, 2022.
[23] Benjamin Recht, Rebecca Roelofs, Ludwig Schmidt, and Vaishaal Shankar. Do imagenet classifiers generalize to imagenet? In Proceedings of the 36th International Conference on Machine Learning, volume 97, pages 5389-5400. PMLR, 2019.
11

---

[24] Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. volume abs/191 108731, 2019.
[25] Shiori Sagawa, Pang Wei Koh, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization, 2020.
[26] Shibani Santurkar, Dimitris Tsipras, Mahalaxmi Elango, David Bau, Antonio Torralba, and Aleksander Madry. Editing a classifier by rewriting its prediction rules. Advances in Neural Information Processing Systems, 34:23359-23373, 2021.
[27] Ramprasaath Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: Visual explanations from deep networks via gradient-based localization. In Proceedings of the IEEE international conference on computer vision, pages 618-626, 2017.
[28] Joshua Vendrow, Saachi Jain, Logan Engstrom, and Aleksander Madry. Dataset interfaces: Diagnosing model failures using controllable counterfactual generation. ArXiv preprint arXiv:2302.07865 2023.
[29] Catherine Wah, Steve Branson, Peter Welinder, Pietro Perona, and Serge Belongie. The Caltech- UCSD Birds-200-201/ Dataset. Jul 2011.
[30] Haohan Wang, Songwei Ge, Zachary Lipton, and Eric P Xing. Learning robust global represen- tations by penalizing local predictive power. In Advances in Neural Information Processing Systems, pages 10518, 2019.
[31] Haotao Wang. Tianlong Chen, Zhangyang Wang, and Kede Ma. I am going mad: Maximum discrepancy competition for comparing classifiers adaptively. International Conference on Learning Representationss 2020.
[32] Yunqiao Yang. Long-Kat Huang, Shengzhuang Chen, Kede Ma, and Ying Wei. Learning where to edit vision transformers. In Advances in Neural Information Processing Systems, 2024.
[33] Yunzhi Yao, Peng Wang, Bozhong Tian, Siyuan Cheng, Zhoubo Li, Shumin Deng, Huajun Chen, and Ningyu Zhang. Editing large language models: Problems, methods, and opportunities. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 10240- 2023.
[34] Michael Zhang and Christopher Ré. Contrastive adapters for foundation model group robustness. Advances in Neural Information Processing Systems, 35:21682-21667,
[35] Renrui Zhang, Wei Zhang, Rongyao Fang, Peng Gao, Kunchang Li, Jifeng Dai, Yu Qiao, and Hongsheng Li. Tip-adapter: Training-free adaption of CLIP for few-shot classification. In European Conference Computer Vision, volume 13695, pages 493-510. Springer, 2022.
[36] Bolci Zhou, Aditya Khosla, Agata Lapedriza, Antonio Torralba, and Aude Oliva. Places: An image database for deep scene understanding, 2016.
12

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims made in the abstract and introduction accurately reflect the paper's theoretical and experiments results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitations are discussed in the limitations section in appendix ??.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
13

---

The paper does not include theoretical results. The formulas are numbered and cross-referenced.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
The implementation details are shown in Appendix B
Guidelines:
The answer NA means that the paper does include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
14

---

Answer: [No]
Justification:
The code will be open-sourced upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https guwdes/ddeeeeee Cooeeeoodeesssseeeeeee 7/nips.com htte:/uulllllliiiee/e//eeeuuaasl ssiniiiiccyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification:
The implementation details are shown in Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Standard deviations multiple random seeds are reported for experiments where applicable.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
15

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The computer resources are provided in section 4
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
The research meets the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA].
Justification: This paper presents work whose goal is to advance the field of Model Editing for ViTs. There are many potential societal consequences our work, none of which feel must be specifically highlighted here.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
16

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA].
Justification: Societal impacts are not applicable to this paper
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The work is solely based open-sourced dataset, models and code.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
17

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not introduce any new datasets or models.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details of the datasettcodeemmodll as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve any crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA].
Justification: This paper does not involve any crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
18

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: This paper does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy btteps://neuriie /t/CCooffreeccss//2222 for what should or should be described.
19

---

A Head Identification Using Mean Ablation
The utility function in Eq. (18), which is based on zero-ablation, is designed for the typical model editing scenario where only a single error sample is available. To extend this approach to cases where multiple error or correctly classified samples are present, we introduce mean-ablation strategy, Specifically, we assume access to three sets of data: (1) a set of misclassified samples with ground- truth label y and predicted label ý, denoted as g.g.i (2) set of correctly classified samples with ground-truu ttunn label y, denoted as Cy; and (3) a set of correctly classified samples with ground label 9. denoted as Cg. For each set, we compute the average feature for each attention head, denoted as h W,,"" h.hh and C, respectively. The mean-ablated representation of an input x is then defined as
((x, h(*h = ((x) - 1 + h.t.
(13)
We propose four utility functions to assess the contribution of each attention head to model predictions, each leveraging different combinations of the available data sets.
Utility A. To evaluate the influence of attention head h we first consider replacing its feature in misclassified samples Wy, with the average feature from correctly classified samples Cy. i.c., fcc, 1.b We then measure whether this replacement shifts the model's prediction towards the correct label y and away from the incorrect label ý:
A = E, (i(az - (a(r,, ( bi
(
(14)
large value of U indicates that the head in question plays significant role in the model's misclassification, as its replacement leads to a notable correction in the prediction.
A
Utility B. While Utility A focuses on misclassified samples, Utility B takes the opposite perspective by considering correctly classified samples in Cy. Here, we replace the attention head feature with the
average from misclassified samples h and measure the extent to which this substitution causes the prediction to deteriorate, i.e., to shift towards the incorrect label ÿ::
Tha = Fwec, [1 kit,,))
(15)
A higher value of samples, pushes the model towards making the same error, highlighting its contribution to in incorrect predictions.
B suggests that the head, when replaced with features from misclassified
Utility C. Utilities A and B focus on samples with the same ground-trute truth label but different predictions. In contrast, Utilities C and D examine samples that share the same predicted label but have different ground-truth labels. For Utility C, we consider misclassified samples in Wy.g and replace the attention
head feature with the average from correctly classified samples of class g. i.c., L.h.
Udhh = Euew,. [ (ā(z,)
mE"),, - (tt
her))
(16)
Intuitively, if the attention head is causal for class ÿ. this replacement will further reinforce the incorrect prediction, resulting in a small U.A Conversely, if the head is misleading, the replacement will have limited effect, leading Utility D. Finally, Utility D evaluates the effect of replacing attention head features in correctly classified samples of class g. i.c.. Cg. with the average from misclassified samples L. h
larger U than that of causal attention head.
Uh h = Ewec, [ (0(2z,
i.b hw,,,) ) - (0( h
W...) ))
(17)
Since the data in Wy. is misclassified, the attention head features associated with the incorrect class y. which are spurious features, tend to dominate the image representation, while the contribution from causal features relevant to the true class remains weak. Consequently, replacing the misleading heads (those aligned with zj further increases the model's prediction confidence for class ý. whereas replacing the causal heads (those aligned with y) enhances the prediction for the correct class y. As a result, the misleading heads exhibit large values of Ul-h while the causal heads correspond to smaller Uh values.
20

---

Summary of Utility Scores, We propose four utility scores to analyze the role of attention heads. The first two, Uhh and U compare samples with the same ground-truth label but different predictions. These scores help identify attention heads that contribute to inconsistent or unstable predictions. In contrast, U and U compare samples with different ground- labels but identical predictions, highlighting attention heads that may cause the model to confuse distinet classes.
Figure 6 illustrates these concepts using the Waterbird dataset [25,, where the goal is to classify images waterbirds or landbirds. In this dataset, the background (water or land) is common source of spurious correlation. Since the four utility scores are designed to capture different types of features, but their specific behavior is difficult validate in general, we leverage this known spurious feature to test them. By selecting failed samples in which background plays a key role, we can examine whether the scores respond expected.
Specifically, using UAA and U examine types of samples: correctly predicted waterbird with a water background, and another incorrectly predicted as a landbird, despite being a waterbird with land background. If attention head focuses the background, it may help in the first case but harm in the second, indicating unstable behavior. U and Ul.h proposed identify these type of attention heads.
On the other hand, U and U focus on cases where the model predicts the same label for two semantically different samples, such waterbird on land and landbird water, both predicted as landbirds. Attention heads that rely on background features can mislead the model both cases, suggesting they encode spurious correlations that negatively affect generalization. U/ and U/ are proposed identify these attention heads.
Waterbird
Model
X
Landbird background
Model
background
andhiid
Landbird
Waterbied
Model
Model
Water background
Waserbird
background
Landbird
UA & UB
Uc & UD
Figure 6: Illuttatioon of utility scores on the Waterbird dataset. The computation of the four utility scores does require carefully selected samples or prior domain knowledge. The examples shown here are chosen solely to illustrate the properties of the scores.
Ranking and Ensembling of Utilities. Each utility score (ULA Ulh U.h U.lA)) produces ranked list of attention heads. For each utility, we select the top T heads, and for each from to T. we define S, as the set of the top f heads. To determine the optimal set of heads to ablate, we evaluate the validation utility of each candidate set S; as follows:
U(S.) = E
(0(z))
Σ (a))
Σ {((aa)
E Σn (m))
(18)
D'Fuul
WINES,
where D denotes the data set consisting of all available samples and Fet refers to the ground truth label for each corresponding image. The candidate set Sabtate with the highest validation utility is selected for ablation.
Summary of the Procedure. For model editing tasks with multiple available samples, the first stage of Refine ViT identifies the attention heads most responsible for prediction failures through the following systematic process:
1. Score Calculation: For each attention head, compute all four utility scores (Uhh U UtA UbA)) permitted by the available data.
21

---

2. Candidate List Generation: For each utility, rank the attention heads in descending order and select the top T heads to form candidate list.
3. Ablation Candidate Generation: For each candidate list, iteratively select the top heads (for t= 1 to T) as ablation candidates, denoted as St.
4. Validation Utility Evaluation: For each ablation set St, evaluate the validation utility (Eq. [8) on all available samples. The set St with the highest validation utility is chosen as the final set of attention heads for ablation, i.e., Sablate
This strategy enables Refine to systematically identify and ablate the attention heads that most contribute prediction errors, leveraging all available data for effective model editing.
B
Experiment Details and Extra Results
B.1
More Details in Section 6.1
The benchmark introduced by 32 comprises 16 types of misclassification for natural images and 22 types of misclassification for Al-generated oil painting images on ViT-B/16. Each sample serves as a single reference point for model editing. simulating scenario where only one misclassified sample is available for correction. Other samples exhibiting the same type of prediction error as the reference sample are treated as neighboring samples within the editing scope for the generalization rate (GR) computation. Furthermore, 2,071 carefully curated images near the decision boundary of ViT-B/16, sourced from the validation sets of ImageNet-1K K (3, ImageNet-R (13), and ImageNet-Skect Sketch (30), are used as unrelated samples outside the editing scope for the calculation of the locality rate (LR).
For Refine VIT, we fix the number of training epochs to 50 and use the Adam optimizer with learning rate of 0.00002. For simplicity, the hyperparameter /3 is set O. We evaluate the method using a E {10, 50, 100). For each value of a, we apply our approach to each sample in the benchmark-ttratingg as the sole available failed sample-and report the average performance across three evaluation metrics.
B.2 More Details in Section 6.2: Performance Enhancement and Debiasing Scenarios
We randomly select 10 samples per Slass--attrbbidd and landbirds for the Waterbirds dataset 25). and young and old celebrities for the CelebA dataset [19]-incllding both correctly and incorrectly predicted instances from CLIP-ViT. These samples are categorized into four groups based on their ground-truth and predicted labels. We then apply the four utility scores described in Appendix A skipping any scores that are infeasible due insufficient data in the comparison set. the comparison reference set lacks sufficient data, we fall back to zero ablation rather than performing mean ablation.
Next, for each valid utility score, we obtain the top T 15 attention heads, candidate list. Then we compute the validation utility score defined in Eq. 18 repeat for all candidate lists and select the attention head set with largest utility validation score the final list.
For our method enhanced with validation set, we perform ablation on the top t (t = 1, 2..,, T) heads in each ordered candidate list and select the one that delivers the best performance the validation set, rather than relying on the validation utility score estimation.
Both our method and Tip-Adapter are evaluated in zero-shot setting, leveraging CLIP-VIT" zero- shot capabilities without training additional classifiers. We show the results in Table D and Table
2
Stability Analysis To evaluate the stability of Refine VIT, we experiments the Binary Water- birds dataset with n = 10, 20, and 30 samples per class, using four random seeds for each setting. 33 the mean and standard deviation of accuracies. Rif consistently demonstrates robust performance improvements across varying sample sizes and different random initializations, indicating its reliability and stability in data-scarce scenarios.
B.3 More Details in Section 6.2: Model Editing Scenarios
Datasets and Settings ImageNet-R typically classify with high accuracy. We select the 18 classes with the most cartoon-style images
K contains a diverse set of real-world images that CLIP-VIT
22

---

Table Seerage-group accuracy (%) and worst- worst-group accuracy (%) for classifying 'Young' or 'Old' on the CelebA dataset using CLIP-ViT-B/LE
Method
Avg. (f)
Wst. (1)
Base
70.1
43.8
Tip-Adapter
73.7
52.2
Ours
74.1
56.4
Table 3: Stability Analysis Results Across Varying Sample Sizes on the Binary Waterbirds Dataset
ViT-B/16
ViT-L/14
VIT-H114
Method
Avg. (1)
Wst.(f)
Avg. (t)
Wst. (t)
Avg.(t)
Wst. (1)
Base
72.8
45.6
75.5
47.7
68.6
37.2
Ours 10)
79.8 1.2
59.79
85.5
71.9 0.8
74.5 1.0
46.5 3.3
Ours (n = 20)
80.8 + 0.2
68.5 2.3
0.4
1.1
0.6
48.1 + 1.0
Ours (n = 30)
80.4 + 0.2
2.3
8.4
2.1
73.5
47.7 + 3.2
from ImageNet-R and combine them with the Binary Waterbirds dataset [25] resulting in new dataset with 20 classes. Our objective is to improve performance on the bird classes while preserving locality- minimiini unintended side effects on the ImageNet-R classes using only 10 samples per bird class.
Similarly, we also evaluate ViT on ImageNet-A [14, dataset of real-world images that are commonly misclassified by ResNet models. We select the 10 most populous classes and target the two worst-performing ones for improvement, while minimizing degradation on the remaining eight. This is done using only 4 samples from each of the two target classes.
Implementation Detail All experiments use the Adam optimizer with a fixed learning rate of 2 x 10-5 We set both hyperparameters a and / to 1000. For the standard fine-tuning baseline, freeze the ViT backbone and train only an appended projection matrix-mirroiie the architecture of our proposed Refine ViT. All experiments are repeated with three random seeds reduce variance due to random initialization.
C Experimental validation of attention heads identified
To validate whether Refine can effectively identify for prediction failures in the first stage, test four mean-ablation-basee utility scores the Binary Waterbirds Dataset 255 This dataset combines thousands of waterbird and landbird images from the CUB dataset 229] with water or land backgrounds from the Places dataset [36]. Since the classification task focuses on bird type, the background introduces a significant known spurious correlation that can lead prediction failures.
features responsible
The core idea of our evaluation is that, while each utility score captures different aspects of attention heads' behavior in prediction failures, we design them to specifically identify heads that focus spurious background cues. This is achieved by selecting samples in which the background expected to be most prominent source of spurious correlation. By demonstrating that the attention regions identified by these methods are largely consistent, we validate that the scores behave as intended. This consistency is expected, as background-assoede associated attention regions fixed and should be detectable by all score variants.
Specifically, for U-A.A and ULA", we use images where the background aligns with the bird type in correctly predicted samples and mismatches it in misclassified ones. For Ul.b and U.A to ensure that the background consistently contributes to prediction errors, we include only samples with mismatched backgrounds in both sets. An example of this selection is shown in Fig. 6
23

---

We select 5 samples in each set, resulting in total of 30 samples, with different correctly predicted samples used across the groups. We evaluate our method on the CLIP-ViT-L/IE model [T8] and select the top T = 15 attention heads that induce the largest expected model shift for each utility score. The selected attention heads for each score are reported in Table while the jointly selected attention heads and their corresponding T 5] textual descriptions are presented in Table
The results, presented in Table show that 0 and B as well as Ut and UD produce identical outcomes. This is expected, as each pair compares the same sets of samples and based on the same underlying principle, leading to approximately the same results.
Notably, 8 of the 15 selected attention heads are shared across four utility scores. We present these shared attention heads along with their corresponding nextSpan-generate descriptions in Table 5 Furthermore, we use Grad-CAM 27 to visualize the regions these attention heads focus on, as shown in Figure 7
The TextSpan- -erarrred textual descriptions and Grad-CAM visualizations consistently show that the shared attention heads predominantly focus on background features. This confirms that the utility scores in RefineViT-ssage -ta-stage--n effectively identify attention heads associated with prediction failures, thereby addressing Q3.
Table 4: Attention heads identified by each utility score. Those jointly selected attention heads are highlighted in bold. (L22,HO) represents the (0+ 1)-th head in the (22+1)-th layer.
Utility A
Utility B
Utility C
Utility D
(I.23, H2)
(L.23, H2)
(L.23, H2)
(I.23, H2)
(I.23, H5)
(I.23, H5)
(L223, H6
(I.23, H6)
(L.22, H6)
(L.22, H6)
(L.22, H1)
(L.22, H1)
(L22, H4)
(L22, H4)
(L23, H5)
(L23, H5)
(L23, H14)
(L23, H14)
(L23, H8)
(L23, H8)
(L23, H3)
(L23,
(L23, HO)
(L23, HO)
(L22, H2)
(L22, H2)
(L22, H5)
(L22, H5)
(L21, HO)
(L21, HO)
(L23, H3)
(L23, H3)
(L23, H12)
(L23, H12)
(L23, H9)
(L23, H9)
(L23, H8)
(L23, H8)
(L23, HI)
(L23, H1)
(L.22, H12)
(L.22, H12)
(L,21, H9)
(L.21, H9)
(I.23, H9)
(L.23,
(I.22, H9)
(L.22, H9)
(L.22, H1)
(L22,, H1)
(L23, H12)
(L223, H12)
(L23, H6)
(L233 H6)
(L20, H10)
(L20, H10)
(L21, H15)
(L21, H15) H15)
(L22, HO)
(L22, HO)
D Ablation Study and Sensitivity Analysis
D.1
Sensitivity Analysis of hyper-parameter T in Stage One
For model editing tasks with multiple available samples, the hyper-parameter T' in Refine stage one does not directly determine the number of attention heads to be ablated; rather, it defines the size of the pool from which we select attention heads based on their utility scores. Consequently. the performance improvement from ablation tends to plateau once T is sufficiently large. To support this claim, we conduct a sensitivity analysis on the Waterbirds dataset (25) using both VVT-B-1 and CLIPP ViT-L/4. Each experiment is repeated with three random seeds to reduce variability. As shown in Fig. g performance gains stabilize when T exceeds 10.
D.2 Ablation Study of Attention Heads Identified in Stage One
we conduct an ablation study to evaluate the contribution of the whole stage one in Refine VIT. Specifically, we remove the influence of stage one by treating all available samples equally during stage two. In this setting, the model is trained to minimize the MSE with the initial model across
24

---

Original Image
Initial Model
Identified by RefineViT
Identified TextSpan
Figure 7: Grad-Cam visualization. We present four examples from the Waterbirds dataset, each illustrating, from left to right: (1) the original image, (2) a heatmap showing the focus of the initial model, (3) heatmap highlighting the attention heads identified by Refine ViT correlated with error sources based on all four utility scores, and (4) a heatmap highlighting the attention heads selected by TextSpan related to the background. In the task of distinguishing waterbirds from landbirds, domain knowledge suggests that bird claws and beaks are causal features, whereas the background often introduces spurious correlations that lead prediction errors. As shown in the examples, the attention heads identified by ViT predominantly focus on background regions, Moreover, compared TextSpan, Refine ViT selects attention heads that exhibit stronger emphasis the background and a weaker focus on causal features, despite TextSpan incorporating prior domain knowledge and requiring substantially more human effort and computational resources.
all available samples, rather than partitioning the samples into two groups one aligned with the output of the ablated model and the other with that of the initial model. This variant is referred to as Refine (ablated). As shown Fig. O the results demonstrate that stage one, where attention heads responsible for errors are identified, effectively guides the error correction process for target classes. This guidance leads to smoother and more robust learning curves across different random seeds.
D.3 Ablation Study of Csuccess (0) and Liocality (0)
In accordance with Eq. 9 10 and 12 the fine-tuning loss function is defined as combination of the edit success loss, Lonccces. (0), the edit locality loss, Liocatiy (0), and the cross-entropy loss, (en(
To evaluate the contribution of the first two components to edit success and locality, we perform
25

---

Table 5: Common attention head with their top-5 results of TextSpan.
Layer 23, Head 5
Layer 23, Head 3
Intertwined tree branches
Bustling city square
Flowing water bodies
Serene park setting
A meadow
Warm and cozy indoor scene
smoky plume
Modern airport terminal
Blossoming springtime blooms
Remote hilltop hut
Layer 23, Head 8
Layer 23, Head 6
Photograph with red color palette
Picture taken in Sumatra
An image with cold green tones Timeless black and white Image with a yellow color Photograph with blue color palette
Picture taken in Alberta, Canada
Picture taken in the geographical location of Spain Image taken in New England Photo captured in the Arizona desert
Layer 23, Head
Layer 23, Head 12
Image showing prairie grouse
Image with polka dot patterns Striped design
Image with penguin magnolia
Checkered design
An image with dogs
Artwork with pointillism technique
image with cats
Photo taken Galapagos Islands
Layer 23, Head
Layer 22, Head
ornate cathedral
A semicircular arch
detailed reptile close-up Image with seagull A clover
An isosceles triangle
An oval
Rectangular object
Futuristic space exploration
A sphere
Avg. (CLIP-VII-/16)
(CLIP-ViT-B/1))
Ave.Ace (CLIP-ViT-L/14)
Ace (CLI-VVT--U1/4)
80
Acemncy (9))
30
60.0
60
50
Figure 8: Sensitivity analysis for T
ablation study on the 'Waterbirds Imagenet-R' dataset by isolating each loss term. For clarity, the weight of the cross-entropy loss is fixed at 1, while the weights of Luucesss (0) and Liocatit (0) are donated as a and B3 respectively. All of our experiments are based on CLTP-VIT-L/14
The ablation study for Ciecality (0) is performed by fixing the weights of LCE(O) and Laccess (0) to
and 0, respectively, while varying the weight 3 of Lucaaity (0) over the set {0, 10, 100, 1000). Each hyperparameter configuration is evaluated using 20 samples selected under three fixed random seeds.
As shown in Fig 10 increasing 3 generally leads to improvements in the average accuracies of both unrelated and target classes. Moreover, variance across the three seeds diminishes with larger 3 values, suggesting more stable and robust training process. These results indicate that an appropriately chosen 3 serves as an effective regularizer, guiding the model toward more consistent and generalizable representations.
26

---

Avg. Acc. (target classes)
Avg. Ace. (other classes)
82.5
92
(%)
80.0
91
77.5
75.0
72.5
89
20
40
80
100
20
40
80
100
epoch
epoch
RefineViT (ablated)
RefineViT
Figure 9: Performance comparison for the ablation study. The hyperparameters of our method are set a = 3 1000. For fair comparison, the weight of the MSE loss RefineVTT (ablated) is also
set to 1000.
Avg. Acc. (target classes)
Avg. (och(other classcs)
85
Accuracy
40
75
20
70
0
20
40
60
80
100
0
20
40
60
80
100
epoch
epoch
a = 0, B = 0
0 = 0, 8 = 100
a
B 10
a 0, B = 1000
Figure 10: Ablation study for Lincality (0).
Similarly, perform ablation study Lsuccess (0) by varying its weight over the set (0, 10, 100, 1000), while keeping the weights (fn(f) and Liucality (0) fixed at and 0, respec- tively. As shown in Fig. I1 increasing yields notable gains in the average accuracy of target classes, accompanied by reduced variance across the three random seeds. In addition, the accuracy of unrelated classes also improves. This suggests that, although the success loss Luxcess (8) may conflict with the original locality objective Liscality (0), it helps mitigate overfitting, providing more generalizable alternative to standard fine-tuning.
D.4 Sensitivity Analysis of and
Based on the ablation study of Lsuccess (8) and Lieeatyy (f) conducted on the WWaterbirds-ImagNeee.k dataset using CLIP-ViT-L/14, we observe that when the weights for Luscess (0) and Ciocality (0) are set to approximately 103, our method achieves optimal performance in terms of both edit success and edit locality. assess whether this weight range consistently yields stable performance across different datasets on CLIP-ViT-L/14, we performed a sensitivity analysis of the coefficients and B on Waterbbrdss+ImageNee-- and 'ImageNerAA' Specifically, we fixed the training epochs to 80 for "Waerbirdss-ImaeeNeee and 20 for 'ImageNet-A'. We used 3 fixed random seeds and computed
27

---

Avg. Acc. (target classes)
Avg. (other classes)
85
Apcuracy (%)
60
40
75
20
70
0
25
50 epoch
75
100
0
25
50 epoch
75
100
a = 0, ==0
0 = 100,
8 0
a 10, B=0
a 1000, ==
Figure 11: Ablation study for Lanceess (0).
Avg.Acc, (largel classes)
Avg.Ave. (other classes)
100
8
1000
1000
2000
3000
3000
Figure 12: Sensitivity analysis 'Waterbirds ImageNet-R"
the average accuracy (in percentage) on the target classes and the unseen unrelated classes for each combination of a and B.
As illustrated in Figures 12 and 13 observe a consistent improvement in accuracy on both the
target and unrelated classes as a and 3 increase. Notably, the accuracy on the target classes reaches its peak when is approximately 1000 across both datasets. Furthermore, between Cuucess (9) and Liocalit (0). Specifically, when the coefficient a for Lsucesss (0) exceeds 1000, increasing the coefficient for Liocality (0) results in degradation of accuracy on the target classes. This phenomenon observed in both datasets.
results reveal trade-off
D.5 Qualitative Analysis of Individual Utility Score Effectiveness
In all experiments, we observe that when selecting the final list of ablated attention heads based on their performance available samples, Utility C and Utility D generally outperform Utility A and Utility in identifying heads for ablation. Notably, on the CelebA [19] dataset, ablating only the heads selected by Utility A and Utility B can significantly degrade performance, often leading the model to predict all samples as belonging to a single class. This failure arises from the implicit assumption behind Utility A and Utility B that the model is capable of fair prediction without systemic bias.
In practice, however, models often rely heavily on narrow set of features and can exhibit strong bias toward dominant classes. For instance, when predicting whether a celebrity young or old in CelebA [19), the ViT-B/6 model tends predict most individuals as "young". Although
28

---

Avg.Acc, (largel classes)
Ave.Ace. (other classes)
100
8
1000
34
1000
2000
56
3000
3000
-1
3000
Figure 13: Sensitivity analysis on ImageNet-A
this yields high accuracy for genuinely young samples, it results in substantially lower accuracy for older individuals. A plausible explanation is that the model relies on simplistic visual cues-such as the presence of wrinkles-to to make predictions: if wrinkles are detected, the model predicts "old"; otherwise, "young." " Although this heuristic may be partially effective, it is insufficient, as not all older individuals exhibit visible wrinkles.
Utility A and Utility B tend to identify attention heads associated with features that are important but exhibit inconsistent behavior across correctly and incorrectly classified samples. In the case of CelebA, this can lead to the unintended selection of heads tied to wrinkle-related features, which are mistakenly flagged as harmful due to their differing influence groups. As result, ablating these heads removes features that are essential but not always correct, which further amplifies the model's bias and ultimately leads it to predict a single class for all inputs.
the model's 's predictions across these
In contrast, Utility C and Utility D aim identify attention heads that negatively impact both correctly and incorrectly classified samples. This leads to more conservative and balanced modifications, making them more effective challenging cases like CelebA.
D.6 Update Strategies
In Section 3.2 we demonstrate that spurious correlations are key factor contributing to prediction failures computer vision tasks. To mitigate the influence of misleading features the learned representations, ViT employs trainable projection matrix after the transformer blocks as an update strategy, rather than directly fine-tuning the later MSA modules, which may distort extracted features.
To further validate the effectiveness of this design choice, we compare it with LoRA, widely used update strategy for MSA modules. Specifically, we apply LoRA to the last two MSA layers during the second stage of Refin as these later layers have been shown to play crucial role. Due to its significantly higher computational cost compared to the original Refine this experiment is conducted only on one group from the Natural ViT Benchmark [32
The rank in LoRA serves as a hyperparameter that balances fine-tuning capacity and the risk of overfitting by controlling the number of trainable parameters. In hyperparameter search, find that values of E [0. 1] generally yield better performance. Accordingly, for each rank, we evaluate LoRA with a e {0, 0.1, 0.5) and compare the results to the original Refine using the projection matrix. As shown in Figure [4] using LoRA in the last two MSA layers results in substantially worse performance, despite requiring much greater computational resources. These observations further support the effectiveness of our proposed update strategy.
D.7 Extension to Segmentation Task
To examine the adaptability of our method beyond image classification, we conduct preliminary experiment CLIP-based segmentation task following the setup in [5]. In contrast classification,
29

---

Group 890-430
100
95
90
GR (%)
85
80
RefineVIT (Ours)
75
RefineVIT (LoRA: Rank = 1)
RefineVIT (LoRA:
70
70
75
80
85
90
95
100
LR (%)
Figure 14: Comparison of update strategies: LoRA vs. projection matrix
segmentation requires the model to process all patch tokens to predict pixel-level outputs, which introduces additional challenges for attribution and decomposition during model editing. Using single edit sample, apply our method to the segmentation model and evaluate the resulting performance in terms of pixel-wise accuracy.
The edited model achieves higher accuracy on images similar to the edit sample and a slight improve- ment in overall performance, as summarized below:
Model
Similar Data
Overall
Original CLIP-ViT
72.29%
76.78% 77.01 %
CLiP-ViT-Edited (Ours)
%
These results suggest that our framework can be adapted to more complex vision tasks such as segmentation, showing its potential generality beyond classification.
30


Reviews:
[{"review_id": "XFyW5VNRAj", "reviewer": "Reviewer_TBhL", "summary": "This paper investigates the problem of model editing for Vision Transformers (ViTs), a largely underexplored topic compared to model editing in language models (LMs). While prior work on LMs often targets modifications in multi-layer perceptron (MLP) modules, the authors observe that ViT predictions are more influenced by multi-head self-attention (MSA) than by MLPs. Motivated by this insight, they propose RefineViT, a two-stage framework for editing ViTs. In Stage 1, RefineViT identifies spurious features by locating attention heads most responsible for misclassifications. In Stage 2, it corrects predictions by either ablating these attention heads (in simple cases) or learning a trainable projection matrix to adjust the image representation while preserving prediction locality. The method is validated across a broad range of benchmarks and ViT variants, including CLIP-ViT, demonstrating state-of-the-art performance in both edit success and edit locality.", "questions": "1. Can the head attribution process be made more robust to ambiguous or adversarial examples? The current process assumes that the utility scores effectively highlight spurious heads. Are there adversarial cases where this could be misleading?\n\n2. How would the method generalize to multi-modal architectures like Flamingo or BLIP-2? Since RefineViT relies on decomposing image representations, could it be adapted to edit vision-language embeddings in multi-modal transformers?\n\n3. How is performance affected by changing the number of attention heads ablated (T)? An ablation study on the sensitivity of the method to T (beyond the top-k selection) would be helpful, especially in low-resource or overfitting-prone regimes.\n\n4. What are the limitations of applying the projection matrix universally across inputs? Since P(θ) is learned from a small set of error instances, could it harm performance on rare or edge cases not covered by the original edit scope?\n\n5. Can RefineViT serve as a debugging tool to discover spurious concepts? It seems that the utility functions reveal which features lead to error. Could this be repurposed for interpretability or dataset bias detection?\n\n6. Can the authors also report results on dense prediction tasks? Would the method generalize? E.g., the authors can study object detection or segmentation as a downstream task.", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\n\n1. The work is technically solid, with clear empirical grounding, rigorous derivations (e.g., utility functions for attention head attribution), and extensive experimental validation.\n\n2. The paper is well-organized and well-written. The motivation is compelling, the diagrams (e.g., Fig. 1) are intuitive, and the theoretical analyses are easy to follow.\n\n3.  Tackling model editing in ViTs is novel and timely, especially as ViTs become prevalent in real-world systems where fast and local error correction is valuable. The method has wide implications for robust deployment and interpretability of vision models.\n\n4. The insight that editing MSA is more effective than MLPs in ViTs is novel. The formulation of utility-based attention head attribution and subsequent representation rectification via a projection matrix is both elegant and original.\n\nWeaknesses\n\n1. Reliance on utility-based heuristics: The head attribution stage relies on heuristics and utility scores that may be brittle in edge cases or less interpretable across tasks. A discussion on potential failure cases of this heuristic selection would be beneficial.\n\n2. Assumes access to small labeled samples for editing: Although framed as data-efficient, the method still requires at least one labeled misclassified instance per error type, which may be non-trivial in open-set or online deployment settings.\n\n3. Evaluation scope: The experiments are comprehensive, but limited to classification. It remains unclear how well this method generalizes to dense prediction tasks such as object detection or segmentation.", "was_revised": true, "final_justification": "Based on the questions raised by the reviewers and the feedback from authors I am raising my score to 4.", "rebuttal": "We sincerely thank the reviewer for the positive feedback, particularly recognizing the technical soundness, clear presentation, and the novelty and timeliness of our ViT model editing approach. Below, we address all questions and weaknesses, supplemented with new supporting experiments.\n\n**Q1 & W1:** Given its reliance on utility-based heuristics. **(1)** Are there adversarial cases where this could be misleading? **(2)** Can the head attribution process be made more robust? **(3)** How interpretable?\n\n> **A:**  Point‑by‑point response:\n> **(1) Potential misleading cases.** Our method is inherently data-driven and thus sensitive to the quality of the editing sample, especially in the one-sample setting. Ambiguous or mislabeled samples can reduce the reliability of head attribution and the resulting edits.\n>\n>  To illustrate this, we conducted a new Waterbirds experiment comparing edits using a correct (“real”) misclassified sample versus an incorrect (“fake”) one. Editing with a real sample improves both average and worst-group accuracy, whereas a fake sample significantly degrades performance:\n> \n> |                   | Avg. Acc (%) | Wst. Acc (%)  |\n> | ----------------- | ------------ | ------------- |\n> | Original CLIP-ViT | 92.06        | 45.71         |\n> | One 'Real' sample     | 92.8 ± 1.18  | 53.9 ± 3.10   |\n> | One 'Fake' sample     | 78.53 ± 3.52 | 14.04 ± 12.36 |\n> \n> This demonstrates the **expected limitations under weak supervision** while confirming that the method behaves sensibly—**improving with informative signals and degrading with misleading ones.**\n \n> **(2) Robustness.** In scenarios where multiple samples are available, we mitigate brittleness **by leveraging four complementary utility scores** (Utility A–D, Appendix A, p.13) capturing different aspects of attention behavior. In the Waterbirds experiment (Appendix C, p.16), Grad-CAM visualizations and TextSpan-generated descriptions validate that the selected heads align with the intended utility functions.\n\n> **(3) Interpretability.** On Waterbirds, which contains strong spurious correlations (water/land backgrounds for bird-species classification), our utility scores consistently identify heads focusing on these background cues, as confirmed by TextSpan descriptions (Tables 4 & 5) and Grad-CAM visualizations (Fig. 7). These analyses demonstrate that our utility-based heuristics are both effective and interpretable.\n\n\n\n**W2.** The method requires at least one labeled misclassified instance per error type, which may be non-trivial in open-set or online deployment settings.\n\n>**A:** We acknowledge the need for at least one labeled misclassified instance to guide editing.\n>- However, **this is a minimal data requirement in the field of model editing.** As far as we know, there are no zero-shot model editing methods. Our method achieves effective editing using just a single such instance. \n>- In practice, model errors are typically discovered through observed prediction failures, making such instances naturally available. Hence, this requirement is reasonable and **aligns with real‑world deployment scenarios.**\n\n**W3. & Q6.** Can the method generalize to tense prediction tasks such as object detection or segmentation?\n\n> **A:** Thank you for the suggestion. **Yes, it can**.To assess generalization beyond classification, we follow the setup in [4] and apply our method to a **segmentation task** using a single edit sample.\n>- As shown in the table below, our method improves pixel-wise accuracy on similar images and even slightly boosts overall performance on the full dataset:\n> > \n> >| Avg. Pixel-Wise Accuracy | Similar Data | Overall |\n> >| ------------------------ | ------------ | ------- |\n> >| Original CLIP-ViT        | 72.29%       | 76.78%  |\n> >| CLIP-ViT-Edited (Ours)   | 74.50%       | 77.01%  |\n> \n> - To our knowledge, there are currently no established model editing methods for segmentation. This result serves as an initial demonstration of our method in this setting. In future work, we plan to build a benchmark for model editing in segmentation and explore more effective methods tailored to dense prediction tasks.\n\n**Q2.** Can the proposed method be generalized to multi-modal architectures like Flamingo or BLIP-2?\n\n>**A:** Our method focuses on editing visual transformers with standalone visual encoders. For multi-modal models like CLIP-ViT, which include an independent visual encoder, our method can be directly applied to the visual component.\n\n>In contrast, models such as BLIP-2 and Flamingo deeply fuse vision and language features via cross-attention, and their visual modules are typically frozen during training. Editing these tightly integrated architectures aligns more closely with large language model (LLM) editing than with visual transformer editing.\n\n>Therefore, adapting our method to these models is **non-trivial** and represents an interesting direction for future research.\n\n**Q3.** How is performance affected by changing the number of attention heads ablated?\n\n\n>**A:** We evaluated the effect of ablating different numbers of attention heads using the ViT Natural Image benchmark (class 954_582). As shown below, GR (target-class performance) **improves** as more heads are ablated, **peaks** around 6 heads, and then **slightly declines**, showing that ablating the right heads matters more than simply ablating more heads:\n>   \n>   | N_ablated | 3   | 6   | 9   | 12  |\n>   | --- | --- | --- | --- | --- |\n>   | GR(%) | 73.78±13.08 | 86.67±2.30 | 84.89±2.03 | 84.44±2.78 |\n>   \n> Importantly, the number of heads ablated is **not a fixed hyperparameter** in our method.\n>   \n>   - In single-sample settings, $T$ refers to the number of heads actually ablated. We select heads whose utility scores exceed a fixed threshold.\n>   - In multi-sample settings (Appendix A, Page 13), **$T$ denotes the size of the candidate pool**, i.e., the number of top-ranked heads selected per utility score. The final ablation set is determined adaptively by applying a validation utility (as described in Appendix, Page 14, “Summary of the Procedure”), so the number ablated is typically smaller than $T$.\n>\n>Furthermore, the **sensitivity analysis** in Figure 8 (Appendix) shows that performance stabilizes when $T ≥ 10$ in multi-sample settings, indicating that beyond this candidate pool size, increasing $T$ has minimal effect as the most impactful heads are already included.\n\n\n**Q4.** The limitations of applying the projection matrix universally? Could it harm performance?\n\n\n>**A:** Applying a projection matrix universally may risk unintended changes, but our method is **specifically designed to minimize such effects while maintaining strong edit performance.**\n\n>- In model editing, three primary goals are typically considered: (1) edit success, (2) generalization, and (3) **locality**—the latter reflecting how well the edit preserves performance on unrelated data. In practice, **edits often involve trade-offs**, and to our knowledge, no method guarantees zero degradation on out-of-scope data.\n>- To address this, we introduce a **locality loss** to reduce such side effects explicitly. We report Locality Rate (LR) in Figures 4 and 5, showing that our method achieves a more favorable balance between edit success and preservation of unrelated behavior than prior approaches.\n>- Additionally, to assess the impact of using a projection matrix as the update mechanism, we compare it to LoRA, a commonly used alternative. As shown in Figure 14 (Appendix, Page 22), our projection-based strategy preserves both effectiveness and locality better than the LoRA-version baseline.\n\n\n**Q5.** Can RefineViT serve as a debugging tool to discover spurious concepts?\n\n>**A:** Yes, it can. As shown in Appendix C (Page 16), we evaluate our method on the Waterbirds dataset, which contains strong spurious correlations between bird type and background. Tables 4 and 5 list the attention heads identified by our utility scores, along with descriptions from TextSpan. Grad-CAM visualizations in Figure 7 further confirm that these heads attend to background regions, indicating our method’s ability to identify spurious features driving incorrect predictions."}, {"review_id": "B7DRhqwyIk", "reviewer": "Reviewer_qQpH", "summary": "This paper introduces a two-stage framework, in the context of model editing for Vision Transformers (ViTs), to correct prediction errors without requiring expensive retraining. The proposed method is called RefineViT. The authors first present an empirical analysis arguing that, unlike in language models where MLP layers are the primary target for editing, the Multi-Head Self-Attention modules are the dominant drivers of predictions in ViTs. Based on this observation, the first stage of RefineViT identifies specific attention heads responsible for a given misclassification by ablating their contributions and measuring the impact on the model's output logits. The second stage, termed \"representation rectification,\" learns a lightweight projection matrix that modifies the ViT's final image representation. This matrix is trained to alter the representation of erroneous samples towards their corrected versions, while preserving the representations of correctly classified samples, thus balancing edit success with locality. The method is shown to achieve state-of-the-art performance on a ViT editing benchmark and demonstrates effectiveness in debiasing tasks for CLIP-ViT models.", "questions": "q1. The L_locality term is a central part of your proposed method for balancing success and locality. Yet, for the main results in Figure 4, it was disabled ($\\beta=0$). Please provide a clear justification for this experimental choice. \n\nq2. Your identification stage relies on a linear decomposition of the ViT representation (Eq. 4). Could you please discuss the potential impact of this simplification? Specifically, can you characterize scenarios where this approximation might fail, leading to the misidentification of faulty heads? Acknowledging this would strengthen the paper's scientific integrity.\n\nq3. Does this method generalize to scenarios where multiplex edits are required?\n\n\nI am open to increasing the score provided that the issues are reasonably addressed.", "limitations": "The authors have discussed the limitations, but it is expected to be expanded, for example, the discussion of core methodological assumptions should be in the main text.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n\n1. This paper finds that ViT editing should target attention modules instead of MLP modules. This is an interesting conceptual contribution, which challenges the direct transfer of paradigms from the language model editing literature and instead proposes a solution tailored to the specific architectural properties of ViTs. I think this could guide future studies, upon further verification.\n2. The two-stage approach is reasonable and is connected to the identified gap. The use of a projection matrix for rectification helps to avoid invasive changes to the model's backbone and is computationally efficient, making it practical for targeted applications.\n3. The method shows substantial improvements over existing methods on a benchmark. The hypothesis is further supported by ablation studies.\n\n---\n\n### Weaknesses\n\n1. The entire method is built on a simplified linear decomposition of the ViT's representation (Eq. 4), which treats attention head contributions as independent and additive. While I understand that the linearity assumptions are widely adopted in the model editing literature, in my opinion, this is a significant simplification of the true, non-linear, and even sequential nature of the Transformer architecture. The paper does not acknowledge or analyze the limitations of this core assumption.\n2. There is a discrepancy between the described method and the main experiments. The `L_locality` term is a key component of the proposed objective function (Eq. 12). But it was disabled (setting $\\beta=0$) for the main benchmark results (according to Appendix B.1). Doesn’t this make the locality preservation effort invalid?\n3. There is limited discussion of when the proposed method might fail or its assumptions might be violated.", "was_revised": true, "final_justification": "The authors provide more details to justify some of their design choices, which triggered my initial concerns. Based on this, I think these issues can be considered addressed , though these design choice are still the results of compromise to some extent. Thus I increase my score to 4 but not an even higher one.", "rebuttal": "We sincerely thank the reviewer for recognizing our conceptual contribution, practical two‑stage design, and strong empirical results supported by ablations. Below, we address all questions and weaknesses, supplemented with new experiments for further support.\n\n**W1 & Q2.** The method relies on a linear decomposition of ViT representations. What are the limitations of this assumption, and in what cases might it lead to misidentification of faulty heads?\n>**A:** We acknowledge this simplification and discuss it in the Limitations section (page 23). Our decomposition focuses on the first‑order (direct) contributions of attention heads to the final representation, omitting higher‑order interactions across subsequent layers. While this does not fully capture the non‑linear and sequential dynamics of transformers, it provides a **tractable and effective** editing process.\n>\n\n> We understand the reviewer’s concern that some heads may contribute to incorrect predictions primarily through higher‑order effects, or vice versa. However, our method modifies the final representation by removing only the direct (first‑order) contributions of identified heads, leaving their indirect effects intact. Therefore, the approach is generally robust to potential misidentification caused by omitting higher‑order effects.\n> \n\n> Investigating higher‑order contributions is a promising future direction for achieving **more fine‑grained and accurate interventions**.\n\n\n**W2 & Q1.** The L_locality term is a central part of the method for balancing success and locality. Why was it disabled for main benchmark in Figure 4?\n\n>**A:** We clarify that $L_{\\text{locality}}$ is indeed a key component for balancing edit success and locality. However, in the experiments shown in Figure 4, only one misclassified sample is available, and no correctly predicted samples are present. Under such conditions, the standard form of $L_{locality}$ (Eq. 10) is not applicable.\n>\n\n>Although an alternative regularization form (Eq. 11) could be used to reduce overfitting, we observed that our method already maintains strong locality without it. This aligns with our sensitivity analyses (Figures 10–13, Appendix), which shows that the success loss $L_{\\text{success}}$ alone implicitly helps mitigate overfitting. Moreover, the small learning rates and few training epochs used in these experiments also help limit overfitting.\n>\n\n>In practice, we find the optimal $\\beta$ is significantly smaller than $\\alpha$ in these specific experiments. Given the minor performance gap and for simplicity, we set $\\beta = 0$ in these experiments.\n\n**W3.** Limited discussion of when the proposed method might fail or its assumptions might be violated.\n\n>**A:** First, as we acknowledged in the Limitations section (Page 23), our method models **only the first‑order**, direct contributions of individual heads to the final representation and does not capture higher‑order or indirect effects; in scenarios where higher-order effects dominate, the proposed method cannot repair the corresponding errors. \n> \n\n> Second, in the single‑sample editing setting, our method depends heavily on the quality of the provided misclassified sample. If the sample is ambiguous or mislabeled, the effectiveness of the edit will degrade.\n> \n\n> We will include these discussions in the Limitation section.\n\n**Q3.** Does this method generalize to scenarios where multiplex edits are required?\n\n>**A:** Yes, our method supports multiplex editing and shows strong performance in such scenarios.\n>> - Evaluation setup: We test on the ViT benchmark by sequentially editing three target classes, each with a single misclassified sample.\n>>     - We compare our approach against the state-of-the-art method (WTE) and standard fine-tuning with $l_{2}$ regularization.\n>>     - At each step, we measure both Generalization Rate (GR), i.e. accuracy on target classes, and Locality Rate (LR), i.e. performance on unrelated data. **Importantly, we track GR for both the current (e.g., CLS_3 for the 3rd Edit 3) and all previously edited target classes (e.g., CLS_1 and CLS_2) to assess edit retention.** The results are shown in the following tables.\n>> \n>> |Generalization Rate (GR)|Edit 1|Edit 2 |Edit 2 | Edit 3 | Edit 3| Edit 3|\n>> |--|----|---|--|-----|---|--|\n>> ||CLS_1|CLS_1|CLS_2|CLS_1|CLS_2|CLS_3|\n>> |RefineViT(ours)| 92.0| 92.0| 98.54|90.67|98.54|96.1|\n>> |WTE|73.61| 73.61| 84.21 |73.61| 84.21 |96.54\n>> |Std.F.T|38.89| 0| 37.84 |5.55| 5.26 |98.77\n>>\n>> |Locality Rate (LR)|Edit 1|Edit 2|Edit 3|\n>> |--|----|---|--|\n>> |RefineViT(ours)|91.79|88.12|84.69|\n>> |WTE|90.48 |86.52 |85.176 |\n>> |Std.F.T|95.99| 83.84| 83.88 \n> \n> - Result: RefineViT **consistently** maintains higher GR and LR than WTE and standard fine-tuning, indicating it preserves previous edits while maintaining locality, which demonstrates strong effectiveness in multiplex editing scenarios."}, {"review_id": "1qG1LTF4EU", "reviewer": "Reviewer_x93k", "summary": "This paper proposes a new method for editing Vision Transformers in classification tasks by localizing the self-attention heads most responsible for incorrect predictions. The authors then introduce a technique to neutralize the effect of these heads by applying a linear transformation to the embeddings—designed to minimally affect unrelated samples while producing outputs that approximate those obtained when the localized heads are ablated.", "questions": "Please read weaknesses.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\nThe paper is well written, and the idea of editing the model by localizing important MSA layers is clearly presented. The experiments are extensive, demonstrating improvements over strong baselines and highlighting the effectiveness of the proposed method. Figure 5 is particularly insightful, showcasing that a lightweight projection matrix applied to localized MSAs can be highly impactful.\n\nWeaknesses:\nFigures 3 and 4 could benefit from clearer and more informative visualizations. Additionally, it would strengthen the analysis to include a comparison where a random subset of MSA layers—rather than the localized ones—is used to optimize the projection matrix. This would help assess whether the identified layers are truly critical. Furthermore, the decomposition of residual tokens into contributions from MSA and MLP layers is not novel, as it has been explored in several prior works [1], limiting the novelty of this aspect.\n\n\n\n\n-----\n1. Balasubramanian, Sriram, Samyadeep Basu, and Soheil Feizi. \"Decomposing and interpreting image representations via text in vits beyond CLIP.\" arXiv preprint arXiv:2406.01583 (2024).", "was_revised": true, "final_justification": "I still think borderline accept could be a good score for this paper, as I see the application of the approach a bit limited, and also the core idea of localization explored in the literature.", "rebuttal": "We sincerely thank the reviewer for recognizing our clear presentation, extensive experiments, and the effectiveness of the lightweight projection on localized MSA layers. Below, we address all weaknesses, supplemented with new experiments for further support.\n\n**w1.** Figures 3 and 4 could benefit from clearer and more informative visualizations.\n>**A:** Thanks for your suggestion. We will revise them.\n\n**W2.** Include a comparison where a random subset of MSA layers is used to optimize the projection matrix?\n\n> **A:** Thank you for the suggestion. We conducted this comparison using the Waterbirds dataset and ViT-L-14. The results are summarized below:\n> \n> |GR |Epoch 0|Epoch 20|Epoch 40| Epoch 60|Epoch 80|\n> |--|----|---|--|-----|---|\n> |RefineViT| 72.64| 80.62±1.20| 82.82±1.25|83.35±1.41|83.58±1.25|\n> |Random Subset|72.64|79.41±1.72| 81.65±1.56|82.47±1.81|82.59±2.09|\n> \n> \n> |LR |Epoch 0|Epoch 20|Epoch 40| Epoch 60|Epoch 80|\n> |--|---|---|--|-----|---|\n> |RefineViT|92.65|91.43±0.25|91.36±0.27|90.96±0.54|90.18±0.89\n> |Random Subset|92.65|91.32±0.17|90.59±0.16|90.97±0.47|90.62±0.46\n> \n> We find that randomly selecting MSA layers for projection matrix optimization leads to consistently lower GR  across training epochs compared to our structured selection, underscoring the importance of the identified heads for effective model editing.\n\n\n> Furthermore, **Figure 9 (page 20) presents additional ablations on stage 1 of RefineViT,**  further demonstrating the value of our head selection mechanism. Taken together, these results confirm that our design choices are essential for achieving strong generalization while preserving locality.\n\n**W3.** The decomposition of residual tokens into contributions from MSA and MLP layers is not novel.\n\n>**A:** We clarify that the decomposition of ViT representations is **not presented as a novel contribution** of our work. It appears in the **Preliminaries section** with proper citations, serving only to establish the foundation for our method. Our main contribution lies in the editing framework and targeted refinement strategy, which go beyond these prior analyses.\n\n> We will also include a citation to the paper you mentioned in the related work."}, {"review_id": "mY3M1kayYa", "reviewer": "Reviewer_uxyq", "summary": "This paper proposes RefineViT, a model editing method tailored for Vision Transformers (ViTs).   Unlike prior approaches that adapt language model (LM) editing strategies by modifying MLP layers, the authors show that ViT predictions are more influenced by multi-head self-attention (MSA) modules.  RefineViT identifies problematic attention heads responsible for prediction errors and selectively ablates or projects their outputs to correct mistakes while preserving performance on unrelated data.   The method is validated on several image classification benchmarks, showing superior correction and locality compared to prior ViT editing approaches.", "questions": "1. How does the approach perform if many edits are required, e.g., in dynamic or continuously evolving datasets?  Are there cumulative effects or degradation as the number of edits increases?\n\n2. Can the proposed framework be adapted for more complex vision tasks such as object detection, instance/semantic segmentation, or video tasks?  What are the main challenges or required modifications?", "limitations": "See the weaknesses.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# Strengths\n\n1. The paper studies the practical problem of model editing for ViTs, which may be important as models are increasingly used in real-world applications where errors can occur after deployment.\n\n2. The main contribution is to target the MSA modules for editing, rather than the traditional approach of editing MLP layers.  This idea is supported by some experimental analysis suggesting that ViT predictions are more sensitive to attention modules.\n\n3. The method RefineViT consists of two stages: identifying which attention heads contribute to incorrect predictions and then adjusting or projecting their outputs to correct the errors.  This process seems relatively lightweight, as it does not require retraining the entire model.\n\n4. Experiments are carried out on several classification datasets, and results show improvement over some existing model editing techniques on these benchmarks.\n\n\n# Weaknesses\n\n1. All of the experiments and analyses are limited to image classification tasks. There is no experimental evidence for whether the approach works for other vision tasks such as object detection or segmentation.\n\n2. Most experiments are conducted on standard or synthetic classification datasets, such as ImageNet-R and ImageNet-A, which may not fully reflect large-scale or highly complex real-world scenarios.", "was_revised": true, "final_justification": "The author’s response addressed my concerns. I have no further questions and will keep my original score. I look forward to seeing the author extend model editing to more complex tasks in the future.", "rebuttal": "We sincerely thank the reviewer for recognizing the practical importance, MSA‑focused design, and strong experimental results of our RefineViT framework. Below, we address all weaknesses and questions, supplemented with new experiments for further support.\n\n**Q1.** Does this method generalize to scenarios where multiplex edits are required? dynamic and continuously? cumulative effects or degradation?\n> Yes, our method supports multiple editing and demonstrates strong performance in such scenarios.\n> - We evaluate this on the ViT benchmark by sequentially editing three target classes, each involving a single misclassified example. \n>     - We compare our approach against the state-of-the-art method (WTE) and standard fine-tuning with $l_{2}$ regularization.\n>     - At each step, we measure both Generalization Rate (GR), i.e., accuracy on target classes, and Locality Rate (LR), i.e., performance on unrelated data. Importantly, we track GR for both the current (e.g., CLS_3 for the Editing 3) and all previously edited target classes (e.g., CLS_1 and CLS_2) to assess edit retention. The results are shown in the following tables.\n> \n> |GR|Edit 1|Edit 2 |Edit 2 | Edit 3 | Edit 3| Edit 3|\n> |--|----|---|--|-----|---|--|\n> ||CLS_1|CLS_1|CLS_2|CLS_1|CLS_2|CLS_3|\n> |RefineViT(ours)| 92.0| 92.0| 98.54|90.67|98.54|96.1|\n> |WTE|73.61| 73.61| 84.21 |73.61| 84.21 |96.54\n> |Std.F.T|38.89| 0| 37.84 |5.55| 5.26 |98.77\n> \n> |LR|Edit 1|Edit 2|Edit 3|\n> |--|----|---|--|\n> |RefineViT(ours)|91.79|88.12|84.69|\n> |WTE|90.48 |86.52 |85.176 |\n> |Std.F.T|95.99| 83.84| 83.88 \n> \n> - Compared to the SOTA method (WTE) and standard fine-tuning with \\$l\\_2\\$ regularization, our approach consistently maintains higher GR and LR across all edits. This indicates that **RefineViT better preserves previous edits while maintaining locality**, showcasing its effectiveness in multiple editing settings.\n\n**Q2 & W1.** Can the method be adapted for more complex vision tasks such as object detection, instance/semantic segmentation? What are the main challenges or required modifications? \n> Extending model editing to complex tasks like **object detection or segmentation** presents several challenges:\n> 1. These tasks involve **different error types** and outputs (e.g., bounding boxes or masks rather than class labels).\n> 2. Unlike classification, which primarily uses the **CLS token**, detection/segmentation rely on **all patch tokens**, making attribution and decomposition more difficult.\n> 3. There is currently **no established benchmark or protocol** for model editing in these settings, making standardized evaluation difficult.\n> \n> Despite these challenges, we demonstrate the **potential of our method beyond classification** by applying it to a **CLIP-based segmentation task**, following the setup in \\[4]. Using a single edit sample, we observe improved pixel-wise accuracy on similar images and a slight improvement in overall performance:\n> \n> | Avg. Pixel-Wise Accuracy | Similar Data | Overall |\n> | ------------------------ | ------------ | ------- |\n> | Original CLIP-ViT        | 72.29%       | 76.78%  |\n> | CLIP-ViT-Edited (Ours)   | 74.50%       | 77.01%  |\n> \n> This shows initial promise for adapting our framework to segmentation tasks. We plan to explore this direction further and develop benchmarks and methods for model editing in segmentation.\n>\n>[4] Yossi Gandelsman, Alexei A. Efros, and Jacob Steinhardt. Interpreting clip’s image represen- tation via text-based decomposition."}, {"review_id": "52wMLvVjY5", "reviewer": "Reviewer_yWzz", "summary": "This paper presents RefineViT, a novel two-stage framework for model editing in Vision Transformers (ViTs). The authors make an observation that, unlike Language Models, ViT predictions are more sensitive to perturbations in the Multi-Head Self-Attention (MSA) layers than the MLP layers. Capitalizing on this, RefineViT identifies and ablates erroneous attention heads to correct misclassifications. To preserve performance on unrelated inputs, the method learns a projection matrix that aligns the edited representation with the original, unperturbed one. The experiments demonstrate that this MSA-centric approach is more effective for ViTs than directly adapting LM editing techniques.", "questions": "* The paper's central claim is that MSA layers are more critical than MLP layers for storing factual knowledge in ViTs. Could the authors provide more intuition or a theoretical argument for why this might be the case?\n* If you repeat Stage 1 using two different misclassified examples of the same class, how consistent is the top-T list?\n* Figure 4: it looks like ROME, which is a method thats put strong constraints on MLP, performs really well on locaility. It seems to be a little contradicting to the core finding of this paper (MSA plays a more important role in preserving the model's capability).\n* L202: \"This allows updates in under 0.3 seconds for 50 epochs\". Does that include the time to find the ablation matrix A?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**Strengths:**\n\nFundamental Insight - The core contribution of this paper is the empirical finding that MSA layers are the primary locus of factual knowledge in ViTs. This is a significant departure from the established wisdom in the NLP domain, where MLP layers are considered the key-value memories. This insight alone is a valuable contribution to our understanding of vision transformers.\n\nMethodological Soundness - The two-stage design of RefineViT follows logically from the empirical findings. The selective ablation of attention heads combined with projection learning is principled and effective.\n\n**Weaknesses:**\n\nAblation Needed - There needs to be experiments to understand that which component in this approach is more critical. In multi-label settings, can you ablate the identified heads without learning P, and conversely train P without zeroing heads, to quantify each stage’s individual impact on GR–LR performance?\n\nComputation Complexity - Computing head utilities requires O(L\\!\\times\\!H) forward passes per sample. Should A be computed for each sample independently?", "was_revised": true, "final_justification": "I would like to thank the authors for providing detailed and well-structured responses. The rebuttal addresses my concerns and I will keep my rating.", "rebuttal": "We sincerely thank the reviewer for recognizing our insight that MSA layers are the primary knowledge locus in ViTs, as well as the principled and effective two‑stage design of RefineViT. Below, we address all weaknesses and questions, supplemented with new experiments for further support.\n\n**W1**. Ablation studies to clarify the individual impact of each component: 1.ablate without learning P in multi-label setting. 2. training P without zeroing heads\n> We conducted two ablation experiments to assess the individual impact of the two stages:\n> \n> 1. **Ablation without learning P**: We disable Stage Two and directly ablate the identified heads. On the ViT Natural Image benchmark (class 954\\_582), this results in a **sharp drop in Locality Rate (LR)** as more heads are ablated, confirming that learning P is critical for preserving locality (see table below).\n> \n> | N\\_ablated | 0   | 3           | 6          | 9          | 12          |\n> | ---------- | --- | ----------- | ---------- | ---------- | ----------- |\n> | LR(%)      | 100 | 46.16±2.09  | 38.08±4.91 | 32.94±8.55 | 29.03±10.37 |\n> | GR(%)      | 0   | 73.78±13.08 | 86.67±2.30 | 84.89±2.03 | 84.44±2.78  |\n> \n> 2. **Training P without zeroing heads**: This is already shown in Figure 9 (Appendix, Page 20). Without Stage One (head zeroing), the GR gains largely vanish, highlighting its role in guiding P's learning.\n> \n> Together, these results demonstrate that both stages are essential and complementary.\n\n\n**W2**.Computation Complexity: $\\mathcal{O}(L! \\times H)$ forward passes per sample?\n> In practice, we perform **only one forward pass per sample**. The required components for computing head utilities (i.e., the decomposition of the backbone output) are stored and reused. This design avoids the $\\mathcal{O}(L! \\times H)$ cost and keeps our method computationally efficient. A detailed description of the head‑selection procedure is provided in Appendix (p.14, Summary of the Procedure).\n\n**Q1**.More intuition or theoretical argument why MSA are more critical?\n\n> The intuition is that in ViT architectures, the lower layers primarily focus on extracting local features from individual patches (e.g., parts of an object like a wing or a head), while higher layers integrate these local features into composite concepts (e.g., recognizing a bird as a combination of head, wings, claws, and tail).\n>\n>Since MSA layers capture relationships between patches, errors in these layers can significantly disrupt the formation of composite concepts, leading to a larger impact on prediction accuracy. In contrast, MLP layers mainly refine the information extracted by MSA but do not directly capture inter‑patch relationships, making their impact comparatively less critical.\n\n**Q2**. If repeating stage 1 using different samples, how consistent are the resulting top-K list?\n\n> We conduct an experiment where we randomly selected two images from the ViT editing benchmark (class 954_582) and measured the overlap in their top-K head selections. We repeated this process 100 times and averaged the results to minimize randomness. The results show that the top-K lists have an average overlap of over 65%, indicating good consistency. Specifically, the overlap percentages for different values of K are as follows:\n> \n> | Top-K | K=1 | K=3 | K=6 | K=9 | K=12 | K=15 |\n> |-------|-----|-----|-----|-----|------|------|\n> | Overlap (%) | 73.0 | 81.33 | 68.67 | 65.11 | 66.92 | 69.87 |\n> \n> These findings demonstrate that our method produces consistent top-T lists even when using different misclassified examples from the same class.\n\n**Q3.** Why ROME (update last layer of MLP) performs really well on locality? Does it contradicts the core finding of this paper (MSA plays a more important role in preserving the model's capability).\n\n> **No, it does not contradict.** \n> - Locality measures how well predictions are preserved on unrelated data. ROME edits only the MLP layer and leaves the MSA modules untouched, so the model's core representations are largely preserved, leading to a higher Locality Rate (LR). \n> - However, its high locality can come at the cost of weak edits. As shown in Figure 4, ROME has a relatively low Generalization Rate (GR), indicating limited editing success. In fact, if no editing were performed, the Locality Rate would remain at 100%, highlighting the trade-off between locality and effective editing. \n>\n> Thus, our core insight remains valid: MSA modules are central to knowledge preservation and effective editing.\n\n**Q4.L202**: \"This allows updates in under 0.3 seconds for 50 epochs\". Does that include the time to find the ablation matrix A?\n> The reported update time of under 0.3 seconds for 50 epochs does not include the time to compute the ablation matrix A. However, as noted in our response to W2, Stage One requires only a single forward pass per sample, after which the backbone’s output is stored. Subsequent steps involve only matrix multiplications (for CLIP-ViT) or classifier-head-only forward passes, making them extremely fast with negligible time cost. Thus, the overall process remains highly efficient."}]

Author rebuttal:

Meta-review: The paper proposes a method for model editing for vision transformers. The authors address classification tasks, and identify the self-attention heads that lead the most to incorrect predictions. The authors then propose to counter the effect of these heads by applying a linear transformation to the embeddings that minimally affected unrelated samples, whilst also approximating the output obtained when these localised heads are removed.

Reviewers appreciated the insight from the authors that editing self-attention layers is more effective for vision transformers, as opposed to MLPs in transformers for language. Moreover, reviewers felt that this is a relevant, well-motivated problem and the paper is clear and well-written.

The main concerns of the reviewers were well-addresed during the rebuttal. Please update the final camera-ready with this.
Decision: Accept (poster)