Paper: Sparse Model Soups: A Recipe for Improved Pruning via Model Averaging (ICLR-2024)

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Evaluation):
It would be valuable to conduct empirical analyses to investigate why performance degradation occurs in regions of extreme sparsity.

Author response:
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.

Reviewer comment (Experiments):
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].

Author response:
Thanks for the quick response and clarification!

Our experiments align with our intuition that the case of low pruning ratio (as in e.g. what you call early stages) is not nearly as detrimental to the overall performance of SMS as the high or extreme pruning ratios. In fact, SMS deals perfectly with low sparsities, the issue of batch randomness causing divergence (what we call 'instability to randomness') is more of a problem in the high sparsity regime. Please see for example Figure 4a: in the low/medium sparsity regime (70\%), SMS yields an immediate over the individual models and with increasing sparsity, this benefits starts to decrease. Further, see Figure 4b, where the green line corresponds to a sparsity of 90\% (which is high, but not extreme in that setting): Even with as little as retraining for a single epoch, the soup is better than performing $IMP_{m\times}$. We refer also to the ablation study in Appendix C, showing that these observations similarly hold for various retraining schedules and lengths.

To sum up: Instability to batch randomness is more of an issue for high sparsities. That is also intuitive and it is not that clear how to empirically analyze why SMS performs well in that scenario. Our argument in the paper is more of a qualitative than quantitative nature: Frankle et al. have shown that there is a point at which sufficient pretraining has been done such that the model can be split into two individually trained branches, which end up in the same loss basin and are thus averageable. That means, starting from a randomly initialized model is not enough to obtain stability, but after some amount of pretraining, the network is stable. In our case, the network has been fully pretrained to convergence. Suppose we would not prune (i.e. prune to a sparsity of of 0\%), then with the results of Frankle et al. in mind, it is clear that the retrained models are averageable (assuming a reasonable learning rate schedule). With increasing sparsity, the model gets more and more damaged and at some point the individual branches start to diverge. However, in the low sparsity regime, the pruned model is still very stable and batch randomness is not able to let the retrained models diverge.

Does that answer your question, or is further clarification needed? Please let us know and we are happy to clarify further. Thanks!

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

Author response:
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?

Reviewer comment (Reproducibility):
I suggest that the authors consider adding an Ethics Statement and a Reproducibility Statement immediately following the main paper.

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

Reviewer comment (Theory):
While the proposed SMS framework incorporates the model soup algorithm in the context of neural network pruning, it does not provide specific insights into the unique factors that are especially pertinent to sparse network training.

Author response:
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.

Reviewer comment (Experiments):
Alongside model soups, another prominent weight-averaging strategy is Stochastic Weight Averaging (SWA). The fact that SMS requires m training runs at each cycle makes SWA, which performs weight averaging within a single SGD trajectory, somewhat appealing. Are there any baseline results using SWA instead of model soups?

Author response:
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.

Reviewer comment (Novelty):
I am not sure if this paper contributes new ideas or analysis. The proposal is to replace the training portion of prune-retrain with model soups. 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.

Author response:
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.

Reviewer comment (Reproducibility):
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.

Author response:
$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:

| Method | $m=3$ | $m=5$ | $m=10$                 |
|---|---|---|------------------------|
| $IMP_{m\times}$ | 74.34\% ($\pm$ 0.09\%) | 74.56\% ($\pm$ 0.24\%) | 74.50\% ($\pm$ 0.09\%) |
| IMP with $m$ phases | 73.69\% ($\pm$ 0.10\%) | 74.08\% ($\pm$ 0.04\%) | 74.70\% ($\pm$ 0.02\%) |

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.

Reviewer comment (Novelty):
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? I felt that there were no new challenges here.

Author response:
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.

Reviewer comment (Novelty):
I have significant novelty concerns with the draft. 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).

Author response:
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?

Reviewer comment (Novelty):
The authors have failed to detail how their method contrasts with existing sparse model soup papers in their related work section. I feel it is just an incremental work over the existing literature. The benefits of averaging the sparse masks is already known.

Author response:
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.

Reviewer comment (Experiments):
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).

Author response:
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.

Reviewer comment (Experiments):
I feel auxiliary benefits of model soups like OOD robustness, fairness, etc are good directions to explore.

Author response:
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?