Unveiling and Mitigating Memorization in Text-to-image Diffusion Models through
Cross Attention

Denoising Diffusion Probabilistic Models (DDPM) [10] typically involves a forward diffusion process and a reverse diffusion process. The forward process is a $T$-step Markov chain which transforms a data point $x_{0}$ from the target image distribution to a random Gaussian distribution. It introduces a small Gaussian noise at each step into the data point following

where $\beta_{t}$ is the variance schedule. The reverse process generates new images following

where each reverse step follows a Gaussian distribution with $\mu_{\theta}$ and $\Sigma_{\theta}$ as some mean and variance. With a parameterized denoising network $\epsilon_{\theta}$, following [10], the generation process can be explicitly expressed as

where $\alpha_{t}=1-\beta_{t}$, $\bar{\alpha}_{t}=\prod_{i=1}^{t}\alpha_{t}$, $\sigma_{t}$ can be $\sqrt{\beta_{t}}$ or $\sqrt{\frac{1-\bar{\alpha}_{t-1}}{1-\bar{\alpha}_{t}}\beta_{t}}$, and $w\sim\mathcal{N}(\mathbf{0},\mathbf{I})$. To connect Eq. (2) with (3), $\mu_{\theta}(x_{t},t)=(x_{t}-\frac{\beta_{t}}{\sqrt{1-\bar{\alpha}_{t}}}\epsilon_{\theta}(x_{t},t))/\alpha_{t}$, and $\Sigma_{\theta}(x_{t},t)$ is determined by $\sigma_{t}$. Note that when generating new images in reverse process, we start from a random noise $x_{T}$ and end at step $0$ with the output as $x_{0}$.

To guide the image generation process using extra conditions, Rombach et al. [23] further propose Latent Diffusion Models (LDMs). Two changes are applied in LDMs compared to a vanilla diffusion model. First, instead of directly generating images, the diffusion process proceeds with a low-dimensional latent representation of the image via the compression using a variational autoencoder (VAE) [12]. In this case, $I_{0}$ is the image, $\mathcal{E}$ is the image encoder of VAE, and we feed the latent representation $x_{0}=\mathcal{E}(I_{0})$ in diffusion models. Second, to introduce an extra condition $c$, $c$ works as an input of $\epsilon_{\theta}$ to guide the output of $\epsilon_{\theta}\left({x}_{t},t,c\right)$.

Extending from the basic framework of LDMs, Stable Diffusion (SD) takes textual prompts as condition $c$ to generate image described by given texts. SD first uses the text encoder of CLIP [21] to get the prompt’s embeddings, $e_{c}$, then uses the cross-attention mechanism to select the information and feed it into each of the hidden layers, $z_{t}$, in the U-Net backbone of diffusion models. In the cross-attention architecture, $Q(z_{t})$, $K(e_{c})$ and $V(e_{c})$ are the three components connecting the prompt with the image. The terms $Q(z_{t})$ and $K(e_{c})$ decide the attention score allocated on each token embedding by

The term $V(e_{c})$ is the textual information to be selected, and is further multiplied with $A$ in the cross attention and passed to the U-Net of the diffusion model. Besides, instead of a single-head cross attention, one may create multiple groups of $(Q,K,V)$ and merge them to a multi-head cross attention.

When training the diffusion model, the training data consists of images and their paired textual prompts. In addition to the U-Net in the diffusion model, the cross attention is also trained to match the textual input and the generated image. The only components fixed in the training are the image encoder (VAE) and the text encoder, using which we can obtain the image latent representation, $x_{0}$, and the prompt embeddings, $e_{c}$.

In this paper, we focus on the attention score $A$ to understand and mitigate the memorization. Each row of $A$ is a set of attention score for a hidden dimension of $z_{t}$ which assigns weights on the embeddings of different tokens of $V(e_{c})$. Then the weighted information is aggregated to guide the generation process. Thus, the sum of attention on each row, i.e. for each hidden dimension of $z_{t}$, is 100%. For multi-head attention, each head calculates its corresponding attention score matrix respectively, and the sum of attention in each row of each head is 100%.

To facilitate the analysis of the connection between cross attention and memorization, we divide the tokens of textual prompts into three categories: beginning token, prompt tokens, and summary tokens. For example, in the prompt, “<begin> A picture of mountains covered by snow <end> <padding> <padding> … <padding>”, the text encoder adds <begin> before the prompt tokens, and <end> <padding> after the prompt tokens. Given that the text encoder operates causally, the beginning token does not hold any semantic information derived from the textual prompt, as it exists before the entirety of the prompt. In contrast, for <end> and <padding>, we refer them as “summary tokens” since the embeddings of these tokens encapsulate the semantic information of all preceding tokens.

From the attention maps of Fig. 1 in Sec. 1, we find that the beginning token shows different patterns for samples with and without memorization. In non-memorization, although the attention distribution is disperse, the beginning token holds a high attention score, especially in Fig. 1(d). Instead, in memorization, the beginning token is assigned with very little attention. To facilitate the understanding of memorization in the following sections, we first demonstrate the pattern of the beginning token’s attention score in this subsection. We demonstrate the averaged attention score of the beginning token through all the diffusion steps in Fig. 3. It can be seen that the beginning token has an attention score higher than 80% in all the diffusion steps. Since the beginning token contains no semantic information, this implies that a small portion of the attention score is sufficient to collect the semantics from other tokens, especially in the late steps with a small $t$. Besides, in Fig. 3, the attention of beginning tokens increases from step $T$ to step $0$. Intuitively, after the main body of the image is generated with a large $t$, the diffusion model mainly focuses on denoising when $t$ is small, which does not need much textual information from prompts.

Fig. 1 in Sec. 1 demonstrates that the memorized samples tend to allocate most of the attention to the embeddings of specific tokens. In this subsection, we further quantify this finding by defining attention entropy.

Recall that in cross attention, each hidden dimension in $z_{t}$ will allocate attention scores on the tokens. The sum of attention scores for each hidden dimension in each attention head is 100%. Thus, we can consider it as a discrete probabilistic distribution. Meanwhile, entropy is usually used to measure the uncertainty and dispersion of probabilistic distributions. Thus, the entropy of attention can measure the dispersion of attention [37, 31]. It is defined as follows

where $N$ is the number of tokens, $t$ is diffusion step, and $\overline{a}_{i}$ is the average attention score on the $i$-th token. Unless otherwise stated, $\overline{a}_{i}$ is averaged across all U-Net layers, attention heads and hidden dimensions. Higher entropy indicates more disperse attention distribution.

In Sec. 3.4, we find that the cross-attention score of non-memorization will gradually concentrate on the beginning token. In contrast, as shown in Fig. 1, for the prompts of memorization, more attention will be concentrated on the embeddings of specific prompt tokens and summary tokens throughout all the steps in the generation process, referred to as trigger tokens. Meanwhile, the memorized samples will distract from the beginning token. This leads to a disperse attention distribution. Therefore, we expect that the attention entropy is low for non-memorization samples, and high for memorized ones.

We demonstrate the entropy of different diffusion steps in Fig. 3 to verify our intuitions. In non-memorization samples, the entropy reduces fast from step $T$ to step $0$, which is consistent with the fact that the attention is gradually concentrated into the beginning token. On the other hand, for memorized samples, the entropy becomes higher than non-memorization especially when $t$ is small, which also aligns with our observation that the model pays high averaged attention scores on trigger tokens. One exception is that, when $t$ is large, e.g. $t=T$ in Fig. 1(c), compared to the memorization case, the non-memorization prompts also allocate higher attentions on prompt and summary tokens, resulting in a higher entropy in Fig. 3. To understand this, although non-memorization in general imposes the highest weight on the meaningless beginning token, to generate proper images, it still needs to correctly understand the input prompt. As a result, for large $t$, the model focuses on collecting semantic information from different tokens, and the corresponding entropy is high. In contrast, with memorization, the model only needs to focus on the trigger tokens, and it is unnecessary to further collect information from other tokens, resulting in a lower entropy for larger $t$.

With the quantification of the attention score via entropy, we verify our finding about the disproportionate attention on trigger tokens.

Existing studies suggest that there exist different types of memorized prompts in diffusion models [32]. In this subsection, we discuss cross attention behaviors on different types of memorized prompts, and provide a deeper understanding about the connection between cross attention and memorization.

In literature, Webster [32] provides the dataset of memorized prompts and divides them into three types:

• •

  Matching memorization (MM): one memorized prompt generates one image that is exactly matching with the original paired training image.

• •

  Retrieval memorization (RM): one memorized prompt generates images that match with a subset of training images.

• •

  Template memorization (TM): a group of prompts generate images highly align with training images, but may have variations in colors or styles.

We demonstrate the sum of attention score on prompt and summary tokens respectively in Fig. 4 to examine to what extent the model focuses on these two types of tokens, and also give prompt examples of each types in Fig. 4. We make the following two observations from Fig. 4.

First, MM pays more attention to summary tokens compared with RM and TM. In MM, the memorized prompts are totally different from each other. Since the summary tokens contain the semantics of the whole sentence, overfitting to the summary tokens can memorize the unique prompts and connect them with the paired training images easier. In contrast, for RM and TM, the memorized prompts share the same sentence template and a few identical tokens in a group of prompts. The model overfits these shared tokens and overlooks the difference.

Second, all three types of memorization have a significantly slower reduction of summary-token attention score than non-memorization from step $T$ to step $0$. It implies that non-memorization distracts from summary tokens faster. We conjecture that this is because different parts of the non-memorized generated images focus on different semantics of the prompt rather than the summary of the prompt. For example, in the prompt “Two dogs playing on the grass”, some pixels focus on “dogs”, and some focus on “grass”. But for memorization, it does not distinguish prompt tokens or summary tokens like non-memorization since the overfitting of trigger tokens can simply guide each pixel in the image. This characteristic can distinguish between memorization and non-memorization generations.

According to Finding 1 in Sec. 4.1, the entropy is better distinguished between memorization and non-memorization in step $0$ than step $T$. This means the whole generation is involved since step $0$ is the last step of the generation. However, in this subsection, we find that different U-Net layers have different behaviors for memorization, and in some active U-Net layers, the memorization is already distinguished at the first step, i.e. when $t=T$.

In SD, there are 16 U-Net layers including down-sampling, middle and up-sampling layers. The cross-attention mechanism has different sets of multi-head $(Q,K,V)$ for each layer. These sets collect different semantic information and feed them into each layer. We find that the concentration on special tokens is not uniform on all the layers. We show the entropy of each layer in the first diffusion step, i.e. $t=T$, in Fig. 5. In the first diffusion step, the entropy of distinct U-Net layers have different overlapping between memorization and non-memorization. The fourth and fifteenth layers have clearer separation, which can distinguish memorization and non-memorization better. This shows that using the first diffusion step has the potential to detect the memorized samples in one step.

Finding 1 to 3 in Section 4 demonstrate different attention distributions between memorization and non-memorization. In this subsection, we introduce two metrics guided by our findings to detect the memorization.

The first metric is designed based on Finding 1 and Finding 2 as:

where $E_{t}^{\text{summary}}=\sum_{i=N-S+1}^{N}-\bar{a}_{i}\log\left(\bar{a}_{i}\right)$ is the entropy calculated only on the summary tokens and $S$ is the number of summary tokens. The term $T_{D}$ is the number of generation steps used in calculating $D$. For SD, we choose $T_{D}=\frac{T}{5}$, which means we use the last $\frac{T}{5}$ steps of the reverse diffusion process to calculate $D$. The first term is inspired by Finding 1 that trigger tokens will cause higher entropy especially than non-memorization when $t$ is close to $0$. The second term is from Finding 2 that the sum of attention score of summary tokens in memorization will reduce slower than non-memorization. In Eq. 6, we use $E_{t}^{\text{summary}}$ to replace sum to make its numerical values comparable with the first term, otherwise it becomes a negligible term. In Appendix 0.B, we show that the sum and $E_{t}^{\text{summary}}$ have consistent conclusion with Finding 2.

The second metric is motivated by Finding 3, which is defined as

where $\bar{a}_{i}^{l}$ is the averaged attention on $l$-th layer of U-Net of the first step in the reverse diffusion process, i.e. $t=T$. The major difference between the two metrics is that the second metric only requires the first step of diffusion process. Although the second requires an additional hyper-parameter $l$, it allows the model builders to detect the memorization instantly, which can save time and computation before adapting the follow-up process like mitigation.

In this subsection, we propose a method to mitigate memorization at the inference stage by reducing the weight on trigger tokens. This method consists of masking out the summary tokens and increasing the logits of the beginning token.

First, from Finding 2, we know that memorization is highly related to the summary tokens, while attention score on summary tokens drops quickly in non-memorization. Consequently, to control memorization, we mask out these summary tokens.

Second, based on Finding 1, trigger tokens tend to have the large attention scores except for the beginning token. As a result, we reversely enlarge the attention score for the beginning token via adjusting the input logits of the softmax operator in the cross attention. To be specific, denote the original input logits as $\boldsymbol{s}=(s_{1},s_{2},...,s_{N})$ where $s_{i}$ is the logit of the $i$-th token, the re-scaled logit vector $\boldsymbol{s}^{\prime}$ is

where $C$ is a factor to be applied on the beginning token. When $C>1$, the attention score of beginning token is increased and the attention scores of other tokens become smaller.

While our primary interest on $C>1$ is to enlarge the attention score of the beginning token, we also note that the re-scaling operation in Eq. 8 will also give a larger reduction on the tokens with larger attention scores. To explain this, when taking the gradient of each attention score w.r.t. $C$, we obtain

Since ${s_{1}e^{Cs_{1}}}/{(e^{Cs_{1}}+\sum_{j=2}^{N}e^{s_{j}})^{2}}$ is the same for $\forall i\in[2,N]$, the reduction ratio depends on $-e^{s_{i}}$, which means token with larger attention will be reduced more.

By combining the above two strategies, our final adjusted logits become

This method only requires applying a mask and a re-scaling factor to the logits, and does not include any extra computation. Nonetheless, it can effectively mitigate memorization during inference and has little impact on generation quality.

Remark. In the experiments in Fig. 6(a), we show that, although detection has almost no extra computation, it is not necessary to detect before mitigation since the mitigation does not compromise the generation quality for both memorization and non-memorization.

During training, we can mitigate memorization by removing the samples whose attention entropy is higher than a pre-defined threshold from the mini-batch. This idea is similar to [36], in which it filters out the samples if their outputs of the diffusion model are significantly different from these of the empty prompt. However, the method in [36] uses additional inference operation of the empty prompt, which requires more computation. In contrast, our approach eliminates the need for extra inference. We only compute the attention entropy as defined by Eq. 5 and remove the samples with high-entropy, which incurs a negligible computational cost. This alone suffices to effectively reduce memorization.

Diffusion models and datasets. We conduct the experiments using SD v1.4 and SD v2.0. We use the dataset of memorized prompts extracted by Webster [32] as memorized samples, and use 500 prompts generated by ChatGPT-4 [1] as non-memorized samples. For the fine-tuning data in training-time mitigation, since duplicated data can cause memorization, we use 200 text-image pairs duplicated 50 times as memorized data and 20,000 captioned images from LAION [26] as non-memorized data, following the procedure of [36].

Baselines. For detection, we use two baselines for comparison. The first is the method from Carlini et al. [5]. They observe that the model generates the same image for the memorized prompt when using different random seeds. Thus they propose to detect memorization by the $l_{2}$ distance of the images generated using different $n$ seeds. The second is from Wen et al. [36]. Their detection method is built upon the observation where the difference of the output between memorized prompt and empty prompt is larger than the difference between non-memorized and empty prompt. To stabilize the method, they generate the image for $n$ times. For mitigation, we compare with a the method from Wen et al. [36]. More details can be found at Sec. 0.C.1.

Evaluation metrics. The detection metrics include the area under the receiver operating characteristic curve (AUROC), the true positive rate at the false positive rate of 3% (TPR@3%FPR) and the time cost. The mitigation metrics include Similarity Score [20, 28] that measures the degree of memorization by the similarity between generated images and original training images (higher similarity means more severe memorization); Fréchet Inception Distance (FID) [9] that measures the generation quality in realism and diversity (lower FID indicates better quality); CLIP score [21] that measures fine-tuning performance (higher CLIP score means better fine-tuning performance); and the time cost. Details can be found at Sec. 0.C.2.

In this subsection, we present the detection performance of the two proposed methods, $D$ and $E^{l}_{t=T}$, in both accuracy and effeciency with Tab. 1. For all the detection, we use $l=4$.

To compare our method $D$ with the benchmarks, we can see that it costs 1.745 seconds on SD v1.4 and 4.220 seconds on SD v2.0, which is more efficient than most of the methods except the two fast settings $(s=1,n=1)$ and $(s=1,n=4)$ from Wen et al. [36]. However, comparing to these fast settings of Wen et al. [36], our AUROC and TPR are much better, especially on SD v1.4 with almost perfect AUROC of 0.9997 and TPR of 0.997.

In terms of our method $E^{l=4}_{t=T}$, it is the most efficient one and achieves good performance with AUROC higher than 0.993 and TPR higher than 0.977 for both SD v1.4 and v2.0. On SD v2.0, our method $E^{l=4}_{t=T}$ is the best one in both AUROC and detection speed. It is even the only method that achieves an AUROC higher than 0.99 on SD v2.0. On SD v1.4, although some of the other methods, e.g., $(s=50,n=32)$ of Wen et al. [36], are better in AUROC and TPR than $E^{l=4}_{t=T}$, their time cost is 20 to 400 times of the method $E^{l=4}_{t=T}$.

In this subsection, we show that our mitigation methods can effectively reduce the memorization without sacrifice of the speed and the generation quality.

For inference-time mitigation, Fig. 6(a) presents the results of similarity score and FID with and without mitigation. The FID is calculated between generated images (including both memorized prompts from [32] and 500 non-memorized prompts) and a subset of LAION [26] with 10,000 images. Similarity Score is calculated between the generated images by memorized prompts and the training images paired with the memorized prompts. When there is no mitigation, Similarity Score is 0.7. Our method can significantly reduce Similarity Scores from 0.7 to the range of 0.25 to 0.3 (by setting $C$ from 1.1 to 1.25). Interestingly, FID is also improved after mitigation. This is likely attributable to the fact that RM and TM tend to produce memorized images that are similar to each other. Our mitigation introduces a greater diversity, thereby contributing to a reduced FID. Although the method of Wen et al. [36] has similar mitigation effect, it is slower than ours by 20% as shown in Fig. 6(b). Notably, our method maintains the same computational efficiency as performing inference without any mitigation.

For training-time mitigation, we compare the mitigation effect and the fine-tuning performance with the baseline method from Wen et al. [36] in Fig. 6(c). The fine-tuning performance is measured by CLIP score which estimates the alignment between prompts and images. Higher CLIP score after fine-tuning indicates that the fine-tuned model can generate images that aligns the prompts better. Both our method and the baseline effectively reduce similarity to around 0.2, with negligible impact on CLIP score. However, the method of Wen et al. [36] incurs a 10% increase in training time. In contrast, our approach maintains the same training duration as scenarios without mitigation as shown in Fig. 6(d).

In this subsection, we conduct ablation study on the two components of inference-time mitigation, i.e. logits re-scaling and summary token mask. We present a set of generated examples in Fig. 7. Our method is shown in Fig. 7(f) with summary tokens masked out and re-scaling factor $C=1.25$. It can effectively mitigate the memorization compared with no mitigation in Fig. 7(b). However, if we only apply attention mask or logits re-scaling as shown in Fig. 7(c) and Fig. 7(d), some memorized samples cannot be eliminated. If we increase $C$ to 2, although the memorization can be prevented, it might cause the loss of some semantic information like the identities and the appearance of the specific celebrities in Fig. 7(e). Due to the difference of formulations between Finding 1 and Finding 2, the two components of our inference-time mitigation method may handle different memorized samples as shown above. On the other hand, when applying them together, they can mitigate more memorization samples without lost of information from prompts. This ablation study implies the benefits of combing the two components of our method.