\documentclass[english]{jnlp_3.8}
\usepackage{jnlpbbl_2.2}
\usepackage[dvipdfmx]{graphicx}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{url}
\usepackage{array}
\allowdisplaybreaks

\usepackage{hangcaption_jnlp}

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% author's definition

%%%%%\usepackage[utf8]{inputenc}

\usepackage{bm}
%%%%%\usepackage{amsfonts}
\usepackage{mathtools}
\usepackage{booktabs}
%%%%%\usepackage{adjustbox}
%%%%%\usepackage{multirow}
%%%%%\usepackage{afterpage,array,rotating}
\usepackage{algorithm}
\usepackage{algorithmic}

\newcommand{\argmin}{\mathop{\rm arg~min}\limits}

\newcommand{\citep}[1]{\cite{#1}}
\newcommand{\citet}[1]{\citeA{#1}}

%%%%%\renewcommand{\algorithmicrequire}{\textbf{Input:}}
%%%%%\renewcommand{\algorithmicensure}{\textbf{Output:}}

\DeclareRobustCommand*{\rciteA}{\citeA}
%%%%%\DeclareRobustCommand*{\rcite}{\cite}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

\Volume{32}
\Number{1}
\Month{March}
\Year{2025}
%%%%\Issuetitle{「　　　　」特集号}

\received{2024}{7}{30}
\revised{2024}{11}{8}
\accepted{2024}{12}{16}



\setcounter{page}{283}



%%%%   Title, Author, etc.            %%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\esubject{General Paper}
\etitle{Dataset Distillation with Attention Labels \\for Fine-tuning BERT}
\eauthor{Aru Maekawa\affiref{TokyoTech} \and Naoki Kobayashi\affiref{LegalOn} \and Kotaro Funakoshi\affiref{TokyoTech} \and Manabu Okumura\affiref{TokyoTech}}
\eabstract{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 study, we focus on constructing distilled few-shot datasets for natural language processing (NLP) tasks to fine-tune pre-trained transformers. Specifically, we propose introducing attention labels, which can efficiently distill knowledge from the original dataset and transfer it to transformer models via attention probabilities. We evaluated our dataset distillation methods in four NLP tasks and demonstrated that it is possible to create distilled few-shot datasets with attention labels, yielding an impressive performance for fine-tuning BERT. Specifically, in AGNews, which is a four-class news classification task, our distilled few-shot dataset achieved up to 93.2\% accuracy, which is 98.5\% that of the original dataset, even with only one sample per class and only one gradient step.
\footnotetext{This paper was published in the Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics \cite{maekawa-etal-2023-dataset}.}%
}
\ekeywords{Dataset Distillation, Dataset Condensation, Attention Labels}

\headauthor{Maekawa et al.}
\headtitle{Dataset Distillation with Attention Labels for Fine-tuning BERT}



%%%%     Affiliation     %%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\affilabel{TokyoTech}{}{Institute of Science Tokyo}
\affilabel{LegalOn}{}{LegalOn Technologies, Inc.}


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{document}

\maketitle



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S1
\section{Introduction}

Deep learning models have achieved state-of-the-art performance in various fields, including computer vision and natural language processing (NLP), using large-scale neural networks trained with large datasets.
However, their successful performance have come with massive training costs, including training time, GPU resources, and energy consumption.
To reduce training costs, the current research has focused on constructing a small training dataset such that models trained with it can achieve a comparable performance to models trained with the entire original dataset.

One classical method for compressing a training dataset is data selection.
Data selection methods select a subset of effective training samples based on several heuristic measures, such as cluster centers \citep{sener2018active}, diversity \citep{aljundi2019gradient}, and the likelihood of models \citep{moore-lewis-2010-intelligent}.
Although data selection methods are effective for efficient model training and several applications, such as active learning \citep{sener2018active} and continual learning \citep{aljundi2019gradient}, 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 to reducing the training dataset, \citet{wang2018dataset} proposed \emph{dataset distillation}, which aims to create few synthetic samples optimized to effectively train models.
Dataset distillation has garnered significant attention in machine learning \citep{wang2018dataset,zhao2021dataset,pmlr-v139-zhao21a,sucholutsky-2021-soft-label,bohdal-2020-flexible-dataset,wang-2022-cafe,Cazenavette_2022_CVPR} for both the theoretical interest and various applications, such as neural architecture/hyper-parameter search \citep{pmlr-v119-such20a}, continual learning \citep{Masarczyk_2020_CVPR_Workshops,rosasco-2022-distilled-replay}, federated learning \citep{goetz-2020-federated,zhou-2020-distilled-one-shot}, and preserving data privacy \citep{li2020soft,pmlr-v162-dong22c}.

However, most of the existing studies on dataset distillation focus on image datasets, and only a few studies involve NLP tasks.
\citet{sucholutsky-2021-soft-label} and \citet{li-2021-data-distillation} extended dataset distillation to text datasets using embedding vectors as an input of the distilled dataset instead of discrete text.
Although these studies applied dataset distillation to model architectures based on convolutional and recurrent neural networks, to the best of our knowledge, no research exists that addressed dataset distillation for pre-trained transformers, such as BERT \citep{devlin-etal-2019-bert}, which have become the de facto standard for various types of NLP tasks.
Therefore, in this study, we aim to obtain distilled few-shot datasets to fine-tune pre-trained transformers for NLP tasks.

To this end, we focus on the attention mechanism, which is a core component of transformers \citep{NIPS2017_3f5ee243}.
Several recent studies have utilized the supervision of attention probabilities to effectively train the model \citep{liu-etal-2016-neural,mi-etal-2016-supervised}.
Moreover, it is used for model distillation to efficiently transfer the knowledge of a transformer model to another via attention probabilities \citep{Aguilar_Ling_Zhang_Yao_Fan_Guo_2020,jiao-etal-2020-tinybert,sun-etal-2020-mobilebert,wang-etal-2020-minilm,wang-etal-2021-minilmv2}.
Inspired by this, we propose distilled attention labels, which are optimized as part of the distilled dataset to supervise attention probabilities, to enhance the effectiveness of the distilled dataset for training transformer models.

In our experiments, we constructed distilled few-shot datasets to fine-tune BERT \citep{devlin-etal-2019-bert} for various types of NLP tasks: AGNews (text classification), SST-2 (sentiment analysis), QNLI (QA/NLI), and MRPC (paraphrase identification).

The main contributions of this study are as follows:
%%%%
\begin{itemize}
\item To the best of our knowledge, this is the first study 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.
\item We present the distilled attention labels, which can easily be applied to dataset distillation for transformer architectures. Experimental results demonstrate that they consistently improved the performance with the distilled datasets in various types of NLP tasks.
\item We open our source code and the distilled datasets obtained through our experiments to facilitate further research.\footnote{\url{https://github.com/arumaekawa/dataset-distillation-with-attention-labels}}
\end{itemize}
%%%%%



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S2
\section{Methodology}



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S2.1
\subsection{Overview}

We presented a novel and effective dataset distillation method for training transformer models for NLP tasks.
The proposed approach extends existing dataset distillation techniques based on the meta-learning algorithm (Section \ref{sec:method_dd}) and soft label distillation (Section \ref{sec:method_sl}) by introducing attention label distillation, which focuses on the attention mechanism of the transformer (Section~\ref{sec:method_al}).
A step-by-step overview of our distillation algorithm is presented in Algorithm \ref{alg1}, and detailed explanations are provided in the following sections.



%%%%%%%%%%%%%%%%%%%%%%%%%%%A1
\begin{algorithm}[t]
    \caption{~Dataset Distillation with attention labels}
    \label{alg1}
\input{10algo1.tex}
\end{algorithm}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%




%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S2.2
\subsection{Dataset Distillation}
\label{sec:method_dd}
In this section, we explain the basic approach for dataset distillation \citep{wang2018dataset}, which aims to optimize a synthetic dataset using a gradient method similar to the current meta-learning approach \citep{finn2017maml}.

Let the original training dataset $\mathcal{D}=\left\{(x_i, y_i)\right\}_{i=1}^N$, where $(x_i, y_i)$ denotes a pair of an input and its class label, and the distilled dataset $\tilde{\mathcal{D}}=\left\{(\tilde{x}_i, \tilde{y}_i)\right\}_{i=1}^M$, which is randomly initialized at first, with $M\ll N$.
The goal of dataset distillation is to optimize $\tilde{\mathcal{D}}$ to train models with performance similar to that of the model trained on $\tilde{\mathcal{D}}$.

The model parameters $\theta$ are updated with a mini-batch of the distilled dataset $(\bm{\tilde{x}}_t, \bm{\tilde{y}}_t)$ using gradient descent (GD) steps as follows:
%%%%%
\begin{equation}
    \begin{split}
        \theta_{t+1} = \theta_t - \tilde{\eta} \nabla_{\theta_t} \mathcal{L}_{task}\\
        \text{s.t.} \quad \mathcal{L}_{task} = L(\bm{\tilde{x}}_t, \bm{\tilde{y}}_t, \theta_t) \label{equation:gd_step},
    \end{split}
\end{equation}
%%%%%
where $L()$ denotes a twice-differentiable loss function and $\tilde{\eta}$ denotes the trainable learning rate of the model, which is optimized together with $\tilde{\mathcal{D}}$.
Given the initial model parameters $\theta_0$, we can represent the model trained with the distilled dataset $\tilde{\mathcal{D}}$, with the number of GD steps $T$, as follows:
%%%%%
\begin{equation}
    \theta_T = F(\theta_0; \tilde{\mathcal{D}}, \tilde{\eta}, T) \label{equation:training_procedure},
\end{equation}
%%%%%
where $F()$ denotes the training procedure for the $T$ steps of the GD updating (Eq. \ref{equation:gd_step}).

Because the goal of dataset distillation is for $\theta_T$ to perform well on the original dataset, the optimization objective of the distilled dataset $\tilde{\mathcal{D}}$ is calculated as follows:
\begin{align}
    \mathcal{L}_{distill}(\tilde{\mathcal{D}}, \tilde{\eta};\theta_0) &\coloneqq L(\bm{x}_t, \bm{y}_t,\theta_T) \label{equation:objective_of_dd_1}\\
    &= L(\bm{x}_t, \bm{y}_t, F(\theta_0; \tilde{\mathcal{D}}, \tilde{\eta}, T)) \label{equation:objective_dd_2},
\end{align}
where $(\bm{x}_t, \bm{y}_t)$ is a minibatch of the original training dataset.

Therefore, the optimization problem for dataset distillation is formulated as
%%%%%
\begin{equation}
    \tilde{\mathcal{D}}^*, \tilde{\eta}^* = \argmin_{\tilde{\mathcal{D}}, \tilde{\eta}} \mathbb{E}_{\theta_0\sim p(\theta_0)} \left[\mathcal{L}_{distill}(\tilde{\mathcal{D}}, \tilde{\eta};\theta_0)\right] \label{equation:objective_of_dd},
\end{equation}
%%%%%
where $p(\theta_0)$ denotes the distribution of $\theta_0$.

We optimize the distilled dataset $\tilde{\mathcal{D}}$ with this objective using current gradient-based optimization techniques, for instance, Adam \citep{kingma-etal-2015-adam}.
However, the discrete nature of text data makes it difficult to directly apply gradient methods.
Inspired by previous work \citep{sucholutsky-2021-soft-label,li-2021-data-distillation}, we use a sequence of embedding vectors for inputs of the distilled dataset instead of text, as is.
Using embeddings makes the loss $\mathcal{L}_{distill}$ differentiable with respect to $\tilde{\mathcal{D}}$, and we can thus optimize the distilled dataset $\tilde{\mathcal{D}}$ using gradient methods.



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S2.3
\subsection{Distilled Soft Labels}
\label{sec:method_sl}
The class labels of the original dataset are typically discrete hard labels (i.e., one-hot labels representing only a single class).
Instead of hard labels, soft labels can be used for the distilled datasets and optimized with input embeddings.
The use of soft labels enables the distilled datasets to contain more information.
Following previous work \citep{sucholutsky-2021-soft-label,bohdal-2020-flexible-dataset}, we initialize the soft labels with one-hot values and enable them to take real values.
We can now optimize soft labels together with input embeddings using the gradient method.



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S2.4
\subsection{Distilled Attention Labels}
\label{sec:method_al}
For efficient knowledge transfer to the transformer models via training with the distilled dataset, we propose attention labels that are optimized to guide the multi-head attention module of the transformer models.

Inspired by previous work \citep{Aguilar_Ling_Zhang_Yao_Fan_Guo_2020,wang-etal-2020-minilm,wang-etal-2021-minilmv2}, we compute the Kullback--Leibler (KL) divergence $D_\text{KL}$ between the self-attention probabilities of model $a(\theta)$ and the distilled attention labels $\tilde{a}$ assigned to each example in the distilled dataset.
Letting $\tilde{a}_{k,h}$ and $a_{k,h}(\theta)$ represent the attention maps for the $h$-th head of the $k$-th layer of the distilled attention labels and the model, respectively, the attention loss $\mathcal{L}_{attn}$ is computed across all layers and heads as follows:
%%%%%
\begin{equation}
    \mathcal{L}_{attn} = \frac{1}{K}\sum_{k=1}^K \frac{1}{H}\sum_{h=1}^H D_\text{KL}(\tilde{a}_{k,h}||a_{k,h}(\theta)) \label{equation:atten_loss},
\end{equation}
%%%%%
where $K$ denotes the number of layers and $H$ denotes the number of heads. 
Considering the data size, we use the attention probabilities only for the first input token (\verb|[CLS]|).

We train the model to minimize $\mathcal{L}_{task}$ and $\mathcal{L}_{attn}$ simultaneously.
Therefore, the GD update of the model (Eq. \ref{equation:gd_step}) is modified as follows:
%%%%%
\begin{equation}
    \theta_{t+1} = \theta_t -\tilde{\eta}\nabla_{\theta_t} (\mathcal{L}_{task}+\lambda \mathcal{L}_{attn}) \label{equation:gd_step_with_al},
\end{equation}
%%%%%
where $\lambda$ denotes the balance weight of $\mathcal{L}_{attn}$.

The attention labels $\tilde{a}$ are first initialized randomly and restricted to 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 soft labels using the gradient method and obtain a distilled dataset $\tilde{\mathcal{D}}=\left\{(\tilde{x}_i, \tilde{y}_i, \tilde{a}_i)\right\}_{i=1}^M$.



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S3
\section{Experiments}



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S3.1
\subsection{Settings}

%%%%%
\noindent\textbf{Datasets.}
We evaluated our dataset distillation method for various NLP tasks.
We used a text classification task (AGNews \citep{DBLP:conf/nips/ZhangZL15}) and three natural language understanding tasks (SST-2, QNLI, and MRPC) from the GLUE benchmark \citep{wang-etal-2018-glue}.
AGNews is a news classification dataset containing four categories: World, Sports, Business, and Sci/Tech.
SST-2 is a dataset for sentiment analysis, and each sentence is labeled as either positive or negative.
QNLI is a question-answering dataset designed to determine whether a sentence provides relevant information for answering a given question.
MRPC is a paraphrase detection dataset that comprises sentence pairs labeled as either paraphrase or not.
For the evaluation metrics, we used the accuracy of AGNews. For the other three tasks, we followed the GLUE evaluation settings  \citep{wang-etal-2018-glue}.
The statistics of each benchmark dataset are summarized in Table \ref{tab:datasets}.



%%%%%%%%%%%%%%%%%%%%%%%%%%%T1
\begin{table}[b]
\input{10table01.tex}
    \hangcaption{Summary of evaluation datasets. $C$ represents the number of classes in each task. For the three GLUE tasks, since the test set is not available, we report the evaluation results on the development set.}
    \label{tab:datasets}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%



%%%%%
\noindent\textbf{Network Architecture.}
To evaluate the dataset distillation methods, we constructed distilled few-shot datasets to fine-tune BERT \citep{devlin-etal-2019-bert}, which is the first pre-trained transformer model on which all subsequent models are based.
We utilized a pre-trained BERT$_\text{BASE}$ model.\footnote{\url{https://huggingface.co/bert-base-uncased}}
Following the fine-tuning procedure in \citet{devlin-etal-2019-bert}, we introduced additional classification layer weights $W\in\mathbb{R}^{C\times D}$ on the last hidden state of the \verb|[CLS]| token, where $D$ denotes the hidden dimension of BERT and $C$ denotes the number of classes.


%%%%
\noindent\textbf{Implementation.}
For all distilled datasets, we used the Adam optimizer \citep{kingma-etal-2015-adam} with a learning rate $\alpha \in \{1e^{-3}, 1e^{-2}, 1e^{-1}\}$ and trained the distilled datasets for $30$ epochs.
We initialized the learnable learning rate $\tilde{\eta} \in \{1e^{-2}, 1e^{-1}\}$.
The input embedding sequence length was set to 512 tokens, corresponding to the maximum input length of BERT$_\text{BASE}$.
For the attention labels, we set $\lambda = 1.0$, which performed well in the preliminary experiments.
We report the results for the best performing combination of $\alpha$ and $\tilde{\eta}$.
Further details on our implementation are provided in Appendix \ref{appendix:implementaiton}.


%%%%%
\noindent\textbf{Evaluation.}
To evaluate the distilled datasets, we fine-tuned the BERT model 100 times, where additional parameters $W$ were randomly initialized each time.
In all our experiments, we reported the mean and standard deviation of 100 evaluation results.



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S3.2
\subsection{Results for One-shot and One-step Setting}
\label{sec:single_step}

First, we evaluated the dataset distillation methods with a one-shot and one-step setting, where the distilled dataset includes only one sample per class, and the labels of the distilled dataset are initialized with one-hot values for class,\footnote{Note that when using soft labels, distilled datasets are not necessarily one sample per class as they can converge to any real-valued vectors.} and BERT was fine-tuned with it only one GD step.
We compared the performance of hard/soft labels with and without attention labels for each task.
For the baselines of the data selection methods, we provide the performance of random selection and K-centers \citep{sener2018active}, which use the centroid examples of each class.

Table \ref{tab:single-step} summarizes the evaluation results.
The distilled datasets with hard labels, that is, 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\% of the performance of the full dataset, and consistently outperformed the data selection baselines.
Furthermore, the use of soft labels further improved these performances, particularly by approximately eight points for QNLI.
However, for MRPC, the distilled dataset achieved the same performance as the majority class baseline, regardless of the use of soft labels.


%%%%%%%%%%%%%%%%%%%%%%T2
\begin{table}
\input{10table02.tex}
    \hangcaption{Experimental results for the 1-shot and 1-step setting. `Majority' is the majority class baseline. 
    Scores for the full dataset with * are cited from \rciteA{devlin-etal-2019-bert}.
    Bold scores show the best results for each task.}
    \label{tab:single-step}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%



When attention labels were applied, the performance of the distilled dataset was significantly improved for all tasks, and their effect was much greater than that of soft labels.
Specifically, our distilled dataset with attention labels yielded up to 98.5, 97.2, 94.1, and 88.9\% of the performances of the full dataset for AGNews, SST-2, QNLI, and MRPC, respectively.
These results indicate that using attention labels enables the extraction of information from the original dataset as attention probabilities and the efficient transfer of this information to the model.

When comparing the performances of the four tasks, dataset distillation performed very well on relatively simple classification tasks such as AGNews and SST-2, whereas the performance was somewhat limited on QNLI and MRPC, which require an understanding of the relationship between two sentences.
In particular, for MRPC, although the performance was improved by applying 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 training the distilled dataset more difficult.
We can assume that room for performance improvement still exists by addressing this issue (e.g., upsampling or downsampling).

The stability of the dataset distillation process was also investigated.
We generated five distilled datasets using both soft and attention labels with five different random seeds, and trained the BERT model 50 times on each of them.
Table \ref{tab:stability} lists the average, best, and worst performances for the five datasets.
These results demonstrate the high stability of the distillation process. Notably, the worst-performing distilled dataset across the five runs consistently outperformed the baseline without attention labels, as summarized in Table \ref{tab:single-step}.



%%%%%%%%%%%%%%%%%%%%%%%T3
\begin{table}[t]
\input{10table03.tex}
    \hangcaption{Experimental results for the stability of the dataset distillation process across the five runs with different random seeds.}
    \label{tab:stability}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%





%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S3.3
\subsection{Results for Multiple-shot and Multiple-step Setting}
\label{sec:multi_step}

We also evaluated the distilled datasets with more than one shot and more than one GD step to fine-tune 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 distilled datasets that used soft labels and attention labels for different numbers of GD steps $T\in\{1, 3, 5\}$.


%%%%%%%%%%%%%%%%%%%%%%%%%T4
\begin{table}[t]
\input{10table04.tex}
    \hangcaption{Experimental results for the multiple-shot and multiple-step setting. Bold scores show the best results for each task.}
    \label{tab:multi-step}
\end{table}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%



Table \ref{tab:multi-step} summarizes the results for multiple-shot and multiple-step settings.
In the single-step setting, the overall performance improved with an increase in the number of shots of the distilled data. We believe that this is simply owing to the improvement in the size of the distilled data.
When using the same distilled data for all steps in the multiple-step setting, the performance of the distilled datasets was degraded, even when compared with that in the single-step setting.
In contrast, the performance was improved by separating the distilled data for each step and was slightly better than that with the same number of shots in the single-step setting.
These results suggest that the role of distilled data differs between the earlier and later steps, and obtaining distilled data that are generally useful for all GD steps is difficult.


In addition, the basic dataset distillation algorithm we used requires computing backpropagation through all GD steps for the optimization of the distilled dataset, which increases the memory and computational costs linearly with $T$.
Therefore, increasing $T$ to be greater than $5$ was difficult in our experiments.
This is a limitation of our dataset distillation method, which requires further improvement to scale to more complex tasks or train models from scratch.




%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S3.4
\subsection{Qualitative Analysis of Distilled Dataset}

For the qualitative analysis, we present examples of the distilled dataset generated in the experiments described in Section \ref{sec:single_step}. Because the inputs of the distilled dataset are sequences of continuous embedding vectors that are difficult to analyze, we converted the input embedding sequences into human-readable text by retrieving the nearest neighbor word of each embedding vector in the input embedding space of the BERT$_\text{BASE}$ model.

Table \ref{tab:example_agnews} lists the nearest neighbor words of the first 100 tokens of each input embedding sequence in the distilled dataset of AGNews.
First, several words related to each class label appeared in each example. For instance, ``landmark,'' ``jews,'' ``building,'' and ``church'' in the ``World'' category, and ``laboratory,'' ``syndrome,'' ``indonesian,'' and ``algorithms,'' in the ``Sci/Tech.'' category.
However, when considered as word sequences, they do not form grammatically coherent texts, and almost no interpretability is present.
Notably, even the first token, which corresponds to the \verb|[CLS]| token in the real data, is replaced with a different word in all examples.
This is a limitation of the current dataset distillation for text data when considering the generalizability across models and interpretability, which are important for practical applications.



%%%%%%%%%%%%%%%%%%%%%%%T5
\begin{table}
\input{10table05.tex}
    \hangcaption{Examples in the distilled dataset of AGNews generated in the experiments described in Section \ref{sec:single_step}. The ``Input embeddings'' column shows the nearest neighbor words of the first 100 tokens of each input embedding sequence in the distilled dataset.}
    \label{tab:example_agnews}
\end{table}
%%%%%%%%%%%%%%%%%%




%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S4
\section{Conclusion}

In this study, we explored dataset distillation in NLP tasks to fine-tune the pre-trained transformers.
We proposed attention labels, which supervise attention probabilities distilled as part of the distilled datasets.
The experimental results across various tasks demonstrated that our distilled few-shot datasets achieved a successful performance, even with only one sample per class.
Notably, attention labels significantly improved the performance of the distilled datasets, even for tasks in which dataset distillation was difficult without them.



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%S5
\section*{Limitations}

The limitations of this study are as follows: (i) As mentioned in Section \ref{sec:multi_step}, the computational cost of our distillation approach increases linearly with the number of GD steps and distilled data size. In future studies, exploring efficient distillation algorithms to scale the proposed method to more complex tasks or full-scratch training will be necessary. (ii) To optimize the distilled dataset using the gradient method, we utilized word embedding vectors rather than directly optimizing the text, as in existing studies. Therefore, the obtained distilled dataset cannot be applied to models with different word embeddings, such as other pre-trained models or full-scratch training. However, the proposed approach can be applied to any dataset distillation method for models that adopt attention mechanisms, including transformers. Therefore, we believe that it has high potential to contribute to performance improvements in distillation methods that may emerge as more broadly applicable in future work. (iii) In our experiments, we only evaluated the proposed approach for text classification tasks. However, the proposed approach can also be applied to text generation tasks by applying attention labels to all input tokens (not only [CLS]) and using vocabulary-wise soft labels. In future studies, we will investigate the performance and explore more effective approaches.



% Start bibliography
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\bibliographystyle{jnlpbbl_1.7}
\bibliography{10refs}



\appendix

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%SA
\section{Implementation Details}
\label{appendix:implementaiton}

In our experiments, we trained the distilled datasets using the Adam optimizer \citep{kingma-etal-2015-adam} with a linear warmup and linear decay learning rate schedule and a gradient clipping of 1.0.
Following the implementation in \citet{wang2018dataset}, we disabled dropout layers to avoid randomness during model training.
We used RTX 3090 or RTX A6000, depending on the required memory size for each experiment.
To obtain the performance of the full dataset for AGNews, which was used as the upper bound of the distilled datasets, we fine-tuned the BERT$_\text{BASE}$ model with a learning rate $\eta = 1e^{-5}$ for epochs $\in \{2, 3, 4\}$ and adopted the best performance.
Additional information regarding our implementation is provided in the source code (\url{https://github.com/arumaekawa/dataset-distillation-with-attention-labels}).



% Start biography
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\begin{biography}

\bioauthor[:]{Aru Maekawa}{%
Aru Maekawa received his B.E. and M.E. degrees from the Tokyo Institute of Technology in 2022 and 2024, respectively. He is currently a Ph.D. student at Institute of Science Tokyo. His Research interests include natural language processing and machine learning, particularly, dataset distillation.
}

\bioauthor[:]{Naoki Kobayashi}{%
Naoki Kobayashi received his Dr.\ Eng.\ degree from the Tokyo Institute of Technology, Japan. His research interests include natural language processing, particularly, discourse parsing.
}

\bioauthor[:]{Kotaro Funakoshi}{%
Kotaro Funakoshi received his Dr.\ Eng.\ degree from the Tokyo Institute of Technology, Japan, and is currently an associate professor at the Institute of Integrated Research at Institute of Science Tokyo. Formerly with Honda Research Institute Japan Co., Ltd., he served as a program-specific associate professor with the Cooperative Intelligence Joint Chair at Kyoto University since 2017. His research interests include natural language processing, particularly computational linguistics, multi-modal dialogue systems, and human--machine interaction.
}

\bioauthor[:]{Manabu Okumura}{%
Manabu Okumura was born in 1962. He earned his B.E., M.E., and Dr.\ Eng.\ from Tokyo Institute of Technology in 1984, 1986, and 1989, respectively. He served as an assistant at the Department of Computer Science, Tokyo Institute of Technology, from 1989 to 1992, and as an associate professor at the School of Information Science, Japan Advanced Institute of Science and Technology, from 1992 to 2000. Currently he is a professor at the Institute of Integrated Research at Institute of Science Tokyo, Japan. His research interests include natural language processing, particularly text summarization, computer-assisted language learning, sentiment analysis, and text data mining.
}
\end{biography}

\biodate   %%%% Following settings will be edited by editors.


\end{document}






