DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation

Our goal is to “implant” a new (key, value) pair into the diffusion model’s “dictionary” such that, given the key for our subject, we are able to generate fully-novel images of this specific subject with meaningful semantic modifications guided by a text prompt. One way is by few-shot fine-tuning of the model. The question becomes: how to supervise this process? Generally, text prompts are written by humans and sourced from large online datasets. The main limitations include the overhead of writing detailed image descriptions for a given image set, and the high variance and subjectivity of human-written captions.

We opt for a simpler approach and label all input images of the subject “a [identifier] [class noun]”, where [identifier] is a unique identifier linked to the subject and [class noun] is a coarse class descriptor of the subject (e.g. cat, dog, watch, etc.). The class descriptor can be obtained using a classifier. We specifically use a class descriptor in the sentence in order to tether the prior of the class to our unique subject. We found that using only a unique identifier, without a class noun, as a key for our subject increased training time and decreased performance. In essence, we want to leverage the diffusion model’s prior of the specific class and entangle it with the embedding of our subject’s unique identifier. In this way, it can leverage the visual prior to generate new poses and articulations of the subject in different contexts.

A naive way of constructing an identifier for our subject is to use an existing word. For example, using the words like “unique” or “special”. One problem is that existing English words tend to have a stronger prior due to occurrence in the training set of text-to-image diffusion models. We generally find increased training time and decreased performance when using such generic words to index our subject, since the model has to both learn to disentangle them from their original meaning and to re-entangle them to reference our subject. This approach can also fail by entangling the meaning of the word with the appearance of our object, for example in the extreme case if the identifier chosen is the adjective “blue” and our subject is grey, colors will be entangled at inference and we will sample a mix of grey and blue subjects (as well as mixes of both). This motivates the need for an identifier that has a weak prior in both the language model and the diffusion model. A hazardous way of doing this is to select random characters in the English language and concatenate them to generate a rare identifier (e.g. “xxy5syt00”). In reality, the tokenizer might tokenize each letter separately, and the prior for the diffusion model is strong for these letters. Specifically, if we sample the model with such an identifier before fine-tuning we will get pictorial depictions of the letters or concepts that are linked to those letters. We often find that these tokens incur the same weaknesses as using common English words to index the subject.

In a nutshell, our approach is to find relatively rare tokens in the vocabulary, and then invert these rare tokens into text space. In order to do this, we first perform a rare-token lookup in the vocabulary and obtain a sequence of rare token identifiers $f\left(\hat{V}\right)$, where $f$ is a tokenizer; a function that maps character sequences to tokens and $\hat{V}$ is the decoded text stemming from the tokens $f\left(\hat{V}\right)$. This sequence can be of variable length $k$ with $k$ being a hyperparameter of our method. We find that relatively short sequences of $k=\{1,...,3\}$ work well. Then, by inverting the vocabulary using the de-tokenizer on $f\left(\hat{V}\right)$ we obtain a sequence of characters that define our unique identifier $\hat{V}$. We observe that using uniform random sampling without replacement of tokens that correspond to 3 or fewer Unicode characters (without spaces) and using tokens in the T5-XXL tokenizer range of $\{5000,...,10000\}$ works well.

Given a small set of images depicting the target subject $X_s\text{coloneqq}\{x_s^i;i\in \{0,...,N\left\}\right\}$, and with the same conditioning vector $c_s$ obtained from the text prompt “a [identifier] [class noun]”, we fine-tune the text-to-image model using the classic denoising loss presented in Equation 1, with the same hyperparameters as the original diffusion model. Two key issues arise with such a naive fine-tuning strategy: Overfitting and Language-drift.

Since our input image set is quite small, fine-tuning the large image generation models can overfit to both the context and the appearance of the subject in the given input images (e.g., subject pose). Figure 12 (top) shows some sample generated images with naive fine-tuning where we clearly see that both the subject dog’s appearance and context are overfitted to those in the input images. There are many techniques that can be used to address these problems, such as regularization or selectively fine-tuning certain parts of the model. There is uncertainty on which parts of the model need to be frozen to both obtain good subject fidelity and semantic modification flexibility. In our experience, the best results that achieve maximum subject fidelity are achieved by fine-tuning all layers of the model. Nevertheless, this includes fine-tuning layers that are conditioned on the text embeddings, which gives rise to the problem of language drift.

The phenomenon of language drift has been an observed problem in the language model literature Lee2019CounteringLD; lu2020countering, where a language model that is pre-trained on a large text corpus and later fine-tuned for a specific task progressively loses syntactic and semantic knowledge of the language as it learns to improve in the target task. To the best of our knowledge, we are the first to find a similar phenomenon affecting diffusion models. Since our text prompt contains both the [identifier] and [class noun], when a diffusion model is fine-tuned on a small set of subject images, we observe that it slowly forgets how to generate subjects of the same class and progressively forgets the class-specific prior and can not generate different instances of the class in question. Figure 13 (middle) shows some sample generated images of “a dog” after fine-tuning the model on specific dog images. The results clearly show that the model loses the capability of generating generic dog images with naive fine-tuning.

We propose an autogenous class-specific prior-preserving loss to counter both the overfitting and language drift issues. In essence, our method is to supervise the model with its own generated samples, in order for it to retain the prior once the few-shot fine-tuning begins. Specifically, we generate data $x_{\text{pr}}=\hat{x}(z_{t_1},c_{\text{\% pr}})$ by using the ancestral sampler on the frozen pre-trained diffusion model with random initial noise $z_{t_1}\sim N(0,I)$ and conditioning vector $c_{\text{pr}}\text{coloneqq}\Gamma \left(f\right(\text{"a [class noun]"}\left)\right)$. The loss becomes:

where $\lambda$ controls for the relative weight of the prior-preservation term. Figure 4 illustrates the model fine-tuning with the class-generated samples and prior-preservation loss. Despite being simple, we find this prior-preservation loss is effective in overcoming the overfitting and language-drift issues. We find that $\sim$ 200 epochs at learning rate ${10}^{-5}$ with $\lambda =1$ is enough to achieve good results. During this process, $\sim 200\times N$ “a [class noun]” samples are generated, with $N$ being the size of the subject dataset, usually ranging from 3-5 images. The training process takes about 15 minutes on one TPUv4.

We find results to be suboptimal if the training recipes and test parameters of Saharia et al. saharia2022photorealistic are used to fine-tune the SR models with the given few shots of a subject instance. Specifically, we find that maintaining the original level of noise augmentation used to train the SR networks leads to the blurring of high-frequency patterns of the subject and of the environment. See Figure 14 (middle row) for sample generations. In order to faithfully reproduce the subject instance, we reduce the level of noise augmentation from ${10}^{-3}$ to ${10}^{-5}$ during fine-tuning of the $256\times 256$ SR model. With this small modification, We are able to recover fine-grained details of the subject instance.

Given a customized model ${\hat{x}}_{\theta }$, we can generate novel images for a specific subject instance by prompting the trained model with a sentence containing the unique identifier [V] and the class noun. Specifically, for recontextualization we generally form sentences in the following form: “a [V] [class noun] [context description]”. For example, if the subject instance is a clock, we can prompt the model using “a [V] clock on top of snow” to generate our selected clock in a snowy scene. Another idea is to recontextualize our subject in famous landmarks, for example “a [V] clock with the Eiffel Tower in the background”. We show examples of this capability in Figure 5. Importantly, as opposed to background modification, we are able to generate the subject in new poses and articulations and with new, previously unseen scene structure. Also, we note the detail in the integration of the subject in the scene, including realistic contact with other objects (e.g. partially planted in the snow, manipulated by human hands, etc.) and realistic shadows and reflections. This shows that our method not only allows for interpolation or recovery of subject details - but for extrapolation, answering questions in the form of: “what would happen if this subject were partially planted in the snow?”, given no initial data of the subject in that context.

Given a prompt “a painting of a [V] [class noun] in the style of [famous painter]” or “a statue of a [V] [class noun] in the style of [famous sculptor]” we are able to generate original artistic renditions of our subject instance. In particular, this task does not equate to style transfer in which the semantics of a source scene are preserved and the style of another image is transferred onto the source scene. Instead, we are able to get meaningful changes in the scene, with subject instance details and identity preservation, depending on the artistic style. This means some of the images we generate have the subject in different poses that were not seen in the subject ground truth data and scenes that were also not in the ground truth data. We show examples of original artistic renditions of a personalized model in Figure 6. We use different famous artists to generate the original renditions. The artistic style of the famous artist is well represented in our images, and there is a certain novelty to the renditions. Note that the dog was never seen in many of the poses shown, for example, there are no real images of this dog in the poses shown in the Michelangelo rendition.

Our method allows for new image generation of the subject with modified expressions that are not seen in the original set of subject images. We show examples in Figure 7. The range of expressiveness is high, ranging from negative to positive valence emotions and different levels of arousal. In all examples, the uniqueness of the subject dog is preserved - specifically, the asymmetric white streak on its face remains in all generated images.

We are able to render the subject under different novel viewpoints. In Figure 8 we show all real images of our subject cat - and we are able to generate new images of the same cat (with consistent complex fur patterns) under new camera viewpoints. We highlight that the model has not seen this specific cat from behind, from below, or from above - yet it is able to extrapolate knowledge from the class prior to generating these novel views given only 4 frontal images of the subject.

An interesting capability stemming from the strong compositional prior of the generation model is the ability to accessorize subjects. In Figure 9 we show examples of accessorization of a Chow Chow dog. We prompt the model with a sentence of the form: “a [V] [class noun] wearing [accessory]”. In this manner, we are able to fit different accessories onto this dog - with aesthetically pleasing results. Note that the identity of the dog is preserved in all frames, and subject-accessory contact and articulation are realistic.

We are able to modify subject instance properties. For example we can include a color adjective in the prompt sentence “a [color adjective] [V] [class noun]”. In that way, we can generate novel instances of our subject with different colors. The generated scene can be very similar to the original scene, or it can be changed given a descriptive prompt. We show color changes of a car in the first row of Figure 10. We select similar viewpoints for effect, but we can generate different viewpoints of the car with different colors in different scenarios. This is a simple example of property modification, but more semantically complex property modifications can be achieved using our method. For example, we show crosses between a specific Chow Chow dog and different animal species in the bottom row of Figure 10. We prompt the model with sentences of the following structure: “a cross of a [V] dog and a [target species]”. In particular, we can see in this example that the identity of the dog is well preserved even when the species changes - the face of the dog has certain individual properties that are well preserved and melded with the target species. Other property modifications are possible, such as material modification (e.g. a dog made out of stone). Some are harder than others and depend on the prior of the base generation model.

We show the results of using no class noun, a wrong class noun, and the correct class noun for text supervision of our subject images. Results are shown in Figure 11. We observe that the class prior of the wrong class (dog) remains entangled and we are not able to generate new images of our subject when the model is trained in this manner. Given a longer fitting time, we are able to disentangle the prior of the class noun “dog”, but at the cost of losing the ability to generate dogs with the model and ultimately with decreased performance. If we are to train without a class noun, the model has difficulty learning the subject instance and does not easily entangle the class prior with the instance. The model takes longer to converge and can generate erroneous samples.

Our proposed loss acts as a powerful regularizer. A naively fine-tuned network can quickly overfit to the small set of subject images. In order to explore this phenomenon, we train two models for 200 epochs, one using naive fine-tuning (i.e. using the loss in Equation 1), and a network using our prior-preservation loss shown in Equation 2. We show results for different context captions in Figure 12. We observe that the regularization effect of our loss allows us to capture a wider range of poses for our subject dog without sacrificing subject fidelity. Importantly, we observe that using naive fine-tuning the dog usually lies on a fabric-type material similar to the training images, whereas this is avoided using our method.

Further, we evaluate how our prior preservation loss described in Section 4.2 conserves variability in the prior and show sample results in Figure 13. We verify that a vanilla model is able to generate a large variety of dogs, while a naively fine-tuned model on the subject dog exhibits language drift and generates our subject dog given the prompt “a dog”. Our proposed loss preserves the variability of the prior and the model is able to generate new instances of our dog given a prompt of the style “a [V] dog” but also varied instances of dogs given a “a dog” prompt.

We show how using lower noise to train the super-resolution models improves fidelity. Specifically, we show in Figure 14 that if the super-resolution models are not fine-tuned, we observe hallucination of high-frequency patterns on the subject which hurts identity preservation. Further, if we use the ground-truth noise augmentation level used for training the Imagen $256\times 256$ model (${10}^{-3}$), we obtain blurred and non-crisp details. If the noise used to train the SR model is reduced to ${10}^{-5}$, then we conserve a large amount of detail without pattern hallucination or blurring.