Neighbourhood Distillation: On the benefits of non end-to-end distillation

1 INTRODUCTION . As Deep Neural Networks improve on challenging tasks , they also become deeper and bigger . Image classification convolutional neural networks grew from 5 layers in LeNet ( LeCun et al. , 1998 ) to more than a 100 in the latest ResNet models ( He et al. , 2016 ) . However , as models grow in size , training by back propagating gradients through the entire network becomes more challenging and computationally expensive . Convergence in a highly non-convex space can be slow and requires the development of sophisticated optimizers to escape local optima ( Kingma & Ba , 2014 ) . Gradients vanish or explode as they get passed through an increasing number of layers . Very deep neural networks that are trained end-to-end also require accelerators , and time to train to completion . Our work seeks to overcome the limitations of training very deep networks by breaking away from the end-to-end training paradigm . We address the procedure of distilling knowledge from a teacher model and propose to break a deep architecture into smaller components which are distilled independently . There are multiple benefits to working on small neighbourhoods as compared to full models : training a neighbourhood takes significantly less compute than a larger model ; during training , gradients in a neighbourhood only back-propagate through a small number of layers making it unlikely that they will suffer from vanishing or exploding gradients . By breaking a model into smaller neighbourhoods , training can be done in parallel , significantly reducing wall-time for training as well as enabling training on CPUs which are cheaper than custom accelerators but are seldom used in Deep Learning as they are too slow for larger models . Supervision to train the components is provided by a pre-trained teacher architecture , as is commonly used in Knowledge Distillation ( Hinton et al. , 2015 ) , a popular model compression technique that encourages a student architecture to reproduce the outputs of the teacher . For this reason , we call our method Neighbourhood Distillation . In this paper , we explore the idea of Neighbourhood Distillation on a number of different applications , demonstrate its benefits , and advocate for more research into non end-to-end training . Contributions • We provide empirical evidence of the thresholding effect , a phenomenon that highlights deep neural networks ’ resilience to local perturbations of their weights . This observation motivates the idea of Neighbourhood Distillation . • We show that Neighbourhood Distillation is up to 4x faster than Knowledge Distillation while producing models of the same quality . We demonstrate this on model compression and sparsification . • Then , we show that neighbourhoods trained independently can be used in a search algorithm that efficiently explores an exponential number of possibilities to find an optimal student architecture . • Finally , we show applications of Neighbourhood Distillation to zero-data settings . Shallow neighbourhoods model less complex functions which we can distill using only Gaussian noise as a training input . 2 RELATED WORK . Non end-to-end training Before the democratization of deep learning , machine learning methods relied on multi-stage pipelines . For example , the face detection algorithm designed by Viola & Jones ( 2001 ) is a multi-stage pipeline relying first on handcrafted feature extraction and then on a classifier trained to detect faces from the features . Then came the idea of directly learning classification from the input image , leaving the model to learn all parts of the pipeline through a series of hidden layers ( LeCun et al. , 1998 ; Fukushima & Miyake , 1982 ) that could be trained with end-to-end with gradient back-propagation ( Rumelhart et al. , 1986 ) or layerwise training ( Vincent et al. , 2008 ) . Endto -end deep learning gained traction with the success of the AlexNet model ( Krizhevsky et al. , 2012 ) in image classification . It is now the main component of various state-of-the art approaches in object detection ( Redmon et al. , 2016 ; Ren et al. , 2015 ) , image segmentation ( He et al. , 2017 ) , speech processing ( Senior et al. , 2012 ) , machine translation ( Seo et al. , 2016 ; Vaswani et al. , 2017 ) . However , gradient-based end-to-end learning comes with a cost . Highly non-convex losses are harder to optimize ; models of bigger sizes also require more data to fully train ; they suffer from vanishing and exploding gradients ( Hochreiter , 1998 ; Pascanu et al. , 2012 ) . Approaches to overcome these issues can be broken down into three categories . First , several methods have been introduced to ease the training of deep models , such as residual connections ( He et al. , 2016 ) , gated recurrent units ( Cho et al. , 2014 ) , normalization layers ( Ioffe & Szegedy , 2015 ; Ba et al. , 2016 ; Salimans & Kingma , 2016 ) , and more powerful optimizers ( Kingma & Ba , 2014 ; Hinton et al . ; Duchi et al. , 2011 ) . Second , engineering best practices have adapted to rise to the challenges raised by deep learning : pre-trained models trained on large-datasets can be reused for transfer learning , only requiring the fine-tuning of a portion of the model for specific tasks ( Devlin et al. , 2018 ; Dahl et al. , 2011 ) . Distributed training ( Krizhevsky et al. , 2012 ; Dean et al. , 2012 ) and custom hardware accelerators ( Jouppi et al. , 2017 ) were also crucial in accelerating training . The last category , which our work falls into , investigates non end-to-end training methods for deep neural networks . One class of non end-to-end learning method relies on splitting a deep network into gradient-isolated modules trained with local objectives ( Löwe et al. , 2019 ; Nøkland & Eidnes , 2019 ) . Layerwise training ( Belilovsky et al. , 2018 ; Huang et al. , 2017 ) also divides the target network into modules that are sequentially trained in a bottom-up approach . Difference Target Propagation ( Lee et al. , 2015 ) seeks to optimize each layer to output activations close to a given target value . These values are computed by propagating inverses from downstream layers while ours are provided by a pre-trained teacher model . All of these approaches also differ from ours as modules still depend on each other , while our neighbourhoods are distilled independently . Knowledge Distillation Our work specifically draws from Knowledge Distillation ( Hinton et al. , 2015 ) , a general-purpose model compression method that has been successfully applied to vision ( Crowley et al. , 2018 ) and language problems ( Hahn & Choi , 2019 ) . Knowledge Distillation transfers knowledge from a teacher in the form of its predicted soft logits . Various variations have been developed to improve distillation . One direction is to transfer additional knowledge in the form of intermediate activations ( Romero et al. , 2014 ; Aguilar et al. , 2019 ; Zhang et al. , 2017 ) , attention maps ( Zagoruyko & Komodakis , 2016 ) , weight projections ( Hyun Lee et al. , 2018 ) or layer interactions ( Yim et al. , 2017 ) . Other methods also seek to directly address the capacity gap between a teacher and student ( Cho & Hariharan , 2019 ) by distilling from a series of intermediate teachers ( Mirzadeh et al. , 2019 ; Jin et al. , 2019 ) . These methods all distill the student end-to-end . Neural Architecture Search Recent papers study how to combine Knowledge Distillation with Neural Architecture Search methods , which automate the design process by exploring a given search space ( Liu et al. , 2018 ; Pham et al. , 2018 ; Liu et al. , 2017 ; Tan & Le , 2019 ; Zoph et al. , 2017 ) . These methods have successfully been applied to find better suited students for a given teacher ( Kang et al. , 2019 ; Liu et al. , 2020 ) . Closely related to our work , Li et al . ( 2020 ) divide a supernet into blocks and use Knowledge Distillation to train it . However , they only focus on extracting the architectural knowledge from the teacher , ignoring the parameters learned during the search process . 3 THRESHOLDING EFFECT . Due to their size , modern neural networks are usually overparameterized . Their learned representations are redundant and recent empirical studies conducted on ResNets ( Zhang et al. , 2019 ; Veit et al. , 2016 ) showed that it is possible to drop or reset layers in a trained network without hurting their performance . We hypothesize further that less drastic modifications in a network , such as replacing part or all of their sub-components by imperfect approximations , will not result in dramatic error accumulation . In the following section , we provide empirical evidence of an interesting property that supports this : sub-components of a trained model may be perturbed without damaging the model ’ s accuracy , as long as individual local errors remain under a certain threshold . We call this phenomenon the thresholding effect . First , we introduce the notion of neighbourhood that will be used throughout the rest of the paper . Deep neural networks are built by stacking a succession of blocks of operations such as convolutions and non-linear layers . We express this by defining a network T as a composition of n sub-networks : ∀i ∈ { 1 . . . n } , Ti : RFi −→ RFi+1 T = Tn ◦ Tn−1 ◦ · · · ◦ T1 ( 1 ) We call neighbourhood any portion of the network that is delimited by one sub-network Ti . This neighbourhood represents an arbitrary logical construction block in the network and may be replaced by variants of its architecture that have the same input and output shapes . For any given netwrok , one can define multiple ways to break it up into neighbourhoods . The question we set out to answer is the following . Imagine we want to replace part or all the neighbourhoods by imperfect approximations , how good do these approximations need to be to prevent a drastic loss of performance in the modified model ? We consider different pre-trained models and estimate how their accuracy is impacted by perturbations of the network ’ s intermediate features . To do so , we perturb each neighbourhood by artificially introducing some gaussian noise of amplitude to each activation output . Si = Ti + δ ( 2 ) δ ∼ N ( 0 , 2 ) The students are then composed into a full network S = Sn ◦ Sn−1 · · · ◦ S1 which we evaluate . On CIFAR-10 ( Krizhevsky et al. , 2009 ) , we train a ResNetV1-20 ( He et al. , 2016 ) model and define a neighbourhood as one bottleneck block which consists of a two-layer convolutional network and a skip connection . We reiterate a similar experiment on a large-scale dataset . On ImageNet ( Deng et al. , 2009 ) , we train several EfficientNet ( Tan & Le , 2019 ) models and define a neighbourhood as one mobile inverted bottleneck block . Figure 1 shows how errors of different amplitudes accumulate across networks when replacing some or all neighbourhoods by approximations . In particular , we consistently witness a thresholding effect in all networks . When the amplitude of the noise is small enough , the final accuracy of the network is not impacted by accumulated perturbations . This threshold appears to depend on the number of neighbourhoods . ( a ) ResNetV1-20 with 91.6 % accuracy on CIFAR-10 . Each line represents a different number of perturbed blocks . ( b ) EfficientNet models trained on ImageNet . Each line represents a different EfficientNet model for which all neighbourhoods have been perturbed . Figure 1 : We perturb the intermediate outputs at regular locations in the network using a gaussian noise of amplitude and measure the effect of these perturbations on the accuracy of the model . We show that accuracy remains stable ( accuracy change close to 0 ) as long as remains under a certain threshold . The threshold differs between network architectures and number of perturbed neighbourhoods . Note that for all models , even when perturbing all neighborhoods , there is still a range for which there is virtually no loss in accuracy . ( a ) Computational Graphs ( b ) Neighbourhoods for parameter reduction . Figure 2 : ( a ) Computation graphs for Neighbourhood Distillation . Top : the teacher and the student neighbourhoods receive activations from the root of the teacher network . The student neighbourhood is trained to reproduce the output of the teacher . Bottom : the teacher and student outputs are propagated to the head of the teacher network and additional activations between teacher and student networks are compared . The look-ahead loss gives an additional training signal for the student to reproduce the teacher . ( b ) Example of teacher and student neighbourhoods . Our experiments on the thresholding effect show that it is possible to locally replace sub-components of a model without hurting the performance of the reconstructed model . This observation is what motivates Neighbourhood Distillation : neighbourhoods trained to approximate their teacher outputs can be used to reconstruct student networks with no or limited accuracy drop . In the appendix , we also present preliminary results on understanding the thresholding effect and show how regularizing the teacher network can impact the threshold .

Summary: This paper introduces Neighbourhood Distillation (ND), a new training pipeline for knowledge distillation (KD), which splits the student network into smaller neighbourhoods and trains them independently. The authors breaks away from the end-to-end paradigm in previous KD methods and provides empirical evidence to reveal feasibility and effectiveness of ND. Specially, ND can: 1) speed up convergence, 2) reuse in neural architecture search and 3) adapt to the synthetic data.