Compression-aware Training of Neural Networks using Frank-Wolfe

TL;DR

This paper introduces a compression-aware training framework using the Stochastic Frank-Wolfe algorithm and norm constraints, producing dense neural networks that are robust to various compression techniques without retraining.

Contribution

It presents a versatile framework combining norm constraints and SFW for robust neural network training, outperforming existing methods and reducing computational costs for low-rank approximations.

Findings

Outperforms existing compression-aware methods

Requires less computation for low-rank matrix decomposition

Dynamic learning rate adjustment is crucial for convergence

Abstract

Many existing Neural Network 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 convolutional filter pruning and low-rank matrix decomposition. Our method is able to outperform 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. (2020), is crucial for convergence and robustness of SFW-trained models and we establish a theoretical foundation for that practice.