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 proposes ParaDiS, a slimmable network that can be executed across multiple devices, and be transferred across different devices without re-training. To train ParaDis, the authors propose to distill from a wider network, and combine activation knowledge distillation to improve the performance. While the idea of parallel execution is promising, there is little detailed discussion on that point.