Sparse Model Soups: A Recipe for Improved Pruning via Model Averaging
Sparse Model Soups (SMS) improves pruning by averaging models, enhancing generalization and OOD performance while maintaining sparsity.
Key Findings
Methodology
Sparse Model Soups (SMS) extends Iterative Magnitude Pruning (IMP) by averaging models after each prune-retrain cycle, ensuring shared sparse connectivity and leveraging model averaging benefits.
Key Results
- On CIFAR-100, SMS achieved 73.05% accuracy at 98% sparsity with WideResNet-20, over 1% higher than single models.
- On ImageNet, SMS reached 74.96% accuracy at 90% sparsity with ResNet-50, outperforming IMP and IMPm×.
- On OOD dataset ImageNet-C, SMS improved robustness by 2.5%, demonstrating adaptability to distribution shifts.
Significance
SMS addresses the challenge of sparsity loss during model averaging, significantly enhancing generalization and OOD robustness. Its modular and parallelizable design makes it practical for deploying sparse models in real-world applications.
Technical Contribution
SMS introduces a novel pruning-averaging framework, combining shared sparse connectivity with iterative averaging. It significantly improves IMP's performance and adapts to state-of-the-art pruning methods.
Novelty
SMS is the first to integrate model averaging with sparsity, solving the sparsity degradation issue and achieving performance gains through iterative averaging.
Limitations
- Performance gains diminish at extreme sparsity levels due to instability caused by randomness.
- SMS requires additional computational resources to train multiple models, though it is parallelizable.
- Performance improvements may be limited in certain extreme pruning scenarios.
Future Work
Future research could explore SMS's stability under extreme sparsity, optimize averaging strategies, and extend its applicability to tasks like reinforcement learning and large language models.
AI Executive Summary
Sparse neural networks reduce storage and computation costs but face challenges in maintaining sparsity during model averaging. Sparse Model Soups (SMS) addresses this by starting each prune-retrain cycle with an averaged model, preserving sparse connectivity while enhancing generalization and OOD performance.
Built on Iterative Magnitude Pruning (IMP), SMS generates multiple models with shared sparse connectivity, retrains them with varied hyperparameters, and averages their parameters into a single model. Experiments on CIFAR-100 and ImageNet show that SMS outperforms traditional IMP and extended IMPm×, with significant improvements in accuracy and robustness on OOD datasets.
While SMS's performance declines at extreme sparsity levels, its modular and parallelizable design offers a promising direction for deploying efficient sparse models. Future work could focus on improving stability under high sparsity and extending SMS to other tasks and architectures.
Deep Analysis
Background
Overparameterized neural networks demand significant storage and computation. Pruning removes redundant parameters to create sparse models, reducing these costs. Model averaging (e.g., Model Soups) improves generalization but disrupts sparsity in sparse models.
Core Problem
Averaging sparse models often reduces overall sparsity, increasing computational costs and potentially degrading performance. Balancing sparsity and averaging benefits is a key challenge.
Innovation
SMS starts each prune-retrain cycle with an averaged model, ensuring consistent sparse connectivity and leveraging averaging benefits. This novel integration of sparsity and model averaging addresses sparsity degradation.
Methodology
- �� Use IMP to prune a pretrained model, creating sparse models
- �� Generate multiple models with shared sparse connectivity
- �� Retrain models with varied hyperparameters (e.g., batch order, weight decay)
- �� Average retrained models into a single model
- �� Begin the next prune-retrain cycle with the averaged model
Experiments
Experiments on CIFAR-100, ImageNet, and OOD datasets evaluate SMS's performance at various sparsity levels. Baselines include IMP, IMPm×, and IMP-RePrune. Metrics include accuracy, OOD robustness, and fairness.
Results
SMS achieved 73.05% accuracy at 98% sparsity on CIFAR-100 and 74.96% at 90% sparsity on ImageNet, outperforming IMP and IMPm×.
Applications
SMS enables efficient sparse model deployment, suitable for edge devices and resource-constrained scenarios. Its modular design integrates easily into existing pruning frameworks.
Limitations & Outlook
SMS's performance gains are limited at extreme sparsity levels and require additional computational resources. Stability under high sparsity remains a challenge.
Plain Language Accessible to non-experts
Imagine you're cooking in a kitchen (training models) and want to reduce food waste (pruning). SMS acts like a smart recipe book, learning from past meals (model averaging) to ensure every dish uses minimal ingredients (sparsity) while tasting great.
ELI14 Explained like you're 14
Think of SMS like a gaming teammate! After every boss fight, it reviews your moves, upgrades your gear (model averaging), and helps you play smarter with fewer resources (sparsity). Cool, right?
Glossary
Iterative Magnitude Pruning (IMP)
A pruning method that removes low-magnitude weights iteratively to create sparse models.
SMS builds on IMP to ensure shared sparse connectivity.
Model Soups
A method that averages parameters of multiple models into a single one.
SMS integrates Model Soups to enhance sparse model performance.
Out-of-Distribution (OOD)
Data that differs from the training distribution.
SMS improves robustness on OOD datasets like ImageNet-C.
Sparsity
The proportion of non-zero parameters in a model.
SMS maintains sparsity to reduce computational costs.
Weight Decay
A regularization technique that penalizes large weights to prevent overfitting.
SMS uses varied weight decay during retraining.
Open Questions Unanswered questions from this research
- 1 SMS's stability under extreme sparsity remains unexplored.
- 2 Optimizing model averaging strategies for diverse tasks is an open challenge.
- 3 SMS's applicability to large-scale language models is yet to be validated.
Applications
Immediate Applications
Edge Device Deployment
SMS generates efficient sparse models, reducing computation and storage needs for edge devices.
OOD Robustness
SMS's strong OOD performance makes it suitable for safety-critical applications like autonomous driving.
Long-term Vision
Universal Sparse Framework
SMS could evolve into a universal framework for various tasks and architectures, advancing sparse model research.
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.