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

TL;DR

PTEE addresses the practical constraint of high inference cost by attaching post-trained early-exit branches to pre-trained backbones, avoiding changes to backbone weights. The method trains only the added branches using pseudo-labels from the backbone and optimizes a joint loss that balances classification accuracy and computational cost, with a single threshold $T$ serving as a real-time knob for speed-accuracy tradeoffs. Across SVHN and CIFAR-10 with ResNet, DenseNet, and VGG backbones, PTEE achieves meaningful reductions in average FLOPs (up to about 42%) with small accuracy losses, illustrating its practicality for edge and resource-constrained deployments. The work highlights the potential to combine post-training exits with other optimization techniques, and suggests extensions such as per-branch thresholds and incremental training to further enhance performance and efficiency.

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.

Figures (4)

• Figure 1: ResNet20 model attached by 3 early exit branches. Each branch allows for early termination of samples propagation by applying confidence head threshold.
• Figure 2: ConvX branch architecture. Using both block0 and block1 defines Conv2 branch architecture. Confidence head output is used for early termination decision making by applying threshold T on its output value. Higher that threshold T values terminate sample propagation and the use of the attached classifier prediction.
• Figure 3: Computational cost reduction and validation accuracy pairs generated from increasing levels of $\mathbf{\lambda}$, each used to train ResEEnet110, VGGEEnet19 and DenseEEnet121 models with various number of branches attached. $\mathbf{\lambda}$ values correspond to the blue curve of ResEEnet110.
• Figure 4: Computational cost reduction and validation accuracy pairs generated from varied levels of confidence threshold T, for ResEEnet110 model with 10 branches attached. T=0.5 maintain accuracy of 97.15% while gaining computational cost reduction of 27%