Finding the Optimal Network Depth in Classification Tasks
This addresses the need for efficient, lightweight models in resource-constrained applications, though it appears incremental as it builds on existing pruning and multi-head classifier techniques.
The paper tackles the problem of reducing neural network size and speeding up inference by developing a method to find the optimal depth, resulting in significant parameter reduction and faster inference across hardware.
We develop a fast end-to-end method for training lightweight neural networks using multiple classifier heads. By allowing the model to determine the importance of each head and rewarding the choice of a single shallow classifier, we are able to detect and remove unneeded components of the network. This operation, which can be seen as finding the optimal depth of the model, significantly reduces the number of parameters and accelerates inference across different hardware processing units, which is not the case for many standard pruning methods. We show the performance of our method on multiple network architectures and datasets, analyze its optimization properties, and conduct ablation studies.