Zero Time Waste: Recycling Predictions in Early Exit Neural Networks

1 Introduction . Deep learning models achieve tremendous successes across a multitude of tasks , yet their training and inference often yield high computational costs and long processing times [ 11 , 22 ] . For some applications , however , efficiency remains a critical challenge , e.g . to deploy a reinforcement learning ( RL ) system in production the policy inference must be done in real-time [ 7 ] , while the robot performances suffer from the delay between measuring a system state and acting upon it [ 34 ] . Similarly , long inference latency in autonomous cars could impact its ability to control the speed [ 13 ] and lead to accidents [ 10 , 17 ] . Typical approaches to reducing the processing complexity of neural networks in latency-critical applications include compressing the model [ 24 , 26 , 46 ] or approximating its responses [ 21 ] . For instance , Livne & Cohen [ 26 ] propose to compress a RL model by policy pruning , while Kouris et al . [ 21 ] approximate the responses of LSTM-based modules in self-driving cars to accelerate their inference time . While those methods improve processing efficiency , they still require samples to pass ∗equal contribution †Corresponding author : maciej.wolczyk @ doctoral.uj.edu.pl 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . ( a ) Comparison of the proposed ZTW ( bottom ) with a conventional early-exit model , SDN ( top ) . ( b ) Detailed scheme of the proposed ZTW model architecture . Figure 1 : ( a ) In both approaches , internal classifiers ( ICs ) attached to the intermediate hidden layers of the base network allow us to return predictions quickly for examples that are easy to process . While SDN discards predictions of uncertain ICs ( e.g . below a threshold of 75 % ) , ZTW reuses computations from all previous ICs , which prevents information loss and waste of computational resources . ( b ) Backbone network fθ lends its hidden layer activations to ICs , which share inferred information using cascade connections ( red horizontal arrows in the middle row ) and give predictions pm . The inferred predictions are combined using ensembling ( bottom row ) giving qm . through the entire model . In contrast , biological neural networks leverage simple heuristics to speed up decision making , e.g . by shortening the processing path even in case of complex tasks [ 1 , 9 , 18 ] . This observation led a way to the inception of the so-called early exit methods , such as Shallow-Deep Networks ( SDN ) [ 19 ] and Patience-based Early Exit ( PBEE ) [ 47 ] , that attach simple classification heads , called internal classifiers ( ICs ) , to selected hidden layers of neural models to shorten the processing time . If the prediction confidence of a given IC is sufficiently high , the response is returned , otherwise , the example is passed to the subsequent classifier . Although these models achieve promising results , they discard the response returned by early ICs in the evaluation of the next IC , disregarding potentially valuable information , e.g . decision confidence , and wasting computational effort already incurred . Motivated by the above observation , we postulate to look at the problem of neural model processing efficiency from the information recycling perspective and introduce a new family of zero waste models . More specifically , we investigate how information available at different layers of neural models can contribute to the decision process of the entire model . To that end , we propose Zero Time Waste ( ZTW ) , a method for an intelligent aggregation of the information from previous ICs . A high-level view of our model is given in Figure 1 . Our approach relies on combining ideas from networks with skip connections [ 41 ] , gradient boosting [ 3 ] , and ensemble learning [ 8 , 23 ] . Skip connections between subsequent ICs ( which we call cascade connections ) allow us to explicitly pass the information contained within low-level features to the deeper classifier , which forms a cascading structure of ICs . In consequence , each IC improves on the prediction of previous ICs , as in gradient boosting , instead of generating them from scratch . To give the opportunity for every IC to explicitly reuse predictions of all previous ICs , we additionally build an ensemble of shallow ICs . We evaluate our approach on standard classification benchmarks , such as CIFAR-100 and ImageNet , as well as on the more latency-critical applications , such as reinforcement-learned models for interacting with sequential environments . To the best of our knowledge , we are the first to show that early exit methods can be used for cutting computational waste in a reinforcement learning setting . Results show that ZTW is able to save much more computation while preserving accuracy than current state-of-the-art early exit methods . In order to better understand where the improvements come from , we introduce Hindsight Improvability , a metric for measuring how efficiently the model reuses information from the past . We provide ablation studies and additional analysis of the proposed method in the Appendix . To summarize , the contributions of our work are the following : • We introduce a family of zero waste models that quantify neural network efficiency with the Hindsight Improvability metrics . • We propose an instance of zero waste models dubbed Zero Time Waste ( ZTW ) method which uses cascade connections and ensembling to reuse the responses of previous ICs for the final decision . • We show how the state-of-the-art performance of ZTW in the supervised learning scenario generalizes to reinforcement learning . 2 Related Work . The drive towards reducing computational waste in deep learning literature has so far focused on reducing the inference time . Numerous approaches for accelerating deep learning models focus on building more efficient architectures [ 15 ] , reducing the number of parameters [ 12 ] or distilling knowledge to smaller networks [ 14 ] . Thus , they decrease inference time by reducing the overall complexity of the model instead of using the conditional computation framework of adapting computational effort to each example . As such we find them orthogonal to the main ideas of our work , e.g . we show that applying our method to architectures designed for efficiency , such as MobileNet [ 15 ] , leads to even further acceleration . Hence , we focus here on methods that adaptively set the inference time for each example . Conditional Computation Conditional computation was first proposed for deep neural networks in Bengio et al . [ 2 ] and Davis & Arel [ 5 ] , and since then many sophisticated methods have been proposed in this field , including dynamic routing [ 27 ] , cascading with multiple networks [ 40 ] and skipping intermediate layers [ 41 ] or channels [ 42 ] . In this work , we focus on the family of early exit approaches , as they usually do not require special assumptions about the underlying architecture of the network and the training paradigm , and because of that can be easily applied to many commonly used architectures . In BranchyNet [ 38 ] a loss function consisting of a weighted sum of individual head losses is utilized in training , and entropy of the head prediction is used for the early exit criterion . Berestizshevsky & Guy [ 4 ] propose to use confidence ( maximum of the softmax output ) instead . A broader overview of early exit methods is available in Scardapane et al . [ 32 ] . Several works proposed specialized architectures for conditional computation which allow for multiscale feature processing [ 16 , 44 , 43 ] , and developed techniques to train them more efficiently by passing information through the network [ 29 , 25 ] . However , in this paper , we consider the case of increasing inference speed of a pre-trained network based on an architecture which was not built with conditional computation or even efficiency in mind . We argue that this is a practical use case , as this approach can be used to a wider array of models . As such , we do not compare with these methods directly . Shallow-Deep Networks ( SDN ) [ 19 ] is a conceptually simple yet effective method , where the comparison of confidence with a fixed threshold is used as the exit criterion . The authors attach internal classifiers to layers selected based on the number of compute operations needed to reach them . The answer of each head is independent of the answers of the previous heads , although in a separate experiment the authors analyze the measure of disagreement between the predictions of final and intermediate heads . Zhou et al . [ 47 ] propose Patience-based Early Exit ( PBEE ) method , which terminates inference after t consecutive unchanged answers , and show that it outperforms SDN on a range of NLP tasks . The idea of checking for agreement in preceding ICs is connected to our approach of reusing information from the past . However , we find that applying PBEE in our setting does not always work better than SDN . Additionally , in the experiments from the original work , PBEE was trained simultaneously along with the base network , thus making it impossible to preserve the original pre-trained model . Ensembles Ensembling is typically used to improve the accuracy of machine learning models [ 6 ] . Lakshminarayanan et al . [ 23 ] showed that it also greatly improves calibration of deep neural networks . There were several attempts to create an ensemble from different layers of a network . Scardapane et al . [ 31 ] adaptively exploit outputs of all internal classifiers , albeit not in a conditional computation context . Phuong & Lampert [ 29 ] used averaged answers of heads up to the current head for anytimeprediction , where the computational budget is unknown . Besides the method being much more basic , their setup is notably different from ours , as it assumes the same computational budget for all samples no matter how difficult the example is . Finally , none of the ensemble methods mentioned above were designed to work with pre-trained models . 3 Zero Time Waste . Our goal is to reduce computational costs of neural networks by minimizing redundant operations and information loss . To achieve it , we use the conditional computation setting , in which we dynamically select the route of an input example in a neural network . By controlling the computational route , we can decide how the information is stored and utilized within the model for each particular example . Intuitively , difficult examples require more resources to process , but using the same amount of compute for easy examples is wasteful . Below we describe our Zero Time Waste method in detail . In order to adapt already trained models to conditional computation setting , we attach and train early exit classifier heads on top of several selected layers , without changing the parameters of the base network . During inference , the whole model exits through one of them when the response is likely enough , thus saving computational resources . Formally , we consider a multi-class classification problem , where x ∈ RD denotes an input example and y ∈ { 1 , . . . , K } is its target class . Let fθ : RD → RK be a pre-trained neural network with logit output designed for solving the above classification task . The weights θ will not be modified . Model overview Following typical early exit frameworks , we add M shallow Internal Classifiers , IC1 , . . . , ICM , on intermediate layers of fθ . Namely , let gφm , for m ∈ { 1 , . . . , M } , be the m-th IC network returning K logits , which is attached to hidden layer fθm of the base network fθ . The index m is independent of fθ layer numbering . In general , M is lower than the overall number of fθ hidden layers since we do not add ICs after every layer ( see more details in Appendix A.1 ) . Although using ICs to return an answer early can reduce overall computation time [ 19 ] , in a standard setting each IC makes its decision independently , ignoring the responses returned by previous ICs . As we show in Section 4.2 , early layers often give correct answers for examples that are misclassified by later classifiers , and hence discarding their information leads to waste and performance drops . To address this issue , we need mechanisms that collect the information from the first ( m− 1 ) ICs to inform the decision of ICm . For this purpose , we introduce two complementary techniques : cascade connections and ensembling , and show how they help reduce information waste and , in turn , accelerate the model . Cascade connections directly transfer the already inferred information between consecutive ICs instead of re-computing it again . Thus , they improve the performance of initial ICs that lack enough predictive power to classify correctly based on low-level features . Ensembling of individual ICs improves performance as the number of members increases , thus showing greatest improvements in the deeper part of the network . This is visualized in Figure 1 where cascade connections are used first to pass already inferred information to later ICs , while ensembling is utilized to conclude the IC prediction . The details on these two techniques are presented in the following paragraphs . Cascade connections Inspired by the gradient boosting algorithm and literature on cascading classifiers [ 39 ] , we allow each IC to improve on the predictions of previous ICs instead of inferring them from scratch . The idea of cascade connections is implemented by adding skip connections that combine the output of the base model hidden layer fθm with the logits of ICm−1 and pass it to ICm . The prediction is realized by the softmax function applied to gφm ( the m-th IC network ) : pm = softmax ( gφm ( fθm ( x ) , gφm−1 ◦ fθm−1 ( x ) ) ) , for m > 1 , ( 1 ) where g ◦ f ( x ) = g ( f ( x ) ) denotes the composition of functions . Formally , pm = pm ( x ; φm ) , where φm are trainable parameters of ICm , but we drop these parameters in notation for brevity . IC1 uses only the information coming from the layer fθ1 which does not need to be the first hidden layer of fθ . Figure 1 shows the skip connections as red horizontal arrows . Each ICm is trained in parallel ( with respect to φm ) to optimize the prediction of all output classes using an appropriate loss function L ( pm ) , e.g . cross-entropy for classification . However , during the backward step it is crucial to stop the gradient of a loss function from passing to the previous classifier . Allowing the gradients of loss L ( pm ) to affect φj for j ∈ 1 , .. , m− 1 leads to a significant performance degradation of earlier layers due to increased focus on the features important for ICm , as we show in Appendix C.3 . Ensembling Ensembling in machine learning models reliably increases the performance of a model while improving robustness and uncertainty estimation [ 8 , 23 ] . The main drawback of this approach is its wastefulness , as it requires to train multiple models and use them to process the same examples . However , in our setup we can adopt this idea to combine predictions which were already pre-computed in previous ICs , with near-zero additional computational cost . To obtain a reliable zero-waste system , we build ensembles that combine outputs from groups of ICs to provide the final answer of the m-th classifier . Since the classifiers we are using vary significantly in predictive strength ( later ICs achieve better performance than early ICs ) and their predictions are correlated , the standard approach to deep model ensembling does not work in our case . Thus , we introduce weighted geometric mean with class balancing , which allows us to reliably find a combination of pre-computed responses that maximizes the expected result . Let p1 , p2 , . . . , pm be the outputs of m consecutive IC predictions ( after cascade connections stage ) for a given x ( Figure 1 ) . We define the probability of the i-th class in the m-th ensemble to be : qim ( x ) = 1 Zm bim ∏ j≤m ( pij ( x ) ) wjm , ( 2 ) where bim > 0 and w j m > 0 , for j = 1 , . . . , m , are trainable parameters , and Zm is a normalization factor , such that ∑ i q i m ( x ) = 1 . Observe that w j m can be interpreted as our prior belief in predictions of ICj , i.e . large weight wjm indicates less confidence in the predictions of ICj . On the other hand , b i m represents the prior of i-th class for ICm . The m indices in wm and bm are needed as the weights are trained independently for each subset { ICj : j ≤ m } . Although there are viable potential approaches to setting these parameters by hand , we verified that optimizing them directly by minimizing the cross-entropy loss on the training dataset works best . Out of additive and geometric ensemble settings we found the latter to be preferable . In this formulation , a low class confidence of a single IC would significantly reduce the probability of that class in the whole ensemble . In consequence , in order for the confidence of the given class to be high , we require all ICs to be confident in that class . Thus , in geometric ensembling , an incorrect although confident IC answer has less chance of ending calculations prematurely . In the additive setting , the negative impact of a single confident but incorrect IC is much higher , as we show in Appendix C.2 . Hence our choice of geometric ensembling . Direct calculation of the product in ( 2 ) might lead to numerical instabilities whenever the probabilities are close to zero . To avoid this problem we note that bim ∏ j≤m ( pij ( x ) ) wjm = bim exp ( ∑ j≤m wjm ln p i j ( x ) ) , and that log-probabilities ln pij can be obtained by running the numerically stable log softmax function on the logits gφm of the classifier . Both cascade connections and ensembling have different impact on the model . Cascade connections primarily boost the accuracy of early ICs . Ensembling , on the other hand , improves primarily the performance of later ICs , which combine the information from many previous classifiers . This is not surprising , given that the power of the ensemble increases with the number of members , provided they are at least weak in the sense of boosting theory [ 33 ] . As such , the two techniques introduced above are complementary , which we also show empirically via ablation studies in Appendix C. The whole training procedure is presented in Algorithm 1 . Conditional inference Once a ZTW model is trained , the following question appears : how to use the constructed system at test time ? More precisely , we need to dynamically find the shortest processing path for a given input example . For this purpose , we use one of the standard confidence Algorithm 1 Zero Time Waste Input : pre-trained model fθ , cross-entropy loss function L , training set T . Initialize M shallow models gφm at selected layers fθm . For m = 1 , . . . , M do in parallel . Cascade connection ICs Set pm according to ( 1 ) . minimize E ( x , y ) ∈T [ L ( pm ( x ) , y ) ] wrt . φm by gradient descent For m = 1 , . . . , M do . Geometric Ensembling Initialize wm , bm and define qm ( x ) according to ( 2 ) . minimize E ( x , y ) ∈T [ L ( qm ( x ) , y ) ] wrt . wm , bm by gradient descent scores given by the probability of the most confident class . If the m-th classifier is confident enough about its prediction , i.e . if max i qim > τ , for a fixed τ > 0 , ( 3 ) where i is the class index , then we terminate the computation and return the response given by this IC . If this condition is not satisfied , we continue processing x and go to the next IC . Threshold τ in ( 3 ) is a manually selected value , which controls the acceleration-performance trade-off of the model . A lower threshold leads to a significant speed-up at the cost of a possible drop in accuracy . Observe that for τ > 1 , we recover the original model fθ , since none of the ICs is confident enough to answer earlier . In practice , to select its appropriate value , we advise using a held-out set to evaluate a range of possible values of τ .

Summary: This paper proposes to use all the outputs from the internal classifiers in an early-exit network. The paper presents two novel ways of making use of the previous prediction and concludes with the weighted geometric mean as a better choice empirically. By taking care of the details, i.e., stop gradient and handling numerical issues during training, the authors report superior results compared to early-exit networks that do not make use of previous classification results in the later classifiers. The paper is well-written with comprehensive empirical results on CIFAR, TinyImageNet with some ImageNet, and Atari.