Dataset Distillation with Attention Labels for Fine-tuning BERT (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 few-shot 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 fine-tuning BERT. Specifically, in AGNews, a four-class news classification task, our distilled few-shot 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.

## introduction 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 @xcite , diversity @xcite , and likelihood of models @xcite . although the data selection methods effectively work for efficient model training and several applications, such as active learning @xcite and continual learning @xcite , 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, @xcite 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 @xcite @xcite @xcite @xcite for both the theoretical interest and various applications, such as neural architecture/hyper-parameter search @xcite , continual learning @xcite , federated learning @xcite , and preserving data privacy @xcite . however, most of the existing research on dataset distillation mainly focuses on image datasets, and only a few studies involve nlp tasks. @xcite and @xcite 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 @xcite , 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. to this end, we focus on the attention mechanism, which is the core component of transformers @xcite . several current studies utilized supervision of the attention probabilities to effectively train the model @xcite . moreover, it is also used for the model distillation to efficiently transfer the knowledge of a transformer model to another one via attention probabilities @xcite @xcite @xcite . 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 @xcite 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. foot_0 2 methodology ## conclusion 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.