E$^2$CM: Early Exit via Class Means for Efficient Supervised and Unsupervised Learning

State-of-the-art neural networks with early exit mechanisms often need considerable amount of training and fine-tuning to achieve good performance with low computational cost . We propose a novel early exit technique , E2CM , based on the class means of samples . Unlike most existing schemes , E2CM does not require gradient-based training of internal classifiers . This makes it particularly useful for neural network training in low-power devices , as in wireless edge networks . In particular , given a fixed training time budget , E2CM achieves higher accuracy as compared to existing early exit mechanisms . Moreover , if there are no limitations on the training time budget , E2CM can be combined with an existing early exit scheme to boost the latter ’ s performance , achieving a better trade-off between computational cost and network accuracy . We also show that E2CM can be used to decrease the computational cost in unsupervised learning tasks . 1 INTRODUCTION . Modern deep learning models require a vast amount of computational resources to effectively perform various tasks such as object detection ( Szegedy et al. , 2015 ) , image classification ( He et al. , 2016 ) , machine translation , ( Vaswani et al. , 2017 ) and text generation ( Brown et al. , 2020 ) . Deploying deep learning models to the edge , such as to mobile phones or the Internet of Things ( IoT ) , thus becomes particularly challenging due to device computation and energy limitations ( Li et al. , 2021a ; b ) . Moreover , the law of diminishing returns applies to the computation-performance tradeoff ( Bolukbasi et al. , 2017 ) : The increase in a deep learning model ’ s performance is often marginal as compared to the increase in the amount of computation . One of the primary reasons behind traditional deep learning models ’ high computation demand is their tunnel-like design . In fact , traditional models apply the same sequence of operations to any given input . However , in many real world datasets , certain inputs may consist of much simpler features as compared to other inputs ( Bolukbasi et al. , 2017 ) . In such a scenario , it becomes desirable to design more efficient architectures that can exploit the heterogeneous complexity of dataset members . This can be achieved by introducing additional exit points to the models ( Panda et al. , 2016 ; Teerapittayanon et al. , 2016 ; Kaya et al. , 2019 ) . These exit points prevent simple inputs to traverse the entire network , reducing the computational cost of inference . Despite reduced inference time , existing early exit neural network architectures require additional training and fine-tuning for the early exit points , which increases the training time ( Teerapittayanon et al. , 2016 ; Bolukbasi et al. , 2017 ; Kaya et al. , 2019 ) . This side-effect is undesirable for scenarios in which the training has to be done in a low-power device . An ideal solution is a plug-and-play approach that does not require gradient-based training and performs well . In this work , we propose such an early-exit mechanism , Early Exit Class Means ( E2CM ) , based on the class means of input samples for the image classification task . By taking the mean of layer outputs for each class at every layer of the model , class means are obtained . During inference , output of a layer is compared with the corresponding class means using Euclidean distance as the metric . If the output of the layer is close enough to a class mean , the execution is stopped and the sample exits the network . In fact , as seen in Fig . 1 , some samples can be classified easily at early stages of the network by just considering a “ nearest class mean ” decision rule , suggesting the potential effectiveness of our method for reducing computational cost . A practical use case for E2CM is when a large , expensive-to-train model is broadcast to edge devices with limited and heterogeneous computation capabilities . In such a scenario , different devices may train the E2CM model at different FLOP operating points depending on their computation capabilities or power limitations . Another practical scenario for E2CM may be transfer learning , where the training has to be done on a low power edge device with a local dataset that is different from the dataset the base network was trained on . In order to keep the battery usage minimum on the edge device , the training time may be limited . The advantage of E2CM is that it does not need any gradient-based training like Shallow-Deep Networks ( Kaya et al. , 2019 ) , therefore it is more suitable for transfer learning on low power devices . One other practical scenario for E2CM may be transfer learning , where the training has to be done on low power edge device with a local dataset that is different from the dataset the base network was trained on . In order to keep the battery usage minimum on the edge device , the training time may be limited . The advantage of E2CM is that it does not need any gradient-based training like Shallow-Deep Networks ( Kaya et al. , 2019 ) . To the best of our knowledge , E2CM is the first early exit mechanism that does not require a gradientbased training and does not modify the original network by any means . Furthermore , it does not have a hyper-parameter for the early exit locations unlike existing schemes . These features make E2CM simpler to use and easier to deploy on low-power devices . While using class means as the only early exit mechanism requires just a single feed forward pass , existing early exit methods reaches the same performance after training for multiple epochs , which suggests our method is more agile and powerful yet simpler . Moreover , combining E2CM with the existing mechanisms that require gradient-based training achieves a better trade-off in terms of computation cost and network accuracy . In addition to its benefits in the supervised learning setting , E2CM is also the first technique that shows the feasibility of early exits in the unsupervised learning setting . We show the effectiveness of E2CM on CIFAR-10 , CIFAR-100 ( Krizhevsky et al. , 2009 ) and KMNIST ( Clanuwat et al. , 2018 ) datasets using ResNet-18 , ResNet-152 ( He et al. , 2016 ) and WideResNet-101 models ( Zagoruyko & Komodakis , 2016 ) . Using class means as the sole decision mechanism for early exiting results in 35 % better accuracy or 50 % faster inference time compared to the current state-of-the-art early exit technique . When combined with the state-of-the-art , we increase the accuracy by 9 % without doing further computation ; or decrease the inference time by 25 % without a loss in accuracy . We also show that it is possible to decrease the computational cost while doing clustering with autoencoders on MNIST ( LeCun et al. , 1998 ) . In particular , E2CM saves on computation by 60 % , while the loss in unsupervised clustering accuracy ( Xie et al. , 2016 ) is marginal . 2 RELATED WORK . Our work is related to the area of conditional computation ( Bengio et al. , 2013 ) , where several small networks are trained to control the computation flow of one deep neural network . For this purpose , adding gates between the blocks of residual networks have been proposed ( Veit & Belongie , 2018 ; Wang et al. , 2018 ) . During inference , these gates allow the input to skip unnecessary blocks , thus saving computation time . However , unlike our proposed method , the gated network has to be trained from scratch . Conditional computation policies can also be learned through reinforcement learning ( Bengio et al. , 2015 ; Jie et al. , 2019 ) . However , this latter approach forces the inputs to go through the entire network as it does not incorporate early exit points ( Bengio et al. , 2015 ) . One of the earliest works that explicitly propose the idea of early exiting is Panda et al . ( 2016 ) , where the authors consider adding a cascade of linear layers after convolutional layers as control blocks . Rather than just linear layers , adding branches consisting of convolutional layers to the original model has also been studied ( Teerapittayanon et al. , 2016 ) . A significant drawback of this method is that the branches may increase the computational cost due to convolutional layers . In addition , this idea requires branches to be trained with the original model jointly , from scratch . In a more recent study , internal classifiers ( ICs ) consisting of a feature reduction layer and a single linear layer are added after certain layers in the network ( Kaya et al. , 2019 ) . Hence , the methods presented in these studies modify the original network by adding linear or convolutional layers . Moreover , they require gradient updates to train those layers . Also , an implicit hyper-parameter is the locations of the early exit points . E2CM is better suited for low-power applications compared to existing studies since we do not modify the original model , do not require gradient based training and additional hyper-parameters . In addition to layer level early exits , a network level early exit mechanism has been introduced in Bolukbasi et al . ( 2017 ) . Both the layer level exit and the network level exit require decision functions to be inserted between the layers and networks . This type of architecture freezes the weights of the original model , and then trains the decision functions one by one using weighted binary classification . The drawback of this approach is such an alternative optimization , which may consume a lot of time and energy when there are many decision functions to optimize . The idea of early exit neural networks shows promising results in natural language processing domain too . Adding early exit points between the transformer layers reduce the computation time with marginal loss in performance ( Zhou et al. , 2020 ; Xin et al. , 2020 ) . Applications for neural network based early exit mechanisms to computer fault management have also been considered ( Biasielli et al. , 2020 ) . Intermediate layer outputs have been used for classification in few-shot and one-shot learning settings in the past ( Koch et al. , 2015 ; Vinyals et al. , 2016 ; Snell et al. , 2017 ; Sung et al. , 2018 ) . These studies are closely related with the area of metric learning . The closest work to E2CM is prototypical networks , in which prototypes for each class are computed ( Pan et al. , 2019 ) . However , none of those approaches aim to reduce the computation cost . E2CM is also related to the phenomenon of neural collapse ( Papyan et al. , 2020 ) . It is known that as the inputs go deeper in a neural network , the classes are separated better from each other as a result of multiple nonlinearities , and the samples begin to concentrate ( Cohen et al. , 2020 ; Zarka et al. , 2020 ) . E2CM exploits this phenomenon with the main idea of stopping the execution as soon as the sample is close enough to a concentration point , i.e. , a class mean . One other idea to reduce the computational cost of neural networks during inference time in the literature is the usage of multiresolution inputs and networks ( Huang et al. , 2017 ; Yang et al. , 2020 ; McGill & Perona , 2017 ) . While this idea works well , it needs subnetworks and the usage of multiresolution inputs , which means it requires more hyperparameters compared to E2CM .

Summary: The paper proposes a method for early exit during inference of DNNs. The method does not require fine-tuning or gradient computation. Rather, it uses the learned embeddings from intermediate model layers to decide where to terminate the execution. Specifically, for a trained model, the method calculates the per-class means of intermediate activation. The distance between the activations of new samples is then computed with the per-class means and inference is exited once the distance becomes lower than a set threshold.