The Iterative Optimal Brain Surgeon: Faster Sparse Recovery by Leveraging Second-Order Information

The rising footprint of machine learning has led to a focus on imposing \emph{model sparsity} as a means of reducing computational and memory costs. For deep neural networks (DNNs), the state-of-the-art accuracy-vs-sparsity is achieved by heuristics inspired by the classical Optimal Brain Surgeon (OBS) framework~\citep{lecun90brain, hassibi1992second, hassibi1993optimal}, which leverages loss curvature information to make better pruning decisions. Yet, these results still lack a solid theoretical understanding, and it is unclear whether they can be improved by leveraging connections to the wealth of work on sparse recovery algorithms. In this paper, we draw new connections between these two areas and present new sparse recovery algorithms inspired by the OBS framework that comes with theoretical guarantees under reasonable assumptions and have strong practical performance. Specifically, our work starts from the observation that we can leverage curvature information in OBS-like fashion upon the projection step of classic iterative sparse recovery algorithms such as IHT. We show for the first time that this leads both to improved convergence bounds under standard assumptions. Furthermore, we present extensions of this approach to the practical task of obtaining accurate sparse DNNs, and validate it experimentally at scale for Transformer-based models on vision and language tasks.

Paper

References (45)

Scroll for more · 33 remaining

Similar papers

Peer review

Reviewer 3y7g7/10 · confidence 4/52024-07-10

Summary

This paper proposes a theoretically convergent iterative Optimal Brain Surgeon (OBS) algorithm, which generalizes the classic Iterative Hard Thresholding (IHT)-based algorithms by incorporating approximate second-order information in the sparse projection step. The author also provides practical variants of these algorithms for solving sparse linear regression and model pruning. The experiments show that the proposed algorithm leads to faster convergence than traditional algorithms and improves accuracy.

Strengths

The proposed algorithm is simple yet effective.

Weaknesses

Theoretical analysis does not hold in general for practical problems.

Questions

Detailed comments are as follows: 1. The author elaborates on the theoretically convergent I-OBS method and provides specific calculations for both practical and theoretical schemes. What is the gap between the two schemes? Can it be verified through experiments or illustrated with examples? What is the specific operator T in the practical scheme? 2. In lines 158 and 169, the specific form of the model \( \phi_t(\theta) \) differs. 3. In the model pruning experiment, the author provides the specific algorithm Alg. 2 for the practical use of I-OBS. By constructing a constrained optimization problem for the learnable parameters \( W \) of each layer to fine-tune the model, the author has verified the method's feasibility through experiments. Is it theoretically feasible?

Rating

7

Confidence

4

Soundness

3

Presentation

4

Contribution

3

Limitations

The author objectively analyzed the limitations of this article.

Reviewer T9TW6/10 · confidence 3/52024-07-11

Summary

This work combines second-order curvature information with sparse recovery algorithms to demonstrate, both theoretically and empirically, that the curvature information leads to improved convergence rates and generalization performance in post-training iterative pruning and sparse recovery

Strengths

* This work offers a rigorous analysis of the proposed method, I-OBS, complete with detailed derivations and proofs. * The authors ground their contributions within the context of the existing literature but clearly identify how it differs from prior works by leveraging second order information. * Despite the intractability of the theoretical approach, the authors offer a practical formulation that yields improved performance across a diverse range of data domains. * The paper is rather dense; however, the authors provide all the necessary notation required to follow their derivations. * This is a well motivated line of inquiry as post-training compression is increasingly important in the era of LLMs. Further, despite requiring additional iterations compared to one-shot pruning methods, I-OBS converges in a small number of iterations (<100 for DeiT and <3 for LLMs).

Weaknesses

* Grounding this work in some practical considerations would improve the overall impact of the paper and may help the reader assess whether this technique is suitable for a desired use case. For instance, adding some overall runtime characteristics for the empirical results would establish an order-of-magnitude estimate of the overall computational requirements. * In a similar vein, conducting experiments with fine-grained sparsity such as 2:4 would be a nice extension to produce sparse neural networks that can actually be accelerated in practice. As it currently stands, the unstructured networks learned by I-OBS do not offer much in the way of an immediate practical application. * For the LLM experiments, perplexity has been shown to be a somewhat misleading metric when evaluating compressed LLMs [1]. Ideally, the LLM experiments should be evaluated on downstream tasks such as GLUE or better yet the LLM-KICK benchmark [1]. * Several post-training compression schemes have been proposed recently. Comparisons with methods such as [2-6] would improve my confidence in the significance of this work. * Several small typos, see suggestions below. [1] A. Jaiswal, Z. Gan, X. Du, B. Zhang, Z. Wang, and Y. Yang, “Compressing LLMs: The Truth is Rarely Pure and Never Simple.” arXiv, Oct. 02, 2023. doi: 10.48550/arXiv.2310.01382. [2] T. Dettmers et al., “SpQR: A Sparse-Quantized Representation for Near-Lossless LLM Weight Compression.” arXiv, Jun. 05, 2023. [Online]. Available: http://arxiv.org/abs/2306.03078 [3] M. Sun, Z. Liu, A. Bair, and J. Z. Kolter, “A Simple and Effective Pruning Approach for Large Language Models.” arXiv, Jun. 20, 2023. doi: 10.48550/arXiv.2306.11695. [4] Y. Zhang, H. Bai, H. Lin, J. Zhao, L. Hou, and C. V. Cannistraci, “Plug-and-Play: An Efficient Post-training Pruning Method for Large Language Models,” presented at the The Twelfth International Conference on Learning Representations, Oct. 2023. [Online]. Available: https://openreview.net/forum?id=Tr0lPx9woF [5] Y. Zhang et al., “Dynamic Sparse No Training: Training-Free Fine-tuning for Sparse LLMs.” arXiv, Oct. 17, 2023. doi: 10.48550/arXiv.2310.08915. [6] Y. Ma et al., “AffineQuant: Affine Transformation Quantization for Large Language Models.” arXiv, Mar. 19, 2024. doi: 10.48550/arXiv.2403.12544.

Questions

## Questions * In terms of wall clock time, what was the duration required for the pruning results in Tables 1 and 2 and what hardware configuration was used? * How do the compressed LLMs compare on downstream tasks such as LLM-KICK or GLUE? * Can I-OBS be extended to fine-grained sparsity types such as N:M sparsity? * Are the results in Table 2 obtained with 32- or 16-bit weights? How does a quantized dense model with half the precision as used in Table 2 compare? ## Suggestions * Use vector graphic formats for Figure 1. * Potential typos to address: * L281 & Figure 1 Caption: Refers to topk-WoodFisher, I believe this should be **topk-I-OBS**? Perhaps a prior naming convention? * Table 2 caption: Phi-1.5M -> Phi-1.5B * L332: pwer-layer -> per-layer * L340: mode -> model

Rating

6

Confidence

3

Soundness

3

Presentation

4

Contribution

2

Limitations

In general the limitations listed are appropriate. However, I encourage the authors to add some discussion about the practical requirements for their algorithm in terms of time and compute.

Reviewer Dt8W6/10 · confidence 2/52024-07-12

Summary

The paper presents a new family of algorithms called Iterative Optimal Brain Surgeon (I-OBS), extending the post-training Optimal Brain Surgeon (OBS) framework to an iterative setting commonly used in sparse recovery. I-OBS algorithms utilize second-order information during the sparse projection step, enhancing convergence guarantees compared to classic algorithms like Iterative Hard Thresholding (IHT). Contributions of the paper: * Introduction of I-OBS, which improves classic IHT-based algorithms by incorporating approximate second-order information in the sparse projection step. * Provision of faster analytical convergence rates for I-OBS under standard first- and second-order smoothness and strong convexity assumptions. It also offers theoretical guarantees for existing practical pruning algorithms such as WoodFisher and OBC. * Development of practical versions of the I-OBS algorithms that relax theoretical constraints, making them easy to implement and scalable for large problem instances, such as compressing vision and language models.

Strengths

* The paper introduces an algorithm with theoretical guarantees that its predecessors do not provide. * It also offers a practical computational algorithm that approximates the theoretical one. * The paper shows the benefit of the proposed algorithm in training sparse linear regression. * The paper shows empirically that the proposed algorithm is applicable to prune large models in experiments and obtain promising results.

Weaknesses

* The paper does not analyze the differences between the practical version of the algorithm and the theoretical one. * Since the algorithm uses second-order information, it may have time complexity issues. * The paper only compares the proposed method with SparseGPT. Hence, other existing methods for pruning can be also added.

Questions

* How does the algorithm using second-order information compare with one using first-order information? * Can you provide experiments comparing training time in the pruning process among the proposed method and existing works? * How does the proposed algorithm apply to the Convolution network?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations.

Reviewer JoZw7/10 · confidence 4/52024-07-12

Summary

Having clarified my concerns in their rebuttal, I have updated my score for acceptance. ---- This paper presents a variant of the classic Optimal Brain Surgeon (OBS) method to iteratively prune multiple weights of a neural network at once, with each step consisting of obtaining a pruning mask for the fraction of weights to be removed and then updating the remaining weights. The update step is framed as a sparse recovery problem to be solved using second-order methods based on the Hessian of the loss function. Two approaches are proposed for the selection of weights to prune as well as for the update of the remaining weights, with one having stronger guarantees and the other being more manageable in practice.

Strengths

In my opinion, the core contribution here is the idea of iteratively changing the pruning mask and updating the remaining weights, which implies that the decision about which weights to prune is revisited at every step. This is a subtle departure from how this is typically done with post-training pruning (i.e., prune once, update once; or prune a smaller amount, update, and then repeat by pruning more). Moreover, I appreciate that the authors discuss the two steps by providing methods that either have better guarantees (Option 1) or that are more tractable in practice (Option 2). The connection with sparse recovery is also very interesting, although I wished it was made more explicit. After one paragraph in Section 2, this connection is leveraged for their algorithm but the authors do not provide greater insight on the algorithms used in sparse recovering.

Weaknesses

The authors imply that their work is the first generalization of OBS for simultaneously pruning multiple weights. However, this was already explored before in an ICML 2022 paper as the Combinatorial Brain Surgeon (CBS) [1]. In that paper, the method is also broken down in two steps, the first step selecting which weights to prune (CBS-S), which is equivalent to selecting the support mask in I-OBS, and the second step determining how to update the remaining weights (CBS-U), which is equivalent to optimizing the parameters in I-OBS. Because of the intractability of the problem, CBS-S is also approached in a greedy way, similar to Top$k$ in I-OBS. The main difference would be that CBS-S uses a semi-greedy heuristic to vary the selection of weights, whereas I-OBS updates the weights after fixing a mask and consequently affects which weights would be chosen by the greedy selection in the next iteration. [1] https://arxiv.org/abs/2203.04466 More generally, the authors only cite fours papers about network pruning in the last decade in their literature review, one of those being a survey and another two being from a same group. For the rest of the paper, they just refer back to these two papers from the same group as the state-of-the-art and do not benchmark with anything else. There has been several papers focused on OBS in the last decade, both before and following [1]. Moreover, there are other pruning methodologies with theoretical guarantees, such as those based on coresets. The connection with sparse recovery is quite interesting, but it generalizes the algorithm much more than it would be need for OBS. Because the objective function is quadratic, my understanding is that the weights has a closed-form solution (see discussion in the WoodFisher paper, in [1] regarding CBS-U, and also in lines 211-212 of I-OBS). Finally, the authors only present results for sparsity 50%, which is not a lot of pruning in practice. For that amount of pruning, magnitude pruning (preserving the half of the weights with largest absolute value) is competitive with more refined techniques, such as WoodFisher and CBS. Because the weight selection done by the authors with Top$k$ is basically magnitude pruning, the results do not provide evidence of the strength of the technique. The technique may be strong, but that needs to be shown in more detail. Other minor comments about the writing: 6-7: "lack a solid theoretical understanding": this is a bold statement, perhaps a bit unfair, and I don't think that the theory provided in the paper addresses exactly that (you just prove convergence rates for an update based on an optimization algorithm) Equations 7 and 8: What is $H_t$ used as a subindex? 201: "critierion" -> criterion 201: "[,] which recovers" 210: $X$ not in bold as in previous uses 215: "in the first step one": remove one? 216 and 217: it is not "bruteforcing" if what you are doing takes linear time; bruteforcing is solving an entire problem by exhaustion 222: "As observed Algorithm 1": please rewrite 251, 278, 304-305, 311, 346: these references should have been in parantheses rather than having the authors names mentioned in the sentence without proper wording for it 256: "call[ed]" 269: "the[o]retical" 282: remove comma before "due" 285: remove "to" Algorithm 2, line 1: remove first "for each layer" and remove comma 339: "guarentee" -> guarantee 340: "mode[l] pruning"? 341: "those assumption": either use "this" or use "assumptions" 342: "is" -> to be

Questions

1) Can you please address the overlap with CBS? 2) Can you please frame your contribution over a broader range of recent work on network pruning, besides the four papers that you cited? 3) Can you please discuss the connection with sparse recovery in more detail, and also the relevance of this connection if having a quadratic objective function as in I-OBS leads to an optimization problem that can solved in closed-form? 4) Can you please benchmark your approach with at least magnitude pruning and CBS? 5) If possible, can you please provide results with higher amounts of sparsity?

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors did a reasonable work discussing limitations and what should be done in future work.

Reviewer Dt8W2024-08-09

Response to authors

Thank you for your detailed response. I will keep my initial rating due to the lack of domain knowledge. I believe Iterative Optimal Brain Surgeon (I-OBS) is a novel contribution and can be widely used. Best regards,

Authorsrebuttal2024-08-12

Discussion gentle reminder

Dear reviewer, As the discussion period will end soon, we wanted to respectfully ask if you could please provide feedback on our responses to your concerns, and specifically on the additional experimental results provided. Best regards,\ The authors

Reviewer T9TW2024-08-12

Thank you for the rebuttal and additional experiments, I believe they will improve the manuscript. I have elected to maintain my original score.

Reviewer JoZw2024-08-13

Follow up

I would like to thank the authors for addressing my concerns and extending their computational evaluation.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC