SAFE: Finding Sparse and Flat Minima to Improve Pruning
Authors: Dongyeop Lee
Kwanhee Lee
Jinseok Chung
Namhoon Lee
Venue: ICML 2025 spotlightposter (2025)
Keywords: Pruning
Constrained optimization
Sharpness minimization
Primary area: deep_learning->algorithms
TL;DR: We propose SAFE, an optimization-based pruning method that improves generalization of sparse models by inducing flatness.

Abstract: Sparsifying neural networks often suffers from seemingly inevitable performance degradation, and it remains challenging to restore the original performance despite much recent progress.
Motivated by recent studies in robust optimization, we aim to tackle this problem by finding subnetworks that are both sparse and flat at the same time.
Specifically, we formulate pruning as a sparsity-constrained optimization problem where flatness is encouraged as an objective.
We solve it explicitly via an augmented Lagrange dual approach and extend it further by proposing a generalized projection operation, resulting in novel pruning methods called SAFE and its extension, SAFE$^+$.
Extensive evaluations on standard image classification and language modeling tasks reveal that SAFE consistently yields sparse networks with improved generalization performance, which compares competitively to well-established baselines.
In addition, SAFE demonstrates resilience to noisy data, making it well-suited for real-world conditions.

SAFE: Finding Sparse and Flat Minima to Improve Pruning
Dongyeop Lee' Kwanhee Lee'
Jinseok Chung' Namhoon Lee'
Abstract
processing (Blumensath & Davies, 2009), and statistics (Tib- shirani, 1996; Beck & Teboulle, 2009). Notably, the recent success of highly overparameterized deep neural networks in achieving human- computer vision and natural lan- guage processing tasks has drastically scaled these models to proportions never seen before, which has spurred con- siderable research on sparsifying neural networks (Hoefler et al., 2021). Since the pioneering work of Han et al. (2015), many works have proposed to remove redundant parameters by various tactics ranging from those discussed in the survey work of Hoefler et al. (2021) to their applications to large foundation models (Kwon et al., 2022; Frantar & Alistarh, 2023; Sun et al., 2024). However, it is witnessed that ex- cessive pruning usually results performance decline due to the reduced capacity, potentially impairing deep learning models from handling tasks of high complexity. Is there way to restore this performance loss?
Sparsifying neural networks often suffers from seemingly inevitable performance degradation. and it remains challenging to restore the origi- nal performance despite much recent progress. Motivated by recent studies in robust optimiza- tion, we aim to tackle this problem by finding subnetworks that are both sparse and flat at the same time. Specifically, we formulate pruning as a aparsity-constrainee
optimization problem where flatness is encouraged as an objective. We solve it explicitly via an augmented Lagrange dual approach and extend it further by proposing gen- eralized projection operation, resulting in novel pruning methods called SAFF and its extension, SAFE*. Extensive evaluations on standard im- age classification and language modeling tasks reveal that SAFE consistently yields sparse net- works with improved generalization performance, which compares competitively to well-established baselines. In addition. SAFE demonstrates re- silience to noisy data, making it well-suited for real-world conditions.
We attend to recent studies that have demonstrated that the generalization performance of a model is closely linked to the flatness of the solution landscape (Keskar et al., 2017: Jiang et al., 2020a). This insight has led the development of techniques such as sharpness-awes (Foret et al., 2021), which explicitly regularizes the sharp- ness of the solution during the optimization process. SAM has been shown to deliver exceptional performance across various domains (Chen et al., 2022; Bahri et al., 2022) and has also demonstrated robustness to label noise (Back et al., 2024). light of the SAM's success, there has been grow- ing interest applying these techniques to model pruning. Research by Na et al. (2022) has explored how sharpness- aware training affects model compressibility. while Peste et al. (2022) and Bair et al. (2023) have examined different strategies to alter this process for pruning, with the goal of enhancing performance and robustness. Nevertheless, the exploration of sharpness-aware techniques within the context of sparsification is still an early stage. We believe there is considerable potential to integrate these approaches more effectively into the sparsification process, which could lead to significant advancements in the development of effi- cient and robust deep learning models.
minimization (SAM)
1. Introduction
Over the past decades, the emergence of computers and the accumulation of digital data have made machine learning a viable tool for everyday applications. However, modern machine learning systems have also grown in complexity with the rapid advances in hardware and databases, demand- ing significant computational and memory resources. This has led to a surge of interest in strategies to reduce these substantial computational costs.
One major approach is sparsification, which aims to find solutions with mostly zero entries. This has been studied for many years in various large-scale applications in machine learning (LeCun et al., 1989; Hassibi & Stork, 1992), signal
POSTECH, South Korea. Correspondence to: Dongyeop Lee <dyles23 @ postech.ac.kr>.
In this work, we aim to achieve exactly that, by proposing to frame as sharpness- aware pparsity-constraime opti- mization problem to simultancously consider both sharpness and sparsity during the training process. To tackle this, we
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
use augmented Lagrangian dual-based approach well established in the optimization literature with convergence guarantees, and propose a new optimization-basel pruning method called SAFE. We evaluate SAFE across standard benchmark tasks in image classification and large language model post-train pest-training pruning, and compare with existing al- ternatives to demonstrate that (i) it induces flatness on the sparse solutions, (ii) yielding performance improvements in the resulting sparse models, and (iii) its robustness to label noise, common image noise, and adversarial noise, and its (iv) effectiveness compared to similar sharpness- innimizaation-inspir pruning techniques. These aspects support the effectiveness and generality of our method, and we conclude by discussing current limitations and potential ideas for future work.
istarh, 2023; Sun al., 2024; Kwon et al., 2022) is ideal for enhancing inference efficiency in pre-trained models, particularly when the training process can be too costly due to large-scale data or complex models such as large language models (ILMM. A widely adopted strategy in post- training pruning is layer-wise reconstruction error mini- mization, which ensures that the pruned model maintains accuracy by preserving layer-wise output approximations (Frantar & Alistarh, 2023; Sun et al., 2024; Meng al., 2024). This approach enables efficient pruning of large models by solving smaller subproblems independently for each layer, reducing computational overhead while preserv- ing the signal in model representations. Optimization-base techniques have been proposed to refine this idea further, improving sparsity while minimizing performance degra- dation. Additionally, extensions of these methods explore structured sparsity patterns, such block-wise sparsity, to enhance hardware efficiency while maintaining competi- tive accuracy. Also, recent studies hint at the possibility of finding an initial random sparse network that can be trained to achieve comparable performance to dense net- works, although this generally involves several rounds of expensive training (Frankle & Carbin, 2019). Among var- ious approaches, this work primarily focuses pruning during training (Peste et al., 2021; Zhou et al., 2021; Kusu- pati et al., 2020; Lin et al., 2020; Sanh et al., 2020; Evci et al., 2020), which is known for achieving best results by guiding the model toward desired sparsity during training (Hoefler et al., 2021), and can be ideal for moderately sized models with adequate computational resources for training. Despite this, preserving the original dense performance re- mains challenging due to the complexity of tasks handled by deep learning models, often leading to the use of heuristics to manage these difficulties.
2. Background
2.1. Sparsity
Throughout the years, various techniques to induce spar- sity in machine learning systems have been developed to simplify models for efficiency, interpretability, and gener- alization. This is usually framed as solving the following optimization problem with sparsity constraint:
min ((x), || Sd
(1)
where f is the objective we wish to minimize, x is the opti- mization variable, ||æ||lo denotes the Lo-norm, which counts the non-zero entries within æ, and d is the number of pa- rameters we wish to preserve. The goal is to find a solution with desired sparsity that minimizes f. Exactly solving this optimization problem is challenging due the combi- natorial nature of the Lonnom,, as it requires an exhaustive search over all possible configurations of zeros within x. Consequently, several approaches have been proposed, such as relaxing the Lo norm as in LASSO (Tibshirani, 1996), or employing advanced optimization techniques, including proximal methods like FISTA (Beck& & Teboulle, 2009) and iterative hard thresholding (Blumensath& & Davies, 2009). Additionally, strategies like optimal brain damage and sur- geon (LeCun et al., 1989; Hassibi & Stork, 1992) have been explored to sparsify multi-layer perceptrons through second-order approximations of the objective function.
2.2. Flat Minima
In-depth empirical analyses into the optimization properties of deep neural network training, especially with regards to mini-batch training, have revealed a surprising correla- tion between well-gennraliiing solutions and their flatness (Keskar et al., 2017; Jiang et al., 2020b). This finding has prompted numerous studies aiming understand the pre- cise nature of this relationship (Andriushchenko et al., 2023; Neyshabur et al., 2017: Zhou et al., 2020), and its impact on neural network pruning, as explored by Lee et al. (2021), who link the challenge of training highly sparse networks to sharper loss landscapes based on an analysis of scaling properties under varying sizes of mini-batches and classical optimization theory.
The recent success of increasingly large deep neural net- works has further accelerated this trend, spurring the de- velopment of various methods to sparsify neural networks at different stages of training, each offering distinct advan- tages depending on the scenario (Hoefler et al., 2021). For instance, pruning before training (Lee et al., 2019; Tanaka et al., 2020; Wang et al., 2020) is advantageous for improv- ing computational efficiency during training by enabling sparse training, while post-training pruning (Frantar& & Al-
This also motivated researchers to develop various tech- niques to explicitly induce flat minima during training (Foret et al., 2021; Izmailov et al., 2018; Orvieto et al., 2022: Chaudhari et al., 2017). Among them, sharpness-awer aware aware min-
2

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
imization (SAM) by Foret et al. (2021) aims tackle this by solving the following min-max optimization problem:
where / is the objective function minimize, d is the num- her of parameters to preserve, and p is the radius of the perturbation e. Thus, the goal is to find a sparse solution x* that minimizes the objective function in the whole c- neighborhood, i.e,, seek flat minima.
min max f(x ee) F
where we minimize the objective function over the entire t- neighborhood with radius i.e., seek flat minima. Solving the inner maximization problem for the first-order Taylor approximation gives the following update rule for SAM:
3.2. Augmented Lagrangian Based Approach
A standard approach for solving such constrained optimiza- tion problems is to employ Lagrangian duality or projected gradient descent. However, the discrete nature of Lo -ormorm makes Lagrangian duality infeasible, while projected gra- dient descent, despite its computational efficiency for Lu constraints, can struggle with highly non-convex objectives in neural network optimization. To leverage the smooth optimization of Lagrangian and the efficiency of projection, we leverage augmented Lagrangian as described below.
Vf((e) (v))()))
T++1 It n Tt +
This has been shown to be effective in improving general- ization performance and robustness across various domains (Chen et al., 2022; Bahri et al., 2022; Baek et al., 2024).
The success of sharpness minimization techniques has natu- rally led to exploring their implications for neural network pruning. Na et al. (2022) studied whether flatter loss land- scape can be more compressible by employing SAM during iterative magnitude pruning for fine-tuning BERT models. Shin al. (2025) observed that the generalization benefits of SAM can be leveraged with sparsity for overparameter- ized neural networks. Inspired by SAM, Peste et al. (2022) minimization (CrAM) min- imize the loss increase induced by perturbation from com- pression, which aims to induce robustness to post-training one-shot pruning of any sparsity. Bair et al. (2023) sug- gested performing additional sharpness-aware training to pre-trained models, with larger perturbations given to co- ordinates of low importance score (i.e., parameters to be pruned) incur less loss increase when pruning.
To achieve this, we first employ variable splitting, a widely used trick, usually to separately deal with objective mini- mization and constraint satisfaction (Boyd et al., 2011). Pre- cisely, instead of directly imposing the sparsity constraint on variable x, we first split it into variables and as follows:
min max f(x + <) + 4 [Hbbsd())
s.t. x = 2,
2,2 lellasS
proposed compression-awee
where (al)) is an indicator function for the sparsity constraint:
O if || S d
I [blosa(z) ==
Do else.
We then slightly alter the Lagrangian by adding penalty term x/2//==21 with penalty parameter X, which preserves equivalence to the original problem while also acting as a proximal term for the projection step. This alteration is a form of augmented Lagrangian, which we apply to form the Lagrangian dual problem of the following:
While these efforts represent initial attempts to verify the ef- fectiveness of sharpness minimization or its loosely inspired variants in enhancing pruning, we believe that sharpness minimization can be more effectively integrated into the sparsification process. Thus, in this work, we aim to weave this charpness-minimizztion minimization objective with the sparsifica- tion process to enhance the quality of the sparsified network via principled optimization-based approaches that are well established in the literature.
max, min [(x,, 2, u) == max f(x + + + +I [11l2sa())) u 2,5
|ellasS
Iull + " = + u113)
where u is a scaled dual variable for the equality constraint scaled by 1/A. Here, the projection can be computed ef- ficiently via hard thresholding operation (Blumensath & Davies, 2009), which sets all entries of but the d ele- ments with the largest magnitudes to zero. Applying dual ascent leaves us with the following x, z z-minimization and u- u-maximization:
3. Method
In this section, we present a detailed derivation of our flatness-inducing sparsification algorithm Sparsification via ADMM with Flatness Enforcement or SAFE
Fk+1, 2b+1 = argmin max f(x + c) + [1))5a(())
3.1. Problem Formulation
2,2
|
We begin by proposing the following min-max problem with sparsity constraint:
max optimization
+ 2
|| 2 + ux 11
min max f(x + c).
(3)
us+1 = argmax
|| - =b+1 + uu
-2 11
11
3

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
To minimize each x and separately with iterative first-order optimization and exact projection operation respectively, we compute and in an alternating manner which gives the following iteration:
naturally integrate these into the projection operation, de- sign generalized distance metric that introduces positive- definite diagonal matrix P that provides framework to incorporate these advanced saliencies of the form P p//3 in a principled manner:
[2,]] [
> 21100 =k + ux Ja+1 = argmin max [1 + c) + +-2111
Icclasp
Fk+1 - projff + us)
2++1 projj (#++1 + ux))
== arg min -11== -11=- (IK+1 + uu)ll
0k+1 uk + 1k++1 2k+1,
Iellogd
arg min 22=- -=-11=++ + ux)) P(z (I4+1 + uz))
where projj. ld projection operation onto the spar- sity constraint (i.e., the hard thresholding operator), and u-maximization is solved through applying single step of gradient ascent with step size of À on y. which is to ensure that the iterate stays within feasibility once reached.
I
Geometrically, this can be understood as modifying the underlying distance metric to better represent the local geo- metric structure (e.g., the Hessian) of the original objective function. We call this SAFE+, where we leverage various saliency scores within the projection step, which we de- scribe in detail below.
3.3. z-minimization
For the z-minimization step, we first approximately solve the e-maximization via first-order approximation of f:
We lay out some notable examples of advanced saliency scores and the corresponding P. The simplest case is P I, where the projection reduces thresholding as it corresponds to the Euclidean norm, yield- ing the original SAFE. Taking this further, setting it as the diagonal Hessian, i.e., P = diag( (V2 (xx),, corre- sponds to Optimal Brain Damage (LeCun et al., 1989), a second-order pruning method that aims to remove pa- rameters with minimal impact on the loss function. Also, using P = diag( V fxx ((:)) aligns with the first-order pruning method SNIP (Lee et al., 2019), which removes pa- rameters based on gradient sensitivity. Furthermore, Wanda (Sun et al., 2024), a layer-wise pruning method for language models, corresponds to taking P = diag( (A T A) where A E Ndd is an activation with batch size and feature dimension from a particular layer to prune. This corre- sponds to the diagonal Hessian of the reconstruction error for a single linear layer (Sun al., 2024).
ff(z)
c (z) x argmax f(x) + Vf(=)))))))))))))))))))(())))V
standard hard
which we apply back to the objective
A 211--2 zx 1
Th+1 = argmin f (x + (x)) +
We solve this using gradient descent, where we remove higher-order terms in the gradient as in Foret et al. (2021),
v (x (1x+E" (x)) + 1 2x + uml
=(I+ (z) (x) + +A(x 2k ux)
Vx(x) -V x ( + +X(x 2k + uk):
(4)
This generalized projection allows SAFE to integrate di- verse sparsification strategies all within its constrained opti- mization framework, thus enhancing both effectiveness and robustness in model pruning. Our empirical evaluation in Section 4.3 demonstrates its effectiveness in large language model pruning, though the methodology is not confined to this domain and is widely applicable.
thus leading to the following z z-minimization steps:
V((a())
Kk I
((+1) (c)
(e)
V
(1)
(V 12)
+ Ma (t)
zk + us))
(5)
where f, (1) are the current step of z-minimmzation and its step-size, respectively.
3.5. Final Algorithm: SAFE and SAFE
Our final algorithm summarized in Algorithm We pro- vide intuitive description of how our algorithm performs sparsification. Every few steps of æ-minimization, SAFE observes where the closest point on the sparsity constraint from the current r is and registers it on z. While performing flatness fatness-inducing inducing minimization of the objective function on x, it penalizes the x iterate to slightly move closer to 2, the
3.4. Extension to Generalized Projection
While the Euclidean projection onto an Lo constraint natu- rally yields magnitude- based sparsification, this often yields subpar performance in practice compared to more advanced saliency scores that account for the objective function. To

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
Algorithm/ 1 SAFE and SAFE* algorithms
Definition 3.4. (6-stationary point) We say a point E is 6-stationary point of the optimization problem (1) if x arg minaeA || (
Require: Target parameter count d, total train iteration T, dual-update interval K, learning rate n7t), perturbation radius p. penalty parameter X, importance matrix P.
i.e. the point 3 cannot be locally improved using projected gradient descent with step-size 8-1. With this definition, we ultimately demonstrate that SAFE converges to this 6- stationary point, which is a necessary condition for the optimal solution to problem (1).
1: Initialize (:
2: u= =0
3: for t in T'do do
4: if' mod K = 0 then
5: SAFE then
6:
= ppro H Hl (t++)) + +u)
We first provide the central lemma on the convergence of our sharpness minimizing I iterates:
7:
else if SAFE then
8:
2 = -roj P
Illosaa (llo + u.)
Lemma 3.5. (Convergence of z-minimization) Suppose
9:
end if
that Assumptions 3.1 and 3.2 hold and let {x be generated by Equation (5) in Algorithm with step-size x/(0) and perturbation radius p(0) satisfying Σi= 17())
10:
u u + (((+1)
11:
end if
(((+1/2) = x(1) [1 (z(0) +p. Viia())
12:
00, - ( < xo, lir sup, ,(0) < 1/B. Let L(x)
13:
2((+1)= (1+) -1(²) ((2(1) -2+22
f(z) + 2 || + ull2 and assume that inf ((x(t))
14:
end for
-00. Then vi(x())) - 0.
15: return proj (oad)il)
The detailed proof is provided in Appendix A.1. This shows that running Equation (5) produces a sequence that con- verges to the stationary point of the augmented Lagrangian L with respect xr
latest estimate of the sparse solution. This gradually moves the dynamics of towards sparsity without incurring sud- den change of loss, all while performing flatness induction, yielding sparse and flat minima.
We use this to derive the convergence of SAFE as the fol- lowing corollary:
In practice, particularly for image classification, we intro- duce scheduling to the penalty parameter from zero to the target value in a cosine curve in order to apply less restriction in the initial phases of training, which slightly improves performance. Details of the ablation study on this scheduling strategy can be found in Appendix F.3.
Corollary 3.6. (Convergence of SAFE Suppose that As- sumptions 3. 1-3.3 hold. Assume further that 5 is chosen large enough so that 8-2 (8 p)/2 2 < O. Let (I, ž, ū) be a limit point of SAFE algorithm. Then &t is a 6-stationary point of the optimization problem (1).
This demonstrates that SAFE converges to the stationary point of the connsttyeon optimization problem (1), We note that, while the technical contributions of our anal- ysis might be considered modest, SAFE is built the- oretically rigorous foundation, unlike many other pruning techniques that often rely primarily on ad-hoc intuitions.
3.6. Convergence Analysis
Here present convergence analysis of SAFE. Precisely, we first prove that our proposed iterative sharpness mini- mization in the x-update converges, then build the rest of the proof upon a well-studied result of ADMM (Boyd et al., 2011; Wang et al., 2019; Huang et al., 2021).
4. Experiments
We start with standard assumptions used in the literature: Assumption 3.1. (Lower bounded on constraint) The func- tion f is lower bounded on A. That is, there exists constant fuin == A f fa) and fmin -0.
In this section, we demonstrate that SAFE converges to sparse and flat solutions, leading performance improve- ments over baselines in both image classification and lan- guage modeling tasks. We also show that SAFE is robust to noisy label training and corruptions during inference. The codes to reproduce the results are provided in JAX and PyTorch, with further details provided in Appendix B.5.
Assumption 3.2. (3-smoothness) The function ff is differ- entiable, and its gradient is 3-smooth. That is, [[V[(x)
Vf(y)] < A||== yll
Assumption 3.3. (μ-weak convexity) There exists a con- stant 2 0 such that f is p-weakly convex. i..., f(r) + #
4.1. Convergence to Sparse and Flat Solutions
2 ||x|| 2 is convex.
We first show that SAFE successfully guides training towards sparse and flat solutions compared to naive baselines on a simple neural network model. Specifically, we analyze the
We also define the following notion of stationarity for the optimization problem (1) from Huang et al. (2021):
5

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
SGD weight distribution
SAFE weight distribution
Sharpness: 0.2
Sharpness: 0.09
10000
10000
8000
8000
Count
6000
6000
0.0
0.0
4000
4000
-0.5
-0.5
2000
2000
0.02 Weight value
0.02 Weight value
3.5
0.5
(a) Dense training
SAFE
ADMM
(d) SAFE
Figure 1: (a-b) Weight distributions of densely-trained model and model trained with SAFE, and (c-d) loss landscape and maximum Hessian eigenvalue of minima found by ADMM and SAFE. SAFE yields sparse and flat solutions.
VGG-19 / CIFAR-10
ResNet-20/ CIFAR-10
VGG-19 / CIFAR-100
ResNet:32 CIFAR-100
93
75
75
Accuracy (%6
93
89
68
Accuracy (96)
92
Accuracy
85
61
49
SAFE
SAFE
SAFE
ADMM
81
54
ADHM
36
ADMM
PEW
77
ADMMM
47
PBW
23
PWW
MLPrure
GMP
MEProne
MLPrune
90.0
95.0
96.0
99.0
99.5
90.0
95.0
96.0
99.0
99.5
90.0
95.0
98.0
99.0
99.5
90.0
95.0
96.0
99.0
99.5
Sparsity (%)
Sparsity (9)
Sparsity (%)
Sparsity (%)
Figure 2: Validation accuracy std) of VGG-19 19 and ResNet-20/32 models on CIFAR-10/100 pruned across different sparsity levels and methods. SAFE consistently achieves superior performance across broad range of sparsity levels.
weight distributions of models trained with standard dense training and SAFE to assess its sparsification capability. We also measure sharpnesses of SAFE and compare it to that of ADMM (Zhang al., 2018) a non-sharpness-miniiizzng baseline, by computing maximum Hessian eigenvalues and visualizing loss landscapes. The results are presented in Fig- ure 1. Our findings indicate that SAFE effectively induces sparsity while simultaneously enforcing flatness, as evi- denced by the concentration of weights near zero in contrast dense training and a wider minimum with a lower Hessian eigenvalue compared to ADMM. This result demonstrates the effectiveness of SAFE in tackling the sharpness-aware sparsity-constrainedd optimization problem (3). Further ex- perimental details are provided in Appendix B.2.
2016) using a range of representative pruning methods, in- cluding PBW (Han et al., 2015), GMP (Kurtic Alistarh, 2022; Zhu & Gupta, 2017), LTH (Liu al., 2024; Frankle Carbin, 2019), ADMM (Zhang et al., 2018), and MI.Prune (Zeng & Urtasun, 2018), some of which achieves competi- tive to state-of-th weffthe-ard art performance in image classification (Hoefler et al., 2021). We mostly use standard values for common hyperparameters such as training epochs, learning rate, and weight decay (Zhou et al., 2021) and tune the hy- perparameters unique SAFE, which we report in detail in Appendix B. Notably, we do perform additional training after pruning, and instead perform cost-efficient statistical correction on the batch-norm layers with only few forward passes (batch-nom norm tuning or BNT), which is a common practice in the literature (Hubara et al., 2021; istarh, 2022; Peste et al., 2022). We refer to Appendix B.3 for full experimental details. The final validation accuracies are provided in Figure 2 and Table 7 of Appendix C.
; Frantar & Al-
4.2. Evaluations on Image Classification
In this section, we show that SAFE can achieve outstand- ing generalization performance among various methods for CIFAR-10/100 image classification tasks (Krizhevsky et al., 2009). Specifically, we evaluate pruning performance on VGG-19 19 (Simonyan, 2014) and ResNet-20/32 (He al.,
Our findings show that across most configurations and spar- sity levels, SAFE generally outperforms all baselines. Also, SAFE exhibits greater robustness under extreme sparsity compared to non-sharpness-minin minimized approaches (e.g., 99.5%). Crucially, SAFE achieves these results without requiring costly retraining, whereas PBW and LTH depend
Following the convention of Wang et al. (2020); Zhou et al. (2021), we double the number of channels from standard ResNet models.
6

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
Table I: Perplexities (tea of LLaMa models pruned to various sparsity levels using different methods. SAFE achieves competitive performance, while SAFE outperforms baselines across all settings.
LLaMa-2
LLaMa-3
7B
13B
8B
Sparsity
Method
Wi i /c4
Wi text /C4
Wi k text /C4
Dense
5.47
7.26
4.88
6.72
6,23
9.53
Magnitude
16.03
21.33
6,82
9.37
134.20
273.3
SparseGPT
10.03
1.203
6.0610.03
8.20
9.361..11
13.96 10.00
Wanda
6.92
/
9.23 10,00
5.98 10.01
8.28 10.01
9.71 10.00
/ 14.88 10.04
50%
ALPS
+o.0
8.98
5.96
8.09 +0.04
9.05
13.40
SAFE
6.78+0..1
8.93
5.76+0.02
7..0
9.59
14.60
SAFE
6.56+0.01
+0.00
+0.0
7.74
8.62
13.26
Magnitude
1864
2043
11.81
14.62
5335
7438
SparseGPT
10.19:0.08
12.86=20.05
8.31
10.85
15.4640.40
21.25
Wanda
10.7540.07
13.87
8.43
11.55
22.06
32.228
60%
ALPS
9.550
11.24
0.54
/
+0.06
14.03
/ 18.72+0.15
SAFE
9.20++.04
11,51 +0.04
7.18+0.03
9.59
15.90
22.26++0.16
SAFE
8.30+0.06
10.59+0.00
6.78++0.04
9.02
12.18
17.30+0.02
Magnitude
15.91
31.61
7.32
9.96
212.5
336.3
SparseGPT
8.42+0.05
/ 10.73
7.02+0.06
9.33+0.04
12.1640.20
/ 17.36
Wanda
8.64:0.03
11.35
7.01
9.70+0.08
13.8440.04
21.11420.0
4:8
ALPS
8.11
10.21
6.81
9.33
11.38
/ 16.1020.10
SAFE
8.2110.01
10.61
6.6010.02
8.95
12.15
17.9010.11
SAFE
7.59
9.88
6.37
8..1
10.51
15.67 10.02
Magnitude
37.77
74.70
8.88
11.72
792.8
/ 2245
SparseGPT
11.00
/ 13,5440.05
8.78
/ 11.2644..1
15.87
/ 22,45
Wanda
12.17:0:02
/ 15,60
9.01
/ 12.4000.o0
23.03
/ /3..9140.31
2:4
ALPS
9.99
/ 12.04410.04
8.1610.17
/ 100.35
14.53
/ 19.744
SAFE
10.5310.13
/ 13.20
7.64
/ 10.10.0..1
17.4910.27
/ 24.455
SAFE
8.961
11.31.1..00
10.03
7.2010.04
9.52
13.39 10.23
19.03
on multiple rounds of retraining. These results highlight the effectiveness of SAFE in preserving model accuracy during sparsification, especially under aggressive pruning scenarios.
pruning methods such as SparseGPT (Frantar Alistarh, 2023), Wanda (Sun et al., 2024), ALPS (ADMM-based) (Meng et al., 2024), well as magnitude pruning (Han et al., 2015) and evaluate the perplexity Wikitext2 (Mer- ity et al., 2022) and C4 validation sets. We follow the common practice of randomly selecting 128 samples from the C4 training dataset (Raffel et al., 2020). We refer to Ap- pendix B.4 for experimental details. The results are reported in Table
4.3. Evaluation Large Language Model Pruning
Here we scale our evaluations to modern large-scale settings and demonstrate that SAFE also delivers competitive perfor- mance against state-of-the-art LLM post-training pruning techniques.
We find that SAFE performs competitively state-of-the-art methods, while SAFE* surpasses them across all models and sparsity settings. Considering that these baselines are tailored specifically to pruning LLMs, this demonstrates the flexibility of SAFE in adapting to different scenarios. More- over, our method is more efficient than ALPS, which re- quires x2.54 more runtime than SAFE (see Appendix E.2).
For this purpose, we adapt SAFE and SAFE to sequen- tially optimize the reconstruction error minimization (REM) objective for each transformer block (Shin al., 2024), similarly other L.LM pruning techniques. For SAFET, incorporate Wanda projection z-step, which identifies superior subnetworks compared to naive magnitude-basee based pruning in LLMs without compromising efficiency (Sun et al., 2024).
4.4. Robustness to Noisy Data
Noisy data pose significant challenges real-world see- narios. To address this, we evaluate SAFE on three repre- sentative challenges: incorrect training labels (Song et al., 2022), inference-time input corruption that arises naturally (Hendrycks & Dietterich, 2019), and corruptions that are de-
With this, we prune one of the most widely adopted language model family, LLaMA2-7b/13b (Touvron et al., 2023) and the more recent LLaMA3-8b (Meta, 2024), to 50% and 60% sparsities, as well as structured 4:8 and 2:4 sparsities. We compare SAFE with state-of-the-art LLM post-training

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
Table 2: Noisy label training. Validation accuracy is mea- sured for sparse models trained with ADMM and SAFE under various levels of label noise and sparsity. SAFR is much more robust to label noise.
Table 3: Evaluation on corrupted data. CIFAR-10C is used for common corruptions, and 1x0 and l2 PGD attacks are used to generate adversarial corruption on the validation set of CIFAR-10 10. SAFE improves robustness over naturally and adversarially corrupted images.
Noise ratio
Spursity
Method
25%
50%
75%
Common commuttion (avg.)
Adversarial
70%
ADMM
77.00
59.18
32.62.0.89
Sparsity
Method
intensity=3
intensity=5
PPG
[g-PGD
SAFE
90.58
86.51 #0.36
67.01
ADMM
70.06=a.03
52.01 20.0s
=9.01
49. =9.71=1100
80%
ADMM
76.18
62. 62.67
32.86
90%
73.98.0.mm
55.11 20.27
56.43...08
56.36.1.11
SAFE
91.25 12
86.55 20.07
66.4920.56
ADMM
68.87
50.56
49.84_1.78
49, 49.6821177
90%
ADMM
79.40
66.64go.1a
36.8420.04
95%
SAFE
72.92
54.86
51.40 80
51.36
SAFE
90.68 20:21
86,49 20.00
64.72 20.61
ADMM
65.46.0.2
48.65 10:04
43.33
43.42
95%
ADMM
77.71 20.52
67.10
39.6 19.68
98%
SAFE
68.20.0.47
49.96
23.34_0..0
43.41
SAFE
89,86 20.11
85.1820.15
64.25 20:36
ADMM
29.21
43,81
30.29
30.32..0.e
99%
SAFE
66.02 56
49.34
43.70 -1:28
32.70 78
liberately introduced by adversaries (Szegedy et al., 2014).
99.5%
ADMM
55.72_0.11
56.58_0.30
41.55
42.27 20.63
20.78
23.25
29.48 08
23.25
29.45 74
23,25 13,25.1..0 1.885
Training on noisy labels To assess the robustness of SAFE against label noise, we randomly corrupt (25%, 50%, 75%) of labels in CIFAR-I0 and use it to train ResNet-20 with both SAFE and ADMM. The same hyperparameters from Section 4.2 are used in all experiments. As presented in Table 2, we observe that SAFE consistently outperforms ADMM across all levels of label noise and sparsity, with accuracy improvements ranging from +10% to +30%.
common corruptions, we use CIFAR-IOC (Hendrycks & Dietterich, 2019), a benchmark consisting of CIFAR-10 test images corrupted with 19 types of real-world noise (e.g.. fog, snow, etc.) and distortions (e.g.. jpeg compression, contrast, etc.) at five levels of intensity. We average the performance across all corruption types for intensity levels 3 and 5. For adversarial noise, we follow Zhang et al. (2024) and use a 10-step Projected Gradient Descent (PGD) attack on each sparse model under lx and 12 norm with bound c 1/255 and 3/ 3/255 and step size a e/4, respectively. As shown in Table 3, SAFE enhances robustness to both common and adversarial image corruptions, aligning with previous work on sharpness minimization (Zhang et al., 2024; Wei et al., 2023).
Additionally, we observe that ADMM relies heavily on spar- sity to mitigate label noise, exhibiting an overall trend of increasing accuracy with higher sparsity levels. This de- pendence is further reflected in the sparse double descent pattern reported at the 25% noise ratio (He et al., 2022), where accuracy initially declines up to 80% sparsity, rises sharply to 79% at 90% sparsity, and then drops again to 77% at 95% sparsity. This contributes to the overall decreasing performance gap between ADMM and SAFE, which may be dimin-
4.5. Comparison with Other SAM- pruner
interpreted as the benefit of tharpness-minimizen ishing with fewer parameters (Shin et al., 2025). However, more crucially, the overall under-pen pedrr-performance of ADMM indicates that sparsity alone is poor remedy for label noise, highlighting the effectiveness of SAFE-particnn particularly its flat- ness enforcement-in reducing the impact of label noise. This aligns well with previous observations that sharpness- minimization can enhance robustness toward label noise (Back et al., 2024). Also, the lack of double descent in SAFF suggests that its effectiveness may be attributed to sharpness minimization functioning as effective regular- izer, as supported by the claims of Nakkiran et al. (2021) that 'optimal' regularization can mitigate the double descent phenomenon.
To strengthen the comparison with closely related baselines, we compare SAFE to two pruning baselines inspired by SAM IMM-SMPPSAM (Na et al., 2022) and CrAM (Peste et al., 2022) on ResNet-20/CIFARR 10 across multiple sparsity levels.
IMP+SAM (Na et al., 2022) involves applying SAM dur- ing iterative magnitude pruning (Liu et al., 2024; Frankle & Carbin, 2019). While it was initially introduced for lan- guage model finetuning, we adapt this method to image clas- sification and use the same training epochs and sharpness- minimization hyperparameter search range as SAFE to en- sure fair comparison. Pruning is performed every 10 epochs with sparsity increasing either linearly, following Na et al. (2022), or cubically (Zhu & Gupta, 2017), with the latter yielding better performance.
Evaluation on corrupted image We evaluate the sparse models trained with ADMM and SAFE, as obtained in Sec- tion 4.2, on the CIFAR-10 test set with common image corruptions and adversarial perturbations. Specifically, for
Compression-Aware Minimizer (CrAM) (Peste et al., 2022), on the other hand, extends upon the robust optimization principles of SAM to train compressible models by en-
8

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
Table 4: Comparison with IMP+SAM, CrAM, and CrAM* on
from CrAM+. A similar trend is observed in SGG further supporting this interpretation. SAFE, on the other hand, delivers competitive performance without relying on these additional techniques, highlighting the intrinsic ef- fectiveness of its smooth penalization via the augmented Lagrangian and split-variable structure of the ADMM frame- work to jointly balance sharpness minimization and sparsity. We provide additional comparison with other variants of CrAM in Appendix D.1.
SAFE.SG. which extends SAFE using similar technique as CrAM+, outperforms most baselines at moderate sparsity and all baselines at extreme sparsity.
Sparsity
Method
95%
98%
99%
99.5%
IMP+SAMiner
80.30g0.12
36.03
18.30
13.80
IMP+SAMonnc
92.50gm.0s
89.2410.00
83.74
73.73
CrAM
90.18
69.53
45.17
10.00
5. Conclusion
CrAM*
93.62.20.06
91.75g0044
88.82
81.30+0.56
In this work, we propose an effective and principled ap- proach called SAFE to obtain sparse and flat solutions by solving a constrained optimization problem based on the augmented Lagrangian, which we further extend to SAFE* by proposing generalization for the projection operation. We show that SAFE can be applied to neural network prun- ing, and as result, it not only obtains the desired flatness as well as high sparsity in the given deep model, but also enhances its generalization performance quite significantly, far better than the compared baselines, as validated across standard benchmarks. Interestingly, SAFE preserves its robustness to various data noise during both training and inference, which stems from the original sharpness mini- mization strategy. Finally, we compare with more directly related SAM-inspired baselines, demonstrating the intrinsic effectiivnness of SAFE without much reliance on auxiliary techniques. We believe that this principled approach for obtaining sparse and flat solutions concepts that have of- ten been explored rather separately in the offes significant potential.
SAFE
92.59
89.58
87.47
19.5
SAFE,SG
92.40g0.00
90.09
89.13 10.06
85.85 10.09
couraging the models to maintain strong post-compression performance under the presence of small perturbations as min, max|jel] Sp f(C(x + c)) given some compression op- eration C. This leads to the CrAM update rule Iz+1 = TV + pV f()))). Along with this, we additionally compare with CrAM+, a variant introduced by Peste et al. (2022) that simply adds the original gradient (x) in the update as 1 = 22 - [V[((( + (V))) + (=))) We run these using the optimal hyperparameters as sug- gested in Peste et al. (2022). It should be noted that the additional technique of CrAM* is somewhat auxiliary to the core robust optimization mechanism that connects CrAM to SAFE, and thus, care must be taken when associating their performance gains with the central strategy that de- fines CrAM. For better comparison, we extend this strategy to SAFE by adding the gradient computed at projected point V (C(x)) during iterative z-minimization of SAFE, which we denote as SAFE,SG-
A Acknowledgements
This work was partly supported by the Institute of Infor- mation & communications Technology Planning & Evalua- tion (ITT)) grant funded by the Korean government (MSIT) (ITTP-201 111-2209---010000 Artificial Intelligence Graduate School Program (POSTECH) and RSS-002--12220959 (part2) Few- Shot learning of Causal Inference in Vision and Language for Decision Making), the National Research Foundation of Korea (NRF) grant funded by the Korean government (MSIT) 109221F1AA1666 2023-00265444).
As shown in Table 4, SAFE and SAFE4SG outperforms IMP+SAM and CrAM, where SAFE-SG demonstrates com- petitive performance to CrAM* on moderate sparsity and outperforms it in extreme sparsity. We suspect that the pruning operation in IMP+SAM may be overly abrupt, po- tentially hindering the benefits of sharpness minimization. We additionally observe similarly strong performance over IMP+SAM on language model pruning in Appendix D.2. Conversely, although CrAM achieves competitive results, its benefits fails to extend to extreme sparsity. More cru- cially, while SAFE achieves reliable performance without much additional techniques, CrAM, by itself, performs poorly in all sparsities, depending heavily on auxiliary tech- niques to drastically improve performance. This suggests that caution is warranted when attributing these gains to the effectiveness of their robust optimization formulation inspired by sharpness-minine minimization, which is more likely provided through the use of both the original dense gradi- ent and the sparse gradient computed at the projected point
RS-2023-00210466 RS-
Impact Statement
This paper presents work aimed at advancing the field of Machine Learning, with the potential to influence both the- oretical understanding and practical applications. While our contributions do not directly raise immediate concerns requiring specific emphasis, we acknowledge that advance- ments in this domain can have far-reaching societal impli-

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
cations. We will ensure ongoing discourse the broader impact of our work in diverse contexts if the need is later recognized.
Foret, P., Kleiner, A., Mobahi, H., and Neyshabur, B, Sharpness-aware minimization for efficiently improving generalization. ICLR, 2021.
Frankle, and Carbin, M. The lottery ticket hypothesis:
References
Finding sparse, trainable neural networks. ICLR, 2019.
Andriushchenko, M., Croce, F., Müller, M., Hein, M., and Flammarion, N. A modern look at the relationship be- tween sharpness and generalization. ICMI,, 2023.
Frantar, E. and Alistarh, D. Optimal brain compression: A framework for accurate post-training quantization and pruning. NeurIPS, 2022.
Baek, C., Kolter, J. Z., and Raghunathan, A. Why is sam robust to label noise? ICLR, 2024.
Frantar, E. and Alistarh, D. Sparsegpt: Massive language models can be accurately pruned in one-shot. ICML, 2023.
Bahri, D., Mobahi, H., and Tay, Y. Sharpness-aware mini- mization improves language model generalization. ACL, 2022.
Goodfellow, I. J., Bengio, Y., and Courville, A. Deep Learn- ing. MIT Press, 2016.
Bair, A., Yin, H., Shen, M., Molchanov, P., and Alvarez, J. Adaptive sharpness-aware pruning for robust sparse networks. arXiv, 2023.
Han, S., Pool, J., Tran, J., and Dally, W. Learning both weights and connections for efficient neural network. NeurlPS, 2015.
Hassibi, B. and Stork, D. Second order derivatives for network pruning: Optimal brain surgeon. NeurIPS, 1992.
Beck, A. and Teboulle, M. A fast iterative shrinkage- thresholding algorithm for linear inverse problems. SIAM, 2009.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learning for image recognition. CVPR, 2016.
Benbaki, R., Chen, W., Meng, X., Hazimeh, H., Pono- mareva, N., Zhao, Z., and Mazumder, R. Fast as chita: Neural network pruning with combinatorial optimization. ICML, 2023.
He, Z., Xie, Z.,, Zhu, Q., and Qin, Z. Sparse double descent: Where network pruning aggravates overfitting. ICML, 2022.
Blumensath, T. and Davies, M. E. Iterative hard threshold- ing for compressed sensing. Applied and computational harmonic analysis, 2009.
Heek, J., Levskaya, A., Oliver, A., Ritter, M., Rondepierre, B., Steiner, A., and van Zee, M. Flax: A neural network library and ecosystem for JAX, 2023, URL http:// gi github. ctthub.comm//cc cww/goooleeeoooooooooooooooooooo e ftlfflax..
Boyd, S., Parikh, N., Chu, F.,, Peleato, B., Eckstein, J., al. Distributed optimization and statistical learning via the alternating direction method of multipliers. Foundations and Trends@ in Machine learning. 2011.
Hendryeks, D. and Dietterich, T. Benchmarking neural network robustness common corruptions and perturba- tions. ICLR, 2019.
Hoefler, T., Alistarh, D., Ben-Nun, T., Dryden, N., and Peste, A. Sparsity deep learning: Pruning and growth for efficient inference and training neural networks. JMLR, 2021.
Bradbury, J., Frostig, R., Hawkins, P., Johnson, M. Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J., Wanderman-Milne, S., and Zhang, Q. JAX: composable transformations of Python+NumPy programs, 2018.
Huang, T., Singhania, P., Sanjabi, M., Mitra, P., and Raza- viyayn, Alternating direction method of multipliers for quantization. International Conference on Artificial Intelligence and Statistics, 2021.
Chaudhari, P., Choromanska, A., Soatto, S., LeCun, Y., Bal- dassi, C., Borgs, C., Chayes, J., Sagun, L., and Zecchina, R. Entropy-SGID: Biasing gradient descent into wide valleys. ICLR, 2017.
Hubara, L., Chmiel, B., Island, M., Banner, R., Naor, J., and Soudry, D. Accelerated sparse neural training: provable and efficient method to find m transposable masks. NeurIPS, 2021.
Chen, X., Hsieh, C.-J, and Gong, B. When vision trans- formers outperform resnets without pre-training or strong data augmentations. ICLR, 2022.
Evci, U., Gale, T., Menick, J., Castro, S., and Elsen, E. Rigging the lottery: Making all tickets winners. ICML, 2020.
Izmailoy, P., Wilson, A., Podoprikhin, D., Vetrov, D., and Garipov, T. Averaging weights leads to wider optima and better generalization. UAI, 2018.
10

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S. Fantastic generalization measures and where to find them. ICLR, 2020a.
Meng, X., Behdin, K., Wang, H., and Mazumder, R. Alps: Improved optimization for highly sparse one-shot pruning for large language models. NeurIPS, 2024.
Jiang, Y., Neyshabur, B., Mobahi, H., Krishnan, D., and Bengio, S. Fantastic generalization measures and where to find them. ICLR, 2020b.
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models. ICLR, 2022.
Meta. The llama 3 herd of models. arXiv, 2024.
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, T. P. On large-batch training for deep learning: Generalization gap and sharp minima. ICLR, 2017.
Na, C., Mehta, S. V., and Strubell, E. Train flat, then com- press: Sharpness-aware minimization learns more com- pressible models. EMNLP, 2022.
Khanh, P. D., Luong, H.-C., Mordukhovich, B. S., and Tran, D. B. Fundamental convergence analysis of sharpness- aware minimization. NewrlPS, 2024.
Nakkiran, P., Venkat, P., Kakade, S. M., and Ma, T. Optimal regularization can mitigate double descent. ICLR, 2021.
Neyshabur, B., Bhojanapalli,, S., McAllester, D., and Srebro, N. Exploring generalization in deep learning, NeurIPS, 2017.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv, 2017.
Orvieto, A., Kersting, H., Proske, F., Bach, F., and Lucchi, A. Anticorrelated noise injection for improved general- ization. ICML, 2022.
Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. 2009.
Kurtic, E. and Alistarh, D. Gmp*; Well-tuned gradual magnitude pruning can outperform most bert-pruning methods. arXiv, 2022.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., al. Pytorch: An imperative style, high-performance deep learning library. NeurlPS, 2019.
Kusupati, A., Ramanujan, V., Somani, R., Wortsman, M., Jain, P., Kakade, S., and Farhadi, A. Soft threshold weight reparameterization for learaale sparsity. ICML, 2020.
Peste, A., Iofinova, F.,, Vladu, A., and Alistarh, D. Ac/de: Alternating compressed/tecomprssee training of deep neural networks. NeurIPS, 2021.
Kwon, W., Kim, S., Mahoney, M. W., Hassoun, J., Keutzer, K., and Gholami, A. A fast post-training pruning frame- work for transformers. NeurlPS, 2022.
Peste, A., Vladu, A., Kurtie, E., Lampert, C. H., and Alis- tarh, D. Cram: A compression- aware minimizer. ICLR, 2022.
LeCun, Y., Denker, J., and Solla, S. Optimal brain damage. NeurIPS, 1989.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. Exploring the limits of transfer learning with unified text-to-text transformer. JMLR, 2020.
Lee, N., Ajanthan, T., and Torr, P. Snip: Single-shot network pruning based on connection sensitivity. ICLR, 2019.
Lee, N., Ajanthan, T., Torr, P., and Jaggi, M. Understand- ing the effects of data parallelism and sparsity on neural network training. ICLR, 2021.
Ramanujan, V.. Wortsman, M., Kembhavi, A., Farhadi, A., and Rastegari, M. What's hidden in randomly weighted neural network? CVPR, 2020.
Li, H., Xu, Z., Taylor, G., Studer, C., and Goldstein, T. Visualizing the loss landscape neural nets. NeurIPS, 2018.
Sanh, V., Wolf, T., and Rush, A. Movement pruning: Adap- tive sparsity by fine-tuning. NeurIPS, 2020.
Lin, T., Stich, S. U., Barba, L., Dmitriev, D., and Jaggi, M. Dynamic model pruning with feedback. ICI., 2020.
Shin, S., Park, W., Lee, J., and Lee, N. Rethinking pruning large language models: Benefits and pitfalls of recon- struction error minimization. EMNLP, 2024.
Liu, B., Zhang, 7., He, P., Wang, Z., Xiao, Y., Ye, R., Zhou, Y., Ku, W.S., and Hui, B. A survey of lottery ticket hypothesis. arXiv, 2024.
Shin, S., Lee, D., Andriushchenko, M., and Lee, N. Critical influence of overparameteriizaion on sharpness-aware minimization. UAI, 2025.
Liu, Z., Sun, M., Zhou, T., Huang, G., and Darrell, T. Re- thinking the value of network pruning. International Conference on Learning Representations 2019.
Simonyan, K. Very deep convolutional networks for large- scale image recognition, arXiv, 2014.

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
Song, H., Kim, M., Park, D., Shin, Y., and Lee, J.-G. Learn- ing from noisy labels with deep neural networks: A sur- vey. TNNLS, 2022.
Zhou, X., Zhang, W., Xu, H., and Zhang, T. Effective sparsi- fication of neural networks with global sparsity constraint. CVPR, 2021.
Sun, M., Liu, Z., Bair, A., and Kolter, J. Z. simple and effective pruning approach for large language models. ICLR, 2024.
Zhu, M. and Gupta, S. To prune, or not to prune: exploring the efficacy of pruning for model compression. arXiv, 2017.
Szegedy, C., Zaremba, W., Sutskever, I., Bruna, J., Erhan, D., Goodfellow, I., and Fergus, R. Intriguing properties of neural networks. ICLR, 2014.
Tanaka, H., Kunin, D., Yamins, D. L., and Ganguli, S. Prun- ing neural networks without any data by iteratively con- serving synaptic flow. NeurlPS, 2020.
Tibshirani, R. Regression shrinkage and selection via the lasso. Journal of the Royal Staristical Sociery Series B: Statistical Methodology, 1996.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv, 2023.
Wang, C., Zhang, G., and Grosse, R. Picking winning tickets before training by preserving gradient flow. ICLR, 2020.
Wang, Y., Yin, W., and Zeng, J. Global convergence of admm nonconvex nonsmooth optimization. J. Sci. Comput., 2019.
Wei, Z., Zhu, J., and Zhang, Y. Sharpness-aware minimiza- tion alone can improve adversarial robustness. 2023.
Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., Davison, Shleifer, S., von Platen, Ma. C., Jernite, Y., Plu, J., Xu, C., Scao, T. I.., Gugger, S., Drame, M., Lhoest, Q., and Rush, A. M. Huggingface's transformers: State-of-the-art natural language processing. arXiv, 2020.
Zeng. W. and Urtasun, R. Miprume: Multi-layer pruning for automated neural network compression. arXiv, 2018.
Zhang, T., Ye, S., Zhang, K., Tang, J., Wen, W., Fardad, M., and Wang, Y. A systematic dnn weight pruning frame- work using alternating direction method of multipliers. In ECCV, 2018
Zhang, Y., He, H., Zhu. J., Chen, H., Wang. Y., and Wei, Z. On the duality between sharpness-aware minimization and adversarial training. ICML, 2024.
Zhou, P., Feng, J., Ma, C., Xiong, C., Hoi, S. C. H., et al. Towards theoretically understanding why sgd generalizes better than adam in deep learning. NewrlPS, 2020.
12

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
A. Convergence analysis of SAFE
In this section, we show that SAFE converges to a stationary point of the augmented Lagrangian. Precisely, we first prove that our proposed sharpness minimized z z-minimization iterate converges to the stationary point for the augmented Lagrangian for the original loss function (i.e., [(x) f()) + A/2 - 2 w113) with respect x. With this result, build the rest of the proof upon well-studied result of ADMM (Boyd et al., 2011; Wang et al., 2019; Huang al., 2021).
Prior to providing detailed proofs of core lemmas and corollaries, we first describe the properties of the augmented Lagrangian [(x) based on the Assumptions 3.2 and 3.3.
Defining strong convexity as:
Definition A.1. (a-strong convexity) Let f be differentiable. f is a-strongly convex there exists a > such that f(y) 2 f(z) + (Vf(z), - z) + lly x||?
present the smoothness and strong convexity of ((z) through the following lemma:
Lemma A.2. Let f be 3-smooth and j--weakly convex. Then ((x) ((x) + || ull2 is also (3 + X)-smooth and
(A pt-strongly convex for X > p.
Proof. From Theorem 3.2, we have
VV((x) - vi(x)] [|+[+))+++)++++++++++++)< A(x - V5(w) - A(v x)ll
- fV Vf(wi + +1 - Xyll
< (8 + A]i|= - yll
thus by definition, is (B + A)-smooth.
Also, from first-order condition of convexity of [(x + 5 | from Theorem 3.3, we have
2 f()) "II ((x + 1|= (v)(x) + paz x)
f(v) f(z) + (v)(z) + pr,y z)
H
|
1
|1yl
2
2
f()) 2 ((x) + (v/(x) + pr. 2)
11
(y
y x)
2
X fv
L(y) ((=)+ + (V((x).y 2)
(y
y
z) +
x. y x)
2
2
X H
+ -(w) [(x)+ + (V((x), 2)
11 x
2
Since - fu > 0, by definition is (A- - sttstrlyyy convex.
A.1. Proof of Theorem 3.5
We provide the convergence proof of our sharpness minimizing x iterates. We mostly follow the procedure of Khanh et al. (2024), with the objective and the update rule altered to match our configurations. Before we proceed, we recall prior results from Khanh et al. (2024)
Lemma A.3. (Lemma B.1 of Khanh et al. (2024)) Let {a!"), (6(9), (c()) be sequences of nonnegative numbers satisfying the conditions
(t+1) (0) (6(00(0) (0(0)c+ (l for sufficient large t N,
(a)
{b(t)] is bounded, ((0) 00.
Σ (1) oo, and Σ (l0 <c.
(b)
1=1
Then we have that (t) -> 0 as > 0e
With this, we derive the convergence of our sharpness minimizing x iterates in the following lemma:
13

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
Lemma A.4. (Convergence of x-minimization) Suppose that Theorem 3.1 and 3.2 hold and let {x be generated by Equation (5) in Algorithm with step-size rft) and perturbation radius (t)) satisfying
00,
n p0000 < 00, lim p()) < 1/B.
(6)
Let [(x)= = f(x) + 2 + u|12 and assume that enN ((x(1)) > -o0. Then v - O.
Proof. Let the gradient of our sharpness minimizing x iterate Equation (4) as
Vi(ato) ) - -xx(0)22+)) u), [V(a(o)1]
(0) = Vf
211 + 0
- -2+
where we drop the subscript denoting the outer SAFE iterate for simplicity, We also denote 3 == B -3 n.. We first begin from the 3-smooothness of L. from Theorem A.2
3
((xl(+2) S ()x(0)) (VL(+ (1+1) - (0)) +
2 (1)((t1)) (00
()x ,/00 (Vi(s)(1)) +
1012
2
012
= ((x 1/ (1 (V) g. (00) + (194) Vi(a000)1 (
112((1(0)))
77)
Here, we bound each []g(0) - vi(x(0)) and (0) 1 (z(((0)) using the 3-smoothness as follows
(0 (((x(()))l
4 f (1) (()) V)(((())
((x((() 2 =z+ u. (5(olt)) -(x(10) +
++=))) u))
Vf
V
(1)
(1))
f(0(1>))
3
(1)
V
(z((tt)
V
[(((0)
(8)
and using this result, we have that
()))) v((((0))
= (g()) vi(Z(1), v((a())) + 112
2 ()) vi(x'))] [V( + [1((2=0))1
> -80 (0) / + ()
(9)
Applying Equation (8) and (9) back Equation (7) gives
14

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
- ((x((11) =((x())- (1- (())) + 3y (1910- -
I ))x (1 - ))) (-30 [ + |V((=00)1 88,1122
[V((2(0))1
2
(
33000221002
-((00)-
(2 By ) VV + Byy (0 - (7 [( +
2
2
(1)
2 (2 37 ( 12 ((1))11)) (1- Boy(o) (1 [[
33,1(12
S ((x1))
+
2
2
(1)
(2 (0 -B0 B3 ( | + p1)())110) (1-
((x(0)
(1- ()n +
2
2
2
(t)
(00 ( 1 1n-B) + 12) p())
Br1())
((x
/1 - 30th ,0 / |
2
From here, we find e₁ and 02 (0, and T N such that
1 (2 - Bn - 3p + () 2 e1,
1 3r + 00
c2, and Br < 1 for all t > T,
2
2
where applying this gives us
((x(t+))) S (1 + c2)0)(0)p0))
(10)
Also, defining w(() :- 02 Eie
sim 31 ((9) for t E N, we get that wi()) + 0 as oo and w(() w(+11) (2 for
all t E N. Then Equation (10) can be rewritten as
(1((+))) + w(++1) S - (
(11)
Here we telescope this bound from t ' to oo and combine with inf f(x!10) > -oo and (() - 0 as t - that
(2xx < (2(x)0)) - ((o((+1)) - w(()) - (tt1))
(12)
:=1 s ((x(k))
inf ((x(t)) + (k(
(13)
TEN
We finally employ Theorem A.3 with (*)) == (v)((())) |l, b00 = 0,, and cll) = (1 for all E N to derive
((x(t)) + 0. Here, condition (a) is satisfied due to the estimates
(1+1) (1)
V [(( ||| - (
I I/ ((x(t)) - V (
S (112 - 01 -
(1g""
I Bn (IVL()))))) (1) - v(()(()1
I "n (I 11g (6) (((x
(00 + 1 for all E.N.
Also, the conditions in (b) hold by Equation (6) and Ex= T/ (2) (( < co. Thus from Theorem A.3, (+))))
(! + 0 as t + Do.
This shows that rumning Equation (5) produces a sequence that converges to the stationary point of the augmented Lagrangian [(x, 2, u) = f()) + (allooa(az)) + 2 [|ul[B + as we show in the following section.
||x - 2 + 31 with respect to r. This is crucial for the convergence of SAFE
15

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
A.2. Proof of Theorem 3.6
Here we prove the convergence of SAFE. This is straightforward procedure: given our convergence guarantee of the iterates in Theorem A.4, the convergence properties of SAFE be described in terms of classical ADMM. Thus, in this section, we walk through the convergence analysis of ADMM as provided in Huang et al. (2021) and demonstrate how sharpness minimizing iterates is applied within the proof.
Corollary A.5. (Convergence of SAFF) Suppose that Assumptions 3. 1-3.3 hold. Assume further that ö is chosen large enough so that 8 8- (Ō p)/2 0. Let (,,200 ū) be limit point of SAFE algorithm. Then & is d-stationary point of the optimization problem (1).
Proof. By Theorem 3.5, every #k+1 found by running Equation (5) until convergence is the stationary point of the augmented Lagrangian, i.e. V (I 2R, uk) 0. This gives us the standard update rule of classical ADMM, where the results of Huang et al. (2021) can be adapted directly.
B. Experimental Details
We present various details our experimental setup. All experiments are run across three different seeds, and the results provided as the mean and the standard error.
B.1. Hyperparameters
Table 5: Hyperparameter details used/searched for SAFE and SAFE* Here, perturbation radius and dual interval were searched only in ResNet-20/CIFAR-I0 and LLaMa-2-7b, then applied across all settings and target sparsity.
Vision
Language
Epoch
200 (ResNer20) 300 (other)
30
Base optimizer
SGD
Adam
Batch size
128 for 126 when using GPUs for data paralictism)
Leaming rate
0.1
0.0002
Leaminn rate schedule
cosine
linear
Warm-up epoch
Weight decay
0.0001
Momentum
BNT sample size
10000
Perturbation radius (p)
(0.01, 0.1, 0.011...0.01100020.300
(0.0001, 0.0002. 0.0005, 0.01)
Dual-update interval (K)
11,2,4,8, 16, 32, 64, 128, 256, 512, 1024, 2048)
{16, 32, 64)
Penalty parameter (X)
{10-4,10-2) 10_,,,0,,,,000 1,000",102.,......010111011 10-11
(0.001, 0.005, 0.01, 0.13
Penalty schedule
cosine warmup
constant
Table Best performing penalty parameter of SAFE for VGG-19 and ResNet-20/32.
CIFAR-10
CIFAR-100
Sparsity
19
ResNet-20
VGG-19
ResNet-32
70-95%
10-1
10-3
10-3
98%
10-3
10 3
99%
10-2
10
10-2
99.5%
16

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
Across all experimental settings, the values for hyperparameters remain consistent, with the exception of the penalty parameter X. We report these in Table 5. Basic hyperparameters such as learning rate, batch size, weight decay, and momentum are set to standard values commonly used in the literature (Kusupati et al., 2020; Ramanujan et al., 2020; Liu et al., 2019). For SAFE-specific hyperparameters, including perturbation radius and dual-update interval, values were optimized on ResNet-20/CIFAR- 10, ILaMa-2-7b and applied universally across all settings, with only the penalty parameter for image classification tasks being adjusted for each setting, which we report in Table This demonstrates the general applicability of its hyperparameter values across different tasks.
Also, we use a cosine warmup schedule for the penalty parameter for the vision tasks, which increases the penalty parameter from 0 to the final target value with cosine curve throughout training iterations.
B.2. Experimental Details in Section 4.1
Here train the standard 3-layer MLPs (with hidden layers of 300 and 100) on MNIST. For the sparsity experiment, standard dense training and SAFE with target sparsity of 90% and plot the distributions of weights. For the flatness measurements, we run SAFE and ADMM (Zhang et al., 2018) (a simple non-sharpness-awar baseline) with target sparsity of 90%. We then plot the loss landscape of the found solutions using standard visualization methods (Li et al., 2018), and compute their maximum Hessian eigenvalue (sharpness) using the power iteration method.
B.3. Experimental Details in Section 4.2
We use standard ResNet and VGG architectures, with VGG having batch-norm layers instead of using dropout. For data augmentation, we used standard techniques such as random cropping and flipping. We used SGD for the base optimizer for SAFE and all baselines. All CIFAR experiments are conducted using a single or three NVIDIA RTX 3090, where a batch size of 126 was used this case. For 1001R- we trained ResNet-20 for 200 epochs, and ResNet-32 and VGG-19 for 300 epochs.
B.4. Experimental Details in Section 4.3
Training data is processed following the standard settings in SparseGPT (Frantar & Alistarh, 2023), where we randomly sample 128 data points with sequence length of 2048 from the first shard of C4 (Raffel al., 2020). All experiments were conducted on single GPU (NVIDIA A6000 L40S) or HPU (Intel Gaudi2). We I.LaMa-277-hhf, L.LMA-2-133-hf and A.LMMA3311-4 1-8 B models from the HuggingFace model hub (Wolf et al., 2020), implemented in PyTorch (Paszke et al., 2019). For SAFE and SAFE*, we perform pruning over 30 epochs using the Adam (Kingma & Ba, 2017) optimizer as the base optimizer, with the hyperparameter /3 set to (0.9, 0.95), without weight decay.
B.5. Implementation and Reproduction Details
The code to reproduce the results of the paper is provided in JAX² (Bradbury et al., 2018; Heek et al., 2023) and PyTorch² (Paszke et al., 2019). Specifically, all image classification experiments using SAFE were conducted in JAX, while PyTorch was used for LLM experiments due to better support for official implementations and pretrained checkpoints of widely adopted models such as LLaMA. To obtain baseline performances, we either run our own implementations (ADMM, GMP, Magnitude) or official ones (SParseGPT, Wanda, ALPS), or refer to reported results from prior work (LTH from Wang al. (2020); PBW and MLPrune from Zhou et al. (2021)). For sound comparison, when running our own or official implementations, we align key settings-sut such as model architecture and training withhs-with those used in prior works to produce reported performance.
17

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
C. Detailed Results for Image Classification Tasks
We present precise numeric values for Figure 2 in Table 7.
Table 7: Validation accuracy of VGG-19 and ResNet-22032 models pruned with SAFE and various baseline methods, trained on CIFAR-10 and CIFAR-IO0 100, across different sparsity levels. The results show that SAFE generally outperforms the baseline methods across all evaluated sparsity levels, indicating its robustness and effectiveness in maintaining accuracy even under high sparsity conditions.
Sparsity
Dataset
Model
Method
90%
95%
98%
99%
99.5%
GMP
93.37.20.0
13.13
93.0840000
92.70
90.63
PBW
93.87
93.57
92.83
90.89
10.00
MLPrune
93.70
93.45
92.48
91.44
88.18
VGG-19
LTH
93.51
92.92
92.34
CIFAR-10 10
ADMM
93.86.00110
93.62
93.58
92.54=0.07
88.5340.16
SAFE
94.65
94.44
94.05
93.93
93.56:10.02
GMP
92.94
91.81
89,4 89,42.2017
85,15
75.83
ResNet-20
ADMM
91.88
89,9620.27
86.96.1.0e
82.25 20.12
73.72:22.85
SAFE
93.4440.01
92.59
89.58
87.47=2007
79.55
GMP
72.00
71.81
69.55
66.98
62.77
PBW
72.41
70.53
58.91
1.00
1.00
MLPrune
71.56
70.31
66.77
60.10
50.98
VGG-19
LTH
72.78
71.14
68.95
ADMM
72.93
71.17
70.02
67.23 67.23
43.40
SAFE
73.67 +0.21
42.83
71.73_00.09
70.0240.2
67.23
CIFAR-100
GMP
71.69
69.10
65.15
58.10
42.93
PBW
72.19
68.42
58.23
43.00
20.75
MLPrune
70.33
61.73
37.86
22.38
13.85
ResNet-32
LTH
68.99
65.02
57.37
ADMM
70.85
68. 68.74
63.7540.06
49.13
12.34
SAFE
73.89 +0.24
72.33 #0.08
67.7440.24
62.77 20.11
51.45 40.32
D. Additional Comparison with SAM-based pruners
Table Comparison with other variants of CrAM on ResNet-20/CIFAR 10, where we also apply similar techniques to SAFE
Table 9: Comparison with IMP+SAM on LLaMA2-76 7b for 50% sparsity.
Sparsity
Method
95%
98%
99%
99.5%
CrAM
69.53 18.18
45.17
10.00
CrAM
93.62_0.06
91.75
88.82-t-0.18
81.30
Perplexity
CrAMMali
92.21 =0.70 00.79
91.05 31
88, 88.834+0.19
84,82 27
Method
C4 / WikiText
CrAM Mulli
92.17 0.91
91.06
88.98440.04
84.95
IMP+SAM
18.27 176.00
SAFE
92.59 00.09
89.58
87.47 00007
19.15
SAFE
8.91 / 6.79
SAFE,SG
92.40
90.09 40.13
89.1344.08
85.85 1..0..00
SAFEMali
91.98 1.11
48.89
86.42.10.10
83.00
SAFE
92.17
90.7940.04
88.1 88.1140.14
84.92-10.14
D.I. Other variants of CrAM
Here we compare SAFE with two additional variants of CrAM introduced in Peste et al. (2022): CrAMMuni and CrAM
Multi Specifically, CrAMMati additionally changes the target sparsity at each iteration chosen randomly from a predefined set, whereas CrAM Mull combines this with CrAM+ explained in Section 4.5. Similarily Section 4.5, for better comparison, we
18

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
apply these "+" and "Multi" strategy SAFE by adding the gradient computed compressed point to the sharpness-awand aware gradient of SAFE Vf + /V every V f(C())) in the z-update, which we denote as SAFE.sG. SAFEMutti and Sa. for applying all both.
V)(C(=)), and changing the target sparsity every z-updates and for
As shown in Table 8, SAFE achieve competitive performance until 98% sparsity and outperforms CrAM in exterme sparsities, despite improved performance of CrAM from CrAMMulii and CrAM Multi- In particular, while SAFE achieves reliable performance without much additional techniques, CrAM, by itself, performs poorly all sparsities, depending heavily on various auxiliary techniques to drastically improve performance. This is further supported by SAFE-SG. where similar sort of benefits are observed. However, we find that the strategy of additionally employing gradients from various projected points also generally effective in SAFE. Although these findings intriguing and merit further study, comprehensive analysis lies beyond the scope of this work and is left to future research.
D.2. IMP+SAM Language model pruning
We extend the experiment in Section 4.5 LLM pruning, where prune LLLam2-7bb to 50% sparsity by applying IMP+SAM to block-wise reconstruction error objective similarly to SAFE. Here, We similarly perform pruning every 5 epochs with sparsity increasing linearly or cubically over the same number of epoch as SAFE for fair comparison. As shown in Table 9, SAFE outperforms IMP+SAM in LLM tasks similarly to results in image classification.
E. Computation Cost Analysis LLM Pruning
We provide theoretical and empirical analysis of the computation costs of SAFE and various baselines in LLM pruning.
Table 10: Time complexity analysis for LLM post-pruning techniques
Table 11: Wall-clock time
Method
Time Complexity
Explanation
Hessian computation over samples (Nd?)
SparseGPT O(Ly(Nd + ))))
+ Hessian inverse (d")
for Ly layers in a single block
Method
Time (s)
Activation norm computation (Nd)
Magnitude
0.48
Wanda
O(Lo(Nd+ &P))
+ weight multiplication (d²)
Wanda
3.98
for LB layers in single block
SparseGPT
15.82
Hessian computation over samples (Nd")
ALPS
788.66
ALPS
O(La(N² + k(l
cigendecomposition (d*)
SAFE
310.68
penalized inverse for k ADMM iterations (kd3)
for Ly layers in single block
Backpropagation through L layers in single block for & iterations
SAFE
O(I.zbkdd)
E.1. Theoretical Time Complexity
We compare the time complexity for pruming a single transformer block with SAFE, SparseGPT, Wanda, and ALPS in terms of hidden dimensions number of data N or batch size h, number of iterations k. and number of layers in a single block Lg to observe how the computation cost scales. The results are given in Table 10.
We observe that while the computation cost of SAFE, similarily to ALPS, scales with the iteration, it does not scale with the size of the dataset N. Also, it only scales quadratically with the hidden dimension of the model. In comparison, SparseGPT and ALPS scale cubically. Given that scaling model and data sizes are a central strategy in the development of large language models, this highlights the advantage and adequecy of SAFE in the current era of large-scale models.
E.2. Wall-clock Time
We report the wall clock time required for each pruning method on the LLaMA-2-7B model at 50% sparsity. Specifically, Table 11 shows the time taken to prune the first transformer block consisting of the self-attention and the feed-forward module. All measurements were conducted on a single Nvidia A6000 GPU to ensure a fair comparison across pruning methods.
19

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
F. Ablation Study
We conduct ablation of various hyperparameters of SAFE on ResNet-20 trained on CIFAR-10.
F.1. Effects of Penalty Parameter
aecuracy
0.94
Dict. to constraint
10
Sparsity
0.8
10
0.92
Sparse
10-1
Sparse (BNT)
0.0001
0,001
0.01
0.1 0.0001 0.001 0.01
0.1
0.0001
0.001 0.01
0.1
0.0001
0.001
0.01 Sparsity=0.95
0.1
0.0001
0.001
0.01
0.1
Sparsity-0.7
Sparsity-0.8
Sparsity-0.9
Peralty parameter (A)
Penalty parameter (A)
(a) Dense/sparse validation accuracies
(b) Distance to constraint
Figure 3: Effect of the penalty parameter À on final validation accuracy of dense/sparsified models (a) and the distance from the constraint (b) over various levels of sparsity. Larger À relieves the performance drop in the final projection step while degrading the performance of the original dense model. Also, BNT provides larger benefits for smaller À and the target sparsity
We observe how the penalty parameter À impacts various aspects of the final model in terms of the validation accuracy of the final network before and after projection (denoted as dense/sparse model in the legend, respectively) and the distance to the constraint. We vary À between {10 4, 10 10 2, ") and observe this for target sparsity of {0.7, 0.8, 0.9, 0.95), with is reported in Figure 3. We find that while larger À relieves the performance degradation in the final projection step by pushing the network closer to the sparsity constraint during training, it also degrades the performance of the original dense model. This indicates that a balance between objective minimization and constraint satisfaction would yield the best results. We also find that the model becomes more sensitive to À for higher target sparsity. where the degree to which large À incurs performance degradation while small À results in failure of the model to get sufficiently close to the constraint becomes much more severe. This highlights the challenge of training a model with extreme sparsity, which demands a careful balance between minimizing the objective function and satisfying the sparsity constraint.
F.2. Effects of Batch-norm Tuning
We also observe how re-evaluating the batch statistics (i.e., batch-norm tuning or BNT) for the final projected parameters impacts the performance of sparse networks over various sparsities. Our assumption is that should be helpful when the parameters change greatly after projection, which will cause the hidden features to deviate from the statistics computed during training. The results are presented in Figure 3. We observe that the benefits of BNT are pronounced when À is small, where the distance to the constraint is the largest, fitting our assumption. However, we also observe that it fails to recover the performance in higher sparsity, despite having a similar distance to the constraint. We suspect it is due to the degraded quality of the sparse network having more impact on the performance in high sparsity, rather than the misalignment of the batch-norm statistics.
F.3. Effect of X Scheduling
There are many strategies employed by the community known to boost performance for deep neural network training. One such strategy is scheduling, which has been a de facto for important hyperparameters such as learning rate (Goodfellow et al., 2016). In particular, gradual sparsity schedules (Zhu & Gupta, 2017; Benbaki et al., 2023) have been widely adopted to enforce less sparsity in the initial phases of training. Here we observe whether a similar effect can be transferred the penalty parameter X, an important parameter for controlling how strongly to push toward the sparsity constraint at any point during training. Precisely, we test whether slow increase from zero to the targeted penalty X5 will yield improvements over constant X. We train ResNet-20 on CIFAR-10 to (70%, 80%, 90%, 95%} sparsity with SAFF using linear, cosine schedules, and constant X, and observe how this affects the final accuracy of the sparse network in Table 12. We find that scheduling consistently yields overall higher accuracy, especially in extreme sparsity where it yields +2% increase.
To gain further insight as to how this occurs, we analyze how scheduling affects various aspects of training through the
20

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
Ar
50
Ponalty paraieter
accurasy
0.8
Dist. constriint
40
30
Sparce
5d lduue constart
0.2
0.0
liner
0
50
100 Epochs
150
200
100
150
200
50
100 Epochs
150
200
50
100
150
200
Epochs
Epochs
(a) Schedules
(b) Sparse val. acc.
(c) Dense acc.
(d) Distance to constraint
Figure 4: Effects of different choices of penalty parameter schedules validation accuracy of sparsified/dense network (b-c) and the distance to the target sparsity constraint (d) over the training process of ResNet-20/CIFAR 10 using SAFE on 95% sparsity. It is observed that scheduling yields better performance, seemingly allowing the network to move away from the constraint in the initial phases to focus more on training, which potentially underscores its impact on securing performance.
Table 12: Validation accuracy for various penalty parameter À schedules on ResNet-20 trained on CIFAR-10. The use of scheduling generally improves the final accuracy of the sparse network, especially at extreme sparsity.
Sparsity
70%
80%
90%
95%
Constant
93.79=0.16
93.33
92.1340.13
90.78
Linear
93.7 93.78=0.04
93 93.6640.17
93.0640.04
92.20=o.o
Cosine
93.9820.09
93.67 10.13
93.4440.01
92.5920.09
validation accuracy of the nearest sparse network (sparse val. acc.), the validation accuracy of the dense network (dense val. acc.), and distance to the constraint Figure 4. Here, we observe that while constant penalty pushes the network drastically close to sparsity in the initial stages, scheduling allows the network temporarily stray away from the constraint. This seems highlight that the initial phase of training is important for securing the performance of the final sparse network.
F.4. Effeets of dual-update interval K
Sparsity
Dense val accuracy
0.940
Sparse val accuracy
Dist. to constraint
0.90
0.9
0.935
0.95
Sparsity 0.7
Sparsity
0.85
0.930
0.8
0.8
0.9
0.80
2 & 8
2048
2048
26 525
1024
Dual-update interval (K)
Dual-update interval (K)
Dual-update interval (K)
Dense validation accuracy
(b) Sparse validation accuracy
(c) Distance constraint
Figure 5: Effect of the dual-update interval K on final validation accuracy of densc/sparsified models (a. b) and the final distance from the constraint (c) over various levels of sparsity. In our search range, K has little impact on accuracy and distance to the constraint. However, in target sparsity of 95%, large K fails to sufficiently push the network towards the sparsity constraint, resulting in performance degradation on the sparsified network.
observe how the dual- dual-update interval K impacts the final validation accuracy before and after projection (denoted as
21

---

SAFE: Finding Sparse and Flat Minima to Improve Pruning
dense/sparse val accuracy, respectively) and the distance the constraint in Figure 5. We find that within our search range. K has little impact the final network. However, in the case of target sparsity of 95%, we observe that large K fails sufficiently push the network towards the sparsity constraint, resulting in performance degradation of the sparsified network. We can expect this trend to appear for all sparsity levels under increasing values of K, since this will result in less execution of dual ascent iteration for constraint satisfaction.
22


Reviews:
[{"review_id": "PFxNRclCdx", "reviewer": "Reviewer_FtwJ", "summary": "This paper presents an algorithm to produce networks that are both flat and sparse. Recent work has demonstrated that flat networks have better generalization and have proposed some approaches for producing flat and sparse networks. This paper introduces a theoretically principled approach for producing sparse and flat models. The approach is based on an augmented Lagrange dual approach and is proven to converge. The proposed algorithm outperforms other pruning algorithms across several tasks, networks, and pruning strategies.\n\n## Update after Rebuttal\n\nThe authors have answered my questions and I maintain a score of 4.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for finding our work interesting and giving us constructive suggestions. We make clarifications to specific comments below.\n\n---\n\n**Robustness to different types of noisy data**\n\nThank you for the insightful suggestion. To test the robustness of SAFE against image corruptions and adversarial attacks, we evaluate ResNet-20 models pruned to various sparsities with SAFE and ADMM on CIFAR-10C (due to limited time frame) and adversarial samples created using $l_\\infty-PGD$ attack following [2], respectively.\n\n`CIFAR-10C (intensity=3)`\n| method | 90% | 95% | 98% | 99% | 99.5% |\n|:-|-:|-:|-:|-:|-:|\n| ADMM  | 70.06±0.03 | 68.92±0.26 | 65.45±0.29 | 59.19±0.47 | 55.71±0.45 |\n| SAFE  | **73.98±0.09** | **72.92±0.41** | **68.20±0.47** | **66.02±0.56** | **56.58±0.36** |\n\n\n`Adversarial noise (l∞-PGD)`\n| method| 90% | 95% | 98% | 99% | 99.5% |\n|:-|-:|-:|-:|-:|-:|\n| ADMM | 49.81±1.02  | 49.84±1.78  | 43.33±1.59  | 30.29±0.64  | 23.25±1.92  |\n| SAFE | **56.43±1.03** | **51.40±0.89** | **43.34±0.90** | **43.80±1.27** | **29.48±0.68** |\n\nSimilar to label-noise experiments, SAFE consistently outperforms the baseline ADMM under these corruptions, indicating strong robustness of sharpness minimization toward various types of noise as suggested in [1,2]. We will include these results in the revised version.\n\n---\n\n**Wall clock time**\n\n> Could you provide further details about what leads to the high runtime of your algorithm\n\nThank you for your question. We analyze the main source of SAFE’s runtime overhead from both theoretical and engineering perspectives.\n\nFrom a theoretical perspective, we present the time complexity for pruning a single transformer block : $d$ = hidden dim, $N/b$ = # of data / batch size, $k$ = # of iterations,  $L$ = # of layers per block.\n\n| Method | Time Complexity | Explanation |\n|-|-|-|\n| SparseGPT | $O(L(Nd^2 + d^3))$ | Hessian over $N$ samples ($Nd^2$) + inverse ($d^3$) |\n| Wanda | $O(L(Nd + d^2))$ | Activation norm ($Nd$) + weight multiplication($d^2$) |\n| ALPS | $O(L(Nd^2 + d^3 + kd^3))$ | Hessian over $N$ samples ($Nd^2$) + eigendecomposition ($d^3$) + penalized inverse for $k$ iterations ($kd^3$) |\n| SAFE | $O(Lbkd^2)$ | Backprop through $L$ layers for $k$ iterations |\n\nUnlike one-shot methods (e.g., SparseGPT, Wanda), iterative methods (SAFE, ALPS) incur runtime overhead as iterations($k$) increase. However, SparseGPT scales cubically due to Hessian inversion, while SAFE’s quadratic complexity offers better efficiency. ALPS also suffers from poor scalability due to second-order computation, making SAFE more practical for large models.\n\nFrom an engineering perspective, CPU offloading was introduced for memory efficiency, but it also contributed to runtime overhead.\n\n> …any ways to reduce this (runtime)?\n\nTo improve SAFE’s runtime, we can consider reducing: (1) iterations $k$, and (2) per-iteration cost.\n\n- **Fewer iterations**: Prior work has proposed ways to accelerate ADMM, such as adaptive penalties [3,4], over-relaxation [4,5], preconditioning [6], and variance reduction [7].\n\n- **Cheaper iteration**: Following the reviewer’s suggestion, SAFE can adopt efficient SAM variants [8,9] and further engineering optimizations.\n\nWe applied one such optimization by moving data to GPU memory, reducing per-iteration cost. (see Appendix B.3 and C.1 for configuration details)\n\n| Config | Runtime (s) |\n|-|-|\n| CPU offloading | 460 |\n| GPU | 277 |\n\nRemoving CPU offloading yielded a 40% speedup, with an additional ~8Gb GPU memory usage under SAFE’s configuration. This reflects a speed-memory tradeoff. Note that runtime differences may result from varying hardware. \n\nWhile we have not yet incorporated algorithmic acceleration techniques (e.g., ADMM improvements or SAM variants), we believe they hold promise for further reducing runtime, and we plan to explore them in future work.\n\n---\n\n**typo (Line 370)**\n\nThank you for pointing it out. We will have them fixed in the revised paper.\n\n---\n\n**References** \\\n[1] Wei et al., Sharpness-Aware Minimization Alone can Improve Adversarial Robustness., ICML 2023 workshop. \\\n[2] Zhang et al., On the Duality Between Sharpness-Aware Minimization and Adversarial Training., ICML 2024 \\\n[3]  Yuan et al., Admm For Nonconvex Optimization Under Minimal Continuity Assumption, ICLR 2025 \\\n[4] Xu et al., Adaptive Relaxed ADMM: Convergence Theory and Practical Implementation, CVPR 2017 \\\n[5] Song et al., Optimizing ADMM and Over-Relaxed ADMM Parameters for Linear Quadratic Problems, AAAI 2024 \\\n[6] Ali et al., A Semismooth Newton Method for Fast, Generic Convex Programming, ICML 2017 \\\n[7] Liu et al., Accelerated Variance Reduction Stochastic ADMM for Large-Scale Machine Learning., IEEE TPAMI 2020  \\\n[8] Du et al., Sharpness-Aware Training for Free, NeurIPS 2022 \\\n[9] Liu et al., Towards Efficient and Scalable Sharpness-Aware Minimization, CVPR 2022", "claims_and_evidence": "Yes. The authors demonstrate that their proposed algorithm does indeed lead to flatter and sparser solutions than standard training.", "methods_and_evaluation_criteria": "Yes, the authors compare their method to relevant baselines and recent methods. They compare across a variety of tasks and architectures to comprehensively demonstrate the strong performance of their method. I would be interested in seeing the performance of their method under different types of noisy data (ie image corruptions such as ImageNet-C or adversarial robustness) in addition to label noise but I do not think this is necessary to make their point clear, especially if these experiments cannot be run during the rebuttal period.", "theoretical_claims": "The theoretical analysis makes sense to me and I did not see any issues.", "experimental_designs_or_analyses": "The experimental design seems valid to me and I did not see any issues.", "relation_to_broader_scientific_literature": "This work builds upon a few previous papers in the space that motivate a need for flat and sparse neural networks. This paper proposes a theoretically principled method that also performs well on several tasks and outperforms existing methods.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "This paper explains the method very well, including relevant information in the main paper with necessary supplementary details in the appendix. I appreciate the clarity and exposition of the method and the comprehensiveness of the results. There are no critical weaknesses; I would like to see additional robustness experiments but understand the time constraint, and I would appreciate a further discussion of the complexity of the method but neither of these are necessary for completeness.", "other_comments_or_suggestions": "- line ~370 in the right column has a missing table reference (‘see ??’)", "questions_for_authors": "Here I will restate the two main questions I would like to see addressed\n\n1. Does this method additionally improve robustness to corruptions (ie ImageNet-C) or adversarial examples?\n2. Could you provide further details about what leads to the high runtime of your algorithm and whether there are any ways to reduce this (similar to some of the works building on top of SAM to reduce its complexity via only performing the update every n steps)? SAM pointers: Sharpness-Aware Training for Free by Du et al; Towards Efficient and Scalable Sharpness-Aware Minimization by Liu et al; An Adaptive Policy to Employ Sharpness-Aware Minimization by Jiang et al.", "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": "93IT1uGcoF", "reviewer": "Reviewer_h6df", "summary": "The authors tackle the problem of performance degradation in sparse neural networks at very high sparsity levels due to reduced capacity. The proposed method is motivated by the recent discoveries connecting the flatness of the solution (loss) landscape and the improved generalization performance of neural networks. \n\nThe authors frame the problem as an optimization problem that is aware of the sharpness of the loss landscape, and sparsity is used as a constraint. The set-up follows the standard SAM set-up [Foret2021], that is a min max problem for the loss function and a point in its perturbed surrounding. This causes the optimization process to explore a more flatter minima. An augmented Lagrangian approach is utilized that provides the smoothness of the Lagrangian for constrained optimization and allows for sparsity to be constrained as a projection function minimization. Finally, update equations are derived for each of the parameters. \n\nThe authors also introduce generalized projection that modifies the distance measure between the model weights and the projected weights. This allows for more sophisticated functions to be utilized as the pruning saliency metric. The results show that the proposed method SAFE (+) results in flatter minima as compared to ADMM. Further, SAFE results in better performance as compared to a wide range of baseline methods on tasks from two domains. Finally, the experiments also show that the proposed method is robust to label noise and one-shot sparsification to a higher sparsity ratio.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We really appreciate the reviewer’s insightful comments and constructive feedback. While we address specific comments below, we would be keen to engage in any further discussions.\n\n---\n\n**Extended sharpness measurements**\n\nThank you for the insightful suggestion. We extend our sharpness analysis by computing the largest eigenvalue of the Hessian matrix in vision tasks. We use ResNet-20/32 on CIFAR-10/100 and compute the Hessian of the cross-entropy loss. The resulting sharpness is compared to that of the ADMM baseline.\n\n| Method | ResNet-20 / CIFAR-10 | ResNet-32 / CIFAR-100 |\n|-|-|-|\n| ADMM | 72.17 | 103.20 |\n| SAFE | **33.02**| **50.57** |\n\nBuilding upon the sharpness analysis in Section 4.1, we observe that SAFE continues to produce flatter solutions compared to ADMM when scaling up to deeper convolutional networks such as ResNet-20/32. These results support our claim that SAFE consistently finds sparse and flat solutions. We focused on vision architectures due to the rebuttal time limit and will expand the analysis in the revised version.\n\n---\n\n**Comparison with other flat pruner**\n\nThank you for the valuable suggestion. In response to the reviewer's comment regarding the need to compare with objectives that encourage flat minima, we implemented the method proposed by Na et al., [1], which integrates Sharpness-Aware-Minimization (SAM) into Iterative Magnitude Pruning (referred to as SAM+IMP). We applied SAM+IMP to both vision and LLM pruning tasks under our experimental setup, and compared its performance with that of SAFE. \n\n\nFor both vision and LLM tasks, we apply SAFE and SAM+IMP using the same training budgets and hyperparameters as reported in the paper, to ensure a fair comparison. Pruning is performed periodically (every 10 epochs for vision / every 5 epochs for LLM) for SAM+IMP, with sparsity increasing linearly over time. We also consider a cubic sparsity schedule for SAM+IMP in the vision setting. In the LLM task, both methods employ the block-wise reconstruction error objective introduced in our work.\n\n**Vision**\n| Method | 95% | 98% | 99% | 99.5% |\n|-|-|-|-|-|\n| SAM+IMP (linear) | 80.30 ± 0.12  | 36.03 ± 4.19  | 18.30 ± 2.80  | 13.80 ± 0.52  |\n| SAM+IMP (cubic) | 92.50 ± 0.05  | 89.24 ± 0.06  | 83.74 ± 0.14  | 73.73 ± 0.30  |\n| SAFE | **92.59 ± 0.07**  | **89.58 ± 0.10**  | **87.47 ± 0.07**  | **79.55 ± 0.13**  |\n\n**LLM**\n| Method | PPL (C4) | PPL (WikiText2) |\n|-|-|-|\n| SAM+IMP | 18.27 | 176.00 |\n| SAFE | **8.91** | **6.79** |\n\nAs shown in the results above, SAFE consistently outperforms SAM+IMP across vision and LLM tasks. We conjecture that the pruning operation in SAM+IMP is hard and abrupt, which may hinder effective sharpness minimization. In contrast, SAFE introduces a smooth penalization scheme via the augmented lagrangian, enabling joint optimization of sharpness and sparsity in a more balanced manner—leading to superior performance.\n\n---\n\n**Computational efficiency of SAFE**\n\nThank you for raising this concern. As the reviewer mentioned, ADMM-based approaches such as SAFE and ALPS introduce auxiliary variables ($z,u$), which adds to the overall operation count. However, unlike ALPS, SAFE updates these variables only once every $k$ iterations (we use $k=32$ in the paper), resulting in minimal overhead–less than a 3% increase in total operations. \n\nMoreover, the update cost of $z$ and $u$ is negligible compared to the weight update $w$. Specifically, $z$ requires $O(dlogd)$ for sorting, $u$ involves only element-wise addition $O(d)$, whereas updating $w$ via backpropagation incurs $O(d^2)$ complexity.\n\nWe have also provided an in-depth time complexity comparison of various pruning algorithms in our response to Reviewer `FtwJ` . While SAFE has higher time complexity than one-shot methods due to its iterative nature, it generally achieves better pruning performance. This reflects a meaningful cost-performance tradeoff. We will include this efficiency analysis in the revised version.\n\n---\n\n**Others**\n\n> “The authors report the validation accuracies in the case of classification tasks. Are those accuracies computed at the last iteration ? or are those used during BNT?”\n\nAll the results in the paper for classification task report validation accuracy after final projection followed by batch-normalization training.\n\n> “[Peste2022] can be merged with Figure 2 for a better consolidation.”\n\nThank you for the suggestion. We will incorporate this suggestion in the revised version.\n\n---\n\n**References** \\\n[1] Na et al., Train Flat, Then Compress: Sharpness-Aware Minimization Learns More Compressible Models, EMNLP 2022 findings", "claims_and_evidence": "The claims made are very well supported by the clear and concrete results. \n\nIt would be also great to see the sharpness analysis using the hessian for all the other networks and tasks considered.", "methods_and_evaluation_criteria": "The proposed method, along with the motivation and derivations are sound. The evaluation is also extensive. However, the authors should consider comparison with additional methods that are closely related and have a similar objective of finding flatter minima to establish better performance over other methods with a similar objective. Finally, the results for [Peste2022] can be merged with Figure 2 for a better consolidation.", "theoretical_claims": "The derivations made in the main body of the paper are sound. I was unable to spend a lot of time on the proofs presented in the Appendix section.", "experimental_designs_or_analyses": "The experimental designs and analysis are well thought out and show clear results.", "relation_to_broader_scientific_literature": "The paper builds upon the idea of achieving flatter minima for better generalization. This has been explored extensively in the case of dense networks and more recently for neural network sparsification as well. The major contribution of the work is in the set-up of the optimization process for both obtaining a flatter minima and constraining for sparsity.", "essential_references_not_discussed": "The paper discusses all the relevant ideas and methods from prior work in detail.", "other_strengths_and_weaknesses": "Strengths: \n\n* The paper is very well written with clear motivation, method description and results. \n\nPlease see the sections above for additional strengths. \n\nWeaknesses:\n\n* One major concern is regarding the computational efficiency of the proposed method. Each iteration requires computation of multiple variables, that adds to the operations count. I would ask the authors to please discuss this in the final version of the paper. \n\nPlease see the sections above for other weaknesses.", "other_comments_or_suggestions": "Minor: Figure 2b is missing other baselines.", "questions_for_authors": "* The authors report the validation accuracies in the case of classification tasks. Are those accuracies computed at the last iteration ? or are those used during BNT?\n\n## Update after rebuttal\n\nI thank the authors for their response and for generating the requested results. Although I believe the paper makes several contributions, I will stick with my current rating as the rebuttal response does not make substantial improvements.", "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": "v5tpOF5MtG", "reviewer": "Reviewer_gjwA", "summary": "This paper introduces SAFE and SAFE+ algorithms for finding sparse subnetworks with flat loss landscapes. The methods utilize an augmented Lagrange dual approach, with SAFE+ extending the base algorithm through a generalized projection operation. The authors evaluate their approaches on image classification and language modeling tasks, demonstrating that the resulting sparse networks achieve better generalization performance.\n\n## Update after rebuttal\nI appreciate the author's acknowledgment of my concerns in the future work section. However, since the current version does not directly address these points, I will maintain my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer for taking the time to review our work and providing us with constructive feedback. While we respond to the reviewer’s specific comments below, we would be keen to engage in any further discussion.\n\n---\n\n**No sparsity ratio in our convergence analysis?**\n\nWe appreciate this interesting question. In our algorithm, we operate by putting the iteration within the sparsity constraint through a tractable projection operation regardless of the sparsity ratio, so “whether convergence occurs at the limit of the algorithm” is guaranteed regardless of the sparsity ratio. However, the convergence *rate* may indeed be affected by the sparsity ratio, as certain factors related to it can influence the actual speed of convergence. This has been studied in the constrained optimization literature [1, 2, 3]. We believe exploring this aspect would be a valuable direction for future work to deepen our understanding of SAFE.\n\n\n\n---\n\n**Theoretical analysis or empirical evidence of solution flatness**\n\nWe would first like to clarify that we already provide empirical evidence that SAFE yields flatter solutions through Hessian eigenvalue measurements and loss landscape visualizations in Figure 1(c-d).\n\nHowever, we agree that a theoretical analysis of the flatness of the solution found by SAFE would be a valuable addition to the paper. It is well-established that the sharpness minimization used in SAFE yields flatter solutions [4-7]. Among these analyses, we plan to build on linear stability analysis [7] to derive an upper bound on the Hessian at the limit point of SAFE iterates, which would allow us to characterize the flatness of SAFE’s solution.\n\nWe plan to explore in this direction and will aim to include it in the final version.\n\n---\n\n**Theoretical difference between SAFE and CrAM**\n\nThank you for the interesting question. While SAFE aims to find sparse and flat solution, CrAM aims to obtain a “compressible (dense) model” by enforcing a model to maintain a strong post-compression performance even with small perturbations:\n\n$$\\min_x \\max_{\\|\\epsilon\\| \\leq \\rho}  f( C(x+\\epsilon)  ).$$\n\nBased on this objective, CrAM updates the parameters as follows:\n\n$$x_{t+1} = x_t - \\eta\\nabla f( C(x+\\rho\\nabla f(x)) ),$$\n\nwhere the gradient is computed at a point obtained by applying a small perturbation followed by a hard compression step. However, this introduces an abrupt discontinuity between the iterate and the point where the gradient is evaluated, potentially disrupting the training process.\n\nOn the contrary, SAFE progressively introduces sparsity into the optimization process via  augmented lagrangian and split variables in the ADMM framework, which may offer a more natural fit for neural network training.\n\n---\n\n\n**References** \\\n[1] Liu et al., One-bit compressive sensing with projected subgradient method under sparsity constraints. IEEE Transactions on Information Theory, 2019 \\\n[2] Yuan et al., Dual Iterative Hard Thresholding, JMLR, 2020 \\\n[3] Yuan et al., Smoothing Proximal Gradient Methods for Nonsmooth Sparsity Constrained Optimization: Optimality Conditions and Global Convergence, ICML, 2024 \\\n[4] Wen et al., How Does Sharpness-Aware Minimization Minimize Sharpness?, ICLR 2023 \\\n[5] Zhou et al., Sharpness-Aware Minimization Efficiently Selects Flatter Minima Late In Training \\\n[6] Agarwala and Dauphin, SAM operates far from home: eigenvalue regularization as a dynamical phenomenon \\\n[7] Shin et al., Critical Influence of Overparameterization on Sharpness-aware Minimization", "claims_and_evidence": "I have two concerns regarding the theoretical analysis in this paper:\n1. While the authors claim that SAFE converges to the stationary point of the sparsity-constrained optimization problem under Assumptions 3.1-3.3, it is unclear why the convergence analysis appears to be independent of the sparsity ratio. This seems counterintuitive given the significant role of sparsity constraints in the optimization problem.\n2. Despite the paper's emphasis on finding flat minima, which is explicitly incorporated into the learning objective, there appears to be no theoretical analysis or empirical evidence demonstrating the flatness of the obtained solutions.", "methods_and_evaluation_criteria": "1. The core methodology employs an augmented Lagrangian dual approach to address sparsity-constrained optimization. Through alternating minimization, this formulation effectively decomposes complex optimization into tractable subproblems, demonstrating a sound implementation.\n2. The SAFE+ variant incorporates an importance matrix P to enhance the sparsification process, enabling more informed pruning decisions by considering parameter significance.\n3. What are the theoretical differences between SAFE and CrAM, and what mechanisms enable SAFE to achieve superior performance?", "theoretical_claims": "While the proof demonstrates convergence to stationary points under specified assumptions, the relationship between convergence properties and different sparsity constraints remains unclear.", "experimental_designs_or_analyses": "The experimental designs are sound.", "relation_to_broader_scientific_literature": "This work investigates neural network sparsification in conjunction with SAM, proposing an augmented Lagrangian approach that offers a theoretically grounded alternative.", "essential_references_not_discussed": "The paper has appropriately cited and discussed the key related works.", "other_strengths_and_weaknesses": "N/A", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "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": "YAWX1DaJaC", "reviewer": "Reviewer_BN7C", "summary": "This work tried to learn sparse and flat minima during training, so such sparse and flat minima can be more suitable for network pruning. The idea is simple but effective. This work formulates pruning as a sparsity-constrained optimization problem where flatness is encouraged as an objective. This optimization problem can be solved explicitly via an augmented Lagrange dual approach, which naturally resulted in novel pruning methods called SAFE and its extension, SAFE+. Convergence analysis is presented. Empirical analysis include vision models and large language models. The improvements are significant.\n\n\n## update after rebuttal\n\nThanks for the responses. It help addressing my concerns. I maintain the rating as Accept.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate the reviewer’s recognition and constructive feedback. We have addressed the reviewer’s specific comments below, while we remain open to any additional suggestions.\n\n\n---\n\n**Convergence for nonconvex case?**\n\nWe appreciate the reviewer’s insight. We believe that the convergence of SAFE without the weak-convexity assumption can potentially be established by extending existing non-convex convergence analyses of sharpness-minimization [1, 2] and ADMM [3, 4] to SAFE’s x-minimization step and its sparsity-constrained alternating dual updates, respectively. We plan to explore this direction and aim to include the results in the revised version.\n\n---\n\n**Lack of efficiency analysis**\n\nWe appreciate the reviewer’s concern regarding the lack of efficiency analysis. \n\nFirst, we already provide wall-clock time comparisons between LLM pruning methods in Appendix C.1.\n\nIn addition, we provide an in-depth time complexity analysis in our response to Reviewer `FtwJ` (“Wall clock time”) and discuss the operational overhead in response to Reviewer `h6df` (“Computational efficiency of SAFE”). While both SAFE and ALPS are iterative methods that introduce additional computation, SAFE is significantly more scalable than ALPS, and this cost is worthwhile trade-off given the favorable improvements in pruning quality.\n\n Lastly, for vision tasks, all baseline methods except MLPrune(a one-shot method) rely on stochastic gradient-based optimization, and therefore share the same asymptotic time complexity as SAFE. While SAFE incurs additional cost compared to GMP and ADMM due to sharpness minimization, this overhead is offset by the consistently strong performance it has shown in our experiments across various tasks and sparsity levels.\nWe will include this analysis into a dedicated cost analysis section in the revised version.\n\n---\n\n**Essential references not discussed**\n\nThank you for bringing this work to our attention. We present detailed discussions below.\n\n[5] presents an extensive empirical study on the generalization performance of various pruning approaches under label noise, revealing a double-desent phenomenon [6] across varying levels of sparsity. Intriguingly, in our label-noise experiments (Table 2), this double-descent behavior isn’t observed in SAFE, whereas it appears prominently in the ADMM baseline, which also consistently shows poor performance under label noise. We suspect this to be attributed to the strong regularization of sharpness minimization in mitigating label noise, as such effective regularization is known to mitigate double descent.\n\nAdditionally, [5] makes a fascinating observation that higher sparsity tends to yield solutions of higher sharpness, which could potentially support our sharpness-minimization approach to improving pruning.\n\nWe will make sure to include these discussions in the final version of the paper.\n\n\n---\n\n**References** \\\n[1] Khanh et al., Fundamental Convergence Analysis of Sharpness-Aware Minimization, NeurIPS 2024 \\\n[2] Oikonomou and Loizou, Sharpness-Aware Minimization: General Analysis and Improved Rates, ICLR 2025 \\\n[3] Ganzhao Yuan, Admm For Nonconvex Optimization Under Minimal Continuity Assumption, ICLR 2025 \\\n[4] Huang and Chen, Mini-Batch Stochastic ADMMs for Nonconvex Nonsmooth Optimization, CoRR 2018 \\\n[5] He et al., Sparse Double Descent: Where Network Pruning Aggravates Overfitting, ICML 2022 \\\n[6] Nakkiran et al., Deep Double Descent: Where Bigger Models and More Data Hurt, ICLR 2020 \\\n[7] Nakkiran et al., Optimal Regularization Can Mitigate Double Descent, ICLR 2021", "claims_and_evidence": "The claim and evidence are basically convincing.", "methods_and_evaluation_criteria": "Yes. The method and its motivation are clear and make sense. The evaluation criteria is common.", "theoretical_claims": "This work include a few theoretical results, particularly including deriving SAFE and its convergence analysis.\n\nI can basically understand the proof sketch, but I did not carefully check the proof in the appendix.\n\nIn convergence analysis, Assumption 3.3 is a bit strong for deep loss landscape. Is it possible to prove the convergence for nonconvex optimization?", "experimental_designs_or_analyses": "I checked the experimental design and analyses. They make sense to me.", "relation_to_broader_scientific_literature": "This paper make a key contribution to network pruning. Because, to my knowledge, it is the first successful method that can explicitly find sparse and flat minima in the training phase.", "essential_references_not_discussed": "An ICML2022 work Sparse Double Descent reveal some insights on generalization of network pruning, which may be helpful for understanding generalization of sparse minima. \n\n[1] He, Z., Xie, Z., Zhu, Q., & Qin, Z. (2022, June). Sparse double descent: Where network pruning aggravates overfitting. In International Conference on Machine Learning (pp. 8635-8659). PMLR.", "other_strengths_and_weaknesses": "Strengths:\n\n- SAFE that explicitly constrain sparsity and encourage flatness is novel to me.\n- I appreciate the theoretical results for SAFE, while some minor concern exists.\n- The empirical analysis is convincing.\n\n\nWeaknesses:\n- This work lacks efficiency analysis. Is SAFE computationally expensive compared with the baselines?\n- In convergence analysis, Assumption 3.3 is a bit strong for deep loss landscape. Is it possible to prove the convergence for nonconvex optimization?", "other_comments_or_suggestions": "No.", "questions_for_authors": "Please see the weakness.", "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:

Meta-review: This paper introduces SAFE and its extension SAFE+, algorithms designed to jointly promote sparsity and flatness in neural networks. The methods are motivated by the strong empirical link between flatter minima and better generalization. Built upon an augmented Lagrangian dual formulation, the approach explicitly enforces sparsity constraints while incorporating sharpness-aware objectives, leading to the construction of sparse subnetworks that generalize well. Extensive empirical evaluations and theoretical grounding are provided.
The reviewers largely agree that this is a strong paper with theoretical depth, practical impact, and sound experimentation. While there are some open questions regarding efficiency and deeper flatness theory, the core idea is novel and well executed. The rebuttal was comprehensive and addressed reviewer concerns effectively, leading all reviewers to maintain or confirm high scores. Therefore, the AC suggests acceptance and encourages the authors to include the rebuttal discussion in the final version of the paper.
Decision: Accept (spotlight poster)