PaperlayerView exact revision on arXiv

Source-linked AI summary

Dataset Distillation by Matching Training Trajectories

George Cazenavette, Tongzhou Wang, Antonio Torralba, Alexei A. Efros, Jun-Yan Zhu

arXiv:2203.11932v1cs.CVcs.AIcs.LG

TL;DR

Dataset distillation seeks to preserve the test performance of full-dataset training using a very small synthetic dataset. The paper matches long-range synthetic and real training trajectories using stored expert checkpoints, outperforming existing methods and extending distillation to higher-resolution images.

  • Problem

    Dataset distillation has mainly remained theoretical because prior methods often struggle with large-scale, higher-resolution datasets and can require costly or unstable optimization.

  • Method

    The method pre-computes expert trajectories from real-data training, initializes students from trajectory states, and optimizes synthetic data to match subsequent parameter trajectories.

  • Results

    The method outperforms existing dataset-distillation and coreset methods on standard datasets and produces 128 × 128 distilled images from ImageNet.

  • Takeaways & Limitations

    The approach balances short-range matching against full-trajectory optimization and scales dataset distillation to higher-resolution visual data.

  • Takeaways & Limitations

    Pre-computed expert trajectories impose substantial additional computational and storage costs, including 8 GPU hours for 200 CIFAR experts and 15 GPU hours for 100 ImageNet experts.

Abstract

from arXiv · show

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.

1. Introduction

Dataset distillation compresses a large training dataset into a very small synthetic set while preserving test performance. The paper addresses prior methods’ computational, stability, and short-range limitations by matching long-range network training trajectories.

  • Dataset distillation transfers knowledge from a large training dataset into a very small set of synthetic images, potentially one image per class.
  • Prior methods struggle with higher-resolution images, while end-to-end optimization can require substantial compute and memory or become unstable.
  • The method matches segments of synthetic-data parameter trajectories to pre-recorded trajectories from networks trained on real data.
  • Expert trajectories are collected by training models on the real dataset, initializing models from trajectory checkpoints, and optimizing synthetic data against subsequent deviations.
  • 46.3% with one image per class and 71.5% with 50 images per class are achieved on CIFAR-10, versus prior results of 28.8% / 63.0% and 36.1% / 46.5%.

2. Related Work

Related work spans dataset-distillation methods based on optimization, gradient or distribution matching, and expert-trajectory imitation. The paper positions its approach as simultaneously reducing memory costs and improving performance on standard and higher-resolution datasets.

  • Dataset Distillation: Prior dataset-distillation methods learn soft labels, amplify learning signals through gradient matching, adopt augmentations, and optimize with response matching.
  • Dataset Distillation: Distribution matching reduces memory requirements and supports higher-resolution datasets, but the concurrent method attains inferior performance in most cases.
  • Dataset Distillation: The paper describes long-range parameter matching between synthetic-data training and real-data training, starting from the same initial parameters.
  • Dataset Distillation: The method simultaneously reduces memory costs while outperforming existing works and the concurrent method on standard benchmarks and higher-resolution datasets.
  • Imitation Learning: The approach can be viewed as imitation learning over optimization trajectories obtained from networks trained on real datasets.

3. Method

The method distills a synthetic dataset by matching multi-step student parameter trajectories to pre-recorded expert trajectories from real-data training. It uses random trajectory initialization, differentiable updates, weight matching, and memory-aware batching to guide larger synthetic datasets.

  • 3.1. Expert Trajectories: Expert trajectories are pre-recorded parameter sequences from networks trained on the full real dataset and serve as guidance for synthetic-data distillation.The trajectories are generated by saving network parameters at every epoch across many real-data-trained networks.
  • 3.2. Long-Range Parameter Matching: At each distillation step, a student starts from a randomly sampled expert state, trains for N synthetic-data updates, and is compared with the expert state M updates later.The starting epoch is bounded to avoid later trajectory regions where parameters change little.
  • 3.2. Long-Range Parameter Matching: Differentiable augmentation is applied to synthetic images during student updates so gradients can be back-propagated through the augmentation layer.The learning rate α is trainable and is optimized together with the distilled images.
  • 3.2. Long-Range Parameter Matching: The distilled dataset is optimized using normalized squared L2 distance between the ending student parameters and future expert parameters.Normalization preserves signal at later epochs and calibrates magnitude differences across neurons and layers; L2 performed better empirically than cosine distance.
  • 3.3. Memory Constraints: Per-update minibatch sampling reduces memory use while exposing all distilled images before weight matching and preserving heterogeneity within classes.Sampling a new minibatch only at outer-loop steps can redundantly encode information and cause catastrophic mode collapse.

4. Experiments

Experiments evaluate the method across low- and high-resolution datasets, compare it with dataset-distillation and selection baselines, and examine matching range, architecture robustness, and visual structure. The method consistently outperforms baselines, supports larger images, and benefits from longer-range trajectory matching.

  • Evaluation setup: The evaluation covers CIFAR-10, CIFAR-100, Tiny ImageNet, and 128 × 128 ImageNet subsets, using standard dataset-distillation baselines and instance-selection methods.The comparisons include DD, LD, DC, DSA, KIP, DM, CAFE, random selection, herding, and forgetting.
  • Low-resolution results: The method significantly outperforms all baselines in every low-resolution setting, nearly doubling the next-best test accuracy with one image per class on CIFAR-10 and CIFAR-100.CIFAR-10 distilled images are visualized in Figure 4, while CIFAR-100 images appear in the supplementary material.
  • Low-resolution results: The method generally outperforms KIP, including when KIP is evaluated on wider networks, except for one setting using a much wider 1024-width network.KIP uses the infinite-width network limit, whereas this method does not have that limitation in practical finite-width networks.
  • Low-resolution results: 46.3% and 71.5% classification accuracy are achieved on CIFAR-10 with 1 and 50 images per class, respectively, while gains diminish from 10 to 50 images per class.Accuracy increases from 46.3% to 65.3% between 1 and 10 images per class, then from 65.3% to 71.5% between 10 and 50.
  • High-resolution results: The method produces meaningful 128 × 128 ImageNet-subset images and consistently outperforms viable higher-resolution baselines, while distilled image structure varies with task-specific discriminative cues.For ImageSquawk, information is concentrated in textures; for ImageYellow, structure or secondary color is accentuated.

5. Discussion and Limitations

The method directly matches synthetic and real-data training dynamics, balancing short-range matching with full-trajectory optimization while scaling to higher-resolution images. Pre-computed trajectories save memory but require substantial storage and expert-training computation.

  • The method optimizes synthetic data to induce network training dynamics similar to those produced by real data.
  • It balances short-range single-step matching against unstable and compute-intensive optimization over the full training process.
  • The method is the first reported to scale to 128×128 ImageNet images, supporting insights into the dataset and potential practical applications.
  • Pre-computed trajectories provide significant memory savings but add disk-storage requirements and computational cost for expert-model training.
  • CIFAR experts required 8 GPU hours across 200 experts, while ImageNet subset experts required 15 GPU hours across 100 experts.

A.1. Additional Visualizations

The appendix provides additional visualizations of distilled datasets across CIFAR-100, Tiny ImageNet, and 128×128 ImageNet subsets, including best- and worst-performing classes.

  • Additional visualizations cover CIFAR-100 with one image per class and all Tiny ImageNet classes at one image per class.
  • The appendix shows the 10 best- and worst-performing distilled classes.
  • It also includes visualizations of all 128×128 ImageNet subsets with 10 images per class.

A.2. Additional Quantitative Results

Additional quantitative analyses examine learned synthetic step sizes and the effects of ZCA normalization and distillation time.

  • Analysis of learned learning rates α: Learned synthetic step sizes vary with the numbers of real and synthetic steps, allowing synthetic updates to cover comparable parameter-space distances.
  • Analysis of learned learning rates α: The adaptive step size avoids searching for an additional fixed learning-rate hyper-parameter.
  • The appendix reports ablations for ZCA normalization and distillation time.

A.2.1 Additional Ablation Studies

The ablations examine initialization, normalization, augmentation, expert-trajectory count, expert time-step range, and synthetic-step scaling, while reporting continued state-of-the-art performance under categorical ablations.

  • Initialization, normalization, and augmentation: Initialization, normalization, and augmentation are studied for CIFAR-100 with one image per class, including initialization from Gaussian noise instead of real samples.
  • The appendix includes visualizations for Gaussian-noise initialization and a distilled set trained with one expert trajectory.
  • Synthetic step sizes scale inversely with the number of synthetic steps when matching a comparable parameter-space distance.
  • The categorical ablations retain state-of-the-art performance on CIFAR-100 with one image per class, compared with DSA at 13.9%.
  • Performance w.r.t. the number of expert trajectories: Performance improves sharply as expert trajectories increase from 1 to 20 and nearly saturates by 200 trajectories.
  • Figure 11 presents CIFAR-100 distilled images using one image per class.
  • Performance w.r.t. expert time-step range: The expert starting-time upper bound must avoid being too high or too low to provide informative gradients and expose later training dynamics.

A.3. Experiment Details

The experiments evaluate distilled datasets across datasets, support sizes, architectures, and training configurations, while reporting compute and distillation-time details. Additional analyses examine image quality, initialization, preprocessing, augmentation, expert trajectories, and comparisons with KIP.

  • Experiment Details: The reported experiments use hyperparameters selected for the paper’s best-performing distillation results, with synthetic batch size varying by dataset.Table 6 distinguishes full-support synthetic steps from batched synthetic data and reports expert epochs for cross-dataset simplicity.
  • Experiment Details: Distillation time is generally comparable to DC/DSA, whereas KIP requires over 1,000 GPU hours.The comparison was run on RTX3090 GPUs in the 10 images-per-class setting.
  • Experiment Details: Most improvement occurs within the first 1k distillation iterations, although training continues for 10k iterations.Distillation times for CIFAR10/100 and TinyImageNet are reported in Table 7.
  • Experiment Details: The KIP-to-neural-network comparison uses width-128 ConvNets with a grid search over checkpoints, regularization, augmentation, whitening, label learning, and normalization.Instance normalization was added because the smaller-width architecture otherwise failed to converge on CIFAR-100 synthetic training data.
  • Experiment Details: The visual analyses include Tiny ImageNet images with at least 30% and at most 4% correctness, plus CIFAR-100 ablations removing whitening, differentiable augmentation, or multiple expert trajectories.Additional figures cover random-noise initialization and one image per class for Tiny ImageNet class subsets.
  • Experiment Details: Further figures evaluate one-image-per-class Tiny ImageNet settings and ten-images-per-class ImageNette, ImageWoof, ImageSquawk, ImageMeow, ImageFruit, and ImageYellow settings.These figures extend the experimental coverage across datasets and support sizes.
Loading 2203.11932v1…