A Bregman Learning Framework for Sparse Neural Networks

Abstract

We propose a learning framework based on stochastic Bregman iterations, also known as mirror descent, to train sparse neural networks with an inverse scale space approach. We derive a baseline algorithm called LinBreg, an accelerated version using momentum, and AdaBreg, which is a Bregmanized generalization of the Adam algorithm. In contrast to established methods for sparse training the proposed family of algorithms constitutes a regrowth strategy for neural networks that is solely optimization-based without additional heuristics. Our Bregman learning framework starts the training with very few initial parameters, successively adding only significant ones to obtain a sparse and expressive network. The proposed approach is extremely easy and efficient, yet supported by the rich mathematical theory of inverse scale space methods. We derive a statistically profound sparse parameter initialization strategy and provide a rigorous stochastic convergence analysis of the loss decay and additional convergence proofs in the convex regime. Using only 3.4% of the parameters of ResNet-18 we achieve 90.2% test accuracy on CIFAR-10, compared to 93.6% using the dense network. Our algorithm also unveils an autoencoder architecture for a denoising task. The proposed framework also has a huge potential for integrating sparse backpropagation and resource-friendly training.

43 pages, 5 figures, some minor modifications, weakened assumptions

References in corpus (7)

• To prune, or not to prune: exploring the efficacy of pruning for model compression
• Group Sparse Regularization for Deep Neural Networks
• The Expressive Power of Neural Networks: A View from the Width
• Carbontracker: Tracking and Predicting the Carbon Footprint of Training Deep Learning Models
• On Sparsity in Overparametrised Shallow ReLU Networks
• A General Family of Stochastic Proximal Gradient Methods for Deep Learning
• Fast Stochastic Bregman Gradient Methods: Sharp Analysis and Variance Reduction

Cited by in corpus (1)

• Neural Architecture Search via Bregman Iterations