Neuron-based Pruning of Deep Neural Networks with Better Generalization using Kronecker Factored Curvature Approximation

TL;DR

This paper introduces a neuron pruning method that trains and prunes simultaneously, using Kronecker-factored curvature approximation to find flatter solutions, leading to better generalization and smaller networks.

Contribution

It presents a novel neuron pruning algorithm that improves generalization by exploring spectral radius of Hessian during training, without pre-training, outperforming existing methods.

Findings

Achieves smaller networks with minimal accuracy loss.

Outperforms state-of-the-art neuron compression techniques.

Enhances model generalization on unseen data.

Abstract

Existing methods of pruning deep neural networks focus on removing unnecessary parameters of the trained network and fine tuning the model afterwards to find a good solution that recovers the initial performance of the trained model. Unlike other works, our method pays special attention to the quality of the solution in the compressed model and inference computation time by pruning neurons. The proposed algorithm directs the parameters of the compressed model toward a flatter solution by exploring the spectral radius of Hessian which results in better generalization on unseen data. Moreover, the method does not work with a pre-trained network and performs training and pruning simultaneously. Our result shows that it improves the state-of-the-art results on neuron compression. The method is able to achieve very small networks with small accuracy degradation across different neural network models.