Dataset Distillation using Neural Feature Regression

Yongchao Zhou, Ehsan Nezhadarya, Jimmy Ba

NeurIPS 2022

[ OpenReview | ArXiv | GitHub | Video | Slides]

CIFAR100

TinyImageNet

ImageNette and ImageWoof

Abstract

Dataset distillation aims to learn a small synthetic dataset that preserves most of the information from the original dataset. Dataset distillation can be formulated as a bi-level meta-learning problem where the outer loop optimizes the meta-dataset and the inner loop trains a model on the distilled data. Meta-gradient computation is one of the key challenges in this formulation, as differentiating through the inner loop learning procedure introduces significant computation and memory costs. In this paper, we address these challenges using neural Feature Regression with Pooling (FRePo), achieving the state-of-the-art performance with an order of magnitude less memory requirement and two orders of magnitude faster training than previous methods. The proposed algorithm is analogous to truncated backpropagation through time with a pool of models to alleviate various types of overfitting in dataset distillation. FRePo significantly outperforms the previous methods on CIFAR100, Tiny ImageNet, and ImageNet-1K. Furthermore, we show that high-quality distilled data can greatly improve various downstream applications, such as continual learning and membership inference defense.

Method

To compute the meta gradient efficiently, FRePo only trains the last layer of a neural network to convergence while keeping the feature extractor fixed. In this case, computing the prediction on the real data using the model trained on the distilled data can be expressed as a kernel ridge regression and computing the meta-gradient is simply back-propagating through the kernel and a fixed feature extractor. As shown in the animation, FRePo is analogous to 1-step TBPTT as it computes the meta-gradient at each step while performing the online model update. However, instead of backpropagating through the inner optimization, FRePo computes the meta-gradient through a kernel and feature extractor.

To alleviate overfitting, FRePo maintains a diverse pool of models instead of periodically training and resetting a single model. From the meta-learning perspective, we maintain a diverse set of meta-tasks to sample from and avoid sampling very similar tasks at each consecutive gradient computation to avoid overfitting to a particular setup.

Results

FRePo achieves state-of-the-art results on various benchmark datasets with a 100x reduction in training time and a 10x reduction in GPU memory requirement.

Performance vs Training Time

GPU Memory Requirement

We demonstrate that high-quality distilled data can significantly improve various downstream applications, such as continual learning and membership inference defense.

Continual Learning

Membership Inference Defense

Distilled Images

JPEG

Checkpoint

Distillation Animation

MNIST

FashionMNIST

CIFAR10

CIFAR100

TinyImageNet

ImageNet-1K

ImageNette

ImageWoof

BibTeX

@inproceedings{zhou2022dataset,

title={Dataset Distillation using Neural Feature Regression},

author={Zhou, Yongchao and Nezhadarya, Ehsan and Ba, Jimmy},

booktitle={Proceedings of the Advances in Neural Information Processing Systems (NeurIPS)},

year={2022}

}

Acknowledgement

We would like to thank Harris Chan, Andrew Jung, Michael Zhang, Philip Fradkin, Denny Wu, Chong Shao, Leo Lee, Alice Gao, Keiran Paster, and Lazar Atanackovic for their valuable feedback. Jimmy Ba was supported by NSERC Grant [2020-06904], CIFAR AI Chairs program, Google Research Scholar Program and Amazon Research Award. This project was supported by LG Electronics Canada. Resources used in preparing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute for Artificial Intelligence.