Compression-Aware training of neural networks using Frank-Wolfe

Abstract

Many existing Neural Network (NN) pruning approaches rely on either retraining or inducing a strong bias in order to converge to a sparse solution throughout training. A third paradigm, "compression-Aware"training, aims to obtain state-of-The-Art dense models that are robust to a wide range of compression ratios using a single dense training run while also avoiding retraining. We propose a framework centered around a versatile family of norm constraints and the Stochastic Frank-Wolfe (SFW) algorithm that encourage convergence to well-performing solutions while inducing robustness towards filter pruning and low-rank matrix decomposition. Our method outperforms existing compression-Aware approaches and, in the case of low-rank matrix decomposition, it also requires significantly less computational resources than approaches based on nuclear-norm regularization. Our findings indicate that dynamically adjusting the learning rate of SFW, as suggested by Pokutta et al. [57], is crucial for convergence and robustness of SFW-Trained models and we establish a theoretical foundation for that practice.