Learning Where to Edit Vision Transformers

Model editing aims to data-efficiently correct predictive errors of large pre-trained models while ensuring generalization to neighboring failures and locality to minimize unintended effects on unrelated examples. While significant progress has been made in editing Transformer-based large language models, effective strategies for editing vision Transformers (ViTs) in computer vision remain largely untapped. In this paper, we take initial steps towards correcting predictive errors of ViTs, particularly those arising from subpopulation shifts. Taking a locate-then-edit approach, we first address the where-to-edit challenge by meta-learning a hypernetwork on CutMix-augmented data generated for editing reliability. This trained hypernetwork produces generalizable binary masks that identify a sparse subset of structured model parameters, responsive to real-world failure samples. Afterward, we solve the how-to-edit problem by simply fine-tuning the identified parameters using a variant of gradient descent to achieve successful edits. To validate our method, we construct an editing benchmark that introduces subpopulation shifts towards natural underrepresented images and AI-generated images, thereby revealing the limitations of pre-trained ViTs for object recognition. Our approach not only achieves superior performance on the proposed benchmark but also allows for adjustable trade-offs between generalization and locality. Our code is available at https://github.com/hustyyq/Where-to-Edit.

Paper

References (63)

Scroll for more · 38 remaining

Similar papers

Peer review

Reviewer MKpn6/10 · confidence 2/52024-07-07

Summary

The paper presents a novel approach for editing Vision Transformers (ViTs) to correct predictive errors, specifically addressing subpopulation shifts. It introduces a learning-to-learn methodology utilizing a hypernetwork that identifies and modifies a small set of critical parameters in response to erroneous samples. This method is distinguished by its focus on localized edits, ensuring minimal disruption to unrelated model functions while generalizing corrections to similar errors.

Strengths

The concept of utilizing a hypernetwork to identify editing locations in ViTs for image processing tasks is innovative, addressing a gap in the literature concerning efficient and localized model editing. The proposed method is rigorously validated through the introduction of new benchmarks, showing significant improvements over existing techniques. The paper is well-organized and clearly written, with technical details and methodologies elaborately explained, making it accessible to readers familiar with the field. This research is highly significant as it enhances the practical utility of ViTs in real-world applications, particularly where predictive reliability is crucial, such as in autonomous driving and medical imaging.

Weaknesses

The success of the proposed method heavily relies on the hypernetwork's ability to accurately predict edit locations. Misidentifications could lead to suboptimal edits, affecting model reliability.

Questions

Could a strategy similar to the one proposed by [1] be adapted to apply a low-rank update to the weight matrix for memory editing in ViTs? [1] Meng, Kevin, et al. "Locating and editing factual associations in GPT." Advances in Neural Information Processing Systems 35 (2022): 17359-17372.s:

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

2

Limitations

The CutMix technique employed may not adequately represent real-world shifts, which could limit the applicability of the findings to practical scenarios.

Reviewer sCvZ7/10 · confidence 3/52024-07-11

Summary

The paper explores and discussed the VIT editing, which is the first work in this field. The paper proposes a learning-to-learn approach that identifies a small set of critical parameters to edit to correct the prediction of an error example. The paper considers the reliability and generalization during the ViT editing. In experiments, the paper curates two benchmarks to validate that the framework can correct the error and provide customization.

Strengths

1. The problem formulation and method are clearly defined and explained, and the method is easy to follow. 2. The problem of ViT and the proposed method based on bi-level optimization are new in this field and meaningful. 3. The paper contributes the benchmark for the problem including the natural data and generated data which is beneficial to later research. 4. The experiments are intensive and validated on both real data and generated data. The ablation and analysis cover various metrics that form strong evidence.

Weaknesses

1. The motivation of CutMix and its influence are not clearly discussed and demonstrated. For example, there are also other data augmentation methods, are they the same as CutMix? Can they generate the same 2. The paper mention that the editing should not influence the irrelevant examples. However, this is not clearly discussed in the paper. I have questions about the metric Locality. Specifically, if the predicted probability of the irrelevant examples changed, does this indicate that the model change?

Questions

Please refer to the weakness part.

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

The author address the limitations.

Reviewer 91WR5/10 · confidence 4/52024-07-12

Summary

The paper explores a model-editing task for ViT that is similar to model-editing in LLMs. For ViT model-editing, the paper proposes a meta-learning-based approach to selecting parameters for fine-tuning updates. The proposed selecting method trains a hyper-network that selects learning parameters using episodes made by CutMix augmentation. The proposed method exhibits superior performance to similar methods on newly made Editing Benchmark.

Strengths

- Exploring model editing for ViT with a meta-learning approach is new and interesting. - The paper proposes a new benchmark for model editing, including data set and evaluation metric.

Weaknesses

- Explanation and justification for the necessity of model editing in the vision domain are not enough. - Although model editing is a powerful and important tool for LLM tuning, it doesn't explain the need for model editing in vision models. - Vision models have relatively small parameters, and FLOPs are considered to be an important cost rather than memory and the number of parameters. Model editing might be effective in reducing parameters and memory, but I'm not sure it is also beneficial in the vision domain. - Hypernetwork requires too high cost to select parameters for model editing. - To the best of my knowledge, the major goal of model editing is to reduce costs for full fine-tuning like LoRA. But, the proposed method requires huge computation, and the training process might be larger than full finetuning. - For 12 blocks ViT-B, the proposed method trains a hyper-network with 5 transformer blocks using multiple training episodes with inner- and outer-loop optimization, which might require huge computation costs compared to the original network. - The hyper-network is trained only for a single pre-trained network. Thus, it is required to train other hyper-networks for other pre-trained networks. This whole process doesn't look efficient.

Questions

- The authors argue the importance of model editing based on LLM scenarios. But, I think vision domains are different from LLM scenarios. What are the benefits of general model editing in the vision domain? Without performance improvement made by the proposed method, does model editing have a unique role in the vision domain? It is important to evaluate contribution as a new editing benchmark and the first model editing in vision. - The proposed method requires training a hyper-network with two-loop optimization, which seems to be a high-cost computation. Does this method offer any computational cost advantages compared to other existing methods? - Can hyper-network be used to select parameters for another pre-trained network? - I think hyper-network training with multiple episodes is a significant advantage for a fair comparison. Do other methods also use similar processes? - The proposed method uses 5 transformer blocks to select just 3 FFNs. It is not parameter efficient and can't be expanded to large-scale networks over a billion scales. Is there any way to expand the proposed method for larger select space?

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

I didn't find a limitation section in the paper.

Authorsrebuttal2024-08-13

Discussion with Reviewer 91WR (2/2)

##### [D2] Extension to large-scale networks > We understand the reviewer's remaining concern, which seems to arise from two main aspects: the **(I) computation cost** and **(II) performance** of the proposed method on larger-scale networks. We would like to humbly clarify the following. > - **(I)** The computational cost of our method on larger-scale networks > - Above all, we humbly clarify again that the primary motivation for model editing is to enable data-efficient updates, not necessarily to reduce computational cost. *If full fine-tuning offers the best generalization-locality trade-off, it should indeed remain the preferred approach*. > - When editing ViT/B-16, our method **demonstrates greater efficiency compared to some SOTA baselines**: ours consumes 50% FLOPs of full fine-tuning while KN consumes 77%. > - When **editing SwinV2-g (1B)** which has 1134% more parameters than ViT/B-16 (86 MB), the FLOPs of our method only **marginally increases** by 435% (from 26.4G to 141.5 G), while the FLOPs of full fine-tuning increases by 976 % (from 52.8 G to 567.9 G). > - **(II)** The performance of our method on larger-scale networks > The reviewer's concern seems to stem from the following hypotheses: (1) larger networks require more FFNs to edit, (2) more FFNs to edit would necessitate significantly larger hypernetworks, and (3) the relatively small performance improvement in Fig r1 (b-c) is due to an in sufficiently large hypernetwork and, consequently, not enough FFNs being edited. However, we humbly clarify each point below. > - **Larger networks do not necessarily require more FFNs to edit.** > - We focus on localizing and editing **a visual concept**; the neurons that describe a visual concept are (1) **concentrated in the middle to late layers** [3], as the early layers encode distributed representations of very primitive concepts (e.g., geometric patterns) [1]; and (2) **compact** -- the number of neurons associated with a single concept does not largely increase as the model size grows [2], with the additional model capacity likely benefiting the representation of more diverse concepts. > - During the discussion period, we conducted experiments by increasing the number of FFNs edited in SwinV2-g. The results, showing only marginal improvement when editing more FFNs, also corroborate that editing just 3 FFNs is sufficient to balance generalization and locality. > > |Number of FFNs |Identified layers by natrual 933-923 group |Generalization|Locality| > |----|---|---|---| > |3|19~21|94.77|75.95| > |5|18~22|94.39|77.69| > |8|17~24|95.92|75.24| > |Number of FFNs|Identified layers by AI oil painting|Generalization|Locality| > |3|19~21|87.99|72.53| > |5|19~23|91.30|72.36| > |8|17~24|88.18|71.23| > - Our practice of editing 3 FFNs has also been validated successful and scalable in the context of LLMs by the prior hypernetwork-based method MEND [4], which successfully edits 3 FFNs across various scales of LLMs, including distilGPT-2 (82M), GPT-Neo (2.7B) and GPT-J (6B). > - **Increasing the number of FFNs to edit requires only a marginal increase in the size of the hypernetwork.** > - To accommodate larger channel sizes, the changes in the hypernetwork are limited to the input and output layers: (1) adding an input projection layer (e.g., 3584x768 for SwinV2-g) that projects the encoded image features to a lower dimension (which is the input dimension of the hypernetwork) and (2) increasing the output dimension of the output projection layer. For editing SwinV2-g, the size increase resulting from these two changes amounts to 29.3% of the original size of the hypernetwork. > - To accommodate more FFNs for editing, the only change required in the hypernetwork is the addition of one more trainable input token for each FC layer. > - **The relatively small performance improvement in Fig. r1(b-c) is attributed to the difference between SwinV2-g and ViT across various classes.** > - We observed significant differences in the error samples produced by SwinV2-g and ViT for each class, reflecting distinct model behaviors. > - The two groups reported in Fig. r1(b-c) were selected because our method for editing ViT showed the largest improvements (by up to 18% GR at the similar TRR) over the baselines (including FT-L2) on these groups. Although the improvement for SwinV2-g on these specific groups appears modest, during the discussion period, we found that our method for editing SwinV2-g achieved a 16% GR improvement at the similar TRR on other groups, including Vase in scene light dataset, 407-654 in the natural dataset. [1] Identifying interpretable subspaces in image representations, ICML 2023. [2] Labeling neural representations with inverse recognition, NeurIPS 2023. [3] Interpreting CLIP's image representation via text-based decomposition, ICLR 2024. [4] Fast model editing at scale, ICLR 2022.

Reviewer 91WR2024-08-14

Thank you for your additional comments. They have addressed my concerns, and I have adjusted my rating to borderline accept. The paper offers valid contributions. However, I still question the value of model editing in the vision domain compared to its impact on LLMs. Additionally, the novelty of the hyper-network design is limited and not particularly surprising. This is why I don't rate higher than borderline accept. Nonetheless, I appreciate the authors' active discussion.

Authorsrebuttal2024-08-14

We sincerely thank the reviewer for the valuable feedback and for increasing the score! We greatly appreciate the opportunity to discuss and refine our research work. Regarding the value of model editing in the vision domain, we anticipate its broad impact for several reasons: (1) correcting subpopulation shifts, which is prevalently encountered by vision foundation models, is an urgent need; (2) re-training a whole vision pre-trained model typically requires access to the full pre-training data (e.g., JFT-3B [1]), which is often inaccessible; and (3) our proposed editing method can be readily applied to diffusion models (e.g., SD-XL) for visual generation, which we are actively exploring. Meantime, we would also like to highlight the novelty of our hypernetwork design: (1) Our approach is the first to learn where to edit, while previous hypernetwork-based methods [2, 3] focus on how to edit; this is reflected in the completely different input/output of our hypernetwork compared to earlier methods. (2) Optimizing such a hypernetwork poses significant challenges, which motivates us to propose (a) construction of pseudo editing episodes, which is the key to learning how to localize, and (b) decoupling of hypernetwork optimization (see Appendix B.3.2), which is crucial for avoiding trivial mask solutions. Thank you once again for your thoughtful evaluation and for recognizing our contributions. [1] Scaling Vision Transformers, CVPR 2022. [2] Editing factual knowledge in language models, EMNLP 2021. [3] Fast model editing at scale, ICLR 2022.

Reviewer wreN7/10 · confidence 3/52024-07-12

Summary

The paper investigates a novel method for editing Vision Transformers (ViTs) to enhance their performance by rectifying predictive errors. Specifically, it proposes training an additional ViT to locate specific rows in the weight matrices of several feedforward network (FFN) layers for efficient fine-tuning. The authors evaluate their approach using benchmarks such as MAD-searched natural images and AI-generated image datasets. Compared with several state-of-the-art (SOTA) baselines, this work achieves the best locality/generalization tradeoff, offering flexible editing choices and accurate localization of where to edit.

Strengths

Originality: 1. The paper introduces a novel approach to model editing in vision tasks, which is an emerging area with significant potential. 2. It combines existing techniques in a new way to address specific challenges in model performance enhancement. 3. The method's ability to identify and edit specific parts of models is innovative and could be valuable for various applications. Quality: The submission is technically sound with comprehensive question-raising and experimental support. Clarity: The paper is well-written and organized, making it easy to follow the methodology and results. Significance: 1. The results are important as they demonstrate a new way to enhance vision models, which could have significant implications for both research and practical applications. 2. The method addresses a challenging task and advances the state of the art in model editing and vision tasks.

Weaknesses

Originality: The main objective, learning where to edit, is not a brand-new idea in NLP. The novelty could be better emphasized by comparing more thoroughly with recent advancements in model editing and vision tasks. Quality: 1. The application of the proposed method is somewhat limited to plain ViTs. It would be interesting to see the application to hierarchical ViTs like Swin and PVT. 2. The evaluation is somewhat limited to specific datasets. It would be better to explore other datasets like CUB-200-2011, NABirds, Oxford Flowers, Stanford Dogs, and Stanford Cars Clarity: No further issues. Significance: The impact of the work would be more convincing with broader validation across diverse datasets and tasks.

Questions

1. For the comparison with FT-L2, is it possible to compare with FT-L1 or apply L2 in the proposed method to have a clearer comparison on the regularization term? 2. Can the authors provide more detail on the computational complexity and resource requirements of their method? How does the additional complexity of the hypernet compare to the saving in fine-tuning? 3. In line 150, is there a detailed description of the challenge that the authors focus on? The flow seems to be interrupted here.

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

Limitations are addressed in the originality and quality in the weakness section.

Reviewer 91WR2024-08-09

Thank you for your response. I have a few additional requests and need clarification on rebuttal comments. - [W2 & Q2 & Q3] Computational efficiency concerns regarding the hypernetwork - It is hard to estimate how much `9 hours on a single RTX A6000 GPU (48G)` is. Could you give me some reference numbers, such as the computation times for ViT-B 1 epoch training time on a single RTX A6000? It might help me to understand how much the hypernetwork training costs. - [Q5] Extention to large-scale networks - I understand that your method's computation can be reduced while maintaining computation. But, it is still not enough to resolve my concern. - At first, I need numbers that make computation costs look small. `9 hours on a single RTX A6000` and `52.8G FLOPs per editing iteration` can be a precise report, but it doesn't look small. I expect comments like `it just takes 0.n% of finetuning computation` or `smaller than computation for single inference` that look small without analyzing and calculating computation costs. - My concern about large-scale networks is scaling. 3 FFNs might be enough for ViT-B because it is 25% of 12 FFNs. But, for SwinV2-G as an example, 3 FFNs are only 6% of 50 FFNs. Also, the hypernetwork should be bigger to cover large channel sizes. Then, the impact of your method might be reduced on a large network, which means it is not scalable. - In Fig r1, improvement compared to FT-L2 is much smaller than that of ViT-B. So, it is not enough to solve my concern about large-scale networks. Reducing the block size can reduce the total computation but still the hypernetwork has to be larger when channel-size increases and impacts of 3 FFNs also be restricted when the depth increases. So, I think this method can't be expanded to large-scale networks since it costs more while the effect is reduced, which can be a major weakness of the method.

Authorsrebuttal2024-08-13

Discussion with Reviewer 91WR (1/2)

We appreciate the reviewer for raising your remaining concerns. Please find our response below, and kindly let us know if it satisfactorily addresses your concern. ##### [D1] Computational efficiency regarding the hypernetwork > - For reference, in response to the reviewer's suggestion, we provide the computation time for training a ViT-B on one RTX A6000 (48G) with ImageNet-21k here. The training time is 25.55 hours per epoch, given the image size of 224x224 and the batch size of 256. > - Thus, training the hypernetwork accounts for 35% of the time required for 1 epoch of ViT-B training, which is fairly acceptable. > - Last but not least, as detailed in our general response [R2], hypernetwork training is **performed only once** for a pre-trained model **prior to editing**. It can be viewed as a brief (35% of 1 epoch) extended stage of pre-training, being both computationally acceptable and worthwhile.

Reviewer MKpn2024-08-12

I appreciate the authors response, which addressed all my questions. I increase my score to 6.

Authorsrebuttal2024-08-13

Thank you very much for the positive feedback and increasing the score! We reiterate our deep appreciation for the reviewer's dedicated time and effort in reviewing our paper and providing invaluable feedback.

Reviewer wreN2024-08-13

Thank you for the detailed response. I have updated my vote to accept.

Authorsrebuttal2024-08-13

We are pleased that the concerns raised by the reviewer have been addressed. Once again, we would like to express our gratitude for your constructive comments and positive feedback on our work.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC