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

Thank you for the valuable questions and positive comments. We address your questions below:

Q1, Regarding the gap between the practical and theoretical schemes:

We note that for the theoretical scheme we have $|| \theta\_{t+1} - \theta^* ||\_2 \leq \left(1+ \sqrt{\frac{L}{\mu}\frac{d-k^*}{d-k}} \right) \frac{M}{2 \mu} || \theta\_t - \theta^*||\_2^2.$ (Theorem 1) and for the practical scheme we have $|| \theta\_{t+1} - \theta\_* ||\_2 \leq \left( 1+ \sqrt{\frac{k^*}{k}} \right) \frac{M}{2 \mu} || \theta\_t - \theta^* ||\_2^2$ ( Lemma 3). In both cases, the upper is in the form of $|| \theta_{t+1} - \theta_* ||_2 \leq C|| \theta_t - \theta^* ||_2^2$, where $C$ is a constant. Thus, both scheme gives a super-exponential local convergence rate, namely both algorithms achieve $\epsilon$-error with initialization $||\theta_0 - \theta^*||_2 \leq \frac{1}{2 C}$ in $O( \log \log \frac{1}{\epsilon})$ steps. Thus, there is no gap in terms of convergence rate. There is indeed a difference in the convergence radius $\frac{1}{2 C}$ since the constant $C$ is different, but we remark that this difference is minor in practice.

Q2, In lines 158 and 169, the specific form of the model ( $\phi_t (\theta)$ ) differs

We thank the reviewer for the detailed review, however, there are no $\phi_t(\theta)$ in line 169. Do you mean the $\phi_t(\theta)$ in line 164? In line 164, we replace the gradient $\nabla f(\theta_t)$ term with the stochastic gradient term$g_t(\theta_t)$, so that if one solves the proximal problem induced by the $\phi_t(\theta)$ in line 164, one will get the update of stochastic gradient descent rather than gradient descent. Please let us know if there are any further questions regarding this point.

Q3, Regarding the theoretical analysis of the practical implementation of I-OBS

While we believe it would be interesting to obtain a theoretical guarantee for the practical implementation of I-OBS, we find it is difficult to make it end-to-end rigorous, as the underlying constrained optimization problem is NP-hard. In particular, to efficiently solve the constraint optimization problem in line 7 of Algorithm 2, we applied existing heuristic quadratic sparsity solvers, such as OBC or SparseGPT. (In preliminary experiments, we have also investigated applying algorithms with approximation guarantees, such as Natarajan’s algorithm [Natarajan95], but found that these are much less scalable and do not provide better solutions.) These solvers use greedy heuristics to efficiently compute the masks, since obtaining the theoretical optimal masks shown in line 7 of Algorithm 1 is practically infeasible. Despite this approximation, we note that our approach obtains high practical performance.

[Natarajan95] Natarajan, Balas Kausik. "Sparse approximate solutions to linear systems." SIAM journal on computing 24.2 (1995): 227-234.

Thank you for the detailed review, as well as your valuable questions and comments. We address your questions and concerns below

Q1 and W1 Regarding the training time of the pruning results

For the experiments on ViTs (Table 1 in the paper), it took about 4 hours to run 100 iterations and the hardware we used was Nvidia A100 GPU with 80GB memory. For LLMs experiments, we ran I-OBS on Llama-2 (7B) and Llama-3 (8B) models, and it takes around 1.5 hours per iteration (15h for 10 iterations). Unfortunately, we haven't recorded the time of running the experiments on OPT-125M and Phi-1.5 (Table 2 in the paper).

Q2 and W3 Evaluating the performance on LLM-KICK or GLUE

To address your concern and study the scalability of I-OBS, we ran I-OBS on Llama-2 (7B) and Llama-3 (8B) models. Specifically, we apply SparseGPT (50% sparsity) and finetune on the same calibration set used for Hessian estimates. We evaluate the performance on 5-shot MMLU, one of the tasks from LLM-KICK benchmark–we did this due to time constraints but will expand to the full benchmark for the next revision.

Experimental results are provided in Table 1 and Table 2 of the PDF file we attached in the global rebuttal.

One can observe that the quality of sparse solution significantly improves after first finetuning iteration, thus validating our approach. Moreover, there is small improvement during the next few iterations, followed by gradual performance deterioration afterward, which we explain due to overfitting on the calibration set. Specifically, I-OBS improves accuracy on MMLU by >1 point, a 20% reduction in error relative to the original dense model.

W4 Comparison to other works in pruning

Thank you for the suggestion. As also suggested by reviewer JoZw, we provide a more detailed literature review. Please refer the answer to Q2 of reviewer JoZw. We apologize that we can not copy the whole response here due to the character limit.

Also, as suggested by reviewer JoZw, we conducted experiments on MobileNetV1 to compare to the CBS methods. The results are shown in Table 3 of the attached file in global rebuttal. Our experiments imply that, starting with 60% sparsity, our I-OBS pruner outperforms CBS method on MobileNetV1 by a large margin: 2% for 60% sparsity, 5% for 70% sparsity and 12% for 80% sparsity. For low sparsities (30% to 50%), the two methods are comparable, since the accuracy difference is less than 0.5%

Q3 and W1 Extend to fine-grained sparsity type

The practical implementation of I-OBS indeed extends to N:M sparsity. In particular, we apply OBC or SparseGPT as sparsity solvers for the layerwise pruning problem, and OBC or SparseGPT applies to N:M sparsity. However, for the theoretical results, we are unable to extend the analysis to the masks selected based on N:M sparsity, we will leave this to future work.

Q4 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?

The results in Table 2 are obtained with bfloat16 (half precision) weights, which is standard. Generally, we found the weight precision to be orthogonal to our results, as sparsification can be applied independently of the baseline weight representation, and always yields similar speedups (the maximum speedup due to 2:4 sparsity is 2x for INT8 and FP16). More broadly, we believe our iterative approach should be generalizable to compression via quantization as well: specifically, we could modify the projection step to perform quantization rather than sparsification.

Suggestions

Thank you for the suggestions. We will change the figure format and correct the typos in the revision

References [CBS] Yu, X., Serra, T., Ramalingam, S., & Zhe, S.. “The combinatorial brain surgeon: pruning weights that cancel one another in neural networks.” ICML 2022

Thank you for your detailed reviews, as well as the insightful comments and questions. We address your comments and questions below

Q1. CBS assumes the model gradient is zero for post-training pruning, while I-OBS does not, making I-OBS more general. We retain the gradient term because models may not be fully trained and subsequent I-OBS iterations can increase gradient norms for sparsity. Additionally, I-OBS is an iterative algorithm, in contrast to CBS's one-shot approach, allowing CBS-S to be used as a projection step in our iterative process by taking the gradient term into consideration.

The objectives for selecting the masks are also different. Specifically, in the IQP formulation of CBS-S procedure for mask selection, the remaining weights and their further updates are ignored. This means that the updates of the unpruned weights are not considered during the mask selection procedure. Consequently, solving each subproblem (CBS-S and CBS-U) to optimality does not guarantee optimal solution to CBS.

In the theoretical I-OBS algorithm, masks and weight updates are solved jointly, as stated in Lemma 2. Optimization in equation (8) handles both new weights and masks. Although the practical I-OBS algorithm first selects a mask and then updates weights (similar to CBS-U), it optimally finds the mask by solving an Integer Programming (IP) problem, unlike CBS-S's IQP. This makes I-OBS mask selection optimal but computationally challenging, so the practical version uses a top-k mask for efficiency, maintaining the same local convergence rate in strongly convex cases.

Q2. We will discuss more literature review over broader range of model pruning: Our work follows the salience-based weight pruning approach, which evaluates the impact of removing weights on the model's loss or output. Among these, methods based on second-order information are most relevant, particularly those following the [OBS] framework. [L-OBS] uses a second-order approximation of the loss function, assuming a zero gradient, and introduces a layerwise strategy to approximate the Hessian. [WoodFisher] scales this idea with the empirical Fisher approximation of the Hessian, improving performance and considering non-zero gradients. However, WoodTaylor methods in [WoodFisher] focus on pruning one weight at a time, whereas our work extends this to multiple weights.

Other methods addressing multiple weight pruning include [CBS], which formulates mask selection as an integer programming problem and proposes two heuristics to approximate its solution. [OBC] tackles layerwise pruning with a quadratic problem and introduces a greedy heuristic for efficient layerwise problem-solving. Similarly, [CHITA] formulates layerwise pruning as an $L0L2$-regularized quadratic problem and proposes an IHT-like iterative algorithm with a line search strategy. [Net-trim] minimizes the $L1$​ norm of weight matrices while ensuring similar activations in the pruned layer. While prior iterative algorithms like [CHITA] focus on specific problems, our methods apply to general loss functions beyond quadratic problems and offer convergence guarantees. Practically, within our iterative framework, using a cost-effective projection step (e.g., TopK or SparseGPT) yields competitive results compared to more complex one-shot solvers.

Q3. We consider sparse optimization problems, aiming to find the sparse solution $\theta_*$ for an objective function $f(\theta)$ with a $k^*$-sparse global optimum $\theta_*$. This includes the classic sparse recovery problem: recovering a $k$-sparse signal $\theta_*$ from a noisy observation $A \theta_* + \epsilon$ with a sensing matrix $A$. I-OBS is a sparse optimization algorithm for sparse recovery. Besides finding $\theta_*$, it ensures each iteration's weights are $k$-sparse (with $k \ge k^*$), guaranteeing a $k$-sparse solution even if $\theta_*$ is not found. Such type of algorithms are known as proper learners in sparse recovery. Both theoretical and practical versions of I-OBS are sparse optimization algorithms resembling sparse Newton's methods. Similar first-order algorithms are studied in [AC/DC], but I-OBS leverages second-order information to potentially speed up convergence (both I-OBS versions achieve super-exponential local convergence rates).

For quadratic objectives, consider two settings: strongly convex (Hessian is positive definite) and convex (Hessian is positive semi-definite but not positive definite). Strongly convex: The theoretical I-OBS solves the problem in one step as the second-order approximation equals the function itself. The challenge is efficiently computing the optimal mask. The practical I-OBS doesn't converge in one step but shows super-exponential local convergence. Convex: I-OBS cannot be directly applied due to the non-invertible Hessian, similar to issues in Newton's method for convex optimization. Adding a small L2 regularization can address invertibility but complicates obtaining a convergence rate, as the global optimum is no longer sparse. This is left for future work.

Q4 and Q5. Thank you for the suggestion. In Table 3 of the attached file in global rebuttal, we provide results for I-OBS applied to MobileNetV1 model used in STR, which is the same model studied in Table 4 of the CBS paper. We skip the depth-wise convolutions having shape $(C, 1, K, K)$ when we apply our pruning algorithm, as is standard. Starting with 60% sparsity, our I-OBS pruner outperforms CBS method on MobileNetV1 by a large margin: 2% for 60% sparsity, 5% for 70% sparsity and 12% for 80% sparsity. For low sparsities (30% to 50%), the two methods are comparable, since the accuracy difference is less than 0.5%.

Thank you for the valuable questions and comments, we address the questions and comments below:

W1. The difference between the practical and theoretical version of the algorithm We compare the practical and theoretical version of Algorithm 1 below: (1) The way of choosing the mask is different: For the theoretical version of the algorithm, we choose the mask in an optimal way, which is equivalent to solving the integer programming problem in Line 7 of Algorithm 1; for the practical version, we simply use the top-k mask to replace the optimal mask (2) The complexity of the two versions is different. For the theoretical one, solving the integer programming problem in Line 7 of Algorithm 1 is NP-hard; while for the practical version, computing a top-k mask only has $O(d)$ complexity (3) The local convergence rate of the two algorithms is the same: in both case we obtain a local convergence in the form of $||\theta\_{t+1} - \theta^*||\_2 \leq C||\theta\_t - \theta^*||\_2^2$ implies a $\mathcal{O}(\log \log(1/\epsilon))$ iteration complexity for achieving an $\epsilon$-error with initialization $||\theta\_0 - \theta^*||\_2 \leq \frac{1}{2C}$.

W2. The complexity issue by using second-order information While second-order information is historically hard to apply at the scale of deep networks, notice that there have recently been several efficient variants that work in linear time and space (in the model dimension) and can therefore be scaled. Specifically, we apply the recent scalable M-FAC method [M-FAC] to approximate the inverse of the Hessian.

W3. Comparing I-OBS with other existing methods

We provide extra experiments that compare our methods with [CBS]. In Table 3 of the attached file in global rebuttal, we provide results for I-OBS applied to MobileNetV1 model used in STR. Our experiment results implies that: starting with 60% sparsity, our I-OBS pruner outperforms CBS method on MobileNetV1 by a large margin: 2% for 60% sparsity, 5% for 70% sparsity and 12% for 80% sparsity. For low sparsities (30% to 50%), the two methods are comparable, since the accuracy difference is less than 0.5%

Q1. How does the algorithm using second-order information compare with one using first-order information?

In I-OBS, we solve a sparse optimization problem at each step with the objective function being the second-order approximation of the loss function, in particular, we have $\theta\_{t+1} = \arg \min_{\theta: ||\theta||\_0 \le k} \; \phi\_t(\theta) + \tfrac{1}{2}||\theta-\theta\_t||^2_{H_t}$. The Hessian of the objective funcion appears in the $\tfrac{1}{2}||\theta-\theta\_t||^2\_{H\_t}$ term. While in the first-order method such as k-IHt studied in [AC/DC], the algorithm solves a sparse optimization problem at each step with the objective function being the first-order approximation of the loss function, in particular $\theta\_{t+1} = \arg \min\_{\theta: ||\theta||\_0\le k} \; \phi\_t(\theta) + \tfrac{1}{2\eta}||\theta-\theta\_t||^2$, the Hessian is replaced by a quadratic term.

Regarding the convergence rate, typical first-order methods such as k-IHt studied in [AC/DC] have achieved $\epsilon$-error in $O(\log \frac{1}{\epsilon})$ iterations for any initialization. The I-OBS proposed in this paper, achieved $\epsilon$-error in $O(\log\log \frac{1}{\epsilon})$ iterations with initialization $||\theta\_0 - \theta^*||\_2 \leq \frac{1}{2 C}$. In short, I-OBS have a faster local convergence rate but we are unable to provide a global convergence rate. The difference resembles the one between gradient descent and Newton's methods.

Q2 Comparing training time in the pruning process among the proposed method and existing works

Thank you for the suggestion. However, it is a bit difficult to directly compare the pruning time since it depends on the implementation of the algorithms and the device. We need to adapt the implementation of methods in other existing works on our device, which is a bit time-consuming. We report the pruning time of our methods below:

For the experiments on ViTs (Table 1 in the paper), it took about 4 hours for running 100 iterations and the hardware we used was Nvidia A100 GPU with 80GB memory. For LLMs experiments, as suggested by reviewer T9TW, we ran I-OBS on Llama-2 (7B) and Llama-3 (8B) models, and it takes around 1.5 hours per iteration (15h for 10 iterations). Unfortunately, we haven't recorded the time of running the experiments on OPT-125M and Phi-1.5 models (Table 2 in the paper).

Q3 Whether the proposed algorithm applies to CNNs

The proposed methods indeed apply to CNNs. In Algorithm 2, we use OBC as the quadratic sparse solver for the problem defined in line 7 of Algorithm 2, and OBC applies to convolution layers. In fact, the experiments we are doing to baseline our methods with CBS are on MobileNet_v1, and we find I-OBS improves the performance of CBS on pruning MobileNet_v1.

References

[M-FAC] Frantar, Elias, Eldar Kurtic, and Dan Alistarh. "M-fac: Efficient matrix-free approximations of second-order information." Neurips 2021.

[AC\DC] Peste, Alexandra, et al. "Ac/dc: Alternating compressed/decompressed training of deep neural networks." Neurips 2021

[CBS] Yu, X., Serra, T., Ramalingam, S., & Zhe, S.. “The combinatorial brain surgeon: pruning weights that cancel one another in neural networks.” ICML 2022