Dataset Condensation with Distribution Matching

Introduction

Computational cost for training a single state-of-the-art model in various fields, including computer vision and natural language processing, doubles every 3.4 months in the deep learning era due to larger models and datasets. The pace is significantly faster than the Moore’s Law that the hardware performance would roughly double every other year . While training a single model can be expensive, designing new deep learning models or applying them to new tasks certainly require substantially more computations, as they involve to train multiple models on the same dataset for many times to verify the design choices, such as loss functions, architectures and hyper-parameters . For instance, Ying et al. spend 100 TPU years of computation time conducting an exhaustive neural architecture search on CIFAR10 dataset , while training the best-performing architectures take only dozens of TPU minutes. Hence, there is a strong demand for techniques that can reduce the computational cost for training multiple models on the same dataset with minimal performance drop. To this end, this paper focuses on lowering the training cost by reducing the training set size.

The traditional solution to reduce the training set size is coreset selection. Typically, coreset selection methods choose samples that are important for training based on heuristic criteria, for example, minimizing distance between coreset and whole-dataset centers , maximizing the diversity of selected samples , discovering cluster centers , counting the mis-classification frequency and choosing samples with the largest negative implicit gradient . Although coreset selection methods can be very computationally efficient, they have two major limitations. First most methods incrementally and greedily select samples, which are short-sighted. Second their efficiency is upper bounded by the information in the selected samples in the original dataset.

An effective way of tackling the information bottleneck is synthesizing informative samples rather than selecting from given samples. A recent approach, dataset condensation (or distillation) aims to learn a small synthetic training set so that a model trained on it can obtain comparable testing accuracy to that trained on the original training set. Wang et al. pose the problem in a learning-to-learn framework by formulating the network parameters as a function of synthetic data and learning them through the network parameters to minimize the training loss over the original data. An important shortcoming of this method is the expensive optimization procedure that involves updating network weights for multiple steps for each outer iteration and unrolling its recursive computation graph. Zhao et al. propose to match the gradients w.r.t. the network weights giving real and synthetic training images that successfully avoids the expensive unrolling of the computational graph. Another efficiency improvement is a closed form optimizer by posing the classification task as a ridge regression problem to simplify the inner-loop model optimization . In spite of the recent progress, the dataset condensation still requires solving the expensive bi-level optimization which jeopardizes their goal of reducing training time due to the expensive image synthesis process. For instance, the state-of-the-art requires 15 hours of GPU time to learn 500 synthetic images on CIFAR10 which equals to the cost of training 6 deep networks on the same dataset. In addition, these methods also require tuning multiple hyper-parameters, e.g. the steps to update synthetic set and network parameters respectively in each iteration, that can be different for different settings such as sizes of synthetic sets.

In this paper, we propose a novel training set synthesis technique that combines the advantages of previous coreset and dataset condensation methods while avoiding their limitations. Unlike the former and like the latter, our method is not limited to individual samples from original dataset and can synthesize training images. Like the former and unlike the latter, our method can very efficiently produce a synthetic set and avoid expensive bi-level optimization. In particular, we pose this task as a distribution matching problem such that the synthetic data are optimized to match the original data distribution in a family of embedding spaces by using the maximum mean discrepancy (MMD) measurement (see Figure 1). Distance between data distributions is commonly used as the criterion for coreset selection , however, it has not been used to synthesize training data before. We show that the family of embedding spaces can be efficiently obtained by sampling randomly initialized deep neural networks. Hence, our method is significantly faster (e.g. 45$\times$ in CIFAR10 when synthesizing 500 images) than the state-of-the-art and involves tuning only one hyperparameter (learning rate for synthetic images), while obtaining comparable or better results. In addition, the training of our method can be independently run for each class in parallel and its computation load can be distributed. Finally, our method provides a different training cost/performance tradeoff for large-scale settings. In contrast to prior works that are limited to learning small synthetic sets on small datasets, our method can be successfully applied in more realistic settings, i.e. synthesizing 1250 images per class for CIFAR10 , and larger datasets, i.e. TinyImageNet and ImageNet-1K . We also validate these benefits in two downstream tasks by producing more data-efficient memory for continual learning and generating more representative proxy dataset for accelerating neural architecture search.

Methodology

The goal of dataset condensation is to condense the large-scale training set $\mathcal{T}=\{(\bm{x}_{1},y_{1}),\dots,(\bm{x}_{|\mathcal{T}|},y_{|\mathcal{T}|})\}$ with $|\mathcal{T}|$ image and label pairs into a small synthetic set $\mathcal{S}=\{(\bm{s}_{1},y_{1}),\dots,(\bm{s}_{|\mathcal{S}|},y_{|\mathcal{S}|})\}$ with $|\mathcal{S}|$ synthetic image and label pairs so that models trained on each $\mathcal{T}$ and $\mathcal{S}$ obtain comparable performance on unseen testing data:

Previous works formulate the dataset condensation as a learning-to-learn problem, pose the network parameters $\bm{\theta}^{\mathcal{S}}$ as a function of synthetic data $\mathcal{S}$ and obtain a solution for $\mathcal{S}$ by minimizing the training loss $\mathcal{L}^{\mathcal{T}}$ over the original data $\mathcal{T}$:

Recently the authors of show that a similar goal can be achieved by matching gradients of the losses over the synthetic and real training data respectively w.r.t. the network parameters $\bm{\theta}$, while optimizing $\bm{\theta}$ and the synthetic data $\mathcal{S}$ in an alternating way:

where $P_{\bm{\theta}_{0}}$ is the distribution of parameter initialization, $T$ is the outer-loop iteration for updating synthetic data, $\varsigma_{\bm{\theta}}$ is the inner-loop iteration for updating network parameters, $\eta_{\bm{\theta}}$ is the parameter learning rate and $D(\cdot,\cdot)$ measures the gradient matching error. Note that all the training algorithms have another loop of sampling $\bm{\theta}_{0}$ over the bi-level optimization.

Dilemma.

The learning problems in eq. 2 and eq. 3 involve solving an expensive bi-level optimization: first optimizing the model $\bm{\theta}^{\mathcal{S}}$ in eq. 2 or $\bm{\theta}_{t}$ in eq. 3 at the inner loop, then optimizing the synthetic data $\mathcal{S}$ along with additional second-order derivative computation at the outer loop. For example, training 50 images/class synthetic set $\mathcal{S}$ by using the method in requires 500K epochs of updating network parameters $\bm{\theta}_{t}$ on $\mathcal{S}$, in addition to the 50K updating of $\mathcal{S}$. Furthermore, Zhao et al. need to tune the hyper-parameters of the outer and inner loop optimization (i.e. how many steps to update $\mathcal{S}$ and $\bm{\theta}_{t}$) for different learning settings, which requires cross-validating them and hence multiplies the cost for training synthetic images.

2 Dataset Condensation with Distribution Matching

Our goal is to synthesize data that can accurately approximate the distribution of the real training data in a similar spirit to coreset techniques (e.g. ). However, to this end, we do not limit our method to select a subset of the training samples but to synthesize them as in . As the training images are typically very high dimensional, estimating the real data distribution $P_{\mathcal{D}}$ can be expensive and inaccurate. Instead, we assume that each training image $\bm{x}\in\Re^{d}$ can be embedded into a lower dimensional space by using a family of parametric functions $\psi_{\bm{\vartheta}}:\Re^{d}\rightarrow\Re^{d^{\prime}}$ where $d^{\prime}\ll d$ and $\bm{\vartheta}$ is the parameter. In other words, each embedding function $\psi$ can be seen as providing a partial interpretation of its input, while their combination provides a complete one.

Now we can estimate the distance between the real and synthetic data distribution with commonly used maximum mean discrepancy (MMD) :

where $\mathcal{H}$ is reproducing kernel Hilbert space. As we do not have access to ground-truth data distributions, we use the empirical estimate of the MMD:

where $P_{\bm{\vartheta}}$ is the distribution of network parameters.

Following , we also apply the differentiable Siamese augmentation $\mathcal{A}(\cdot,\omega)$ to real and synthetic data that implements the same randomly sampled augmentation to the real and synthetic minibatch in training, where $\omega\sim\Omega$ is the augmentation parameter such as the rotation degree. Thus, the learned synthetic data can benefit from semantic-preserving transformations (e.g. cropping) and learn prior knowledge about spatial configuration of samples while training deep neural networks with data augmentation. Finally, we solve the following optimization problem:

We learn the synthetic data $\mathcal{S}$ by minimizing the discrepancy between two distributions in various embedding spaces by sampling $\bm{\vartheta}$. Importantly eq. 6 can be efficiently solved, as it requires only optimizing $\mathcal{S}$ but no model parameters and thus avoids expensive bi-level optimization. This is in contrast to the existing formulations (see eq. 2 and eq. 3) that involve bi-level optimizations over network parameters $\bm{\theta}$ and the synthetic data $\mathcal{S}$.

Note that, as we target at image classification problems, we minimize the discrepancy between the real and synthetic samples of the same class only. We assume that each real training sample is labelled and we also set a label to each synthetic sample and keep it fixed during training.

3 Training Algorithm

We depict the mini-batch based training algorithm in Algorithm 1. We train the synthetic data for $K$ iterations. In each iteration, we randomly sample the model $\psi_{\bm{\vartheta}}$ with parameter $\bm{\vartheta}\sim P_{\bm{\vartheta}}$. Then, we sample a pair of real and synthetic data batches ($B^{\mathcal{T}}_{c}\sim\mathcal{T}$ and $B^{\mathcal{S}}_{c}\sim\mathcal{S}$) and augmentation parameter $\omega_{c}\sim\Omega$ for every class $c$. The mean discrepancy between the augmented real and synthetic batches of every class is calculated and then summed as loss $\mathcal{L}$. The synthetic data $\mathcal{S}$ is updated by minimizing $\mathcal{L}$ with stochastic gradient descent and learning rate $\eta$.

4 Discussion

The family of embedding functions $\psi_{\bm{\vartheta}}$ can be designed in different ways. Here we use a deep neural network with different random initializations rather than sampling its parameters from a set of pre-trained networks which is more computationally expensive to obtain. We experimentally validate that our random initialization strategy produces better or comparable results with the more expensive strategy of using pretrained networks in Section 3.4. However, one may still question why randomly initialized networks provide meaningful embeddings for distribution matching. Here we list two reasons based on the observations from previous work. First, randomly initialized networks are reported to produce powerful representations for multiple computer vision tasks . Second, such random networks are showed to perform a distance-preserving embedding of the data, i.e. smaller distances between samples of same class and larger distances across samples of different classes . In addition, the combination of many weak embeddings provides a complete interpretation of the inputs.

Connection to Gradient Matching.

While we match the mean features of the real and synthetic image batches, Zhao et al. match the mean gradients of network weights over the two batches. We find that, given a batch of data from the same class, the mean gradient vector w.r.t. each output neuron in the last layer of a network is equivalent to a weighted mean of features where the weights are a function of classification probabilities predicted by the network and proportional to the distance between prediction and ground-truth. In other words, while our method weighs each feature equally, Zhao et al. assign larger weights to samples whose predictions are inaccurate. Note that these weights dynamically vary for different networks and training iterations. We provide the derivation in the appendix.

Generative Models.

The classic image synthesizing techniques, includes AutoEncoders and Generative Adversarial Networks (GANs) , aim to synthesize real-looking images, while our goal is to generate data-efficient training samples. Regularizing the images to look real may limit the data-efficiency. Previous work showed that the images synthesized by cGAN are not better than the randomly selected real images for training networks. We further provide the comparison to state-of-the-art VAE and GAN models and GMMN method in the appendix. Although generative models can be trained to produce data-efficient samples with suitable objectives, e.g. and ours, it is not trivial work to build it and achieve state-of-the-art results . We leave it as the future work.

Experiments

We evaluate the classification performance of deep networks that are trained on the synthetic images generated by our method. We conduct experiments on five datasets including MNIST , CIFAR10, CIFAR100 , TinyImageNet and ImageNet-1K . MNIST consists of 60K $28\times 28$ gray-scale training images of 10 classes. CIFAR10 and CIFAR100 contain 50k $32\times 32$ training images from 10 and 100 object categories respectively. TinyImageNet and ImageNet-1K have 100K training images from 200 categories and 1.3M training images from 1K categories respectively. We resize these ImageNet images with $64\times 64$ resolution. These two datasets are significantly more challenging than MNIST and CIFAR10/100 due to more diverse classes and higher image resolution.

Experimental Settings.

We first learn 1/10/50 image(s) per class synthetic sets for all datasets by using the same ConvNet architecture in . Then, we use the learned synthetic sets to train randomly initialized ConvNets from scratch and evaluate them on real test data. The default ConvNet includes three repeated convolutional blocks, and each block involves a 128-kernel convolution layer, instance normalization layer , ReLU activation function and average pooling. Note that four-block ConvNets are used to adjust to the larger input size ($64\times 64$) of TinyImageNet and ImageNet-1K images. In each experiment, we learn one synthetic set and use it to test 20 randomly initialized networks. We repeat each experiment for 5 times and report the mean testing accuracy of the 100 trained networks. We also do cross-architecture experiments in Section 3.3. where we learn the synthetic set on one network architecture and use them to train networks with different architectures.

Hyper-parameters.

Like the standard neural network training, dataset condensation also involves tuning a set of hyperparameters. Our method needs to tune only one hyper-parameter, i.e. learning rate for the synthetic images, for learning different sizes of synthetic sets, while existing methods have to tune more hyper-parameters such as the steps to update synthetic images and network parameters respectively. We use a fixed learning rate 1 for optimizing synthetic images for all 1/10/50 images/class learning on all datasets. When learning larger synthetic sets such as 100/200/500/1,000 images per class, we use larger learning rate (i.e. 10) due to the relatively smaller distribution matching loss. We train synthetic images for 20,000 iterations on MNIST, CIFAR10/100 and 10,000 iterations on TinyImageNet and ImageNet-1K respectively. The mini-batch size for sampling real data is 256. We initialize the synthetic images using randomly sampled real images with corresponding labels. All synthetic images of a class are used to compute the class mean. We use the same augmentation strategy as .

2 Comparison to the State-of-the-art

We compare our method to three standard coreset selection methods, namely, Random Selection, Herding and Forgetting . Herding method greedily adds samples into the coreset so that the mean vector is approaching the whole dataset mean. Toneva et al. count how many times a training sample is learned and then forgotten during network training. The samples that are less forgetful can be dropped. We also compare to four state-of-the-art training set synthesis methods, namely, DD , LD , DC and DSA . Note that we are aware of concurrent works that largely improves the existing bilevel optimization based dataset condensation solutions. Unlike them, we contribute the first solution that has neither bi-level optimization nor second-order derivative, and provide a different training cost/performance tradeoff. Compared to them, our method is significantly simpler and faster. Thus, it is able to scale to large settings i.e. learning 1250 images per class for CIFAR10 and large datasets i.e. ImageNet-1K, while these concurrent works can’t. More detailed comparison and discussion to other methods , MMD baseline and generative baselines including DC-VAE , BigGAN and GMMN can be found in the appendix.

Performance Comparison.

Here we evaluate our method on MNIST, CIFAR10 and CIFAR100 datasets and report the results in Table 1. Among the coreset selection methods, Herding performances the best in most settings. Especially, when small synthetic sets are learned, Herding method performs significantly better. For example, Herding achieves 8.4% testing accuracy when learning 1 image/class synthetic set on CIFAR100, while Random and Forgetting obtains only 4.2% and 4.5% testing accuracies respectively.

Training set synthesis methods have clear superiority over coreset selection methods, as the synthetic training data are not limited to a set of real images. Best results are obtained either by DSA or our method. While DSA produces more data-efficient samples with a small number of synthetic samples (1/10 image(s) per class), our method outperforms DSA at 50 images/class setting in CIFAR10 and CIFAR100. The possible reason is that the inner-loop model optimization in DSA with limited number of steps is more effective to fit the network parameters on smaller synthetic data (see eq. 3). In case of bigger learned synthetic data, the solution obtained in the inner-loop becomes less accurate as it can use only limited number of steps to keep the algorithm scalable. In contrast, our method is robust to increasing synthetic data size, can be efficiently optimized significantly faster than DSA.

TinyImageNet and ImageNet-1K.

Due to higher image resolution and more diverse classes, prior bilevel optimization based methods do not scale to TinyImageNet and ImageNet-1K. Our method takes 27 hours with one Tesla V100 GPU to condense TinyImageNet into three condensed sets (1/10/50 images/class synthetic sets), and it takes 28 hours with ten GTX 1080 GPUs to condense ImageNet-1K into these three sets. As shown in Table 1, our method achieves 3.9%, 12.9% and 24.1% testing accuracies when learning 1, 10 and 50 images/class synthetic sets for TinyImageNet, and recovers 60% classification performance of the baseline that is trained on the whole original training set with only 10% of data. Our method significantly outperforms the best coreset selection method - Herding, which obtains 2.8%, 6.3% and 16.7% testing accuracies. On ImageNet-1K dataset, our method achieves 1.3%, 5.7% and 11.4% testing accuracies when learning 1, 10 and 50 images/class synthetic sets, which outperforms random selection (0.52%, 1.94% and 7.54%) by large margins.

Visualization.

The learned synthetic images of MNIST and CIFAR10 are visualized in Figure 2. We find that the synthetic MNIST images are clear and noise free, while the number images synthesized by previous methods contain obvious noise and some unnatural strokes. The synthetic images of CIFAR10 dataset are also visually recognizable and diverse. It is easy to distinguish the background and foreground object.

Figure 3 depicts the feature distribution of the (50 images/class) synthetic sets learned by DC, DSA and our method (DM). We use a network trained on the whole training set to extract features and visualize the features with T-SNE . We find that the synthetic images learned by DC and DSA cannot cover the real image distribution. In contrast, our synthetic images successfully cover the real image distribution. Furthermore, fewer outlier synthetic samples are produced by our method.

Learning with Batch Normalization.

Zhao et al. showed that instance normalization works better than batch normalization (BN) when learning small synthetic sets because the synthetic data number is too small to calculate stable running mean and standard deviation (std). When learning with batch normalization, they first pre-set the BN mean and std using many real training data and then freeze them for synthetic data. Thus, the inaccurate mean and std will make optimization difficult . In contrast, we estimate running mean and std by inputting augmented synthetic data from all classes. Hence, our method benefits from the true mean and std of synthetic data. Table 2 show that using ConvNet with BN can further improve our performance. Specifically, our method with BN achieves 65.2%, 48.0% and 28.2% testing accuracies when learning 50 images/class synthetic sets on CIFAR10, CIFAR100 and TinyImageNet respectively, which means $2.2\%$, $4.4\%$ and 4.1% improvements over our method with the default instance normalization, and also outperforms DSA with BN by 5.3% and 3.4% on CIFAR10 and CIFAR100 respectively.

Training Cost Comparison.

Our method is significantly more efficient than those bi-level optimization based methods. Without loss of generality, we compare the training time of ours and DSA in the setting of learning 50 images/class synthetic data on CIFAR10. Figure 4 shows that our method needs less than 20 minutes to reach the performance of DSA trained for 15 hours, which means less than $2.2\%$ training cost. Note that we run the two methods in the same computation environment with one GTX 1080 GPU.

Learning Larger Synthetic Sets

We show that our method can also be used to learn larger synthetic sets, while the bi-level optimization based methods typically requires more training time and elaborate hyper-parameter tuning for larger settings. Figure 7 compares our method to random selection baseline in CIFAR10 in terms of absolute and relative performance w.r.t. whole dataset training performance. Clearly our method outperforms random baseline at all operating points which means that our synthetic set is more data-efficient. The advantage of our method is remarkable in challenging settings, i.e. settings with small data budgets. Our method obtains $67.0\pm 0.3\%$, $71.2\pm 0.4\%$, $76.1\pm 0.3\%$, $79.8\pm 0.3\%$ and $80.8\pm 0.3\%$ testing accuracies when learning 100, 200, 500, 1000 and 1250 images/class synthetic sets on CIFAR10 dataset respectively, which means we can recover $79\%$, $84\%$, $90\%$, $94\%$ and $95\%$ relative performance using only $2\%$, $4\%$, $10\%$, $20\%$ and $25\%$ training data compared to whole dataset training. We see that the performance gap between the two methods narrows when we learn larger synthetic set. This is somewhat expected, as randomly selecting more samples will approach the whole dataset training which can be considered as the upper-bound. As we initialize synthetic images from random real images, the initial distribution discrepancy becomes tiny when the synthetic set is large.

3 Cross-architecture Generalization

verified the cross-architecture generalization ability of synthetic data in an easy setting - learning 1 image/class for MNIST dataset. In this paper, we implement a more challenging cross-architecture experiment - learning 50 images/class for CIFAR10 dataset. In Table 3, the synthetic data are learned with one architecture (denoted as C) and then be evaluated on another architecture (denoted as T) by training a model from scratch and testing on real testing data. We test several sophisticated neural architectures namely ConvNet, AlexNet , VGG-11 and ResNet-18 . Batch Normalization is used in all architectures.

Table 3 shows that learning and evaluating synthetic set on ConvNet achieves the best performance 65.2%. Comparing with DSA, the synthetic data learned by our method with ConvNet have better generalization performance than that learned by DSA with the ConvNet. Specifically, our method outperforms DSA by 8.0%, 8.9% and 9.7% when testing with AlexNet, VGG and ResNet respectively. These results indicate that the synthetic images learned with distribution matching have better generalization performance on unseen architectures than those learned with gradient matching. The learning of synthetic set can be worse with more sophisticated architecture such as ResNet. It is reasonable that the synthetic data fitted on sophisticated architecture will contain some bias that doesn’t exist in other architectures, therefore cause worse cross-architecture generalization performance. We also find that the evaluation of the same synthetic set on more sophisticated architectures will be worse. The reason may be that sophisticated architectures are under-fitted using small synthetic set.

4 Ablation Study on Network Distribution

Here we study the effect of using different network distributions while learning 1/10/50 image(s)/class synthetic sets on CIFAR10 with ConvNet architecture. Besides sampling randomly initialized network parameters, we also construct a set of networks that are pre-trained on the original training set. In particular, we train 1,000 ConvNets with different random initializations on the whole original training set and also store their intermediate states. We roughly divide these networks into nine groups according to their validation accuracies, sample networks from each group, learn the synthetic data on them and use learned synthetic data to train randomly initialized neural networks. Interestingly we see in Table 4 that our method works well with all nine network distributions and the performance variance is small. The visualization and analysis about synthetic images learned with different network distributions are provided in the appendix.

5 Continual Learning

We also use our method to store more efficient training samples in the memory for relieving the catastrophic forgetting problem in continual (incremental) learning . We set up the baseline based on GDumb which stores training samples in memory greedily and keeps class-balance. The model is trained from scratch on the latest memory only. Hence, the continual learning performance completely depends on the quality of the memory construction. We compare our memory construction method i.e. training set condensation to the random selection that is used in , herding and DSA . We implement class-incremental learning on CIFAR100 dataset with an increasing memory budget of 20 images/class. We implement 5 and 10 step learning, in which we randomly and evenly split the 100 classes into 5 and 10 learning steps i.e. 20 and 10 classes per step respectively. The default ConvNet is used in this experiment.

As depicted in Figure 7 and Figure 7, we find that our method GDumb + DM outperforms others in both two settings, which means that our method can produce the best condensed set as the memory. The final performances of ours, DSA, herding and random are 34.4%, 31.7%, 28.2% and 24.8% in 5-step learning and 34.6%, 30.5%, 27.4% and 24.8% in 10-step learning. We find that ours and random selection performances are not influenced by how the classes are split namely how many new training classes and images occur in each learning step, because both two methods learn/generate the sets independently for each class. However, DSA and herding methods perform worse when the training classes are densely split into more learning steps. The reason is that DSA and herding needs to learn/generate sets based on the model(s) trained on the current training data, which is influenced by the data split. More details can be found in the appendix.

6 Neural Architecture Search

The synthetic sets can also be used as a proxy set to accelerate model evaluation in Neural Architecture Search (NAS) . Following , we implement NAS on CIFAR10 with the search space of 720 ConvNets varying in network depth, width, activation, normalization and pooling layers. Please refer to for more details. We train all architectures on the learned 50 images/class synthetic set, i.e. 1% size of the whole dataset, from scratch and then rank them based on the accuracy on a small validation set. We compare to random, DSA and early-stopping methods. The same size of real images are selected as the proxy set in random. DSA means that we use the synthetic set learned by DSA in the same setting. In early-stopping, we use the whole training set to train the model but with the same training iterations like training on the proxy datasets. Therefore, all these methods have the same training time. We train models on the proxy sets for 200 epochs and whole dataset for 100 epochs. The best model is selected based on validation accuracies obtained by different methods. The Spearman’s rank correlation between performances of proxy-set and whole-dataset training is computed for the top 5% architectures selected by the proxy-set.

The NAS results are provided in Table 5. Although the architecture selected by early-stopping achieves the best performance (84.3%), its performance rank correlation (0.11) is remarkably lower than DSA (0.68) and DM (0.76). In addition, early-stopping needs to use the whole training set, while other proxy-set methods need only 500 training samples. The performance rank correlation of Random (-0.04) is too low to provide a reliable ranking for the architectures. Our method (DM) achieves the highest performance rank correlation (0.76), which means that our method can produce reliable ranking for those candidate architectures while using only around $\frac{1}{25}$ training time of whole dataset training. Although our method needs 72 min to obtain the condensed set, it is negligible compared to whole-dataset training (3580.2 min). More implementation details and analysis can be found in the appendix.

Conclusion

In this paper, we propose an efficient dataset condensation method based on distribution matching. To our knowledge, it is the first solution that has neither bi-level optimization nor second-order derivative. Thus, the synthetic data of different classes can be learned independently and in parallel. Thanks to its efficiency, we can apply our method to more challenging datasets - TinyImageNet and ImageNet-1K, and learn larger synthetic sets - 1250 images/class on CIFAR10. Our method is ${45}$ times faster than the state-of-the-art for learning 50 images/class synthetic set on CIFAR10. We also empirically prove that our method can produce more informative memory for continual learning and better proxy set for speeding up model evaluation in NAS. Though remarkable progress has been seen in this area since the pioneering work released in 2018, dataset condensation is still in its early stage. We will extend dataset condensation to more complex vision tasks in the future.

This work is funded by China Scholarship Council 201806010331 and the EPSRC programme grant Visual AI EP/T028572/1.

References

Appendix A Implementation details

As didn’t report 50 images/class learning performance on CIFAR100, we obtain the result in Table 1 by running their released code and coarsely searching the hyper-parameters (outer and inner loop steps). Then, we set both outer and inner loop to be 10 steps. The rest hyper-parameters are the default ones in their released code. To obtain the DSA results with batch normalization in Table 2 and Table 3, we also run DSA code and set batch normalization in ConvNet.

ResNet with Batch Normalization.

We follow the modification of ResNet in . They replace the $stride=2$ convolution layer with $stride=1$ convolution layer followed by an average pooling layer in the ResNet architecture that is used to learn the synthetic data. This modification enables smooth error back-propagation to the input images. We directly use their released ResNet architecture.

A.2 Continual Learning

Prabhu et al. use cutmix augmentation strategy for training models. Different from them, we follow and use the default DSA augmentation strategy in order to be consistent with other experiments in this paper.

DSA and Herding Training.

Without loss of generality, we run DSA training algorithm on the new training classes and images only in every learning step. It is not a easy work to take old model and memory into DSA training and achieve better performance. The synthetic data learned with old model can also be biased to it, and thus perform worse. Similarly, we train the embedding function (ConvNet) for herding method on the new training classes and images only.

A.3 Neural Architecture Search

We randomly select 10% training samples in CIFAR10 dataset as the validation set. The rest are the training set. The batch size is 250, then one training epoch on the small (50 images/class) proxy sets includes 2 batches. The DSA augmentation strategy is applied to all proxy-set methods and early-stopping. We train each model 5 times and report the mean accuracies. We do NAS experiment on one Tesla v100 GPU.

We visualize the performance rank correlation between proxy-set and whole-dataset training in Figure F8. The top 5% architectures are selected based on the validation accuracies of models trained on each proxy-set. Each point represents a selected architecture. The horizontal and vertical axes are the testing accuracies of models trained on the proxy-set and the whole dataset respectively. The figure shows that our method can produce better proxy set to obtain more reliable performance ranking of candidate architectures.

Appendix B Comparison to More Baselines and Related Works

In this subsection, we compare the data-efficiency of samples generated by our dataset condensation method to those generated by traditional generative models, namely VAE and GAN. Specifically, we choose the state-of-the-art DC-VAE and BigGAN . The BigGAN model is trained with the differentiable augmentation . In addition, we also compare to a related generative model GMMN which aims to learn an image generator that can map a uniform distribution to real image distribution. Our method differs from GMMN in many ways significantly. First, GMMN aims to generate real-looking images, while our goal is to condense a training set by synthesizing informative training samples that can be used to efficiently train deep networks through MMD. Second, our method learns pixels directly, while GMMN learns a generator network. Third, our method learns a few synthetic samples to approximate the distribution of large real training set in any embedding space with any augmentation, while GMMN learns to map a uniform distribution to real image distribution which is an easier task.

We train these generative models on CIFAR10 dataset. ConvNets are trained on these synthetic images and then evaluated on real testing images. The results in Table T6 verify that our method outperforms them by large margins, indicating that our synthetic images are more informative for training deep neural networks. The comparison to random baseline indicates that the images generated by traditional generative models are not more informative than randomly selected real images.

B.2 Comparison to MMD Baseline

Another baseline is to learn synthetic images by distribution matching with vanilla MMD in the pixel space. This baseline can also been considered as the ablation study of the embedding function and differentiable augmentation in our method. We try this baseline with linear, polynomial, RBF and Laplacian kernels and with various kernel hyper-parameters. We find that only MMD with linear kernel can achieve better synthetic images, i.e. better than randomly selected real images. The performance of MMD with linear kernel in the pixel space is presented in Table T6, which outperforms all generative models while is inferior to our method. This result also verifies that the distribution matching mechanism enables learning more informative synthetic samples.

B.3 Comparison to GTN and KIP Methods

We notice the recent works Generative Teaching Networks (GTN) and Kernel Inducing Point (KIP) on dataset condensation. Such et al. propose to learn a generative network that outputs condensed training samples by minimizing the meta-loss on real data. They report the performance of 4,096 synthetic images learned on MNIST which is comparable to our 50 images/class synthetic set (i.e. 500 images in total) performance.

Nguyen et al. propose to replace the neural network optimization in the bi-level optimization with kernel ridge regression which has a closed-form solution. Zero Component Analysis (ZCA) is applied for pre-processing images. Although Nguyen et al. report the results on 1024-width neural networks while we train and test 128-width neural networks, our results still outperform theirs in some settings, for example $98.6\pm 0.1\%$ v.s. $98.3\pm 0.1\%$ when learning 50 images/class on MNIST and $29.7\pm 0.3\%$ v.s. $28.3\pm 0.1\%$ when learning 10 images/class on CIFAR100. Note that they achieve those results by leveraging distributed computation environment and training for thousands of GPU hours. In contrast, our method can learn synthetic sets with one GTX 1080 GPU in dozens of minutes, which is significantly more efficient.

Appendix C Extended Visualization and Analysis

We visualize the 10 images/class synthetic sets learned on CIFAR10 dataset with different network parameter distributions in Figure F9. It is interesting that images learned with “poor” networks that have lower validation accuracies look blur. We can find obvious checkerboard patterns in them. In contrast, images learned with “good” networks that have higher validation accuracies look colorful. Some twisty patterns can be found in these images. Although synthetic images learned with different network parameter distributions look quite different, they have similar generalization performance. We think that these images are mainly different in terms of their background patterns but similar in semantics. It means that our method can produce synthetic images with similar network optimization effects while significantly different visual effects. Our method may have promising applications in protecting data privacy and federated learning .

Appendix D Connection to Gradient Matching

In this section, we show the connection between gradient matching and our method. Both and our training algorithm sample real and synthetic image batches from one class in each iteration, which is denoted as class $y$. We embed each training sample $(\bm{x}_{i},y)$ and obtain the feature $\bm{e}_{i}$ using a neural network $\psi_{\bm{\vartheta}}$ followed a linear classifier $\mathbf{W}=[\bm{w}_{0},...,\bm{w}_{C-1}]$, where $\bm{w}_{j}$ is the weight vector connected to the $j^{\text{th}}$ output neuron and $C$ is the number of all classes. Note that the weight and its gradient vector are organized in the same way in . We focus on the weight and gradient of the linear classification layer (i.e. the last layer) of a network in this paper. The classification loss $J_{i}$ of each sample is denoted as

Then, we compute the partial derivative w.r.t. each weight vector,

This equation can be simplified using the predicted probability $p_{i,j}=\frac{\exp{(\bm{w}_{j}^{T}\cdot\bm{e}_{i})}}{\Sigma_{k}\exp{(\bm{w}_{k}^{T}\cdot\bm{e}_{i})}}$ that classifies sample $\bm{x}_{i}$ into category $j$:

Eq. 9 shows that the last-layer gradient vector $\bm{g}_{i,j}$ is equivalent to a weighted feature vector $\bm{e}_{i}$ and vice versa. The weight is a function of classification probability. Generally speaking, the weight is large when the difference between predicted probability $p_{i,j}$ and ground-truth one-hot label ($1$ or ) is large.

As the real and synthetic samples in each training iteration are from the same class $y$, we can obtain the mean gradient over a data batch by averaging the corresponding gradient components:

$N$ is the batch size. Thus, last-layer mean gradient is equivalent to the weighted mean feature, and the mean gradient matching is equivalent to the matching of weighted mean feature.

Our method can learn synthetic images with randomly initialized networks. Given networks with random parameters, we assume that the predicted probability is uniform over all categories, i.e. $p_{i,j}=\frac{1}{C}$. Then, the mean gradient is

which is equivalent to the mean feature with a constant weight. Thus, with randomly initialized networks, the last-layer mean gradient matching is equivalent to mean feature matching multiplied by a constant.