SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Wenhu Chen\* Hexiang Hu\* Yandong Li\* Nataniel Ruiz Xuhui Jia Ming-Wei Chang William W. Cohen
$\spadesuit$  Google Deepmind  $\text{♥}$  Google Research {wenhuchen,hexiang,mingweichang,wcohen}@google.com
Figure 1: We train a single SuTI model to generate novel scenes faithfully reflecting given subjects (unseen in training, specified only by 3-5 in-context text→image demonstrations), without any optimization.
Recent text-to-image generation models like DreamBooth have made remarkable progress in generating highly customized images of a target subject, by fine-tuning an "expert model" for a given subject from a few examples. However, this process is expensive, since a new expert model must be learned for each subject. In this paper, we present SuTI, a Subject-driven Text-to-Image generator that replaces subject-specific fine tuning with in-context learning. Given a few demonstrations of a new subject, SuTI can instantly generate novel renditions of the subject in different scenes, without any subject-specific optimization. SuTI is powered by apprenticeship learning, where a single apprentice model is learned from data generated by a massive number of subject-specific expert models. Specifically, we mine millions of image clusters from the Internet, each centered around a specific visual subject. We adopt these clusters to train a massive number of expert models, each specializing in a different subject. The apprentice SuTI model then learns to imitate the behavior of these fine-tuned experts. SuTI can generate high-quality and customized subject-specific images 20x faster than optimization-based SoTA methods. On the challenging DreamBench and DreamBench-v2, human evaluation shows that SuTI significantly outperforms other existing models.
*Core Contribution
Recent text-to-image generation models [1] have shown great progress in generating highly realistic, accurate, and diverse images from a given text prompt. These models are pre-trained on web-crawled image-text pairs like LAION [2] with autoregressive backend models [3, 4] or diffusion backend models [5, 1]. Though achieving unprecedented success in generating highly accurate images, these models are not able to customize to a given subject, like a specific dog, shoe, backpack, etc. Therefore, subject-driven text-to-image generation, the task of generating highly customized images with respect to a target subject, has attracted significant attention from the community. Subject-driven image generation is related to text-driven image editing but often needs to perform more sophisticated transformations to source images (e.g., rotating the view, zooming in/out, changing the pose of subject, etc.) so existing image editing methods are generally not suitable for this new task.
Current subject-driven text-to-image generation approaches are slow and expensive. While different approaches like DreamBooth [6], Imagc [7], and Textual Inversion [8] have been proposed, they all require fine-tuning specific models for a given subject on one or a few demonstrated examples, which typically takes at least 10-20 minutes² to specialize the text-to-image model checkpoint for the given subjects. These approaches are time-consuming as they require back-propagating gradients over the entire model for hundreds or even thousands of steps per customization. Moreover, they are space-consuming as they require storing a subject-specific checkpoint per subject. To avoid the excessive cost, Re-Imagen [9] proposed a retrieval-augmented text-to-image framework to train a subject-driven generation model in a weakly-supervised fashion. Since the retrieved neighbor images are not guaranteed to contain the same subjects, the model does not perform as good as DreamBooth [6] for the task of subject-driven image generation.
To avoid excessive computation and memory costs, we propose to train a single subject-driven text-to-image generation model that can perform on-the-fly subject customization. Our method is dubbed Subject-driven Text-to-Image generator (SuTI), which is trained with a novel apprenticeship learning algorithm. Unlike standard apprenticeship learning which only focuses on learning from one expert, our apprentice model imitates the behaviors of a massive number of specialized expert models. After such training, SuTI can instantly adapt to unseen subjects and unseen or even compositional descriptions with only 3-5 in-context demonstrations within 30 seconds (on a Cloud TPU v4).
Figure 2: Conceptual Diagram of the Learning Pipeline
Figure 2 presents a conceptual diagram of the learning and data preparation pipeline. We first group the images in WebLI [10] by their source URL to form tiny image clusters because images from the same URL are likely to contain the same subject. We then performed extensive image-to-image and image-to-text similarity filtering to retain image clusters that contain highly similar content. For each subject image cluster, we fine-tuned an expert model to specialize in the given subject. Then, we use the fine-tuned experts to synthesize new images given unseen creative captions proposed by large language models. However, the tuned expert models are not perfect and prone to errors, therefore, we adopt a quality validation metric to filter out a large portion of degraded outputs. The remaining high-quality images are provided as a training signal to teach the apprentice model SuTI to perform subject-driven image generation with high fidelity. During inference, the trained SuTI can attend to a few in-context demonstrations to synthesize new images on the fly.
We evaluate SuTI on various tasks such as subject re-contextualization, attribute editing, artistic style transfer, and accessorization. We compare SuTI with existing models on DreamBench [6], which contains diverse subjects from wide categories accompanied by some prompt templates. We compute the CLIP-I/CLIP-T and DINO scores of SuTI's generated images on this dataset and compare them with DreamBooth. The results indicate that SuTI can outperform DreamBooth while having 20x faster inference speed and significantly less memory footprint.
<sup>2</sup>Running on A100 according to public colab: https://huggingface.co/sd-dreambooth-library and https://huggingface.co/docs/diffusers/training/text_inversion.
Further, we manually created 220 diverse and compositional prompts regarding the subjects in DreamBench for human evaluation, which is dubbed the DreamBench-v2 dataset. We then comprehensively compare with other baselines like InstructPix2Pix [11], Null-Text Inversion [12], Imagic [7], Textual Inversion [8], Re-Imagen [9], and DreamBooth [6] on DreamBench-v2. Our human evaluation results indicate that SuTI is  $5\%$  higher than DreamBooth and at least  $30\%$  better than the other baseline in terms of human evaluation. We conduct detailed fine-grained analysis and found that SuTI's textual alignment is significantly better than DreamBoothm while its subject alignment is slightly better than DreamBooth. However, DreamBooth's outputs are still better in the photorealism aspect, especially in terms of fine-grained detail preservation.
We summarize our contributions in the following aspects:
- We introduce the SuTI model, a subject-driven text-to-image generator that performs instant and customized generation for a visual subject with few (image, text) exemplars, all in context.
- We employ the apprenticeship learning to train one single apprentice SuTI model to imitate half a million fine-tuned subject-specific experts on a large-scale seed dataset, leading to a generator model that generalizes to unseen subjects and unseen compositional descriptions.
- We perform a comprehensive set of automatic and human evaluations to show the capability of our model on generating highly faithful and creative images on DreamBench and DreamBench-v2.
To facilitate the reproducibility of our model performance, we release the SuTI model API as a Google Cloud Vertex AI model service, under the production name 'Instant tuning'.
In this section, we introduce the key concepts and notations about subject-driven image-text data, then discuss the basics of text-to-image diffusion models.
Diffusion Models. Diffusion models [13] are latent variable models, parameterized by  $\Theta$ , in the form of  $p_{\Theta}(\boldsymbol{x}_0) \coloneqq \int p_{\Theta}(\boldsymbol{x}_{0:T}) d\boldsymbol{x}_{1:T}$ , where  $\boldsymbol{x}_1, \dots, \boldsymbol{x}_T$  are "noised" latent versions of the input image  $\boldsymbol{x}_0 \sim q(\boldsymbol{x}_0)$ . Note that the dimensionality of both latents and the image is the same throughout the entire process, with  $\boldsymbol{x}_{0:T} \in \mathbb{R}^d$  and  $d$  equals the product of <height, width, # of channels>. The process that computes the posterior distribution  $q(\boldsymbol{x}_{1:T} | \boldsymbol{x}_0)$  is also called the forward (or diffusion) process, and is implemented as a predefined Markov chain that gradually adds Gaussian noise to the data according to a schedule  $\beta_t$ :
$$
q \left(\boldsymbol {x} _ {1: T} \mid \boldsymbol {x} _ {0}\right) = \prod_ {t = 1} ^ {T} q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}\right); \quad q \left(\boldsymbol {x} _ {t} \mid \boldsymbol {x} _ {t - 1}\right) := \mathcal {N} \left(\boldsymbol {x} _ {t}; \sqrt {1 - \beta_ {t}} \boldsymbol {x} _ {t - 1}, \beta_ {t} \boldsymbol {I}\right) \tag {1}
$$
Diffusion models are trained to learn the image distribution by reversing the diffusion Markov chain. Theoretically, this reduces to learning to denoise  $\pmb{x}_t \sim q(\pmb{x}_t | \pmb{x}_0)$  into  $\pmb{x}_0$ , with a time re-weighted square error loss—see [14] for the complete proof:
$$
\mathbb {E} _ {\left(\boldsymbol {x} _ {0}, \boldsymbol {c}\right) \sim D} \left\{\mathbb {E} _ {\epsilon , t} \left[ w _ {t} \cdot \left| \left| \hat {\boldsymbol {x}} _ {\theta} \left(\boldsymbol {x} _ {t}, \boldsymbol {c}\right) - \boldsymbol {x} _ {0} \right| \right| _ {2} ^ {2} \right] \right\} \tag {2}
$$
where  $D$  is the training dataset containing (image, condition)  $= (x_0, c)$  pairs, the condition normally refers to the input text prompt. In practice,  $w_t$  can be simplified as 1 according to [14, 15].
Subject-Driven Text-to-Image Generation. Existing subject-driven generation models [6, 16, 8] often fine-tune a pre-trained text-to-image diffusion model on a set of provided demonstrations  $\mathbb{C}_s$  about a specific subject  $s$ . Formally, such demonstration contains a set of text and image pairs  $\mathbb{C}_s = \{(x_k, c_k)\}_{k}^{\mathrm{K}_s}$ , centered around the subject  $s$ . Images  $\boldsymbol{x}_k$  contains images of the same subject  $s$ , while  $\boldsymbol{c}_s$  is a short description of images  $\boldsymbol{x}_k$ . DreamBooth [6] also requires an additional  $\bar{\mathbb{C}}_s$ , which contains images about different subjects of the same category as  $s$  for prior preservation. To obtain a customized diffusion model  $\hat{x}_{\theta_s}(\boldsymbol{x}_t, \boldsymbol{c})$ , we need to optimize the following loss function:
$$
\theta_ {s} = \underset {\theta} {\arg \min } \mathbb {E} _ {\left(\boldsymbol {x} _ {0}, \boldsymbol {c}\right) \sim \mathbb {C} _ {s} \cup \bar {\mathbb {C}} _ {s}} \left\{\mathbb {E} _ {\boldsymbol {\epsilon}, t} [ \| \hat {\boldsymbol {x}} _ {\theta} \left(\boldsymbol {x} _ {t}, \boldsymbol {c}\right) - \boldsymbol {x} _ {0} \| _ {2} ^ {2} ] \right\} \tag {3}
$$
The customized diffusion model  $\hat{\pmb{x}}_{\theta_s}(\pmb{x}_t,\pmb{c})$  has shown impressive capabilities to generate highly faithful images of the specified subject  $s$ .
3Generally available at https://cloud.google.com/vertex-ai/docs/generative-ai/image/fine-tune-model
Figure 3: Overview of the apprentice learning pipeline for SuTI. Left part shows the customization procedure for expert models, and the right parts show the SuTI model that imitates the behaviors of differently customized experts. Note that this framework can cope with expert models of arbitrary architecture and model family.
3 Apprenticeship Learning from Subject-specific Experts
Notation. Figure 3 presents the concrete workflow of learning. Our method follows apprenticeship learning [17] with two major components, i.e., the expert diffusion models  $\hat{\pmb{x}}_{\theta_s}(\pmb{x}_t,\pmb{c})$  parameterized by  $\theta_s$  regarding subject  $s\in \mathbb{S}$  and apprentice diffusion model  $\hat{\pmb{x}}_{\Theta}(\pmb{x}_t,\pmb{c},\mathbb{C}_s)$  parameterized by  $\Theta$ . The apprentice model takes an additional set of image-text demonstrations  $\mathbb{C}_s$  as input. We use  $\mathbb{S}$  to denote the superset of subjects we include in the training set.
Dataset. The training set  $\mathcal{D}_{\mathbb{S}}$  contains a collection of  $\{\mathbb{C}_s, \pmb{p}_s\}_{s \in \mathbb{S}}$ , where each entry contains an image-text cluster  $\mathbb{C}_s$  accompanied by an unseen prompt  $\pmb{p}_s$ . The image-text cluster  $\mathbb{C}_s$  contains a set of 3-10 image-text pairs. The unseen prompt is an imaginary caption proposed by PaLM [18]. For example, if  $c$  is 'a photo of berry bowl', then  $\pmb{p}_s$  would be an imaginary caption like 'a photo of berry bowl floating on the river'. We describe the dataset construction process to section 4.
Learning. To obtain an expert  $\hat{\pmb{x}}_{\theta_s}(\pmb{x}_t,\pmb{c})$  on a subject  $s$ , we fine-tune a pre-trained diffusion model [1] on the image cluster  $\mathbb{C}_s$  with the denoising loss as:
$$
\theta_ {s} = \underset {\theta} {\arg \min } \mathbb {E} _ {\left(\boldsymbol {x} _ {s}, \boldsymbol {c}\right) \sim \mathbb {C} _ {s}} \left\{\mathbb {E} _ {\boldsymbol {\epsilon}, t} [ \| \hat {\boldsymbol {x}} _ {\theta} \left(\boldsymbol {x} _ {t}, \boldsymbol {c}\right) - \boldsymbol {x} _ {s} \| _ {2} ^ {2} ] \right\} \tag {4}
$$
where  $\pmb{x}_t \sim q(\pmb{x}_t | \pmb{x}_s)$ . The training is similar to Eqn. 3 except that we do not have negative examples for prior preservation because finding the negative examples from the same class is expensive.
Once an expert model is trained, we use it to sample images  $\mathbf{y}_s$  for the unseen text description  $\mathbf{p}_s$  to guide the apprentice SuTI model. We gather the outputs from the massive amount of expert models and then use CLIP filtering to construct a dataset  $G$ . Similarly, we fine-tune the apprentice model  $\hat{\mathbf{x}}_{\Theta}(\mathbf{x}_t,\mathbf{p}_s,\mathbb{C}_s)$  with the denoising loss on the pseudo target generated by the expert:
$$
\Theta = \underset {\Theta} {\arg \min } \mathbb {E} _ {\left(\boldsymbol {y} _ {s}, \boldsymbol {p} _ {s}, \mathbb {C} _ {s}\right) \sim G} \left\{\mathbb {E} _ {\epsilon , t} [ \| \hat {\boldsymbol {x}} _ {\Theta} \left(\boldsymbol {x} _ {t}, \boldsymbol {p} _ {s}, \mathbb {C} _ {s}\right) - \boldsymbol {y} _ {s} \| _ {2} ^ {2} ] \right\} \tag {5}
$$
where  $\pmb{x}_t \sim q(\pmb{x}_t | \pmb{y}_s)$ , and the training triples  $(\pmb{y}_s, \pmb{p}_s, \mathbb{C}_s)$  are drawn from  $G$ .
Algorithm. We formally introduce our learning algorithm in the Algorithm 1. To improve the training efficiency, we use distributed training algorithm to accelerate the training process. At each training step, we randomly sample a batch  $\{B_{s_i}\}_{i = 1}^{\mathrm{K}}$  of size  $\mathsf{K}$  from the dataset  $\mathcal{D}_{\mathbb{S}}$ , with  $B_{s_i} = (\mathbb{C}_{s_i},p_{s_i})$ . We then fine-tune  $\mathsf{K}$  expert models separately w.r.t. Eqn. 4 in parallel, across  $\mathsf{K}$  different TPU cores. For every subject  $s$  inside the batch  $B_{s}$ , we use the corresponding expert model  $\theta_{s}$  to synthesize the image  $y_{s}$  given the unseen prompt  $p_s$ . As not all expert models can generate highly faithful images, we introduce a quality assurance step to validate the synthesized images. Particularly, we measure the quality of an expert's generation by the delta CLIP score [19]  $\Delta (y_s,\mathbb{C}_s,p_s)$ , which is used to decide whether a sample should be included in the dataset  $G$ . This ensures the high quality of the text-to-image training signal for SuTI. Specifically, the delta CLIP score is computed as the increment of CLIP score of  $y_{s}$  over the demonstrated images  $x\in \mathbb{C}_s$ :
$$
\Delta \left(\boldsymbol {y} _ {s}, \mathbb {C} _ {s}, \boldsymbol {p} _ {s}\right) = \operatorname {C L I P} \left(\boldsymbol {y} _ {s}, \boldsymbol {p} _ {s}\right) - \max  _ {\boldsymbol {x} \in \mathbb {C} _ {s}} \operatorname {C L I P} \left(\boldsymbol {x}, \boldsymbol {p} _ {s}\right) \tag {6}
$$
We then feed  $G$  as a training batch to update the parameter  $\Theta$  of the apprentice model using Eqn. 5. In all our experiments, we set  $K = 400$ , with each TPU core training an expert model.
Inference. To perform subject-driven text-to-image generation, the trained SuTI takes 3-5 image-text pairs as the demonstration to generate new images based on the given text description. No
Algorithm 1 Apprenticeship Learning from a Large Crowd of Specialized Expert Models
1: Input: Dataset  $\mathcal{D}_{\mathbb{S}} = \{(\mathbb{C}_s,\pmb {p}_s)\}_{s\in \mathbb{S}}$  containing subject image cluster  $\mathbb{C}_s$  and unseen prompt  $\pmb{p}_{s}$    
2: Input: Pre-trained diffusion model parameterized by  $\theta$    
3: Output: Apprentice diffusion model parameterized by  $\Theta$    
4: Initialize a buffer  $G = \varnothing$    
5: while  $\mathcal{D}_{\mathbb{S}}\neq \emptyset$  do   
6:  $\{B_{s_i}\}_{i = 1}^K = \mathrm{Dequeue}(\mathcal{D}_{\mathbb{S}},\mathbb{K})$  , where  $B_{si} = (\mathbb{C}_{si},\pmb{p}_{si})$    
7: Fine-tune K expert models  $\theta_{s_1},\dots,\theta_{s_K}$  on  $\{B_{s_i}\}_{i = 1}^K$  in parallel, based on Eqn. 4   
8: for  $i = 1$  to K do   
9: Sample a subject-specific generation  $\pmb{y}_{si}$  with DDPM using  $\hat{\pmb{x}}_{\theta_{s_i}}(\pmb {x}_t,\pmb {p}_{s_i})$    
10: if  $\Delta (\pmb {y}_{s_i},\mathbb{C}_{s_i},\pmb {p}_{s_i}) > \lambda$  then   
11:  $G = \mathrm{Enqueue}(G,(y_s,\mathbb{C}_{s_i},p_{s_i}))$    
12: end if   
13: end for   
14: end while   
15: Train  $\hat{\pmb{x}}_{\Theta}$  on the generated demonstration  $G$  , based on the Eqn. 5
optimization is needed during inference time. The only overhead of SuTI is the cost of encoding these 3-5 image-text pairs and the attention computation, which is more affordable. Our inference speed is roughly in the same order as the original text-to-image generator [1].
4 Mining and Generating Subject-driven Text-to-Image Demonstrations
In this section, we discuss how we created the seed dataset  $\mathcal{D}_{\mathbb{S}}$  by mining images and text over the Internet. We construct the seed dataset using the WebLI [10, 20] dataset. Particularly, we derive our initial image cluster via subsampling the Episodic WebLI data [20], which grouped Web images from the same URL. Then we filter the clusters to ensure high intra-cluster visual similarity, using image matching models. The filtered set of image-text clusters is then denoted  $\{\mathbb{C}_s\}_{s\in \mathbb{S}}$ .
After obtaining the subject-driven image clusters, we further prompt a large language model [18] to generate a description about the subject, with the goal of creating descriptions of plausible imaginary visual scenes. The generating instances of the descriptions will require skills like subject re-contextualization, attribute editing, artistic style transfer, and accessionation. We denote the generated unseen captions as  $p_s$ . Together with  $\mathbb{C}_s$ , this forms the final dataset  $D_{\mathbb{S}}$ . More details regarding the grouping and filtering are provided in the Appendix.
The dataset  $\mathcal{D}_{\mathbb{S}}$  contains a total of 2M  $(\mathbb{C}_s, p_s)$  pairs. Using the aforementioned delta CLIP score filtering (using a high threshold  $\lambda = 0.02$ ), we remove low-quality synthesized images  $\pmb{y}_s$  from the expert model, finally obtaining a dataset  $G$  with  $\sim 500\mathrm{K}$ $(\mathbb{C}_s, p_s)$  effective training pairs for the following apprenticeship learning.
In this paper, we only train SuTI on the text  $\rightarrow 64\mathrm{x}64$  diffusion model and retain the original  $256\mathrm{x}256$  and  $1024\mathrm{x}1024$  super-resolution as it is from Imagen [1].
Expert Models. The expert model is initialized from the original 2.1B Imagen 64x64 model. We tune each model on a single TPU core (32 GB) for 500 steps using Adafactor optimizer with a learning rate of 1e-5, which only takes 5 minutes to finish. We use classifier-free guidance to sample new images, where the guidance weight is set to 30. To avoid excessive memory costs, we use fine-tuned experts to sample pseudo-target images and then write the samples as separate files. SuTI will read these files asynchronously to maximize the training speed. Our expert models have a few distinctions from the DreamBooth [6]: 1) we adopt Adafactor instead of Adam optimizer, 2) we do not include any class word token like 'DOG' dog' in the prompt. 3) we do not include in-class negatives for prior preservation. Though our expert model is weaker than DreamBooth, such design choices significantly reduce time/space costs to enable us to train millions of experts with reasonable resources.
Figure 4: Comparison with other Image Editing and Image Personalization Models.
**Apprentice Model.** The apprentice model contains 2.5B parameters, which is 400M parameters larger than the original 2.1B Imagen 64x64 model. The added parameters are coming from the extra attention layers over the demonstrated image-text inputs. We adopt the same architecture as Re-Imagen [9], where the additional image-text pairs are encoded by re-using the UNet DownStack, and the attention layers are added to the down UNet DownStack and UpStack at different resolutions.
We initialize our model fromImagen's checkpoint. For the additional attention layers, we use random initialization. The apprentice training is performed on 128 Cloud TPU v4 chips. We train the model for a total of 150K steps. We use an Adafactor optimizer with a learning rate of 1e-4. We use 3 demonstrations during training, while the model can generalize to leverage any number of demonstrations during inference. We show our ablation studies in the following section.
Inference. We normally provide 4 demonstration image-text pairs to SuTI during inference. Increasing the number of demonstrations does not improve the generation quality much. We use a lower classifier-free guidance weight of 15 with DDPM [14] sampling strategy.
5.1 Datasets and Metrics
DreamBench. In this paper, we use the DreamBench dataset proposed by DreamBooth [6]. The dataset contains 30 subjects like backpacks, stuffed animals, dogs, cats, clocks, etc. These images are downloaded from Unsplash. The original dataset contains 25 prompt templates covering different skills like recontextualization, property modification, accessorization, etc. In total, there are a total of 750 unique prompts generated by the template. We follow the original paper to generate 4 images for each prompt to form the 3000 images for robust evaluation. We follow DreamBooth to adopt DINO, CLIP-I to evaluate the subject fidelity, and CLIP-T to evaluate the text fidelity.
DreamBench-v2. To further increase the difficulty and diversity of DreamBench, we annotate 220 prompts for the 30 subjects in DreamBench as DreamBench-v2. We gradually increase the compositional levels of the prompt to increase the difficulty, like 'back view of [dog]  $\rightarrow$  'back view of [dog] watching TV'  $\rightarrow$  'back view of [dog] watching TV about birds'. This enables us to perform a breakdown analysis to understand the model's compositional capabilities.
We use human evaluation to measure the generation quality in DreamBench-v2. Specifically, we aim at measuring the following three aspects: (1) the subject fidelity score  $s_s$  measures whether the subject is being preserved, (2) the textual fidelity score  $s_t$  measures whether it is aligned with the text description, (3) the photorealism score  $s_p$  measures whether the image contains artifacts or blurry subjects. These are all binary scores, which are averaged over the entire dataset. We combine them as an overall score  $s_o = s_s \wedge s_t \wedge s_p$ , which is the most stringent score.