Fast as CHITA: Neural Network Pruning with Combinatorial Optimization
This work addresses the problem of efficient model compression for neural network deployment, offering significant gains in speed, memory, and performance, though it is incremental as it builds on the Optimal Brain Surgeon framework.
The authors tackled the computational challenge of serving large neural networks by proposing CHITA, a novel optimization-based pruning framework that improves sparsity-accuracy tradeoffs, achieving a 63% relative accuracy improvement over state-of-the-art methods for MLPNet with only 2% of weights retained.
The sheer size of modern neural networks makes model serving a serious computational challenge. A popular class of compression techniques overcomes this challenge by pruning or sparsifying the weights of pretrained networks. While useful, these techniques often face serious tradeoffs between computational requirements and compression quality. In this work, we propose a novel optimization-based pruning framework that considers the combined effect of pruning (and updating) multiple weights subject to a sparsity constraint. Our approach, CHITA, extends the classical Optimal Brain Surgeon framework and results in significant improvements in speed, memory, and performance over existing optimization-based approaches for network pruning. CHITA's main workhorse performs combinatorial optimization updates on a memory-friendly representation of local quadratic approximation(s) of the loss function. On a standard benchmark of pretrained models and datasets, CHITA leads to significantly better sparsity-accuracy tradeoffs than competing methods. For example, for MLPNet with only 2% of the weights retained, our approach improves the accuracy by 63% relative to the state of the art. Furthermore, when used in conjunction with fine-tuning SGD steps, our method achieves significant accuracy gains over the state-of-the-art approaches.