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

We thank you for your positive feedback and constructive remarks.

Weaknesses

Although they can parallelize the training process, performing $m \times k$ training epochs still imposes a substantial computational burden. And if the training cost becomes small, the overall performance gain significantly drops for the extreme sparsity cases.

We acknowledge that all model-averaging methods that rely on training multiple models in parallel (as e.g. the original model soups) induce some computational overhead in order to increase the final performance. Approaches that average models along the training trajectory (such as e.g. SWA) do not have that drawback, but fall behind in performance, at least in our particular setting (please see the response to Reviewer fESa for a full account on methods like SWA).

However, please note that this is only true for extreme levels of sparsity, that would require iterative pruning in any case to reach the full performance of classical IMP. In fact, as visible in Figure 4b, for reasonable levels of sparsity SMS is able to outperform its prolonged IMP counterpart after a minimum number of epochs and with significant improvements. In Appendix C, we further study that SMS is able to improve among its individual averaging candidates independent of the retraining length and sparsity - there just exists a certain tradeoff at which retraining is too short such that averaging $m$ models gives better results than retraining $m$ times as long.

It would be valuable to conduct empirical analyses to investigate why performance degradation occurs in regions of extreme sparsity.

Thank you for that remark - we have addressed this point from a qualitative perspective, but we agree that further empirical substantiation improves our work. In the paper, we argue that extremely high sparsity levels (the problem occurs at 99% pruned in One Shot and above) lead to a model that is not stable to randomness anymore, i.e., two retrained models do not lie in the same basin and thus cannot be averaged. To substantiate this claim, we have run further experiments and tracked several metrics to see how the candidates for averaging behave. One important empirical observation lies in tracking the $L_2$-norm distance between the candidates: we observe that for sparsities in the range 90%-98%, the mean and maximal $L_2$-distance between the 5 candidate models as in the experiment for Figure 4a are relatively stable among sparsities. Increasing the sparsity to 99% and 99.5% however leads to a much increased distance between the retrained model. At this sparsity, the models are driven further apart, supporting our hypothesis of instability to randomness - they do not converge to the same basin.

The upcoming revision will include an ablation study including these experiments in a more detailed form. We will include them in the manuscript as soon as possible and will let you know.

Similarly, it would be beneficial to empirically analyze why SMS performs well in situations of early sparsity, where batch randomness can lead to divergence between averaged weights [3].

Could you elaborate what you mean by 'early sparsity'? We are not entirely sure if we understand you correctly here.

It would be advantageous to include an ablation study exploring different combinations of averaging coefficients where the $\lambda_i$ values differ from each other.

We have experimented a lot with what Wortsman et al. call LearnedSoup, where the $\lambda_i$ coefficients are learned through optimization on a held-out validation set. In our experiments, LearnedSoup was never able to reach the performance of UniformSoup or GreedySoup. We are happy to include these results in the appendix in the next revision - do you have any particular suggestion on how to perform an ablation study on the effect of differing $\lambda_i$? How should the values be chosen in such a study, if not automatically as in LearnedSoup?

Recommendations

I suggest that the authors consider adding an Ethics Statement and a Reproducibility Statement immediately following the main paper.

Thanks for that remark, we will add both to the revision.

We thank you again for your insightful review and hope to have addressed your concerns. Please let us know in case there are further doubts or questions.

Thank you for your review! We are happy to address the individual points raised, as they help us to improve our work.

Weaknesses

What attributes of the model soup algorithm contribute to its effectiveness in the neural network pruning regime? Is it the same reason why weight-averaging methods have succeeded in conventional dense network training?

To clarify, let us re-emphasize the key insight of our work. First of all, addressing your second question, weight-averaging methods that "have succeeded in conventional dense network training" (such as e.g. SWA) are of a different nature: they average models along the optimization path to explore different optima but always train a single model with fixed hyperparameters.

In contrast, the less conventional model soup approach is specific to transfer learning and involves training and averaging multiple models with varied hyperparameters. Its effectiveness in transfer learning has been thoroughly validated by Wortsman et al., with their empirical evidence (Figure 2 in Wortsman et al., 2022a) showing the best generalizing solution often resides 'between' two finetuned networks.

Returning to our work, we maintain the same downstream task without domain shifts. However, we think and have argued that pruning and retraining a model strongly resembles the finetuning paradigm in transfer learning: we interpret pruning as feature perturbation and retraining as adapting these features to the original task, as if the original pretraining would have been done on a different task (i.e. the one leading to the perturbed features). While we think that this provides a good intuition, it is far from obvious that the model soup mechanism also holds when a large portion of weights are pruned, reducing the model's performance to that of a random classifier. The model undergoes significant perturbation/damage, raising questions about its 'stability to randomness' (as per Frankle et al. 2020). It's uncertain if two versions of such a heavily pruned model would be averageable, particularly as aggressive learning rate schedules during retraining could drive the models apart.

Our work is the first to empirically demonstrate that this is indeed possible and significant performance gains, both in-distribution and out-of-distribution, are achievable across diverse architectures and datasets. Returning to your first question, we think that the insights from transfer learning can be transferred to the network pruning regime, motivated by the resemblance between the two paradigms as outlined above, and substantiated by extensive results.

Questions

Are there any baseline results using SWA instead of model soups?

Thank you for suggesting the comparison between SWA and SMS, it's an interesting improvement to our work. We conducted experiments, but before reporting and discussing the results, let us make some general remarks.

1. SWA is only beneficial if models of the same sparsity level and pattern are averaged, as differing sparsities will densify the model (see Figure 1). We apply SWA separately in each phase, starting each phase with the averaged model from the previous phase and reinitializing SWA accordingly.
2. SWA and SMS are not excluding each other, they can be used in conjunction, potentially further improving the effect of SMS.
3. SWA requires either a cyclic or high constant learning rate to explore multiple optima for beneficial averaging. However, retraining after pruning uses specific translated learning rate schedules (such as FT, LRW, CLR, LLR or ALLR) to maximize performance.

In our experiments using ResNet-50 on ImageNet, following the same setup as in Table 1 for a sparsity of 90% in three cycles,we observed slightly inferior results with SWA:

• IMP: Test Accuracy 73.72% ($\pm$ 0.00%)
• IMP+SWA: Test Accuracy 73.05% ($\pm$ 0.02%)

SWA is not able to improve the results of classical IMP (and hence also falls behind SMS by a large margin, cf. Table 1). We think that this is mostly due to the specific retraining schedules used for IMP, which stand in conflict with the requirements for SWA. We will add an ablation study dissecting these factors as soon as possible.

[...] despite the potential fair comparison issue, presenting additional GreedySoup results might offer valuable insights into the benefits of selectively using soup ingredients at high sparsity levels.

We are happy to create a separate validation split for these two specific tasks and will add the experiments as soon as possible, however we do not think that GreedySoup will be able to offer much of an improvement - apart from extreme levels of sparsity, the uniform approach lead to the most consistent improvements.

Thanks for your review, we hope to have addressed your concerns and answered your questions. Please let us know if further clarification is needed.

We thank you for reviewing our work and acknowledging its strengths! Let us respond to your concerns and questions in detail.

Weaknesses

I do not have background on transfer learning, but as a general machine learning person, it is not surprising that it improves the accuracy of the model given the backdrop of model soups paper . Since in both the cases it holds that m copies of model start from the same initialization.

We regret to hear your doubts about the novelty of our work and wish to clarify. To begin with, the averageability of models trained from identical initializations is not obvious. Neyshabur et al. (2020) observed that models with the same random initialization but varying batch selections diverge and are not averageable, whereas Frankle et al. (2020) established averageability after starting from a sufficiently pretrained model. Wortsman et al. (2022a) found that finetuning models from a single pretrained model are averageable, despite task changes (i.e. transfer learning) and the pretrained model potentially being bad at the new task.

In our research, we avoid domain shifts, focusing on the same downstream task. In non-pruning scenarios, it's somewhat clear that extended training with a reasonable learning rate might yield models whose average is not worse than individual models. However, this is far from obvious in heavily pruned models, where performance drops significantly. The substantial damage these models undergo raises questions about their 'stability to randomness' (as per Frankle et al. 2020). It's not evident that two heavily pruned models are averageable, especially under aggressive retraining schedules.

Our work is the first to empirically show that averaging such models leads to notable performance improvements in- and out-of-distribution across various architectures and datasets. We link pruning and retraining to finetuning in transfer learning, treating pruning as feature perturbation and retraining as features adaption to the original task (as if pretraining would have lead to the perturbed features). A key challenge was overcoming the decreased sparsity from naive weight averaging due to different sparsity patterns (cf. Figure 1), requiring retraining at the expense of performance (cf. IMP-RePrune, Table 1). Our solution, SMS, effectively addresses this issue. Therefore, we respectfully disagree that our papers does not contribute new ideas, as outlined.

Questions

How is $IMP_{m\times}$ implemented? Is the pruning rate for each prune step reduced? or is the training portions increased m times. The latter, I suspect, will not be very useful.

$IMP_{m\times}$ increases training epochs per prune-retrain-cycle by a factor of $m$. We believe this, arguably most natural baseline is useful; could you elaborate on why you think it isn't?. Nevertheless, we agree that reducing the pruning rate accordingly (i.e. for $m=3$, triple the amount of prune-retrain cycles to reach a final sparsity) is an interesting baseline, thank you for that remark! In our experience, lower pruning rates do not necessarily enhance IMP's final product; in fact, empirical evidence from Bartoldson et al. indicates that the instability from significant pruning can actually aid generalization.

In response to your question, we conducted experiments with ResNet-50 on ImageNet, aiming for a 90% sparsity goal, replicating the setting in Table 1. We compared the final test accuracies of $IMP_{m\times}$, which extends retraining length per cycle, with a variant increasing the number of cycles by $m$ but maintaining a 10-epoch retraining length per cycle. The results are as follows:

This approach mostly falls behind $IMP_{m\times}$ and in consequence also behind SMS. We however thank you for suggesting this additional baseline, we will include it in the manuscript.

Are there any challenges that are specific to using model soups for training portion of prune-retrain algorithm which differentiate it from applying model soups to finetuning of pretrained models?

We hope that this concern is sufficiently addressed by our explanations regarding the novelty and contribution of our work - the setting is different and it is a priori not clear whether aggressive pruning yields models that are averageable, especially given the existing (aggressive) learning rate schemes for retraining. Further, it is a priori not clear how to mitigate the problem of reduced sparsity when averaging arbitrary sparse models; a problem that does not exist in the setting of transfer learning.

Thanks again for your review - please let us know if further clarification is needed!

Thank you for reviewing our manuscript. In the following, let us address your concerns in detail.

Weaknesses

1. The idea of sparse model averaging has been widely explored including the model soups (eg. https://arxiv. org/abs/2205.15322 https://arxiv.org/abs/2208.10842 https://arxiv.org/abs/2306.10460 etc).

We are surprised to see the novelty of our work being questioned, especially since we cite all the papers you mention. We regret if this misunderstanding stems from poor communication from our side, see also the response to your second concern below. Let us clarify how our work differs from the cited papers:

2205.15322: Yin et al. focus on Dynamic Sparse Training (DST), which is different from our domain (prune-after- and -during-training). They average models within a single training run using fixed hyperparameters, not across multiple runs. Their averaging destroys the sparsity (cf. Figure 1), requiring re-pruning: this is inherent to their prune-and-grow DST approach, requiring to explore different sparsity patterns. To mitigate the negative impacts of re-pruning, they propose averaging methods like CIA and CAA. In contrast, our work deliberately avoids re-pruning by keeping consistent sparsity patterns. Note that we explicitly compare to the re-pruning approach (cf. Table 1, IMP-RePrune), using both their methods CIA and CAA, outperforming their re-pruning schemes significantly.

2208.10842: Yin et al. use IMP with weight rewinding (IMP-WR) and concentrate on generating lottery tickets, a distinct goal from ours of creating high-performing sparse models. Weight rewinding rewinds the networks to the point of stability (as per Frankle et al.), the issue whether adjacent models are averageable is thus entirely different to whether they are averageable without rewinding. While the authors average models with varying sparsity (those adjacent in the lottery ticket generation process) necessitating re-pruning, our method averages parallely trained models, avoiding this detrimental step in our setting.

2306.10460: Jaiswal et al. do not average model parameters at all, but rather average early pruning masks to quickly generate subsequent masks. This approach doesn't overlap with our contributions, focusing on fast mask generation rather than averaging sparse models.

If there are other publications you believe align closely with our work, especially those implied by "etc." in your list, could you please specify them?

2. I feel it is just an incremental work over the existing literature. The benefits of averaging the sparse masks is already known.

We were unaware that our manuscript could be misunderstood in that regard, but we take your concern seriously and will clarify the differences with the mentioned papers in our revised manuscript, which will be available soon. We hope that this resolves the issue.

We respectfully disagree that our work is incremental. To our knowledge, no other publication has explored sparse model averaging in the context of prune-after-training or during-training; in fact, "averaging the sparse masks" is entirely different from what we do (cf. our explanations regarding Jaiswal et al.). Secondly, previous works on averaging sparse models involve a single training run and fixed hyperparameters, unlike our model soup approach that averages models from multiple runs with varying hyperparameters. It has been far from obvious that this was possible given that pruning drastically alters the model and reduces stability to randomness. Finally, our approach uniquely preserves the sparsity pattern, avoiding the need for repruning, a step we demonstrate as harmful and which we significantly improve upon.

3. Although I appreciate the extensive experiments by authors, I still feel the experiments are limited to small-scale datasets and models (maybe ViT scale or OPT models-based experiments will add value).

We respectfully note that our experiments extend beyond "small-scale datasets and models." In fact, they encompass models of "ViT scale", including extensive experimentation with the vision-transformer MaxViT on ImageNet and the T5-transformer on WMT16, detailed in Appendix B.

4. I feel auxiliary benefits of model soups like OOD robustness, fairness, etc are good directions to explore.

We have conducted thorough experiments regarding the benefits of SMS in the OOD (Appendix B.2.3) and fairness (Appendix B.2.4) setting, demonstrating a consistent and significant improvement over the baseline, particularly in the OOD-case. Could you elaborate what additional directions you think are worth exploring?

We hope our responses have clarified the contributions and novelty of our work. If you have any more questions or concerns, or suggestions for improvement, please let us know. Otherwise, we kindly ask you to reconsider your initial evaluation.