Dataset Condensation with Gradient Matching

1 INTRODUCTION . Large-scale datasets , comprising millions of samples , are becoming the norm to obtain state-ofthe-art machine learning models in multiple fields including computer vision , natural language processing and speech recognition . At such scales , even storing and preprocessing the data becomes burdensome , and training machine learning models on them demands for specialized equipment and infrastructure . An effective way to deal with large data is data selection – identifying the most representative training samples – that aims at improving data efficiency of machine learning techniques . While classical data selection methods , also known as coreset construction ( Agarwal et al. , 2004 ; Har-Peled & Mazumdar , 2004 ; Feldman et al. , 2013 ) , focus on clustering problems , recent work can be found in continual learning ( Rebuffi et al. , 2017 ; Toneva et al. , 2019 ; Castro et al. , 2018 ; Aljundi et al. , 2019 ) and active learning ( Sener & Savarese , 2018 ) where there is typically a fixed budget in storing and labeling training samples respectively . These methods commonly first define a criterion for representativeness ( e.g . in terms of compactness ( Rebuffi et al. , 2017 ; Castro et al. , 2018 ) , diversity ( Sener & Savarese , 2018 ; Aljundi et al. , 2019 ) , forgetfulness ( Toneva et al. , 2019 ) ) , then select the representative samples based on the criterion , finally use the selected small set to train their model for a downstream task . Unfortunately , these methods have two shortcomings : they typically rely on i ) heuristics ( e.g . picking cluster centers ) that does not guarantee any optimal solution for the downstream task ( e.g . image classification ) , ii ) presence of representative samples , which is neither guaranteed . A recent method , Dataset Distillation ( DD ) ( Wang et al. , 2018 ) goes beyond these limitations by learning a small set of informative images from large training data . In particular , the authors model the network parameters as a function of the synthetic training data and learn them by minimizing the training loss over the original training data w.r.t . synthetic data . Unlike in the coreset methods , the synthesized data are directly optimized for the downstream task and thus the success of the method does not rely on the presence of representative samples . Inspired from DD ( Wang et al. , 2018 ) , we focus on learning to synthesize informative samples that are optimized to train neural networks for downstream tasks and not limited to individual samples in original dataset . Like DD , our goal is to obtain the highest generalization performance with a model trained on a small set of synthetic images , ideally comparable performance to that of a model trained on the original images ( see Figure 1 ( a ) ) . In particular , we investigate the following 1The implementation is available at https : //github.com/VICO-UoE/DatasetCondensation . questions . Is it possible to i ) compress a large image classification dataset into a small synthetic set , ii ) train an image classification model on the synthetic set that can be further used to classify real images , iii ) learn a single set of synthetic images that can be used to train different neural network architectures ? To this end , we propose a Dataset Condensation method to learn a small set of “ condensed ” synthetic samples such that a deep neural network trained on them obtains not only similar performance but also a close solution to a network trained on the large training data in the network parameter space . We formulate this goal as a minimization problem between two sets of gradients of the network parameters that are computed for a training loss over a large fixed training set and a learnable condensed set ( see Figure 1 ( b ) ) . We show that our method enables effective learning of synthetic images and neural networks trained on them , outperforms ( Wang et al. , 2018 ) and coreset methods with a wide margin in multiple computer vision benchmarks . In addition , learning a compact set of synthetic samples also benefits other learning problems when there is a fixed budget on training images . We show that our method outperforms popular data selection methods by providing more informative training samples in continual learning . Finally , we explore a promising use case of our method in neural architecture search , and show that – once our condensed images are learned – they can be used to train numerous network architectures extremely efficiently . Our method is related to knowledge distillation ( KD ) techniques ( Hinton et al. , 2015 ; Buciluǎ et al. , 2006 ; Ba & Caruana , 2014 ; Romero et al. , 2014 ) that transfer the knowledge in an ensemble of models to a single one . Unlike KD , we distill knowledge of a large training set into a small synthetic set . Our method is also related to Generative Adversarial Networks ( Goodfellow et al. , 2014a ; Mirza & Osindero , 2014 ; Radford et al. , 2015 ) and Variational AutoEncoders ( Kingma & Welling , 2013 ) that synthesize high-fidelity samples by capturing the data distribution . In contrast , our goal is to generate informative samples for training deep neural networks rather than to produce “ real-looking ” samples . Finally our method is related to the methods that produce image patches by projecting the feature activations back to the input pixel space ( Zeiler & Fergus , 2014 ) , reconstruct the input image by matching the feature activations ( Mahendran & Vedaldi , 2015 ) , recover private training images for given training gradients ( Zhu et al. , 2019 ; Zhao et al. , 2020 ) , synthesize features from semantic embeddings for zero-shot learning ( Sariyildiz & Cinbis , 2019 ) . Our goal is however to synthesize a set of condensed training images not to recover the original or missing training images . In the remainder of this paper , we first review the problem of dataset condensation and introduce our method in section 2 , present and analyze our results in several image recognition benchmarks in section 3.1 , showcase applications in continual learning and network architecture search in section 3.2 , and conclude the paper with remarks for future directions in section 4 . 2 METHOD . 2.1 DATASET CONDENSATION . Suppose we are given a large dataset consisting of |T | pairs of a training image and its class label T = { ( xi , yi ) } ||T |i=1 where x ∈ X ⊂ Rd , y ∈ { 0 , . . . , C − 1 } , X is a d-dimensional input space and C is the number of classes . We wish to learn a differentiable function φ ( i.e . deep neural network ) with parameters θ that correctly predicts labels of previously unseen images , i.e . y = φθ ( x ) . One can learn the parameters of this function by minimizing an empirical loss term over the training set : θT = argmin θ LT ( θ ) ( 1 ) where LT ( θ ) = 1|T | ∑ ( x , y ) ∈T ` ( φθ ( x ) , y ) , ` ( · , · ) is a task specific loss ( i.e . cross-entropy ) and θT is the minimizer of LT . The generalization performance of the obtained model φθT can be written as Ex∼PD [ ` ( φθT ( x ) , y ) ] where PD is the data distribution . Our goal is to generate a small set of condensed synthetic samples with their labels , S = { ( si , yi ) } ||S|i=1 where s ∈ Rd and y ∈ Y , |S| |T | . Similar to eq . ( 1 ) , once the condensed set is learned , one can train φ on them as follows θS = argmin θ LS ( θ ) ( 2 ) where LS ( θ ) = 1|S| ∑ ( s , y ) ∈S ` ( φθ ( s ) , y ) and θ S is the minimizer of LS . As the synthetic set S is significantly smaller ( 2-3 orders of magnitude ) , we expect the optimization in eq . ( 2 ) to be significantly faster than that in eq . ( 1 ) . We also wish the generalization performance of φθS to be close to φθT , i.e . Ex∼PD [ ` ( φθT ( x ) , y ) ] ' Ex∼PD [ ` ( φθS ( x ) , y ) ] over the real data distribution PD . Discussion . The goal of obtaining comparable generalization performance by training on the condensed data can be formulated in different ways . One approach , which is proposed in ( Wang et al. , 2018 ) and extended in ( Sucholutsky & Schonlau , 2019 ; Bohdal et al. , 2020 ; Such et al. , 2020 ) , is to pose the parameters θS as a function of the synthetic data S : S∗ = argmin S LT ( θS ( S ) ) subject to θS ( S ) = argmin θ LS ( θ ) . ( 3 ) The method aims to find the optimum set of synthetic images S∗ such that the model φθS trained on them minimizes the training loss over the original data . Optimizing eq . ( 3 ) involves a nested loop optimization and solving the inner loop for θS ( S ) at each iteration to recover the gradients for S which requires a computationally expensive procedure – unrolling the recursive computation graph for S over multiple optimization steps for θ ( see ( Samuel & Tappen , 2009 ; Domke , 2012 ) ) . Hence , it does not scale to large models and/or accurate inner-loop optimizers with many steps . Next we propose an alternative formulation for dataset condensation . 2.2 DATASET CONDENSATION WITH PARAMETER MATCHING . Here we aim to learn S such that the model φθS trained on them achieves not only comparable generalization performance to φθT but also converges to a similar solution in the parameter space ( i.e . θS ≈ θT ) . Let φθ be a locally smooth function2 , similar weights ( θS ≈ θT ) imply similar mappings in a local neighborhood and thus generalization performance , i.e . Ex∼PD [ ` ( φθT ( x ) , y ) ] ' Ex∼PD [ ` ( φθS ( x ) , y ) ] . Now we can formulate this goal as min S D ( θS , θT ) subject to θS ( S ) = argmin θ LS ( θ ) ( 4 ) where θT = argminθ LT ( θ ) and D ( · , · ) is a distance function . In a deep neural network , θT typically depends on its initial values θ0 . However , the optimization in eq . ( 4 ) aims to obtain an optimum set of synthetic images only for one model φθT with the initialization θ0 , while our actual goal is to generate samples that can work with a distribution of random initializations Pθ0 . Thus we modify eq . ( 4 ) as follows : min S Eθ0∼Pθ0 [ D ( θ S ( θ0 ) , θ T ( θ0 ) ) ] subject to θS ( S ) = argmin θ LS ( θ ( θ0 ) ) ( 5 ) where θT = argminθ LT ( θ ( θ0 ) ) . For brevity , we use only θS and θT to indicate θS ( θ0 ) and θT ( θ0 ) respectively in the next sections . The standard approach to solving eq . ( 5 ) employs implicit differentiation ( see ( Domke , 2012 ) for details ) , which involves solving an inner loop optimization for θS . As the inner loop optimization θS ( S ) = argminθ LS ( θ ) can be computationally expensive in 2Local smoothness is frequently used to obtain explicit first-order local approximations in deep networks ( e.g . see ( Rifai et al. , 2012 ; Goodfellow et al. , 2014b ; Koh & Liang , 2017 ) ) . case of large-scale models , one can adopt the back-optimization approach in ( Domke , 2012 ) which re-defines θS as the output of an incomplete optimization : θS ( S ) = opt-algθ ( LS ( θ ) , ς ) ( 6 ) where opt-alg is a specific optimization procedure with a fixed number of steps ( ς ) . In practice , θT for different initializations can be trained first in an offline stage and then used as the target parameter vector in eq . ( 5 ) . However , there are two potential issues by learning to regress θT as the target vector . First the distance between θT and intermediate values of θS can be too big in the parameter space with multiple local minima traps along the path and thus it can be too challenging to reach . Second opt-alg involves a limited number of optimization steps as a tradeoff between speed and accuracy which may not be sufficient to take enough steps for reaching the optimal solution . These problems are similar to those of ( Wang et al. , 2018 ) , as they both involve parameterizing θS with S and θ0 .

Summary: This paper tackles the challenging dataset condensation problem. The goal is to learn to synthesize a small dataset, so that a neural network trained on the small synthetic dataset can have similar performance as a network trained on the full dataset. The proposed method tackles the problem by gradient matching. The proposed method achieves state-of-the-art performance, and shows promising results on two other downstream tasks, continual learning and neural architecture search.