Soft-Label Dataset Distillation and Text Dataset Distillation

Introduction

The increase in computational requirements for modern deep learning presents a range of issues. The training of deep learning models has an extremely high energy consumption (Strubell et al., 2019), on top of the already problematic financial cost and time requirement. One path for mitigating these issues is to reduce network sizes. Hinton et al. (2015) proposed knowledge distillation as a method for imbuing smaller, more efficient networks with all the knowledge of their larger counterparts. Instead of decreasing network size, a second path to efficiency may be to decrease dataset size. Dataset distillation (DD) has recently been proposed as an alternative formulation of knowledge distillation to do exactly that (Wang et al., 2018).

Dataset distillation is the process of creating a small number of synthetic samples that can quickly train a network to the same accuracy it would achieve if trained on the original dataset. It may seem counter-intuitive that training a model on a small number of synthetic images coming from a different distribution than the training data can result in comparable accuracy, but Wang et al. (2018) have shown that for models with known initializations this is indeed feasible; they achieve 94% accuracy on MNIST, a hand-written digit recognition task (LeCun et al., 1998), after training LeNet on just 10 synthetic images.

We propose to improve their already impressive results by learning ‘soft’ labels as a part of the distillation process. The original dataset distillation algorithm uses fixed, or ‘hard’, labels for the synthetic samples (e.g. the ten synthetic MNIST images each have a label corresponding to a different digit). In other words, each label is a one-hot vector: a vector where all entries are set to zero aside from a single entry, the one corresponding to the correct class, which is set to one. We relax this one-hot restriction and make the synthetic labels learnable. The resulting distilled labels are thus similar to those used for knowledge distillation as a single image can now correspond to multiple classes. An example comparing a ‘hard’ label to a ‘soft’ label is shown in Figure 2. A ‘hard’ label can be derived from a ‘soft’ label by applying the softmax function and setting the element with the highest probability to one, while the remaining elements are set to zero. Our soft-label dataset distillation (SLDD) not only achieves over 96% accuracy on MNIST when using ten distilled images (as seen in Figure 1), a 2% increase over the state-of-the-art (SOTA), but also achieves almost 92% accuracy with just five distilled images, which is less than one image per class. In addition to soft labels, we also extend dataset distillation to the natural language/sequence modeling domain and enable it to be used with several additional neural network architectures. For example, we show that Text Dataset Distillation (TDD) can train a custom convolutional neural network (CNN) (LeCun et al., 1999) with known initialization up to 90% of its original accuracy on the IMDB sentiment classification task (Maas et al., 2011) using just two synthetic sentences.

The rest of this work is divided into four sections. In Section 2, we discuss related work in the fields of knowledge distillation, dataset reduction, and example generation. In Section 3, we propose improvements and extensions to dataset distillation and associated theory. In Section 4, we empirically validate SLDD and TDD in a wide range of experiments. Finally, in Section 5, we discuss the significance of SLDD and TDD, and our outlook for the future.

Related Work

Dataset distillation was originally inspired by network distillation (Hinton et al., 2015) which is a form of knowledge distillation or model compression (Buciluǎ et al., 2006). Network distillation has been studied in various contexts including when working with sequential data (Kim and Rush, 2016). Network distillation aims to distill the knowledge of large, or even multiple, networks into a smaller network. Similarly, dataset distillation aims to distill the knowledge of large, or even multiple, datasets into a small number of synthetic samples. ‘Soft’ labels were recently proposed as an effective way of distilling networks by feeding the output probabilities of a larger network directly to a smaller network (Hinton et al., 2015), and have previously been studied in the context of different machine learning algorithms (El Gayar et al., 2006). Our soft-label dataset distillation (SLDD) algorithm also uses ‘soft’ labels but these are persistent and learned over the training phase of a network (rather than being produced during the inference phase as in the case of network distillation).

2 Learning from ‘small’ data

Deep supervised learning generally requires a very large number of examples to train on. For example, MNIST and CIFAR10 both contain thousands of training images per class. Meanwhile, it appears that humans can quickly generalize from a tiny number of examples (Lake et al., 2015). Getting machines to learn from ‘small’ data is an important aspect of trying to bridge this gap in abilities. Dataset distillation provides a method for researchers to generate synthetic examples that are optimized for allowing machines to learn from a small number of them. Studying the distilled images produced by dataset distillation may enable us to identify what allows neural networks to generalize so quickly from so few of them. In some sense, dataset distillation can be thought of as an algorithm for creating dataset summaries that machines can learn from.

3 Dataset Reduction, Prototype Generation, and Summarization

There are a large number of methods that aim to reduce the size of a dataset with varying objectives. Active learning aims to reduce the required size of the labeled portion of a dataset by only labeling examples that are determined to be the most important (Cohn et al., 1996; Tong and Koller, 2001). Several methods aim to ‘prune’ a dataset, or create a ‘core-set’, by leaving in only examples that are determined to be useful (Angelova et al., 2005; Bachem et al., 2017; Sener and Savarese, 2017; Tsang et al., 2005). In general, all of these methods use samples from the true distribution, typically subsets of the original training set. By lifting this restriction and, instead, learning synthetic samples, dataset distillation requires far fewer samples to distill the same amount of knowledge.

In the field of nearest-neighbor classification, these dataset reduction techniques are typically referred to as ‘prototype selection‘ and ‘prototype generation‘, and are studied extensively as methods of reducing storage requirements and improving the efficiency of nearest-neighbor classification (Garcia et al., 2012; Triguero et al., 2011). As with the methods above, prototype selection methods use samples from the true distribution, typically just subsets of the original training dataset. Prototype generation methods typically create samples that are not found in the training data; however, these methods are designed specifically for use with nearest-neighbor classification algorithms.

All the dataset reduction methods discussed above also share another restriction. They all use fixed labels. Soft-label dataset distillation removes this restriction and allows the label distribution to be optimized simultaneously with the samples (or prototypes) themselves.

4 Generative Adversarial Networks

Generative Adversarial Networks (GANs) have recently become a widely used method for image generation. They are primarily used to produce images that closely mimic those coming from the true distribution (Ledig et al., 2017; Goodfellow et al., 2014; Choi et al., 2018; Radford et al., 2015). For dataset distillation, we instead set knowledge distillation as the objective but do not attempt to produce samples from the true distribution. Using the generator from a trained GAN may be a much faster way of producing images than the gradient-based method employed by dataset distillation. However, since the number of distilled images we aim to produce is very small, solving the objective directly through gradient-based optimization is sufficiently fast, while also more straightforward. Additionally, while some GANs can work with text (Reed et al., 2016; Yu et al., 2017), they are primarily intended for image generation.

5 Measuring Problem Dimensionality

We may intuitively believe that one deep learning task is more difficult than another. For example, when comparing the digit recognition task MNIST, to the image classification task CIFAR10 (Krizhevsky et al., 2009), CIFAR10 appears to be the more difficult problem though it is difficult to measure the extent of this increase in difficulty. One approach is to compare error rates for state-of-the-art (SOTA) results on these datasets. For example, the near-SOTA ‘dropconnect’ model on MNIST achieves a 0.21% error rate, while on CIFAR10 it achieves an error rate of 9.32% (Wan et al., 2013). However, this approach reveals less as deeper networks approach perfect accuracy on multiple tasks. Li et al. (2018) instead derive a fairly model-independent metric for comparing the dimensionality of various problems based on the minimum number of learnable parameters needed to achieve a good local optimum. Similarly, dataset distillation aims to find the minimum number of synthetic samples needed to achieve a good local optimum. The difference is that Li et al. (2018) constrain the number of searchable dimensions within the network weight space, while dataset distillation constrains them within the data space.

Extending Dataset Distillation

As mentioned above, nearest-neighbors classification often involves data reduction techniques known as prototype selection and prototype generation. We can use these concepts, along with the k-Nearest Neighbors (kNN) classification algorithm, to visualize the difference between classical dataset reduction methods, dataset distillation, and soft-label dataset distillation. When we fit a kNN model, we essentially divide the entire space into classes based on the location of points in the training set. However, the cost of fitting a kNN model increases with the number of training points. Prototype selection and generation are methods for reducing the number of points in the training set while trying to maintain the accuracy of the original model.

Prototype selection methods use subsets of the training set to construct the reduced training set. In the first column of Figure 3, we visualize attempts to reduce the training set for a three-class problem, the Iris flower dataset (Fisher, 1936), by selecting one point from each class and then fitting the kNN on these three selected points. It is clear from this visualization that only being able to use a subset of the original points, limits the ability to finely tune the resulting separation of the space. Prototype generation methods relax this restriction and create synthetic points whose placement can be optimized for resulting kNN performance. This is effectively what dataset distillation does for neural networks. In the second column of Figure 3, we generate one point for each class and then optimize the location of one of these points, before fitting the kNN on all three. Using synthetic, optimizable points allows for finer-grained tuning of the resulting class separation. In both these cases, each of the three resulting points had a single class assigned to it.

We now propose that the three points instead be assigned an optimizable distribution of classes, effectively a ‘soft’ label as described above. In the third column of Figure 3, in order to visualize the effect of changing a point’s label distribution, we arbitrarily fix one point for each class, but we change the label distribution of one of these points, increasingly making it a mixture of the other classes, and then fit the kNN. In the final column of Figure 3, we combine the prototype generation method with our soft-label modification, to visualize the effect of simultaneously changing a point’s location and label distribution. This last case is the kNN counterpart to our proposed soft-label distillation algorithm, and from the visualization, it is clear that it provides the finest-grained tuning for the resulting class separation. In fact, by using soft labels with kNN, we can separate three classes using just two points, as seen in Figure 4.

Animated versions of both Figure 3 and Figure 4 are in the online appendix.

2 Basic Approach

Our basic approach is the same as Wang et al. (2018). We summarize it here in a slightly modified way to explicitly show the labels of the distilled dataset. This additional notation becomes useful once we enable label learning in the next section.

In general, training with stochastic gradient descent (SGD) involves repeatedly sampling mini-batches of training data and updating network parameters by their error gradient scaled by learning rate $\eta$.

3 Learnable Labels

As mentioned above, one formulation of knowledge distillation proposes that a smaller network be trained on the outputs of a larger network rather than the original training labels. Unlike the training labels, the output labels of the larger network are not ‘hard’ labels. Because they are generally the outputs of a softmax layer, the output labels form a probability distribution over the possible classes. The idea is that any training image contains information about more than one class (e.g. an image of the digit ‘3’ looks a lot like other digits ‘3’ but it also looks like the digit ‘8’). Using ‘soft’ labels allows us to convey more information about the associated image.

4 Text and Other Sequences

The original dataset distillation algorithm was only shown to work with image data, but intuitively, there is no reason why text or other sequences should not be similarly distillable. However, it is difficult to use gradient methods directly on text data as it is discrete. In order to be able to use SLDD with text data, we need to first embed the text data into a continuous space. This is a common practice when working with many modern natural language processing models, though the embedding method itself can vary greatly (Ma and Hovy, 2016; Devlin et al., 2018; Peters et al., 2018). Any popular embedding method can be used; in our experiments, we used pre-trained GloVe embeddings (Pennington et al., 2014). Once the text is embedded into a continuous space, the problem of distilling it becomes analogous to soft-label image distillation. If all sentences are padded/truncated to some pre-determined length, then each sentence is essentially just a one-channel image of size [length]$*$[embedding dimension]. When working with models that do not require fixed-length input, like recurrent neural networks (RNN), the distilled sentences do not have to be padded/truncated to the same length. It is also important to note that the embedding is performed only on sentences coming from the true dataset; the distilled samples are learned directly as embedded representations. Since the distilled data produced by this algorithm would still be in the embedding space, it may be of interest to find the nearest sentences that correspond to the distilled embeddings. To compute the nearest sentence to a distilled embedding matrix, for every column vector in the matrix, the nearest embedding vector from the original dictionary must be found. These embedding vectors must then be converted back into their corresponding words, and those words joined into a sentence. The resulting algorithm for text dataset distillation (TDD) is detailed in Algorithm 1b which is a modification of the SLDD Algorithm 1a.

5 Random initializations and multiple steps

The procedures we described above make one important assumption: network initialization $\theta_{0}$ is fixed. The samples created this way do not lead to high accuracies when the network is re-trained on them with a different initialization as they contain information not only about the dataset but also about $\theta_{0}$. In the distilled images in Figures 1 and 5, this can be seen as what looks like a lot of random noise. Wang et al. (2018) propose that the method instead be generalized to work with network initializations randomly sampled from some restricted distribution.

The resulting images, especially for MNIST, appear to have much clearer patterns and much less random noise, and the results detailed in Section 4 suggest that this method generalizes fairly well to other randomly sampled initializations from the same distribution.

Additionally, Wang et al. (2018) suggest that the above methods can work with multiple gradient descent (GD) steps. If we want to perform multiple gradient descent steps, each with a different mini-batch of distilled data, we simply backpropagate the gradient through every one of these additional steps. Finally, it may also be beneficial to train the neural networks on the distilled data for more than one epoch. The experimental results suggest that multiple steps and multiple epochs improve distillation performance for both image and text data, particularly when using random network initializations.

Experiments

The simplest metric for gauging distillation performance is to train a model on distilled samples and then test it on real samples. We refer to the accuracy achieved on these real samples as the ‘distillation accuracy’. However, several of the models we use in our experiments do not achieve SOTA accuracy on the datasets they are paired with, so it is useful to construct a relative metric that compares distillation accuracy to original accuracy. The first such metric is the ‘distillation ratio’ which we define as the ratio of distillation accuracy to original accuracy. The distillation ratio is heavily dependent on the number of distilled samples so the notation we use is $r_{M}=100\%*\frac{[\text{distillation accuracy}]}{[\text{original accuracy}]},M=[\text{number of distilled samples}]$. We may refer to this metric as the ‘$M$-sample distillation ratio’ when clarification is needed. It may also be of interest to find the minimum number of distilled images required to achieve a certain distillation ratio. To this end we can define a second relative metric that we call the ‘$A$% distillation size’, and we write $d_{A}=M$ where $M$ is the minimum number of distilled samples required to achieve a distillation ratio of $A$%.

2 Image Data

The LeNet model we use with MNIST achieves nearly SOTA results, 99% accuracy, so it is sufficient to use distillation accuracy when describing distillation performance with it. However, AlexCifarNet only achieves 80% on CIFAR10 so it is helpful to use the 2 relative metrics when describing this set of distillation results.

It is useful to compare dataset distillation against several other methods of dataset reduction. We use the following baselines suggested by Wang et al. (2018).

Random real images: We randomly sample the same number of real images per class from the training data. These images are used for two baselines: training neural networks and training K-Nearest Neighbors classifiers.

Optimized real images: We sample several sets of random real images as above, but now we choose the 20% of these sets that have the best performance on training data. These images are used for one baseline: training neural networks.

$k$-means: We use $k$-means to learn clusters for each class, and keep the resulting centroids. These images are used for two baselines: training neural networks and training K-Nearest Neighbors classifiers.

Average real images: We compute the average image for each class and use it for training. These images are used for one baseline: training neural networks.

Each of these baseline methods produces a small set of images that can be used to train models. All four of the baseline methods are used to train and test LeNet and AlexCifarNet on their respective datasets. Additionally, two of the baseline methods are used to also train K-Nearest Neighbor classifiers to compare performance against neural networks. The results for these six baselines, as determined by Wang et al. (2018), are shown in Table 3.

Fixed initialization.

When the network initialization is fixed between the distillation and training phases, synthetic images produced by dataset distillation result in high distillation accuracies. The SLDD algorithm produces images that result in equal or higher accuracies when compared to the original DD algorithm. For example, DD can produce 10 distilled images that train a LeNet model up to 93.76% accuracy on MNIST (Wang et al., 2018). Meanwhile, SLDD can produce 10 distilled images that train the same model up to 96.13% accuracy (Figure 1). The full distilled labels for these 10 images are laid out in Table 1. SLDD can even produce a tiny set of just 5 distilled images that train LeNet to 91.56% accuracy. As can be seen in Figure 6, the 90% distillation size (i.e. the minimum number of images needed to achieve 90% of the original accuracy) of MNIST with fixed initializations is $d_{A}=5$, and while adding more distilled images typically increases distillation accuracy, this begins to plateau after five images. Similarly, SLDD provides a 7.5% increase in 100-sample distillation ratio (6% increase in distillation accuracy) on CIFAR10 over DD. Based on these results, detailed further in Table 3, it appears that SLDD is even more effective than DD at distilling image data into a small number of samples. This intuitively makes sense as the learnable labels used by SLDD increase the capacity of the distilled dataset for storing information.

Random initialization.

It is also of interest to know whether distilled data are robust to network initialization. Specifically, we aim to identify if distilled samples store information only about the network initializations, or whether they can store information contained within the training data. To this end, we perform experiments by sampling random network initializations generated using the Xavier Initialization (Glorot and Bengio, 2010). The distilled images produced in this way are more representative of the training data but generally result in lower accuracies when models are trained on them. Once again, images distilled using SLDD lead to higher distillation accuracies than DD when the number of distilled images is held constant. For example, 100 MNIST images learned by DD result in accuracies of 79.5 $\pm$ 8.1%, while 100 images learned by SLDD result in accuracies of 82.75 $\pm$ 2.75%. There is similarly a 3.8% increase in 100-sample distillation ratio (3% increase in distillation accuracy) when using SLDD instead of DD on CIFAR10 using 100 distilled images each. These results are detailed in Table 3. It is also interesting to note that the actual distilled images, as seen in Figures 7 and 8, appear to have much clearer patterns emerging than in the fixed initialization case. These results suggest that DD, and even more so SLDD, can be generalized to work with random initializations and distill knowledge about the dataset itself when they are trained this way. All the mean and standard deviation results for random initializations in Table 3 are derived by testing with 200 randomly initialized networks.

3 Text Data

As mentioned above, TDD does not work in the space of the original raw data, but rather produces synthetic samples from the embedding space. Because each distilled ‘sentence’ is actually a matrix, the embedding layer is applied only to real sentences from the training data and not the matrices coming from the distilled data. For text experiments, we use the IMDB sentiment analysis dataset, the Stanford Sentiment Treebank 5-class task (SST5) (Socher et al., 2013), and the Text Retrieval Conference question classification tasks with 6 (TREC6) and 50 (TREC50) classes (Voorhees et al., 1999). The text experiments are performed with three different networks: a fairly shallow but wide CNN (TextConvNet), a bi-directional RNN (Bi-RNN) (Schuster and Paliwal, 1997), and a bi-directional Long Short-Term Memory network (Bi-LSTM) (Hochreiter and Schmidhuber, 1997). These models do not all achieve the same accuracies on the text datasets we use in distillation experiments, so the models’ original accuracies are detailed in Table 4. The distillation ratios appearing in this section are calculated based on these accuracies. The exact architectures of these models are detailed in the online appendix.

We consider the same six baselines as in the image case but modify them slightly so that they work with text data.

Random real sentences: We randomly sample the same number of real sentences per class, pad/truncate them, and look up their embeddings. These sentences are used for two baselines: training neural networks and training K-Nearest Neighbors classifiers.

Optimized real sentences: We sample and pre-process different sets of random real sentences as above, but now we choose the 20% of the sets that have the best performance. These sentences are used for one baseline: training neural networks.

$k$-means: First, we pre-process the sentences. Then, we use $k$-means to learn clusters for each class, and use the resulting centroids to train. These sentences are used for two baselines: training neural networks and training K-Nearest Neighbors classifiers.

Average real sentences: First, we pre-process the sentences. Then, we compute the average embedded matrix for each class and use it for training. These sentences are used for one baseline: training neural networks.

Each of these baseline methods produces a small set of sentences, or sentence embeddings, that can be used to train models. All four of the baseline methods are used to train and test the TextConvNet on each of the text datasets. Additionally, two of the baseline methods are used to also train K-Nearest Neighbor classifiers to compare performance against neural networks. The baseline results are shown in Table 3.

Fixed initialization.

When the network initialization is fixed between the distillation and training phases, synthetic text produced by text dataset distillation also results in high model accuracies. For example, TDD can produce 2 distilled sentences that train the TextConvNet up to a distillation ratio of 89.88% on the IMDB dataset. Even for far more difficult language tasks, TDD still has impressive results but with larger distilled datasets. For example, for the 50-class TREC50 task, it can produce 1000 distilled sentences that train the TextConvNet to a distillation ratio of 79.86%. Some examples of TDD performance are detailed in Table 3 and Table 5. The distilled text embeddings from the six-sentence Trec6 experiments are visualized in Figure 9. However, since these distilled text embeddings are still in the GloVe embedding space, it may be difficult to interpret them visually. We provide a more natural method for analyzing distilled sentences by using nearest-word decoding to reverse the GloVe embedding. We find the nearest word to each distilled vector based on Euclidean distances. The result of this decoding is an approximation of the distilled sentence in the original text space. We list the decoded distilled sentences corresponding to the matrices from Figure 9 in Table 6, along with their respective label distributions. These sentences can contain any tokens found in the TREC6 dataset, including punctuation, numbers, abbreviations, etc. The sentences do not have much overlap. This is consistent with the distilled labels which suggest that each sentence corresponds strongly to a different class. It appears that the TDD algorithm encourages the separation of classes, at least when there are enough distilled samples to have one or more per class. Additional results and visualizations for TDD with fixed initialization can be found in the online appendix.

Random initialization.

When using random initialization the performance decrease for TDD is similar to that for SLDD. TDD can produce two distilled sentences that train the TextConvNet with random initialization up to a distillation ratio of 79.96% on IMDB, or 20 distilled sentences that train it up to a distillation ratio of 85.22%. This is only slightly lower performance than in the fixed initialization case. However, there is a larger difference in performance between fixed and random initializations for the recurrent networks. For TREC6, TDD can produce six distilled sentences that train a randomly initialized Bi-LSTM to a distillation ratio of 69.33%, or 120 distilled sentences that train it to a distillation ratio of 78.87%. All the mean and standard deviation results for random initializations in Table 3 and Table 5 are derived by testing with 200 randomly initialized networks. The distilled text embeddings from the IMDB experiment with two distilled sentences are visualized in Figure 10. We list the decoded distilled sentences corresponding to these matrices in Table 7. These sentences can contain any tokens found in the IMDB dataset, including punctuation, numbers, abbreviations, etc. Since this is a binary sentiment classification task, each label is a scalar. If a probability is needed, a sigmoid function can be applied to the scalar soft labels. In this case, the distillation algorithm appears to have produced one sentence with a positive associated sentiment, and one with a negative sentiment. Curiously, the model appears to have overcome the challenge of having to describe these long sentences with a single scalar by using duplication. For example, in the second sentence, corresponding to the negative label, negative words like ‘dump’, ‘stupid’, and ‘shoddy’ are all repeated several times. In such a way, the algorithm is likely assigning lower sentiment scores to these words than other ones, while using only a single label. Additional results for TDD with random initialization can be found in the online appendix.

Conclusion

By introducing learnable distilled labels we have increased distillation accuracy across multiple datasets by up to 6%. By enabling text distillation, we have also greatly increased the types of datasets and architectures with which distillation can be used.

However, even with SLDD and TDD, there are still some limitations to dataset distillation. The network initializations used for both SLDD and TDD all come from the same distribution, and no testing has yet been done on whether a single distilled dataset can be used to train networks with different architectures. Further investigations are needed to determine more precisely how well dataset distillation can be generalized to work with more variation in initializations, and even across networks with different architectures.

Interestingly, the initialization of distilled labels appears to affect the performance of dataset distillation. Initializing the distilled labels with ‘hard’ label values leads to better performance than with random initialization, possibly because it encourages class separation earlier on in the distillation process. However, it is not immediately clear whether it is better to separate similar classes (e.g. ‘3’ and ‘8’ in MNIST), thereby increasing the network’s ability to discern between them, or to instead keep those classes together, thereby allowing soft-label information to be shared between them. It may be interesting to explore the dynamics of the distillation process when using a variety of label initialization methods.

We have shown dataset distillation works with CNNs, bi-directional RNNs, and LSTMs. There is nothing in the dataset distillation algorithm that would limit it to these network types. As long as a network has a twice-differentiable loss function and the gradient can be back-propagated all the way to the inputs, then that network is compatible with dataset distillation.

Another promising direction is to use distilled datasets for speeding up Neural Architecture Search and other very compute-intensive meta-algorithms. If distilled datasets are a good proxy for performance evaluation, they can reduce search times by multiple orders of magnitude. In general, dataset distillation is an exciting new branch of knowledge distillation; improvements may help us not only better understand our datasets but also enable several applications related to efficient machine learning.

We would like to thank Dr. Sebastian Fischmeister for providing us with the computational resources that enabled us to perform many of the experiments found in this work.

References