Assessing Open-world Forgetting in
Generative Image Model Customization

Diffusion models are a class of generative models that generate data by gradually denoising a sample from a pure noise distribution. The process is modeled in two stages: a forward process and a reverse process. In the forward process, Gaussian noise is iteratively added to data samples, typically over $T$ steps, forming a Markov chain. At each step, the transition is defined as:

$$q(\mathbf{x}_{t}|\mathbf{x}_{t-1})=\mathcal{N}(\mathbf{x}_{t};\sqrt{1-\beta_{t}}\mathbf{x}_{t-1},\beta_{t}\mathbf{I}),$$

(1)

where $\beta_{t}$ controls the noise schedule.

The reverse process denoises the data by learning the conditional probability $p_{\theta}(\mathbf{x}_{t-1}|\mathbf{x}_{t})$, typically parameterized by a neural network $\epsilon_{\theta}(\mathbf{x}_{t},t)$ that predicts the added noise at each step. The model is trained by minimizing a simple mean-squared error between the true and predicted noise:

$$L(\theta)=\mathbb{E}_{t,\mathbf{x}_{0},\epsilon}\left[\left\|\epsilon-\epsilon_{\theta}(\mathbf{x}_{t},t)\right\|^{2}\right].$$

(2)

Text-to-image diffusion models employ an additional conditioning vector $\mathbf{c}=\mathcal{E}(P)$ generated using a text encoder $\mathcal{E}$ and a text prompt $P$. These models have gained prominence for their ability to generate high-quality, diverse samples, often outperforming other generative models like GANs and VAEs in terms of mode coverage and sample quality (Ho et al., 2020; Song et al., 2021).

Diffusion models often require finetuning for specific domains or user needs. This involves introducing new conditioning mechanisms or retraining on specialized datasets. This paper applies two adaptation methods to evaluate finetuning’s impact on image generation models.

Dreambooth (Ruiz et al., 2023b) enables personalization of diffusion models by finetuning them with a small set of images. It reuses an infrequent token of the vocabulary to represent a unique subject, allowing the model to generate images of the subject in varied contexts or styles. This approach induces language drift and reduced output diversity in the model, which is mitigated by replaying class-specific instances alongside the subject training, called prior preservation loss. The final training objective reads

$$\mathbb{E}_{\mathbf{\epsilon},\mathbf{x},\mathbf{c},t}[w_{t}\lVert\mathbf{\epsilon}-\mathbf{\epsilon}_{\theta}(\mathbf{x}_{t},\mathbf{c},t)\rVert+\lambda w_{t^{\prime}}\lVert\mathbf{\epsilon}^{\prime}-\mathbf{\epsilon}_{\theta}(\mathbf{x}_{t^{\prime}}^{\text{pr}},\mathbf{c}^{\text{pr}},{t^{\prime}})\rVert],$$

(3)

where $\lambda$ is a weighting parameter, and $\mathbf{x}_{t}^{\text{pr}}$ and $\mathbf{c}^{\text{pr}}$ come from the prior dataset. DreamBooth is especially useful for personalized content generation where subject fidelity is critical.

Custom diffusion (Kumari et al., 2023) is another approach aimed at efficiently finetuning diffusion models with minimal data and compute. This method observes that the cross-attention layer parameters undergo the most change during personalization, so they propose to only update the key and value projections in these layers. It introduces a token into the text-encoder representing a unique subject, rather than reusing an old one. By freezing the majority of the model’s parameters and focusing updates on a few key layers, Custom Diffusion facilitates rapid customization with less degradation in image quality. Prior preservation loss is maintained, since language drift is still experienced otherwise.

Customized Model Set In our experiments, we will evaluate Dreambooth and Custom Diffusion. We adapt both these models to ten different concepts based on 5 images per concept. The concepts are ‘lamp’, ‘vase’, ‘person2’,‘person3’,‘cat’,‘dog’,‘lighthouse’,‘waterfall’,‘bike’ and ‘car’ taken from CustomConcept101 (Kumari et al., 2023). We will refer to these ten models for both DreamBooth and Custom Diffusion as the Customized Model Set.

The two studied approaches, Dreambooth (Ruiz et al., 2023b) and Custom Diffusion (Kumari et al., 2023), apply finetuning to adapt to the new data: they mainly focus on how good the learned model is on the target data, and do not study the possible detrimental effects for other classes. The Dreambooth method includes a method called prior regularization, which by replaying general instances of the concept being learned (see Eq. 3), helps to prevent the model from overfitting to the new data and ensures that the representation of the superclass remains stable. This same mitigation strategy is also applied in custom diffusion (Kumari et al., 2023).

In this paper, we propose another regularization technique that can be applied during new concept learning. The method is remarkably simple and is motivated from continual learning literature. This field has proposed a variety of methods to counter forgetting during the learning of new concepts (De Lange et al., 2021). Regularization methods aim to regularize the learning of new concepts in such a way that it does not change weights which were found relevant for previous tasks. The field differentiates between parameter regularization methods, like EWC (Kirkpatrick et al., 2017) which directly learn an importance weight for all the network parameters, or functional (or data) regularization, like Learning-without-Forgetting (Li & Hoiem, 2017; Pan et al., 2020) which regularizes the weights indirectly by imposing a penalty on changes between the (intermediate) outputs of a previous and current model.

We propose to apply a functional regularization loss to the network during the training of new concepts. Our loss, called drift correction loss, constrains the difference between the outputs of the pre-trained and fine-tuned models when the new concept is not present in the prompt. It has the following form:

$$\mathbb{E}_{\mathbf{\epsilon},\mathbf{x},\mathbf{c},t}[w_{t}\lVert\mathbf{\epsilon}-\mathbf{\epsilon}_{\theta}(\mathbf{x}_{t},\mathbf{c},t)\rVert+\lambda w_{t^{\prime}}\lVert\mathbf{\epsilon}_{\theta^{*}}(\mathbf{x}_{t^{\prime}}^{\text{pr}},\mathbf{c}^{\text{pr}},{t^{\prime}})-\mathbf{\epsilon}_{\theta}(\mathbf{x}_{t^{\prime}}^{\text{pr}},\mathbf{c}^{\text{pr}},{t^{\prime}})\rVert],$$

(4)

where the second term is the distillation loss, $\lambda$ is a relative weighting parameter and $\mathbf{\epsilon}_{\theta^{*}}$ is the base model. This loss helps to maintain consistency in the model’s internal representations while allowing it to learn new information effectively. For the training process, we choose instances from the same class as the concept being learned, similar to those used by prior regularization. The change between our proposed drift correction method (Eq. 4) and the existing prior regularization (Eq. 3) is that we do not require the finetuned network to estimate the true forward noise, but instead we want it to estimate the same noise as the original starting network. We will see that this small change significantly improves stability and mitigates forgetting.

In our evaluation, we provide results for DreamBooth (DB) which includes the prior regularization, for Dreambooth with Drift Correction (DB-DC) which also includes the prior regularization and for DreamBooth with Drift Correction without the prior regularization (DB-DC$\backslash$pr). Similarly, we show results for the various variants of Custom Diffusion (CD, CD-DC, and CD-DC$\backslash$pr).

Semantic drift refers to alterations in a model’s representation that cause the generation of semantically divergent content following customization. In the experiment depicted in Figure 2a, almost all prompts exhibit some level of drift, with a notable long tail of highly dissimilar generations. Many of these pronounced deviations have resulted in the generation of content that semantically no longer align with the input prompt.

To evaluate how semantic drift affects generative models, we use a straightforward approach: we utilize the model’s internal representations on different classification tasks (Mittal et al., 2023; Tang et al., 2023). It is based on a recent insight that showed that diffusion models can be directly applied for zero-class classification, by leveraging the conditional likelihood estimation of the model. Concretely, we use Diffusion Classifier (Li et al., 2023), where a posterior distribution over classes $\{\mathbf{c}_{i}\}_{i=1}^{N}$ is calculated as:

$$p_{\theta}(\mathbf{c}_{i}\mid\mathbf{x})=\frac{\exp\{-\mathbb{E}_{t,\epsilon}[\|\epsilon-\epsilon_{\theta}(\mathbf{x}_{t},\mathbf{c}_{i})\|^{2}]\}}{\sum_{j=1}^{N}\exp\{-\mathbb{E}_{t,\epsilon}[\|\epsilon-\epsilon_{\theta}(\mathbf{x}_{t},\mathbf{c}_{j})\|^{2}]\}}.$$

(5)

Monte Carlo sampling is performed over $t_{i}$ and $\epsilon$ to obtain a classifier from a text-to-image model $\epsilon_{\theta}$.

While this method offers a simple, parameter-free approach to evaluating semantic drift, it is worth noting that alternative techniques have been proposed to assess the representation space of diffusion models. These include linear probing on activations (Xiang et al., 2023), analysis of hierarchical features (Mukhopadhyay et al., 2023), and methods requiring a preliminary likelihood maximization stage (Chen et al., 2023a). However, these alternatives often involve additional computational steps or are subject to specific settings, potentially limiting their applicability or introducing complexity to the evaluation process.

We conducted zero-shot classification experiments across multiple datasets spanning diverse domains to quantify the semantic drift of the models. We perform two measurements. First, we measure the average zero-shot classification score for the various models (the results are averaged over the 10 models of the Customized Model Set). Second, we establish the performance of the original pretrained model as the baseline, and measure the presence of semantic drift by calculating the drop in accuracy from the baseline. We also report the worst class drop which is the drop in accuracy of the class that has suffered the largest deterioration due to the adaptation. For further details on the classification method, please refer to Appendix B.2.

The results in Table 1 are surprising, average zero-shot classification accuracy drop significantly on all the datasets: adapting a huge generative image foundation model to just five images of a new concept has a vast impact throughout the latent space of the diffusion models. When applying DreamBooth, average zero-shot performance drops by over 4% on CIFAR10, Pets, Food and Aircraft. If we look at individual classes, the impact can be much larger. As indicated by the worst class drop, for some classes, zero-shot performance drops by over 60% (e.g. ‘vacuum cleaner’ gets recognized as ‘microwave’, ‘drill’ or ‘laptop’). We show that these drops in performance are mitigated to a large extent by our alternative Drift Correction results (see DB-DC and CD-DC results) and their average zero-shot classification scores are in general within 1% of the base model. Removing the prior regularization from our method (see DB-DC$\backslash$pr and CD-DC$\backslash$pr) leads to only slightly lower results, showing the impact of our proposed regularization method. Also, worst class drop significantly reduces when applying DC, but for some datasets remains still high.

Furthermore, to measure generation quality, diversity, and alignment between image and text of the learned concepts, we employ metrics adapted from personalization methods: DINO, CLIP-I, and CLIP-T. In Table 2 we can see that the proposed regularization method DC does not negatively impact the image generation quality of the learned concepts.¹¹1The results with standard deviations for Table 1 and 2 are provided in Appendix B.

While open-world forgetting does not always result in significant changes to the core content of the image, as shown in Figure 2a, it notably affects intra-class variation, color distribution, and texture characteristics. We define these collective changes as appearance drift, a phenomenon that alters the model’s representation space in subtle yet impactful ways. Figure 3 demonstrates two key aspects of appearance drift; intra-class and contextual variation (first row), where different customizations of the base model lead to changes in car brand and background, while maintaining the overall concept of ‘car’. Color shift (second row), where the color palette of the generated images changes significantly, even when the intra-class characteristics and background remain relatively constant.

Appearance drift manifests through alterations in visual attributes at varying degrees of intensity. Finetuning can cause the model to reinterpret these visual features, leading to inconsistencies between original and newly generated outputs. Although initially subtle, appearance drift can substantially impact customized models. For example, when attempting to learn and generate a set of new concepts within the same context (e.g., for synthetic dataset creation or advertising purposes), each customization of the base model may result in color and content changes. This variability makes it challenging to achieve consistent results across multiple iterations. Moreover, as the customization process alters the model’s manifold, the resulting model becomes less reliable in domains outside the scope of the customization training images. This limitation highlights the importance of understanding and mitigating appearance drift in applications of fine-tuned text-to-image models.

In this paper, we have focused on drift throughout the whole diffusion model manifold. Previous works, especially those in the machine unlearning community (Gandikota et al., 2023), have concentrated on local drift. When removing a concept from a model, it is believed to mainly impact the representation of closely related concepts (hence the name local drift). Our findings suggest that the effects of finetuning are more pervasive than previously thought, potentially influencing the model’s understanding and representation of far-away categories as well as close-by (local) categories.

Here, we repeated our experiments from Section 4.1 and 4.2 to measure the local semantic and local appearance drift. For this setup, we generated 1,000 samples of the closest concepts (superclasses) to each trained model (see Appendix B.1 for the details) and evaluated the CLIP-I, CDI, FID, and KID metrics. As Figure 5a shows, the semantic drift is showing a significant shift towards the left, indicating that local drift is more pronounced. For appearance drift, Figure 5c depicts a more uniform color and KID shift over all the models; this shows that related concepts are affected with a similar magnitude by appearance drift. Again the application of our proposed Drift Correction method greatly reduces both the local semantic drift (as measured in Figure 5b and it almost removes the local appearance drift as measured by KID, even though some color drift remains (see Figure 5c).