SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Anton Voronov*
MIPT, Yandex
HSE University, Yandex
Mikhail Khoroshikh*
Max Ryabinin*
Text-to-image generation models represent the next step of evolution in image synthesis, offering a natural way to achieve flexible yet fine-grained control over the result. One emerging area of research is the fast adaptation of large text-to-image models to smaller datasets or new visual concepts. However, many efficient methods of adaptation have a long training time, which limits their practical applications, slows down experiments, and spends excessive GPU resources. In this work, we study the training dynamics of popular text-to-image personalization methods (such as Textual Inversion or DreamBooth), aiming to speed them up. We observe that most concepts are learned at early stages and do not improve in quality later, but standard training convergence metrics fail to indicate that. Instead, we propose a simple drop-in early stopping criterion that only requires computing the regular training objective on a fixed set of inputs for all training iterations. Our experiments on Stable Diffusion for 48 different concepts and three personalization methods demonstrate the competitive performance of our approach, which makes adaptation up to 8 times faster with no significant drops in quality.
Large text-to-image synthesis models have recently attracted the attention of the research community due to their ability to generate high-quality and diverse images that correspond to the user's prompt in natural language [1, 2, 3, 4, 5]. The success of these models has driven the development of new problem settings that leverage their ability to draw objects in novel environments. One particularly interesting task is personalization (or adaptation) of text-to-image models to a small dataset of images provided by the user. The goal of this task is to learn the precise details of a specific object or visual style captured in these images: after personalization, the model should be able to generate novel renditions of this object in different contexts or imitate the style that was provided as an input.
Several recent methods, such as Textual Inversion [6], DreamBooth [7], and Custom Diffusion [8], offer easy and parameter-efficient personalization of text-to-image models. Still, a major obstacle on the path to broader use of such methods is their computational inefficiency. Ideally, models should adapt to user's images in real or close to real time. However, as noted by the authors of the first two papers, the training time of these methods can be prohibitively long, taking up to two hours for a single concept. The reported training time for Custom Diffusion is 12 minutes per concept on a single GPU, which is much faster but still outside the limits of many practical applications.
*Equal contribution. Correspondence to mryabinin0@gmail.com
Both images are taken from the same training run.
Using the deterministic loss  $\mathcal{L}_{\mathrm{det}}$  as a convergence indicator, we can complete Textual Inversion much earlier without drops in quality.
Figure 1: A summary of our key findings: the quality of methods like Textual Inversion saturates early on, but the training loss does not indicate that. Evaluating the same loss on a batch of data fixed throughout the run makes the training dynamics significantly more informative.
In this work, we seek answers to the following questions: are text-to-image customization techniques inherently time-consuming and can we decrease their runtime without major quality tradeoffs? Focusing on the training dynamics, we observe that the CLIP [9] image similarity score (often used to assess image quality in such tasks [10]) grows sharply only in the early stages of Textual Inversion and hardly improves after that. However, as shown in Figure 1, neither the training loss nor the gradient norm indicate the convergence of the concept embedding, which prevents us from stopping the adaptation process earlier. While it is possible to score samples from the model with CLIP during training, generating them can take quite a long time.
With that in mind, we study the training objective itself and attempt to understand the reasons of its non-informative dynamics. As we demonstrate, the primary cause lies in several sources of stochasticity (such as diffusion time steps or the diffusion noise) introducing noise to the loss function. If we sample random variables only once and fix them across training steps, the loss for these inputs reflects convergence better even when training with the original (fully stochastic) objective.
Motivated by this finding, we propose Deterministic VARiance Evaluation (DVAR), an early stopping criterion for text-to-image customization methods that generates a single fixed batch of inputs at the beginning of training and evaluates the model on this batch after each optimization step. This criterion is straightforward to implement, has two interpretable hyperparameters, and its stopping time corresponds to convergence in terms of the CLIP image score. We validate DVAR by comparing it with a range of baselines on three popular adaptation methods, showing that it is possible to run methods like Textual Inversion and DreamBooth much faster with up to 8x speedups and little to no decline in quality. For Custom Diffusion, our method recovers the optimal number of training iterations, which makes it useful to avoid empirical tuning of the step count for every specific dataset.
- We investigate the difficulties in detecting the convergence of text-to-image adaptation methods from training-time metrics. As we demonstrate, the objective becomes much more interpretable across iterations if computed on the same batch of inputs without resampling any random variables.
- We propose DVAR, a simple early stopping criterion for text-to-image personalization techniques. This criterion is easy to compute and use, does not affect the training process and correlates with convergence in terms of visual fidelity.
- We compare this criterion with multiple baselines on Stable Diffusion v1.5, a popular text-to-image model, across three methods: Textual Inversion, DreamBooth, and Custom Diffusion. DVAR offers a significant decrease in runtime while having quality comparable to both original methods and other early stopping baselines.
<sup>1</sup>The code of our experiments is available at github.com/yandex-research/DVAR.
2.1 Denoising Diffusion Probabilistic Models
Diffusion models [11] are a class of generative models that has become popular in recent years due to their ability to generate both diverse and high-fidelity samples [12]. They approximate the data distribution through iterative denoising of a variable  $z_{t}$  sampled from Gaussian distribution. In simple words, the model  $\epsilon_{\theta}$  is trained to predict the noise  $\epsilon$ , following the objective below:
$$
\min  _ {\theta} \mathbb {E} _ {z _ {0}, \epsilon \sim N (0, I), t \sim U [ 1, T ]} | | \epsilon - \epsilon_ {\theta} (z _ {t}, c, t) | | _ {2} ^ {2}. \tag {1}
$$
Here,  $z_{t}$  corresponds to a Markov chain forward process  $z_{t}(z_{0},\epsilon) = \sqrt{\alpha_{t}} z_{0} + \sqrt{1 - \alpha_{t}}\epsilon$  that starts from a sample  $z_{0}$  of the data distribution. For example, in the image domain,  $z_{0}$  corresponds to the target image, and  $z_{T}$  is its version with the highest degree of corruption. In general,  $c$  can represent the condition embedding from any other modality. The inference (reverse) process occurs with a fixed time step  $t$  and starts from  $z_{t}$  equal to a sample from the Gaussian distribution.
2.2 Text-to-image generation
The most natural and well-studied source of guidance for generative models is the natural language [13, 14, 15] because of its convenience for the user, ease of collecting training data, and significant improvements in text representations over the past several years [5]. The condition vector  $c$  is often obtained from Transformer [16] models like BERT [17] applied to the input text.
State-of-the-art text-to-image models perform forward and reverse processes in the latent space of an autoencoder model  $z_0 = \mathcal{E}(x), x = \mathcal{D}(z_0)$ , where  $x$  is the original image, and  $\mathcal{E}, \mathcal{D}$  are the encoder and the decoder, respectively. We experiment with Stable Diffusion v1.5 [3], which uses a variational autoencoder [18] for latent representations. Importantly, this class of autoencoder models is not deterministic, which makes inference even more stochastic but leads to a higher diversity of samples. In total, given image and caption distributions  $X, Y$ , the training loss can be formulated as:
$$
\mathcal {L} _ {\text {t r a i n}} = \mathbb {E} _ {y, x, \epsilon} | | \epsilon - \epsilon_ {\theta} \left(z _ {t} (\mathcal {E} (x), \epsilon), c (y), t\right) | | _ {2} ^ {2}, \tag {2}
$$
$$
y \sim Y, x \sim X, \epsilon \sim \mathcal {N} (0, I), t \sim U [ 0, T ] \tag {3}
$$
2.3 Adaptation of text-to-image models
While text-to-image generation models are flexible because of the natural language input, it is often difficult to design a prompt that corresponds to an exact depiction of an object of choice. Hence, several methods for adapting such models to a given collection of images were developed. In this work, we focus on the most well-known methods, discussing others in Appendix A.
Most approaches inject the target concept into the text-to-image model by learning the representation of a new token  $\hat{v}$  inserted into the language encoder vocabulary. The representation is learned by optimizing the reconstruction loss from Equation 2 for a few (typically  $3 - 5$ ) reference images  $I$  with respect to the embedding of  $\hat{v}$ , while other parts of the model are frozen. The main advantage of these methods is the ability to flexibly operate with the "pseudo-word"  $\hat{v}$  in natural language, for example, by placing the concept corresponding to that word into different visual environments.
Textual Inversion [6] is the simplest method for adaptation of text-to-image models that updates only the token embedding for  $\hat{v}$ . While this method is parameter-efficient, its authors report that reaching acceptable inversion quality requires 6000 training iterations, which equals  $\approx 2$  GPU hours on most machines. This number is a conservative estimate of the number of steps sufficient for all concepts: authors of the method propose selecting the earliest checkpoint that exhibits sufficient train image reconstruction quality $^2$ .
While Textual Inversion only learns the embedding of the target token, DreamBooth [7] does the same with a fully unfrozen U-Net component, and Custom Diffusion [8] updates the projection matrices of cross-attention layers that correspond to keys and values. These methods use a significantly smaller number of iterations, 1000 and 500 respectively, which is also not adaptive and can lead to overfitting.
$^{2}$ github.com/rinongal/textual_inversion/issues/34#issuecomment-1230831673
3 Understanding adaptation dynamics
As we explained previously, the goal of our study is to find ways of speeding up text-to-image without significantly degrading the quality of learned concept representations. To accomplish this goal, we analyze the optimization process of running Textual Inversion on a given dataset. As we demonstrate in Section 4 and in Appendix B, the same findings hold for DreamBooth and Custom Diffusion, which makes our analysis broadly applicable to different methods for customized text-to-image generation.
We apply Textual Inversion to concepts released by Gal et al. [6], using Stable Diffusion v1.5<sup>3</sup> as the base model. We monitor several metrics during training:
1. First, one would hope to observe that optimizing the actual training objective would lead to its convergence, and thus we track the value of  $\mathcal{L}_{train}$ .
2. Second, we monitor the gradient norm, which is often used for analyzing convergence in non-convex optimization. As the model converges to a local optimum, the norm of its gradient should also decrease to zero or stabilize at a point determined by the gradient noise.
3. Lastly, every 50 training iterations, we generate 8 samples from the model using the current concept embedding and score them with the CLIP image similarity score using the training set as references. In the original Textual Inversion paper, this metric is named the reconstruction score and is used for quantitative evaluation.
Note that we do not rely on the CLIP text-image score for captions: in our preliminary experiments, we observed no identifiable dynamics for this metric when using the entire set of CLIP caption templates. Writing more specific captions and choosing the most appropriate ones for each concept takes substantial manual effort; hence, we leave caption alignment out of the scope for this study.
3.1 Initial observations
First, we would like to view the training dynamics in terms of extrinsic evaluation: by measuring how the CLIP image score changes throughout training, we can at least estimate how fast the samples begin to resemble the training set. For this, we perform inversion of all concepts released by [6, 7, 8] (a total of 48 concepts): an example of such an experiment for one concept is available in Figure 2.
From these experiments, we observe that the CLIP image score exhibits sharper growth at an early point of training (often within the first 1000 iterations) and stabilizes later. This finding agrees with the results of our own visual inspection: the generated samples for most concepts undergo the most drastic changes at the beginning and do not improve afterwards. Practically, this observation means that we can interrupt the inversion process much earlier without major drawbacks if we had a criterion for detecting its convergence. What indicators can we use to create such a criterion?
The most straightforward idea is to consider the training loss  $\mathcal{L}_{train}$ . Unfortunately, it is not informative by default: as we demonstrate in Figure 2, the loss exhibits too much noise and has no signs of convergence. The gradient norm of the concept embedding is also hardly informative: in the same experiment, we can see that it actually increases during training instead of decreasing. As we show in Appendix C, these findings generally hold even for much larger training batches, which means that the direct way of making the loss less stochastic is not practical. Still, as reported in the original work and shown by samples and their CLIP scores, the model successfully learns the input concepts. Curiously, we see no reflection of that in the dynamics of the training objective.
Another approach to early stopping is to leverage our observations about the CLIP image score and measure it during training, terminating when the score fails to improve for a specific number of iterations. However, there are two downsides to this approach. First, frequently sampling images during training significantly increases the runtime of the method. Second, this criterion can be viewed as directly maximizing the CLIP score, which is known to produce adversarial examples for CLIP instead of actually improving the image quality [2].
<sup>3</sup>We chose this version because it was given as a default in the Diffusers [19] example for Textual Inversion.
Figure 2: An overview of the convergence process for Textual Inversion with an example concept. Figures 6 and 7 in Appendix B present more examples.
3.2 Investigating the sources of randomness
We hypothesize that the cause of excessive noise in  $\mathcal{L}_{train}$  is several factors of randomness injected at each training step, as we mentioned previously in Section 2.2. Thus, we aim to estimate the influence of the following factors on the dynamics of the inversion objective:
1. Input images  $x$
2. Captions corresponding to images  $y$
3. VAE Latent representations for images  $\mathcal{E}(x)$
4. Diffusion time steps  $t$
5. Gaussian diffusion noise  $\epsilon$
Now, our goal is to identify the factors of stochasticity that affect the training loss. Importantly, we do not change the training batches, as it alters the objective of Textual Inversion and might affect its outcome. Thus, we train the model in the original setting (with batches of entirely random data) but evaluate it on batches with some sources of randomness fixed across all iterations. Note that  $\mathcal{E}(x)$  depends on  $x$ : if we resample the input images, we also need to resample their latent representations.
First, we try the most direct approach of making everything deterministic: in other words, we compute  $\mathcal{L}_{\text{det}}$ , which is the same as  $\mathcal{L}_{\text{train}}$ , but instead of the expectation over random data and noise, we compute it on the same inputs after each training step. Formally, we can define it as
$$
\mathcal {L} _ {\text {d e t}} = \left| \left| \epsilon - \epsilon_ {\theta} \left(z _ {t} (\mathcal {E} (x), \epsilon), c (y), t) \right| \right| _ {2} ^ {2}, \right. \tag {4}
$$
with  $x, y, \mathcal{E}(x), t,$  and  $\epsilon$  sampled only once in the beginning of inversion. Essentially, this means that the only argument of this function that changes across training iterations is  $c(y)$  that depends on the trained concept embedding.
As we show in Figure 2, this version of the objective becomes informative, indicating convergence across a broad range of concepts. Moreover, it displays approximately the same behavior as the CLIP score and is much less expensive to compute, which makes  $\mathcal{L}_{det}$  particularly useful as a metric for the stopping criterion.
For the next step, we aim to find if any of the above sources of stochasticity have negligible impact on the noise in  $\mathcal{L}_{train}$  or can be compensated with larger batches. We evaluate them separately and provide results in Section 4.5. Our key findings are that (1) resampling captions and VAE encoder noise still preserves the convergence trend, (2) using random images or resampling diffusion noise reveals the training dynamics only for large batches, and (3) sampling different diffusion timesteps leads to a non-informative training loss regardless of the batch size. Still, for the sake of simplicity and efficiency, we compute  $\mathcal{L}_{det}$  on a batch of 8 inputs sampled only once at the beginning of training for the rest of our experiments.
def DVAR(losses, window_size, threshold):
    running_var = losses[-window_size].var()
    total_var = losses.var()
    ratio = running_var / total_var
    return ratio < threshold
Figure 3: An example NumPy/PyTorch implementation of DVAR. See Appendix D for an example of its usage in the training code.
3.3 Deterministic Variance Evaluation
The results above show that fixing all random components of the textual inversion loss makes its dynamics more interpretable. To achieve our final goal of decreasing the inversion runtime, we need to design an early stopping criterion that leverages  $\mathcal{L}_{det}$  to indicate convergence.
We propose Deterministic Variance Evaluation (DVAR), a simple variance-based early stopping criterion. It maintains a rolling variance estimate of  $\mathcal{L}_{det}$  over the last  $N$  steps, and once this rolling variance becomes less than  $\alpha$  of the global variance estimate  $(0 < \alpha < 1)$ , we stop training. A pseudocode implementation of DVAR is shown in Figure 3.
This criterion is easy to implement and has two hyperparameters that are easy to tune: the window size for local variance estimation  $N$  and the threshold  $\alpha$ . In our experiments, we found  $\{N = 310, \alpha = 0.15\}$  for Textual Inversion,  $\{N = 440, \alpha = 0.4\}$  for DreamBooth, and  $\{N = 180, \alpha = 0.15\}$  for Custom Diffusion to work relatively well across all concepts we evaluated.
Importantly, we use this criterion while training in the regular fully stochastic setting: our goal is not to modify the objective, and using fixed random variables and data can affect the model's generalization capabilities. As we demonstrate in Section 4, our approach demonstrates significant improvements when compared to baselines, even when all sources of randomness are fixed.
Along with DVAR, we tested other early stopping strategies that use  $\mathcal{L}_{det}$  and are based on different notions of loss value stabilization, such as estimating the trend coefficient or tracking changes in the mean instead of variance. As we show in Appendix E, most of them result in less reliable convergence indicators and have hyperparameters that do not transfer as well between different image collections.
In this section, we compare DVAR with several baselines in terms of sample fidelity for the learned concept, output alignment with novel prompts, and training time. Our goal is to verify that this early stopping criterion is broadly applicable and has little impact on the outcome of training.
4.1 Setup and data
We evaluate approaches to early stopping on three popular text-to-image personalization methods: Textual Inversion, DreamBooth with LoRA [20], and Custom Diffusion, all applied to Stable Diffusion v1.5<sup>4</sup>. For Textual Inversion, we use the implementation from the Diffusers library<sup>5</sup> [19] and take the hyperparameters from the official repository<sup>6</sup>. We use the official code and hyperparameters<sup>7</sup> for Custom Diffusion, and we use the implementation of DreamBooth-LoRA from the Diffusers library<sup>8</sup>.
For evaluation, we combine the datasets published by authors of the three techniques above that were available as of March 2023, which results in a total of 48 concepts. Our main results are obtained by applying each method to all of these concepts, which in total took around 80 GPU hours. Each experiment used a single NVIDIA A100 80GB GPU. Although the methods we study can be used with several concepts at a time, we train on the images of each concept separately. Following Equation 2, the training batches are generated from a small set of original images with randomly sampled augmentations (central/horizontal crop), captions, timesteps and diffusion noise.
huggingface.co/runwayml/stable-diffusion-v1-5  
github.com/huggingface/diffusers/tree/main/examples/textual_inversion  
github.com/rinongal/textual_inversion  
github.com/adobe-research/custom-diffusion  
github.com/huggingface/diffusers/tree/main/examples/dreambooth-lora
To automatically assess the quality of generated images, we employ CLIP image-to-image similarity for images generated from prompts used during training ("Train CLIP img"), which allows us to monitor the degree to which the target concept is learned. However, relying solely on this metric can be risky, as it does not capture overfitting. To measure generalization, we utilize CLIP text-to-image similarity, evaluating the model's ability to generate the concept in new contexts with prompts that are not seen at training time ("Val CLIP txt"). For each method, we report both the number of iterations and the average adaptation runtime, since the average duration of one iteration may differ due to intermediate image sampling or additional loss computation. In addition, we evaluate the identity preservation and diversity of all methods in Appendix F.
We compare DVAR with several baselines: the original training setup of each adaptation method (named as "Baseline"), early stopping based on the CLIP similarity of intermediate samples to the training set ("CLIP-s"), as well as the original setup with the reduced number of iterations and no intermediate sampling ("Few iters").