Jointly-Learned Exit and Inference for a Dynamic Neural Network : JEI-DNN
This work addresses the resource inefficiency of large pretrained models for practical applications, offering an incremental improvement in dynamic neural network training methods.
The paper tackles the challenge of training early-exiting dynamic neural networks (EDNNs) by addressing the decoupling of gating mechanisms and inference modules, which causes train-test mismatch and bias in uncertainty estimation. It proposes a novel architecture that jointly learns these components, resulting in significant performance improvements on classification datasets and enhanced uncertainty characterization.
Large pretrained models, coupled with fine-tuning, are slowly becoming established as the dominant architecture in machine learning. Even though these models offer impressive performance, their practical application is often limited by the prohibitive amount of resources required for every inference. Early-exiting dynamic neural networks (EDNN) circumvent this issue by allowing a model to make some of its predictions from intermediate layers (i.e., early-exit). Training an EDNN architecture is challenging as it consists of two intertwined components: the gating mechanism (GM) that controls early-exiting decisions and the intermediate inference modules (IMs) that perform inference from intermediate representations. As a result, most existing approaches rely on thresholding confidence metrics for the gating mechanism and strive to improve the underlying backbone network and the inference modules. Although successful, this approach has two fundamental shortcomings: 1) the GMs and the IMs are decoupled during training, leading to a train-test mismatch; and 2) the thresholding gating mechanism introduces a positive bias into the predictive probabilities, making it difficult to readily extract uncertainty information. We propose a novel architecture that connects these two modules. This leads to significant performance improvements on classification datasets and enables better uncertainty characterization capabilities.