Dataset Distillation by Matching Training Trajectories (2203.11932v1)

Abstract: Dataset distillation is the task of synthesizing a small dataset such that a model trained on the synthetic set will match the test accuracy of the model trained on the full dataset. In this paper, we propose a new formulation that optimizes our distilled data to guide networks to a similar state as those trained on real data across many training steps. Given a network, we train it for several iterations on our distilled data and optimize the distilled data with respect to the distance between the synthetically trained parameters and the parameters trained on real data. To efficiently obtain the initial and target network parameters for large-scale datasets, we pre-compute and store training trajectories of expert networks trained on the real dataset. Our method handily outperforms existing methods and also allows us to distill higher-resolution visual data.

• The paper's main contribution is a novel technique that matches synthetic data-induced parameter trajectories with expert trajectories to replicate full-data performance.
• The methodology leverages backpropagation through multiple training updates to align long-range dynamics, markedly improving accuracy on CIFAR-10, CIFAR-100, and Tiny ImageNet.
• Empirical results demonstrate significant accuracy gains, establishing state-of-the-art performance and enabling the distillation of high-resolution synthetic datasets.

Dataset Distillation by Matching Training Trajectories

Abstract

The paper "Dataset Distillation by Matching Training Trajectories" (2203.11932) introduces a novel approach to dataset distillation, which is the task of creating a small synthetic dataset that allows a model trained on this dataset to achieve comparable test accuracy to a model trained on the full dataset. This method leverages long-range training dynamics by matching synthetic data induced parameter trajectories with expert trajectories derived from networks trained on real data. The approach significantly improves upon previous methods by distilling higher-resolution visual data, with substantial empirical results demonstrating superior performance across various datasets including CIFAR-10, CIFAR-100, and Tiny ImageNet.

Introduction

Dataset distillation aims to compress a large dataset into a smaller set of high-information synthetic images while preserving task-specific features necessary for model generalization. Unlike model distillation, which focuses on compressing model complexity, dataset distillation compresses the training data itself. This paper proposes optimizing synthetic data to emulate long-range training characteristics of real data by matching training parameter trajectories (Figure 1).

Figure 1: Dataset distillation aims to generate a small synthetic dataset for which a model trained on it can achieve a similar test performance as a model trained on the whole real train set.

The method employs pre-computed expert trajectories recorded from networks trained on the full dataset, using these as a gold standard to guide the distillation process. This long-range trajectory matching addresses challenges such as optimization difficulty and accumulation of error when matching only short-range behavior.

Methodology

The distillation method is centered around matching the synthetic dataset's influence on model training dynamics with that of pre-recorded expert trajectories derived from networks trained on real data (Figure 2). This involves initializing a model from expert parameters at a random time step and training it on the synthetic dataset, followed by penalizing deviations from the expert trajectory.

Figure 2: We perform long-range parameter matching between training on distilled synthetic data and training on real data. Starting from the same initial parameters, we train distilled data $\mathcal{D}_\mathsf{syn}$ such that $N$ training steps on them match the same result (in parameter space) from much more $M$ steps on real data.

The key aspect of the methodology lies in backpropagating through multiple training updates on the synthetic dataset, aligning the final synthetic-trained parameter proximity to the parameter location attained by the expert trajectory after several steps on real data. This trajectory-based distillation approach contrasts with short-range approaches, displaying improvements due to its emphasis on long-term learning features.

Experimental Results

The experiments conducted on CIFAR-10, CIFAR-100, and Tiny ImageNet demonstrate compelling results, with substantial accuracy improvements over existing methods:

• CIFAR-10: Achieved 46.3% accuracy with a single image per class, compared to the previous state-of-the-art of 28.8% [dsa], and 65.3% accuracy with 10 images per class.
• CIFAR-100: Enhanced performance to 24.3% accuracy with one image per class, increasing from prior results of 13.9% [dsa].
• Tiny ImageNet: The method improved accuracy to 8.8% with one image per class, outperforming concurrent work DM [dm].

Furthermore, the approach scaled to distilling higher-resolution $128 \times 128$ ImageNet images for the first time, exploring various subsets such as ImageNette and ImageWoof with substantial classification performance gains.

Implications and Future Work

The method's ability to accurately distill datasets and generalize over a long trajectory provides deeper insights into model training dynamics and dataset composition. While promising, challenges remain regarding memory consumption and computational costs associated with expert trajectory training. Future work could explore adaptive selection of expert trajectory segments, further optimizing distillation processes for different model architectures and datasets, potentially facilitating practical applications in privacy-preserving ML and efficient neural architecture search.

Conclusion

This paper presents a dataset distillation framework leveraging long-range trajectory matching, markedly enhancing distillation efficacy and dataset understanding. Through comprehensive experiments, authors illustrate the feasibility of synthesizing compact, high-resolution representative datasets, thereby opening avenues for practical real-world implementation and advancing theoretical discourse in dataset distillation methodologies.