Enhancing Early Exit Performance With Uncertainty-Aware Training in Convolutional Neural Networks for Image Classification

Deep Neural Networks have become more and more complex over the past years delivering efficient results but demanding more requirements for resources. The traditional cloud-based DNN inference produces poor real-time performance due to its high latency requirement in the network. Nowadays, fog and edge have been popular, where computation occurs near the data source or between the source and the cloud. High usage of computational resources by DNN incurs limited room for its deployment in such resource-constrained devices. Deep Neural Networks with multiple exit architecture are an optimal approach to preserve time and computational resources by predicting results at an early stage using multiple early exits to support Edge and Fog Intelligence. Early exit, although lightweight and energy efficient, mostly exhibits degraded performance compared to later exits due to the count of samples that are misclassified. To overcome the downfall, we propose uncertainty estimation in the loss function as an additional parameter for classification. Observational results on CIFAR, MNIST, human lung CT scan, and brain MRI datasets using VGG-16, Lenet-5, and Resnet architectures exhibit that our proposed model is better than the existing models to improve performance as well as prediction confidence in Deep Learning Models with early exits.