End-to-End Neural Network Compression via L1/l2 Regularized Latency Surrogates

Abstract

Neural network (NN) compression via techniques such as pruning, quantization requires setting compression hyperparameters (e.g., number of channels to be pruned, bitwidths for quantization) for each layer either manually or via neural architecture search (NAS) which can be computationally expensive. We address this problem by providing an end-to-end technique that optimizes for model’s Floating Point Operations (FLOPs) via a novel $\frac{{{\ell _1}}}{{{\ell _2}}}$ latency surrogate. Our algorithm is versatile and can be used with many popular compression methods including pruning, low-rank factorization, and quantization, and can optimize for on-device latency. Crucially, it is fast and runs in almost the same amount of time as a single model training run; which is a significant training speed-up over standard NAS methods. For BERT compression on GLUE fine-tuning tasks, we achieve 50% reduction in FLOPs with only 1% drop in performance. For compressing MobileNetV3 on ImageNet-1K, we achieve 15% reduction in FLOPs without drop in accuracy, while still requiring 3× less training compute than SOTA NAS techniques. Finally, for transfer learning on smaller datasets, our technique identifies 1.2× - 1.4× cheaper architectures than standard MobileNetV3, EfficientNet suite of architectures at almost the same training cost and accuracy.