ParaDiS: Parallelly Distributable Slimmable Neural Networks

1 INTRODUCTION . Neural networks are more and more frequently run on end user devices such as mobile phones or Internet of Things ( IoT ) devices instead of the cloud . This is due to clear advantages in terms of better processing latency and privacy preservation . However , those devices have often very limited ( runtime and memory ) resources . While manually designing lightweight networks ( Howard et al. , 2017 ; Zhang et al. , 2018 ) or using model compression schemes ( Cheng et al. , 2018 ) ( e.g. , pruning ( LeCun et al. , 1990 ) or knowledge distillation ( Hinton et al. , 2015 ; Xie et al. , 2020 ) ) allow satisfying particular constraints , these approaches are not well suited when those resources vary across devices and over time . Indeed , a particular model is handcrafted or compressed for given resource constraints , and once those constraints change , a new model needs to be created . To this end , a new family of approaches that are able to instantly trade off between the accuracy and efficiency was introduced . We call these approaches ( or models ) flexible ( Ozerov & Duong , 2021 ) , and they include among others early-exit models ( Huang et al. , 2018 ) , slimmable neural networks ( Yu et al. , 2019 ; Yu & Huang , 2019 ) , Once-for-All ( OFA ) networks ( Cai et al. , 2019 ; Sahni et al. , 2021 ) , and neural mixture models ( Ruiz & Verbeek , 2020 ) . An alternative path to gain in efficiency under constrained resources setting , consists in distributing the neural network inference over several devices ( if available ) ( Teerapittayanon et al. , 2017 ) . The simplest way is to distribute a neural network sequentially , e.g. , in case of two devices by executing early layers on the first device and the remaining layers on the second one , while transmitting the intermediate features between the two devices over a network ( see Fig . 1 ( B ) ) . Depending on the transmission network bandwidth , the features might undergo a lossless or a more or less severe lossy compression ( Choi et al. , 2018 ) . Neural network ( A ) ( see Fig . 1 ) might be either distributed as it is or , in case the features are distorted by compression , after a fine-tuning or retraining ( Choi et al. , 2018 ) . Though sequential distribution might overcome issues of limited memory and com- puting power , it does not decrease the processing latency and might only increase it because of the additional communication delay . This might be overcome via a parallel distribution ( Hadidi et al. , 2019 ) , as illustrated on Fig . 1 ( C ) . However , to maintain the functionality of the original neural network all convolutional and fully connected layers need to communicate , which requires a considerable communication burden and increases processing latency . To overcome this the model may be distributed parallelly without intermediate communication ( Bhardwaj et al. , 2019 ; Asif et al. , 2019 ) , as on Fig . 1 ( D ) , where the only communication needed is to transmit the data to each device and the final result from each device for a fusion . In this work we chose to focus on this parallel distribution without communication and call it simply parallel distribution hereafter , since it leads to a much smaller processing latency , as compared to sequential distribution ; does not suffer from high communication burden of communicating parallel distribution , and is also easy to deploy . The price to be paid for these advantages consists in a potentially decreasing performance and a need of retraining due to the loss of connections between the intermediate layers . The question we are trying to answer in this work is : Can we have one neural network that may instantly and near-optimally 1 be parallelly distributed over several devices , regardless of the number of devices and their capacities ? In this paper we introduce a new framework called parallelly distributable slimmable neural networks or ParaDiS in short that allows such functionality . As its name suggests , ParaDiS is inspired by slimmable neural networks framework ( Yu et al. , 2019 ) , where the full model variants ( or so-called switches ) consist of the model sub-networks of different widths , and all variants are trained jointly while sharing the parameters . For example , the switch denoted as 0.25× is a sub-network with each layer composed of the first 25 % of the channels of the full network . Similarly , in ParaDiS framework we consider several variants or switches . However , ParaDiS switches include configurations distributable on several devices as illustrated on Figure 2 . As such , we represent each configuration or switch as a list of widths . For example , [ 0.5 , 0.5 ] × denotes a configuration consisting of two parallel networks extracted from the first and the second halves of channels of the full network , respectively . All ParaDiS model switches are trained jointly , while strongly sharing most of their parameters . Once a set of available devices is known , a suitable configuration may be selected and instantly deployed . Moreover , if a copy of the full model is available on each device , the configuration may be changed instantly without re-loading a new model with a new configuration . For example , if one device has got more computational resources for some reason ( e.g. , some other process has finished ) configuration [ 0.5 , 0.25 , 0.25 ] × may be instantly changed to [ 0.5 , 0.5 ] × by simply changing the sub-models executed on each device . Finally , similar to slimmable framework , ParaDiS framework is applicable for most modern convolutional neural network ( CNN ) architectures and for different tasks ( e.g. , classification , detection , identification , image restoration and super-resolution ) . 1The term “ near-optimally ” means here that each distributed configuration makes profit of all available resources and that it performs on par with a similar configuration trained specifically for this setup . It has been noticed in ( Yu et al. , 2019 ) that slimmable network switches may share all the parameters except the global statistics of batch normalization ( BN ) blocks . Individual BN statistics are adopted instead , which does not lead to any serious drawback . Indeed , individual BN statistics may be computed within a simple calibration phase after the model is trained ( Yu & Huang , 2019 ) , and they represent a very tiny overhead in terms of the total number of parameters . Since ParaDiS is an extension of the slimmable framework , we adopt in this work the same strategy of individual BN statistics . As for training , while it is inspired by the training procedure for universally slimmable ( US ) networks ( Yu & Huang , 2019 ) , where all switches are trained jointly , the full model is trained from the data and all other switches are distilled via knowledge distillation ( KD ) ( Hinton et al. , 2015 ) from the full model . We introduce two important differences . First , instead of distilling the knowledge from the full model ( switch [ 1.0 ] × ) we distill it from a wider model ( e.g. , switch [ 1.2 ] × ) called simply wide model hereafter . Second , we change the KD procedure by distilling the feature activation maps before the final fully connected layer in addition to the output predicted by the wide model . We have shown that these new ingredients are important for training ParaDiS . We investigate the proposed ParaDiS framework on ImageNet classification task using MobileNet v1 ( Howard et al. , 2017 ) and ResNet-50 ( He et al. , 2016 ) as underlying architectures and on image super-resolution task using WDSR architecture ( Yu et al. , 2020a ) . We compare ParaDiS models with the corresponding parallelly distributed configurations trained individually and with the US models ( Yu & Huang , 2019 ) that are not distributable . First , we find that ParaDiS model performs as good as and in many cases better than the individually trained distributed configurations . Second , we observe that distributable ParaDiS switches perform almost as good as non-distributable US model switches of the same overall complexity . Third , we show that , once distributed overs several devices , ParaDiS outperforms greatly the US models . Finally , we conduct an exhaustive ablation study to show the importance of knowledge distillation , using a wide model and distilling activations . 2 RELATED WORK . Compressed and flexible models . Model compression schemes ( Cheng et al. , 2018 ) ( e.g. , pruning ( LeCun et al. , 1990 ) or knowledge distillation ( Hinton et al. , 2015 ) ) allow satisfying particular memory and processing constraints , though they do only allow producing a fixed model for every particular setting . To go further , a new family of so-called flexible approaches allowing to instantly trade off between the accuracy and efficiency was introduced . These approaches include among others early-exit models ( Huang et al. , 2018 ) , slimmable neural networks ( Yu et al. , 2019 ; Yu & Huang , 2019 ) , Once-for-All ( OFA ) networks ( Cai et al. , 2019 ; Sahni et al. , 2021 ) , and neural mixture models ( Ruiz & Verbeek , 2020 ) . Moreover , flexible models like OFA networks have been shown to be useful for efficient neural architecture search ( Yu et al. , 2020b ) and for developing dynamic inference approaches ( Li et al. , 2021 ) . However , unlike our approach , all these approaches are within one device non-distributed settings . Sequentially distributed models . In Neurosurgeon ( Kang et al. , 2017 ) , the authors propose a dynamic partitioning scheme of an existing deep neural network between a client ( end-user device , for example ) and a distant server ( edge or cloud ) ( see also Fig . 1 ( B ) ) . The neural network can only be split between layers to form a head and a tail and the optimal partition is obtained by estimating the execution time of each part on the different devices . The main drawback of this approach is the added communication time occurring between the different partitions . To this end other optimization strategies have been developed and intermediate feature compression schemes have been proposed ( Matsubara et al. , 2021 ) . However , as we have discussed above , distributing a model sequentially does not allow making use of a plurality of devices to reduce processing latency . Parallelly distributed models . In contrast to sequential distribution , distributing model parallelly with no communication ( see Fig . 1 ( D ) ) allows reducing processing latency without increasing communication burden at the same time . Several recently proposed approaches ( Bhardwaj et al. , 2019 ; Asif et al. , 2019 ) suggest creating such efficient parallel models via knowledge transfer from a bigger single pre-trained model . Wang et al . ( 2020 ) revisit model ensembles that are parallellizable as well , and show that they are almost as efficient as individual non-distributable models of the same complexity . However , in contrast to our proposal , none of those models are sharing parameters .

Summary: The paper proposed an algorithm that distill from a larger teacher to a set of parallely distributable student models each of which only compute a fraction of channels in the network. The combined output of the student models can be used to compute the final class prediction.  In this way, student models can be deployed in a set of mobile or edge devices.