Paper: Localizing and Editing Knowledge In Text-to-Image Generative Models (ICLR-2024)

Reviewer comment (Novelty):
The main toolbox of the interpretation method, Causal Mediation Analysis, is borrowed from previous works. There is a limited novelty in terms of the interpretation framework.

Author response:
**The main toolbox of the interpretation method, Causal Mediation Analysis, is borrowed from previous works. There is a limited novelty in terms of the interpretation framework.** : We highlight that Causal Mediation Analysis [1] is a tool from the causality literature which has been used in well-cited and impactful papers ([2]). While we borrow causal mediation analysis from the existing causality framework, to the best of our knowledge we are the first to use it to understand the inner workings of text-to-image models in detail.
We also emphasize that it is not trivial to adapt causal mediation analysis to diffusion models (especially large text-to-image models such as Stable-Diffusion). There are differences in not only architecture (e.g., the UNet in SD has several modules such as self-attention, cross-attention, residual blocks at various granularities), but also in sampling objectives where diffusion models have an added complexity of sampling over multiple time-steps. 
Moreover, our adaptation of causal mediation analysis for text-to-image models has revealed many interesting findings for e.g., Knowledge is distributed in the UNet which makes it difficult to edit concepts in the UNet, while it is localized in the text-encoder which makes it easier to perform closed-form edits.  These are significantly novel scientific findings and can help improve the understanding of large text-to-image models. 
We believe that using well-accepted and published tools from the community to reveal new scientific findings is a more constructive way to perform research, as it leads to greater adoption of existing research works by the community. Had we introduced an altogether different toolbox, the reviewer would need to evaluate this new toolbox in its entirety, in addition to our other findings. We sincerely hope that the reviewer appreciates this perspective! 

[1]. Judea Pearl. Direct and indirect effects. In Proceedings of the Seventeenth Conference on Uncertainty in Artificial Intelligence, UAI’01, pp. 411–420, San Francisco, CA, USA, 2001. Morgan Kaufmann Publishers Inc. ISBN 1558608001

[2]. Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in GPT. Advances in Neural Information Processing Systems, 36, 2022

Reviewer comment (Experiments):
The experiments presented in the paper all use Stable-Diffusion. The results would be more convincing if other classes of diffusion models could be investigated, which would provide important cues on whether the observations are specific to the Stable-Diffusion architecture or can be transferred to other models that adopt diffusion-based training.

Author response:
**The experiments presented in the paper all use Stable-Diffusion. The results would be more convincing if other classes of diffusion models could be investigated, which would provide important cues on whether the observations are specific to the Stable-Diffusion architecture or can be transferred to other models that adopt diffusion-based training.** : In our updated version of the paper, we have added new experiments using our framework on other Stable-Diffusion versions (v2-1, v1-5), OpenJourney, Kandinsky and SD-XL architectures.  Our current results are consistent across the newer SD versions too. Across all these models, we find that knowledge is distributed in the UNet (see Sec. R) - Appendix), whereas it is localized in the text-encoder (see Sec.Q). Also we find that our method can effectively edit concepts in all these models (Sec Q.1 - Appendix and Sec. Q.2, Sec. Q.2.1). 

For SD-XL[1] (see Sec. (S)) which has a slightly different architecture than Stable-Diffusion, we have the following observations : In SD-XL, which is a recent text-to-image model consisting of a much larger UNet (consisting of 2.6B parameters) – we find that causal states are distributed like the Stable-Diffusion checkpoints. However for SD-XL we find that only 12 causal layers exist out of 227 layers (~5% of the total layers).  This shows that our interpretability framework can be effectively used to probe other text-to-image generative models. As the Reviewer TjyF noted that for each of the models, laborious experiments are needed – a full study of SD-XL is attributed for future work.  In Sec. (R.2) - we show successful edits on OpenJourney[2]. 

In Kandinsky[3] text-to-image generative model (see Sec. (U)) – we find only one localized causal state in the text-encoder, thereby enabling fast edits using DiffQuickFix. Therefore each component of our paper can be used first to probe any text-to-image models and then edit them if localized causal states exist. 

[1].Podell et al. 2023, arXiv. https://arxiv.org/pdf/2307.01952.pdf

[2]. OpenJourney: https://huggingface.co/prompthero/openjourney

[3]. Kandinsky: https://huggingface.co/docs/diffusers/api/pipelines/kandinsky

Reviewer comment (Evaluation):
The causal states are retrieved solely relying on CLIP-Score, but it's possible that such a score is sensitive to only some prominent visual attributes such as colors. The paper claims to identify the knowledge source of general visual attributes but does not provide an investigation on the efficacy of CLIP-Score to discriminate different kinds of visual attributes.

Author response:
**The causal states are retrieved solely relying on CLIP-Score, but it's possible that such a score is sensitive to only some prominent visual attributes such as colors. The paper claims to identify the knowledge source of general visual attributes but does not provide an investigation on the efficacy of CLIP-Score to discriminate different kinds of visual attributes.** : While CLIP-Score is not the perfect metric it is still widely accepted by the community and still used for evaluating text-to-image generative models. For the attributes which we study in our paper, we find that CLIP-Score is a good metric to discriminate between different kinds of visual attributes. We have added a new section in the paper on this (see Sec. (P) - Appendix). 
In particular, we use the captions in our prompt-dataset to generate images corresponding to different visual attributes.  We then annotate each image based on the visual attribute which was used to generate the image. Then within the category of a given visual attribute, we use CLIP-Score to classify the generated images into the right categories. Overall, we find that CLIP-Score is effectively able to perform well in all the categories of visual-attributes, thereby highlighting its efficacy in our specific use-case. 
However, we also highlight that there are potential open future directions which can be explored to improve upon CLIP-Score (e.g., using a VQA model for color or a specialized style detector for style) for text-to-image models in general – however that warrants a separate investigation in itself.

Reviewer comment (Experiments):
The paper highlights the difference in causal state locations between Stable-Diffusion and GPT. Is such a difference specific to the CLIP encoder or is general to the diffusion model class?

Author response:
**The paper highlights the difference in causal state locations between Stable-Diffusion and GPT. Is such a difference specific to the CLIP encoder or is general to the diffusion model class?** : We highlight that we perform causal tracing in two components of the diffusion model : (i) UNet and (ii) Text-Encoder. In the UNet, we find that the causal states are distributed across the whole model, while in GPT causal components are localized within a window of layers. This distribution is primarily due to the diffusion model class, as the UNet is optimized during pre-training. For the text-encoder, we find only one causal state, as opposed to GPT’s multiple causal states within a window. Given that the text-encoder is frozen during pre-training, we believe this surprising observation is more due to CLIP’s pre-training objective. Our general hypothesis is that language models trained with different pre-training objectives will contain causal states in a different way. A full study of such phenomena across different language models is a good direction for future study!

Reviewer comment (Experiments):
The scope of the claims seems too broad. The title and introduction claim to locate the knowledge of text-to-image diffusion models, while in the paper, only one stable diffusion model checkpoint is investigated. Given that most of the findings on this model are through laborious experiments, it is unclear if these findings can be generalized to even other versions of stable diffusion models, not to mention other types of text-to-image models. If the findings are only applicable to the studied checkpoint, the impact of the method may be significantly restricted, as the studied model is not considered state-of-the-art.

Author response:
**The scope of the claims seems too broad. The title and introduction claim to locate the knowledge of text-to-image diffusion models, while in the paper, only one stable diffusion model checkpoint is investigated. Given that most of the findings on this model are through laborious experiments, it is unclear if these findings can be generalized to even other versions of stable diffusion models...** :  The reviewer raises an interesting question about the generalizability of our framework and results for other checkpoints of Stable-Diffusion.  We first want to highlight that the current version of Stable-Diffusion in our paper is also used in other fine-tuning based concept ablation works (e.g., Concept-Ablation[1], Concept-Erasure[2]).  Therefore to maintain consistency and for fair comparison, we use the current Stable-Diffusion checkpoint. We also note that each component of our paper (causal tracing for UNet and causal tracing for text-encoder) can be independently used to probe any text-to-image model and if the localization property arises, then DiffQuickFix can be used to further edit the model. We have also experimented with our causal tracing + editing framework, in other Stable-Diffusion versions (v1-5 and v2-1). In Sec.(R) and Sec.(Q) - we show that our interpretability and editing method can be used with other SD versions. In fact, we find similar observations in other Stable-Diffusion checkpoints – Knowledge about various visual attributes is localized in the CLIP text-encoder (see Sec. (Q.1)), whereas it is distributed in the UNet (see Sec. (R)). In fact, we show that our editing method works well across the latest Stable-Diffusion checkpoints (see Fig. (60) and Fig. (61) - Appendix). 

We have also experimented with other diffusion models such as SD-XL [3](Sec S) and OpenJourney[4] (Sec. R.1.3, R.2). In SD-XL, which is a recent text-to-image model consisting of a much larger UNet (consisting of 2.6B parameters) – we find that causal states are distributed in the UNet like the Stable-Diffusion checkpoints. However for SD-XL we find that only 12 causal layers exist out of 227 layers (~5% of the total layers).  This shows that our interpretability framework can be effectively used to probe other text-to-image generative models. As the reviewer noted that for each of the models, laborious experiments are needed – a full study of SD-XL is attributed for future work. For OpenJourney, we were able to perform successful edits, as shown in Sec. R.2. 

In Kandinsky[5] text-to-image generative model (see Sec. (U)) – we find only one localized causal state in the text-encoder, thereby enabling fast edits using DiffQuickFix. 

Overall our paper is based on the idea that if there exists localized causal components in some part of the model, then fast edits can be made to remove or edit concepts. Each component of our paper can be used for other text-to-image generative models as shown in the recent experiments. For e.g., the interpretability framework using Causal Mediation Analysis can be used to identify if localized knowledge exists or not. Once localized knowledge is found, then our editing method DiffQuickFix can be used to remove or edit concepts effectively. 

[1]. Nupur Kumari, Bingliang Zhang, Sheng-Yu Wang, Eli Shechtman, Richard Zhang, and Jun-Yan Zhu. Ablating concepts in text-to-image diffusion models. In ’International Conference on Computer Vision (ICCV)’, 2023.

[2]. Rohit Gandikota, Joanna Materzynska, Jaden Fiotto-Kaufman, and David Bau.
Erasing concepts from diffusion models. In Proceedings of the 2023 IEEE International Conference on Computer Vision, 2023.

[3].Podell et al. 2023, arXiv.  https://arxiv.org/pdf/2307.01952.pdf

[4]. OpenJourney:  https://huggingface.co/prompthero/openjourney

[5]. Kandinsky: https://huggingface.co/docs/diffusers/api/pipelines/kandinsky

Reviewer comment (Experiments):
My main question is how the results of this study generalize to other text-to-image generation models, and what are the components that can be reused if we want to apply the method to another model.

Author response:
We thank the reviewer for their response! We want to highlight that our methodology is general but its applications to different text-to-image generative models can reveal their inherent model specific behavior. We have added a new section highlighting these points (see updated Sec. (W)) -- that although our interpretability framework can be used to find knowledge distribution in different text-to-image models, their inference about the distribution ratio about different visual attributes might be different due to different UNet architectures. We believe that this in itself is a novel finding and it's an artifact of the model behaviour which our causal tracing framework is able to articulate.

Although the UNet in SD-XL is almost 3 times bigger than SDv-1.4 and SDv-2.1 -- we still find that knowledge is distributed in the UNet,  but the ratio of distribution is different from Stable-Diffusion versions (as emphasized in Sec. (S)).  For e.g., down-1-resnets.1 is causal as well as mid-blocks.attention.0.transformer-blocks.0.attn1 , up-blocks.0.resnets.2 are causal and these layers are located very far from one another in the Unet. This signifies that knowledge is indeed distributed in the UNet for SD-XL.   

We also highlight that our work is more comprehensive than existing published works such as [1,2] which show editing performance on only one Stable-Diffusion variant.  During the rebuttal we have shown the generalizability of our editing framework to multiple Stable-Diffusion checkpoints, OpenJourney and Kandinsky! 

[1]. Nupur Kumari, Bingliang Zhang, Sheng-Yu Wang, Eli Shechtman, Richard Zhang, and Jun-Yan Zhu. Ablating concepts in text-to-image diffusion models. In ’International Conference on Computer Vision (ICCV)’, 2023.

[2]. Rohit Gandikota, Joanna Materzynska, Jaden Fiotto-Kaufman, and David Bau. Erasing concepts from diffusion models. In Proceedings of the 2023 IEEE International Conference on Computer Vision, 2023.

We sincerely hope that these additional points in the paper and in this comment, will serve as a pointer to provide a full picture of causal tracing to the reader across various Unet architectures and alleviate the reviewer's concerns.

Reviewer comment (Experiments):
Other than optimizing the W_out, an alternative approach would be to directly replace the token embeddings, e.g., from embeddings for "Van Gogh" to those for "painting". I wonder what the authors' view is on the effectiveness of this approach.

Author response:
**Other than optimizing the W_out, an alternative approach would be to directly replace the token embeddings, e.g., from embeddings for "Van Gogh" to those for "painting". I wonder what the authors' view is on the effectiveness of this approach**: The reviewer brings up a good suggestion about alternative mechanisms for model editing, apart from the one we have provided. Intuitively, there are four pathways to go from text-prompt to generated image in a diffusion model. These pathways are : from user's prompt input to text-encoder, from text-encoder to UNet, from UNet to Image-Decoder, and from output of Image-Decoder back to the user. Disrupting any of these pathways (in a multitude of ways) can be effective in concept editing.
For example, one could intercept the user's original prompt and replace words or phrases using a look-up table (but a simple look-up table fails to handle mis-spellings or paraphrases [1]). One could similarly run an image classifier on the output image (eg : a "Pikachu detector", where "Pikachu" is a trademarked object owned by Nintendo Inc), and decline to show the user this output. In reality, one might need to combine such interventions to create a truly robust concept-ablation system.  With respect to the reviewer’s suggestion, it is not immediately clear to us how replacing token embeddings would work for concepts with multiple subwords or mis-spellings (we refer reviewer to Section K, which shows our model is robust to mis-spellings, a common real-world attack used to bypass safety features in diffusion models[1]). It might be possible to do this using a look-up table, but vendors would now need to ship this look-up table/ban-list along with the model, whereas our method can just ship the model. In summary, we agree with the reviewer in spirit. An analogy we have often used in internal discussions has been : one can cure poor eyesight using different mechanisms - wearing glasses, wearing contact lenses, performing LASIK surgery, or perhaps even performing neuro-surgery in the brain. Each of these methods has pros and cons, appropriate to the patient's circumstances, ease of technological availability, recovery period, etc. It helps to have these multiple options at our disposal, and combinations of these options may be required at times. We provide one such novel option in our submission.

[1] Hongcheng Gao, Hao Zhang, Yinpeng Dong, and Zhijie Deng. Evaluating the robustness of text-toimage
diffusion models against real-world attacks, 2023.

Reviewer comment (Experiments):
I don't see any results on removing trademarked objects as claimed at the top of page 3. Were they removed during the preparation of the manuscript?

Author response:
**I don't see any results on removing trademarked objects as claimed at the top of page 3. Were they removed during the preparation of the manuscript?**:  We refer the reviewer to Section F.2 and Section H. R2-D2 is a robot from star-wars, trademarked by Lucasfilm Ltd LLC. Snoopy the dog is trademarked by Peanuts LLC. Nemo is an animated fish trademarked by Disney Enterprises, Inc. Grumpy cat is trademarked by Grumpy Cat Limited. We show their individual concept ablation in Section F.2 - i.e., four edited models, each unable to generate the trademarked concept that was erased from it. In Section H, we show that all of these concepts can be simultaneously erased by our model, i.e.,the same model is unable to generate all four trademarked objects. We have added this information to the Figures in Section F.2 to add more context to the reader.

Reviewer comment (Theory):
While it seems reasonable to use CLIPScore as AIE, it does not seem that |CLIPScore(x_0^{restored}, c) - CLIPScore(x_{0}^{corr}, c)| can be approximated by CLIPScore(x_0^{restored}, c). In particular, both increases and decreases of CLIPScore(x_0^{restored}, c) by the same amount around CLIPScore(x_{0}^{corr}, c) lead to the same absolute difference, but the CLIPScore(x_0^{restored}, c) is obviously different.

Author response:
**While it seems reasonable to use CLIPScore as AIE, it does not seem that |CLIPScore(x_0^{restored}, c) - CLIPScore(x_{0}^{corr}, c)| can be approximated by CLIPScore(x_0^{restored}, c).** : The reviewer raises a valid point. We have modified the approximation (by removing the absolute |.| operation), where now (CLIPScore(x_0^{restored}, c) - CLIPScore(x_{0}^{corr}, c)) can be approximated by CLIPScore(x_0^{restored}, c), as CLIPScore(x_{0}^{corr}, c) is common across all the layer interventions for a given caption c. We have updated this detail in the paper. Note that even with this modified approximation, the results don’t change in our paper, as AIE is still used as CLIPScore(x_0^{restored}, c).