SAFE: Finding Sparse and Flat Minima to Improve Pruning

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.

Convergence for nonconvex case?

We 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.

Lack of efficiency analysis

We appreciate the reviewer’s concern regarding the lack of efficiency analysis.

First, we already provide wall-clock time comparisons between LLM pruning methods in Appendix C.1.

In 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.

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. We will include this analysis into a dedicated cost analysis section in the revised version.

Essential references not discussed

Thank you for bringing this work to our attention. We present detailed discussions below.

[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.

Additionally, [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.

We will make sure to include these discussions in the final version of the paper.

References
[1] Khanh et al., Fundamental Convergence Analysis of Sharpness-Aware Minimization, NeurIPS 2024
[2] Oikonomou and Loizou, Sharpness-Aware Minimization: General Analysis and Improved Rates, ICLR 2025
[3] Ganzhao Yuan, Admm For Nonconvex Optimization Under Minimal Continuity Assumption, ICLR 2025
[4] Huang and Chen, Mini-Batch Stochastic ADMMs for Nonconvex Nonsmooth Optimization, CoRR 2018
[5] He et al., Sparse Double Descent: Where Network Pruning Aggravates Overfitting, ICML 2022
[6] Nakkiran et al., Deep Double Descent: Where Bigger Models and More Data Hurt, ICLR 2020
[7] Nakkiran et al., Optimal Regularization Can Mitigate Double Descent, ICLR 2021

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.

No sparsity ratio in our convergence analysis?

We 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.

Theoretical analysis or empirical evidence of solution flatness

We 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).

However, 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.

We plan to explore in this direction and will aim to include it in the final version.

Theoretical difference between SAFE and CrAM

Thank 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:

$\min_x \max_{\|\epsilon\| \leq \rho} f( C(x+\epsilon) ).$

Based on this objective, CrAM updates the parameters as follows:

$x_{t+1} = x_t - \eta\nabla f( C(x+\rho\nabla f(x)) ),$

where 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.

On 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.

References
[1] Liu et al., One-bit compressive sensing with projected subgradient method under sparsity constraints. IEEE Transactions on Information Theory, 2019
[2] Yuan et al., Dual Iterative Hard Thresholding, JMLR, 2020
[3] Yuan et al., Smoothing Proximal Gradient Methods for Nonsmooth Sparsity Constrained Optimization: Optimality Conditions and Global Convergence, ICML, 2024
[4] Wen et al., How Does Sharpness-Aware Minimization Minimize Sharpness?, ICLR 2023
[5] Zhou et al., Sharpness-Aware Minimization Efficiently Selects Flatter Minima Late In Training
[6] Agarwala and Dauphin, SAM operates far from home: eigenvalue regularization as a dynamical phenomenon
[7] Shin et al., Critical Influence of Overparameterization on Sharpness-aware Minimization

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.

Extended sharpness measurements

Thank 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.

Building 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.

Comparison with other flat pruner

Thank 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.

For 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.

Vision

LLM

As 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.

Computational efficiency of SAFE

Thank 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.

Moreover, 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.

We 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.

Others

“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?”

All the results in the paper for classification task report validation accuracy after final projection followed by batch-normalization training.

“[Peste2022] can be merged with Figure 2 for a better consolidation.”

Thank you for the suggestion. We will incorporate this suggestion in the revised version.

References
[1] Na et al., Train Flat, Then Compress: Sharpness-Aware Minimization Learns More Compressible Models, EMNLP 2022 findings

We sincerely thank the reviewer for finding our work interesting and giving us constructive suggestions. We make clarifications to specific comments below.

Robustness to different types of noisy data

Thank 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.

CIFAR-10C (intensity=3)

Adversarial noise (l∞-PGD)

Similar 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.

Wall clock time

Could you provide further details about what leads to the high runtime of your algorithm

Thank you for your question. We analyze the main source of SAFE’s runtime overhead from both theoretical and engineering perspectives.

From 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.

Unlike 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.

From an engineering perspective, CPU offloading was introduced for memory efficiency, but it also contributed to runtime overhead.

…any ways to reduce this (runtime)?

To improve SAFE’s runtime, we can consider reducing: (1) iterations $k$, and (2) per-iteration cost.

• 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].

• Cheaper iteration: Following the reviewer’s suggestion, SAFE can adopt efficient SAM variants [8,9] and further engineering optimizations.

We applied one such optimization by moving data to GPU memory, reducing per-iteration cost. (see Appendix B.3 and C.1 for configuration details)

Removing 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.

While 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.

typo (Line 370)

Thank you for pointing it out. We will have them fixed in the revised paper.

References
[1] Wei et al., Sharpness-Aware Minimization Alone can Improve Adversarial Robustness., ICML 2023 workshop.
[2] Zhang et al., On the Duality Between Sharpness-Aware Minimization and Adversarial Training., ICML 2024
[3] Yuan et al., Admm For Nonconvex Optimization Under Minimal Continuity Assumption, ICLR 2025
[4] Xu et al., Adaptive Relaxed ADMM: Convergence Theory and Practical Implementation, CVPR 2017
[5] Song et al., Optimizing ADMM and Over-Relaxed ADMM Parameters for Linear Quadratic Problems, AAAI 2024
[6] Ali et al., A Semismooth Newton Method for Fast, Generic Convex Programming, ICML 2017
[7] Liu et al., Accelerated Variance Reduction Stochastic ADMM for Large-Scale Machine Learning., IEEE TPAMI 2020
[8] Du et al., Sharpness-Aware Training for Free, NeurIPS 2022
[9] Liu et al., Towards Efficient and Scalable Sharpness-Aware Minimization, CVPR 2022