An Improved Method for Personalizing Diffusion Models

With the advancements in multi-modal models and large language models, the application of generative models to these frameworks has made text-to-image tasks feasible. In recent years, text-to-image generation models have undergone rapid evolution. There are models based on the Autoregressive architecture like DALL-E[10] and Parti[15], as well as those trained using GANs like Lafite[17]. Text-to-image models based on diffusion models have also garnered significant attention. DALL-E2[9] utilizes CLIP[8] to transform textual descriptions into image embeddings, which are then decoded to generate images that align with the text. Imagen[13] introduces a cascaded architecture that initially generates 64x64 resolution images and subsequently employs a two-stages text-conditioned super-resolution diffusion model to upscale them. Stable Diffusion[11] is among the earliest open-source text-to-image diffusion models. Unlike other diffusion models, it doesn’t directly train on pixel images. Instead, it first compresses images into low dimensional representations and then train on the latent space.

While text-to-image models exhibit a high level of semantic consistency between images and text descriptions, they still face limitations due to the inherent constraints of textual descriptions. Particularly, when dealing with objects that possess intricate details, the model may struggle to accurately convey the nuances, resulting in generated images that do not align with expectations.

The concept of personalized image generation within the diffusion model framework was first introduced by DreamBooth[12]. In their paper, they coined this concept as “subject-driven generation”. Subsequent papers refer to this task as “personalization”. In the paper [4], they use highly personalized (HiPer) text embeddings to replacing the uninformative embeddings and optimize only the HiPer embeddings. During inference, replacing the end of new embeddings with the learned HiPer embeddings, enabling the generation of object images based on different descriptions. Another paper employed Apprenticeship Learning[2]. Initially, the diffusion model was fine-tuned to produce expert models. Subsequently, a dataset was curated using imaginary captions proposed by PaLM, along with images generated by the expert model based on these captions. Finally, the apprentice model was trained using this dataset. The [7] suggested training only on parameters of Cross-attention layers in the U-net¹¹1The main part of diffusion model is composed of resnet blocks, self-attention blocks, cross-attention blocks in a U-net architecture. while freezing others. It also introduced the idea of training a single model for multiple concepts.

We present an approach that combines textual inversion and Dreambooth. While combining the two is a straightforward idea, their proper integration leads to the resolution of each method’s issues, as will be shown in our experimental results.

Specifically, the proposed approach comprises of two stages. In the first stage, the optimization of a token embedding similar to textual inversion is performed, and in the second stage, the optimization (i.e. fine-tuning) of the diffusion model similar to Dreambooth, but without prior preservation loss, is performed. The two stages have differences from textual inversion and Dreambooth, as explained below.

In the first stage, we incorporate a new token “$\langle rare\rangle$” which is used as an adjective to modify a noun, representing the target object, e.g.,

• “a photo of $\langle rare\rangle$ backpack”.

Note that this differs from textual inversion, which utilizes a new word $S_{*}$ of a noun representing the target object. We optimize the embedding $v_{\mathrm{rare}}$ of “$\langle rare\rangle$” by

$$\min_{v_{\mathrm{rare}}}\mathbb{E}_{z\sim\mathcal{E}(x),y,{\epsilon}\sim\mathcal{N}(0,I),t}\Big{[}\|{\epsilon}-{\epsilon}_{\Theta}(z_{t},t,\tau_{\theta}(y)\|_{2}^{2}\Big{]},$$

(3)

with $v_{\mathrm{rare}}$ initialized with the embedding of the word “rare”.

Another difference from textual inversion is that we aim at making “$\langle rare\rangle$” only coarsely represent the appearance of the target object. Thus, we conduct updating

$v_{\mathrm{rare}}$ for only about 100 steps. Note that textural inversion aims at making the new token represent the target object as perfectly as possible, requiring 3,000-5,000 updating steps.

In the second stage, we fine-tune the parameters $\Theta$ of the diffusion model by

$$\min_{\Theta}\mathbb{E}_{z\sim\mathcal{E}(x),y,{\epsilon}\sim\mathcal{N}(0,I),t}\Big{[}\|{\epsilon}-{\epsilon}_{\Theta}(z_{t},t,\tau_{\theta}(y)\|_{2}^{2}\Big{]}.$$

(4)

This loss uses the input images of the target object and the above prompt containing “$\langle rare\rangle$” whose embedding $v_{\mathrm{rare}}$, optimized in the first stage, is frozen in this stage; the parameter $\theta$ of the text encoder is frozen, too. Note that we do not use the prior preservation loss as above here. The aim is to train the model together with the token embedding $v_{\mathrm{rare}}$ learned in the first stage to be able to represent the detailed appearance of the target object.

There is another difference from Dreambooth, which is that we update the model’s parameters $\Theta$ for only 200-400 steps in the fine-tuning. This is in contrast with Dreambooth, which requires 1,000 steps of parameter updates. Note also that it needs the generation of an extra 200 to 1,000 images of common class objects for the prior preservation loss.

To summarize, our method offers two advantages over existing methods. First, it enhances the quality of generated images, often significantly, as demonstrated in the subsequent section. Second, it considerably reduces training time. Fewer parameter updates also decrease the chances of overfitting and language drift. The removal of the prior preservation loss mitigates potential complications. These two also lead to superior image generation quality.

It should be noted that the approach employed by Imagic [6] has some similarity to our method; it employs a similar two-stage approach of optimizing the embedding of a prompt followed by optimizing the model. However, Imagic does not explicitly extract the features of the target object; a separate model needs to be trained for each different editing prompt. Our method allows users to freely modify the image based on different text inputs, providing more flexibility in editing.

We experimentally evaluate the proposed method using the dataset introduced in [12]. The dataset consists of 30 sets of images, each featuring four to six images containing the same specific object. For the base text-to-image synthesis model, we use a pre-trained latent diffusion model by huggingface, i.e., “stabilityai/stable-diffusion-2-1-base” which is available from: https://huggingface.co/stabilityai/stable-diffusion-2-1-base.

As explained above, the proposed method consists of two stages. The parameters are updated for 100 iterations with a learning rate of $5\times 10^{-4}$ in the first stage, and 800 iterations with a learning rate of $5\times 10^{-6}$ in the second stage. We save a checkpoint every 200 steps in the second stage for evaluation.

To evaluate personalization methods, it’s essential to assess how well the model preserves the details of the target object in the generated images. However, designing a metric that aligns with human judgment to measure the similarity of objects in two images is a challenging task.

While it remains an open problem, we follow previous studies [12] for comparative evaluation of the proposed and existing methods.

Specifically, we employ the CLIP score and the DINO score for evaluation metrics; see [1][8] for their definition. The evaluation process typically unfolds as follows. Initially, we employ the fine-tuned model to produce images of the target object. Subsequently, both the generated and input sample images of the target object are encoded into their respective embeddings, either through the CLIP image encoder or the DINO vision transformer. We then compute the cosine similarity between these embeddings to gauge the resemblance between the generated and sample images. For this evaluation, we utilize two methods for prompting: one sourced from Dreambooth [12] and the other from textual inversion [3].

The first method from Dreambooth uses 25 diverse prompts⁵⁵5The compelete prompts are available here https://github.com/google/dreambooth/blob/main/dataset/prompts_and_classes.txt with extra descriptions, each generating four images, totaling 100 images. This aims to test whether the model can produce diverse images. Among all the prompts, some of them might change the object’s shape and color, e.g. “a red $\langle rare\rangle$ dog” or “a cube shaped $\langle rare\rangle$ dog”, which can lead to a decrease in the CLIP score.

The second method from textural inversion uses the simple prompt that contains only the target object,e.g. “a photo of $\langle rare\rangle$ dog”, to generate 64 images.

Table 1 shows the results. The number of optimal training steps can fluctuate based on target objects and their corresponding sample images. As a result, we opted for the highest-quality images that most closely matched the training images from checkpoints between 200 and 800 steps. For our method, optimal results emerged from either 200 or 400 steps; for Dreambooth, the best outcomes were between 400 and 800 steps. It is seen from Table 1 that across all metrics and testing prompts, our method consistently outperforms. Additionally, using simple prompts results in higher scores compared to diverse prompts.

We next show several examples to qualitatively compare the results of the different methods.