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 Zero Time Waste (ZTW), a novel approach to leverage previous intermediate classifiers (ICs) for early-exit neural networks. ZTW consists of two simple extensions of early-exits NNs. The first is to add cascaded connections from the result of previous ICs to the next IC. The second is to ensemble previous results for confidence measurement. The author conducts extensive experiments on image classification benchmarks and shows improved performance compared to vanilla early-exit neural networks and PABEE, a previous work investigating using results from previous ICs for confidence measurement.