summaries: The paper proposes to use infinitely wide CNNs to perform dataset distillation – obtaining a tiny dataset that enables high accuracy. The approach is an extension of a prior work (Kernel Inducing Points (KIP), etc.). Results are shown on MNIST, Fashion-MNIST, SVHN, CIFAR-10/100 datasets. In this work the authors extend the Kernel Inducing Points (KIP) method for dataset distillation to a new set of kernel functions and achieve new state of the art results on dataset compression. Since the KIP algorithm is a key algorithmic mechanism in this work, I will summarize it: The KIP method optimizes what can be seen as a test residual of a target dataset fitted by a kernel regression model that is created using a support dataset. This means that given a target dataset distribution, one can differentiate through the kernel computations to evolve the support data points in the support dataset freely in the input space and optimize them. After optimization the support dataset (Which is usually taken to be much smaller than the target distribution) can be used as a distilled version of the dataset. This compressed dataset enables classification and training without training with extensively large datasets. In this work the authors use the KIP algorithm using a family of kernels corresponding to infinitely wide neural networks that were discovered in the realm of NTK. Such infinite NTK convolution kernels are known to provide state of the art results (when compared with other kernels) on classic classification tasks such as MNIST, CIFAR, etc. The biggest issue of such infinite-network kernels is that they require immense amounts of computation to construct the kernel. In this work, the authors construct such kernels at each training step (!) and then back-propagate through the kernels (!!) to optimize the datapoint. This impressive feat is made possible by considering 3 layer convolutional neural networks and working with a relatively small number of data-points (5K as opposed to the whole dataset). This is still highly computationally challenging and the authors create a careful implementation including an orchestration system working with 100s of GPUs for optimizing the support dataset. With this new kernel family the authors achieve state of the art results on dataset compression and analyze the compressed dataset using high dimensional statistics tools. Paper mainly extends the dataset distillation algorithms KIP and LS (Nguyen et al. 2021) to infinitely wide neural networks. To that end, authors present a distributed framework that draws upon huge hardware resources and show improvement on the distillation performance. Paper also analyzes the synthesized data samples (images, labels) via multiple studies. The paper performs thorough analysis and exploration on the implication of the proposed KIP and LS methods by Nyugen 2021 for dataset distillation. KIP and LS are based on the inf-width network limit as kernels, and devised based on kernel Ridge regression, and has been shown in prior work to have good dataset distillation performance. This paper greatly extends the analysis in Nyugen 2021, by investigating how KIP and LS can be used to achieve SOTA dataset distillation results in a kernel setting, and also sometimes in a neural network (NN) transfer setting. Furthermore, the authors provide impressive and insightful analyses on (1) how the dataset distillation performance behaves as knobs of the algorithms are tuned (2) what are the properties (spectral, intrinsic dim., visual) of the distilled data. Not every paper needs to propose a new method or new theory. Here, thorough analysis and understanding of an existing good method should still be appreciated. This paper performs dataset distillation to improve the deep learning training efficiency. By using the algorithm made by [1] and some expanded methods, the authors achieved the state-of-the-arts accuracy on CIFAR-10, MNIST, Fashion-MNIST, CIFAR10, CIFAR100, and SVHN datasets.  [1] DATASET META-LEARNING FROM KERNEL RIDGE REGRESSION, ICLR 2021
abstractText: The effectiveness of machine learning algorithms arises from being able to extract useful features from large amounts of data. As model and dataset sizes increase, dataset distillation methods that compress large datasets into significantly smaller yet highly performant ones will become valuable in terms of training efficiency and useful feature extraction. To that end, we apply a novel distributed kernel-based meta-learning framework to achieve state-of-the-art results for dataset distillation using infinitely wide convolutional neural networks. For instance, using only 10 datapoints (0.02% of original dataset), we obtain over 65% test accuracy on CIFAR10 image classification task, a dramatic improvement over the previous best test accuracy of 40%. Our state-of-the-art results extend across many other settings for MNIST, Fashion-MNIST, CIFAR-10, CIFAR-100, and SVHN. Furthermore, we perform some preliminary analyses of our distilled datasets to shed light on how they differ from naturally occurring data.
authors.name: Timothy Nguyen
authors.name: Roman Novak
authors.name: Lechao Xiao
authors.name: Jaehoon Lee
references.authors: Alex Krizhevsky Ilya Sutskever Geoffrey E Hinton
references.title: Imagenet classification with deep convolutional neural networks
references.venue: Advances in neural information processing systems,
references.authors: Jacob Devlin Ming-Wei Chang Kenton Lee Kristina Toutanova. Bert
references.title: Pre-training of deep bidirectional transformers for language understanding
references.venue: arXiv preprint arXiv:1810.04805,
references.authors: Tom B Brown Benjamin Mann Nick Ryder Melanie Subbiah Jared Kaplan Prafulla Dhariwal Arvind Neelakantan Pranav Shyam Girish Sastry Amanda Askell
references.title: Language models are few-shot learners
references.authors: Alexey Dosovitskiy Lucas Beyer Alexander Kolesnikov Dirk Weissenborn Xiaohua Zhai Thomas Unterthiner Mostafa Dehghani Matthias Minderer Georg Heigold Sylvain Gelly
references.title: An image is worth 16x16 words: Transformers for image recognition at scale
references.venue: arXiv preprint arXiv:2010.11929,
references.authors: Tongzhou Wang Jun-Yan Zhu Antonio Torralba Alexei
references.title: A Efros. Dataset distillation
references.venue: arXiv preprint arXiv:1811.10959,
references.authors: Ondrej Bohdal Yongxin Yang Timothy Hospedales
references.title: Flexible dataset distillation: Learn labels instead of images
references.venue: arXiv preprint arXiv:2006.08572,
references.authors: Timothy Nguyen Zhourong Chen Jaehoon Lee
references.title: Dataset meta-learning from kernel ridge-regression
references.venue: In International Conference on Learning Representations,
references.authors: Bo Zhao Hakan Bilen
references.title: Dataset condensation with differentiable siamese augmentation
references.venue: In International Conference on Machine Learning,
references.authors: Andrew Ilyas Shibani Santurkar Dimitris Tsipras Logan Engstrom Brandon Tran Aleksander Madry
references.title: Adversarial examples are not bugs, they are features, 2019
references.authors: Minyoung Huh Pulkit Agrawal Alexei A Efros
references.title: What makes imagenet good for transfer learning
references.venue: arXiv preprint arXiv:1608.08614,
references.authors: Katherine L Hermann Andrew K Lampinen
references.title: What shapes feature representations? exploring datasets, architectures, and training
references.venue: arXiv preprint arXiv:2006.12433,
references.authors: Zalán Borsos Mojmír Mutnỳ Andreas Krause
references.title: Coresets via bilevel optimization for continual learning and streaming
references.venue: arXiv preprint arXiv:2006.03875,
references.authors: Bo Zhao Konda Reddy Mopuri Hakan Bilen
references.title: Dataset condensation with gradient matching
references.venue: In International Conference on Learning Representations,
references.authors: Roman Novak Lechao Xiao Jaehoon Lee Yasaman Bahri Greg Yang Jiri Hron Daniel A. Abolafia Jeffrey Pennington Jascha Sohl-Dickstein
references.title: Bayesian deep convolutional networks with many channels are gaussian processes
references.venue: In International Conference on Learning Representations,
references.authors: Adrià Garriga-Alonso Laurence Aitchison Carl Edward Rasmussen
references.title: Deep convolutional networks as shallow gaussian processes
references.venue: In International Conference on Learning Representations,
references.authors: Sanjeev Arora Simon S Du Wei Hu Zhiyuan Li Russ R Salakhutdinov Ruosong Wang
references.title: On exact computation with an infinitely wide neural net
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Zhiyuan Li Ruosong Wang Dingli Yu Simon S Du Wei Hu Ruslan Salakhutdinov Sanjeev Arora
references.title: Enhanced convolutional neural tangent kernels
references.authors: Vaishaal Shankar Alex Chengyu Fang Wenshuo Guo Sara Fridovich-Keil Ludwig Schmidt Jonathan RaganKelley Benjamin Recht
references.title: Neural kernels without tangents
references.venue: In International Conference on Machine Learning,
references.authors: Alberto Bietti
references.title: Approximation and learning with deep convolutional models: a kernel perspective
references.venue: arXiv preprint arXiv:2102.10032,
references.authors: Sanjeev Arora Simon S. Du Zhiyuan Li Ruslan Salakhutdinov Ruosong Wang Dingli Yu
references.title: Harnessing the power of infinitely wide deep nets on small-data tasks
references.venue: In International Conference on Learning Representations,
references.authors: Jaehoon Lee Samuel S Schoenholz Jeffrey Pennington Ben Adlam Lechao Xiao Roman Novak Jascha Sohl-Dickstein
references.title: Finite versus infinite neural networks: an empirical study
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Roman Novak Lechao Xiao Jiri Hron Jaehoon Lee Alexander A. Alemi Jascha Sohl-Dickstein Samuel S. Schoenholz
references.title: Neural tangents: Fast and easy infinite neural networks in python
references.venue: In International Conference on Learning Representations,
references.authors: James Bradbury Roy Frostig Peter Hawkins Matthew James Johnson Chris Leary Dougal Maclaurin Skye Wanderman-Milne
references.title: JAX: composable transformations of Python+NumPy programs, 2018
references.venue: URL http://github.com/google/jax
references.authors: Yann LeCun Corinna Cortes CJ Burges
references.title: Mnist handwritten digit database
references.venue: ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist,
references.authors: Han Xiao Kashif Rasul Roland Vollgraf
references.title: Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017
references.authors: Yuval Netzer Tao Wang Adam Coates Alessandro Bissacco Bo Wu Andrew Y Ng
references.title: Reading digits in natural images with unsupervised feature learning
references.authors: Alex Krizhevsky
references.title: Learning multiple layers of features from tiny images
references.venue: Technical report,
references.authors: Arthur Jacot Franck Gabriel Clement Hongler
references.title: Neural tangent kernel: Convergence and generalization in neural networks
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Jaehoon Lee Lechao Xiao Samuel S. Schoenholz Yasaman Bahri Roman Novak Jascha Sohl-Dickstein Jeffrey Pennington
references.title: Wide neural networks of any depth evolve as linear models under gradient descent
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Jiaoyang Huang Horng-Tzer Yau
references.title: Dynamics of deep neural networks and neural tangent hierarchy
references.venue: In International Conference on Machine Learning,
references.authors: Ethan Dyer Guy Gur-Ari
references.title: Asymptotics of wide networks from feynman diagrams
references.venue: In International Conference on Learning Representations,
references.authors: Anders Andreassen Ethan Dyer
references.title: Asymptotics of wide convolutional neural networks
references.venue: arxiv preprint arXiv:2008.08675,
references.authors: Sho Yaida
references.title: Non-Gaussian processes and neural networks at finite widths
references.venue: In Mathematical and Scientific Machine Learning Conference,
references.authors: R. Bennett
references.title: The intrinsic dimensionality of signal collections
references.venue: IEEE Transactions on Information Theory,
references.authors: Francesco Camastra Antonino Staiano
references.title: Intrinsic dimension estimation: Advances and open problems
references.venue: Information Sciences,
references.authors: Elena Facco Maria d’Errico Alex Rodriguez Alessandro Laio
references.title: Estimating the intrinsic dimension of datasets by a minimal neighborhood information
references.venue: Scientific reports,
references.authors: A. Ansuini A. Laio J. Macke D. Zoccolan
references.title: Intrinsic dimension of data representations in deep neural networks
references.venue: In NeurIPS,
references.authors: Ilia Sucholutsky Matthias Schonlau
references.title: Soft-label dataset distillation and text dataset distillation
references.venue: arXiv preprint arXiv:1910.02551,
references.authors: Radford M. Neal
references.title: Priors for infinite networks (tech
references.venue: rep. no. crg-tr-94-1). University of Toronto,
references.authors: Jaehoon Lee Yasaman Bahri Roman Novak Sam Schoenholz Jeffrey Pennington Jascha Sohl-dickstein
references.title: Deep neural networks as gaussian processes
references.venue: In International Conference on Learning Representations,
references.authors: Alexander G. de G. Matthews Jiri Hron Mark Rowland Richard E. Turner Zoubin Ghahramani
references.title: Gaussian process behaviour in wide deep neural networks
references.venue: In International Conference on Learning Representations,
references.authors: Greg Yang
references.title: Scaling limits of wide neural networks with weight sharing: Gaussian process behavior, gradient independence, and neural tangent kernel derivation
references.venue: arXiv preprint arXiv:1902.04760,
references.authors: Greg Yang
references.title: Wide feedforward or recurrent neural networks of any architecture are gaussian processes
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Jiri Hron Yasaman Bahri Jascha Sohl-Dickstein Roman Novak
references.title: Infinite attention: NNGP and NTK for deep attention networks
references.venue: In International Conference on Machine Learning,
references.authors: Lechao Xiao Jeffrey Pennington Samuel S Schoenholz
references.title: Disentangling trainability and generalization in deep learning
references.venue: In International Conference on Machine Learning,
references.authors: Ben Adlam Jeffrey Pennington
references.title: The neural tangent kernel in high dimensions: Triple descent and a multi-scale theory of generalization
references.venue: In International Conference on Machine Learning,
references.authors: Aitor Lewkowycz Yasaman Bahri Ethan Dyer Jascha Sohl-Dickstein Guy Gur-Ari
references.title: The large learning rate phase of deep learning: the catapult mechanism
references.venue: arXiv preprint arXiv:2003.02218,
references.authors: Aitor Lewkowycz Guy Gur-Ari
references.title: On the training dynamics of deep networks with l_2 regularization
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Ben Adlam Jaehoon Lee Lechao Xiao Jeffrey Pennington Jasper Snoek
references.title: Exploring the uncertainty properties of neural networks’ implicit priors in the infinite-width limit
references.venue: In International Conference on Learning Representations,
references.authors: Daniel S Park Jaehoon Lee Daiyi Peng Yuan Cao Jascha Sohl-Dickstein
references.title: Towards nngp-guided neural architecture search
references.venue: arXiv preprint arXiv:2011.06006,
references.authors: Wuyang Chen Xinyu Gong Zhangyang Wang
references.title: Neural architecture search on imagenet in four {gpu} hours: A theoretically inspired perspective
references.venue: In International Conference on Learning Representations,
references.authors: Amir Zandieh Insu Han Haim Avron Neta Shoham Chaewon Kim Jinwoo Shin
references.title: Scaling neural tangent kernels via sketching and random features
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Joel Hestness Sharan Narang Newsha Ardalani Gregory F. Diamos Heewoo Jun Hassan Kianinejad Md. Mostofa Ali Patwary Yang Yang Yanqi Zhou
references.title: Deep learning scaling is predictable
references.venue: empirically. CoRR,
references.authors: Jonathan S. Rosenfeld Amir Rosenfeld Yonatan Belinkov Nir Shavit
references.title: A constructive prediction of the generalization error across scales
references.venue: In International Conference on Learning Representations,
references.authors: Jared Kaplan Sam McCandlish Tom Henighan Tom B. Brown Benjamin Chess Rewon Child Scott Gray Alec Radford Jeffrey Wu Dario Amodei
references.title: Scaling laws for neural language models, 2020
references.authors: Yasaman Bahri Ethan Dyer Jared Kaplan Jaehoon Lee Utkarsh Sharma
references.title: Explaining neural scaling laws
references.venue: arXiv preprint arXiv:2102.06701,
references.authors: Ekin D. Cubuk Barret Zoph Dandelion Mane Vijay Vasudevan Quoc V. Le
references.title: Autoaugment: Learning augmentation strategies from data
references.venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),
references.authors: Jascha Sohl-Dickstein Roman Novak Samuel S Schoenholz Jaehoon Lee
references.title: On the infinite width limit of neural networks with a standard parameterization
references.venue: arXiv preprint arXiv:2001.07301,
sections.heading: 1 Introduction
sections.text: Deep learning has become extraordinarily successful in a wide variety of settings through the availability of large datasets [Krizhevsky et al., 2012, Devlin et al., 2018, Brown et al., 2020, Dosovitskiy et al., 2020]. Such large datasets enable a neural network to learn useful representations of the data that are adapted to solving tasks of interest. Unfortunately, it can be prohibitively costly to acquire such large datasets and train a neural network for the requisite amount of time.
  One way to mitigate this problem is by constructing smaller datasets that are nevertheless informative. Some direct approaches to this include choosing a representative subset of the dataset (i.e. a coreset) or else performing a low-dimensional projection that reduces the number of features. However, such methods typically introduce a tradeoff between performance and dataset size, since what they produce is a coarse approximation of the full dataset. By contrast, the approach of dataset distillation is to synthesize datasets that are more informative than their natural counterparts when equalizing for dataset size [Wang et al., 2018, Bohdal et al., 2020, Nguyen et al., 2021, Zhao and Bilen, 2021]. Such resulting datasets will not arise from the distribution of natural images but will nevertheless capture features useful to a neural network, a capability which remains mysterious and is far from being well-understood [Ilyas et al., 2019, Huh et al., 2016, Hermann and Lampinen, 2020].
  The applications of such smaller, distilled datasets are diverse. For nonparametric methods that scale poorly with the training dataset (e.g. nearest-neighbors or kernel-ridge regression), having a reduced dataset decreases the associated memory and inference costs. For the training of neural networks, such distilled datasets have found several applications in the literature, including increasing the
  ∗Work done while at Google Research.
  35th Conference on Neural Information Processing Systems (NeurIPS 2021).
  effectiveness of replay methods in continual learning [Borsos et al., 2020] and helping to accelerate neural architecture search [Zhao et al., 2021, Zhao and Bilen, 2021].
  In this paper, we perform a large-scale extension of the methods of Nguyen et al. [2021] to obtain new state-of-the-art (SOTA) dataset distillation results. Specifically, we apply the algorithms KIP (Kernel Inducing Points) and LS (Label Solve), first developed in Nguyen et al. [2021], to infinitely wide convolutional networks by implementing a novel, distributed meta-learning framework that draws upon hundreds of accelerators per training. The need for such resources is necessitated by the computational costs of using infinitely wide neural networks built out of components occurring in modern image classification models: convolutional and pooling layers (see §B for details). The consequence is that we obtain distilled datasets that are effective for both kernel ridge-regression and neural network training.
  Additionally, we initiate a preliminary study of the images and labels which KIP learns. We provide a visual and quantitative analysis of the data learned and find some surprising results concerning their interpretability and their dimensional and spectral properties. Given the efficacy of KIP and LS learned data, we believe a better understanding of them would aid in the understanding of feature learning in neural networks.
  To summarize, our contributions are as follows:
  1. We achieve SOTA dataset distillation results on a wide variety of datasets (MNIST, FashionMNIST, SVHN, CIFAR-10, CIFAR-100) for both kernel ridge-regression and neural network training. In several instances, our results achieve an impressively wide margin over prior art, including over 25% and 37% absolute gain in accuracy on CIFAR-10 and SVHN image classification, respectively, when using only 10 images (Tables 1, 2, A11).
  2. We develop a novel, distributed meta-learning framework specifically tailored to the computational burdens of sophisticated neural kernels (§2.1).
  3. We highlight and analyze some of the peculiar features of the distilled datasets we obtain, illustrating how they differ from natural data (§4).
  4. We open source the distilled datasets, which used thousands of GPU hours, for the research community to further investigate at https://github.com/google-research/ google-research/tree/master/kip.
sections.heading: 2 Setup
sections.text: Background on infinitely wide convolutional networks. Recent literature has established that Bayesian and gradient-descent trained neural networks converge to Gaussian Processes (GP) as the number of hidden units in intermediary layers approaches infinity (see §5). These results hold for many different architectures, including convolutional networks, which converge to a particular GP in the limit of infinite channels [Novak et al., 2019, Garriga-Alonso et al., 2019, Arora et al., 2019]. Bayesian networks are described by the Neural Network Gaussian Process (NNGP) kernel, while gradient descent networks are described by the Neural Tangent Kernel (NTK). Since we are interested in synthesizing datasets that can be used with both kernel methods and common gradient-descent trained neural networks, we focus on NTK in this work.
  Infinitely wide networks have been shown to achieve SOTA (among non-parametric kernels) results on image classification tasks [Novak et al., 2019, Arora et al., 2019, Li et al., 2019, Shankar et al., 2020, Bietti, 2021] and even rival finite-width networks in certain settings [Arora et al., 2020, Lee et al., 2020]. This makes such kernels especially suitable for our task. As convolutional models, they encode useful inductive biases of locality and translation invariance [Novak et al., 2019], which enable good generalization. Moreover, flexible and efficient computation of these kernels are possible due to the Neural Tangent library [Novak et al., 2020].
  Specific models considered. The central neural network (and corresponding infinite-width model) we consider is a simple 4-layer convolutional network with average pooling layers that we refer to as ConvNet throughout the text. This architecture is a slightly modified version of the default model used by Zhao and Bilen [2021], Zhao et al. [2021] and was chosen for ease of baselining (see §A for details). In several other settings we also consider convolutional networks without pooling
  layers ConvVec,2 and networks with no convolutions and only fully-connected layers FC. Depth of architecture (as measured by number of hidden layers) is indicated by an integer suffix.
  Background on algorithms. We review the Kernel Inducing Points (KIP) and Label Solve (LS) algorithms introduced by Nguyen et al. [2021]. Given a kernel K, the kernel ridge-regression (KRR) loss function trained on a support dataset (Xs, ys) and evaluated on a target dataset (Xt, yt) is
  L(Xs, ys) = 1 2 ÂÂÂÂÂyt −KXtXs(KXsXs + λI) −1 ys ÂÂÂÂÂ 2 2 , (1)
  where if U and V are sets, KUV is the matrix of kernel elements (K(u, v))u∈U,v∈V . Here λ > 0 is a fixed regularization parameter. The KIP algorithm consists of minimizing (1) with respect to the support set (either just the Xs or along with the labels ys). Here, we sample (Xt, yt) from a target dataset D at every (meta)step, and update the support set using gradient-based methods. Additional variations include augmenting theXt or sampling a different kernelK (from a fixed family of kernels) at each step.
  The Label Solve algorithm consists of solving for the least-norm minimizer of (1) with respect to ys. This yields the labels
  y ∗ s = (KXtXs(KXsXs + λI) −1) + yt, (2)
  where A+ denotes the pseudo-inverse of the matrix A. Note that here (Xt, yt) = D, i.e. the labels are solved using the whole target set.
  In our applications of KIP and Label Solve, the target dataset D is always significantly larger than the support set (Xs, ys). Hence, the learned support set or solved labels can be regarded as distilled versions of their respective targets. We also initialize our support images to be a subset of natural images, though they could also be initialized randomly.
  Based on the infinite-width correspondence outlined above and in §5, dataset distillation using KIP or LS that is optimized for KRR should extend to the corresponding finite-width neural network training. Our experimental results in §3 validate this expectation across many settings.
sections.heading: 2.1 Client-Server Distributed Workflow
sections.text: We invoke a client-server model of distributed computation3, in which a server distributes independent workloads to a large pool of client workers that share a queue for receiving and sending work. Our distributed implementation of the KIP algorithm has two distinct stages:
  Forward pass: In this step, we compute the support-support and target-support matrices K(Xs, Xs) and K(Xt, Xs). To do so, we partition Xs ×Xs and Xt ×Xs into pairs of images (x, x′), each with batch size B. We send such pairs to workers compute the respective matrix block K(x, x′). The server aggregates all these blocks to obtain the K(Xs, Xs) and K(Xt, Xs) matrices. Backward pass: In this step, we need to compute the gradient of the loss L (1) with respect to the support set Xs. We need only consider ∂L/∂Xs since ∂L/∂ys is cheap to compute. By the chain rule, we can write
  ∂L
  ∂Xs =
  ∂L ∂(K(Xs, Xs)) ∂K(Xs, Xs) ∂Xs +
  ∂L ∂(K(Xt, Xs)) ∂K(Xt, Xs) ∂Xs .
  The derivatives of L with respect to the kernel matrices are inexpensive, since L depends in a simple way on them (matrix multiplication and inversion). What is expensive to compute is the derivative of the kernel matrices with respect to the inputs. Each kernel element is an independent function of the inputs and a naive computation of the derivative of a block would require forwardmode differentiation, infeasible due to the size of the input images and the cost to compute the individual kernel elements. Thus our main novelty is to divide up the gradient computation into backward differentiation sub-computations, specifically by using the built-in function jax.vjp in JAX [Bradbury et al., 2018]. Denoting K = K(Xs, Xs) or K(Xt, Xs) for short-hand, we divide the
  2The abbreviation “Vec” stands for vectorization, to indicate that activations of the network are vectorized (flattened) before the top fully-connected layer instead of being pooled.
  3Implemented using Courier available at https://github.com/deepmind/launchpad.
sections.heading: 1 12.8±0.3 13.9±0.3 - 23.8 34.9±0.1 33.3±0.3CIFAR-100 10 25.2±0.3 32.3±0.3 - 39.2 47.9±0.2 49.5±0.3
sections.text: matrix ∂L/∂K, computed on the server, into B ×B blocks corresponding to ∂L/∂K(x, x′), where x and x′ each have batch size B. We send each such block, along with the corresponding block of image data (x, x′), to a worker. The worker then treats the ∂L/∂K(x, x′) it receives as the cotangent vector argument of jax.vjp that, via contraction, converts the derivative of K(x, x′) with respect to x into a scalar. The server aggregates all these partial gradient computations performed by the workers, over all possible B ×B blocks, to compute the total gradient ∂L/∂Xs used to update Xs.
sections.heading: 3 Experimental Results
sections.heading: 3.1 Kernel Distillation Results
sections.text: We apply the KIP and LS algorithms using the ConvNet architecture on the datasets MNIST [LeCun et al., 2010], Fashion MNIST [Xiao et al., 2017], SVHN [Netzer et al., 2011], CIFAR-10 [Krizhevsky, 2009], and CIFAR-100. Here, the goal is to condense the train dataset down to a learned dataset of size 1, 10, or 50 images per class. We consider a variety of hyperparameter settings (image preprocessing method, whether to augment target data, and whether to train the support labels for KIP), the full details of which are described in §A. For space reasons, we show a subset of our results in Table 1, with results corresponding to the remaining set of hyperparameters left to Tables A3-A10. We highlight here that a crucial ingredient for our strong results in the RGB dataset setting is the use of regularized ZCA preprocessing. Note the variable effect that our augmentations have on performance (see the last two columns of Table 1): they typically only provides a benefit for a sufficiently large support set. This result is consistent with Zhao et al. [2021], Zhao and Bilen [2021], in which gains from augmentations are also generally obtained from larger support sets. We tried varying the fraction (0.25 and 0.5) of each target batch that is augmented at each training step and found that while 10 images still did best without augmentations, 100 and 500 images typically did slightly better with some partial augmentations (versus none or full). For instance, for 500 images on CIFAR-10, we obtained 81.1% test accuracy using augmentation rate 0.5. Thus, our observation is that given that larger support sets can distill larger target datasets, as the former increases in size, the latter can be augmented more aggressively for obtaining optimal generalization performance.
  Remarkably, our results in Table 1 outperform all prior baselines across all dataset settings. Our results are especially strong in the small support size regime, with our 1 image per class results for KRR outperforming over 100 times as many natural images (see Table A1). We also obtain a significant margin over prior art across all datasets, with our largest margin being a 37% absolute gain in test accuracy for the SVHN dataset.
sections.heading: 3.2 Kernel Transfer
sections.text: In §3.1, we focused on obtaining state-of-the-art dataset distillation results for image classification using a specific kernel (ConvNet). Here, we consider the variation of KIP in which we sample from a family of kernels (which we call sampling KIP). We validate that sampling KIP adds robustness to the learned images in that they perform well for the family of kernels sampled during training.
  In Figure 1, we plot test performance of sampling KIP when using the kernels ConvNet, Conv-Vec3, and Conv-Vec8 (denoted by “all”) alongside KIP trained with just the individual kernels. Sampling KIP performs well at test time when using any of the three kernels, whereas datasets trained using a single kernel have a significant performance drop when using a different kernel.
sections.heading: 3.3 Neural Network Transfer
sections.text: In this section, we study how our distilled datasets optimized using KIP and LS transfer to the setting of finite-width neural networks. The main results are shown in Table 2. The third column shows the best neural network performance obtained by training on a KIP dataset of the corresponding size with respect to some choice of KIP and neural network training hyperparameters (see §A for details). Since the datasets are optimized for kernel ridge-regression and not for neural network training itself, we expect some performance loss when transferring to finite-width networks, which we record in the fourth column. Remarkably, the drop due to this transfer is quite moderate or small and sometimes the transfer can even lead to gain in performance (see LS for SVHN dataset with 10 images per class).
  Overall, our transfer to finite-width networks outperforms prior art based on DC/DSA [Zhao et al., 2021, Zhao and Bilen, 2021] in the 1 image per class setting for all the RGB datasets (SVHN, CIFAR-10, CIFAR-100). Moreover, for CIFAR-10, we outperform DC/DSA in all settings.
  Figures 2 and 3 provide a closer look at KIP transfer changes under various settings. The first of these tracks how transfer performance changes when adding additional layers as function of the number of KIP training steps used. The normalization layers appear to harm performance for MSE loss, which can be anticipated from their absence in the KIP and LS optimization procedures. However they appear to provide some benefit for cross entropy loss. For Figure 3, we observe that as KIP training progresses, the downstream finite-width network’s performance also improves in general. A notable exception is observed when learning the labels in KIP, where longer training steps lead to deterioration of information useful to training finite-width neural networks. We also observe that as predicted by infinite-width theory [Jacot et al., 2018, Lee et al., 2019], the overall gap between KIP or LS performance and finite-width neural network decreases as the width increases. While our best performing transfer is obtained with width 1024, Figure 3 (middle) suggest that even with modest width of 64, our transfer can outperform prior art of 60.6% by Zhao and Bilen [2021].
  Finally, in Figure 4, we investigate the performance of KIP images over the course of training of a single run, as compared to natural images, over a range of hyperparameters. We find the outperformance of KIP images above natural images consistent across hyperparameters and checkpoints. This suggests that our KIP images may also be effective for accelerated hyperparameter search, an application of dataset distillation explored in Zhao et al. [2021], Zhao and Bilen [2021].
  Altogether, we find our neural network training results encouraging. First, it validates the applicability of infinite-width methods to the setting of finite width [Huang and Yau, 2020, Dyer and Gur-Ari, 2020, Andreassen and Dyer, 2020, Yaida, 2020, Lee et al., 2020]. Second, we find some of the transfer results quite surprising, including efficacy of label solve and the use of cross-entropy for certain settings (see §A for the full details).
sections.heading: 4 Understanding KIP Images and Labels
sections.text: A natural question to consider is what causes KIP to improve generalization performance. Does it simplify support images, removing noise and minor sources of variation while keeping only the core features shared between many target images of a given class? Or does it make them more complex by producing outputs that combine characteristics of many samples in a single resulting collage? While
  properly answering this question is subject to precise definitions of simplicity and complexity, we find that KIP tends to increase the complexity of pictures based on the following experiments:
  • Visual analysis: qualitatively, KIP images tend to be richer in textures and contours. • Dimensional analysis: KIP produces images of higher intrinsic dimensionality (i.e. an
  estimate of the dimensionality of the manifold on which the images lie) than natural images.
  • Spectral analysis: unlike natural images, for which the bulk of generalization performance is explained by a small number of top few eigendirections, KIP images leverage the whole spectrum much more evenly.
  Combined, these results let us conclude that KIP usually increases complexity, integrating features from many target images into much fewer support images.
  Visual Analysis. A visual inspection of our learned data leads to intriguing observations in terms of interpretability. Figure 5 shows examples of KIP learned images from CIFAR-100. The resulting images are heterogeneous in terms of how they can be interpreted as distilling the data. For instance, the distilled apple image seems to consist of many apples nested within a possibly larger apple, whereas the distilled bottle image starts off as two bottles and before transforming into one, while other classes (like the beaver) are altogether visually indistinct. Investigating and quantifying aspects that make these images generalize so well is a promising avenue for future work. We show examples from other datasets in Figure A2.
  In Figure 6, we compare MNIST KIP data learned with and without label learning. For the latter case with images and labels optimized jointly, while labels become more informative, encoding richer inter-class information, the images become less interpretable. This behavior consistently leads to superior KRR results, but appears to not be leveraged as efficiently in the neural network transfer setting (Table 2). Experimental details can be found in §A.
  Dimensional Analysis. We study the intrinsic dimension of KIP images and find that they tend to grow. Intrinsic dimension (ID) was first defined by Bennett [1969] as “the number of free parameters required in a hypothetical signal generator capable of producing a close approximation to each signal in the collection”. In our context, it is the dimensionality of the manifold embedded into the image space which contains all the support images. Intuitively, simple datasets have a low ID as they can be described by a small number of coordinates on the low-dimensional data manifold.
  ID can be defined and estimated differently based on assumptions on the manifold structure and the probability density function of the images on this manifold (see [Camastra and Staiano, 2016] for review). We use the “Two-NN” method developed by [Facco et al., 2017], which makes relatively few assumptions and allows to estimate ID only from two nearest-neighbor distances for each datapoint.
  Figure 7 shows that the ID is increasing for the learned KIP images as a function of the training step across a variety of configurations (training with or without augmentations/label learning) and datasets. One might expect that a distillation procedure should decrease dimensionality. On the other hand, Ansuini et al. [2019] showed that ID increases in the earlier layers of a trained neural network. It remains to be understood if this latter observation has any relationship with our increased ID. Note that ZCA preprocessing, which played an important role for getting the best performance for our RGB datasets, increases dimensionality of the underlying data (see Figure A4).
  Spectral Analysis. Another distinguishing property of KIP images is how their spectral components contribute to performance. In Figure 8, we spot how different spectral bands of KIP learned images affect test performance as compared to their initial natural images. Here, we use the FC2, Conv-Vec8, and ConvNet architectures. We note that for natural images (light bars), most of their performance is captured by the top 20% of eigenvalues. For KIP images, the performance is either more evenly distributed across the bands (FC and Conv-Vec8) or else is skewed towards the tail (ConvNet).
sections.heading: 5 Related Work
sections.text: Dataset distillation was first studied in Wang et al. [2018]. The work of Sucholutsky and Schonlau [2019], Bohdal et al. [2020] build upon it by distilling labels. Zhao et al. [2021] proposes condensing a training set by harnessing a gradient matching condition. Zhao and Bilen [2021] takes this idea further by applying a suitable augmentation strategy. Note that while Zhao and Bilen [2021] is limited in augmentation expressiveness (they have to apply a single augmentation per training iteration), we can sample augmentations independently per image in our target set per train step. Our work together with Nguyen et al. [2021] are, to the best of our knowledge, the only works using kernel-based methods for dataset distillation on image classification datasets.
  Our use of kernels stems from the correspondence between infinitely-wide neural networks and kernel methods [Neal, 1994, Lee et al., 2018, Matthews et al., 2018, Jacot et al., 2018, Novak et al., 2019, Garriga-Alonso et al., 2019, Arora et al., 2019, Yang, 2019a,b, Hron et al., 2020] and the extended correspondence with the finite width corrections [Dyer and Gur-Ari, 2020, Huang and Yau, 2020, Yaida, 2020]. These correspondences underlie the transferability of our KRR results to neural networks , and have been utilized in understanding trainability [Xiao et al., 2020], generalizations [Adlam and Pennington, 2020], training dynamics [Lewkowycz et al., 2020, Lewkowycz and Gur-Ari, 2020], uncertainty [Adlam et al., 2021], and demonstrated their effectiveness for smaller datasets [Arora et al., 2020] and neural architecture search [Park et al., 2020, Chen et al., 2021].
sections.heading: 6 Conclusion
sections.text: We performed an extensive study of dataset distillation using the KIP and LS algorithms applied to convolutional architectures, obtaining SOTA results on a variety of image classification datasets. In some cases, our learned datasets were more effective than a natural dataset two orders of magnitude larger in size. There are many interesting followup directions and questions from our work:
  First, integrating efficient kernel-approximation methods into our algorithms, such as those of Zandieh et al. [2021], will reduce computational burden and enable scaling up to larger datasets. In this direction, the understanding of how various resources (e.g. data, parameter count, compute) scale when optimizing for neural network performance has received significant attention as machine learning models continue to stretch computational limits [Hestness et al., 2017, Rosenfeld et al., 2020, Kaplan et al., 2020, Bahri et al., 2021]. Developing our understanding of how to harness smaller, yet more useful representations data would aid in such endeavors. In particular, it would be especially interesting to explore how well datasets can be compressed as they scale up in size.
  Second, LS and KIP with label learning shows that optimizing labels is a very powerful tool for dataset distillation. The labels we obtain are quite far away from standard, interpretable labels and we feel their effectiveness suggests that understanding of how to optimally label data warrants further study.
  Finally, the novel features obtained by our learned datasets, and those of dataset distillation methods in general, may reveal insights into interpretability and the nature of sample-efficient representations. For instance, observe the bee and bicycle images in Figure 5: the bee class distills into what appears to be spurious visual features (e.g. pollen), while the bicycle class distills to the essential contours of a typical bicycle. Additional analyses and explorations of this type could offer insights into the perennial question of how neural networks learn and generalize.
  Acknowledgments We would like to acknowledge special thanks to Samuel S. Schoenholz, who proposed and helped develop the overall strategy for our distributed KIP learning methodology. We are also grateful to Ekin Dogus Cubuk and Manuel Kroiss for helpful discussions.
title: Dataset Distillation with Infinitely Wide Convolutional Networks