Dataset Distillation with Attention Labels for Fine-tuning BERT
Year: 2023

Abstract:
Dataset distillation aims to create a small dataset of informative synthetic samples to rapidly train neural networks that retain the performance of the original dataset. In this paper, we focus on constructing distilled fewshot datasets for natural language processing (NLP) tasks to fine-tune pre-trained transformers. Specifically, we propose to introduce attention labels, which can efficiently distill the knowledge from the original dataset and transfer it to the transformer models via attention probabilities. We evaluated our dataset distillation methods in four various NLP tasks and demonstrated that it is possible to create distilled few-shot datasets with the attention labels, yielding impressive performances for finetuning BERT. Specifically, in AGNews, a fourclass news classification task, our distilled fewshot dataset achieved up to 93.2% accuracy, which is 98.5% performance of the original dataset even with only one sample per class and only one gradient step.

heading: 1 Introduction
text: Deep learning models have achieved state-ofthe-art performance in various fields, including computer vision and natural language processing (NLP), using large-scale neural networks trained with huge datasets. Unfortunately, their successful performances have come with massive training costs, including training time, GPU resources, and energy consumption. To reduce the training costs, current research has been focusing on constructing a small training dataset such that models trained with it can achieve comparable performances to models trained with the whole original dataset.
One classical way to compress the training dataset is data selection. Data selection methods choose a subset of effective training samples on the basis of a number of heuristic measures, for example, cluster centers (Sener and Savarese, 2018), diversity (Aljundi et al., 2019), and likelihood of
models (Moore and Lewis, 2010). Although the data selection methods effectively work for efficient model training and several applications, such as active learning (Sener and Savarese, 2018) and continual learning (Aljundi et al., 2019), their performance is clearly restricted because they rely on the existence of representative samples that are effective for model training in the original dataset.
As an alternative approach for reducing the training dataset, Wang et al. (2018b) proposed dataset distillation, which aims to create a small number of synthetic samples optimized to effectively train models. Dataset distillation has attracted much attention in machine learning (Wang et al., 2018b; Zhao et al., 2021; Zhao and Bilen, 2021; Sucholutsky and Schonlau, 2021; Bohdal et al., 2020; Wang et al., 2022; Cazenavette et al., 2022) for both the theoretical interest and various applications, such as neural architecture/hyper-parameter search (Such et al., 2020), continual learning (Masarczyk and Tautkute, 2020; Rosasco et al., 2022), federated learning (Goetz and Tewari, 2020; Zhou et al., 2020), and preserving data privacy (Li et al., 2020; Dong et al., 2022).
However, most of the existing research on dataset distillation mainly focuses on image datasets, and only a few studies involve NLP tasks. Sucholutsky and Schonlau (2021) and Li and Li (2021) extended dataset distillation to text datasets by using embedding vectors as an input of the distilled dataset instead of discrete text. While these studies applied dataset distillation to those model architectures based on convolutional neural networks (CNNs) and recurrent neural networks (RNNs), we cannot find any research that tackles dataset distillation for pre-trained transformers, such as BERT (Devlin et al., 2019), which have become the de-facto standard for various kinds of NLP tasks. Therefore, in this paper, we aim to obtain distilled few-shot datasets to fine-tune the pre-trained transformers for NLP tasks.
119
To this end, we focus on the attention mechanism, which is the core component of transformers (Vaswani et al., 2017). Several current studies utilized supervision of the attention probabilities to effectively train the model (Liu et al., 2016; Mi et al., 2016). Moreover, it is also used for the model distillation to efficiently transfer the knowledge of a transformer model to another one via attention probabilities (Aguilar et al., 2020; Jiao et al., 2020; Sun et al., 2020; Wang et al., 2020, 2021). Inspired by this, we propose distilled attention labels, which are the supervision of attention probabilities optimized as a part of the distilled dataset, to enhance the effectiveness of the distilled dataset for training the transformer models.
In our experiments, we constructed distilled fewshot datasets to fine-tune BERT (Devlin et al., 2019) in various types of NLP tasks: AGNews (text classification), SST-2 (sentiment analysis), QNLI (QA/NLI), and MRPC (paraphrase identification).
Our main contributions are as follows: (i) To the best of our knowledge, this is the first work to explore dataset distillation for pre-trained transformers. Specifically, we demonstrate that our distilled datasets effectively fine-tune BERT even with only one sample for each class and only one gradient step. (ii) We present the distilled attention labels, which can easily be applied to dataset distillation for transformer architectures. Experimental results show that they consistently improved the performance with the distilled datasets in various types of NLP tasks. (iii) We open our source code and the distilled datasets obtained through our experiments to facilitate further research.1

heading: 2 Methodology

heading: 2.1 Dataset Distillation
text: In this section, we explain the basic approach of dataset distillation (Wang et al., 2018b), which aims to optimize a synthetic dataset through the gradient method similar to the current meta-learning approach (Finn et al., 2017).
Let the original training dataset D = {(xi, yi)}Ni=1, where (xi, yi) is a pair of an input and its class label. Our goal is to optimize a distilled dataset D̃ = {(x̃i, ỹi)}Mi=1, which is randomly initialized at first, with M ≪ N .
The model parameters θ are updated with a minibatch of the distilled dataset (x̃t, ỹt) by gradient
1https://github.com/arumaekawa/ dataset-distillation-with-attention-labels
descent (GD) steps as follows:
θt+1 = θt − η̃∇θtLtask s.t. Ltask = L(x̃t, ỹt, θt),
(1)
where L() is a twice-differentiable loss function and η̃ is the learnable learning rate of the model, which is optimized together with D̃. Given initial model parameters θ0, we can represent the model trained with the distilled dataset D̃, with the number of GD steps T , as
θT = F (θ0; D̃, η̃, T ), (2)
where F () is the training procedure of the T steps for the GD updating (Eq. 1).
As the goal of dataset distillation is that θT performs well on the original dataset, the optimization objective of the distilled dataset D̃ is calculated as follows:
Ldistill(D̃, η̃; θ0) := L(xt,yt, θT ) (3) = L(xt,yt, F (θ0; D̃, η̃, T )),
(4)
where (xt,yt) is a mini-batch of the original training dataset.
Therefore, the optimization problem for dataset distillation is formulated as
D̃∗, η̃∗ = arg min D̃,η̃
Eθ0∼p(θ0) [ Ldistill(D̃, η̃; θ0) ] ,
(5) where p(θ0) is the distribution of θ0.
We optimize the distilled dataset D̃ with this objective by using current gradient-based optimization techniques, e.g., Adam (Kingma and Ba, 2015). However, the discrete nature of text data makes it difficult to apply the gradient methods directly. Inspired by previous work (Sucholutsky and Schonlau, 2021; Li and Li, 2021), we use a sequence of embedding vectors for inputs of the distilled dataset instead of text as it is. Using the embeddings makes the loss Ldistill differentiable with respect to D̃, and we can thus optimize the distilled dataset D̃ by the gradient methods.

heading: 2.2 Distilled Soft Labels
text: The class labels of the original dataset are usually discrete hard labels (i.e., one-hot labels representing only a single class). Instead of hard labels, we can use soft labels for distilled datasets and optimize them with the input embeddings. Using soft
labels enables the distilled datasets to contain more information. Following previous work (Sucholutsky and Schonlau, 2021; Bohdal et al., 2020), we first initialize the soft labels with one-hot values and enable them to take any real values. We can now optimize the soft labels through the gradient method as well as the input embeddings.

heading: 2.3 Distilled Attention Labels
text: For efficient knowledge transfer to transformer models via training with the distilled dataset, we propose attention labels, which are optimized to guide the multi-head attention module of the transformer models.
Inspired by previous work (Aguilar et al., 2020; Wang et al., 2020, 2021), we compute the KullbackLeibler (KL) divergence DKL between the selfattention probabilities of the model a(θ) and the distilled attention labels ã across all layers and heads. The attention loss Lattn is computed as follows:
Lattn = 1
K
K∑
k=1
1
H
H∑
h=1
DKL(ãk,h||ak,h(θ)), (6)
where ãk,h and ak,h(θ) are the attention maps for the h-th head of the k-th layer of the distilled attention labels and the model, respectively, K is the number of layers, and H is the number of heads. Due to the data size, we consider the attention probabilities only for the first input token ([CLS]).
We train the model to minimize Ltask and Lattn at the same time. Thus, the GD updating of the model (Eq. 1) is modified as
θt+1 = θt − η̃∇θt(Ltask + λLattn), (7)
where λ is the balance weight for Lattn. The attention labels ã are first initialized randomly and restricted to being a valid probability distribution (i.e., non-negative and the sum equals 1) by applying the softmax function to real-valued vectors. We optimize the attention labels together with the input embeddings and the soft labels by the gradient method. The details of the step-by-step procedure of our distillation algorithm are shown in Appendix A.

heading: 3 Experiments

heading: 3.1 Settings
text: Datasets. We evaluated our dataset distillation methods in various types of NLP tasks. We used
a text classification task (AGNews (Zhang et al., 2015)) and three different natural language understanding tasks (SST-2, QNLI, and MRPC) from the GLUE benchmark (Wang et al., 2018a). For the evaluation metrics, we used accuracy for AGNews. For the other three tasks, we followed the evaluation settings of GLUE (Wang et al., 2018a). The statistics of each benchmark dataset are summarized in Table 1. Network Architecture. To evaluate the dataset distillation methods, we constructed distilled fewshot datasets to fine-tune BERT (Devlin et al., 2019), which is the first pre-trained transformer model, that all subsequent models are based on. We utilized the pre-trained BERTBASE model. Following the fine-tuning procedure in Devlin et al. (2019), we introduced additional classification layer weights W ∈ RC×D on the last hidden state of the [CLS] token, where D is the hidden dimension of BERT and C is the number of classes. Implementation. For all our distilled datasets, we used Adam optimizer (Kingma and Ba, 2015) with a learning rate α ∈ {1e−3, 1e−2, 1e−1} and trained the distilled datasets for 30 epochs. We initialized the learnable learning rate η̃ ∈ {1e−2, 1e−1}. For the attention labels, we set λ = 1.0, which performed well in our preliminary experiments. We report the results for the best performing combination of α and η̃. Note that due to the coarse granularity of the search, there is no need to care about overfitting to the test set. More details of our implementation are shown in Appendix B. Evaluation. To evaluate the distilled datasets, we fine-tuned the BERT model with them for 100 times, where the additional parameters W were randomly initialized each time. In all our experiments, we report the mean and standard deviation over the 100 evaluation results.

heading: 3.2 Results for 1-shot and 1-step Setting
text: We first evaluated the dataset distillation methods with a 1-shot and 1-step setting, where the distilled
dataset includes only one sample per class, and BERT was fine-tuned with it by only one GD step. We compared the performance for hard/soft labels and with/without attention labels for each task.
Table 2 shows the evaluation results. The distilled datasets with the hard labels, i.e., only optimizing the input embeddings and not applying the attention labels, still achieved 87.4, 81.6, and 68.6 for AGNews, SST-2, and QNLI, respectively, which is 92.4, 88.0, and 74.7% performance of the full dataset. Furthermore, using the soft labels further improved these performances, especially by almost 8 points for QNLI. However, for MRPC, the distilled dataset achieved only the same performance as the majority class baseline regardless of the use of the soft labels.
When applying the attention labels, the performance of the distilled dataset was significantly improved for all tasks, and their effect is much greater than the soft labels. Specifically, our distilled dataset with the attention labels yielded up to 98.5, 97.2, 94.1, and 88.9% performance of the full dataset for AGNews, SST-2, QNLI, and MRPC, respectively. These results indicate that using the attention labels enables to extract the information from the original dataset as the attention probabilities and to efficiently transfer it to the model.
When comparing the performance between the four tasks, dataset distillation performed very well on relatively simple classification tasks such as AGNews and SST-2, while the performance was somewhat limited on QNLI and MRPC, which require understanding the relationship between two sentences. In particular, for MRPC, although the performance was improved by applying the attention labels, the gap from the full dataset was still larger than that in the other three tasks. The class
imbalance in the original training dataset (68% positive) may make the training of the distilled dataset more difficult. We can say there is still room for performance improvement by dealing with this issue (e.g., by upsampling or downsampling).

heading: 3.3 Results for Multiple-shot and Multiple-step Setting
text: We also evaluated the distilled datasets with more than one shot and more than one GD step to finetune BERT. For the multiple-step setting, we considered two different scenarios: using the same distilled data in all steps and using different distilled data for each step. In these experiments, we evaluated the distilled datasets that use soft labels and attention labels for different numbers of GD steps T ∈ {1, 3, 5}.
Table 3 shows the results for the multiple-shot and multiple-step setting. In the single-step setting, overall performance improved with the number of shots of the distilled data. We believe that this is simply due to the expressiveness of the distilled data improved with the size of them. When using the same distilled data for all steps in the multiple-step setting, the performance of the distilled datasets degraded even compared with that in the single-step setting. In contrast, the performance was improved by separating the distilled data for each step and slightly but better than that with the same number of shots in the single-step setting. These results suggest that the role of the distilled data is different between the earlier and later steps, and it is difficult to obtain the distilled data that are generally useful for all GD steps.
In addition, the basic dataset distillation algorithm we used requires computing the back propagation through all GD steps for the optimization of
the distilled dataset, which increases memory and computational costs linearly with T . Therefore, it was difficult to increase T to be larger than 5 in our experiments. This is the limitation of our dataset distillation method, and it needs further improvement to scale to more complex tasks or to train models from scratch.

heading: 4 Conclusion
text: In this paper, we explored dataset distillation in NLP tasks to fine-tune pre-trained transformers. We proposed attention labels, which are the supervision of attention probabilities distilled as a part of the distilled datasets. Experimental results across various tasks demonstrate that our distilled fewshot datasets achieved successful performances even with only one sample per class. Notably, the attention labels significantly improved the performance of the distilled datasets even for the tasks where dataset distillation is difficult without them.
Limitations
We think the following three points are the limitations of this work. (i) As mentioned in Section 3.3, the computational cost of our distillation approach increases linearly with the number of GD steps and the distilled data size. It is necessary to explore efficient distillation algorithms to scale our method to more complex tasks or full-scratch training in future work. (ii) To optimize the distilled dataset through the gradient method, we utilized word embedding vectors instead of directly optimizing the text as in the existing work. Therefore, the distilled dataset we obtained cannot be applied to models with different word embeddings, such as other pretrained models or full-scratch training. (iii) In our experiments, we evaluated our approach only on text classification tasks. However, our approach can also be applied to text generation tasks as well by applying the attention labels to all input tokens (not only [CLS]) and using vocabulary-wise soft labels. In future work, we should investigate its performance and explore more effective approaches.

heading: A Overview of Proposed Method
text: Algorithm 1 illustrates an overview of our distillation algorithm.
B Implementation details
In our experiments, we trained the distilled datasets using Adam optimizer (Kingma and Ba, 2015) with linear warmup and linear decay learning rate schedule and gradient clipping with 1.0. Following the implementation in Wang et al. (2018b), we disabled dropout layers to avoid the randomness of the model training. We used a RTX 3090 or a RTX A6000, depending on the required memory size for each experiments. To obtain the performance of the full dataset for AGNews, which is used as the upper-bound of the distilled datasets, we fine-tuned BERTBASE model with learning rate η = 1e−5 for epochs ∈ {2, 3, 4}, and adopted the best performance. More information about our implementation can be found in our source code1.
ACL 2023 Responsible NLP Checklist

heading: A For every submission:

heading: A1. Did you describe the limitations of your work?
text: Left blank.

heading: A2. Did you discuss any potential risks of your work?
text: Left blank.

heading: A3. Do the abstract and introduction summarize the paper’s main claims?
text: Left blank.

heading: A4. Have you used AI writing assistants when working on this paper?
text: Left blank.

heading: B Did you use or create scientific artifacts?
text: Left blank.

heading: B1. Did you cite the creators of artifacts you used?
text: Left blank.

heading: B2. Did you discuss the license or terms for use and / or distribution of any artifacts?
text: Left blank.

heading: B3. Did you discuss if your use of existing artifact(s) was consistent with their intended use, provided
text: that it was specified? For the artifacts you create, do you specify intended use and whether that is compatible with the original access conditions (in particular, derivatives of data accessed for research purposes should not be used outside of research contexts)? Left blank.

heading: B4. Did you discuss the steps taken to check whether the data that was collected / used contains any
text: information that names or uniquely identifies individual people or offensive content, and the steps taken to protect / anonymize it? Left blank.

heading: B5. Did you provide documentation of the artifacts, e.g., coverage of domains, languages, and
text: linguistic phenomena, demographic groups represented, etc.? Left blank.
B6. Did you report relevant statistics like the number of examples, details of train / test / dev splits, etc. for the data that you used / created? Even for commonly-used benchmark datasets, include the number of examples in train / validation / test splits, as these provide necessary context for a reader to understand experimental results. For example, small differences in accuracy on large test sets may be significant, while on small test sets they may not be. Left blank.

heading: C Did you run computational experiments?
text: Left blank.

heading: C1. Did you report the number of parameters in the models used, the total computational budget
text: (e.g., GPU hours), and computing infrastructure used? Left blank.
The Responsible NLP Checklist used at ACL 2023 is adopted from NAACL 2022, with the addition of a question on AI writing assistance.

heading: C2. Did you discuss the experimental setup, including hyperparameter search and best-found
text: hyperparameter values? Left blank.

heading: C3. Did you report descriptive statistics about your results (e.g., error bars around results, summary
text: statistics from sets of experiments), and is it transparent whether you are reporting the max, mean, etc. or just a single run? Left blank.
C4. If you used existing packages (e.g., for preprocessing, for normalization, or for evaluation), did you report the implementation, model, and parameter settings used (e.g., NLTK, Spacy, ROUGE, etc.)? Left blank.

heading: D Did you use human annotators (e.g., crowdworkers) or research with human participants?
text: Left blank.

heading: D1. Did you report the full text of instructions given to participants, including e.g., screenshots,
text: disclaimers of any risks to participants or annotators, etc.? Left blank.

heading: D2. Did you report information about how you recruited (e.g., crowdsourcing platform, students)
text: and paid participants, and discuss if such payment is adequate given the participants’ demographic (e.g., country of residence)? Left blank.

heading: D3. Did you discuss whether and how consent was obtained from people whose data you’re using/curating? For example, if you collected data via crowdsourcing, did your instructions to
text: crowdworkers explain how the data would be used? Left blank.

heading: D4. Was the data collection protocol approved (or determined exempt) by an ethics review board?
text: Left blank.

heading: D5. Did you report the basic demographic and geographic characteristics of the annotator population that is the source of the data?
text: Left blank.

heading: Future Work
text: It is necessary to explore efficient distillation algorithms to scale our method to more complex tasks or full-scratch training in future work. (ii) To optimize the distilled dataset through the gradient method, we utilized word embedding vectors instead of directly optimizing the text as in the existing work. Therefore, the distilled dataset we obtained cannot be applied to models with different word embeddings, such as other pretrained models or full-scratch training. (iii) In our experiments, we evaluated our approach only on text classification tasks. However, our approach can also be applied to text generation tasks as well by applying the attention labels to all input tokens (not only [CLS]) and using vocabulary-wise soft labels. In future work, we should investigate its performance and explore more effective approaches