Complexity-Aware Training of Deep Neural Networks for Optimal Structure Discovery

We propose a novel algorithm for combined unit and layer pruning of deep neural networks that functions during training and without requiring a pre-trained network to apply. Our algorithm employs three user-defined hyper-parameters that are easy to interpret as i) controlling the general pruning level, ii) regulating layer vs. unit pruning and iii) trading-off computational load vs. the number of parameters of the network and therefore easy to tune. Optimal network structures for a given set of these hyperparameters are found as the solution to a stochastic optimization problem over the network weights and the parameters of variational Bernoulli distributions for binary Random Variables taking values either 0 or 1 and scaling the units and layers of the network. Pruning occurs when a variational parameter converges to 0 rendering the corresponding structure permanently inactive, thus saving computations both during training and prediction. A key contribution of our approach is to define a cost function that combines the objectives of prediction accuracy and network pruning in a computational/parameter complexity-aware manner and the automatic selection of the many regularization parameters. We show that the proposed algorithm converges to solutions of the optimization problem corresponding to deterministic networks. We analyze the ODE system that underlies our stochastic optimization algorithm and establish domains of attraction for the dynamics of the network parameters. These theoretical results lead to practical pruning conditions avoiding the premature pruning of units on each layer as well as whole layers during training. We evaluate our method on the CIFAR-10/100 and ImageNet datasets using ResNet and VGG architectures and demonstrate that it gives improved results with respect to pruning ratios and test accuracy over layer-only or unit-only pruning and favorably competes with combined unit and layer pruning algorithms requiring pre-trained networks.