UCBEE: A Multi Armed Bandit Approach for Early-Exit in Neural Networks

Abstract

Deep Neural Networks (DNNs) have demonstrated exceptional performance in diverse tasks. However, deploying DNNs on resource-constrained devices presents challenges due to energy consumption and delay overheads. To mitigate these issues, early-exit DNNs (EE-DNNs) incorporate exit branches within intermediate layers to enable early inferences. These branches estimate prediction confidence and employ a fixed threshold to determine early termination. Nonetheless, fixed thresholds yield suboptimal performance in dynamic contexts, where context refers to distortions caused by environmental conditions, in image classification, or variations in input distribution due to concept drift, in NLP. In this article, we introduce Upper Confidence Bound in EE-DNNs (UCBEE), an online algorithm that dynamically adjusts early exit thresholds based on context. UCBEE leverages confidence levels at intermediate layers and learns without the need for true labels. Through extensive experiments in image classification and NLP, we demonstrate that UCBEE achieves logarithmic regret, converging after just a few thousand observations across multiple contexts. We evaluate UCBEE for image classification and text mining. In the latter, we show that UCBEE can reduce cumulative regret and lower latency by approximately 10%–20% without compromising accuracy when compared to fixed threshold alternatives. Our findings highlight UCBEE as an effective method for enhancing EE-DNN efficiency.