Detecting, Explaining, and Mitigating Memorization in Diffusion Models

TL;DR

This work tackles memorization in diffusion-based image generation by proposing a lighting-fast detector that uses the magnitude of text-conditional noise predictions to identify memorized prompts, even from a single generation. It provides an explainable token-level analysis to rank trigger tokens and offers two mitigation strategies: inference-time perturbation of the prompt embedding and training-time exclusion of memorized samples, both designed to preserve generation quality. Empirical results show high detection performance, with $\text{AUC}$ approaching $0.999$ and $\text{TPR}@1\%\text{FPR}$ near $0.99$ under favorable settings, while achieving substantial speedups over prior methods. The practical impact lies in enabling model owners to detect and counter memorization efficiently, with interactive guidance for users to adjust prompts and robust defenses during inference and training. Code is provided to reproduce the results and apply the methods to real-world diffusion models.

Abstract

Recent breakthroughs in diffusion models have exhibited exceptional image-generation capabilities. However, studies show that some outputs are merely replications of training data. Such replications present potential legal challenges for model owners, especially when the generated content contains proprietary information. In this work, we introduce a straightforward yet effective method for detecting memorized prompts by inspecting the magnitude of text-conditional predictions. Our proposed method seamlessly integrates without disrupting sampling algorithms, and delivers high accuracy even at the first generation step, with a single generation per prompt. Building on our detection strategy, we unveil an explainable approach that shows the contribution of individual words or tokens to memorization. This offers an interactive medium for users to adjust their prompts. Moreover, we propose two strategies i.e., to mitigate memorization by leveraging the magnitude of text-conditional predictions, either through minimization during inference or filtering during training. These proposed strategies effectively counteract memorization while maintaining high-generation quality. Code is available at https://github.com/YuxinWenRick/diffusion_memorization.

Figures (9)

• Figure 1: Memorization vs. non-memorization generation. We display the magnitude of text-conditional noise prediction at each time-step, as described in \ref{['sec:detection_method']}, for all four generations distinctly (with 4 different random seeds) for each prompt. As illustrated in the first two rows, the metric typically indicates a higher value when memorization occurs. On the other hand, the normal generations, represented in the third row, consistently exhibit significantly lower metric values.
• Figure 2: Statistics of the magnitude of text-conditional noise predictions.
• Figure 3: By modifying the trigger tokens, memorization can be effectively mitigated. The significance score for each token is illustrated in a histogram in Appendix \ref{['fig:token_significance_extended']}. The two most significant tokens are highlighted in red and blue. A green arrow indicates modifications made to the top-2 tokens, while a white arrow represents changes to less significant tokens.
• Figure 4: Mitigation results. A lower similarity score suggests reduced memorization, whereas a higher clip score denotes a better alignment between the generation and the prompt.
• Figure 5: Mitigation results with different mitigation strategies during inference and training phase.