D i LM : Distilling Dataset into Language Model for Text-level Dataset Distillation (2024)

Abstract: Dataset distillation aims to compress a training dataset by creating a small number of informative synthetic samples such that neural networks trained on them perform as well as those trained on the original training dataset. Current text dataset distillation methods create each synthetic sample as a sequence of word embeddings instead of a text to apply gradient-based optimization; however, such embedding-level distilled datasets cannot be used for training other models whose word embedding weights are different from the model used for distillation. To address this issue, we propose a novel text dataset distillation approach, called Distilling dataset into Language Model (DiLM), which trains a language model to generate informative synthetic training samples as text data, instead of directly optimizing synthetic samples. We evaluated DiLM on various text classification datasets and showed that distilled synthetic datasets from DiLM outperform those from current coreset selection methods. DiLM achieved remarkable generalization performance in training different types of models and in-context learning of large language models. Our code will be available at https://github.com/arumaekawa/DiLM.

## introduction the successful advancements in machine learning in a wide range of fields are due to the scaling-up of deep neural networks and large training datasets. in the natural language processing (nlp) field, large language models (llms), which are pre-trained with a huge amount of text, such as bert-and gpt-family models @xcite @xcite , have shown remarkable capabilities for various nlp tasks. however, training such large-scale models requires large computational resources and a long time, which makes it difficult to develop new llms, and even to fine-tune them. to address this issue, dataset distillation @xcite has attracted much attention in the machine learning community, which aims to reduce training costs by compressing training datasets. in contrast to traditional coreset selection approaches @xcite @xcite , which heuristically select a small subset of representative training samples from the original dataset, dataset distillation creates more informative synthetic samples by distilling the knowledge from the original dataset. with this approach, synthetic samples are optimized with gradient descent according to objective functions for dataset distillation, including meta-learning @xcite , gradient matching @xcite , training trajectory matching @xcite , and feature distribution matching @xcite . the recent remarkable performance of dataset distillation, especially in the computer vision (cv) field, has also led to studies of its various applications, including neural architecture search @xcite , federated learning @xcite , continual learning @xcite , and privacy preservation @xcite . while most previous studies applied dataset distillation only to image classification datasets, some studies focused on text dataset distillation @xcite @xcite . in contrast to the image, which can be applied gradient-based optimization by considering it as a pixel-wise continuous data, the discrete nature of text makes dataset distillation challenging @xcite . to address this issue, all existing text dataset distillation methods used the widely used neural nlp technique called embedding, i.e., optimizing a synthetic dataset as continuous input word embeddings instead of discrete text. however, such embedding-level distilled synthetic datasets cannot be used for training other models that have different word embedding weights, which is a crucial issue in terms of practical applications. furthermore, distilled word embedding sequences are also completely unreadable to humans, which makes it difficult to interpret and analyze the original training dataset by observing distilled synthetic samples. motivated by these shortcomings, this paper explores the text dataset distillation to obtain distilled synthetic datasets at the text-level as the first study. we propose the first text-level dataset distillation approach called "distilling dataset into language model (dilm)". to overcome the optimization difficulty of discrete text, dilm uses a language model as a surrogate continuous optimization target instead of directly optimizing a synthetic sample's text. specifically, dilm trains a language model to minimize the gradient matching loss @xcite of generated synthetic samples as a dataset distillation objective. to enable back-propagating the gradient matching loss to the language model, we design a differentiable backward pass via loss weighting with generation probabilities to bypass the non-differentiable generated text (figure 1 ). in our experiments, we applied dilm to distill three text classification datasets from the glue benchmark @xcite , sst-2, qqp, and mnli-m. the results indicate that the synthetic datasets distilled with dilm outperformed representative real samples selected from the original datasets with current coreset selection methods. our distilled datasets also achieved remarkable generalization performance not only for training different types of pre-trained models but also for in-context learning of llms as few-shot prompts. our main contributions are as follows: ## methodology in this section, we introduce dilm, which distills text datasets into text data, not word embeddings, for the model-agnostic applicability and interpretability of the distilled synthetic datasets. the main idea of dilm is to avoid the optimization difficulties of discrete text by instead training continuous parameters of a language model as a surrogate optimization target of dataset distillation. ## experimental settings datasets. we evaluated dilm in distilling three major text classification datasets, sst-2, qqp, and mnli-m, from the glue benchmark @xcite . following wang et al. (2018a), we report accuracy for sst-2 and mnli-m, and the average of accuracy and f1 score for qqp as our results. more details about each dataset are shown in appendix b. baselines. following previous studies on dataset distillation in the cv field, we compared the performance of dilm with three coreset selection methods, random, k-centers @xcite , and herding @xcite , as well as tdd @xcite , which is a recent embedding-level distillation method. note that tdd also trains the learnable soft-labels and learning rates for each training step together with the input word embeddings. we also evaluated the vanilla lm, which skips training for gradient matching (in section 3.3), to validate its effectiveness. note that we applied k-center representative sample selection (in section 3.4) to the vanilla lm as well. the details of each baseline are given in appendix c. ## conclusion we proposed the first text-level dataset distillation approach, called dilm, which trains a language model to generate informative synthetic samples as text data for model-agnostic applicability and interpretability of distilled datasets. experimental results across various text classification datasets indicated that the distilled datasets from dilm achieve successful performance for training various types of models beyond the source model used for distillation, even for in-context learning of llms.