Meta-Learning with Network Pruning for Overfitting Reduction

1 INTRODUCTION . The ability of adapting to a new task with several trials is essential for artificial agents . The goal of few-shot learning ( Santoro et al. , 2016 ) is to build a model which is able to get the knack of a new task with limited training samples . Meta-learning ( Schmidhuber , 1987 ; Bengio et al. , 1990 ; Thrun & Pratt , 2012 ) provides a principled way to cast few-shot learning as the problem of learning-to-learn , which typically trains a hypothesis or learning algorithm to memorize the experience from previous tasks for a future task learning with very few samples . The practical importance of meta-learning has been witnessed in many vision and online/reinforcement learning applications including image classification ( Ravi & Larochelle , 2016 ; Li et al. , 2017 ) , multi-arm bandit ( Sung et al. , 2017 ) and 2D navigation ( Finn et al. , 2017 ) . Among others , one particularly simple yet successful meta-learning paradigm is first-order optimization based meta-learning which aims to train hypotheses that can quickly adapt to unseen tasks by performing one or a few steps of ( stochastic ) gradient descent ( Ravi & Larochelle , 2016 ; Finn et al. , 2017 ) . Reasons for the recent increasing attention to this class of gradient-optimization based methods include their outstanding efficiency and scalability exhibited in practice ( Nichol et al. , 2018 ) . Challenge and motivation . A challenge in the existing meta-learning approaches is their tendency to overfit ( Mishra et al. , 2018 ; Yoon et al. , 2018 ) . When training an over-parameterized meta-learner such as very deep and/or wide convolutional neural networks ( CNN ) which are powerful for representation learning , there are two sources of potential overfitting at play : the inter-task overfitting of meta-learner ( or meta-overfitting ) to the training tasks and the inner-task overfitting of task-specific learner to the task training data . There have been recent efforts put to deal with inner-task overfitting ( Lee et al. , 2019 ; Zintgraf et al. , 2019 ) . The study on the inter-task meta-overfitting , however , still remains under explored . Since in principle the optimization-based meta-learning is designed to learn fast from small amount of data in new tasks , we expect the meta-overfitting to play a more important role in influencing the overall generalization performance of the trained meta-learner . For an example , it can be observed from the curves in Figure 1 ( a ) for MiniImageNet that although a 4- layer CNN with 128 channels achieves significantly higher training accuracy than the same network architecture with 32 channels , the two networks have quite comparable testing accuracy . It is thus crucial to develop robust methods that can well handle meta-overfitting especially when the number of training tasks is substantially smaller than the number of meta-parameters . Sparsity model is a promising tool for high-dimensional machine learning with guaranteed statistical efficiency and robustness to overfitting ( Maurer & Pontil , 2012 ; Yuan et al. , 2018 ; Abramovich & Grinshtein , 2019 ) . It has been theoretically and numerically justified by Arora et al . ( 2018 ) that sparsity benefits considerably the generalization performance of deep neural networks . In the regime of compact deep learning , the so called network pruning technique has been widely studied and evidenced to work favorably in generating sparse subnetworks without compromising generalization performance ( Frankle & Carbin , 2018 ; Jin et al. , 2016 ; Han et al. , 2016 ) . Inspired by these remarkable success of sparsity models , it is reasonable to conjecture that sparsity would also be of help for enhancing the robustness of optimization based meta-learning to meta-overfitting . Our contribution . In this paper , we present a novel sparse optimization based approach for metalearning with over-parameterized neural networks . The problem is formulated as learning a sparse meta-initialization network from training tasks such that in a new task the learned subnetwork can quickly converge to the optimal solution via gradient descent . The core idea is to reduce metaoverfitting by controlling the counts of the non-zero parameters in the meta-learner . Theoretically , we have established a set of generalization gap bounds for the proposed sparse meta-learner showing the benefit of sparsity in improving its generalization performance . Practically , we have implemented our approach in a joint algorithmic framework of Reptile ( Nichol et al. , 2018 ) with network pruning , along with two specifications to network pruning methods , which are Dense-Sparse-Dense ( DSD ) ( Han et al. , 2016 ) and Iterative Hard Thresholding ( IHT ) ( Jin et al. , 2016 ) , respectively . The actual performance of our approach has been extensively evaluated on few-shot classification tasks with over-parameterized deep/wide CNNs including ResNet-18 ( He et al. , 2016 ) . Numerical results demonstrate that our method can effectively alleviate overfitting and achieve similar or even superior generalization performance to the conventional dense models . As illusrated in Figure 1 ( b ) , using 128-channel network on 5-way 1-shot tasks , our DSD and IHT based Reptile algorithms are effective in closing the gap between the training accuracy and testing accuracy . 2 RELATED WORK . Optimization-based meta-learning . The family of optimization-based meta-learning approaches usually learns a good hypothesis which can fast adapt to unseen tasks ( Ravi & Larochelle , 2016 ; Finn et al. , 2017 ; Nichol et al. , 2018 ; Khodak et al. , 2019 ) . Compared to the metric ( Koch et al. , 2015 ; Snell et al. , 2017 ) and memory ( Weston et al. , 2014 ; Santoro et al. , 2016 ) based meta learning algorithms , optimization based meta-learning algorithms are gaining increasing attention due to their simplicity , versatility and effectiveness . As a recent leading framework for optimizationbased meta-learning , MAML ( Finn et al. , 2017 ) is designed to estimate a meta-initialization network which can be well fine-tuned in an unseen task via only one or few steps of minibatch gradient descent . Although simple in principle , MAML requires computing Hessian-vector product for back-propagation , which could be computationally expensive when the model is big . The first-order MAML ( FOMAML ) is therefore proposed to improve the computational efficiency by simply ignoring the second-order derivatives in MAML . Reptile Nichol et al . ( 2018 ) is another approximated first-order algorithm which works favorably since it maximizes the inner product between gradients from the same task yet different minibatches , leading to improved model generalization . In ( Lee et al. , 2019 ) , the meta-learner is treated as a feature embedding module of which the output is used as input to train a multi-class kernel support vector machine as base learner . To deal with overfitting , the CAVIA method ( Zintgraf et al. , 2019 ) decomposes the meta-parameters into the so called context parameters and shared parameters . The context parameters are updated for task adaption with limited capacity while the shared parameters are meta-trained for generalization across tasks . Network pruning . Early network weight pruning algorithms date back to Optimal Brain Damage ( LeCun et al. , 1990 ) and Optimal Brain Surgeon ( Hassibi et al. , 1993 ) . A dense-to-sparse algorithm was developed by Han et al . ( 2015 ) to first remove near-zero weights and then fine tune the preserved weights . As a serial work of dense-to-sparse , the dense-sparse-dense ( DSD ) method ( Han et al. , 2016 ) was proposed to re-initialize the pruned parameters as zero and retrain the entire network after the dense-to-sparse pruning phase . The iterative hard thresholding ( IHT ) method ( Jin et al. , 2016 ) shares a similar spirit with DSD to conduct multiple rounds of iteration between pruning and retraining . Srinivas & Babu ( 2015 ) proposed a data-free method to prune the neurons in a trained network . In ( Louizos et al. , 2017 ) , an L0-norm regularized risk minimization framework was proposed to learn sparse networks during training . More recently , Frankle & Carbin ( 2018 ) introduced and studied the “ lottery ticket hypothesis ” which assumes that once a network is initialized , there should exist an optimal subnetwork , which can be learned by pruning , that performs as well as the original net or even superior . Despite the remarkable success achieved by both meta-learning and network pruning , it still remains largely open to investigate the impact of network pruning on alleviating the meta-overfitting of optimization based meta-learning , which is of primal interest to our study in this paper . 3 META-LEARNING WITH SPARSITY . In this section , we propose a sparse optimization based meta-learning approach for few-shot learning with over-parameterized neural networks . Our hope is to alleviate the overfitting of meta-learner by reducing its parameter counts without compromising accuracy . 3.1 PROBLEM SETTING . We consider theN -wayK-shot problem as defined in ( Vinyals et al. , 2016 ) . Tasks are sampled from a specific distribution p ( T ) and will be divided into meta training set Str , meta validation set Sval , and meta testing set Stest . Classes in different datasets are disjoint ( i.e. , the class in Str will not appear in Stest ) . During training , each task is made up of support set Dsupp and query set Dquery . Both Dsupp and Dquery are sampled from the same classes of Str . Dsupp is used to for training while Dquery is used for evaluation . For a N -way K-shot classification task , we sampled N classes from dataset , and then K samples are sampled from each of these classes to form Dsupp , namely Dsupp = { ( xkc , ykc ) , k = 1 , 2 , ... , K ; c = 1 , 2 , ... , N } . For example , for a 5-way 2-shot task , we sample 2 data-label pairs from each of 5 classes , thus , such a task has 10 samples . Usually , several other samples of the same classes will be sampled to compose Dquery . For example , Dquery is used in Reptile ( Nichol et al. , 2018 ) in evaluation steps . We use the loss function ` ( ŷ , y ) to measure the discrepancy between the predicted label ŷ and the true label y . Notation . For an integer n , we denote [ n ] as the abbreviation of the index set { 1 , ... , n } . We use to denote the element-wise product operator . We say a function g : Rp 7→ R is G-Lipschitz continuous if |g ( θ ) − g ( θ′ ) | ≤ G‖θ − θ′‖2 , and g is H-smooth if it obeys ‖∇g ( θ ) −∇g ( θ′ ) ‖2 ≤ H‖θ − θ′‖2 . For any sparsity level 1 ≤ s ≤ p , we say a function g is restricted µs-strongly convex if there exists µs > 0 such that g ( θ ) − g ( θ′ ) − 〈∇g ( θ′ ) , θ − θ′〉 ≥ µs2 ‖θ − θ ′‖2 , ∀‖θ − θ′‖0 ≤ s . 3.2 META-LEARNING WITH SPARSITY . Our ultimate goal is to learn a good initialization of parameters for a convolutional neural network fθ : X 7→ Y , where θ is the model parameters set , from a set of training tasks such that the learned initialization network generalizes well to future unseen tasks . Inspired by the recent remarkable success of MAML ( Finn et al. , 2017 ) and the strong generalization capability of sparse deep learning models ( Frankle & Carbin , 2018 ; Arora et al. , 2018 ) , we propose to learn from previous task experience a sparse subnetwork started from which the future task-specific networks can be efficiently learned using first-order optimization methods . To this end , we introduce the following layer-wise sparsity constrained stochastic first-order meta-learning formulation : min θ R ( θ ) : = ET∼p ( T ) [ LDqueryT ( θ − η∇θLDsuppT ( θ ) ) ] , s.t . ‖θl‖0 ≤ kl , l ∈ [ L ] , ( 1 ) where LDsuppT ( θ ) = 1 NK ∑ ( xkc , y k c ) ∈D supp T ` ( fθ ( x k c ) , y k c ) is the empirical risk for task T , η is a stepsize constant , ||θl||0 denotes the number of non-zero entries in the parameters of l-th layer θl , kl controls the sparsity level of the l-th layer , and L is the total number of layers in the network . In general , the mathematical formulation of task distribution p ( T ) is unknown but we usually have access to a set of i.i.d . training tasks S = { Ti } Mi=1 sampled from p ( T ) . Thus the following empirical version of the population form in equation 1 is alternatively considered for training : min θ RS ( θ ) : = 1 M M∑ i=1 [ LDqueryTi ( θ − η∇θLDsuppTi ( θ ) ) ] , s.t . ‖θl‖0 ≤ kl , l ∈ [ L ] . ( 2 ) To compare with MAML , our model shares an identical objective function but with the layer-wise sparsity constraints ‖θl‖0 ≤ kl imposed for the purpose of enhancing learnability of the overparameterized meta-initialization network . In view of the “ lottery ticket hypothesis ” ( Frankle & Carbin , 2018 ) , the model in equation 2 can be interpreted as a first-order meta-learner for estimating a subnetwork , or a “ winning ticket ” , for future task learning . Given the strong statistical efficiency and generalization guarantees of sparsity models ( Yuan et al. , 2018 ; Arora et al. , 2018 ) , such a subnetwork is expected to achieve advantageous generalization performance over the dense initialization networks learned by vanilla MAML .

Summary: The paper proposes to use the sparse network to mitigate the task-overfitting. The algorithm is based on an iterative training over the pruning phase and retraining phase. A theoretical analysis of the sparsity effects for the Reptile-based first-order meta-learning is provided which indicates if the sparsity level k satisfies k^2 << parameter dimension p, the sparsity can help with the generalization. The paper studies the Reptile + DSD pruning and Reptile + IHT pruning on two standard few-shot classification benchmark.