Swap: Sparse Entropic Wasserstein Regression for Robust Network Pruning

Lei You*, Hei Victor Cheng*

*Corresponding author for this work

Research output: Contribution to conferencePaperResearch

61 Downloads (Pure)


This study tackles the issue of neural network pruning that inaccurate gradients exist when computing the empirical Fisher Information Matrix (FIM). We introduce SWAP, an Entropic Wasserstein regression (EWR) network pruning formulation, capitalizing on the geometric attributes of the optimal transport (OT) problem. The “swap” of a commonly used standard linear regression (LR) with the EWR in optimization is analytically showcased to excel in noise mitigation by adopting neighborhood interpolation across data points, yet incurs marginal extra computational cost. The unique strength of SWAP is its intrinsic ability to strike a balance between noise reduction and covariance information preservation. Extensive experiments performed on various networks 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.
Original languageEnglish
Publication date2024
Number of pages23
Publication statusPublished - 2024
EventThe Twelfth International Conference on Learning Representations - Vienna, Austria
Duration: 7 May 202411 May 2024
Conference number: 12


ConferenceThe Twelfth International Conference on Learning Representations


Dive into the research topics of 'Swap: Sparse Entropic Wasserstein Regression for Robust Network Pruning'. Together they form a unique fingerprint.

Cite this