SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Enneng Yang $^{1}$ , Li Shen $^{2*}$ , Zhenyi Wang $^{3*}$ , Tongliang Liu $^{4}$ , Guibing Guo $^{1*}$
$^{1}$ Northeastern University, China  $^{2}$ JD Explore Academy, China
<sup>3</sup>University of Maryland, College Park, USA <sup>4</sup>The University of Sydney, Australia
ennengyang@stumail.neu.edu.cn; mathshenli@gmail.com; zwang169@umd.edu
tongliang.liu@sydney.edu.au; guogb@swc.neu.edu.cn
Dataset condensation (DC) distills a large real-world dataset into a small synthetic dataset, with the goal of training a network from scratch on the latter that performs similarly to the former. State-of-the-art (SOTA) DC methods have achieved satisfactory results through techniques such as accuracy, gradient, training trajectory, or distribution matching. However, these works all perform matching in the high-dimension pixel space, ignoring that natural images are usually locally connected and have lower intrinsic dimensions, resulting in low condensation efficiency. In this work, we propose a simple-yet-efficient dataset condensation plugin that matches the raw and synthetic datasets in a low-dimensional manifold. Specifically, our plugin condenses raw images into two low-rank matrices instead of parameterized image matrices. Our plugin can be easily incorporated into existing DC methods, thereby containing richer raw dataset information at limited storage costs to improve the downstream applications' performance. We verify on multiple public datasets that when the proposed plugin is combined with SOTA DC methods, the performance of the network trained on synthetic data is significantly improved compared to traditional DC methods. Moreover, when applying the DC methods as a plugin to continual learning tasks, we observed that our approach effectively mitigates catastrophic forgetting of old tasks under limited memory buffer constraints and avoids the problem of raw data privacy leakage.
Modern machine learning techniques utilize large-scale real-world datasets and advanced deep neural networks (DNNs) to achieve amazing success in various fields, such as models like SAM [24] and GPT [6, 38] in the fields of computer vision and natural language processing, both of which have surpassed classical models [51, 42, 43] trained on small datasets. However, training a well-performing model in the machine learning community requires repeated tuning of various aspects of the model [14], such as the number of layers, learning rate, and other important hyper-parameters. When the dataset is large, the cost of data management and repeated model training becomes unacceptable. As stated in [65], NAS-Bench-101 [63] spent 100 TPU-years of computing time on the CIFAR10 dataset [26] for an exhaustive neural architecture search. In contrast, it only takes dozens of TPU minutes to train the best model from scratch using the optimal configuration discovered [65]. Additionally, in continual learning [41, 2, 60, 61], to prevent forgetting old tasks while learning new tasks, a certain amount of old task data is typically stored for replay [58, 57, 2, 41]. Nevertheless, storing large old task datasets is unrealistic due to strict memory and privacy constraints. Therefore, reducing the data size becomes a valuable and emerging research direction.
*Corresponding authors.
A direct way to reduce data size is to select a representative subset from the original dataset [17]. This paradigm calculates the importance score for each sample based on specific importance criteria and selects a subset to replace the entire training dataset [62]. The calculation criteria include the distance between the sample and the class center [10], the gradient norm of sample [39], and the diversity among samples [47, 57], etc. However, the selection-based method is not always effective, particularly when the task-condition data information is uniformly distributed in the original data [30]. In such cases, the heuristic coreset selection method can only capture a limited amount of information. Recent research [56] has shown that generating a small dataset that performs similarly to the original dataset, instead of selecting a subset, is a promising direction. Therefore, dataset condensation (DC) or dataset distillation (DD) [56, 66, 64, 65, 37, 11, 15] has recently received increasing attention. This novel paradigm distills information from a large real-world dataset into a compact synthetic dataset that produces comparable results to the original dataset. The earliest DD [56] method uses accuracy value as the direct distillation objective and describes the distillation process as a bi-level optimization problem. However, matching accuracy directly through bi-level optimization involves high computational costs and memory overhead [15]. Therefore, state-of-the-art (SOTA) DC methods perform condensation through surrogate objectives, such as gradient [66, 64, 22], training trajectory [8], feature [55], or distribution [65] matching, and achieve more satisfactory results.
However, existing DC methods all optimize parameterized condensed images in the original high-dimensional pixel space, overlooking the fact that natural images typically exhibit local connectivity and have low intrinsic dimensionality [31]. More specifically, in the visual self-supervision task, Masked Autoencoder [18] divides the image into multiple patches and reconstructs the complete image by randomly masking a certain proportion of patches. Experimental results show that masking  $75\%$  of the patches can still reconstruct the original image. Also, Maximum Likelihood Estimation [31, 40] of the image's intrinsic dimensions on ImageNet [12] dataset shows that although each image contains 150, 528 pixels, its intrinsic dimension is only between 26 and 43. We further verify in the appendix that both the original images and the images generated by traditional dataset distillation (e.g., DSA [64], DM [65]) are low rank. These results imply that compressing data in high-dimensional pixel space is an inefficient approach.
In this work, we propose a simple yet efficient plugin for dataset condensing that compresses a large dataset into a compact synthetic dataset on a low-dimensional manifold. Specifically, unlike existing DC methods that train a compact dataset  $S$  of size  $N \times D \times H \times W$ , where  $N$ ,  $D$ ,  $H$ , and  $W$  are the number, channels, height, and width of images, respectively, we decompose an image representation under each channel  $ch \in \{1, \dots, D\}$  into a low-dimensional space, and learn two low-rank matrices  $\mathcal{A} \in \mathbb{R}^{N \times D \times H \times r}$  and  $\mathcal{B} \in \mathbb{R}^{N \times D \times r \times W}$ , where  $r << \min\{H, W\}$  represents the rank size. Obviously, our low-rank method is orthogonal to existing DC methods and can be integrated into SOTA DC methods as a flexible plugin to improve the learning efficiency of DC. Without loss of generality, we conduct extensive experiments on multiple publicly real-world datasets, integrating the proposed plugin into gradient matching-based [66, 64] and distribution matching-based [65] DC methods. Experimental results show that using our low-rank plugin significantly reduces storage consumption for compact datasets and achieves comparable performance as SOTA DC methods based on high-dimensional pixel spaces, with the same number of images. Furthermore, under the same memory consumption, our plugin can effectively store more information from the large dataset, thus significantly improving the performance of condensed data in downstream tasks. Lastly, when applying low-rank DC methods to continual learning (CL) tasks, we observe that our approach effectively mitigates catastrophic forgetting of old tasks under the constraints of limited memory buffers and avoids data privacy issues by storing condensed rather than raw data.
The main contributions of this paper are as follows:
- We propose a simple yet effective dataset distillation plugin that condenses a large dataset into a compact synthetic dataset on a low-dimensional manifold, offering an orthogonal approach to existing DC methods.
- Experiments on deep learning tasks demonstrate that the proposed plugin achieves performance comparable to that of existing SOTA DC methods while significantly reducing memory consumption. It achieves significantly better accuracy than existing SOTA DC methods with the same memory consumption.
- We also verify that the dataset condensed in the low-dimensional manifold has good cross-architecture transferability and maintains the excellent characteristics of existing DC methods, such as learning the class distribution information of the large dataset.
- When applying low-rank DC as a plugin for CL tasks, our DC plugin approach effectively mitigates the problem of catastrophic forgetting of old tasks while protecting the data privacy.
In this section, we summarize the most related work to this work as three-fold, including coreset sample selection and dataset distillation; continual learning; and low-rank manifolds.
Coreset Selection and Data Condensation. Coreset selection [17] and dataset condensation [56] are two methods to eliminate data redundancy, which help to improve the model's training efficiency and reduce the cost of data management. Coreset selection has been widely studied in active learning [48] and continual learning [34], which tries to identify the most informative training samples [62, 59, 10, 44, 3]. Unfortunately, these methods fail when the task-condition information is evenly distributed among the original samples [30]. Empirical studies [56, 66] show that the benefit of existing coreset selection methods is marginal compared with random selection across multiple benchmarks. In recent years, dataset condensation [56] has been proposed to distill a large real dataset onto a small synthetic dataset and achieve better performance than coreset selection. DD [56] formulates the dataset distillation as a bi-level optimization. The inner loop utilizes the condensed dataset to train a network, while the outer loop minimizes the prediction error of the trained network on the original dataset. It updates the synthesized dataset pixel by pixel using Back-Propagation Through Time [35] to compute the meta gradient. KRR [36] transforms the distillation problem into a kernel ridge regression problem, simplifying the expensive nested optimization in DD [56] to a first-order optimization. Unlike the aforementioned works that optimize based on accuracy value, recent works believe that the effectiveness of models trained on the larger dataset and the compact dataset can be reflected in their corresponding parameter states or gradient states [15]. Therefore, they choose to optimize more fine-grained surrogate objectives [66, 64, 30, 8, 65]. Notably, DC [66] and DSA [64] minimize the gradient matching loss between the large dataset and the synthetic dataset at each training step, MTT [8] matches the training trajectory of parameters, LCMat [50] matches loss curvature, CAFE [55] aligns layer-wise features, and DM [65] generates synthetic samples that resemble the distribution of real samples in the feature space. There has recently been a small amount of work on parameter-efficient dataset distillation. For example, IDC [23] and IDM [67] take advantage of the local similarity characteristics of images to partition and expand condensed data to generate more images at the same memory cost. HaBa [33] and RememberThePast [13] utilize bases to produce images or share bases among all classes respectively. However, these methods ignore that natural images have low intrinsic dimensionality [40], which leads to inefficient condensation.
Continual Learning. Continual learning (CL) aims to enable a neural network model to learn new tasks continuously without forgetting the old tasks [68, 58]. Inspired by the working mechanism of the human brain, mainstream memory-based CL methods consolidate previously learned knowledge by replaying old data, thereby avoiding catastrophic forgetting. Due to strict memory and privacy constraints, usually only a small portion of old task data can be kept, many above-mentioned coreset selection methods are used to select informative samples for storage. For example, some heuristic method is used to select the most representative samples from each class [59] (e.g., iCaRL [44]) or the sample closest to the decision boundary (e.g., Rwalk [9], MIR [2], Shim [49]). However, the number of new task's samples and stored old samples in memory is often highly unbalanced, leading the model to be biased towards learning new tasks with more data. Recent works [7, 41] addressing this class imbalance in memory-based CL has achieved impressive and satisfactory performance: End-to-End Incremental Learning [7] samples the data after training on the new classes and combines it with the data in memory to create balanced data for fine-tuning. GDumb [41] greedily stores samples in memory and trains the model from scratch on new tasks using only the samples in memory.
Low-rank Manifolds. The low-rank structure is prevalent in machine learning, such as computer vision [40, 18, 32], natural language processing [20, 1], recommendation systems [25, 45], etc. For example, MLE [31, 40] found that natural image datasets containing thousands of pixels can actually be described with fewer variables. LoRA [20] can fine-tune the GPT-3 175B [6] with comparable accuracy to full-parameter fine-tuning while reducing the number of trainable parameters by a factor of 10,000 and GPU memory requirements by a factor of three. In recommender systems, the original highly sparse rating/interaction matrix is usually decomposed into low-rank user embedding matrices and item embedding matrices [25] to capture user interests and item features. However, to the best
of our knowledge, no research has explored the exploitation of low-rank structure to improve the condensation efficiency for the dataset distillation task, and this work is the first to fill this gap.
3 Low-Rank Data Condensation Plugin
In this section, we first define the problem of DC in Sec. 3.1, then introduce our proposed low-rank plugin for DC in Sec. 3.2, and then describe how to integrate our plugin with existing DC methods in Sec. 3.3 and further apply it to continual learning in Sec. 3.4.
The goal of dataset condensation (DC) is to distill a large target dataset  $\mathcal{T} = \{\mathbf{x}_i, y_i\}_{i=1}^{N_{\mathcal{T}}}$  containing  $N_{\mathcal{T}}$  training image  $\mathbf{x}_i \in \mathbb{R}^d$  and its label  $y_i \in \{1, 2, \dots, |Y|\}$  into a small dataset  $\mathcal{S} = \{\mathbf{s}_i, y_i\}_{i=1}^{N_{\mathcal{S}}}$  with  $|S|$  synthetic image  $\mathbf{s}_i \in \mathbb{R}^d$ , where  $N_{\mathcal{S}} \ll N_{\mathcal{T}}$  ( $2 - 3$  orders of magnitude),  $|Y|$  represents the number of classes, and  $\mathbb{R}^d$  defines a  $d$ -dimensional space. We expect a network  $\phi_{\theta^S}$  trained on the small dataset  $\mathcal{S}$  to have similar performance to a network  $\phi_{\theta^T}$  trained on the large training set  $\mathcal{T}$  on the unseen test dataset, that is:
$$
\mathbb {E} _ {\mathbf {x} _ {i} \sim P _ {T}} \left[ \ell \left(\phi_ {\theta^ {T}} (\mathbf {x} _ {i}), y\right) \right] \simeq \mathbb {E} _ {\mathbf {x} _ {i} \sim P _ {T}} \left[ \ell \left(\phi_ {\theta^ {S}} (\mathbf {x} _ {i}), y\right) \right],
$$
$$
\text {s . t .} \theta^ {\mathcal {T}} = \underset {\theta^ {\mathcal {T}}} {\arg \min } \mathcal {L} ^ {\mathcal {T}} \left(\theta^ {\mathcal {T}}\right) = \underset {\theta^ {\mathcal {T}}} {\arg \min } \frac {1}{N _ {\mathcal {T}}} \sum_ {\left(\mathbf {x} _ {i}, y\right) \in \mathcal {T}} \ell \left(\phi_ {\theta^ {\mathcal {T}}} \left(\mathbf {x} _ {i}\right), y\right), \tag {1}
$$
$$
\theta^ {S} = \underset {\theta^ {S}} {\arg \min} \mathcal {L} ^ {S} (\theta^ {S}) = \underset {\theta^ {S}} {\arg \min} \frac {1}{N _ {S}} \sum_ {(\mathbf {x} _ {i}, y) \in S} \ell (\phi_ {\theta^ {S}} (\mathbf {x} _ {i}), y),
$$
where  $P_{\mathcal{T}}$  represents the real distribution of the test dataset,  $\mathbf{x}_i$  represents the input image,  $y$  represents the ground truth, and  $\ell(\cdot)$  represents a loss function such as cross-entropy loss.
To achieve the goal stated in Eq. 1, existing DC methods [56, 66, 64, 8, 65] first initialize the dataset  $S \in \mathbb{R}^{N_S \times D \times H \times W}$  as a set of learnable parameters in high-dimensional pixel space. Here,  $N_S$  denotes the number of synthetic images, and  $C$ ,  $H$ , and  $W$  represents the number of channels, the image's height, and the image's width, respectively. Then, the dataset  $S$  is updated pixel-by-pixel based on accuracy value matching [56, 36] or surrogate objective matching [66, 8, 65] of the condensed dataset  $S$  and the large dataset  $\mathcal{T}$ . For example, in the first dataset distillation work DD [56], dataset  $S$  is treated as a hyperparameter in a bi-level optimization problem as follows:
$$
\mathcal {S} ^ {*} = \underset {\mathcal {S}} {\arg \min } \mathcal {L} ^ {\mathcal {T}} \left(\phi_ {\theta} s\right), \text {s u b j e c t} \theta^ {S} = \underset {\theta} {\arg \min } \mathcal {L} ^ {S} \left(\phi_ {\theta}\right), \tag {2}
$$
where the inner loop trains a randomly initialized network  $\phi$  (parameters denoted as  $\theta$ ) on the synthetic dataset  $S$  until convergence, and the outer loop uses the large target dataset  $\mathcal{T}$  as a validation set to optimize  $S$ . The small dataset  $S$  is updated by solving the meta-gradient [35] of the bile-level optimization in Eq. 2, allowing the trained model  $\phi$  on the dataset  $S$  to perform well on the real dataset  $\mathcal{T}$ . Instead of optimizing directly based on the accuracy value of the distilled data  $S$ , SOTA DC methods are based on surrogate objectives to make the model trained on  $S$  and  $\mathcal{T}$  approximate each other in the parameter space [66], i.e.,  $\theta^{\mathcal{T}} \simeq \theta^{\mathcal{S}}$ , or in the gradient space [64, 66], i.e.,  $\nabla_{\theta} \mathcal{L}^{\mathcal{T}}(\theta) \simeq \nabla_{\theta} \mathcal{L}^{\mathcal{S}}(\theta)$ , or in the feature distribution space [65], i.e.,  $\phi_{\theta}(\mathbf{x}_i) \simeq \phi_{\theta}(\mathbf{s}_i)$ . However, these methods all focus on pixel-by-pixel optimizations of  $S$ , ignoring that natural images are locally connected [18] and usually have low intrinsic dimensionality [32, 40]. In the appendix, we perform principal component analysis on real images and images obtained by traditional dataset distillation methods (such as DSA [64] and DM [65]) and show that they have low-rank properties. Hence, optimizing  $S$  in a high-dimensional pixel space is inefficient.
3.2 Our Low-Rank Data Condensation Plugin
In this work, we introduce a low-rank dataset condensation plugin that distills knowledge from a large dataset  $\mathcal{T}$  to a small synthetic dataset  $S$  in a low-rank manifold. Specifically, instead of directly initializing  $S$  as a variable with shape  $N_{S} \times D \times H \times W$ , we conduct a low-rank decomposition of the content  $\mathbf{x}_{i,ch} \in \mathbb{R}^{H \times W}$  in the channel  $ch \in \{1, \dots, D\}$  of an image  $\mathbf{x}_i$  and employ two variables,  $\mathcal{A}_{i,ch} \in \mathbb{R}^{H \times r}$  and  $\mathcal{B}_{i,ch} \in \mathbb{R}^{r \times W}$ , to reconstruct a high-dimensional image by  $\mathbf{x}_{i,ch} \approx \mathcal{A}_{i,ch} \mathcal{B}_{i,ch}$ ,
Algorithm 1: LoDC: Low-rank Dataset Condensation with Gradient Matching [66]
Input: A large training dataset  $\mathcal{T}$    
Required: Randomly initialized  $\mathcal{A},\mathcal{B}$  of rank  $r$  for  $|Y|$  classes, the probability distribution of randomly initialized weights  $P_{\theta_0}$  , neural network  $\phi_{\theta}$  , number of outer-iterations  $T_{ou}$  and number of inner-iterations  $T_{in}$  , learning rates  $\eta_{dc}$  and  $\eta_{\theta}$  , minibatch size  $N\tau$ $N_{AB}$    
for outer-iteration  $k = 1,\dots ,T_{ou}$  do   
Initialize  $\theta_0\sim P_{\theta_0}$    
for inner-iteration  $t = 1,\dots ,T_{in}$  do   
for class  $c = 1,\dots ,|Y|$  do Sample a minibatch real pair  $\mathbf{B}_c^{\mathcal{T}}\sim \mathcal{T}$  and a minibatch synthetic pair  $\mathbf{B}_c^{\Omega (\mathcal{A},\mathcal{B})}\sim \Omega (\mathcal{A},\mathcal{B})$  Compute  $\mathcal{L}_c^\top (\theta_t|\mathbf{B}_c^\top) = \frac{1}{N_\top}\sum_{(\mathbf{x}_i,y)\in \mathbf{B}_c^\top}\ell (\phi_{\theta_t}(\mathbf{x}_i),y)$  Compute  $\mathcal{L}_c^{\Omega (\mathcal{A},\mathcal{B})}(\theta_t|\mathbf{B}_c^{\Omega (\mathcal{A},\mathcal{B})}) = \frac{1}{N_{\mathcal{A}\mathcal{B}}}\sum_{(\mathcal{A}_i\mathcal{B}_i,y)\in \mathbf{B}_c^{\Omega (\mathcal{A},\mathcal{B})}}\ell (\phi_\theta (\mathcal{A}_i\mathcal{B}_i),y)$  Compute  $\mathcal{L}_c^{DC}(\mathbf{B}_c^{\mathcal{T}},\mathbf{B}_c^{\Omega (\mathcal{A},\mathcal{B})}) = d\left(\nabla_\theta \mathcal{L}_c^{\mathcal{T}}(\theta_t|\mathbf{B}_c^{\mathcal{T}}),\nabla_\theta \mathcal{L}_c^{\Omega (\mathcal{A},\mathcal{B})}(\theta_t|\mathbf{B}_c^{\Omega (\mathcal{A},\mathcal{B}))\right)$  Update  $\mathcal{A}\gets \mathcal{A} - \eta_{dc}\nabla_{\mathcal{A}}\mathcal{L}_{c}^{DC}(\mathbf{B}_{c}^{\top},\mathbf{B}_{c}^{\Omega (\mathcal{A},\mathcal{B})})$  and  $\mathcal{B}\gets \mathcal{B} - \eta_{dc}\nabla_{\mathcal{B}}\mathcal{L}_{c}^{DC}(\mathbf{B}_{c}^{\top},\mathbf{B}_{c}^{\Omega (\mathcal{A},\mathcal{B})})$  Update  $\theta_{t + 1}\gets \theta_t - \eta_\theta \nabla_\theta \mathcal{L}^{\Omega (\mathcal{A},\mathcal{B})}(\theta_t)$    
Output: A small dataset  $\Omega (\mathcal{A},\mathcal{B})$
where the rank  $r \ll \{H, W\}$ . Therefore, the goal of data condensation in the low-rank manifold is to optimize  $\mathcal{A} \in \mathbb{R}^{N_S \times D \times H \times r}$  and  $\mathcal{B} \in \mathbb{R}^{N_S \times D \times r \times W}$  such that the network  $\phi_{\theta^{\Omega}(\mathcal{A}, \mathcal{B})}$ , trained on the small reconstructed data  $\Omega(\mathcal{A}, \mathcal{B})$ , achieves similar performance to the network  $\phi_{\theta^{\mathcal{T}}}$  trained on the high-dimensional large dataset  $\mathcal{T}$ . Therefore, the DC in the low-rank manifold is formalized as:
$$
\mathbb {E} _ {\mathbf {x} _ {i} \sim P _ {T}} \left[ \ell \left(\phi_ {\theta} \tau (\mathbf {x} _ {i}), y\right) \right] \simeq \mathbb {E} _ {\mathbf {x} _ {i} \sim P _ {T}} \left[ \ell \left(\phi_ {\theta^ {\Omega (A, B)}} (\mathbf {x} _ {i}), y\right) \right],
$$
$$
\text {s . t .} \theta^ {\mathcal {T}} = \underset {\theta^ {\mathcal {T}}} {\arg \min } \mathcal {L} ^ {\mathcal {T}} \left(\theta^ {\mathcal {T}}\right) = \underset {\theta^ {\mathcal {T}}} {\arg \min } \frac {1}{N _ {\mathcal {T}}} \sum_ {\left(\mathbf {x} _ {i}, y\right) \in \mathcal {T}} \ell \left(\phi_ {\theta^ {\mathcal {T}}} \left(\mathbf {x} _ {i}\right), y\right), \tag {3}
$$
$$
\theta^ {\Omega (\mathcal {A}, \mathcal {B})} = \underset {\theta^ {\Omega (\mathcal {A}, \mathcal {B})}} {\arg \min } \mathcal {L} ^ {\Omega (\mathcal {A}, \mathcal {B})} \left(\theta^ {\Omega (\mathcal {A}, \mathcal {B})}\right) = \underset {\theta^ {\Omega (\mathcal {A}, \mathcal {B})}} {\arg \min } \frac {1}{N _ {\mathcal {S}}} \sum_ {\left(\mathcal {A} _ {i} \mathcal {B} _ {i}, y\right) \in \Omega (\mathcal {A}, \mathcal {B})} \ell \left(\phi_ {\theta^ {\Omega (\mathcal {A}, \mathcal {B})}} \left(\mathcal {A} _ {i} \mathcal {B} _ {i}\right), y\right),
$$
where  $\Omega (\mathcal{A},\mathcal{B})$  is an operation that reconstructs  $\mathbf{x}_i$  channel-wise from low-rank matrices  $\mathcal{A}$  and  $\mathcal{B}$ , i.e.,  $\mathbf{x}_i = \mathcal{A}_i\mathcal{B}_i = [\mathcal{A}_{i,1}\mathcal{B}_{i,1}|\ldots |\mathcal{A}_{i,D}\mathcal{B}_{i,D}]\in \mathbb{R}^{D\times H\times W}$ , where  $[\cdot |\cdot ]$  means channel-wise stacked the image representation.