SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yifan Zhang $^{1*}$  Daquan Zhou $^{2*}$  Bryan Hooi $^{1}$  Kai Wang $^{1}$  Jiashi Feng $^{2}$ $^{1}$ National University of Singapore  $^{2}$ ByteDance
The power of DNNs relies heavily on the quantity and quality of training data. However, collecting and annotating data on a large scale is often expensive and time-consuming. To address this issue, we explore a new task, termed dataset expansion, aimed at expanding a ready-to-use small dataset by automatically creating new labeled samples. To this end, we present a Guided Imagination Framework (GIF) that leverages cutting-edge generative models like DALL-E2 and Stable Diffusion (SD) to "imagine" and create informative new data from the input seed data. Specifically, GIF conducts data imagination by optimizing the latent features of the seed data in the semantically meaningful space of the prior model, resulting in the creation of photo-realistic images with new content. To guide the imagination towards creating informative samples for model training, we introduce two key criteria, i.e., class-maintained information boosting and sample diversity promotion. These criteria are verified to be essential for effective dataset expansion: GIF-SD obtains  $13.5\%$  higher model accuracy on natural image datasets than unguided expansion with SD. With these essential criteria, GIF successfully expands small datasets in various scenarios, boosting model accuracy by  $36.9\%$  on average over six natural image datasets and by  $13.5\%$  on average over three medical datasets. The source code is available at https://github.com/Vanint/DatasetExpansion.
A substantial number of training samples is essential for unleashing the power of deep networks [14]. However, such requirements often impede small-scale data applications from fully leveraging deep learning solutions. Manual collection and labeling of large-scale datasets are often expensive and time-consuming in small-scale scenarios [52]. To address data scarcity while minimizing costs, we explore a novel task, termed Dataset Expansion. As depicted in Figure 1 (left), dataset expansion aims to create an automatic data generation pipeline that can expand a small dataset into a larger and more informative one for model training. This task particularly focuses on enhancing the quantity and quality of the small-scale dataset by creating informative new samples. This differs from conventional data augmentations that primarily focus on increasing data size through transformations, often without creating samples that offer fundamentally new content. The expanded dataset is expected to be versatile, fit for training various network architectures, and promote model generalization.
We empirically find that naive applications of existing methods cannot address the problem well (cf. Table 1 and Figure 4). Firstly, data augmentation [12, 15, 109], involving applying pre-set transformations to images, can be used for dataset expansion. However, these transformations primarily alter the surface visual characteristics of an image, but cannot create images with novel content (cf. Figure 5a). As a result, the new information introduced is limited, and the performance gains tend to saturate quickly as more data is generated. Secondly, we have explored pre-trained generative models (e.g., Stable Diffusion (SD) [59]) to create images for model training. However,
*Corresponding to: yifan.zhang@u.nus.edu, zhoudaquan21@gmail.com
Figure 1: Dataset Expansion aims to create data with new information to enrich small datasets for training DNNs better (left). As shown, the created images by our method are all class-relevant but diverse (e.g., new sitting postures of cats, new patterns of cushions). This enables ResNet-50 trained on our expanded datasets to perform much better than the one trained on the original datasets (right).
these pre-trained generative models are usually category-agnostic to the target dataset, so they cannot ensure that the synthetic samples carry the correct labels and are beneficial to model training.
Different from them, our solution is inspired by human learning with imagination. Upon observing an object, humans can readily imagine its different variants in various shapes, colors, or contexts, relying on their prior understanding of the world [75, 80]. Such an imagination process is highly valuable for dataset expansion as it does not merely tweak an object's appearance, but leverages rich prior knowledge to create object variants infused with new information. In tandem with this, recent generative models like SD and DALL-E2 [56] have demonstrated exceptional abilities in capturing the data distribution of extremely vast datasets [4, 66] and generating photo-realistic images with rich content. This motivates us to explore their use as prior models to establish our computational data imagination pipeline for dataset expansion, i.e., imagining different sample variants given seed data. However, the realization of this idea is non-trivial and is complicated by two key challenges: how to generate samples with correct labels, and how to ensure the created samples boost model training.
To handle these challenges, we conduct a series of exploratory studies (cf. Section 3) and make two key findings. First, CLIP [54], which offers excellent zero-shot classification abilities, can map latent features of category-agnostic generative models to the specific label space of the target dataset. This enables us to generate samples with correct labels. Second, we discover two informativeness criteria that are crucial for generating effective training data: 1) class-maintained information boosting to ensure that the imagined data are class-consistent with the seed data and bring new information; 2) sample diversity promotion to encourage the imagined samples to have diversified content.
In light of the above findings, we propose a Guided Imagination Framework (GIF) for dataset expansion. Specifically, given a seed image, GIF first extracts its latent feature with the prior (generative) model. Unlike data augmentation that imposes variation over raw images, GIF optimizes the "variation" over latent features. Thanks to the guidance carefully designed by our discovered criteria, the latent feature is optimized to provide more information while maintaining its class. This enables GIF to create informative new samples, with class-consistent semantics yet higher content diversity, for dataset expansion. Such expansion is shown to benefit model generalization.
As DALL-E2 [56] and SD are powerful in generating images, and MAE [24] excels at reconstructing images, we explore their use as prior models of GIF for data imagination. We evaluate our methods on small-scale natural and medical image datasets. As shown in Figure 1 (right), compared to ResNet-50 trained on the original dataset, our method improves its performance by a large margin across various visual tasks. Specifically, with the designed guidance, GIF obtains  $36.9\%$  accuracy gains on average over six natural image datasets and  $13.5\%$  gains on average over three medical datasets. Moreover, the expansion efficiency of GIF is much higher than data augmentation, i.e.,  $5\times$  expansion by GIF-SD outperforms even  $20\times$  expansion by Cutout, GridMask and RandAugment on Cars and DTD datasets. In addition, the expanded datasets also benefit out-of-distribution performance of models, and can be directly used to train various architectures (e.g., ResNeXt [85], WideResNet [97], and MobileNet [64]), leading to consistent performance gains. We also empirically show that GIF is more applicable than CLIP to handle real small-data scenarios, particularly with non-natural image domains (e.g., medicine) and hardware constraints (e.g., limited supportable model sizes). Please note that GIF is much faster and more cost-effective than human data collection for dataset expansion.
Learning with synthetic images. Training with synthetic data is a promising direction [26, 34, 110]. For example, DatasetGANs [42, 105] explore GAN models [17, 33] to generate images for segmentation model training. However, they require a sufficiently large dataset for in-domain GAN training, which is not feasible in small-data scenarios. Also, as the generated images are without labels, they need manual annotations on generated images to train a label generator for annotating synthetic images. Similarly, many recent studies [2, 3, 22, 39, 63, 78, 82, 94] also explored generative models to generate new data for model training. However, these methods cannot ensure that the synthesized data bring sufficient new information and accurate labels for the target small datasets. Moreover, training GANs from scratch [3, 39, 63, 78, 94], especially with very limited data, often fails to converge or produce meaningful results. In contrast, our dataset expansion aims to expand a small dataset into a larger labeled one in a fully automatic manner without involving human annotators. As such, our method emerges as a more effective way to expand small datasets.
Data augmentation. Augmentation employs manually specified rules, such as image manipulation [91], erasing [15, 109], mixup [28, 100, 101], and transformation selection [11, 12] to boost model generalization [69]. Despite certain benefits, these methods enrich images by pre-defined transformations, which only locally vary the pixel values of images and cannot generate images with highly diversified content. Moreover, the effectiveness of augmented data is not always guaranteed due to random transformations. In contrast, our approach harnesses generative models trained on large datasets and guides them to generate more informative and diversified images, thus resulting in more effective and efficient dataset expansion. For additional related studies, please refer to Appendix A.
3 Problem and Preliminary Studies
Problem statement. To address data scarcity, we explore a novel dataset expansion task. Without loss of generality, we consider image classification, where a small-scale training dataset  $\mathcal{D}_o = \{x_i,y_i\}_{i = 1}^{n_o}$  is given. Here,  $x_{i}$  denotes an instance with class label  $y_{i}$ , and  $n_o$  denotes the number of samples. The goal of dataset expansion is to generate a set of new synthetic samples  $\mathcal{D}_s = \{x_j',y_j'\}_{j = 1}^{n_s}$  to enlarge the original dataset, such that a DNN model trained on the expanded dataset  $\mathcal{D}_o\cup \mathcal{D}_s$  outperforms the model trained on  $\mathcal{D}_o$  significantly. The key is that the synthetic sample set  $\mathcal{D}_s$  should be highly related to the original dataset  $\mathcal{D}_o$  and bring sufficient new information to boost model training.
3.1 A proposal for computational imagination models
Given an object, humans can easily imagine its different variants, like the object in various colors, shapes, or contexts, based on their accumulated prior knowledge about the world [75, 80]. This imagination process is highly useful for dataset expansion, as it does not simply perturb the object's appearance but applies rich prior knowledge to create variants with new information. In light of this, we seek to build a computational model to simulate this imagination process for dataset expansion.
Deep generative models, known for their capacity to capture the distribution of a dataset, become our tool of choice. By drawing on their prior distribution knowledge, we can generate new samples resembling the characteristics of their training datasets. More importantly, recent generative models, such as Stable Diffusion (SD) and DALL-E2, have demonstrated remarkable abilities in capturing the distribution of extremely large datasets and generating photo-realistic images with diverse content. This inspires us to explore their use as prior models to construct our data imagination pipeline.
Specifically, given a pre-trained generative model  $G$  and a seed example  $(x,y)$  from the target dataset, we formulate the imagination of  $x'$  from  $x$  as  $x' = G(f(x) + \delta)$ . Here,  $f(\cdot)$  is an image encoder of  $G$  to transform the raw image into an embedding for imagination with the generative model.  $\delta$  is a perturbation applied to  $f(x)$  such that  $G$  can generate  $x'$  different from  $x$ . A simple choice of  $\delta$  would be Gaussian random noise, which, however, cannot generate highly informative samples. In the following subsection, we will discuss how to optimize  $\delta$  to provide useful guidance.
It is worth noting that we do not aim to construct a biologically plausible imagination model that strictly follows the dynamics and rules of the human brain. Instead, we draw inspiration from the imaginative activities of the human brain and propose a pipeline to leverage well pre-trained generative models to explore dataset expansion.
(b)
Figure 2: Effectiveness of the informativeness criteria for sample creation. (a) Comparison between random expansion and our selective expansion on CIFAR100-Subset. (b) Visualization of DALLE expansion with and without diversity guidance.
3.2 How to guide imagination for effective expansion?
Our data imagination pipeline leverages generative models to create new samples from seed data. However, it is unclear what types of samples are effective and how to optimize  $\delta$  accordingly in the pipeline to create data that are useful for model training. Our key insight is that the newly created sample  $x^{\prime}$  should introduce new information compared to the seed sample  $x$ , while preserving the same class semantics as the seed sample. To achieve these properties, we explore the following preliminary studies and discover two key criteria: (1) class-maintained information boosting, and (2) sample diversity promotion.
Class-maintained informativeness boosting. When enhancing data informativeness, it is non-trivial to ensure that the generated  $x'$  has the same label  $y$  as the seed sample  $x$ , since it is difficult to maintain the class semantics after perturbation in the latent space  $f(x)$ . To overcome this, we explore CLIP [54] for its well-known image-text alignment ability: CLIP's image encoder can project an image  $x$  to an embedding space aligned with the language embedding of its class name  $y$  [72, 81]. Therefore, we can leverage CLIP's embedding vectors of all class names as a zero-shot classifier to guide the generation of samples that maintain the same class semantics as seed data. Meanwhile, the entropy of the zero-shot prediction can serve as a measure to boost the classification informativeness of the generated data.
To pinpoint whether the criteria of class-maintained information boosting helps to generate more informative samples, we start with exploratory experiments on a subset of CIFAR100 [41]. Here, the subset is built for simulating small-scale datasets by randomly sampling 100 instances per class from CIFAR100. We first synthesize samples based on existing data augmentation methods (i.e., RandAugment and Cutout [15]) and expand CIFAR100-Subset by  $5 \times$ . Meanwhile, we conduct selective augmentation expansion based on our criteria (i.e., selecting the samples with the same zero-shot prediction but higher prediction entropy compared to seed samples) until we reach the required expansion ratio per seed sample. As shown in Figure 2a, selective expansion outperforms random expansion by  $1.3\%$  to  $1.6\%$ , meaning that the selected samples are more informative for model training. Compared to random augmentation, selective expansion filters out the synthetic data with lower prediction entropy and those with higher entropy but inconsistent classes. The remaining data thus preserve the same class semantics but bring more information gain, leading to better expansion effectiveness.
Sample diversity promotion. To prevent the "imagination collapse" issue that generative models yield overly similar or duplicate samples, we delve further into the criterion of sample diversity promotion. To study its effectiveness, we resort to a powerful generative model (i.e., DALL-E2) as the prior model to generate images and expand CIFAR100-Subset by  $5 \times$ , where the guided expansion scheme and the implementation of diversity promotion will be introduced in the following section. As shown in Figure 2b, the generated images with diversity guidance are more diversified: starfish images have more diverse object numbers, and motorbike images have more diverse angles of view and even a new driver. This leads to  $1.4\%$  additional accuracy gains on CIFAR100-Subset (cf. Table 20 in Appendix F.5), demonstrating that the criterion of sample diversity promotion is effective in bringing diversified information to boost model training.
Figure 3: Illustration of the proposed GIF method based on DALL-E2 [56], which expands small datasets by creating informative new samples with guided imagination. Here, we resort to DALL-E2 as the prior model, in which the image/text encoders are CLIP image/text encoders while the decoder is the diffusion model of DALL-E2. Moreover,  $\oplus$  denotes guided residual multiplicative perturbation. More implementation details of GIF-DALLE, GIF-SD and GIF-MAE are provided in Appendix D.
4 GIF: A Guided Imagination Framework for Dataset Expansion
In light of the aforementioned studies, we propose a Guided Imagination Framework (GIF) for dataset expansion. This framework guides the imagination of prior generative models based on the identified criteria. Given a seed image  $x$  from a target dataset, we first extract its latent feature  $f(x)$  via the encoder of the generative model. Different from data augmentation that directly imposes variations over raw RGB images, GIF optimizes the "variations" over latent features. Thanks to the aforementioned criteria as guidance, the optimized latent features result in samples that preserve the correct class semantics while introducing new information beneficial for model training.
Overall pipeline. To detail the framework, we use DALL-E2 [56] as a prior generative model for illustration. As shown in Figure 3, DALL-E2 adopts CLIP image/text encoders  $f_{\mathrm{CLIP - I}}$  and  $f_{\mathrm{CLIP - T}}$  as its image/text encoders and uses a pre-trained diffusion model  $G$  as its image decoder. To create a set of new images  $x'$  from the seed image  $x$ , GIF first repeats its latent feature  $f = f_{\mathrm{CLIP - I}}(x)$  for  $K$  times, with  $K$  being the expansion ratio. For each latent feature  $f$ , we inject perturbation over it with randomly initialized noise  $z \sim \mathcal{U}(0,1)$  and bias  $b \sim \mathcal{N}(0,1)$ . Here, to prevent out-of-control imagination, we conduct residual multiplicative perturbation on the latent feature  $f$  and enforce an  $\varepsilon$ -ball constraint on the perturbation as follows:
$$
f ^ {\prime} = \mathcal {P} _ {f, \epsilon} ((1 + z) f + b), \tag {1}
$$
where  $\mathcal{P}_{f,\epsilon}(\cdot)$  means to project the perturbed feature  $f^{\prime}$  to the  $\varepsilon$ -ball of the original latent feature, i.e.,  $\| f' - f\|_{\infty} \leq \varepsilon$ . Note that each latent feature has independent  $z$  and  $b$ . Following our explored criteria, GIF optimizes  $z$  and  $b$  over the latent feature space as follows:
$$
z ^ {\prime}, b ^ {\prime} \leftarrow \underset {z, b} {\arg \max } \mathcal {S} _ {i n f} + \mathcal {S} _ {d i v}, \tag {2}
$$
where  $S_{inf}$  and  $S_{div}$  correspond to the class-maintained informativeness score and the sample diversity score, respectively, which will be elaborated below. This latent feature optimization is the key step for achieving guided imagination. After updating the noise  $z'$  and bias  $b'$  for each latent feature, GIF obtains a set of new latent features by Eq. (1), which are then used to create new samples through the decoder  $G$ .
Class-maintained informativeness. To boost the informativeness of the generated data without changing class labels, we resort to CLIP's zero-shot classification abilities. Specifically, we first use  $f_{\mathrm{CLIP - T}}$  to encode the class name  $y$  of a sample  $x$  and take the embedding  $w_{y} = f_{\mathrm{CLIP - T}}(y)$  as the zero-shot classifier of class  $y$ . Each latent feature  $f_{\mathrm{CLIP - I}}(x)$  can be classified according to its cosine similarity to  $w_{y}$ , i.e., the affinity score of  $x$  belonging to class  $y$  is  $\hat{s}_y = \cos (f(x),w_y)$ , which forms a prediction probability vector  $s = \sigma ([\hat{s}_1,\dots ,\hat{s}_C])$  for the total  $C$  classes of the target dataset based on softmax  $\sigma (\cdot)$ . The prediction of the perturbed feature  $s^{\prime}$  can be obtained in the same way. We then design  $S_{inf}$  to improve the information entropy of the perturbed feature while maintaining its class semantics as the seed sample:
$$
\mathcal {S} _ {i n f} = s _ {j} ^ {\prime} + \left(s \log (s) - s ^ {\prime} \log \left(s ^ {\prime}\right)\right), \text {s . t .} j = \arg \max  (s).
$$
Specifically,  $s_j'$  denotes the zero-shot classification score of the perturbed feature  $s'$  regarding the predicted label of the original latent feature  $j = \arg \max(s)$ . Here, the zero-shot prediction of the
original data serves as an anchor to regularize the class semantics of the perturbed features in CLIP's embedding space, thus encouraging class consistency between the generated samples and the seed sample. Moreover,  $s\log (s) - s'\log (s')$  means contrastive entropy increment, which encourages the perturbed feature to have higher prediction entropy and helps to improve the classification informativeness of the generated image.
Sample diversity. To promote the diversity of the generated samples, we design  $S_{div}$  as the Kullback-Leibler (KL) divergence among all perturbed latent features of a seed sample:  $S_{div} = \mathcal{D}_{KL}(f' \| \bar{f})$ , where  $f'$  denotes the current perturbed latent feature and  $\bar{f}$  indicates the mean over the  $K$  perturbed features of the seed sample. To enable measuring KL divergence between features, inspired by [84], we apply the softmax function to transform feature vectors into probability vectors for KL divergence.
Theoretical analysis. We then analyze our method to conceptually understand its benefits to model generalization. We resort to  $\delta$ -cover [62] to analyze how data diversity influences the generalization error bound. Specifically, "a dataset  $E$  is a  $\delta$ -cover of a dataset  $S$  " means a set of balls with radius  $\delta$  centered at each sample of the dataset  $E$  can cover the entire dataset  $S$ . According to the property of  $\delta$ -cover, we define the dataset diversity by  $\delta$ -diversity as the inverse of the minimal  $\delta_{min}$ , i.e.,  $\delta_{div} = \frac{1}{\delta_{min}}$ . Following the same assumptions of [67], we have the following result.
Theorem 4.1. Let  $A$  denote a learning algorithm that outputs a set of parameters given a dataset  $\mathcal{D} = \{x_i, y_i\}_{i \in [n]}$  with  $n$  i.i.d. samples drawn from distribution  $\mathcal{P}_{\mathcal{Z}}$ . Assume the hypothesis function is  $\lambda^{\eta}$ -Lipschitz continuous, the loss function  $\ell(x, y)$  is  $\lambda^{\ell}$ -Lipschitz continuous for all  $y$ , and is bounded by  $L$ , with  $\ell(x_i, y_i; A) = 0$  for all  $i \in [n]$ . If  $\mathcal{D}$  constitutes a  $\delta$ -cover of  $\mathcal{P}_{\mathcal{Z}}$ , then with probability at least  $1 - \gamma$ , the generalization error bound satisfies:
$$
| \mathbb {E} _ {x, y \sim \mathcal {P} _ {\mathcal {Z}}} [ \ell (x, y; A) ] - \frac {1}{n} \sum_ {i \in [ n ]} \ell (x _ {i}, y _ {i}; A) | \stackrel {c} {\leq} \frac {\lambda^ {\ell} + \lambda^ {\eta} L C}{\delta_ {d i v}},
$$
where  $C$  is a constant, and the symbol  $\leq c$  indicates "smaller than" up to an additive constant.