Balancing training time vs. performance with Bayesian Early Pruning
Authors: Mohit Rajpal, Yehong Zhang, Bryan Kian Hsiang Low
Keywords: Efficient Training, Multi-Output Gaussian Process, Gaussian Process, Bayesian, Single-shot network pruning, Dynamic Sparse Reparameterization, Lottery Ticket Hypothesis
Abstract: Pruning is an approach to alleviate overparameterization of deep neural networks (DNN) by zeroing out or pruning DNN elements with little to no efficacy at a given task. In contrast to related works that do pruning before or after training, this paper presents a novel method to perform early pruning of DNN elements (e.g., neurons or convolutional filters) during the training process while preserving performance upon convergence. To achieve this, we model the future efficacy of DNN elements in a Bayesian manner conditioned upon efficacy data collected during the training and prune DNN elements which are predicted to have low efficacy after training completion. Empirical evaluations show that the proposed Bayesian early pruning improves the computational efficiency of DNN training with small sacrifices in performance. Using our approach we are able to achieve a $48.6\%$ faster training time for ResNet-$50$ on ImageNet to achieve a validation accuracy of $72.5\%$.