Diffusion models (DMs) produce very detailed and high-quality images. Their power results from extensive training on large amounts of data---usually scraped from the internet without proper attribution or consent from content creators. Unfortunately, this practice raises privacy and intellectual property concerns, as DMs can memorize and later reproduce their potentially sensitive or copyrighted training images at inference time. Prior efforts prevent this issue by either changing the input to the diffusion process, thereby preventing the DM from generating memorized samples during inference, or removing the memorized data from training altogether. While those are viable solutions when the DM is developed and deployed in a secure and constantly monitored environment, they hold the risk of adversaries circumventing the safeguards and are not effective when the DM itself is publicly released. To solve the problem, we introduce \textsc{NeMo}, the first method to localize memorization of individual data samples down to the level of neurons in DMs' cross-attention layers. Through our experiments, we make the intriguing finding that in many cases, \textit{single neurons} are responsible for memorizing particular training samples. By deactivating these \textit{memorization neurons}, we can avoid the replication of training data at inference time, increase the diversity in the generated outputs, and mitigate the leakage of private and copyrighted data. In this way, our \textsc{NeMo} contributes to a more responsible deployment of DMs.
Paper
Similar papers
Peer review
Summary
This paper introduces NeMo, a straightforward and easy-to-understand method. NeMo identifies neurons responsible for memorization and deactivates them, offering potential applications in the protection of sensitive or copyrighted data.
Strengths
This topic is highly interesting as it enhances interpretability and offers a lightweight solution for addressing issues related to sensitive or copyrighted training images.
Weaknesses
1, There is insufficient comparison with other methods, such as the work mentioned in [1]. The authors state, "While mitigations that rely on preventing the generation of memorized samples are effective when the DM is developed and deployed in a secure environment, they hold the inherent risk of adversaries circumventing them. Additionally, they are not effective when the DMs are publicly released, such that users can freely interact with them." However, this argument is not persuasive enough. Providing more compelling reasons for the unique position of this method in mitigating memorization would enhance my evaluation score. 2, The side effects of this method are not thoroughly investigated. [1] Detecting, Explaining, and Mitigating Memorization in Diffusion Models
Questions
1, I have some reservations about the connection between copyright/sensitive images and memorization. It seems that copyright issues might only be partially triggered by memorization. Could you provide further clarification on this point? This pertains to the application of your method. 2, I find the activation-based localization method reasonable. However, I'm curious if pruning methods such as [1] or other neuron localization methods such as [2] could also be effective here. I would appreciate some analysis on this, though I do not require any additional experiments. 3, In section 4.3, could you explain why deactivating neurons further reduces FID and KID in some cases? 4, I'm also interested in the potential side effects of this method. You mention that deactivation has a minimal impact on non-memorized generation, which suggests that these neurons do not interact significantly with other neurons and function independently. I believe a more detailed analysis from this perspective would be beneficial. [1] Structural Pruning for Diffusion Models [2] Cones: Concept Neurons in Diffusion Models for Customized Generation.
Rating
6
Confidence
3
Soundness
3
Presentation
3
Contribution
2
Limitations
The limitations have been discussed.
Summary
This paper introduces NEMO, a method that identifies and deactivates specific neurons responsible for memorizing training data in DMs. This approach prevents the reproduction of sensitive images, enhances output diversity, and mitigates data leakage, enabling more responsible use of DMs.
Strengths
1. It explains the connection between neurons and memorization, which provides more understanding in this problem. 2. This paper provides two selection stages to select the different neurons. 3. The paper writing is clear.
Weaknesses
1. The author should specify their novelty compared with [1]. The novelty seems limited considering the conclusion in [1]. 2. No practical mitigation method is provided. Deactivating is not possible without knowing the neurons. 3. The author claim that most memorization happens in the value mappings of the cross-attention layers of DMs. However, actually, the author only prove that some neurons in attention layers are necessary in memorization. They didn't show that those neurons are everything about the memorization. It is possible that deactivating other neurons that are not in attention layers can also remove memorization. Localizing should consider how to exclude other neurons.
Questions
The author should address weakness 1 and 3.
Rating
6
Confidence
5
Soundness
3
Presentation
3
Contribution
3
Limitations
The limitations that the author mentioned in Appd A is about how to further improve the method. It is not clear solved but propose the key factors is possible to solve in the future work.
Missing Reference
Dear ACs and PCs, dear Reviewer due to a lack of reference [1] in the review, we were not able to provide the comparison requested by the reviewer during the rebuttal phase. We would have been happy to provide the requested information, but this was not possible during the initial rebuttal phase without further details from the reviewer. We, therefore, kindly ask the Reviewer to provide additional information. Best regards, the authors
Can you specify the missing reference?
Dear Reviewer K7EP, Can you specify the missing reference [1] in your review? It seems to have been dropped in your original review. I'm not sure if there will be enough time for the authors to provide comparison but at least some clarifications regarding the novelty could be discussed during this week. Thanks! AC
Dear Reviewer K7EP, thank you for your quick response and the clarification of the missing reference. **1.) The contribution of this paper seems limited given that Ren et al. show that memorization and cross attention may have a connection**: We first emphasize that the referred paper (Ren et al.) is an unpublished preprint made available on Arxiv about 2 months (March 17th) before the NeurIPS deadline, such that it would fulfill the definition of concurrent work. Additionally, a thorough comparison between the two papers yields that they are conceptually different and that our work has substantially novel contributions: - a) Ren et al. investigate the attention scores, i.e., the distribution of the dot product between key and query vectors, finding that attention concentrates more on the beginning token for non-memorized prompts and is more concentrated on certain trigger tokens for memorized prompts. Whereas this is an interesting finding, it is orthogonal to our work. While Ren et al. investigate the model's input space in terms of CLIP embeddings for memorization, we focus on the U-Net itself and identify the neurons in the value mapping responsible for triggering memorization. - b) The findings of Ren et al. are more coarse grained (layer focus rather than neuron focus), which enables them to make a related finding to ours. Still, they do not consider the most interesting part of our research, namely that most of the memorization is triggered by 1 or a handful neurons in the value layers. - c) The inference-time mitigation strategy by Ren et al. is conceptually different from ours. Based on their findings 1 and 2 (both on the input level), where they look at token entropy, their method consists of masking out the summary tokens and increasing the logits of the beginning token. This strategy is very different from ours, which prunes individual neurons from the value mapping in cross-attention layers. Also, their training time mitigation simply removes samples with high attention entropy, which is somewhat related to our detection of out-of-distribution activation in the value layers. Yet, our mitigation strategy acts at inference time, a more realistic and practical scenario given the amount of computing required for training text-to-image models. We, therefore, kindly disagree that our paper's contribution is limited. We also emphasize that the other three Reviewers all agree on the novelty and importance of our research and findings. **2.) Memorization in Other Layers**: It seems there might have been some misunderstanding between the question raised and our response. We fully acknowledge that layers following the value mappings in the U-Net also play a significant role in memorizing training samples. The convolutional and fully connected layers, being integral to the process of constructing images by denoising the initial input, naturally contribute to the memorization of the appearance of training samples. However, the focus of our paper is on identifying the specific signals that trigger the generation and reproduction of memorized data in the first place. Our findings indicate that certain neurons in the value layers of the cross-attention blocks act as such triggers. We do not claim that these neurons are solely responsible for memorization or that other model components are irrelevant to the process. Our research specifically aims to detect these triggering parts within the model, and our analysis identified them in the value layers of cross-attention blocks. While we also analyzed the convolutional layers, we did not find individual channels exhibiting similar trigger behavior, where removing such channels would mitigate memorization. This does not imply that these layers are unimportant for reconstructing memorized samples. We are happy to clarify this point in our paper.
Thanks for your response. For the second question, I would say before K and V, there are also some components that would play some roles in memorization. For example, transformations are conducted from text embedding space to the input space of K and V. That is to say, the memorization is like a chain, while the attetion neurons are one part in the middle. Since it is in the middle, it is not correct to say they trigger memorization. They are only the one that is found to have special behaviors which are connected with memorization. They are actually triggered by their input.
Thank you for pointing that out. We agree that the input prompt serves as the actual initiation of the memorization process, and without a specifically memorized prompt, the model is unlikely to trigger memorization. In this regard, transforming the prompt into an embedding vector is critical and marks the start of this process. If the text encoder would transform the prompt into a different embedding—for instance, by deactivating parts of the model or through fine-tuning—memorization might not occur in the output of the diffusion model. However, as noted in our previous response, we interpret the text embedding as a fixed conditioning factor since the text encoder (CLIP for Stable Diffusion) was trained separately using a contrastive method (on OpenAI’s internal data) and likely never saw the exact prompts memorized by the diffusion model (LAION 2B) during its training. In this scenario, the processing of the text embedding begins with the key and value layers in the U-Net, marking the start of the memorization chain within the U-Net, which is triggered by the input embedding. We will incorporate the perspective and discussion of the memorization chain, starting from the input prompt, into our final paper. We hope this addresses all your questions and concerns regarding the memorization aspect.
Summary
The paper localize memorized samples to neurons in the cross attention layers in diffusion models. By deactivating neurons responsible for memorization, the proposed method enables models to generate diverse images different from training images.
Strengths
1. The authors propose the first method to localize memorization within the model, which is a novel idea that has not been studied in the context of text-to-image generation models before. 2. The experiments and findings presented in the paper are interesting. 3. The paper is well-written and easy to follow.
Weaknesses
1. I understand that detecting memorized propmts is out of scope of this paper. Previous methods in mitigating memorization construct an entire pipeline in which memorized prompts are first detected with a heuristic/metric, and a mitigation strategy is triggered during inference/training if the metric surpasses a certain threshold. My expectation from "localization of memorization" was that there is a specific subset of neurons that, when blocked, avoid replication of training data. Such a localization would not require the authors to develop a pipeline to detect memorization as those neurons can be pruned. (a) From the paper (Table 1 and Figure 4) it seems that the proposed method operates on a per-prompt basis and all experiments have been conducted by separating the memorized prompts into template and verbatim. (b) I believe that although these findings are interesting, one weakness is that the method is impractical as in practice, it would still require one to know the type of memorization before deciding which neurons to block. For some highly memorized prompts, the highly-iterative neuron discovery process (Alg D.1 to D.4) needs to be repeated. This not only requires interference with the inference pipeline but is also computationally inefficient.
Questions
1. How are tokens handled while calculating neuron activation in Equation 3? Do you consider an average or max over tokens? 2. How faithful is SSIM for detecting memorized prompts? Can it be used to detect if a prompt is memorized or not?
Rating
6
Confidence
4
Soundness
2
Presentation
2
Contribution
2
Limitations
The authors have clearly addressed the limitations of the proposed method in the paper. I believe some more limitations are addressed in the Weaknesses section.
Thank you for your prompt response and your inclination toward accepting our paper. We are still eager to convince you even more of our method. **1.)** We respectfully disagree with the concern that our method might not effectively identify the relevant neurons. For our set of memorized prompts, our method successfully identified 200 neurons in the first value layer out of a total of 320 neurons—constituting 62.5\% of the neurons in this layer. The memorization of prompts is actually distributed across multiple different neurons. Importantly, deactivating the neurons associated with the memorization of one prompt does not affect the memorization of other prompts, provided they do not share the same memorizing neurons. Furthermore, if memorization is triggered by neurons in the first value layer, deactivating neurons in subsequent layers does not prevent the generation of memorized samples. This observation extends to deeper layers as well—deactivating neurons in other value layers similarly does not mitigate memorization. Additionally, we have pinpointed neurons within the value layers of the U-Net’s first attention blocks as the primary trigger points of memorization, which represents a relatively small subset of the model's overall parameters. Regarding the hypothesis that a fixed set of neurons responsible for all memorization might exist, we would like to emphasize that the training objective of diffusion models (U-Nets) differs significantly from that of supervised image classifiers or next-token prediction in language models. Consequently, insights into memorization within these architectures may not directly apply to diffusion models. Without evidence or proof of a fixed set of neurons (which our experiments did not support), we believe this does not constitute a limitation of our method. We would also like to highlight that, similar to [1], "we determine the minimum number of neurons required to flip an example’s predictions" [1]. From our perspective, our method extends the line of research on memorization in supervised learning. The results we present for text-to-image models provide insights that are comparable to those from supervised learning, which identify a set of neurons for individual training images. **3.)** Our method does not necessarily rely on Wen et al.'s approach for detecting memorized prompts. As demonstrated in our previous response, our SSIM-based detection method is also capable of identifying memorization. Since it only requires a single denoising step, it can be applied at inference time to detect memorized prompts. We chose to use the evaluation dataset provided by Wen et al. to establish a benchmark and ensure our method is directly comparable to related work. Overall, we are very pleased that you appreciate our paper and the idea of localizing memorization in diffusion models. We will also expand our discussion in the paper to address the detection of previously undiscovered memorization. We thank you for this fruitful discussion.
Summary
This paper introduces NEMO, a method for localizing memorization in diffusion models (DMs) down to the level of individual neurons. The authors empirically evaluate NEMO on the Stable Diffusion model and demonstrate that deactivating the identified memorization neurons effectively mitigates memorization, increases image diversity, and maintains prompt alignment. The paper makes valuable contributions by providing insights into the memorization mechanisms in DMs and proposing a method for permanently mitigating memorization without compromising overall model performance or the quality of the generated images.
Strengths
The paper presents an innovative and practical methodology, NEMO, for localizing memorization in DMs. This methodology allows model providers to identify and deactivate specific neurons responsible for memorizing training samples, effectively mitigating memorization. The empirical findings presented in the paper provide valuable insights into the distribution of memorization neurons and the impact of deactivating these neurons on image generation. The experiments demonstrate that deactivating memorization neurons increases image diversity. The paper is well-structured and provides a clear explanation of the proposed methodology, including the initial selection and refinement steps of NEMO. The authors also provide algorithmic descriptions for each component of NEMO in the appendix and supplement it with extensive experimental results.
Weaknesses
The paper lacks a thorough evaluation of the proposed method's performance compared to existing baselines (if exist) or state-of-the-art methods for mitigating memorization in DMs. It would be beneficial to include comparisons with other methods to assess the novelty and effectiveness of NEMO.
Questions
Can the authors provide more details on the selection of the memorization threshold $τ_mem$? How was it determined based on the pairwise SSIM scores? Did the authors consider any alternative methods for setting the threshold? Have the authors considered comparing NEMO with other existing methods for mitigating memorization in DMs, such as rescaling attention logits or adjusting prompt embeddings? How does NEMO perform in terms of effectiveness and efficiency compared to these methods?
Rating
6
Confidence
4
Soundness
3
Presentation
4
Contribution
4
Limitations
yes
Missed reference and follow-up response
Sorry about the missed reference. Please compare the novelty with Unveiling and Mitigating Memorization in Text-to-image Diffusion Models through Cross Attention. Since they also find the attention score may show connection with memorization. The contribution of this paper seems limited. I disagree with your reponse to weakness 3. From the end of input to the end of output, the models are structued by layer by layer and module by module. Value may memorize text embedding, but the following layer must also memorize the output of attention layer. Otherwise the final output cannot be memorized. It is like a chain. Memorization will not work without the any one of the layer. The memorization is never something that is totally dependent on one neuron or a few neurons. In other words, these neurons are just some clues or signals of memorization. It is not correct to say they control memorization. It is correct to say value or kay layers has a very special behavior that is related to memorization. But they didn't decides the whole thing. In other layers like the convolutional layers in u-net, they must also memorize their input. That memorization is not probabily about one neuron. It is maybe need many neurons for a memorization. The author cannot say something they didn't find does not exist. If the author acknowledge this and update in their future version, I will increase my score. Otherwise I would possibily decrease my score.
I would suggest to list your response. Otherwise it would increase the cost to write the response for reviewers, which may decrease the possibility to get a response.
Thank you for the rebuttal
Thank you for the effort you put into this rebuttal. I appreciate the authors' detailed clarifications on some things I might have misunderstood but now I have more questions. > While the runtime varies based on the strength of the memorization, for verbatim memorized prompts, which represent the strongest type of memorization, the runtime was only a few seconds (as discussed in our limitations section in Appendix A). 1. Can you quantify the time taken to discover these memorization neurons? I would like to see the time taken to run Alg 1,2,3,4 separately. 2. I am a bit confused about the pipeline now as this is not explicitly stated in the paper. Do you (1) first find memorization neurons for every memorized prompt, and then prune the union of these memorization neurons from the model permanently? or (2) For every prompt, find memorized neurons, prune the model, evaluate, and then reset the model?
Dear Reviewer zaL8, Thank you for your quick response and appreciating our rebuttal. We are happy to provide additional information and clarifications. 1.) We will run some additional experiments for your first question regarding the runtime of individual parts of our algorithm and provide them as soon as they become available. 2.) In our paper, we followed the first approach you sketched. We identify the neurons of each memorized prompt individually and then deactivate/prune the union of these neurons in the model. The metrics in the paper are computed on the U-Net, with all identified neurons removed simultaneously.
Additional Runtime Results
Dear Reviewer zaL8, Thank you for your question. We have rerun the runtime experiments and measured the time for each algorithm D1-D4. Please find the values in seconds in the following table: | | D1 | D2 | D3 | D4 | Total | |----|------|------|------|-------|-------| | VM | 0.21 | 0.29 | 2.31 | 12.85 | 15.68 | | TM | 0.18 | 0.25 | 6.09 | 39.37 | 45.90 | In the table above, one can clearly see that the algorithms D1 and D2 take less than a second to run, while D3 takes roughly 2 seconds for verbatim memorized prompts and 6 seconds for template memorized prompts. As can be seen, D4 has the longest runtime. This is because each candidate neuron is checked individually for memorization by iterating over them and setting their output to zero individually. For template memorized samples, this takes a bit longer since, for this kind of memorization, our initial neuron candidate set is larger than for verbatim memorized samples. However, even though this process takes a bit longer for TM than for VM, the overall runtime is still only 45 seconds. Additionally, we want to emphasize here again that our algorithm only has to be run once. In contrast to existing mitigation approaches, once the memorization neurons are identified, the inference time, as well as the inference memory consumption, will be the same as that of the default diffusion model.
Thank you for the rebuttal
Thank you for clarifying my doubts. I have a much better understanding of the paper. I appreciate that the authors calculated the time required to run their algorithm and I encourage them to add this analysis in the paper. I still have some concerns listed below. 1. The authors said that > While some sets of neurons memorize multiple prompts, no fixed set of neurons is responsible for memorizing all prompts. This is not a weakness of our method, but rather a reflection of the behavior of already trained models. This may be due to the method you proposed for searching neurons not being able to identify a fixed set of neurons. Previous studies have successfully localized memorization to a specific region of the model in image classifiers [1] and language models [2], suggesting that it could also be achievable in a diffusion model. The proposed method may lack the capability to accomplish this. [1] Pratyush Maini, Michael C. Mozer, Hanie Sedghi, Zachary C. Lipton, J. Zico Kolter, and Chiyuan 385 Zhang. Can neural network memorization be localized? [2] Localizing Paragraph Memorization in Language Models, Stoehr et al 3. A significant limitation I see in the paper, which could be quite restricting, is that the proposed method requires observing all memorized prompts extracted for SD using a specific extraction method—in this case, the approach outlined in Wen et al.'s paper. However, this does not represent an exhaustive set of prompts that the model has memorized. Given that their method does not rely on identifying a fixed set of neurons, I am concerned that it may not extend to undiscovered prompts. The authors should consider explicitly stating this limitation in their paper as a disclaimer to avoid giving a false sense of security. Overall, I will lean toward a positive rating (and keep my rating) because I like the idea of localizing memorization in DMs and this paper presents interesting insights.
I appreciate the authors' detailed rebuttal and additional insights. Most of my questions have been addressed, and I will be maintaining my original rating.
Thank you for your detailed rebuttal. We're glad that you are leaning towards accepting the paper. If you have any further questions or comments, we would be happy to provide additional information.
Have concerns been addressed?
Dear Reviewers, We sincerely appreciate your questions and feedback, which have helped us to further improve our paper. We would like to follow up on our responses to ensure that they adequately address your concerns. If there are any additional questions or areas where further clarification is needed, please don't hesitate to reach out. Thank you very much for your time and effort. Best regards, The Authors
Thank you for your detailed rebuttal and for the engaging and fruitful discussion. Your insights will definitely enhance our paper. We're also pleased to see that you have raised your score in favor of accepting the paper.
Decision
Accept (poster)