Learning Sparse Neural Networks through $L_0$ Regularization
This addresses the challenge of model sparsity for efficient neural network training and deployment, representing a novel method for a known bottleneck rather than a paradigm shift.
The paper tackles the problem of learning sparse neural networks by proposing a practical method for L0 regularization that prunes weights to exactly zero during training, which speeds up training/inference and improves generalization. They achieve this by introducing stochastic gates with a hard concrete distribution that makes the expected L0 norm differentiable, allowing joint optimization with network parameters.
We propose a practical method for $L_0$ norm regularization for neural networks: pruning the network during training by encouraging weights to become exactly zero. Such regularization is interesting since (1) it can greatly speed up training and inference, and (2) it can improve generalization. AIC and BIC, well-known model selection criteria, are special cases of $L_0$ regularization. However, since the $L_0$ norm of weights is non-differentiable, we cannot incorporate it directly as a regularization term in the objective function. We propose a solution through the inclusion of a collection of non-negative stochastic gates, which collectively determine which weights to set to zero. We show that, somewhat surprisingly, for certain distributions over the gates, the expected $L_0$ norm of the resulting gated weights is differentiable with respect to the distribution parameters. We further propose the \emph{hard concrete} distribution for the gates, which is obtained by "stretching" a binary concrete distribution and then transforming its samples with a hard-sigmoid. The parameters of the distribution over the gates can then be jointly optimized with the original network parameters. As a result our method allows for straightforward and efficient learning of model structures with stochastic gradient descent and allows for conditional computation in a principled way. We perform various experiments to demonstrate the effectiveness of the resulting approach and regularizer.