PaperlayerView exact revision on arXiv

Source-linked AI summary

Scaling Up Dataset Distillation to ImageNet-1K with Constant Memory

Justin Cui, Ruochen Wang, Si Si, Cho-Jui Hsieh

arXiv:2211.10586v4cs.CVcs.AI

TL;DR

MTT achieves strong dataset-distillation performance but cannot scale to large datasets because unrolled-gradient computation consumes prohibitive memory. The paper introduces constant-memory gradient computation and teacher-based soft label assignment, enabling TESLA to scale to ImageNet-1K with improved performance at high IPCs.

  • Problem

    MTT’s unrolled SGD optimization requires memory that grows with trajectory length, preventing its use on large datasets such as ImageNet-1K.

  • Method

    The paper rearranges trajectory-gradient computation to store one computational graph and uses hyperparameter-free teacher-generated soft labels to guide synthetic-image training.

  • Results

    TESLA reaches ImageNet-1K at 10 and 50 IPCs, with a 5.9% accuracy drop versus full-dataset training at 50 IPCs while using 4.2% of the training data.

  • Takeaways & Limitations

    Combining constant-memory MTT with soft label assignment enables dataset distillation on ImageNet-1K beyond the IPC range reached by prior methods.

  • Takeaways & Limitations

    The method still requires storing teacher-model checkpoints, and removing the teacher-model requirement remains future work.

Abstract

from arXiv · show

Dataset Distillation is a newly emerging area that aims to distill large datasets into much smaller and highly informative synthetic ones to accelerate training and reduce storage. Among various dataset distillation methods, trajectory-matching-based methods (MTT) have achieved SOTA performance in many tasks, e.g., on CIFAR-10/100. However, due to exorbitant memory consumption when unrolling optimization through SGD steps, MTT fails to scale to large-scale datasets such as ImageNet-1K. Can we scale this SOTA method to ImageNet-1K and does its effectiveness on CIFAR transfer to ImageNet-1K? To answer these questions, we first propose a procedure to exactly compute the unrolled gradient with constant memory complexity, which allows us to scale MTT to ImageNet-1K seamlessly with ~6x reduction in memory footprint. We further discover that it is challenging for MTT to handle datasets with a large number of classes, and propose a novel soft label assignment that drastically improves its convergence. The resulting algorithm sets new SOTA on ImageNet-1K: we can scale up to 50 IPCs (Image Per Class) on ImageNet-1K on a single GPU (all previous methods can only scale to 2 IPCs on ImageNet-1K), leading to the best accuracy (only 5.9% accuracy drop against full dataset training) while utilizing only 4.2% of the number of data points - an 18.2% absolute gain over prior SOTA. Our code is available at https://github.com/justincui03/tesla

1. Introduction

Dataset distillation compresses large datasets into learned synthetic samples, but MTT’s unrolled optimization requires prohibitive memory on large datasets. The paper addresses this bottleneck and improves MTT’s ImageNet-1K scalability with constant-memory gradients and soft label assignment.

  • Motivation: Dataset distillation learns a small synthetic set intended to achieve competitive performance with full-dataset training.Unlike core-set selection, synthetic samples are learned in continuous space and can perform better at higher compression rates.
  • Motivation: MTT cannot scale to large datasets because back-propagation stores T gradient computational graphs from its unrolled SGD trajectory.Its memory requirement grows with the number of unrolled optimization steps.
  • Contributions: The proposed memory-efficient MTT stores one gradient computational graph while unrolling T steps, reducing memory complexity from linear to constant with identical performance and marginal overhead.This enables ImageNet-1K experiments at 1, 2, 10, and 50 IPCs.
  • Contributions: Soft label assignment uses teacher-generated soft labels because hard labels converge poorly for MTT on large-class datasets.SLA is train-free, hyperparameter-free, and directly leverages MTT’s reference teacher models.
  • Contributions: TESLA outperforms prior methods at 1 and 2 IPCs and reaches 10 and 50 IPCs, where existing methods fail to scale.The contribution combines constant-memory MTT with SLA.

2. Related Work

Related work groups dataset distillation methods into matching-based and kernel-based approaches. MTT achieves strong medium-dataset performance, while KIP and FrePo reduce the bilevel problem differently but retain substantial computational or memory challenges.

  • Overview: Dataset distillation methods are broadly categorized as matching-based or kernel-based approaches.The field also applies distilled datasets to continuous learning, federated learning, and neural architecture search.
  • Matching-based Approaches: Matching-based methods generate synthetic data by matching gradients or features between models trained on distilled and real datasets.Feature matching was introduced to reduce the memory and computation costs of gradient matching.
  • Matching-based Approaches: MTT matches model training trajectories and achieves state-of-the-art performance on medium-sized datasets such as CIFAR-10 and CIFAR-100.DC-BENCH also reports advantages in transferability and stability under varied augmentations and IPC settings.
  • Kernel-based Approaches: Kernel-based methods formulate distillation as bilevel optimization, using NTK-based kernel ridge regression to obtain a closed-form inner solution.KIP reduces the bilevel problem to single-level optimization but requires thousands of GPU hours for NTK computation.
  • Kernel-based Approaches: FrePo reduces KIP’s cost by learning only last-layer parameters, yet still requires all computational graphs and heavy matrix inversion, limiting larger IPCs.Its approximation keeps other network parameters fixed.

3. Background

MTT trains models on synthetic images for a short trajectory and matches their parameters to models trained longer on real data. Its trajectory-level gradient computation creates memory requirements that become prohibitive as datasets and trajectories grow.

  • Matching Training Trajectories: MTT generates synthetic datasets by matching parameters from synthetic-data training with parameters from real-data training.The target is a short synthetic trajectory that matches a teacher trajectory trained for more steps.
  • Matching Training Trajectories: MTT starts from a teacher parameter state at step t and seeks to match the teacher’s state after M real-data updates using T synthetic-data updates, typically with T ≪ M.Synthetic batches may contain augmented synthetic images.
  • Matching Training Trajectories: The student model is updated with standard SGD using synthetic batches and student learning rate β.The update is applied at each step of the synthetic trajectory.
  • Scalability: Although MTT performs strongly on small datasets, it fails to scale to ImageNet-1K and other large real-world datasets, limiting practicality.The paper identifies scalability as a central weakness of existing condensation methods.
  • Scalability: Unrolling T SGD steps requires differentiating through and storing T high-order gradient computational graphs, making GPU memory grow with trajectory length.Trajectory length increases with dataset size, and simply reducing matching steps can cause suboptimal performance because images may encode redundant information.

4. Our proposed method

TESLA scales trajectory matching to large datasets by combining constant-memory exact gradient computation with train-free soft-label assignment from pretrained teachers.

  • 4.1. MTT with constant memory: O(1) memory replaces O(T) memory for trajectory matching while preserving the original MTT performance.The method stores one computational graph and rearranges gradient computation without requiring all T trajectory graphs.
  • 4.1. MTT with constant memory: Sequential batch-wise computation and gradient accumulation produce constant memory with negligible runtime overhead.The consecutive loss computation has nearly the same total cost as original MTT computation.
  • 4.2. Memory complexity v.s. other methods: O(|˜X|G) memory makes TESLA independent of trajectory length T, unlike MTT’s O(T|˜X_i|G) complexity.The comparison assumes computation-graph memory scales approximately linearly with batch size.
  • 4.3. Soft labels: Soft labels enable information flow across classes, but jointly learning them with images and model weights makes MTT’s inner optimization harder.Hard labels yield poor matching-based performance on ImageNet-1K, whereas the proposed method assigns labels directly from pretrained teacher models.
  • 4.3. Soft labels: SLA assigns each synthetic image soft labels from a sampled teacher trajectory without training additional label parameters.At each iteration, the teacher evaluates synthetic images and the generated labels are used in the student’s MTT updates.
  • 4.3. Soft labels: 7% to 13.4% absolute accuracy gains result from using the target-step teacher model for soft-label assignment across ImageNet-1K IPC settings.The target-step teacher is selected for the main experiments.

5. Experimental Results

TESLA is evaluated on CIFAR-10/100 and ImageNet-1K against prior methods, with experiments covering accuracy, memory, runtime, and soft-label strategies. It achieves strong ImageNet-1K compression and accuracy while substantially reducing memory, and soft labels improve performance at larger class counts.

  • Experimental setup: TESLA is evaluated on CIFAR-10/100 and ImageNet-1K using 1/10/50 IPCs on CIFAR and 1/2/10/50 IPCs on ImageNet-1K.ImageNet-1K is resized to 64×64, while CIFAR-10/100 use ZCA whitening; DSA augmentation is applied during training and evaluation.
  • Experimental setup: The evaluation trains five randomly initialized models on each distilled dataset and reports mean and standard deviation on the real test set.Baseline values are taken from prior papers when available, and unavailable large-scale experiments are marked absent.
  • Empirical results: 5.9% accuracy drop on ImageNet-1K at 50 IPCs is achieved using 4.2% of the training dataset size, an 18.2% accuracy improvement over prior art.At 10 IPCs, TESLA incurs a 16% accuracy drop while using 0.83% of the dataset size; it also outperforms FrePo and DM at 1 and 2 IPCs.
  • Training cost analysis: ∼6x memory reduction on ImageNet-1K requires only ∼2% extra runtime, while CIFAR-100 obtains ∼5x memory reduction with ∼27% runtime overhead.These measurements use batch size 100 and 50 synthetic training or matching steps.
  • Ablation study on soft labels: 7% on IPC 1 and 13.4% on IPC 10 are the reported accuracy gains from soft labels over hard labels on ImageNet-1K.Soft labels are especially effective on ImageNet-1K, whereas they are not needed for simpler datasets with fewer classes such as CIFAR-10/100.
  • Ablation study on soft labels: Target-step teacher models outperform last-epoch models for soft-label generation, while learned labels provide only slight additional improvement.The last-epoch strategy causes the trajectory-matching algorithm to fail to learn effectively.
  • Ablation study on soft labels: 5.3% accuracy is achieved by teacher-generated labels without updating synthetic images, compared with around 7.7% when soft labels are integrated with memory-efficient MTT.The labels are initialized using teacher models from epochs 0 to 9 on ImageNet-1K at IPC 1.

6. Conclusion

The paper reduces MTT’s memory cost to constant complexity and combines this with soft label assignment to scale dataset distillation to ImageNet-1K. The method achieves SOTA performance and transfers across architectures, while remaining subject to teacher-checkpoint and performance-gap limitations.

  • Contributions: O(1) memory complexity and soft label assignment enable dataset distillation on ImageNet-1K with IPC 10 and 50.The method reduces MTT’s memory cost from O(T) to O(1) and guides model-training trajectory matching with soft labels.
  • Results: The distilled datasets achieve SOTA performance on ImageNet-1K and transfer across architectures, including ViT.
  • Limitations: The method still requires storing teacher-model checkpoints despite its constant memory cost with respect to unrolled steps.Removing the teacher-model requirement is identified as a future research direction.
  • Limitations: Dataset distillation remains in an early stage, with a continuing need to close the performance gap to the full model.

A.1. Bi-level optimization

Dataset distillation is formulated as a bi-level optimization problem that seeks a much smaller synthetic dataset whose training behavior approximates that of the original dataset. The surrogate objective can compare synthetic and real data through gradients or training trajectories.

  • Formulation: Dataset distillation seeks a synthetic subset S much smaller than D whose training loss behavior approximates that of the original dataset.The formulation compares expectations under models trained with the original and distilled datasets.
  • Optimization: The bi-level process uses an inner loop to optimize distilled-model parameters and an outer objective to distinguish synthetic data from the original dataset.The surrogate function M identifies the difference between S and D.
  • Surrogate objectives: Examples of surrogate objectives match training gradients or training trajectories generated from the synthetic and real datasets.

A.2. Complexity of other distillation methods

Other dataset-distillation methods use different optimization or matching strategies, but the cited bi-level approaches have lower performance and do not scale to large datasets.

  • Bi-level methods: Earlier bi-level methods achieve lower performance than MTT and FrePo and cannot scale to large datasets.
  • Distribution matching: Distribution-matching methods use single-level optimization, making them fast, but their trade-offs are discussed separately from MTT and FrePo.

A.3. Datasets

The experiments evaluate dataset distillation on CIFAR-10, CIFAR-100, and ImageNet-1K. ImageNet-1K provides the large-scale setting, with 1,000 classes and images resized to 64×64 resolution.

  • Datasets: Experiments cover CIFAR-10, CIFAR-100, and ImageNet-1K.CIFAR-10 and CIFAR-100 contain 10 and 100 classes, respectively, while ImageNet-1K contains 1,000 classes.
  • ImageNet-1K: ImageNet-1K contains 1,281,167 training images and 50,000 testing images across 1,000 classes.
  • Preprocessing: ImageNet-1K images are resized to 64×64 resolution following prior work.

A.4. Data preprocessing

On ImageNet-1K, the work does not apply ZCA preprocessing at any IPC setting, while following MTT’s Kornia ZCA configuration elsewhere.

  • A.4. Data preprocessing: ImageNet-1K experiments omit ZCA preprocessing at all IPCs.For other evaluations, the implementation follows MTT’s Kornia-based ZCA settings.

A.5. Models

The experiments evaluate the proposed memory-efficient MTT across models, datasets, augmentations, labels, learning rates, and resource settings. They emphasize stable optimization and reduced memory usage for scaling dataset distillation, including ImageNet-1K.

  • A.5. Models: The implementation uses the shared ConvNet architecture from DSA, DM, and MTT, while KIP and FrePo use different evaluation models.The common ConvNet has 128 convolutional filters with 3 × 3 kernels, instance normalization, ReLU, and average pooling.
  • A.5. Models: Experiments use one NVIDIA A6000 GPU with 49GB memory, except original MTT memory measurements that use two GPUs when necessary.This setup defines the reported hardware conditions for memory comparisons.
  • A.11. Runtime and memory: The proposed memory cost remains constant with respect to matching steps, whereas original MTT requires storing computational graphs for its unrolled updates.The numerical memory and runtime comparisons are reported in the appendix and use 50 matching steps with batch size 100.
  • Data augmentation: DSA augmentation is applied during both training and evaluation, unlike FrePo’s evaluation-only augmentation setup.The authors report better performance with augmentation, while AutoAugment and RandAugment generally reduce evaluation performance.
  • A.11. Learning learning rate: Learning the student-model learning rate stabilizes training and improves distilled-dataset testing accuracy; on ImageNet-1K IPC 1, it rises from 0.01 to 0.02 after 1000 iterations.The reported learning-rate schedule increases as training progresses.
Loading 2211.10586v4…