Sparse Model Soups: A Recipe for Improved Pruning via Model Averaging

TL;DR

Sparse Model Soups (SMS) tackles the tension between sparsity and model averaging by ensuring identical sparse connectivity across averaged models. By pruning a pretrained model once and retraining multiple copies with varied hyperparameters, SMS creates a set of models that lie in a common loss basin and can be merged into a single sparse model without sacrificing sparsity. After each prune-retrain phase, SMS starts the next phase from the merged soup, preserving the sparsity pattern and enabling parallel retraining, which substantially improves generalization and OOD robustness over standard IMP and naive averaging. The approach extends to pruning-during-training methods and yields practical gains across image classification, segmentation, and machine translation benchmarks, offering a scalable, parallelizable path to highly sparse, high-performing models suitable for deployment. Overall, SMS demonstrates that sparsity-preserving model averaging can enhance performance while maintaining the computational benefits of sparse networks, with notable improvements in robustness and fairness in pruning regimes.

Abstract

Neural networks can be significantly compressed by pruning, yielding sparse models with reduced storage and computational demands while preserving predictive performance. Model soups (Wortsman et al., 2022) enhance generalization and out-of-distribution (OOD) performance by averaging the parameters of multiple models into a single one, without increasing inference time. However, achieving both sparsity and parameter averaging is challenging as averaging arbitrary sparse models reduces the overall sparsity due to differing sparse connectivities. This work addresses these challenges by demonstrating that exploring a single retraining phase of Iterative Magnitude Pruning (IMP) with varied hyperparameter configurations such as batch ordering or weight decay yields models suitable for averaging, sharing identical sparse connectivity by design. Averaging these models significantly enhances generalization and OOD performance over their individual counterparts. Building on this, we introduce Sparse Model Soups (SMS), a novel method for merging sparse models by initiating each prune-retrain cycle with the averaged model from the previous phase. SMS preserves sparsity, exploits sparse network benefits, is modular and fully parallelizable, and substantially improves IMP's performance. We further demonstrate that SMS can be adapted to enhance state-of-the-art pruning-during-training approaches.

Figures (10)

• Figure 1: Creating the average (middle) of two networks with different sparsity patterns (left, right) may lower overall sparsity, changing pruned weights (dashed) to non-zero (solid), with reactivated weights highlighted in orange.
• Figure 2: Sparse Model Soups
• Figure 3: Accuracy of average of two models vs. the maximal individual accuracy. All models are pruned to 70% sparsity (One Shot) and retrained, varying the indicated hyperparameters.
• Figure 4: WideResNet-20 on CIFAR-100: (a) Accuracy difference between the soup ($m=5$) and best averaging candidate after One Shot pruning and retraining for varying sparsity levels. (b) Accuracy difference between the soup ($m=3$) and $\text{IMP}_{3\! \times}$ retrained three times as long as indicated on the x-axis, using One Shot pruning to 90%, 95% and 98% sparsity. Results are averaged over multiple random seeds with min-max bands indicated.
• Figure 5: WideResNet-20 on CIFAR-100: Accuracy of average of two models vs. the maximal individual accuracy. All models are pruned to 90% sparsity (One Shot) and retrained, varying the indicated hyperparameters.