SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zeyuan Yin\*1, Eric Xing1,2, Zhiqiang Shen\*1
<sup>1</sup>Mohamed bin Zayed University of AI  <sup>2</sup>Carnegie Mellon University
{zeyuan.yin,eric.xing,zhiqiang.shen}@mbzuai.ac.ae
We present a new dataset condensation framework termed Squeeze (Squeeze), Recover (Relabel) and Relabel (X) (SRe $^2$ L) that decouples the bilevel optimization of model and synthetic data during training, to handle varying scales of datasets, model architectures and image resolutions for efficient dataset condensation. The proposed method demonstrates flexibility across diverse dataset scales and exhibits multiple advantages in terms of arbitrary resolutions of synthesized images, low training cost and memory consumption with high-resolution synthesis, and the ability to scale up to arbitrary evaluation network architectures. Extensive experiments are conducted on Tiny-ImageNet and full ImageNet-1K datasets $^1$ . Under 50 IPC, our approach achieves the highest  $42.5\%$  and  $60.8\%$  validation accuracy on Tiny-ImageNet and ImageNet-1K, outperforming all previous state-of-the-art methods by margins of  $14.5\%$  and  $32.9\%$ , respectively. Our approach also surpasses MTT [1] in terms of speed by approximately  $52\times$  (ConvNet-4) and  $16\times$  (ResNet-18) faster with less memory consumption of  $11.6\times$  and  $6.4\times$  during data synthesis. Our code and condensed datasets of 50, 200 IPC with 4K recovery budget are available at link.
Figure 1: Left is data synthesis time vs. accuracy on ImageNet-1K with 10 IPC (Images Per Class). Models include ConvNet-4, ResNet-{18, 50, 101}.  $\dagger$  indicates ViT with 10M parameters [2]. Right is the comparison of widely-used bilevel optimization and our proposed decoupled training scheme.
Over the past few years, the task of data condensation or distillation has garnered significant interest within the domain of computer vision [3, 4, 5, 6, 7, 1]. By distilling large datasets into representative, compact subsets, data condensation methods enable rapid training and streamlined storage, while preserving essential information from the original datasets. The significance of data condensation on both research and applications cannot be understated, as it plays a crucial role in the efficient handling
*Equal contribution. Project page: https://zeyuanyin.github.io/projects/SRe2L/.
<sup>1</sup>We discard small datasets like MNIST and CIFAR to highlight the scalability and capability of the proposed method on large-scale datasets for real-world applications.
(1) squeeze
Figure 2: Overview of our framework. It consists of three stages: in the first stage, a model is trained from scratch to accommodate most of the crucial information from the original dataset. In the second stage, a recovery process is performed to synthesize the target data from the Gaussian noise. In the third stage, we relabel the synthetic data in a crop-level scheme to reflect the true label of the data.
and processing of vast amounts of data across numerous fields. Through the implementation of sophisticated algorithms, such as Meta-Model Matching [3, 5], Gradient Matching [8, 9, 7], Distribution Matching [10, 6], and Trajectory Matching [1, 11], data condensation has made remarkable strides. However, prior solutions predominantly excel in distilling small datasets such as MNIST, CIFAR, Tiny-ImageNet [12], downscaled ImageNet [13] featuring low resolution, or a subset of ImageNet [7]. This limitation arises due to the prohibitive computational expense incurred from executing a massive number of unrolled iterations during the bilevel optimization process (comprising an inner loop for model updates and an outer loop for synthetic data updates). In our study, employing a meticulously designed decoupling strategy for model training and synthetic data updating (as illustrated in Fig. 1 left), the proposed method is capable of distilling the entire large-scale ImageNet dataset at the conventional resolution of  $224 \times 224$  while maintaining state-of-the-art performance. Remarkably, our training/synthesis computation outstrips the efficiency of prior approaches, even those utilizing reduced resolution or subsets of ImageNet. A comparison of efficiency is provided in Table 1.
To address the huge computational and memory footprints associated with training, we propose a tripartite learning paradigm comprised of Squeeze, Recover, and Relabel stages. This paradigm allows for the decoupling of the condensed data synthesis stage from the real data input, as well as the segregation of inner-loop and outer-loop optimization. Consequently, it is not restricted by the scale of datasets, input resolution, or the size of network architectures. Specifically, in the initial phase, rather than simultaneous sampling of real and synthetic data and its subsequent network processing to update the target synthetic data, we bifurcate this process into Squeeze and Recover stages to distinguish the relationship and reduce similarity between real and synthetic data.
In the subsequent phase, we exclusively align the batch normalization (BN) [14] statistics derived from the model trained in the first phase to synthesize condensed data. In contrast to feature matching solutions that perform the matching process solely on individual batches within each iteration, the trained BN statistics on original data span the entire dataset, thereby providing a more comprehensive and accurate alignment between the original dataset and synthetic data. Given that real data is not utilized in this phase, the decoupled training can considerably diminish computational costs compared to the preceding bilevel training strategy. In the final phase, we employ the models trained in the first phase to relabel the synthetic data, serving as a data-label alignment process. An in-depth overview of this process is illustrated in Fig. 2.
Advantages. Our approach exhibits the following merits: (1) It can process large resolution condensation during training effortlessly with a justified computational cost. (2) Unlike other counterparts, we can tackle the relatively large datasets of Tiny-ImageNet and ImageNet-1K to make our method more practical for real-world applications. (3) Our approach can directly utilize many off-the-shelf pre-trained large models that contain BN layers, which enables to further save the training overhead.
Extensive experiments are performed on Tiny-ImageNet and ImageNet-1K datasets. On ImageNet-1K with  $224 \times 224$  resolution and IPC 50, the proposed approach obtains a remarkable accuracy of  $60.8\%$ , outperforming all previous methods by a large margin. We anticipate that our research will contribute to the community's confidence in the practical feasibility of large-scale dataset condensation using a decoupled synthesis strategy from the real data, while maintaining reasonable computational costs.
Methods	Condensed Arch	Time Cost (ms)	Peak GPU Memory Usage (GB)
DM [10]	ConvNet-4	18.11	10.7
MTT [1]	ConvNet-4	12.64	48.9
SRe2L (Ours)	ConvNet-4	0.24	4.2
ResNet-18	0.75	7.6
ResNet-50	2.75	33.8
Table 1: Synthesis Time and Memory Consumption on Tiny-ImageNet (64×64 resolution) using a single RTX-4090 GPU for all methods. Time Cost represents the consumption (ms) when generating one image with one iteration update on synthetic data. Peak value of GPU memory usage is measured or converted with a batch size of 200 (1 IPC as the dataset has 200 classes).
- We propose a new framework for large-scale dataset condensation, which involves a three-stage learning procedure of squeezing, recovery, and relabeling. This approach has demonstrated exceptional efficacy, efficiency and robustness in both the data synthesis and model training phases.
- We conduct a thorough ablation study and analysis, encompassing the impacts of diverse data augmentations for original data compression, various regularization terms for data recovery, and diverse teacher alternatives for relabeling on the condensed dataset. The comprehensive specifications of the learning process can offer valuable insights for subsequent investigations in the domain.
- To the best of our knowledge, this is the first work that enables to condense the full ImageNet-1K dataset with an inaugural implementation at a standard resolution of  $224 \times 224$ , utilizing widely accessible NVIDIA GPUs such as the 3090, 4090, or A100 series. Furthermore, our method attains the highest accuracy of  $60.8\%$  on full ImageNet-1K within an IPC constraint of 50 using justified training time and memory cost, outperforming all previous methods by a significant margin.
Data Condensation/Distillation. The objective of dataset condensation is to acquire a small synthetic dataset that retains a substantial amount of information present in the original data. Suppose we have a large labeled dataset  $\mathcal{T} = \left\{(\boldsymbol{x}_1, \boldsymbol{y}_1), \ldots, (\boldsymbol{x}_{|\mathcal{T}|}, \boldsymbol{y}_{|\mathcal{T}|})\right\}$ , we aim to learn a small condensed dataset  $\mathcal{C}_{\mathrm{syn}} = \left\{(\widetilde{\boldsymbol{x}}_1, \widetilde{\boldsymbol{y}}_1), \ldots, (\widetilde{\boldsymbol{x}}_{|\mathcal{C}|}, \widetilde{\boldsymbol{y}}_{|\mathcal{C}|})\right\} (|\mathcal{C}| \ll |\mathcal{T}|)$  that preserves the crucial information in original  $\mathcal{T}$ . The learning objective on condensed synthetic data is:
$$
\theta_ {\mathcal {C} _ {\text {s y n}}} = \underset {\boldsymbol {\theta}} {\arg \min } \mathcal {L} _ {\mathcal {C}} (\boldsymbol {\theta}) \tag {1}
$$
where  $\mathcal{L}_{\mathcal{C}}(\pmb {\theta}) = \mathbb{E}_{(\widetilde{\pmb{x}},\widetilde{\pmb{y}})\in \mathcal{C}_{\mathrm{syn}}}\Bigl [\ell (\phi_{\pmb{\theta}_{\mathcal{C}_{\mathrm{syn}}}}(\widetilde{\pmb{x}}),\widetilde{\pmb{y}})\Bigr ]$ $\widetilde{\pmb{y}}$  is the soft label corresponding to the synthetic data  $\widetilde{\pmb{x}}$
The ultimate goal of data condensation task is to synthesize data for achieving a certain/minimal supremum performance gap on original evaluation set when models are trained on the synthetic data and the original full dataset, respectively. Following the definition of coresets [15] and  $\epsilon$ -approximate [16], the objective of data condensation task can be formulated as achieving:
$$
\left. \sup  \left\{\left| \ell \left(\phi_ {\boldsymbol {\theta} _ {\mathcal {T}}} (\boldsymbol {x}), \boldsymbol {y}\right) - \ell \left(\phi_ {\boldsymbol {\theta} _ {\mathrm {C} _ {\mathrm {s y n}}}} (\boldsymbol {x}), \boldsymbol {y}\right) \right| \right\} _ {(\boldsymbol {x}, \boldsymbol {y}) \sim \mathcal {T}} \leq \epsilon \right. \tag {2}
$$
where  $\epsilon$  is the performance gap for models trained on the synthetic data and the original full dataset. Thus, we aim to optimize the synthetic data  $\mathcal{C}_{\mathrm{syn}}$  through:
$$
\underset {\mathcal {C} _ {\mathrm {s y n}}, | \mathcal {C} |} {\arg \min } \left(\sup  \left\{\left| \ell \left(\phi_ {\boldsymbol {\theta} _ {\mathcal {T}}} (\boldsymbol {x}), \boldsymbol {y}\right) - \ell \left(\phi_ {\boldsymbol {\theta} _ {\mathcal {C} _ {\mathrm {s y n}}}} (\boldsymbol {x}), \boldsymbol {y}\right) \right| \right\} _ {(\boldsymbol {x}, \boldsymbol {y}) \sim \mathcal {T}}\right) \tag {3}
$$
Then, we can learn  $<\text{data}, \text{label}> \in \mathcal{C}_{\text{syn}}$  with the associated number of condensed data in each class.
Decoupling the condensed data optimization and the neural network training: Conventional solutions, such as FRePo [5], CAFE [6], DC [8], typically choose for the simultaneous optimization of the backbone network and synthetic data within a singular training framework, albeit in an iterative fashion. The primary drawback associated with these joint methods is their computational burden due to the unrolling of the inner-loop during each outer-loop update, coupled with bias transference
from real to synthetic data as a result of truncated unrolling. The objective of this study is to devise an efficient learning framework capable of individually decoupling model training and synthetic data optimization. This approach circumvents information bias stemming from real data, concurrently enhancing efficiency in handling diverse scales of datasets, model architectures, and image resolutions, thereby bolstering effective dataset condensation. Our framework is predicated on the assumption that pivotal information within a dataset can be adequately trained and preserved within a deep neural network. The training procedure of our approach is elaborated in the following section.
2.1 Decoupling Outer-loop and Inner-loop Training
Inspired by recent advances in DeepDream [17], Inverting Image [18, 19] and data-free knowledge transfer [20], we propose a decoupling approach to disassociate the traditional bilevel optimization inherent to dataset condensation. This is accomplished via a tripartite process to reformulate it into a unilevel learning procedure.
Stage-1 Squeeze (): During this stage, our objective is to extract pertinent information from the original dataset and encapsulate it within the deep neural networks, evaluating the impact of various data augmentation techniques, training strategies, etc. Deep neural networks typically comprise multiple parametric functions, which transform high-dimensional original data (e.g., pixel space of images) into their corresponding low-dimensional latent spaces. We can exploit this attribute to abstract the original data to the pretrained model and then reconstruct them in a more focused manner, akin to DeepDream [17] and Inverting Images [18, 19]. It's noteworthy that the purpose of this stage is to extract and encapsulate critical information from the original dataset. Hence, excessive data augmentation resulting in enhanced performance does not necessarily lead to the desired models. This approach diverges from previous solutions that sample two data batches from the original large-scale dataset  $\mathcal{T}$  and the learnable small synthetic dataset  $\mathcal{C}$ . The learning procedure can be simply cast as a regular model training process on the original dataset with a suitable training recipe:
$$
\boldsymbol {\theta} _ {\mathcal {T}} = \underset {\boldsymbol {\theta}} {\arg \min } \mathcal {L} _ {\mathcal {T}} (\boldsymbol {\theta}) \tag {4}
$$
where  $\mathcal{L}_{\mathcal{T}}(\pmb{\theta})$  typically uses cross-entropy loss as  $\mathcal{L}_{\mathcal{T}}(\pmb{\theta}) = \mathbb{E}_{(\pmb{x},\pmb{y})\in \mathcal{T}}[\pmb{y}\log (\pmb{p}(\pmb{x}))]$ .
Enabling BN Layers in ViT for Recovering Process: In contrast to distribution matching [10] that aligns the feature distributions of the original and synthetic training data in sampled embedding spaces, thus allowing for the use of a randomly initialized network, our matching mechanism is solely performed on the Batch Normalization (BN) layers using their statistical properties, akin to data-free knowledge transfer [20]. Unlike the feature matching solution which executes the matching process on individual batches within each iteration, the referential BN statistics are calculated over the entirety of the dataset, providing a more comprehensive and representative alignment with the original dataset. Our experiments empirically substantiate that BN-matching can significantly outperform the feature-matching method. BN layer is commonly used in ConvNet but is absent in ViT. To utilize both ConvNet and ViT for our proposed data condensation approach, we engineer a BN-ViT which replaces all LayerNorm by BN layers and adds additional BN layers in-between two linear layers of feed-forward network, as also utilized in [21]. This marks the first instance of broadening the applicability of the data condensation architecture from ConvNets to encompass ViTs as well.
Stage-2 Recover (): This phase involves reconstructing the retained information back into the image space utilizing class labels, regularization terms, and BN trajectory alignment. Unlike conforming to batch feature distributions or comprehensive parameter distributions, we solely track the distribution of BN statistics derived from the original dataset. The pairing of BN and predictive probability distribution restricts the optimization process to a singular level, thereby significantly enhancing scalability. By aligning the final classification and intermediary BN statistics (mean and variance), the synthesized images are compelled to encapsulate a portion of the original image distribution. The learning objective for this phase can be formulated as follows:
$$
\underset {\mathcal {C} _ {\text {s y n}}, | \mathcal {C} |} {\arg \min } \ell \left(\phi_ {\boldsymbol {\theta} _ {\mathcal {T}}} \left(\widetilde {\boldsymbol {x}} _ {\text {s y n}}\right), \boldsymbol {y}\right) + \mathcal {R} _ {\text {r e g}} \tag {5}
$$
where  $\mathcal{R}_{\mathrm{reg}}$  is the regularization term.  $\phi_{\theta_T}$  is the model pre-trained in the first stage and it will be frozen in this stage, we solely optimize the  $\widetilde{\pmb{x}}_{\mathrm{syn}}$  as a single-level training process. Following [18, 20], we discuss three regularizers that can be used and the ablation for them is provided in our experiments. The first two regularizers are image prior regularisers of  $\ell_2$  and total variation (TV) proposed in [18]:
$$
\mathcal {R} _ {\text {p r i o r}} \left(\widetilde {\boldsymbol {x}} _ {\text {s y n}}\right) = \alpha_ {\mathrm {t v}} \mathcal {R} _ {\mathrm {T V}} \left(\widetilde {\boldsymbol {x}} _ {\text {s y n}}\right) + \alpha_ {\ell_ {2}} \mathcal {R} _ {\ell_ {2}} \left(\widetilde {\boldsymbol {x}} _ {\text {s y n}}\right) \tag {6}
$$
where  $\mathcal{R}_{\ell_2} = \| \widetilde{\boldsymbol{x}}_{\mathrm{syn}}\| _2$ , this regularizer encourages image to stay within a target interval instead of diverging.  $\mathcal{R}_{\mathrm{TV}} = \sum_{i,j}\left((\widetilde{\boldsymbol{x}}_{i,j + 1} - \widetilde{\boldsymbol{x}}_{ij})^2 +(\widetilde{\boldsymbol{x}}_{i + 1,j} - \widetilde{\boldsymbol{x}}_{ij})^2\right)^{\frac{\beta}{2}}$  where  $\beta$  is used to balance the sharpness of the image with the removal of "spikes" and smooth the synthetic data. While, this is not necessary for dataset condensation since we care more about information recovery.
Learning Condensed Data with BN Consistency: DeepInversion [20] utilizes the feature distribution regularization term to improve the quality of the generated images. Here, we also leverage this property as our recovering loss term. It can be formulated as:
$$
\begin{array}{l} \mathcal {R} _ {\mathrm {B N}} (\widetilde {\boldsymbol {x}}) = \sum_ {l} \| \mu_ {l} (\widetilde {\boldsymbol {x}}) - \mathbb {E} \left(\mu_ {l} \mid \mathcal {T}\right) \| _ {2} + \sum_ {l} \left\| \sigma_ {l} ^ {2} (\widetilde {\boldsymbol {x}}) - \mathbb {E} \left(\sigma_ {l} ^ {2} \mid \mathcal {T}\right) \right\| _ {2} \\ \approx \sum_ {l} ^ {\iota} \left\| \mu_ {l} (\widetilde {\boldsymbol {x}}) - \mathbf {B N} _ {l} ^ {\mathrm {R M}} \right\| _ {2} + \sum_ {l} ^ {\iota} \left\| \sigma_ {l} ^ {2} (\widetilde {\boldsymbol {x}}) - \mathbf {B N} _ {l} ^ {\mathrm {R V}} \right\| _ {2} \tag {7} \\ \end{array}
$$
where  $l$  is the index of BN layer,  $\mu_l(\widetilde{\pmb{x}})$  and  $\sigma_l^2 (\widetilde{\pmb{x}})$  are mean and variance.  $\mathbf{BN}_l^{\mathrm{RM}}$  and  $\mathbf{BN}_l^{\mathrm{RV}}$  are running mean and running variance in the pre-trained model at  $l$ -th layer, which are globally counted.
Multi-crop Optimization: RandomResizedCrop, a frequently utilized technique in neural network training, serves as a preventative measure against overfitting. Inspired by this, we propose the strategy of multi-crop optimization during the process of image synthesis with the aim of enhancing the informational content of the synthesized images. In practice, we implement it by randomly cropping on the entire image and subsequently resizing the cropped region to the target dimension of  $224 \times 224$ . Under these circumstances, only the cropped region is updated during each iteration. This approach aids in refining the recovered data when viewed from the perspective of cropped regions.
Stage-3 Relabel  $(\times)$ : To match our multi-crop optimization strategy, also to reflect the true soft label for the recovered data. We leverage the pre-generated soft label approach as FKD [22].
$$
\widetilde {\boldsymbol {y}} _ {i} = \phi_ {\boldsymbol {\theta} _ {\mathcal {T}}} \left(\widetilde {\boldsymbol {x}} _ {\mathbf {R} _ {i}}\right) \tag {8}
$$
where  $\widetilde{\pmb{x}}_{\mathbf{R}_i}$  is the  $i$ -th crop in the synthetic image and  $\widetilde{\pmb{y}}_i$  is the corresponding soft label. Finally, we can train the model  $\phi_{\theta_{\mathrm{syn}}}$  on the synthetic data using the following objective:
$$
\mathcal {L} _ {\mathrm {s y n}} = - \sum_ {i} \widetilde {\boldsymbol {y}} _ {i} \log \phi_ {\boldsymbol {\theta} _ {\mathcal {C} _ {\mathrm {s y n}}}} \left(\widetilde {\boldsymbol {x}} _ {\mathbf {R} _ {i}}\right) \tag {9}
$$
We found that this stage is crucial to make synthetic data and labels more aligned and also significantly improve the performance of the trained models.
Discussion: How does the proposed approach reduce compute and memory consumption? Existing solutions predominantly employ bilevel optimization [3, 23, 4] or long-range parameter matching strategies [1, 11], which necessitate the feeding of real data into the network to generate guiding variables (e.g., features, gradients, etc.) for target data updates, as well as for the backbone network training, through an iterative process. These approaches incur considerable computational and memory overhead due to the concurrent presence of real and synthetic data on computational hardware such as GPUs, thereby rendering this training strategy challenging to scale up for larger datasets and models. To this end, a natural idea is to decouple real and synthetic data during the training phase, thereby necessitating only minimal memory during each training session. This is achieved by bifurcating the bilevel training into a two-stage process: squeezing and recovering. Moreover, we can conveniently utilize off-the-shelf pre-trained models for the first squeezing stage.
In this section, we evaluate the performance of our proposed  $\mathrm{SRe}^2\mathrm{L}$  over various datasets, models and tasks. First, we conduct extensive ablation experiments to investigate the effect of each component in three stages. Next, we demonstrate the superior results of  $\mathrm{SRe}^2\mathrm{L}$  in large-scale datasets, cross-architecture generalization, and continual learning application. Finally, we provide the comparison of visualizations on distilled data with other state-of-the-art methods.
Experiment Setting. We evaluate our method  $\mathrm{SRe}^2\mathrm{L}$  on two large-scale datasets Tiny-ImageNet [12] and full ImageNet-1K [24]. Detailed comparison among variants of ImageNet-1K is in appendix. For backbone networks, we employ ResNet-{18, 50, 101} [25], ViT-Tiny [26], and our new constructed BN-ViT-Tiny (Sec. 2.1) as the target model training. For distilling ImageNet-1K, we
Recovery model	Squeezing Budget (iteration)	Data Augmentation
50	100	200	400	Mixup	CutMix
ResNet-18	37.88	37.49	32.18	23.88	35.59	25.90
ResNet-50	39.19	33.63	30.65	26.43	36.95	29.61
Table 2: Ablation of squeezing budget and data augmentation on Tiny-ImageNet.
recover the data from PyTorch off-the-shelf pre-trained ResNet-\{18, 50\} with the Top-1 accuracy of  $\{69.76\%, 76.13\% \}$  for saving re-training computational overhead. For distilling Tiny-ImageNet, ResNet-\{18, 50\} are used as base models, while the first  $7 \times 7$  Conv layer is replaced by  $3 \times 3$  Conv layer and the maxpool layer is discarded, following MoCo (CIFAR) [27]. After that, they are trained from scratch on Tiny-ImageNet. More implementation details are provided in appendix.
Evaluation and baselines. Following previous works [3, 1, 8], we evaluate the quality of condensed datasets by training models from scratch on them and report the test accuracy on real val datasets.
3.1 Squeezing Analysis