Dataset Distillation with Infinitely Wide Convolutional Networks

1 Introduction . 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 . 2 Setup . 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 . 2.1 Client-Server Distributed Workflow . 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 . 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 . 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 .

Summary: 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.