Dataset Distillation using Neural Feature Regression
Abstract
Dataset distillation aims to learn a small synthetic dataset that preserves most of the information from the original dataset.
Dataset distillation can be formulated as a bi-level meta-learning problem where the outer loop optimizes the metadataset and the inner loop trains a model on the distilled data.
Meta-gradient computation is one of the key challenges in this formulation, as differentiating through the inner loop learning procedure introduces significant computation and memory costs.
In this paper, we address these challenges using neural Feature
Regression with Pooling (FRePo), achieving the state-of-the-art performance with an order of magnitude less memory requirement and two orders of magnitude faster training than previous methods.
The proposed algorithm is analogous to truncated backpropagation through time with a pool of models to alleviate various types of overfitting in dataset distillation.
FRePo significantly outperforms the previous methods on CIFAR100, Tiny ImageNet, and
ImageNet-1K. Furthermore, we show that high-quality distilled data can greatly improve various downstream applications, such as continual learning and membership inference defense. Please check out our webpage at https://sites.google.com/view/frepo .
Introduction
Knowledge distillation [1] is a technique in deep learning to compress knowledge for easy deployment.
Most previous works focus on model distillation [2, 3] where the knowledge acquired by a large teacher model is transferred to a small student model.
In contrast, dataset distillation [4, 5] aims to learn a small set of synthetic examples preserving most of the information from a large dataset such that a model trained on it can achieve similar test performance as one trained on the original dataset.
Distilled data can accelerate model training and reduce the cost of storing and sharing a dataset.
Moreover, its highly condensed and synthetic nature can also benefit various applications, such as continual learning [5–8], neural architecture search [5, 7], and privacy-preserving tasks [9, 10].
Dataset distillation was first studied by Maclaurin et al.
[11] in the context of gradient-based hyperparameter optimization and subsequently Wang et al.
[4] formally proposed dataset distillation as a new task.
Dataset distillation can be naturally formulated as a bi-level meta-learning problem.
The inner loop optimizes the model parameters on the distilled data (meta-parameters), while the outer loop refines the distilled data with meta-gradient updates.
One key challenge in dataset distillation is computing the meta-gradient.
Several methods [4, 11–13] compute it by back-propagating through the unrolled computation graph, but they often suffer from cup bottle can plate clock lamp couch bed chair wardrobe telephone television bicycle motorcycle pickup_truck bus tank rocket mountain cloud sea plain road forest orange pear apple maple_tree oak_tree palm_tree bee beetle crab lion tiger wolf (a) CIFAR espresso pizza basketball rugby ball maypole ladybug fly monarch koala
African elephant go-kart sports car school bus bullet train lifeboat obelisk flagpole water tower triumphal arch barn magnetic compass teapot wok backpack rocking chair (b) Tiny ImageNet tench english springer cassette player church french horn garbage truck golf ball parachute english foxhound golden retriever rhodesian ridgeback old English sheepdog australian terrier dingo
Samoyed shih-tzu (c) ImageNette and ImageWoof huge compute and memory requirement [14], training instability [15, 16], and truncation bias [17]. To avoid unrolled optimization, surrogate objectives are used to derive the meta-gradient, such as gradient matching [5, 7, 18], feature alignment [8, 19], and training trajectory matching [20].
Nevertheless, a surrogate objective may introduce its own bias [19], and thus, may not accurately reflect the true objective.
An alternative is using kernel methods, such as Neural Tangent Kernel (NTK) [21], to approximate the inner optimization [22, 23].
However, computing analytical NTK for modern neural network can be extremely expensive [22, 23].
Even with an accurate meta-gradient, dataset distillation still suffers from various types of overfitting.
For instance, the distilled data can easily overfit to a particular learning algorithm [4, 13, 20], a certain stage of optimization [13, 19], or a certain network architecture [5, 7, 20, 22, 23].
Meanwhile, the model can also overfit the distilled data during training, which is the most common cause of overfitting when we train on a small dataset.
All these kinds of overfitting impose difficulties on the training and general-purpose use of the distilled data.
We propose an efficient meta-gradient computation method and a “model pool” to address the overfitting problems.
The bottleneck in meta-gradient computation arises due to the complexity of inner optimization, as we need to know how the inner parameters vary with the outer parameters [24].
However, the inner optimization can be pretty simple if we only train the last layer of a neural network to convergence while keeping the feature extractor fixed.
In this case, computing the prediction on the real data using the model trained on the distilled data can be expressed as a kernel ridge regression (KRR) with respect to the conjugate kernel [25].
Hence, computing the meta-gradient is simply back-propagating through the kernel and a fixed feature extractor.
To alleviate overfitting, we propose to maintain a diverse pool of models instead of periodically training and resetting a single model as in prior work [7, 13, 18].
Intuitively, our algorithm targets the following question: what is the best data to train the linear classifier given the current feature extractor?
Due to the diverse feature extractors we use, the distilled data generalize well to a wide range of model distributions.
Summary of Contributions:
• We propose an effective method for dataset distillation.
Our method, named neural Feature Regression with Pooling (FRePo), achieves state-of-the-art results on various benchmark datasets with a 100x reduction in training time and a 10x reduction in GPU memory requirement.
Our distilled data looks real (Figure 1) and transfers well to different architectures.
• We show that FRePo scales well to datasets with high-resolution images or complex label space.
We achieve 7.5% top1 accuracy on ImageNet-1K [26] using only one image per class.
The same classifier obtains only 1.1% accuracy from a random subset of real images.
The previous methods struggle in this task due to large memory and compute requirements.
• We demonstrate that high-quality distilled data can significantly improve various downstream applications, such as continual learning and membership inference defense.
[Figure]
Method 2.1 Dataset Distillation as Bi-level Optimization
Suppose we have a large labeled dataset T “ ␣ p x 1 , y 1 q , . . .
, ` x | T | , y | T | ˘( with | T | image and label pairs.
Dataset distillation aims to learn a small synthetic dataset S “ ␣ p x 1 , y 1 q , . . .
, ` x | S | , y | S | ˘( that preserves most of the information in T .
We train several neural networks parameterized by θ on the dataset S and then compute the validation loss L p
A lg p θ, S q , T q on the real dataset T , where A lg p θ, S q is the neural network parameters optimized by a learning algorithm
A lg with the model initialization θ and distilled dataset S as its inputs.
The validation loss L p
A lg p θ, S q , T q is a noisy objective with the stochasticity coming from random model initialization and inner learning algorithm.
Thus, we are interested in minimizing the expected value of this loss, which we denote it as F p S q .
We formulate the dataset distillation as the following bi-level optimization problem. outer-level hkkkkkkkkkkkikkkkkkkkkkkj S [Equation] “ argmin S
F p S q , where F p S q
[Equation] E θ [Equation]
P θ
[Equation] inner-level hkkkkikkkkj [Equation] In this bi-level setup, the outer loop optimizes the distilled data to minimize F p S q , while the inner loop trains a neural network using the learning algorithm, A lg , to minimize the training loss on the distilled data S .
From the meta-learning perspective, the task is defined by the model initialization θ , and we want to learn a meta-parameter S that generalizes well to different models sampled from the model distributions P θ .
During learning, we optimize the meta-parameter S by minimizing the meta-training loss F p S q .
In contrast, at meta-test time, we train a new model from scratch on S and evaluate the trained model on a held-out real dataset.
This meta-test performance reflects the quality of the distilled data.
2.2 Dataset Distillation using Neural Feature Regression with Pooling (FRePo)
The outer-level problem can be solved using gradient-based methods of the form S Ð S
´ α ∇ S F p S q , where α is the learning rate for the distilled data and ∇ S F p S q is the meta-gradient [27].
For a particular model θ , the meta-gradient can be expressed as ∇ S L p
A lg p θ, S q , T q .
Computing this meta-gradient requires differentiating through inner optimization.
If A lg is an iterative algorithm like gradient descent, then backpropagating through the unrolled computation graph [14] can be a solution.
However, this type of unrolled optimization introduces significant computation and memory overhead, as the whole training trajectory needs to be stored in memory (Figure 2(b)).
Traditionally, these issues are alleviated with truncated backpropagation through time (TBPTT) [28–30].
Instead of backpropagating through an entire unrolled sequence, TBPTT performs backpropagation for each subsequence separately.
It is efficient because its time and memory complexity scale linearly with respect to the truncation steps.
However, truncation may yield highly biased gradients that severely impact training.
To mitigate this truncation bias [14], we consider training only the top layer of a network to convergence .
The key insight is that the data helpful for training the output layer can also help train the whole network.
Thus, we decompose the neural network into a feature
Require:
T : a labeled dataset; α the learning rate the distilled data
Initialization Initialize
S Initialization Initialize 1: while not converged do 2: Ż Sample θ
M .
3: Ż Sample batch uniformly dataset p X t Y t q
4: Ż Compute the meta training loss L Eq
5: Ż Update the data S :
X
X α ∇ X
Y s Ð Y s ´
α ∇ Y s L 6: Ż Train the θ i on the current distilled data S for one step. 7: Ż Reinitialize the θ P θ has updated more than K steps. 8: end while
Output: Learned distilled dataset S “ p X s , Y s q extractor and a linear classifier. We fix the feature extractor at each meta-gradient computation and train the linear classifier to convergence before updating S . After that, we adjust the feature extractor by training the whole network on the updated distilled data. We note that similar two-phase procedure has been studied in the context of representation learning [31].
Meta-Gradient Computation: If we consider the mean square error loss, then the optimal weights for the linear classifier have a closed-form solution. Moreover, since the feature dimension is typically larger than the number of distilled data, we can use kernel ridge regression (KRR) with a conjugate kernel [25] rather than solving the weights explicitly [12].
The resulting meta-training loss (Eq. 2) is similar to that used in KIP [22, 23], but we use a more flexible kernel rather than NTK.
[Equation]
where p X t , Y t q and p X s , Y s q are the inputs and labels of the real data T and distilled data S respectively.
The Gram matrix between real inputs and distilled inputs is denoted as K θX t X s P R | T |ˆ| S | , while the Gram matrix between distilled inputs is denoted as K θX s X s P R | S |ˆ| S | .
λ controls the regularization strength for KRR.
Let us denote the neural network feature for a given input X and model parameter θ as f p X, θ q P R N ˆ d , where N is the number of input and d is the feature dimension 1 .
The conjugate kernel is defined by the inner product of the neural network features.
Thus, the two
Gram matrices are computed as follows:
[Equation]
Now, computing the meta-gradient ∇ S L p
A lg p θ, S q , T q is just back-propagating through the conjugate kernel and a fixed feature extractor, which is very efficient and takes even fewer operations than computing the gradient for the network’s weights.
Moreover, we decouple the meta-gradient computation from the model online update.
Hence, we can train the online model using any optimizer, and the distilled data will be agnostic to the specific learning algorithm choice.
Our proposed method is similar to 1-step TBPTT in that we compute the meta-gradient at each step while performing the online model update.
Unlike the conventional 1-step TBPTT, we compute the meta-gradient using a
KRR output layer to mitigate truncation bias, illustrated in Figure 2(a).
Model Pool: As discussed in Section 1, there are various types of overfitting in dataset distillation.
Several techniques have been proposed to alleviate such problem, such as random initialization [4], periodic reset [13, 5, 7], and dynamic bi-level optimization [19].
These techniques share the same underlying principle: the model diversity matters.
Thus, we propose to maintain a “model pool” filled with diverse set of parameters obtained from different number of training steps and different random initializations.
Unlike the previous methods that periodically training and resetting a single model, FRePo randomly sample a model from the pool at each meta-gradient computation and update it using the current distilled data.
However, if a model has been updated more than K steps, we reinitialize it with a new random seed.
From the meta-learning perspective, we maintain a diverse set of meta-tasks to sample from and avoid sampling very similar tasks at each consecutive gradient computation to avoid overfitting to a particular setup.
Pool Diversity: We can increase the regularization strength by increasing the diversity of the model pool by setting a larger K , using data augmentation when training the model on the distilled data, or using models with different architectures.
To keep our method simple, we use the same architecture for all models in the pool and do not use any data augmentation when training the model on the distilled data.
Thus, our model pool only contains models with different initialization, at different optimization stages, and trained at different time-step of the distilled data.
Related Work
Unrolling in Bi-Level Optimization: One way to compute the meta-gradient is to differentiate through the unrolled inner optimization [4, 11–13].
However, this approach inherits several difficulties of the unrolled optimization, such as: 1) large computation and memory cost [14]; 2) truncation bias with short unrolls [17]; 3) exploding or vanishing gradients with long unrolls [15]; 4) chaotic and poorly conditioned loss landscapes with long unrolls [16].
In contrast, our method considers approximating the inner optimization with kernel ridge regression instead of unrolled optimization.
Surrogate Objective: To avoid unrolled optimization, several works turn to surrogate objectives.
DC [5], DSA [7], and DCC [18] formulate the dataset distillation as a gradient matching problem between the gradients of neural network weights computed on the real and distilled data.
In contrast,
DM [8] and CAFE [19] consider the feature distribution alignment between the real and distilled data.
Moreover, MTT [20] shows that knowledge from many expert training trajectories can be distilled to a dataset by using a training trajectory matching objective.
Nevertheless, surrogate objectives may introduce new biases and thus, may not accurately reflect the true objective.
For example, gradient matching approaches [5, 7, 18] only focus on short-range behavior and may easily overfit to a biased set of samples that produce dominant gradients [19, 20].
Closed-form Approximation: An alternative way to circumvent unrolled optimization is to find a closed-form approximation to the inner optimization.
Based on the correspondence between infinitelywide neural networks and kernel methods, KIP [22, 23] approximates the inner optimization with
NTK [21].
In this case, the meta-gradient can be computed by back-propagating through the NTK.
However, computing NTK for modern neural networks is extremely expensive.
Thus, using NTK for dataset distillation requires thousands of GPU hours and sophisticated implementation of the distributed kernel computation framework [23].
Similar to ours, Bohdal et al. [12] also decomposes the neural network as a feature extractor and a linear classifier. However, they only learn the label and explicitly solve for the optimal classifier weights rather than perform KRR.
Dataset Distillation 4.1 Implementation Details
We compare our method to four state-of-the-art dataset distillation methods [7, 8, 20, 23] on various benchmark datasets [26, 32–37].
We train the distilled data using Algorithm 1 with the same set of hyperparameters for all experiments except stated otherwise.
Unlike prior work [7, 8, 20], we do not apply data augmentation during training. However, we apply the same data augmentation [7, 20] during evaluation for a fair comparison. We preprocess the data in a similar way as in previous works [20, 23] but use a wider architecture than previous works [7, 8, 20] because the KRR component does not behave well when the feature dimension is low, resulting in a significant performance drop for our method.
Results on the original architecture are included in Appendix C.6.
We evaluate each distilled data using five random neural networks and report the mean and standard deviation. For the baseline method, we report the best of the reported value in the original paper and our reproducing results.
For the sake of brevity, we provide implementation details about data preprocessing, distilled data initialization, and hyperparameters in Appendix A.1 and various ablation studies regarding the model pool, batch size, distilled data initialization, label learning, and model architectures in Appendix C. More distilled [Table] [Figure]
[Figure] [Figure] MTTDSADMFRePo Model Width
T i m e p e r S t e p ( s ) (c) Time
Per Step  64 128 2560.1 0.
MTTDSADMFRePo
Model Width
P ea k G P U
M e m o r
y U s a g e ( G
[Figure] [Figure] 4.2 Standard Benchmarks Distillation Performance: We first evaluate our method on six standard benchmark datasets.
We learn 1, 10, and 50 images per class for datasets with only ten classes, while we learn 1 andimages per class for CIFAR100 [34] with 100 classes, Tiny ImageNet [35] with 200 classes, and
CUB-200 [37] with 200 fine-grained classes.
As shown in Table in 1, we achieve the state-of-the-art performance in most settings despite the hyperparameter may be suboptimal.
Our method performs exceptionally well on datasets with a complex label space when learning few images per class.
For example, we improve the CIFAR100, Tiny ImageNet, and CUB-200 in one image per class setting from 24.3%, 8.8%, and 2.2% to 28.7%, 15.4%, and 12.4%, respectively.
Figure 4 shows that our distilled images look real and natural though we do not directly optimize for this objective.
We observe a strong correlation between the test accuracy and image quality: the better the image quality, the higher the test accuracy.
Our results suggest that a highly condensed dataset does not need to be very different from the real dataset as it may just reflect the most common pattern in a dataset.
We also report the KRR predictor’s test accuracy using the feature extractor trained on the distilled data.
When the dataset is as simple as MNIST [32], the KRR predictor achieves similar performance as the neural network predictor.
In contrast, for more complex datasets, the KRR predictor consistently outperforms the neural network predictor, with the most significant gap being 3.7% for Tiny ImageNet in the one image per class setting.
[Table]
[Figure]
Label Learning: A similar trend can also be observed for label learning.
When the dataset is simple and has only a few classes, label learning may not be necessary.
However, it becomes crucial for complex datasets with many labels, such as CIFAR100 and Tiny-ImageNet (See more details in
Appendix C.4).
Similar to the teacher label in the knowledge distillation [1], we observe that the distilled label also encodes the class similarity.
We identify three typical cases in Figure 4d.
The first group consists of highly confident labels with a much higher value for one class than other classes (large margin), such as sunflower, bicycle, and chair.
In contrast, the distilled labels in the second group are confident but may get confused with some closely-related classes (small margin).
For instance, the learned label for "girl" has almost equally high values for the girl, woman, man, boy, and baby, suggesting that these classes are very similar and may be difficult for the model to distinguish them apart.
The last group contains distilled labels with low values for all classes, such as bear, beaver, and squirrel.
It is often hard for humans to recognize the distilled images in such a group, suggesting that they may be the challenging classes in a dataset.
Training Cost Analysis: Figure 3a, 3b shows that our method is significantly more time-efficient than the previous methods.
When learning one image per class on CIFAR100, FRePo reaches a similar test accuracy (23.4%) to the second-best method (24.0%) in 38 seconds, compared toseconds for MTT, which is roughly two orders of magnitude faster.
Moreover, FRePo achieves 92% of its final test accuracy (26.4% out of 28.7%) in only 385 seconds.
As shown in Figure 3c, our algorithm takes much less time to perform one gradient step on the distilled data.
Thus, we can perform more gradient steps in a fixed time.
Furthermore, Figure 3d suggests that our algorithm has much less GPU memory requirement.
Therefore, we can potentially use a much larger and more complex model to take advantage of the advancement in neural network architecture.
Cross-Architecture Generalization: One desired property of our distilled data is that it generalizes well to architecture it has not seen during the training. Similar to previous works [5, 20], we evaluate the distilled data from CIFAR10 on a wide range of architectures which it has not seen during training, including AlexNet [38], VGG [39], and ResNet [40].
Table 2 shows that our method outperforms previous methods on all unseen architectures.
Instance Normalization (IN) [41], as the vital ingredient in several methods (DSA, DM, MTT), seems to hurt the cross-architecture transfer. The performance
[Table]
degrades a lot when no normalization (NN) is applied (Conv-NN, AlexNet) or using a different normalization, like Batch Normalization (BN) [42].
It suggests that the distilled data generated by those methods encode the inductive bias of a particular training architecture.
In contrast [Table] our distilled data generalize well to various architectures, including those without normalization (Conv-NN, AlexNet).
Note that Figure 1, 4 also indicate that our distilled data encode less architectural bias as the distilled images look natural and authentic.
A simple idea to further alleviate the overfitting of a particular architecture is to include more architectures in the model pool.
However, the training may not be stable as the meta-gradient computed by different architectures can be very different
[Table] 4.3 ImageNet
High Resolution ImageNet Subset To understand how well our method performs on high-resolution images, we evaluate it on ImageNette and ImageWoof datasets [36] with a resolution of 128x128.
We learn 1 and 10 images per class on both datasets and report the performance in Table 3 and visualize some distilled images in Figure 1.
As shown in Table 3, we outperform MTT on all settings and achieve much better performance when we distill ten images per class on a more difficult dataset
ImageWoof.
It suggests that our distilled data is better at capturing the discriminative features for each class.
Figure 1 shows that our distilled images look real and capture the distinguishable feature of different classes.
For the easy dataset (i.e., ImageNette), all images have clear different structures, while for ImageWoof, the texture of each dog seems to be crucial.
Resized ImageNet-1K: We also evaluate our method on a resized version of ILSVRC2012 [26] with a resolution of 64x64 to see how it performs on a complex label space.
Surprisingly, we can achieve 7.5% and 9.7% Top1 accuracy using only 1k and 2k training examples, compared to 1.1% and 1.4% using an equally-sized real subset.
Application 5.1 Continual Learning
Continual learning (CL)
[43] aims to address the catastrophic forgetting problem [43–45] when a model learns sequentially from a stream of tasks.
A commonly used strategy to recall past knowledge is based on a replay buffer, which stores representative samples from previous tasks [46–49].
Since sample selection is an important component of constructing an effective buffer [48–51], we believe distilled data can be a key ingredient for a continual learning algorithm due to its highly condensed nature.
Several works [6–8, 52] have successfully applied the dataset distillation to the continual learning scenario.
Our work shows that we can achieve much better results by using a better dataset distillation technique.
We follow Zhao and Bilen [8] that sets up the baseline based on GDumb [49] which greedily stores class-balanced training examples in memory and train model from scratch on the latest memory only.
In that case, the continual learning performance only depends on the quality of the replay buffer.
We perform 5 and 10 step class-incremental learning [53] on CIFAR100 with an increasing buffer size of 20 images per class.
Specifically, we distill 400 and 200 images at each step and put them into the replay buffer.
We follow the same class split as Zhao and Bilen [8] and compare our method to random [49], herding [54, 55], DSA [7], and DM [8].
We use the default data preprocessing and default model for each method in this experiment as we find it gives the best performance for each method.
We use the test accuracy on all observed classes as the performance measure [8, 48]. [Table]
[Figure] Figure 5 shows that our method performs significantly better than all previous methods.
The final test accuracy for all classes for our method (FRePo) and the second-best method (DM) are 41.6%, 33.9% in 5-step learning, and 38.0%, 34.0% in 10-step learning.
However, we notice that for FRePo, distilling 2000 images in a continual learning setup achieves a similar test accuracy (41.6%) as distilling only 1000 images from the whole dataset (41.3% from Table 1).
In addition, performance drops as we perform more steps.
It suggests that FRePo considers all available classes to derive the most condensed dataset.
Splitting the data into multiple groups and performing independent distillation may generate redundant information or fail to capture the distinguishable features.
5.2 Membership Inference Defense
Membership inference attacks (MIA) aim to infer whether a given data point has been used to train the model or not [56–58].
Ideally, we want a model to learn from the data but not memorize it to preserve privacy.
However, deep neural networks are well-known for their ability to memorize all the training examples, even on large and randomly labeled datasets [59].
Several methods have been proposed to defend against such attacks by either modifying the training procedure [60] or changing the inference workflow [61].
This section shows that the distilled data contain little information regarding sample presence in the original dataset.
Thus, instead of training on the original datasets, training on distilled data allows privacy preservation while retaining model performance.
We consider three distilled data generated by DSA [7], DM [8] and FRePo.
We perform five popular "black box" MIA provided by Tensorflow Privacy [62] on models trained on the real data or the data distilled from it.
The attack methods include a threshold attack and four model-based attacks using logistic regression (LR), multi-layer perceptron (MLP), random forest (RF) and K-nearest neighbor (KNN).
The inputs to those attack methods are ground-truth labels, model predictions, and losses.
To measure the privacy vulnerability of the trained model, we compute the area under the ROC curve (AUC) of an attack classifier.
Following prior work, [56, 63], we keep a balanced set of training examples (member) and test examples (non-member) with 10K each to maximize the uncertainty of
MIA.
Thus, the random guessing strategy results in a 50% MIA accuracy.
We conduct experiments on MNIST and FashionMNIST with a distillation size of 500.
For space reasons, we provide more implementation details and results in appendix.
As shown in Table 4, all models trained on the distilled data preserve privacy as their attack AUCs are closed to random guessing.
However, we observe a small drop in test accuracy compared to the model trained on the full dataset, which is expected as we only distill 500 examples instead of 10,000 examples. Compared to the model trained on an equally sized subset of the original data, the model trained on distilled data results in much better test performance.
Figure 5c, 5d demonstrate the trade-off between test accuracy and attack effectiveness as measured by ROC AUC.
It shows that early stopping can be an effective technique to preserve privacy.
However, we will still be under high
MIA risk if we perform early stopping by monitoring the validation loss.
In contrast, training a model on the distilled data does not have this problem as the attack AUCs keep at a very low level regardless of training steps.
Conclusion
We propose neural Feature Regression with Pooling (FRePo) to overcome two challenges in dataset distillation: meta-gradient computation and various types of overfitting in dataset distillation.
We obtain state-of-the-art performance on various datasets with a 100x reduction in training time and a 10x reduction in GPU memory requirement.
The distilled data generated by FRePo looks real and natural and generalizes well to a wide range of architectures.
Furthermore, we demonstrate two applications that take advantage of the high-quality distilled data, namely, continual learning and membership inference defense.
Broader Impact “Synthetic data”, in the broader sense of artificial data created by generative models, can help researchers understand how an otherwise opaque learning machine “sees” the world.
There have been concerns regarding the risk of fake data.
This paper explores a new research direction in generating synthetic data only for downstream classification tasks.
We believe this work can provide additional interpretability and potentially address the common concerns in machine learning regarding training data privacy. Acknowledgments and Disclosure of Funding We would like to thank Harris Chan, Andrew Jung, Michael Zhang, Philip Fradkin, Denny Wu, Chong
Shao, Leo Lee, Alice Gao, Keiran Paster, and Lazar Atanackovic for their valuable feedback. Jimmy Ba was supported by NSERC Grant [2020-06904], CIFAR AI Chairs program, Google Research Scholar Program and Amazon Research Award.
This project was supported by LG Electronics
Canada.
Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute for
Artificial Intelligence. Checklist The checklist follows the references.
Please read the checklist guidelines carefully for information on how to answer these questions.
For each question, change the default [TODO] to [Yes] , [No] , or [N/A] .
You are strongly encouraged to include a justification to your answer , either by referencing the appropriate section of your paper or providing a brief inline description.
For example:
• Did you include the license to the code and datasets? [Yes] See Section ?? . • Did you include the license to the code and datasets? [No] The code and the data are proprietary.
• Did you include the license to the code and datasets? [N/A] Please do not modify the questions and only use the provided macros for your answers. Note that the
Checklist section does not count towards the page limit. In your paper, please delete this instructions block and only keep the Checklist section heading above along with the questions/answers below. For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? [Yes] (b) Did you describe the limitations of your work?
[Yes] (c) Did you discuss any potential negative societal impacts of your work? [Yes] (d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes] 2. If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[N/A] (b) Did you include complete proofs of all theoretical results?
[N/A] 3. If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes] (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] 4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets? [No] (c) Did you include any new assets either in the supplemental material or as a URL? [Yes] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating? [No] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[No] 5. If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable? [N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
Datasets: We evaluate our methods on the following datasets: i) MNIST [32]: A standard image dataset consists of 10 classes of 28x28 grey-scale images of handwritten digits, including 60, training examples and 10,000 test examples. ii) FashionMNIST [33]: A direct drop-in replacement for MNIST [32] consisting of 10 classes of 28x28 grey-scale images of clothing, including 60, training examples and 10,000 test examples. We denote FashionMNIST as F-MNIST for simplicity.
iii)
CIFAR [34]: A standard image dataset with two tasks: one coarse-grained over 10 classes (CIFAR10) and one fine-grained over 100 classes (CIFAR100).
Both CIFAR10 and CIFAR100 have 50,000 training examples and 10,000 test examples with a resolution of 32x32.
iv) Tiny ImageNet [35]: A higher resolution (64x64) dataset with 200 classes, including 100,000 training examples and 10,000 test examples.
We denote Tiny ImageNet as T-ImageNet for simplicity.
v) ImageNette and ImageWoof [36]: ImageNette (assorted objects) and ImageWoof (dog breeds) are two 10-class subsets from ILSVRC2012 [26] designed to be easy and hard to learn, respectively.
ImageNette consists of 9,469 training examples and 3,925 testing examples, while ImageWoof containstraining examples and 3929 testing examples.
We resize all examples to a resolution of 128x128.
vi) ImageNet [26]: A standard image benchmark dataset consists of 1000 classes from the Large
Scale Visual Recognition Challenge 2012 [26], including 1,281,167 training examples and 50, testing examples.
We resize all examples to a resolution of 64x64.
vii) CUB-200 [37]: A fine-grained classification task, consisting of 200 subcategories belonging to birds.
It has 5,994 data points for training and 5,794 data points for testing.
We resize all examples to a resolution of 32x32.
Data Preprocessing: We use the standard preprocessing for all datasets but add regularized ZCA transformation for RGB datasets as described in KIP [22, 23].
However, unlike KIP, we do not apply layer normalization to all examples.
Moreover, for simplicity, we apply the same regularization strength λ “ 0 . 1 to all datasets rather than tune each dataset.
This regularization strength is tuned for CIFAR in KIP [23].
Besides, for ImageNette and ImageWoof, performing the full-size ZCA transformation is extremely expensive due to the high resolution, so we use a checkboard ZCA instead.
It is the reason why our distilled images in Figure 1(c) have checkboard artifacts. As for the previous method, DSA [7], DM [8] only use standard preprocessing, while MTT [20] and KIP [23] consider both standard preprocessing and ZCA preprocessing, but they implement ZCA preprocessing differently. To account for the difference in data preprocessing, we reproduce each method using our data preprocessing and report the best of the reported value in the original paper and our reproducing results. As shown in Table 1, it turns out that our data preprocessing can improve DSA [7] and DM [8] a lot on RGB datasets but achieves a comparable performance for MTT [20].
When visualizing the distilled images, we directly apply the reverse transformation according to the corresponding data preprocessing without any other modifications.
Models: We use a simple convolutional neural network for all experiments based on the architecture used in previous works [7, 20, 23].
It consists of several blocks of 3x3 convolution, normalization,
RELU, and 2×2 average pooling layer with stride 2.
We use 3, 4, and 5 blocks for datasets with resolutions 32x32, 64x64, and 128x128, respectively.
Unlike previous works, which use the same number of filters at every layer, we double the number of filters if the feature map size is halved, following the modern neural network designs that preserve the time complexity per layer[39, 40].
We observe this to be crucial for our method when we distill thousands of data because we need the feature dimension to be much larger than the distilled size to make the KRR work properly.
Furthermore, we replace the Instance Normalization [41] with Batch Normalization [42] during training and do not use any normalization during evaluation.
For DSA, DM, and MTT, the default normalization for both training and evaluation is the Instance Normalization [41].
In contrast, KIP [23] uses analytical NTK during training and uses a 1024-width network without normalization for evaluation.
Moreover, KIP [23] adds an extra convolution layer at the first layer of the neural network. We initialize our model using Lecun Initialization, which is the default in Flax library [64]. In contrast, DSA, DM, and MTT use Kaiming Initialization, which is the default in PyTorch, while KIP [23] initializes the model using random Gaussian with standard deviations ? 2 and 0.
for weights and biases, respectively.
We denote the default model used in each method as Conv for simplicity.
We reproduce each method using our model and report the best of the reported value in the original paper and our reproducing results. We find our model achieves comparable or worse performance than each method’s default methods, so we use their default model in most settings. Require data do 3: Ż v
4: Ż
5: Ż
6: Ż Ð ´ α ∇ Y L 7: while q
Besides, we notice that the distilled data can encode the architecture’s inductive bias, so we provide an ablation study on how the model architecture affects the distilled data in Section C.6.
Initialization: Staying consistent with previous works, [7, 20, 23], we initialize the distilled image with randomly sampled real images.
We also investigate initializing the distilled image with random
Gaussian noise and observe that initializing does not affect the final performance too much but affects the convergence speed.
We find the real initialization gives a decent convergence speed, so we choose the real initialization for simplicity rather than fine-tune the scale of random Gaussian initialization.
As for labels, we initialize them using a scaled mean-centered one-hot vector for the corresponding image, where the scaling factor (i.e., 1 {p a
C { 10 q ) depends on the number of classes
C . We find that this label initialization scheme speeds up the convergence but has little impact on the final performance. We provide an ablation study regarding the initialization scheme in Section C.3. Label Learning: Whether to learn the label is a Boolean hyperparameter in our experiments.
When true, we optimize it using Algorithm 1.
When false, we stop the gradient so that they remain fixed at its initialization.
We provide an ablation study on label learning in Section C.4. Training and Evaluation: The proposed algorithm has two versions depending on the order of meta-gradient computation and online model update. We present the two versions in Algorithmand highlight the difference in red. We denote computing the meta-gradient after the online model update as v1 and denote computing the meta-gradient before the online model update as v2.
The first implementation (v1) is more similar to the standard back-propagation through time as we first do the forward pass (online model update) and then the backward pass.
However, instead of backpropagating the gradient through the inner optimization, we backpropagate through the kernel.
In contrast, the insight behind the second implementation (v2) is that we want to find good data to train the last layer first, then we use it to train the whole network.
Empirically, we do not observe any difference between these two implementations, and we choose v2 as the default for all experiments. Besides, we use the distilled data and its flipped version in meta-gradient computation to mitigate the mirroring effect for RGB datasets. For evaluation, we follow the standard protocol [7, 20]: training a randomly initialized neural network from scratch on distilled data and evaluating on a held-out test dataset. We apply the same data augmentation as in previous work [7, 20] during evaluation for a fair comparison.
Hyperparameters: We aim to keep our method as simple and efficient as possible.
As a result, our method requires very few hyperparameter tuning efforts than all the previous methods.
We use the same set of hyperparameters for all experiments, except stated otherwise.
Specifically, we use LAMB optimizer [65] with a cosine learning rate schedule starting from 0.0003 for both images and labels. We use a batch size of 1024 and train the distilled data up to 2 million steps to see its long-run behavior. In practice, we observe that most of the convergence (more than 95% of its final test accuracy) is achieved after a few thousand steps with a slow, logarithmic increase with more iterations, as shown in Figure 3b.
{+The model pool contains ten models trained up to K “ 100 steps. Each model is trained using Adam optimizer [66] with a constant learning rate of 0.0003, and no weight decay is applied. We use the same kernel regularizer λ as KIP [23]. Instead of being a fixed constant, the regularizer is adapted to the scale of K θX s X s , λ “ λ 
+}
Tr p K θX s X s q , where λ 0 “ 10 ´ 6 . We note that our hyperparameter choice may be sub-optimal, but it is a good starting point. We conduct some ablation study regarding the hyperparameters in Section C. Moreover, we provide a hyperparameter tuning guideline for practitioners in Section D accompanied by a list of additional tricks that we find to improve the performance but do not include in the current algorithm.
Summary: We implement our method in JAX [67] and reproduce previous methods using their released code. To take into account the differences in data processing and architectures, we try our best to reproduce previous results by varying different data preprocessing and models. Experiments show that our data preprocessing can sometimes improve performance, but our model does not give better performance than previous methods. We report the best of the reported value in the original paper and our reproducing results.
We evaluate each distilled data using five random neural networks and report the mean and standard deviation.
A.2 Experimental Setups
profiler and torch.profiler for the JAX program and PyTorch program to measure the peak GPU memory usage.
Different from Figure 3a, 3b, we use the same model (i.e., Conv used by DSA [7]) and same batch size=256 in this section.
We take an optimistic estimation for the previous methods (i.e., DSA, MTT) by choosing a smaller inner loop or outer loop number to generate more data points before encountering out of memory errors.
Specifically, we use outer_loop=1, inner_loop=1 for DSA and use syn_steps=15 for MTT.
All models are trained on Nvidia Quadro RTX 6000 with 22.17GB memory, and 7.5 compute capability.
To give the best image quality, we use the distilled images provided by the original paper. ( by Besides, we adjust the contrasts to give the best visualization sunflower
Figure -class accuracies all classes observed up to a specific time point. For a fair comparison, we follow the same setup in DM [8], including the class split of five different runs. We use the same set of hyperparameters as other experiments and train the distilled data up to 500K steps.
We report the mean and standard deviation of the five different runs. We observe that the primary source of variance comes from the class split. We use the same set of hyperparameters as other experiments except for data augmentation. We do not apply any data augmentation when we train models on the distilled data since any type of regularization can alleviate the MIA risks, making it hard to see the effects of distilled data.
[Table]
paper and our reproducing results.
As for our method, Table 1, 3 report the best value of learning label or not learning label with our default architecture and data preprocessing.
We use the same set of hyperparameters for all experiments.
The complete results can be found in Table 7.
Besides, we also include the results when training the model on the full dataset using mean square error loss in
Table 5.
the reported value in the original paper.
We evaluate DSA, DM, and MTT in PyTorch and KIP and
FRePo in JAX.
We notice that our reproducing results for KIP are much better than the reproducing results reported in MTT [20].
It may be due to the differences in initialization or hyperparameter choices, such as learning rate.
Table 4,
6 present
MNIST, respectively.
We random sample 10K data points to generate the distilled data and sample another 10K non-overlapped data as non-member data.
We use the same set of hyperparameters as other experiments except for data augmentation.
We do not apply any data augmentation when we train models on the distilled data since any regularization can alleviate the MIA risks, making it hard to see the effects of distilled data.
B Additional Results
Resized ImageNet-1K: We also evaluate our method on a resized version of ILSVRC2012 [26] with a resolution of 64x64 to see how it performs on a complex label space.
As shown in Table 5, we can achieve 7.5% and 9.7% Top1 accuracy using only 1k and 2k training examples, compared to 1.1% and 1.4% using an equally-sized real subset or 19.8% on the full dataset using 1281167 training examples.
We notice that Mean Square Error (MSE) loss may not be suitable for a complex dataset like ImageNet as the same model trained with Cross-Entropy loss achieve a Top1 accuracy of 32.2%.
However, since we use MSE loss to evaluate the distilled data, we report the model trained using
MSE loss for a fair comparison.
Membership Inference Defense: We only show MNIST results in Section 5.2 due to space reasons.
We provide the same experimental results on FashionMNIST in Figure 6 and Table 6.
We observe a similar trend on the two different datasets: training a model on the distilled data can preserve the privacy while achieving a good performance.
We notice that there is a gap between training on the whole data set, which may be closed by distilling more data points and adding more noise when perform the distillation.
C Ablation Study
We conduct several ablation studies to understand the key components of the proposed method, including the role of kernel approximation (Section C.1), model pool (Section C.2), initialization (Section C.3), label learning (Section C.4), scalability (Section C.5), and model architecture C.6.
[Table]
[Figure]
training steps.
AUC keeps increasing when training a model on the real data for more steps.
In contrast, AUC keeps low when training on distilled data.
All experiments are conducted on CIFAR100 when learning 1 Img/Cls using Algorithm 1 with the default hyperparameters except stated otherwise.
C.1 FRePo vs TBPTT
Figure 7 illustrates the computation graph of FRePoand 1-step TBPTT, which is different from
Figure 2 where we compare FRePo with unrolled optimization.
FRePo differs from 1-step TBPTT in meta-gradient computation.
As shown in Figure 7, 1-step TBPTT computes the meta-gradient by backpropagating through the inner optimization, while FRePo uses backpropagating through a kernel and a feature extractor.
Figure 8 compares the training loss, training accuracy, and test accuracy of
FRePo and TBPTT when using the same training and evaluation protocol.
Note that FRePo computes the training loss and accuracy using the KRR head during training, while TBPTT uses the neural network head.
We use the same optimizer to perform the online model update and use the same optimizer to train the same neural network for the same amount of steps on the distilled data during evaluation.
Figure 8c shows that the distilled data generated by TBPTT keeps getting worse test accuracy as the training goes on.
It suggests that the meta-gradient computed by TBPTT is highly biased and does not help learn a generalizable distilled dataset. As a result, the distilled data is overfitted to a k-step learning setup, where the k is the truncation step. This learning scenario is very similar to the dataset distillation setup in DD [4] and MTT [20], where a specific optimizer is learned to take advantage of the distilled data. Though using more truncation steps can alleviate this problem, eliminating the truncation bias needs an infinite unrolled optimization, making it intractable.
In contrast, FRePo alleviates the truncation bias by training a subset of a neural network to convergence. Moreover, FRePo decouples the meta-gradient computation and online model update such that the distilled data will not overfit the inner-level optimization. Conversely, TBPTT needs to fine-tune the inner optimization to get the best performance, which we do not explore further here.
[Figure]
[Figure]
subset of a neural network to convergence.
We investigate how the number of online models, the maximum online updates, and batch size affect the test accuracy. When one hyperparameter is modified, all the other hyperparameters are unchanged at their default value.
As shown in Figure 9a, using more than one model is better than using only one model, which is equivalent to the training and resetting strategy used in the previous methods [5, 7, 13].
Our default choice of using ten models seems not to be the best but gives reasonable performance. Figure 9b shows that both using too few updates and using too many updates hurt the performance, and our default choice of 100 gives a good starting point. If we want to squeeze the performance, it is worth tuning these two hyperparameters. Intuitively, a small regularization strength is needed when we distill a small number of distilled data.
On the contrary, when we distill more distilled data, we may want to use a stronger regularization.
Figure 9c, 9d show that using a large batch size may give slightly better results and converge faster in terms of number of training steps.
However, time per step is also an increasing function of batch size, so a large batch size may not give the best test accuracy and efficiency trade-off.
Our default choice of 1024 may not be optimal, but it is a good starting point.
C.3 Initialization
We initialize the distilled image using real images and initialize the distilled label using a meancentered one-hot vector scaled by 1 {p a
C { 10 q as default.
Is this real initialization that explains why our distilled images look real and natural?
Does random initialization give a similar result?
How does the label initialization affect the performance?
We answer these questions by trying different initialization schemes.
We investigate initializing the distilled image from random Gaussian noise or a combination of real images and random noises.
We also vary the scaling factor of the label initialization and the noise scale to figure out the best setting for dataset distillation.
Figure 10a shows that initializing from random noise is not a problem and the best initialization scheme is the
[Figure]
[Figure]
T e s t
A cc u r ac y ( % )
[Figure]
scale=0.2scale=0.3scale=0.
scale=0.5scale=0.6scale=0.7scale=0.
scale=0.9scale=1.
[Figure]
Steps
T e s t
A cc u r ac y ( % ) (
[Figure]
Random Image 3k 10k 30k 100k 300k   scale=0.0scale=0.1scale=0.2scale=0.
scale=0.4scale=0.5scale=0.6scale=0.
scale=0.8scale=0.9scale=1.
[Figure] T e s t [Figure] cc u r
[Figure]
[Figure]
combination of real images and random noises.
With a properly chosen standard deviation of 0.5, we can achieve 27.8% test accuracy compared to 27.2% using the default hyperparameter.
Figure 10b, 10c,10d show that a right scale of noise or label can significantly improve the convergence speed.
Besides, Figure 10b shows that scaling the mean-centered one-hot vector by 0.3 is a good choice for CIFAR100, and we also find that 1.0 works well for datasets with ten classes. Therefore, we decide to use the mean-centered one-hot vector scaled by 1 {p a C { 10 q as our default label initialization.
As shown in Figure 11, 12, though the distilled images look very different at initialization, they look quite similar after optimization.
We also provide four videos to visualize the evolution of the distilled images in the supplementary.
[Figure] [Figure]
C.4 Label Learning
As discussed in Section 4.2, label learning is an essential component of our algorithm.
We provide a detailed comparison of learning and not learning labels in Table 7.
We observe that when the label space is simple, such as MNIST, F-MNIST, and CIFAR10, label learning may not be necessary (Figure 21, 22, 23). However, it becomes crucial for complex datasets with many classes, such as CIFAR100 (Figure 13), Tiny-ImageNet (Figure 26), and ImageNet (Figure 27, 28). For ImageNet, we can achieve 7.5% test accuracy when we learn the label, compared to 1.6% when we fix the label.
Besides the fact that the distilled label encodes rich information regarding the class similarity, we find that learning labels can make the optimization easier and distill more natural and real-looking images, as shown in Figure 13, 26 and supplementary videos.
[Table]
C.5 Training Cost Analysis
Similar to Figure
[Table] we also investigate how time per step and GPU memory usage vary when we increase the number of distilled data. As shown in Figure [Table]
our method becomes more expensive as we increase the number of distilled data.
It is because 1) we always use all the distilled data for gradient computation rather than sample a batch as in other methods; 2) Matrix inversion with time complexity of O p N 3 q in KRR becomes more and more expensive as we distill more data.
Similar to other kernel methods, distilling tens of thousands of data can be difficult for our method.
We can circumvent this problem by 1) sampling a batch at each gradient computation, 2) performing subset distillation, or 3) distillation by class as in Section 5.1.
However, the performance is expected to drop because the redundant information can be generated in different groups, and the distilled data may not be able to capture all the distinguishable features when learning from a subset.
We leave it for future work to address this scaling challenge.
We provide the numerical values for Figure 3c, 3d, 14a, and 14b in Table 8, 9, 10, and
[Table]
[Table]
[Table]
[Table]
[Table]
[Figure]
T i m e p
[Figure]
S t e p
[Figure]
C.6 Model Architectures
From the Cross-Architecture Generalization experiments in Section 4.2, we observe that the distilled data can encode the architecture’s inductive bias.
Thus, we perform a qualitative and quantitative comparison of the distilled data generated by different architectures to understand how architecture affects the distilled data.
We evaluate across various architectures, including Conv (Our default model), DCConv (the default model of DSA [7], DM [8], MTT [20]), AlexNet [38], VGG [39], and
ResNet [40].
We also consider a wide range range of normalization layers, such as no normalization (NN), Instance Normalization (IN) [41], Batch Normalization (BN) [42], Layer Normalization (LN) [68], and Group Normalization (GN) [69].
Besides, we also vary the depth of Conv and denote twolayer, three-layer, and four-layer Conv as Conv-BN-D2, Conv-BN-D3, and Conv-BN-D4, respectively.
We do not rescale any image for better visualization, so the over-saturated images indicate that the images have different statistics from the real images.
Qualitative Results : Figure 15, 16, 17, 18, and 19 show that 1) the simplest architecture (i.e., Conv) gives the images that look almost like real images; 2) Different normalization layers have different effects on the distilled images, resulting in images with very different brightness and contrasts.
No
Normalization or Batch Normalization seems to generate the most natural-looking images; 3)
The distilled images generated by modern architectures like ResNet and VGG are very different from the natural images.
Thus additional attention is needed to transfer the distilled images to a different architecture.
On the other hand, additional tricks such as image regularization or projection may help make the distilled images more similar to natural ones.
4) Number of average pooling layers (or size of the final activation map) can affect the distilled image quality.
Using a fewer pooling (i.e.,
Conv-BN-D2) will generate blur images or repeated objects in the images, and we find it is more obvious in high resolution such as Tiny-ImageNet.
5) Distilled images may reflect the similarity of the architecture.
For instance, the images generated by Conv and DCConv are almost identical since the only architectural difference is the filter width.
Besides, we observe that the images generated by
AlexNet are very similar to those generated by Conv, which suggests that the inductive bias of those two architectures is very similar.
It may be one reason why Conv’s distilled images work extremely well for AlexNet (Table 2).
[Figure]
[Figure]
[Figure] [Figure] [Figure] Quantitative Results : We evaluate the distilled data trained using different architectures (Figureand 16) to study the role of normalization and width in training and evaluation.
Table 12, 13, andshow that 1) Training with Conv-BN and evaluating using Conv-NN yields the best performance, which is our default choice; 2) Training with Conv-BN yields most generalizable and transferable images as it performs well for other architectures; 3) Evaluating using Conv-GN seems to be the best choice if the training architecture is unknown.
4) Evaluating using Conv-NN is a good way to see whether the inductive bias of architecture has been distilled to the dataset since it is very sensitive to the training architecture.
There, it is not a good idea to use this architecture when the training architecture is unknown; 5) Training the distilled data using a wider network Conv achieves slightly better performance than a narrower network DCConv.
We also perform the same experiments in Table 1 using Conv-IN, and DCConv-IN and summarize the results in Table 16 and 17.
Moreover, we also evaluate the cross-architecture transfer performance of Conv-IN in Table 15. We observe that DCConv-IN works reasonably well when we distill a small number of images ( 100). The performance degrades a lot when distilling 1000 images from CIFAR100 because the KRR component needs a larger feature dimension to perform well. Besides, Conv-IN performs slightly worse than the default Conv-BN.
However, Table 15 suggests that the drawback of the instance norm is the transferability.
The distilled data trained using instance normalization transfer less well to other architecture, especially those without normalization.
[Figure]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
D Hyperparameter Tuning Guideline
We find that several modifications to the current method can improve the test accuracy of the model trained on the distilled data.
We do not include them in our current algorithm for simplicity or fair comparison, and we guess it may be helpful for practitioners.
• Dropout: We find dropout is very effective at alleviating the overfitting when training on a small set of examples.
• Learning Rate Schedule for the online model: Though we use a constant schedule in our experiments, we find that the learning rate schedule, especially the warm-up phase, may be crucial for some architectures.
When NAN is in gradient, adding a learning rate schedule to the online model may solve the problem.
• Data Augmentation during Training:
There are two ways to add data augmentation during training.
One is to add to X s during the online model update.
The other is to add to X t , which can be thought of as distilling data augmentation to the data.
• Tune Max Online Update: We train the online model up to 100 steps in our experiments.
When the distillation size is tiny (e.g., ten examples in total), 100 may be too large.
Setting to a lower value (less regularization) turns out to be better.
On the contrary, if the distillation size is large, setting it to a higher value (more regularization) can give better results.
• Exponential Moving Average (EMA): We find that evaluating on the EMA version of the distilled data or using the EMA version of the model parameters can improve the test accuracy.
• Soft Cross Entropy Loss: We train models on the distilled dataset using MSE loss in all our experiments to take advantage of the distilled label.
An alternative way is to use the soft cross-entropy loss with a fine-tuned temperature, which usually outperforms the MSE loss.
• Image Regularization: Though our method does not have any image regularization, we find a regularization term on the image norm is necessary for some architectures.
Otherwise, the image norm will keep increasing or decreasing.
We find that using an L2 penalty between the distilled image norm and the real image norm is enough in some cases. An alternative way is to project the distilled image to a norm ball every few iterations. • Label Regularization: Our method has a small regularization to ensure the class-balanced distillation where we force the margin between the target label and any other label is greater than 1/C. We believe a better label regularization incorporating prior knowledge of class similarity can improve the performance. • Maximal Update Parametrization ( µP ) [70]: We observe that using a model parameterization proposed by Yang et al.
[70] can give additional test accuracy improvement. Based on our experience, we provide the following hyperparameter tuning guideline for those who want to squeeze the performance of our method or apply our method to a different dataset or use a different model. Besides validation loss and accuracy, we suggest monitoring the norm and gradient norm of the distilled images and labels, which can be good indicators for the final performance.
• Step1 - Online Model: Choose an online model architecture and tune its hyperparameter (e
g., optimizer, weight decay) in the standard way on the whole real dataset or a subset of real data. The same hyperparameters can be used for the online model update and final evaluation.
• Step2 - Distilled Data Optimization and Initialization: Use the default setting for the model pool (i
e., ten models with max online update K “ 100 ) and tune the learning rate and batch size for distilled data optimization and the scale of initialization. • Step3 - Pool Diversity: Tune the model pool diversity by adjusting the max online update, adding models with different architectures, or applying data augmentation. E Additional Visualization E.1 Distilled Image Visualization
We provide some additional distilled images visualization for MNIST (Figure 21), FashionMNIST (Figure 22), CIFAR10 (Figure 23), CIFAR100 (Figure 24), CUB-200 (Figure 25), Tiny ImageNet (Figure 26), ImageNet (Figure 27, 28), ImageNette (Figure 29)
[Figure]
and ImageWoof (Figure 30).
[Figure]
[Figure]
[Figure]
[Figure] [Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]