SWAP: Sparse Entropic Wasserstein Regression for Robust Network Pruning

TL;DR

SWAP introduces a robust network pruning method using Entropic Wasserstein regression that effectively mitigates noise in gradients, leading to improved accuracy especially in large or noisy scenarios.

Contribution

The paper proposes SWAP, a novel pruning technique based on Entropic Wasserstein regression, which enhances noise robustness and preserves covariance information during network pruning.

Findings

SWAP achieves comparable performance to state-of-the-art methods.

SWAP outperforms existing methods with large network sizes or high sparsity.

SWAP improves MobileNetV1 accuracy by 6% with significantly fewer parameters.

Abstract

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.