SWAP: Sparse Entropic Wasserstein Regression for Robust Network Pruning
This work addresses robust pruning for deep learning models, particularly under noisy conditions, but is incremental as it builds on existing pruning methods with a novel optimization formulation.
The study tackled inaccurate gradients in neural network pruning by introducing SWAP, an Entropic Wasserstein regression method that mitigates noise while preserving covariance information, achieving up to 6% accuracy improvement and 8% loss reduction on MobileNetV1 with high sparsity.
This study addresses the challenge of inaccurate gradients in computing the empirical Fisher Information Matrix during neural network pruning. We introduce SWAP, a formulation of Entropic Wasserstein regression (EWR) for pruning, capitalizing on the geometric properties of the optimal transport problem. The ``swap'' of the commonly used linear regression with the EWR in optimization is analytically demonstrated to offer noise mitigation effects by incorporating neighborhood interpolation across data points with only marginal additional computational cost. The unique strength of SWAP is its intrinsic ability to balance noise reduction and covariance information preservation effectively. Extensive experiments performed on various networks and datasets show comparable performance of SWAP with state-of-the-art (SoTA) network pruning algorithms. Our proposed method outperforms the SoTA when the network size or the target sparsity is large, the gain is even larger with the existence of noisy gradients, possibly from noisy data, analog memory, or adversarial attacks. Notably, our proposed method achieves a gain of 6% improvement in accuracy and 8% improvement in testing loss for MobileNetV1 with less than one-fourth of the network parameters remaining.