PTEENet: Post-Trained Early-Exit Neural Networks Augmentation for Inference Cost Optimization

TL;DR

This paper introduces PTEENet, a method for adding early-exit branches to pre-trained deep neural networks, enabling significant inference cost reduction with controllable accuracy trade-offs.

Contribution

It extends existing early-exit methods by attaching branches to pre-trained models without altering original weights and introduces a new branch architecture with confidence heads.

Findings

Reduces inference computational cost on image datasets.

Allows real-time control of speed-accuracy tradeoff.

Effective across various DNN architectures like ResNet, DenseNet, VGG.

Abstract

For many practical applications, a high computational cost of inference over deep network architectures might be unacceptable. A small degradation in the overall inference accuracy might be a reasonable price to pay for a significant reduction in the required computational resources. In this work, we describe a method for introducing "shortcuts" into the DNN feedforward inference process by skipping costly feedforward computations whenever possible. The proposed method is based on the previously described BranchyNet (Teerapittayanon et al., 2016) and the EEnet (Demir, 2019) architectures that jointly train the main network and early exit branches. We extend those methods by attaching branches to pre-trained models and, thus, eliminating the need to alter the original weights of the network. We also suggest a new branch architecture based on convolutional building blocks to allow enough training capacity when applied on large DNNs. The proposed architecture includes confidence heads that are used for predicting the confidence level in the corresponding early exits. By defining adjusted thresholds on these confidence extensions, we can control in real-time the amount of data exiting from each branch and the overall tradeoff between speed and accuracy of our model. In our experiments, we evaluate our method using image datasets (SVHN and CIFAR10) and several DNN architectures (ResNet, DenseNet, VGG) with varied depth. Our results demonstrate that the proposed method enables us to reduce the average inference computational cost and further controlling the tradeoff between the model accuracy and the computation cost.