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