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-based conditional computation for reducing the inference time of a given neural network. The approach works as follows: E2CM saves the mean activation of each class after each layer and then, during inference, compares the activation of a given example to the recorded mean activations to get a prediction and its certainty. If the certainty is high enough, this prediction is returned, otherwise, the next layer is queried instead.  The authors show empirically that the method outperforms other approaches in a setting with a very limited training budget, and is sometimes able to improve the performance of other methods when applied in combination. Finally, the method is also presented in an unsupervised learning setting where it allows for significant computation savings.