Sparse Weight Averaging with Multiple Particles for Iterative Magnitude Pruning
Authors: Moonseok Choi
Hyungi Lee
Giung Nam
Juho Lee
Venue: ICLR 2024 poster (2024)
Keywords: Unstructured pruning
Sparse neural networks
Lottery ticket hypothesis
Weight averaging
Bayesian neural networks
Primary area: general machine learning (i.e., none of the above)
TL;DR: We propose an IMP-based sparse weight averaging technique that leverages the unique properties of the loss landscape.

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 neural network structures and data.

Published as a conference paper at ICLR 2024
SPARSE WEIGHT AVERAGING WITH MULTIPLE PARTI- CLES FOR ITERATIVE MAGNITUDE PRUNING
Moonseok Choi
Hyungi Lee"!
Giung Nam
Juho Lee 1.2
KAIST AI AITRICS
{ms. choi, 1hk2708, giung, juholee]ekaist.ac.
ABSTRACT
Given the ever-increasing size of modern neural networks, the significance of sparse architectures has surged due to their accelerated inference speeds and min- imal memory demands. When it comes global pruning techniques, Iterative Magnitude Pruning (IMP) still stands as state-of-the-art algorithm despite its simple nature, particularly in extremely sparse regimes. In light of the recent find- ing that the two successive matching IMP solutions are linearly connected with- out a loss barrier, we propose Sparse Weight Averaging with Multiple Particles (SWAMP), straightforward modification of IMP that achieves performance com- parable to an ensemble two IMP solutions. For every iteration, we concurrently train multiple sparse models, referred to particles, using different batch orders yet the same matching ticket, and then weight average such models to produce a single mask. We demonstrate that method consistently outperforms existing baselines across different sparsities through extensive experiments on various data and neural network structures.
INTRODUCTION
Deep neural networks are often highly over- prerpparameterized, and the majority of their parameters can be pruned without sacrificing model performance (LeCun al., 1989). The lottery ticket hypothesis proposed by Frankle & Carbin (2019) suggests that there exists a sparse subnetwork at initialization that can be trained to achieve the same level of performance as the original dense network. Such matching subnetworks can be found via Iterative Magnitude Pruning (IMP) with rewinding (Frankle al., 2020), which involves the following three steps: (i) training the network for certain number of iterations, (ii) pruning weights with the smallest magnitudes, and (iii) rewinding the weights back to an early iteration while fixing the pruned weights to zero. This procedure is repeated for several rounds, and the final rewound subnetwork corresponds to the matching ticket that can achieve the performance of the full network. Despite its simplicity, IMP offers the state-of-the-arr performance as to finding sparse mask, especially for extreme sparsity regimes (Renda et al., 2020).
The success of IMP is indeed counter-intuitive considering its simplicity. In this regard, Frankle et al. (2020) revealed underlying connection between the lottery ticket hypothesis and linear mode connectivity, indicating that the effectiveness of IMP is reliant upon its stability to stochastic optimization; IMP solutions reside in the same basin of attraction in the loss landscape. Delving deeper into the subject matter, Paul et al. (2023) found that linear mode connectivity also exists between successive IMP solutions with different sparsity levels. More precisely, they concluded that IMP fails to find matching subnetwork if the solutions from consecutive rounds are disconnected and further highlighted the significance of both the pruning ratio and the rewinding iteration to retain the connectivity between IMP solutions.
Inspired by the connection between IMP and linear mode connectivity, we expand the understanding to the loss landscape perspective. Analyzing the loss landscape of deep neural networks is an effec- tive tool that is widely employed to study mode connectivity (Draxler et al., 2018; Garipov et al., 2018 8; Fort& & Jastrzebski, 2019; Benton et al., 2021), and it also motivates us to find solutions located
Equal contribution

---

Published as a conference paper at ICLR 2024
at the flat region of the loss landscape to enhance generalization (Chaudhari et al., 2017; Izmailov et al., 2018; Foret et al., 2021). Notably, both fields share a common objective of identifying "good" subspaces characterized by low loss value, and this objective aligns with the ultimate goal of neural network pruning- to identify "matching" sparse subspaces within given dense parameter space.
In this paper, study how IMP can benefit from the multiple models connected in loss surfaces. Our contributions are summarized as follows:
We first empirically demonstrate that multiple models trained with different SGD noise yet from the same matching ticket can be weight-averaged, i.e., there exists no loss barrier within the convex hull of the model weights. We further show that taking an average of the particles leads to flat minima, which exhibit superior generalization performance compared to each individual particle.
Building upon prior observations, we propose novel iterative pruning technique, Sparse Weight Averaging with Multiple Particles (SWAMP), tailored specifically for IMP. We ver- ify that SWAMP preserves the linear connectivity of successive solutions, which is crucial feature contributing to the success of IMP.
Through extensive experiments, we provide empirical evidence that supports the superior- ity of the proposed SWAMP algorithm over other baselines.
BACKGROUNDS
2.1 NEURAL NETWORK PRUNING AS CONSTRAINED OPTIMIZATION
Conventional training neural networks aims find an optimal neural network parameter w e RD that minimizes a given loss function C IR + R. for given training dataset D. Such optimization typically employs the Stochastic Gradient Descent (SGD; Robbins Monro, 1951) methods, which we denote as wy +- SGDo-r (Wo, (WWDo---wo.. D) throughout the paper. Here, wr denotes the solution obtained by performing SGD with randomness of & (e.g., mini-batch ordering) over iterations, starting from the initial weight wp- On the other hand, neural network pruning the process of ob- taining a sparse neural network with a desired sparsity level x E [o, 1) from the original dense neural network. The goal is now to find an optimal sparse solution that the number of non-zero elements in the mask m E [0, 1] satisfies ||mm llo I D(1 D1--)).
w mm'" subject to the constraint
2.2 ITERATIVE MAGNITUDE PRUNING WITH REWINDING
Iterative Magnitude Pruning (IMP: Frankle & Carbin, 2019) is an iterative pruning method that is both straightforward and highly effective. Each cycle of IMP involves the following three steps: (i) Train a network parameter w. at the nd cycle is trained until it reaches convergence. (ii) Prune mask m is created by setting the smallest weights to zero based on a predefined pruning ratio of
a. (iii) Reset- the weights are then reverted back to their initial values before the next cycle begins. This train-prune-resel cycle repeated until the desired level of sparsity is achieved.
However, in practical scenarios, the original version of IMP suffers from rapid performance degra- dation as sparsity increases and fails match the performance of the original dense solution. To address this issue, the concept of rewinding is introduced (Frankle et al., 2020; Renda et al., 2020). Rather than resetting the unpruned weights to their initial values, the weights are rewound to an early training point- the matching ticker. The matching ticket is simply the weights obtained after training for few iterations. Refer to Appendix C.I for more details on IMP algorithms.
2.3 LINEAR CONNECTIVITY OF NEURAL NETWORK WEIGHTTS
Consider a one-dimensional path denoted as P: [0, + RD, connecting two neural network weights (u) and ((1)) in D-dimensional space, where the starting and end points are P(O) = ( and P(1) = (!),, respectively. In a simplified sense, we can say that there is a connectivity between (0) and (!) if the condition supxe[o.1] C(P(A) < max (C(P(O)), C(P(1)) + holds, where is a small margin value. While recent advances in deep learning have revealed the existence of non- linear paths between local minima obtained through stochastic optimization (Draxler et al., 2018;
2

---

Published as a conference paper at ICLR 2024
Matching threshold
IMP-1
-
IMP-2
-
IMP-3
SWAMP (ours)
7.50
7.50
Classitication error (%)
Classification error (%)
6.50
6.50
5.50
5.50
4.50
4.50
20
40
60
80
100
3.0
2.5
2.0
1.5
1.0
0.5
0.0
Sparsity (%)
Relative number of parameters
Figure 1: Classification error function of the sparsity (left) and the relative number of param- eters (right). Our proposed SWAMP achieves remarkable performance comparable to an ensemble of IMP solutions, where IMP-n indicates the ensemble of IMP solutions, while maintaining the same inference cost as a single model. Notably, SWAMP demonstrates matching performance even extremely sparse levels, unlike IMP. The results are presented for WRN-28- on CIFAR-10 10, and we refer readers to Appendix C.6 for the same plot for CIFAR- as well as VGG architectures.
Garipov et al., 2018), it is still not straightforward to establish linear connectivity (i.e., connectivity with linear connector P(A) = (1 A)w( (9) + Aw ()) for modern deep neural networks (Lakshmi- narayanan et al., 2017; Fort & Jastrzebski, 2019; Fort et al., 2020).
3 SPARSE WEIGHT AVERAGING WITH MULTIPLE PARTICLES (SWAMP)
3.1 IMP: A LOSS LANDSCAPE PERSPECTIVE
Frankle al. (2020) demonstrated that the matching ticket has significant impact on the stability of neural networks to SGD noise &. Even when two networks are trained with the same random initialization wo, the different SGD noise c()), g(2) disrupts the linear connectivity between the solutions obtained through SGD, i.e., there is no linear connector between
WT + SGDo-rr(wo ( and wT + SGDo-r(wo. D),
(1)
and thus the optimization is rendered unstable SGD noise. They further empirically confirmed that sparse solutions obtained through IMP are matching if and only if they are stable to SGD noise, and diagnosed this instability A simple treatment to ensure the stability sharing the early phase of the optimization trajectory. In other words, there exists linear connector between
failure of the original IMP algorithm (Frankle & Carbin, 2019).
(1) SGDT,-T g(),, and wi + SGDT. ((2) DD, WT
when SGD runs are started from the same initialization of wTs + SGDu-Ia (Wo, E. D). Further- more, Paul et (2023) demonstrated linear connectivity between two consecutive IMP solutions with different sparsity levels and identified it as a crucial factor for the success of IMP.
Nevertheless, the question of whether a low-loss subspace is formed by the convex hull of three or more solutions remains uncertain, despite the presence of linear connectivity between each pair of solutions. If it becomes feasible construct a low-loss volume subspace using IMP solutions, could potentially yield more effective solution with improved generalization at the midpoint of this subspace (Wortsman al., 2021).
3.2 SWAMP: AN ALGORITIM
Inspired by the stability analysis of the matching ticket presented in 3.1, propose Sparse Weight Averaging with Multiple Particles (SWAMP) as a tailored sparse weight averaging technique for IMP. The detailed algorithm is presented in Algorithm 1.
SWAMP differs from vanilla IMP in two main aspects. Firstly, we create multiple copics of the matching ticket (line 5; Algorithm 1) and train them simultaneously with different random seeds
3

---

Published as a conference paper at ICLR 2024
Algorithm Iterative Magnitude Pruning with SWAMP (ours)
Require: Neural network parameter pruning mask m, training dataset DD the number of cycles for iterative magnitude pruning C. the number of iterations for each eycle T. pruning ratio a, SGD noise Ç. the number of iteration for matching ticket To, and the number of particles N. Sparse solution W.T.
Ensure:
D Starts from random dense weights. 2: Train Wo.7% <- SGDo-Io (Wo,o m. Eo, D). Gets matching ticket from the initialization.
1: Randomly initialize woo0 and set mask m <- 1.
3: for с [l,..., C.....,C.... do
4: for (I....,N) N do
5: Rewind w c,0 + Wo,T. m.
(x6)
> Starts from the matching ticket.
6: Train to (m)
0,7 - S SWAo-r 0.0 :00 DD.
c(1)
b Averages weights over trajectory.
7: end for
8: Averaging W.T 2 w (n)
N
2T/N.
D Averages weights of particles.
9:
Prune m + Prune (T, a).
b Updates the mask based on magnitudes.
10:
end for
Sparsity 20%
Sparsity 75%
Sparsity 90%
Wi
W
Who
0.220
0.252
0.284
0.222
0.237
0.252
0.234
0.248
0.262
Figure 2: Visualization of loss surfaces as a function of network weights a two-dimensional subspace, spanned by three particles (marked as white circles). The averaged weight w, (marked by 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, 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 CIFAR-10, and we refer the reader Appendix C.8 for the same plot for CIFAR-
(line 6; Algorithm 1), whereas IMP employs a single particle. Secondly, we replace SGD training with Stochastic Weight Averaging (SWA; Izmailov et al., 2018), method that constructs a moving average of parameters by periodically sampling subset of the learning trajectory, and SWA enables to accumulate virtually before proceeding to the pruning step (line 8; Algorithm 1).
particles throughout the training. We then average all particles
As illustrated in Figure 1, algorithm achieves superior performance, which is par with that of an ensemble consisting of two sparse networks. This is quite remarkable considering that our solu- tion achieves this level of performance while having significantly lower inference the ensembling approach. Further ablation studies presented in Š 4.2 and Appendix C.4 also validate that both ingredients independently contribute our algorithm, with each playing erucial role in achieving superior performance.
compared to
3.3 SWAMP: A LOSS LANDSCAPE PERSPECTIVE
In this section, we explore step-by-step whether the characteristics of IMP introduced 3.1 also hold for SWAMP along with highlighting the strengths of SWAMP. To begin with, we examine the linear connectivity of SWAMP particles in a single cycle. Although Frankle et al. (2020) empirically proves pair-wise linear connectivity, it remains uncertain whether this holds true for the convex combination of more than two particles. Figure 2, we visualize the loss surface of IMP-trained particles along with the veight-averaged particle. We notice that weight averaging fails at the earlier stages of IMP due to the highly non-convex nature of the landscape. However, as sparsity increases, particles tend to locate in the same wide basin which enables weight-averaging. Such
4

---

Published as a conference paper at ICLR 2024
Table 1: Trace of Hessian Tr(H) evaluated across training data. In most cases, SWAMP with mul- tiple particles exhibits smaller trace value, i.e., finds flatter minima, compared to others. Reported values are averaged over three random seeds, and the best and second- best results are boldfaced and underlined, respectively. Refer to Appendix C.7 for the results on CIFAR-100.
Sparsity
Training
20%
50%
75%
90%
CIFAR-10
SGD
1784.64 +152.95
1704.69
144.557
1719.49
1114225
1710.28
1885.84
(WRN-28-2)
SWAMP (N 1)
500.24 1: 50.00
500.70
4: 21.57
543.28
4: 31.27
579.79
20,69
SWAMP (N 4)
463,90 +: 24.84
467.45
+ 35.96
515.54
+: 20.23
533.62
4.27
CIFAR-IO (VGG-13)
SGD
1155.64 +: 51.21
1231.69
+ 55.21
1143,70
+ 65.25
1197.73
21,04
SWAMP (N- = 1)
427.92 1 11.15
424.85 1 6.84
450.44
1. 20.82
473.89
55.69
SWAMP (N
432.45 + 28.87
403.32
15,40
403.42
40.56
449.36
16,30
Sparsity 20%
Sparsity 75%
Sparsity 95%
Accuracy (%)
94.5
Accuracy (%)
95.0
94.5
94.0
94.8
Accuracy
94.2
93.5
94.6
94.0
WA
P1
P2
P3
P4
WA
P1
P2
P3
P4
WA
P1
P2
P3
P4
Figure 3: Bar plots depicting the accuracy of individual particles involved the averaging process of the SWAMP algorithm. While the averaged weight (denoted as WA) may not outperform indi- vidual 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 2 on the test split of CIFAR-10. We refer readers to Appendix C.9 for the same plot for CIFAR-100.
finding is in line with Frankle et al. (2020) that demonstrated the case of finding low-loss curve with a smaller network compared to a larger one, i.e., a sparse network tends to be more stable. Additionally, it further demonstrates that our algorithm benefits more with sparser networks.
Figure provides additional evidence showing that the weight-averaged solution indeed superior its individual members, other than in the cases where the dense network is not yet stabilized. Better, the notable performance gap between individual particles promotes the need for weight- averaging. We further quantify the flatness of the local minima through the trace of Hessian em- ploying the Power Iteration algorithm (Yao et al., 2020). Higher Hessian trace value implies that the converged local minimum exhibits high curvature. The results in Table validate that SWAMP locates a flatter sparse network compared to IMP, only except for earlier cycles.
Finally, we check whether consecutive solutions from SWAMP cycles are linearly connected key to the success of IMP pointed out by Paul et al. (2023) which indeed turns out to be true according to Figure 4. Not only is this true, but our method also exhibits minimal variance and maintains highly stable trajectory throughout the pruning process, suggesting that SWAMP finds flat and well-connected basin. To this end, we provide empirical evidence that our method effectively identifies flat minima while retaining the desirable properties of IMP, resulting in single highly sparse network that outperforms IMP.
EXPERIMENTS
4.1
MAIN RESULTS: IMAGE CLASSIFICATION TASKS
Baseline approaches. In addition to IMP with weight rewinding (Frankle et al., 2020), our method compared to a list of pruning techniques. This includes one-shot pruning methods (SNIP (Lee et al., 2019); GraSP (Wang et al., 2020); SynFlow (Tanaka et al., 2020)), dense- conse-to-sparee training with dynamic masks (DST (Liu et al., 2020); RigL (Evci et al., 2020)), SAM-optimized IMP (Na al., 2022), and Lottery Pools (Yin et al., 2023).
5

---

Published as a conference paper at ICLR 2024
IMP
----.
SWAMP
7.0
7.0
Classification error (%6
Classification error (%)
6.5
6.6
6.0
6.2
5.5
5.8
5.0
5.4
3
8
10
11
12
13
3
10 11
12 13
Sparsity level (IMP cycle)
Sparsity level (IMP cycle)
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 CIFAR-1O, and we refer the reader Appendix C.10 for the same plot for CIFAR-100.
Table 2:
Classification accuracy on residual networks. SWAMP outperforms all the baselines across varying sparsities. Reported values are averaged over three random seeds, and the best and second- best results are boldfaced and underlined, respectively.
Sparsity
Method
50%
75%
90%
95%
SNIP
92.63+0.56
92.13+021
91.02=0011
88.97=0.35
SynFlow
92,9940.46
92.98=00.8
90,98=0.0
89.53=0020
GraSP
92.34-10.26
90.74=+0.77
89.70=0060
88. 38.78:20.2
RigL
93.59+0.15
93,09=00.3
91.81
90.73=00166
CIFAR-IO
DST
94.04-20.37
93.62++0.12
92.59=0.10
92. 10.0.
(WRN-28-2)
IMP
93.97=0.16
94.02++0.5
93,90=0015
93.58=0009
IMP + SAM
94.06-20.0
94. 04.54
93,89=0.53
92.03=1.95
Lottery Pools
94.39+0.16
94.28+0.14
94.16=0.11
93.43=0025
SWAMP (ours)
94.74
94.88
94.73.20.0
94.23=2011
SNIP
71.98
72.16+t077
69.82+1.14
68. 18.10==028
SynFlow
74.22:4060
74.60=0.54
74. 17:777
GraSP
70.75+11.8
48.17-0048
66.96=0005
65. =5.80=089
RigL
75.05
73.37=a.11
71.23
70.23
CIFAR-100
DST
7..7
74.02++1.3
72.79=1.16
71.41=0.3
(WRN-32-4)
IMP
75.40+20.3
75.72+t0:1
25.24=022
74. 00-337
IMP + SAM
75.63+
76.19+0.81
75.85:+0.8
75.07=065
Lottery Pools
76.31+0.51
76.17#1.03
75.84=0.67
75.14=0.99
SWAMP (ours)
77.29+0.5
77.35
77.14
76.48=00.3
Experimental setup. Our method is evaluated on diverse image classification benchmarks, which include CIFAR-10, CIFAR-I00, Tiny-ImageNet and ImageNet datasets. Throughout the experi- ments, residual networks (He et al., 2016) and VGG networks (Simonyan & Zisserman, 2015) as a basis: WRN-28-2 and VGG- 3 for CIFAR- WRN-32-4 and VGG-16 for CIFAR- 100; R18 for TmngeImageNet; and R50 for ImageNet. Unless specified, we set the number of SWAMP particles N = 4 and pruning ratio = 0.2. Refer to Appendix B for further experimental details.
Table 2 presents the performance of SWAMP together with other baseline methods on the CIFAR-10 10 and CIFAR-100 datasets, respectively. Compared to other baseline methods, SWAMP consistently achieves the highest classification accuracy across all sparsity levels and models when evaluated on the CIFAR-10 10 and CIFAR-100 test sets; we defer results on VGG networks and Tiny-ImageNet to Appendix C. To see the uncertainty quantification aspects, in Appendix C.6, we report negative log-liielihoods (NLLs) well. Again, SWAMP achieves the best NLL in all settings. Furthermore, Table 3 highlights that our method consistently outperforms IMP on ImageNet, a large-scale dataset that is known to be hard prune.

---

Published as a conference paper at ICLR 2024
Table Classification accuracy with ResNet-50, which trained on ImageNet-Train, validated on IN-Valid, and tested on ImageNet- V2, ImageNet-Rendite (denoted by IN-R), and ImageNet- Sketch (denoted by IN-S). Reported values are averaged over three trials.
Sparsity
Method
0%
45.9%
68.8%
80.3%
86.0%
88,9 %
IMP
76.25=0044
76.40=a.10 10
76. 13:0.99
7,,54400m7
74.22=00.0
71,66*009
IN-Val
SWAMP
76.562008
70.51
75.69=0.0
74.25=0.03
IMP
64.18
64.11
63.78=0.05
62.77
61.20=a019
59.00
IN-V2
SWAMP
64.34
64.06=0.40
63,43=+024
61.82=0.15
59,44+007
IMP
35.38
15.04
34.71
34.05
32.85
30.92
IN-R
SWAMP
37.12
36.61
35.611022
14.14
32.21 20.73
IMP
23.90=0a2
23.86
23.74=0.32
22.88
21.60=0.24
19.32+4046
IN-S
SWAMP
25.29
24.95=0.17
24.05
22.43=0.02
20.08
Table 4: Ablation study to validate SWAMP's efficacy in two aspects: (i) mask generation and (ii) sparse training. Except for earlier iterations, i.e., low sparsity regime, our method excels both areas relative to vanilla SGD optimization. Reported values are classification accuracy averaged over three random seeds, and the best results are boldfaced. Refer to Appendix C.4 for VGG networks.
Sparsity
Mask
Training
50%
75%
90%
95%
SGD
SGD
93.97=0016
94.02=0023
93.90
93.58#0009
CIFAR-10 (WRN-28-2)
SGD
SWAMP
94.85=0.05
94.91 =0.09
94.48
93.99#0.1
SWAMP
SGD
94.15
94.37
94.26110.0
13.83
SWAMP
SWAMP
94.74=0.04
94.88=0000
94.73..011
94.23
4.2 ABLATION STUDIES
Does SWAMP find better mask? To validate that SWAMP indeed identifies a superior mask compared to IMP, we conduct an experiment with two different masks: (i) mask obtained from IMP, and (ii) a mask obtained from SWAMP. At predetermined fixed sparsity level, we initially train our model using SGD (or SWAMP) and the mask from IMP, At the same time, we train the model using SGD (or SWAMP) and the mask from SWAMP at the same sparsity level. These two processes differ only in the masks utilized, while the training approach and the fixed sparsity level remain the same. Table 4 presents clear evidence that the SWAMP mask consistently outperforms its counterpart in terms of performance, with the exception of the WRN-28-2 model at 50% and 75% sparsity levels when trained using SWAMP, This result shows that SWAMP generates befler sparse mask than IMP.
Does SWAMP offer a better sparse training? In Table we can also verify that SWAMP of- fers better sparse training compared to IMP. By comparing the results between SGD training and SWAMP training using the same mask, it becomes evident that SWAMP consistently outperforms across all masks, sparsity levels, and models. It verifies that SWAMP effectively guides the weight particles towards converging into flat local minima, resulting in improved generalization on the test split. The induced flatness of the local minima through SWAMP's weight distribution contributes to enhanced performance and robustness of the model as we discussed in 3.3.
Two averaging strategies of SWAMP. As described in $ 3.2, we here investigate how stochastic weight averaging and multi-particle averaging contribute to the final performance of SWAMP. In Table 5, throughout all four sparsities, applying only one of the two techniques displays better per- formance than IMP (bottom-row) but clearly lower than SWAMP (top-row). We conclude that two ingredients complement each other, achieving optimal performance when applied together. Further in Table 6, we conduct an empirical analysis to investigate the correlation between the number of particles and the performance of SWAMP. We provide additional ablation studies in Appendix C.4.
7

---

Published as a conference paper at ICLR 2024
Table 5: Ablation study on the impact of the two main components SWAMP: averaging multi- ple particles (denoted by Multi) and averaging learning trajectory (denoted by SWA). Our findings indicate that the best performance achieved when both techniques are employed. Reported clas- sification accuracies are averaged over three random seeds, and the best and second-best results are boldfaced and underlined, respectively. Refer to Appendix C.4 for VGG networks.
Sparsity
Multi
SWA
50%
75%
90%
95%
94.742001
94,88
94.73:0.10
94.23:0.11
CIFAR-10 (WRN-28-2)
X
94,43
94.44
94.37
93.73
94.62+006
94.67+0066
94.35=0.66
93.97=0.10
93.97.10.6
94.02.1..23
93.90.a015
93.58
Table 6: Ablation study the number of SWAMP particles. The finding indicates that the per- formance improves with an increase in the number of particles. Reported values are classification accuracy averaged over three random seeds, and the best and second-best results boldfaced and underlined, respectively. Refer to Appendix C.4 for VGG networks
Sparsity
# particles
50%
75%
90%
95%
94.62 .4.002200066
94.67:20.06
94.35::006
93.97=0.10
CIFAR-10 (WRN-28-2
94.57 +0004
94.59+00.0
94.38+:017
94.05=0.16
4
94.74:10.01
94.88
94.73+10110
94.23
94.80
94.90
94.74:1.10
94.21
4.3 FURTHER REMARKS
Parallelization strategy in distributed training, A notable limitation of the SWAMP algorithm is its training cost, which scales linearly with the number of particles. The training cost is typically not a major issue when working with small datasets like CIFAR-10/100 but it becomes significant when handling large datasets like ImageNet. Consequently, we suggest parallelizing the particles across machines when implemented within distributed training environments, common practice for handling large-scale models and datasets. This strategy incurs virtually no additional costs compared to IMP, except for the extra memory required for storing the averaged parameters. Indeed, we put this strategy into practice during our ImageNet experiments, and SWAMP achieved outstanding results while incurring almost the same training expenses as IMP (cf. Table 3).
Reducing training costs SWAMP. The parallelization strategy mentioned above is exclusively applicable to distributed training setups. Consequently, we further introduce practical methods to reduce the training costs of the SWAMP algorithm. which can be used even in non con-distributed train- ing environments: (I) Employing multiple particles only in the high-sparsity regime mitigates the significant training costs mainly encountered in low-sparsity regimes. Table 7 demonstrates that this approach reduces training costs by a factor of 3 to 4 with minimal performance degradation. Here, we initiate training with a single-particle SWAMP for the first ten IMP cycles, achieving a sparsity level of 90%, and then transition to using four particles afterward. (2) Increasing the pruning ratio decreases the number of pruning cycles necessary to achieve a certain sparsity level and thus sig- nificantly reduces total training costs. Table verifies that SWAMP is proficient in pruning even when using higher pruning ratio, highlighting distinctive advantage of SWAMP compared to IMP. These findings, combined with those in Table 1, align with Lemma 3.1 in Paul al. (2023); a smaller Hessian eigenvalue, indicating flatter minima, enhances robustness to SGD noise and mak- ing it more likely to restore matching performance.
Extension to language tasks and dynamic pruning. Up to this point, we have demonstrated that multi-particle averaging benefits IMP in iamge classification tasks. However, is important to note that SWAMP can be easily applied to different pruning techniques across range of tasks. To clarify, we present two distinct extensions: (1) Table 9 further confirms that SWAMP outperforms IMP in language tasks, where experimental details are available in Appendix C.2. (2) We also illustrate how
8

---

Published as a conference paper at ICLR 2024
Table 7: Further comparison between (i) IMP, (ii) SWAMP, and (iii) the cost-efficient version of the SWAMP algorithm (denoted by SWAMP+) in terms of accuracy and total training FLOPs. Reported values are averaged over three random seeds.
Sparsity 95%
Sparsity 98%
Method
Accuracy
GFLOPs
Accuracy
GFLOPs
IMP
93.58=00.9
1.19
89.05=1.39
1.23
CIFAR-I0
SWAMP
94.23.1.11
4.75
90.85.0.47
4.92
(WRN-28-2)
SWAMP+
94.32:1..4
1.39
90.51
1.56
IMP
74.60
5.25
70.74
5.38
CIFAR-100
SWAMP
76.48st0.7
20.99
72. 0.
(WRN-32-4)
SWAMP+
16.19
5.96
71.90
6.51
Table 8: Results on CIFAR for larger pruning ratio of E 0.3, 0.5) after 13, 8, 6, IMP cycles, respectively. Reported values are averaged over three random seeds. The numbers within parentheses indicate the difference from the IMP baseline.
Method
0.2
0.3
0.4
0.5
CIFAR-10
IMP
93.58.009
93.36.10.35
13.11
93.77.0.12
(WRN-28-2)
SWAMP
94.23 (+0.65)
93,83 (+0.7)
94.41 (-130)
94. 14.68 (-091)
CIFAR-100
IMP
74.60
24.34
73.87
73.85
(WRN-32-4)
SWAMP
76.48 (+1.88)
77.3 11077 (+3.7))
76.26:10.55-233
(-2.39)
76.50.0.40 (-265)
Table 9: F1 score and accuracy on the development set of MRPC and RTE, respectively. SWAMP displays improved performance language tasks as well. Reported values averaged over three random seeds, and the numbers within parentheses indicate the difference from the IMP baseline.
Sparsity
Method
0%
40.9%
52.2%
61.3%
MRPC
IMP
92.6020.0
91.02=a31
88.31
86.81 +0:50
(RoBERTa)
SWAMP
92.81 10.38 (-021)
91.51 1:039 (40.49)
89 29.440.0.40 (-1.9)
18.424 (+131)
RTE
IMP
80. 13
73. 1.55
63 30-337
57.6 16.4..1..9
(RoBERTa)
SWAMP
80,30 (+0.17)
75.40 (+2.15)
68 91 11.08 (-561) 62.63 1.33 (+5.17)
SWAMP has the potential to enhance dynamic sparse training methods, along with additional results with RigL (Evei al., 2020) Appendix C.3.
CONCLUSION
Drawing inspiration from previous research on the relationship between iterative magnitude prun- ing and linear mode connectivity, we extended the angle-particle scenario to incorporate multiple particles. Our initial empirical findings demonstrated that multiple models trained with different SGD noise but sharing the same matching ticket can be veight-averaged without encountering loss barriers. We further observed that the averaged particle results flat minima with improved gen- eralization performance. In light of these insights, we introduced SWAMP, novel iterative global pruning technique. We also established that SWAMP preserves the linear connectivity between con- secutive solutions, a critical factor contributing to the efffectiveness of IMP. Extensive experiments showed that SWAMP generates superior sparse masks and effectively trains sparse networks over other baseline methods. A theoretical analysis exploring why the convex hull of the particles in weight space forms a low-loss subspace would be valuable direction for future research. Further, investigating the underlying principles and mathematical properties of the convex hull of the solu- tion particles and its relationship to the low-loss subspace could provide insights into the behavior and effectivenees of the SWAMP algorithm.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENT
This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (N0.2019-0-00075, Artificial Intel- ligence Graduate School Program (KAIST), No.2022-0-00713. Meta-learning Applicable to Real- world Problems, No.2022-0-00184 Development and Study of Al Technologies to Inexpensively Conform to Evolving Policy on Ethics), and the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (NRF20022RA55A70336666
REPRODUCIBILITY STATEMENT
Our algorithm is built on Pytorch 1.10.2 (Paszke et al., 2019), which is available under BSD-style license'. All experiments are conducted on NVIDIA RTX 2080 and NVIDIA RTX 3090 machines. Basic experimental setup including network and dataset choice is listed in 4.1. In Appendix B, we further provide all detailed hyperparameter settings to ensure the fair comparison between the baselines and our algorithm. We also include our code in the supplementary material.
ETHICS STATEMENT
The paper does not raise any ethical concerns. We only utilize publicly available datasets and python packages adhering the appropriate licenses.
REFERENCES
Arsenii Ashukha, Alexander Lyzhov, Dmitry Molchanov, and Dmitry Vetrov. Pitfalls of in-domain uncertainty estimation and ensembling in deep learning. In International Conference Learning Representations (ICLR), 2020. 18
Roy Bar Haim, Ido Dagan, Bill Dolan, Lisa Ferro, Danilo Giampiccolo, Bernardo Magnini, and
Idan Szpektor. The second PASCAL recognising textual entailment challenge, 2006. 16
Luisa Bentivogli, Ido Dagan, Hoa Trang Dang, Danilo Giampiccolo, and Bernardo Magnini. The fifth PASCAL recognizing textual entailment challenge, 2009. 16
G. W. Benton, W. J. Maddox, S. Lotfi, and A. G. Wilson. Loss surface simplexes for mode con- necting volumes ans fast ensembling. In Proceedings of the 38th International Conference Machine Learning (ICML), 2021.
Pratik Chaudhari, Anna Choromanska, Stefano Soatto, Yann LeCun, Carlo Baldassi, Christian Borgs, Jennifer Chayes, Levent Sagun, and Riccardo Zecchina. Entropy-SGD; Biasing gradi- descent into wide valleys. International Conference on Learning Representations (ICLR), 2017.
Tianlong Chen, Yongduo Sui, Xuxi Chen, Aston Zhang, and Zhangyang Wang. A unified lottery ticket hypothesis for graph neural networks. In Proceedings of the 38th International Conference on Machine Learning (ICML), 2021. 14
Ido Dagan, Oren Glickman, and Bernardo Magnini. The PASCAL recognising textual entailment challenge. In Machine learning challenges. evaluating predictive uncertainty, visual object clas- sification, and recognising tectual entailment, pp. 177-190. Springer, 2006. 16
William B Dolan and Chris Brockett. Automatically constructing corpus of sentential paraphrases.
In Proceedings of the International Workshop on Paraphrasing, 2005.
Xin Dong, Shangyu Chen, and Sinno Pan. Learning to prune deep neural networks via layer-wise optimal brain surgeon. In Advances in Neural Information Processing Systems (NIPS), 2017.
10

---

Published as a conference paper at ICLR 2024
Felix Draxler, Kambis Veschgini, Manfred Salmhofer, and Fred Hamprecht. Essentially no barriers in neural network energy landscape. In Proceedings of the 35th International Conference on Machine Learning (ICML), 2018. 2, 14
Utku Evci, Trevor Gale, Jacob Menick, Pablo Samuel Castro, and Erich Elsen. Rigging the lottery: Making all tickets winners. In Proceedings of the 37th International Conference on Machine Learning (ICML), 2020. 5, 9, 14, 15, 17
Utku Evci, Yani Ioannou, Cem Keskin, and Yann Dauphin. Gradient flow in sparse neural networks and how lottery tickets win. Proceedings of the AAAI Conference Artificial Intelligence, 36(6), Jun. 2022. 14
Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur, Sharpness-aware minimiza- tion for efficiently improving generalization. In International Conference on Learning Represen- tations (ICLR), 2021. 2
Stanislav Fort and Stanislaw Jastrzebski. Large scale structure of neural network loss landscapes. In Advances in Neural Information Processing Systems (NeurIPS), 2019. 1,3
Stanislav Fort, Gintare Karolina Dziugaite, Mansheej Paul, Sepideh Kharaghani, Daniel M Roy, and Surya Ganguli. Deep learning versus kernel learning; an empirical study of loss landscape geometry and the time evolution of the neural tangent kernel. In Advances in Neural Information Processing Systems (NeurIPS), 2020. 3
Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. In International Conference on Learning Representations (ICLR), 2019. 1, 2, 3, 14, 16, 17
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Linear mode con- nectivity and the lottery ticket hypothesis. In Proceedings of the 37th International Conference on Machine Learning (ICML), 2020. 2, 3, 5, 14, 16, 17
Jonathan Frankle, Gintare Karolina Dziugaite, Daniel Roy, and Michael Carbin. Pruning neural net- works at initialization: Why are we missing the mark? In International Conference on Learning Representations (ICLR), 2021. 14
Timur Garipov, Pavel Izmailov, Dmitrii Podoprikhin, Dmitry P Vetrov, and Andrew G Wilson. Loss surfaces, mode connectivity, and fast ensembling of DNNs. In Advances in Neural Information Processing Systems (NeurIPS), 2018. 3, 14
Danilo Giampiccolo, Bernardo Magnini, Ido Dagan, and Bill Dolan. The third PASCAL recognizing textual entailment challenge. In Proceedings of the ACL-PASCAL workshop on textual entailment and paraphrasing, pp. 1-9. Association for Computational Linguistics, 2007. 16
Sharath Girish, Shishira R Maiya, Kamal Gupta, Hao Chen, Larry S Davis, and Abhinay Shrivastava. The lottery ticket hypothesis for object recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2021. 14
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q Weinberger. On calibration of modern neural networks. Proceedings of the 34th International Conference on Machine Learning (ICML), 2017. 18
S. Han, J. Pool, J. Tran, and W. Dally. Learning both weights and connections for efficient neural network. In Advances in Neural Information Processing Systems (NIPS). 2015. 14
Babak Hassibi and David Stork. Second order derivatives for network pruning: Optimal brain surgeon. In Advances Neurul Information Processing Systems (NIPS), 1992. 14
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR). 2016. 6
11

---

Published as a conference paper at ICLR 2024
P. Izmailov, D. Podoprikhin, T. Garipov, D. Vetrov, and A. G. Wilson. Averaging weights leads to wider optima and better generalization. In Proceedings of the 34th Conference on Uncertainty in Artificial Intelligence (UAI), 2018. 2, 4, 14, 16
Pavel Izmailov, Wesley Maddox, Polina Kirichenko, Timur Garipov, Dmitry Vetrov, and An- drew Gordon Wilson. Subspace inference for bayesian deep learning. In Proceedings of the 36th Conference on Uncertainty in Artificial Intelligence (UAI), 2020. 14
Steven A Janowsky. Pruning versus clipping in neural networks. Physical Review A, (91))6600,, 1989. 14
Neha Mukund Kalibhat, Yogesh Balaji, and Soheil Feizi. Winning lottery tickets in deep generative models. Proceedings of the AAAI Conference on Artificial Intelligence, 35(9), May 2021. 14
Diederick P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015. 16
Balaji Lakshminarayanan. Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Advances in Neural Information Processing Systems (NIPS), 2017. 3
Brett W Larsen, Stanislav Fort, Nic Becker, and Surya Ganguli. How many degrees of freedom do we need to train deep networks: a loss landscape perspective. In International Conference on Learning Representations (ICLR), 2022. 14
Yann LeCun, John Denker, and Sara Solla. Optimal brain damage. In Advances in Neural Informa-
tion Processing Systems (NIPS), 1989. 1, 14
Namhoon Lee, Thalaiyasingam Ajanthan, and Philip HS Torr. Snip: Single-shot network pruning based on connection sensitivity. In International Conference on Learning Representations (ICLR), 2019. 5, 14, 15
Junjie Liu, Zhe XU, Runbin SHI, Ray C. C. Cheung, and Hayden K.H. So. Dynamic sparse train- ing: Find efficient sparse network from scratch with trainable masked layers. In International Conference on Learning Representations (ICLR), 2020. 5, 14, 15
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprini arXiv: I 1907. / 1692, 2019. 16
Xiaolong Ma, Geng Yuan, Xuan Shen, Tianlong Chen, Xuxi Chen, Xiaohan Chen, Ning Liu, Ming- hai Qin, Sijia Liu, Zhangyang Wang, et al. Sanity checks for lottery tickets: Does your winning ticket really win the jackpot? In Advances in Neural Information Processing Systems (NeurIPS), 2021. 14
Clara Na, Sanket Vaibhav Mehta, and Emma Strubell. Train flat, then compress: Sharpness Sharpness-aware aware minimization learns more compressible models. In Findings of the Association for Computational Linguistics: EMNLP 2022, 2022. 5, 15
Sejun Park, Jaeho Lee, Sangwoo Mo, and Jinwoo Shin. Lookahead: A far-sighted alternative of magnitude-based pruning. In International Conference on Learning Representations (ICLR), 2020. 14
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-Performance Deep Learning Library. In Advances in Neural Information Processing Systems (NeurIPS), 2019. 10, 15
Mansheej Paul, Feng Chen, Brett W. Larsen, Jonathan Frankle, Surya Ganguli, and Gintare Karolina Dziugaite. Unmasking the lottery ticket hypothesis: What's 's encoded in a winning ticket's mask? In International Conference on Learning Representations (ICIR), 2023. 1, 3, 5, 8, 14, 17, 19
12

---

Published as a conference paper at ICLR 2024
Alex Renda, Jonathan Frankle, and Michael Carbin. Comparing rewinding and fine-tuning in neural network pruning. In International Conference on Learning Representations (ICLR), 2020. 1, 2, 14, 15
Herbert Robbins and Sutton Monro. A stochastic approximation method. The annals of mathemati- cal statistics, pp. 400-407, 1951. 2, 16
Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. International Conference on Learning Representations (ICLR), 2015. 6
Jingtong Su, Yihang Chen, Tianle Cai, Tianhao Wu, Ruiqi Gao, Liwei Wang, and Jason Lec. Sanity-checking pruning methods: Random tickets win the jackpot. In Advances in Neural Information Processing Systems (NeurIPS), 2020. 14
Hidenori Tanaka, Daniel Kunin, Daniel L Yamins, and Surya Ganguli. Pruning neural networks without any data by iteratively conserving synaptic flow. In Advances Neural Information Processing Systems (NeurIPS), 2020. 5, 14, 15
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. In International Conference Learning Representations (ICLR), 2019. 16
Chaoqi Wang. Guodong Zhang, and Roger Grosse. Picking winning tickets before training by preserving gradient flow. In International Conference on Learning Representations (ICLR), 2020. 5, 14, 15
Mitchell Wortsman, Maxwell C Horton, Carlos Guestrin, Ali Farhadi, and Mohammad Rastegari. Learning neural network subspaces. In Proceedings of the 38th International Conference on Machine Learning (ICMI.), 2021. 3, 14
Mitchell Wortsman, Gabriel Ilharco, Samir Ya Gadre, Rebecca Roelofs, Raphael Gontijo-Lopes. Ari S Morcos, Hongseok Namkoong, Ali Farhadi, Yair Carmon, Simon Kornblith, et al. Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing in- ference time. In Proceedings of the 39th International Conference Machine Learning (ICML), 2022. 14
Zhewei Yao, Amir Gholami, Kurt Keutzer, and Michael W Mahoney. Pyhessian: Neural networks through the lens of the hessian. In 2020 IEEE international conference on big data (Big data), 2020. 5
Lu Yin, Shiwei Liu, Meng Fang, Tianjin Huang, Vlado Menkovski, and Mykola Pechenizkiy. Lot- tery pools: Winning more by interpolating tickets without increasing training or inference cost. Proceedings of the AAAI Conference Artificial Intelligence, 37(9), Jun. 2023. 5, 14, 15
Haoran You, Chaojian Li, Pengfei Xu, Yonggan Fu, Yue Wang. Xiaohan Chen, Richard G Baraniuk, Zhangyang Wang, and Yingyan Lin. Drawing early-bird tickets: Towards more efficient training of deep networks. In International Conference on Learning Representations (ICLR), 2020. 14
Hattie Zhou, Janice Lan, Rosanne Liu, and Jason Yosinski. Deconstructing lottery tickets: Zeros, signs, and the supermask. In Advances in Neural Information Processing Systems (NeurIPS), 2019. 14
13

---

Published as a conference paper at ICLR 2024
A RELATED WORKS
Lottery ticket hypothesis. Neural network pruning techniques aim to identify parameters that can be removed from the network without affecting its performance. The simplest yet effective pruning strategy is magnitude-based pruning, which involves removing parameters with small mag- nitudes (Janowsky, 1989; Han et al., 2015). While there is no guarantee that zeroing the weights that are close to zero will mitigate the increase in training loss (LeCun et al., 1989; Hassibi & Stork, 1992), magnitude-basee
pruning can still be effective as it minimizes changes in the output of each layer (Dong et al., 2017: Park et al., 2020). Empirical studies have demonstrated that iterative mag- nitude pruning (IMP)- - appllyng magnitude-based pruning multiple times- can effectively remove a large proportion of weights in neural networks (Frankle & Carbin, 2019; Renda et al., 2020; Frankle et al., 2020).
A large body of work tries to unravel the success of lottery tickets given the increasing interest across many applications (Chen et al., 2021; Kalibhat et al., 2021; Girish et al., 2021). In light of existing samity- checkic methods such as score inversion and layer-wise random shuffling, IMP passes those sanity checks while other one-shot pruning methods fail, which shows the matching tickets do contain useful information regarding network weights (Ma et al., 2021; Frankle et al., 2021; Su et al., 2020). Zhou et al. (2019) focus on the zeroing out weights explaining such operation behaves resembles neural network training along with putting emphasis on preserving the sign of weights. Renda et al. (2020) proposes learning rate rewinding which is empirically shown to achieve superior performance compared to weight rewinding and fine-tuning. Inspired by Gordon's 's escape theorem, Larsen et al. (2022) theoretically relates the matching ticket to burn-in random subspace crossing the loss sublevel set.
More pruning literature. Lottery tickets are primarily hindered by their substantial training costs whereas there exist one-shot pruning methods which only require cost corresponding to single cycle of IMP. As a pioneering work, Lee et al. (2019) proposes SNIP which approximates the sen- sitivity of a connection, i.e.. the impact of removing single connection. GraSP (Wang et al., 2020) suggests preserving the gradient flow as a pruning criterion, and SynFlow (Tanaka et al., 2020) avoids the layer-collapse phenomenon without looking the training data. With the intention of catching up with such computationally efficient algorithms, You et al. (2020) proposes carly-bird tickets which significantly reduces the cost of the mask-searching step. As IMP provides so-called sparse-to-spee training, i.e., the mask is pre-defined, some work investigates dense-to-se to-ssrs train- ing with dynamic masks (Evei et al., 2020; Liu et al., 2020). Liu al. (2020) employs trainable mask layer jointly optimizing masking threshold and weight throughout a single cycle. Evci et al. (2020) proposes a rigged lottery (Rigl.) that regrows the dead weights yet with a large gradient flow. It has been demonstrated empirically that Rigl. can escape local minima and discover improved basins during optimization process.
Loss landscape and weight averaging. Recent studies have shown that the loss landscape of modern deep neural networks exhibits mode-comnectivity. which means that the local minima found through stochastic optimization are connected by continuous low-loss paths (Draxler et al., 2018; Garipov et al., 2018). It motivates researchers to explore subspaces in the weight space that contain high-performing solutions (Izmailov et al., 2020; Wortsman et al., 2021). One promising approach to obtain such solutions is weight averaging, which involves taking the average of weights across multiple points (Izmailov et al., 2018; Wortsman et al., 2022). The main assumption of the weight averaging strategy is that all points are located in the same basin of attraction, which may not hold in the case of sparse networks where different sparse structures may correspond to different areas in the loss landscape. However. recent works (Frankle et al.. 2020; Evci et al., 2022; Paul et al., 2023) suggest that this assumption may hold even for sparse networks, as there is a linear connectivity between the solutions of different sparsities obtained through iterative magnitude pruning. Frankle et al. (2020) show that two neural networks with the same matching ticket yet trained with different SGD noises are linearly connected without a high loss barrier in between. Evei et al. (2022) argues that lottery tickets are biased towards the final pruned solution, i.e., winning tickets enables relearn- ing the previous solution. Paul et al. (2023) identify the existence of a low-loss curve connecting two solutions from consecutive IMP cycles, which guarantees that SGD noise cannot significantly alter the optimization path. Recently, Yin et al. (2023) propose to weight-averal the tickets generated from consecutive IMP iterations in order to obtain single stronger sparse network.
14

---

Published as a conference paper at ICLR 2024
B EXPERIMENTAL DETAILS
Our algorithm is built on Pytorch 1.10.2 (Paszke et al., 2019), which is available under BSD-style license". All experiments are conducted on NVIDIA RTX 2080 and NVIDIA RTX 3090 machines.
Pruning. Throughout our experiments, we exclusively perform weight pruning on the convolu- tional layer, leaving the batch normalization and fully-connected layers unpruned. Additionally, do not utilize any predefined layer-wise pruning ratio. It is worth noting that we empirically checked that including the fully-connected layers part of the prunable parameters does not affect the sults we have obtained. By default, we use the constant pruning ratio of a = 0.2 at all IMP pruning rounds, i.e,, retain 80% of weight. For the convenience of the readers, we provide approximate sparsity values at the 3rd, 6th, 10th, 13th cycles as 50%, 75%, 90%, and 95%, respectively, while the precise sparsity values are 48. 88.80, 73.79%, 89.26%, and 94.50%. Unless otherwise specified. we set the number of SWAMP particles to four.
Baselines. For the fair comparison in terms of training epochs, we train all the baselines listed in $ 4.1 until convergence. Merely increasing the number of training epochs does not bene- fit baselines; the performance of SWAMP cannot be pertained to its extensive training computa- tions. For one-shot pruning methods (Lee et al., 2019; Wang et al., 2020; Tanaka et al., 2020), we post-train for 300 epochs after obtaining the sparse mask and follow the same hyperparam- eter setting in the original paper. For DST (Liu et al., 2020), we explore E [1c - 7, le 5] for all models and datasets to match similar sparsity levels with other baselines. For RigL. (Evei al., 2020), we search plore E {0.005, 0.01, 0.05, 0.1, 0.5} for SAM (Na et al., 2022) with the same hyperparam- eter and optimizer setting of SWAMP. For Lottery Pools (Yin et al., 2023), we explore a e {0.01, 0.05, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.99), and greedily choose among 7 candidate models w.r.t. validation accuracy.
[0.3, 0.5, 0.7]. Here we fixed hyperparameter 5 100. We
Optimization. We utilize the SGD optimizer with a momentum of 0.9 and a learning rate schedule that follows the cosine-decaying method. We choose the cosine-decaying schedule as it yields better results compared to the step decay schedule commonly used in lottery ticket literature.
CIFAR-10/100. To obtain a matching solution, we first train the model for 10 epochs using a constant learning rate of 0.1, weight decay Ic-4, and a batch size of 128. Subsequently, we employ a cosine decaying learning rate schedule over 150 training epochs in each cycle of the IMP algorithm. During the final quarter of the training epochs, we collect particles for SWAMP, resulting in total of 38 particles in a single trajectory. The learning rate for this phase is set to a constant value of 0.05.
Tiny-ImageNer To obtain matching solution, we first train the model for 20 epochs using constant learning rate of 0.1, a weight decay 5e-4, and a batch size of 256. Subsequently, we employ cosine decaying learning rate schedule over 160 training epochs in each eycle of the IMP algorithm. During the final quarter of the training epochs, we collect particles for SWAMP, resulting in total of 40 particles in a single trajectory. The learning rate for this phase is set to a constant value of 0.05.
ImageNet. To obtain a matching solution, we first train the model for 30 epochs using a constant learning rate of 0.8, weight decay le-4, and a batch size of 2048. Subsequently, we employ cosine decaying learning rate schedule over 60 training epochs in each cycle of the IMP algorithm. During the final quarter of the training epochs, we collect particles for SWAMP, resulting in a total of 15 particles in a single trajectory. The learning rate for this phase is set to a constant value of 0.004. We establish distributed training environment comprising eight machines, and consequently, we distribute particles across these eight machines as discussed in $ 4.2. We also employed the learning rate rewinding (Renda al., 2020) instead of weight rewinding in ImageNet experiments.
15

---

Published as a conference paper at ICLR 2024
Algorithm 2 Stochastic Gradient Descent (i.e., SGD {,
Require: An initial neural network parameter wo, training dataset D, the number of training itera- tion T, learning rate 7), a loss function L, and SGD noise Ç. Solution wr.
Ensure:
for (1....,7]) do
2: Sample a mini-batch B C D.
3: Update parameter w. f- Wt-1 - V ww B).
4: end for
Algorithm 3 Stochastic Weight Averaging (i.e., SWAo-r(Wap (,D)
Require: An initial neural network parameter wo, training dataset D, the number of training itera- tion T, learning rate 1), a loss function L, and SGD noise &.
Ensure: Solution WSWA
Initialize USWA + and nswa < 0.
2: for / (1,..., T do
3: Sample a mini-batch D.
4: Update parameter f- Wi-1 17. V www,w B).
5: Periodically update WSWA f- (NSWA WSWA + w/)/((n + 1) and n2swA + nswa +
6: end for
C SUPPLEMENTARY RESULTS
C.I ALGORITHMS
Stochastic Weight Averaging Algorithms and 3 describe a detailed procedure for stochastic gradient descent (SGD; Robbins & Monro, 1951) and stochastic weight averaging (SWA; Izmailov et al., 2018), which are respectively denoted as SGDo main text of the paper for notational simplicity. For SWA, we took moving average over model copies sampled from the last 25% of the training epochs.
D) and SWAD-s E, D) in the
Iterative Magnitude Pruning Algorithm and Algorithm 5 respectively provide detailed pro- cedure for the vanilla iterative magnitude pruning (IMP; Frankle & Carbin, 2019) and IMP with weight rewinding RMP-WR; Frankle et al., 2020). While IMP rewinds the network the initial- ization after each cycle, IMP-WR rewinds the weights to the early epoch the training process. Prune(w, a) returns pruning mask where a)) of the remaining parameters pruned based on their magnitudes.
C.2 EXTENSION TO LANGUAGE TASKS
Table 9 verified the capability of SWAMP language tasks. Specifically, we fine-tuned the RoBERTa-Base model (Liu et al., 2019) on two subtasks from GLUE benchmark (Wang et al., 2019): Microsoft Research Paraphrase Corpus (MRPC; Dolan & Brockett, 2005) and Recognizing Textual Entailment (RTE; Dagan et al., 2006; Bar Haim et al., 2006; Giampiccolo et al., 2007; Ben- tivogli et al., 2009). We mainly followed the experimental configuration outlined the work of Liu et al. (2019), which includes specifications such as a token limit of 512, and utilization of the Adam optimizer (Kingma & Ba, 2015) with first-moment coefficient set 0.9 and second-moment coefficient set to 0.98.
To obtain a matching solution, we first train the model for 5 epochs using a constant learning rate of le-05. Subsequently, each cycle of the IMP algorithm. For SWAMP. we collect particles distributed across eight machines during the final half of the training epochs. Figure 5 summarizes the results for each IMP cycle, where we used the pruning ratio of = 0.9. We can readily find out that SWAMP outperforms the IMP baseline by noticeable margin on both datasets.
employ linear decaying learning rate schdule over 10 training epochs in
16

---

Published as a conference paper at ICLR 2024
Algorithm 4 Iterative Magnitude Pruning (Frankle & Carbin, 2019)
Require: Neural network parameter pruning mask m, training dataset D, the number of cycles for iterative magnitude pruning C, the number of iterations for each eycle T. pruning ratio a, and SGD noise Ç.
Ensure: Sparse solution w..T
1: Randomly initialize wo,0 and set mask m t- 1.
D Starts from random dense weights.
2: for E (1,..., C) do
3: Reset We,0 + Wo.T. m.
> Starts from the matching ticket. D Averages weights over trajectory. b Updates the mask based on magnitudes.
4: Train Wa,T SGDo-r(Wo. Ees D).
5: Prune m f- Prune(W..r.e a).
6: end for
Algorithm Iterative Magnitude Pruning with Rewinding (Frankle al., 2020)
Require:
Neural network parameter w, pruning mask m, training dataset DD the number of cycles for iterative magnitude pruning C, the number of iterations for each cycle T, pruning ratio a, SGD noise ç, and the number of iteration for matching ticket To.
Ensure: Sparse solution w..
Starts from random dense weights. 2: Train Wo.7. & SGDo-TT (Woo0 o m, Eo. D). >> Gets a matching ticket from the initialization.
1: Randomly initialize woo0 and set mask 1.
3:
for E fI,...,00 C.. do
4: Rewind w..0 <- Wo,7, o m.
D> Starts from the matching ticket. > Averages weights over trajectory. Updates the mask based on magnitudes.
5: Train wer SGDo-r(We.D Ec, D).
6: Prune m f- Prune(Wa.r,e
7: end for
C.3 EXTENSION DYNAMIC PRUNING METHODS
In this paper, we extend IMP to SWAMP leveraging particle averaging heavily inspired by Paul al. (2023). Although the primary focus of our paper the loss landscape analysis of IMP, one can readily apply SWAMP to dynamic pruning methods as well. Here we further extend our al- gorithm to encompass dense-to-sparse training, offering a means to mitigate the intensive training costs associated with IMP which includes multiple rounds of retraining. Our approach enhances RigL (Evci et al., 2020) by introducing multiple particles during the final training phase, integrated after completing 75% of the total training epochs. The results presented in Table 10 demonstrate that our method gracefully merges with Rigl., incurring only marginal training overhead. This result shows that SWAMP can be naturally applied to dynamic pruning methods.
C.4 FURTHER ABLATION STUDIES
As described 3.2, we conduct an empirical analysis to investigate the correlation between the number of particles and the performance of SWAMP. Table 11 presents our findings. demonstrating a strong positive correlation between the generalization performance of SWAMP and number of individual particles. These results indicate that maximizing the number of particles used in SWAMP can significantly enhance model performance. Additionally, they provide evidence supporting the mention made in 3.2 that both SWA and particle averaging contribute independently and comple- mentarily to the overall effectiveness of the SWAMP algorithm. We also provide additional ablation results with VGG networks in Table 12 and Table 13.
C.5 FURTHER RESULTS ON TINY-IMAGENET
Figure 6 depicts the results on Tiny-ImageNet. the case of ResNet-18 on Tiny-ImageNer SWAMP outperforms the ensemble performance of two sparse networks.
17

---

Published as a conference paper at ICLR 2024
RoBERTa-Base on MRPC
RoBERTa-Base on RTE
80
92
75
F1 score
Accuracy
70
90
65
.
60
IMP
.
SWAMP
SWAMP
86
55
1.0
0.8
0.6
0.4
1.0
0.8
0.6
0.4
1 Sparsity
1
Sparsity
Figure 5: Results for RoBERT MRPC and RTE. Reported values averaged over three random seeds.
Table 10: Further comparison between (a) RigL and (b) RigL + SWAMP on CIFAR-10. Reported values averaged over three random seeds.
Sparsity
WRN28x2
75%
90%
95%
(a)
93.09
91.81
90.73=0.16
(b)
93.41 +0.04
92.78
92.40:0.46
Sparsity
VGG-13
75%
90%
95%
(a)
93.18
92.4610.16
91.45
(b)
93.53
93.01
93.01
C.6 SPARSITY PLOTS
We provide supplementary results Figure 7, which clearly indicate that SWAMP performs better than IMP in terms of overall performance. Furthermore, report the negative log-likelihood (NLL) values in Tables 15 and 16, where we employed temperature scaling (Guo et al., 2017) for better in- domain uncertainty evaluation, as discussed in Ashukha al. (2020). Notably, SWAMP consistently outperforms other baselines in terms of NLL as well.
C.7 FLATNESS OF LOCAL MINIMA
We further provide the results on the trace of Hessian for CIFAR-100 in Table to show SWAMP with multiple particles finds flatter local minima than IMP leading to better generalization perfor- mance.
C.8 VISUALIZATION LOSS SURFACES
We provide additional loss landscape visualization plots for CIFAR-100 using WRN-32-4 model in Figure 8. We can observe a similar trend that SWAMP averages better as network sparsity grows.
C.9 PARTICLE-WIS PERFORMANCE
We provide the particle-wist wise accuracies for CIFAR-100 in Figure 9. Individual SWAMP particles outperform the averaee particle in lower sparsities, but we can observe the opposite trend in higher sparsities which is the area of our particular interest.
18

---

Published as a conference paper at ICLR 2024
Table 11: Ablation study on the number of SWAMP particles. Performance improves with an increase in the number of particles. Reported values are classification accuracy averaged over three random seeds, and the best and second-best results are boldfaced and underlined, respectively.
Sparsity
# particles
50%
75%
90%
95%
94.62 + 0.06
94.67 = 0.06
94.35 + 0.06
93.97 + 0.10
CIFAR-10 (WRN-28-2)
2
94.57 + 0.04
94.59 0.07
94.38 = 0.17
94.05 0.16
94.74 0.04
94.88 0.09
94.73 0.10
94.23 1 0.11
8
94.80 + 0.04
94.90 0.09
94.74 = 0.10
94.21 + 0.24
94.07 0.03
94.15 = 0.09
94.12 = 0.22
93.99 + 0.17
CIFAR-10
2
94.03 + 0.10
93.93 0.16
94.01 = 0.09
94.03 + 0.07
(VGG-13)
4
94.14 + 0.08
94.39 0.15
94.40 + 0.16
94.34 0.11
94.24 0.18
94.41 0.09
94.54 = 0.09
94.41 + 0.15
Ablation study to validate SWAMP's efficacy in two aspects: (i) mask generation and (ii) sparse training. Except for earlier iterations, i.e., low sparsity regime, method excels in both areas relative to vanilla SGD optimization. Reported values are classification accuracy averaged over three random seeds, and the best results are boldfaced.
Table 12:
Sparsity
Mask
Training
50%
7.%%
90%
95%
SGD
SGD
93.3420.7
93.53:20.15
93.58
93.62:0011
CIFAR-10
SGD
SWAMP
94.01
94.30
94.35
94. 18.013
(VGG-13)
SWAMP
SGD
93.55
93.844t0100
94.07
93.98=0007
SWAMP
SWAMP
14.08
94.39.10115
94.40
94.34=0..1
C.10 CONNECTIVITY BETWEEN SUCCESSIVE SWAMP SOLUTIONS
We present additional results on CIFAR- -10AR-100A regarding the linear connectivity between two consec- utive SWAMP solutions. Following Paul et al. (2023), we find that SWAMP solutions from different cycles are also linearly well-connected.
19

---

Published as a conference paper at ICLR 2024
Table 13: Ablation study the impact of the two main components of SWAMP; averaging multi- ple particles (denoted by Multi) and averaging learning trajectory (denoted by SWA). Our findings indicate that the best performance is achieved when both techniques employed. Reported clas- sification accuracies are averaged over three random seeds, and the best and second-bes results are boldfaced and underlined, respectively.
Sparsity
Multi
SWA
50%
75%
90%
95%
94.14 10.08
94.39.10.15
94.40..0.1
94.34_0.11
CIFAR-10 (VGG-13)
V
93.96=0003
94.09
94.08+0.18
94.05=0.11
94.07 1003
94.15
94. 12.
93,99_0.17
X
93.74
93.53=00115
93.58+0003
93.62=0.11
Matching threshald
IMP-1
-B-
IMP-2
SWAMP (ours)
40.00
1.80
Classification error (96)
36.00
1.60
34.00
1.50
32.00
30.00
1.30
0
20
40
60
80
100
20
40
60
80
100
Sparsity (%)
Sparsity (%)
Classification error (left) and negative log-likelihood as function of the sparsity. SWAMP achieves remarkable performance comparable to an ensemble IMP solutions while maintaining the same inference cost single model. Here in the notation IMP-n indicates the number IMP ensemble members. Notably, SWAMP demonstrates matching performance even at extremely sparse levels, unlike IMP. The results are presented for ResNet-18 Imay-ema
Figure 6:
Table 14: Classification accuracy on VGG networks. SWAMP outperforms all the baselines across varying sparsities, including one-shot and dynamic pruning approaches. Reported values are aver- aged over three random seeds, and the best and second- beooddbbett results are boldfaced and underlined, respectively.
Sparsity
Method
50%
75%
90%
95%
SNIP
92.85 0.19
92.59
0.22
91.30
0.26
90.34 0.25
SynFlow
93.01 + 0.27
93.09 + 0.27
92.74
0.40
91.54 0.21
GraSP
92.20 + 0.06
91.94 +
91.16 01.16+00 0.11
90.68 + 0.08
RigL
93.56 0.18
9..18
92.46
91.45 0.16
CIFAR-10
DST
93.93 0.20
0.09
9.00
93.38 0.06
(VGG-13)
IMP
93.34 + 0.17
0.15
93.58
93.62 + 0.11
IMP SAM
93.73 + 0.12
93.96 +006+00 0.19
94.01 1.01+0.12
93.89 + 0.13
SWAMP (ours)
94.14 + 0.08
94.39 + 0.15
94.40 + 0.16
94.34 + 0.11
SNIP
71.21 0.21
70.41
67.96
66.24 0.40
SynFlow
71.52 + 0.05
7.13
+.22
67.18 0.24
GraSP
69.08 1 0.25
67.26
6..25
63.50
RigL
72.22 0.43
71.38 0.07
0.41
0.05
CIFAR-100 (VGG-16)
DST
72,93 1093:0.11 0.16
72,87 0.13
0.12 0.110
IMP
72.32 + 0.15
72.50 + 0.15
72.64
+.174
IMP SAM
72.11 + 0.45
72.27 + 0.12
72.56
72.50
SWAMP (ours)
73.27 + 0.26
73.54 + 0.36
73.40 + 0.33
73.53 + 0.32
20

---

Published as a conference paper at ICLR 2024
Matching threshold
IMP-1
IMP-2
-
IMP-3
SWAMP (ours)
29.0
29.0
Classinc ation ation error (%)
26.0
26.0
Classifiction
23.0
23.0
20.0
20.0
20
40
60
80
100
3.0
2.5
2.0
1.5
1.0
0.5
0.0
Sparsity (%)
Relative number of parameters
(a) WRN-32-4 on CIFAR-100
Matching threshold
IMP-1
IMP-2
IMP-3
*
SWAMP (ours)
7.20
7.20
Plassiiiation error
Classification error
6.50
5.80
5.10
60
80
2.5
1.5
0.5
Sparsity (%)
Relative number of parameters
(b) VGG-13 CIFAR-10
Matching threshold
IMP-1
IMP-2
IMP-3
SWAMP (ours)
28.8
28.8
Classiication error
27.6
Classificatio error
27.6
26.4
26.4
25.2
25.2
24.0
24.0
20
40
60
80
100
3.0
2.5
2.0
1.5
1.0
0.5
0.0
Sparsity (%)
Relative number of parameters
(c) VGG-16 on CIFAR-100
Figure 7: Supplementary plots for Figure 1. Classification error function of the sparsity (left) and the relative number of parameters (right). Again, SWAMP demonstrates matching performance at extremely sparse levels, unlike IMP.
21

---

Published as a conference paper at ICLR 2024
Negative log-likelihood on CIFAR-10. SWAMP outperforms all the baselines across varying sparsities, including one-shot and dynamic pruning approaches. Reported values are aver- aged over three random seeds, and the best and second-best results are boldfaced and underlined, respectively.
Table 15:
Sparsity
Method
50%
75%
90%
95%
SNIP
0.227 + 0.012
0.242 0.005
0.282
0.003
0.327
0.007
SynFlow
0,219 + 0.007
0.223 0.004
0.280
0.316 0.003
GraSP
0.243 +
0.281 = 0.009
0.308
0.017
0.338
RigL
0.206 + 0.010
0.231 + 0.005
0.261 0.009
0.288 = 0.003
CIFAR-10
DST
0.189 + 0.009
0.214 + 0.008
0.241 0.010
0.251 = 0.008
(WRN-28-2)
IMP
0.198 + 0.001
0.198 + 0.002
0.197 0.014
0.210 0.003
IMP SAM
0.182 + 0.009
0.180 = 0.012
0.187 0.011
0.237 0.053
SWAMP (ours)
0.158 + 0.002
0.156
0.002
0.162 0.002
0.177
0.002
SNIP
0,234 + 0.003
0.244 + 0.008
0.269
0.003
0.291 = 0.008
SynFlow
0,233 + 0.006
0.224 + 0.008
0.230 0.007
0.261
0.002
GraSP
0.252 + 0.002
0.257 + 0.006
0.271 0.001
0.282
0.001
Rigl.
0.216 + 0.002
0.230 + 0.006
0.246 0.007
0.270
0.002
CIFAR-10 (VGG-13)
DST
0,210 + 0.002
0.208 0.004
0.213 0.002
0.220
0.002
IMP
0.220 + 0.006
0.214 0.004
0.211 0.002
0.210
0.004
IMP SAM
0.195 + 0.001
0.191 0.003
0.190 0.003
0.192 0.006
SWAMP (ours)
0.185 + 0.002
0.179 + 0.002
0.177 0.004
0.178 0.001
Table 16: Negative log-likelihood CIFAR-100. SWAMP outperforms all the baselines across varying sparsities, including one-shot and dynamic pruning approaches. Reported values are aver- aged over three random seeds, and the best and second-best results boldfaced and underlined, respectively.
Sparsity
Method
50%
75%
90%
95%
SNIP
1.067 + 0.010
1.064 + 0.022
1.121
0.036
1.17 + 0.007
SynFlow
0.964 + 0.006
0.019
0,994
0.009
1.009 0.041
GraSP
1.112 + 0.043
1.198 + 0.028
1.235 0.014
1.252 0.036
RigL.
1.005 + 0.012
1.047 + 0.016
1.105 0.029
1.125 0.026
CIFAR-100 (WRN-32-4)
DST
1.012 + 0.036
1.029 + 0.028
1.235 0.014
1.252 = 0.036
IMP
0,974 + 0.014
0.973 + 0.007
1.001 0.014
1.010 = 0.011
IMP + SAM
0.924 + 0.031
0.918 + 0.028
0.934 0.024
0,953 0.023
SWAMP (ours)
0.847 + 0.020
0.851
0.016
0.868
0.016
0.893
0.020
SNIP
1.135 + 0.015
1.150 + 0.008
1.216 0.012
1.259 = 0.025
SynFlow
1.123 + 0.005
1.115 - 0.006
1.092 0.010
1.203 0.004
GraSP
1.219 + 0.003
1.242 + 0.013
1.311 0.016
1.351 0.005
Rigl.
1.121 + 0.007
1.142 + 0.007
1.120 0.016
1.238 0.004
CIFAR-100
DST
1.110 + 0.007
1.121 + 0.004
1.138 0.004
1.146 0.014
(VGG-16)
IMP
1.105 + 0.006
1.06 + 0.009
1.090 0.004
1.096 0.002
IMP SAM
1.075 + 0.014
1.067 0.010
1.071 0.004
1.068 0.007
SWAMP (ours)
1.023 + 0.016
1.029 + 0.015
1.033 0.016
1.029 = 0.016
22

---

Published as a conference paper at ICLR 2024
Trace of Hessian Tr(H) evaluated across training data. In most cases, SWAMP with mul- tiple particles exhibits smaller trace value, i.e., finds flatter minima, compared to others. Reported values are averaged over three random seeds, and the best and second- recond-best results boldfaced and underlined, respectively.
Table 17:
Sparsity
Training
20%
50%
75%
90%
SGD
16394.8 +: 3262.1
20335.2 + 4875.3
29337.7
110887.7
25504.9 +: 4370.0
CIFAR-100
SWAMP (N = 1)
2586.9 = 10.7
2669.4 + 21.9
3079.8 + 1508
3159.3 + 147.3
(WRN-32-4)
SWAMP (N =
2457.8
161.8
2556.1
86.9
2896.7
201.0
2968.5
79.4
CIFAR-100
SGD
4776.8
4842.5
5033.1
2336.
5141.7
20.5
(VGG-16)
SWAMP
2072.4
34.5
2189.2
73.9
2472.8
#: 93.1
2570.
+ 88.8
SWAMP (N= =
1995.9
39.6
2201.1
85.5
2341.6
10966
2423.8
+ 48.6
Sparsity 20%
Sparsity 75%
Sparsity 95%
Wh
1.110
1.195
1,280
1.085
1.123
1.160
1.155
1.245
1.335
Figure 8: Visualization of loss surfaces as a function of neural network weights a two- dimensional subspace, spanned by three particles (marked as white circles). The averaged weight (marked by a yellow star) is observed not to be positioned in the flat region of the surface dur- ing the earlier stages of IMP (left; Sparsity 20%). However, as the sparsity increased, the weight averaging technique effectively captures the flat region of the surface. The results are presented for WRN-32-4 on the test split CIFAR-100.
Sparsity 20%
Sparsity 75%
Sparsity 90%
76.0
Accuracy (9)
77.5
Accuracy (89
77.5
Accuracy
77.0
72.0
76.5
76.5
WA
P2
P3
P4
WA
P1
P2
P3
P4
WA
P1
P3
Figure 9: 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 indi- vidual particles (denoted as P1-P4) in the early stages of IMP (left; Sparsity 20%), it achieves high performance at high sparsity levels. The results are presented for WRN-32-4 on the test split of CIFAR-100.
IMP ..G...
SWAMP
28.0
28.0
Classification error (%6
26.5
Classification error
27.5
25.0
27.0
23.5
26.5
22.0
26.0
1
2
5
6
8
9
10
11
12
2
9 Sparsity level (IMP cycle)
5 6
8
10
11
12
Sparsity level (IMP cycle)
Figure 10: Linear connectivity between sparse solutions having different sparsity along with the IMP cycle. The results are presented for WRN-32-4 (left) and VGG-16 on the test split of CIFAR- 100.
23


Reviews:
[{"review_id": "AgLefLDqdr", "reviewer": "Reviewer_MGna", "summary": "SWAMP (Sparse Weight Averaging with Multiple Particles) is a new pruning method that enhances the performance of sparse neural networks by averaging multiple models trained with different stochastic gradients but sharing an identical sparse structure, known as a \"matching ticket.\" This process results in improved generalization due to the creation of flat minima and maintains the important linear connectivity between successive solutions, a key strength of the traditional Iterative Magnitude Pruning (IMP) method. SWAMP has demonstrated its ability to outperform other pruning baselines across various datasets and network structures. The technique's success invites further theoretical investigation into why the convex hull of the weight space of these averaged models forms a beneficial low-loss subspace, which could provide deeper insights into the algorithm's effectiveness.", "questions": "Please refer to Weakness comments", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The motivation behind SWAMP is firmly rooted in robust theoretical frameworks, notably the lottery ticket hypothesis and the concept of linear mode connectivity.\n- The visualization of the loss landscape in Figure 2 provides a clear illustration of the methodology and supports the validation of the claims made.\n- It is clear from the evidence presented in Table 4 that SWAMP is adept at identifying more effective pruning masks.\n- Table 2 and 3 demonstrate that SWAMP achieves superior classification accuracy for a designated target sparsity level.", "weaknesses": "- The study's reliance on demonstrating the process primarily through wide networks such as WRN and VGG-19, which are not the most parameter-efficient architectures, raises questions about the choice of models. An explanation of why these particular, potentially less efficient, models were selected for this research is needed.\n\n- The improvement in accuracy provided by SWAMP over IMP is modest, as shown in Tables 2 and 3, and this increment is even less pronounced for the ResNet model as evidenced in Table 3. This calls for a discussion on the significance of the marginal gains achieved by SWAMP, particularly when benchmarked against other models.\n\n- The feasibility of achieving an optimal sparse structure with SWAMP, especially for pre-trained models which are commonplace, may entail significant computational costs. It is imperative that the authors address the computational overhead, both in terms of space and time complexity, and the practical constraints when applied to large models, including Transformers. A comprehensive discussion on the limitations is warranted, given that IMP—the foundation of SWAMP—may have its own constraints with larger models.\n\n- The applicability of the proposed method to architectures like Transformers needs clarification. In Table 8, the RoBERTa model exhibits a noticeable performance drop even with less than 50% sparsity. The question arises as to whether this decline is attributed to the inherent limitations of IMP, on which SWAMP is based, or if it pertains to the broader challenges of applying pruning techniques to RoBERTa. Additionally, it would be beneficial to understand whether the principles behind SWAMP remain valid for other models, such as GPT-like architectures, and how they compare with alternative pruning strategies for these models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hYmJHT5a9x", "reviewer": "Reviewer_tk5N", "summary": "The authors propose weight averaging of sparse models trained from a checkpoint of a single model, in many ways \"model soups\" for Iterative Maginitude Pruning (IMP). The authors motivate the method for IMP as as model soups in the dense context are, with the loss landscape perspective: we know that LTs lie within the same loss basin, and might expect that weight averaging would find a better generalizing solution. Experiments demonstrate that the approach identifies solutions within a flatter region of the loss basin, and improved generalization over IMP and many other sparse training methods for CIFAR-10/100 and ImageNet models.", "questions": "* While the paper is motivated by comparing the generalization of a SWAMP to an ensemble of two IMP solutions, what is the comparison in generalization when using other sparse training methods, e.g. RiGL, given that these other methods often generalize better than IMP?\n* Is there any reason to believe SWAMP is not a general method that applies to any sparse training method? If so what? If not, why focus on IMP?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "* The paper is overall well-written, with a good organization, clear writing for the most part, and a clear methodology.\n* The experimental analysis is appropriate, using reasonable datasets and models (except VGG), and demonstrates clear knowledge of the sparse training literature appropriate to the methodology.\n* The paper has a clear and well defined motivation: the method is motivated as cheaper than ensembles, much along the same lines of the model soups paper and how it is motivated for dense training.\n* The loss landscape analysis also originally used in model soups is clearly applicable to the sparse domain, especially since much of the linear-mode connectivity methodology comes from the sparse literature to begin with.\n* Hessian Trace analysis also provides some signal that the loss-landscape motivation for weight averaging holds in the sparse realm.", "weaknesses": "* The method comes down to applying the model soup paper to sparse training/IMP. I believe there is sufficient novelty in applying a method only shown on dense training and not necessarily repeatable in the sparse training context, never mind the extensive analysis shown by the authors in this work. Saying that, it's also not the most novel research direction out there compared to many papers.\n* As presented in the main body of the paper, SWAMP is *much* more expensive than most of the compared sparse training methods in e.g. Table 2 at *training time*. This is because IMP with weight rewinding is extraordinarily expensive in practice. However, the authors do demonstrate that the SWAMP methodology applies to other much more efficient sparse training methods in the appendix, notably RiGL, a state-of-the-art sparse training method, and one that is reasonably efficient. I believe the authors should focus their method as being widely applicable to sparse training methods in the main body of the paper, rather than focusing on IMP however - this is especially important given the motivation that SWAMP is better than training an ensemble (which is in fact likely cheaper than SWAMP when using more practical sparse training methods than IMP!).\n* While CIFAR-10/100 results are relatively strong, the ImageNet results (Table 3) are relatively quite weak and not as obviously significant.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HyJB2CPuy6", "reviewer": "Reviewer_YhLg", "summary": "This paper proposes a modification to the Iterative Magnitude Pruning algorithm, SWAMP. The basis of this algorithm is the empirical evidence that different models trained from the same matching tickets can be weight averaged without encountering a loss barrier post certain sparsity levels. SWAMP obtains marginal accuracy improvements with respect to the baselines used.", "questions": "Q1. I would like to understand why is it that the authors choose to average the weights in SWAMP? As demonstrated in Figure 1, there might be individual IMP runs that outperform SWAMP. Why not take the best of multiple pruned weights?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "S1. The manuscript is well written\n\nS2. The method is empirically sound and arguments are well made.\n\nS3. Extensive empirical is provided to justify the merit in this approach.", "weaknesses": "W1. The authors have not empirically justified their choice of using Stochastic Weight Averaging (SWA) as opposed to SGD in the manuscript. It would be important to understand the impact of SWA on the proposed approach by demonstrating two things.\n1. How does IMP perform when it uses SWA as opposed to SGD.\n2. How does SWAMP perform when it uses SGD as opposed to SWA.\n\nW2. Multiple instances of imprecise statements. For example, \"As illustrated in Figure 1, our algorithm achieves superior performance, which is on par with that of an ensemble consisting of two sparse networks.\" It is not clear with respect to what are the authors claiming superior performance? Because in Figure 1, IMP-3 outperforms SWAMP in terms of accuracy.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XaCITnjJeD", "reviewer": "Reviewer_1f6B", "summary": "This paper extends the Iterative Magnitude Pruning (IMP) technique by proposing an approach called SWAMP that trains multiple sparse models called particles in each magnitude pruning iteration of IMP using Stochastic Weighted Averaging (SWA) optimization. The particles in each pruning iteration exhibit the same matching ticket and their diversity is achieved through different batch orders. The trained particle masks are combined in the weighted average fashion to get the single mask of a given pruning iteration. This process of training multiple particles followed by the weighted average of their mask is repeated until desired sparsity or pruning iteration is achieved. The experimentation is conducted on multiple datasets along with different architectures to showcase the effectiveness of the proposed SWAMP model.", "questions": "In Figure 3, why does the proposed technique have a lower performance compared to individual particles in the lower network sparsity?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The authors have done a great job in terms of summarizing their contributions compared to the IMP technique in Section 3.2. \n* The paper is very easy to read and the proposed contribution can be easily understood through Algorithm 1.\n* Extensive experimentation is conducted on multiple tasks (vision and language), multiple datasets, and multiple architectures. \n* A very comprehensive ablation study is conducted to showcase the effectiveness of the proposed components in the paper. For example, Table  5 clearly shows the importance of the SWA optimization along with the weighted average mechanism of the particles to enhance the performance.", "weaknesses": "* In terms of methodology, the proposed technique provides an empirically guided straightforward extension over the IMP technique. The proposed SWAMP therefore has a trivial contribution and therefore lacks novelty.\n* In terms of experimentation, the performance gain over other techniques seems to be marginal and  reduces the significance of their proposed methodology. \n* In Figure 3, for relatively lower sparsity (e.g., sparsity of 20%), the proposed Weighted Average (WA) technique seems to underperform the individual particle performance. Does this mean, the proposed technique  harm the performance on the lower sparsity? The authors may need to provide more extensive justification to explain this phenomenon.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "o3PsZEWatp", "reviewer": "Reviewer_Eb3S", "summary": "The paper presents SWAMP, a novel method using the average of multiple particles' stochastic weight averaging (SWA) to achieve improved model performance. The authors have tested SWAMP’s effectiveness across various tasks, including vision models (CNNs) and language models (RoBerTa finetuning), providing a comprehensive evaluation.", "questions": "1. The results in Figure 4 show that interpolated weights yield even lower errors compared to IMP weights. Could the authors provide a detailed explanation or hypothesis as to why this is the case?\n2. Appendix B states, \"The learning rate for this phase (SWA phase) is set to a constant value of 0.05.\" Does this imply that the minimum learning rate is set at 0.05 for SWAMP, and for other baselines such as IMP?\n3. Do IMP and SWAMP use the same epoch T_0 to rewind weights?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written, and easy-to-follow.\n2. The authors have conducted extensive experiments, covering vision tasks and language model fine-tuning. Additional studies such as mask analysis and efficient implementation SWAMP+ are also provided.", "weaknesses": "While the authors address the computational efficiency of SWAMP+ in section 4.3, stating that it can utilize a single particle for the first few iterations, this claim seem to work because the networks are already very sparse. My concern lies in the computational cost of SWAMP+ at lower sparsity levels. Furthermore, did the authors test SWAMP+ on ImageNet?\n\nSee also my questions.\n\nI am willing to adjust my rating if my questions are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)