Sparse Weight Averaging with Multiple Particles for Iterative Magnitude Pruning

TL;DR

This work tackles pruning-induced performance loss in extreme sparsity by extending Iterative Magnitude Pruning (IMP) with Sparse Weight Averaging using Multiple Particles (SWAMP). SWAMP trains $N$ particles from the same matching ticket using different SGD noise and aggregates them via stochastic weight averaging to produce a single sparse mask, achieving accuracy near that of an ensemble at the same inference cost. The authors show that SWAMP preserves linear connectivity between successive IMP solutions and tends to locate flatter minima, leading to better generalization than standard IMP across image and language tasks. They also discuss practical strategies to reduce training cost through parallelization and selective multi-particle usage, and demonstrate extensions to dynamic pruning methods, highlighting SWAMP’s broad applicability and potential impact on efficient sparse modeling.

Abstract

Given the ever-increasing size of modern neural networks, the significance of sparse architectures has surged due to their accelerated inference speeds and minimal memory demands. When it comes to global pruning techniques, Iterative Magnitude Pruning (IMP) still stands as a state-of-the-art algorithm despite its simple nature, particularly in extremely sparse regimes. In light of the recent finding that the two successive matching IMP solutions are linearly connected without a loss barrier, we propose Sparse Weight Averaging with Multiple Particles (SWAMP), a straightforward modification of IMP that achieves performance comparable to an ensemble of two IMP solutions. For every iteration, we concurrently train multiple sparse models, referred to as particles, using different batch orders yet the same matching ticket, and then weight average such models to produce a single mask. We demonstrate that our method consistently outperforms existing baselines across different sparsities through extensive experiments on various data and neural network structures.

Figures (10)

• Figure 1: Classification error as a function of the sparsity (left) and the relative number of parameters (right). Our proposed SWAMP achieves remarkable performance comparable to an ensemble of IMP solutions, where IMP-$n$ indicates the ensemble of $n$ IMP solutions, while maintaining the same inference cost as a single model. Notably, SWAMP demonstrates matching performance even at extremely sparse levels, unlike IMP. The results are presented for WRN-28-2 on CIFAR-10, and we refer readers to \ref{['app:sec:additional_experiments:sparsity']} for the same plot for CIFAR-100, as well as VGG architectures.
• Figure 2: Visualization of loss surfaces as a function of network weights in a two-dimensional subspace, spanned by three particles (marked as white circles). The averaged weight $\boldsymbol{w}_c$ (marked by a yellow star) is observed not to be positioned in the flat region of the surface during the earlier stages of IMP (left; Sparsity 20%). However, as sparsity increases, the weight averaging technique effectively captures the flat region of the surface. The results are presented for WRN-28-2 on the test split of CIFAR-10, and we refer the reader to \ref{['app:sec:additional_experiments:loss_surface']} for the same plot for CIFAR-100.
• Figure 3: Bar plots depicting the accuracy of individual particles involved in the averaging process of the SWAMP algorithm. While the averaged weight (denoted as WA) may not outperform individual particles (denoted as P1-P4) in the early stages of IMP (left; Sparsity 20%), it achieves high performance at higher sparsity levels. The results are presented for WRN-28-2 on the test split of CIFAR-10. We refer readers to \ref{['app:sec:additional_experiments:particle']} for the same plot for CIFAR-100.
• Figure 4: Linear connectivity between sparse solutions with different sparsity levels gathered from the end of IMP cycles. The results are presented for WRN-28-2 (left) and VGG-13 (right) on the test split of CIFAR-10, and we refer the reader to \ref{['app:sec:additional_experiments:connectivity']} for the same plot for CIFAR-100.
• Figure 5: Results for RoBERT on MRPC and RTE. Reported values are averaged over three random seeds.