ALPS: Improved Optimization for Highly Sparse One-Shot Pruning for Large Language Models

Reply to W1 and W2: Thanks for these references—we will include them in our revised paper. ALPS differs significantly from the mentioned works as follows:

• Comparison with "Progressive weight pruning of deep neural networks using ADMM": This paper applies ADMM to the original loss function, requiring expensive full network training via SGD. In contrast, ALPS performs post-training pruning, applying ADMM to layerwise reconstruction problems (least squares objective) with cardinality constraints. This approach enables us to scale ALPS to 30B parameters on a single V100 GPU.

• Comparison with "Fast and optimal pruning" (Boža): In this interesting work, Boža performs pruning post-training using a layerwise reconstruction loss. They apply ADMM to solve a least squares problem for a given sparsity mask. Boža selects the sparsity mask via iterative magnitude pruning. Using ADMM speeds up the back-solve procedure.

  • ALPS, in contrast, is an end-to-end approach that directly targets the $\ell_0$ constrained least squares problem with theoretical guarantees. ALPS simultaneously optimizes over the weights and the sparsity pattern, unlike the work by Boža where ADMM is used to perform the backsolve operation. Therefore, ALPS achieves lower layerwise reconstruction loss than Boža, as shown in Table 1 in the general rebuttal.

  • Since Boža employs ADMM for backsolving [finding the weights of the least squares problem given a mask], we also compare it with our PCG procedure, which serves as a backsolve method. We tested both approaches for computing weights on a given support. Results show that PCG outperforms Boža’s ADMM procedure in both computational time and objective value (see Table 2 in general rebuttal). The time advantage of PCG stems from its ability to backsolve without explicitly computing matrix inverses.

• Novelty of ALPS: Please note that ALPS has several differences/innovations compared to standard ADMM:

  • A novel penalty parameter updating scheme that enables finding high-quality support and accelerates convergence. As we show in our response to Q1 and Q4, in practice, our proposed ADMM has better convergence properties compared to the usual version of ADMM for $\ell_0$-constrained least squares problems.
  • As far as we know, our theory is the first theoretical convergence result for ADMM applied to $\ell_0$-constrained problems.
  • A PCG-based post-processing technique optimized with vectorization and GPU parallelism, achieving up to 200x speedup compared to direct backsolve for updating weights on a given support. These contributions, taken together, make our work significantly different from existing ADMM methods especially in the area of LLM pruning.

Reply to W3: Please note that our study demonstrates ALPS's effectiveness as a pruning method across all sparsity levels, including 60% sparsity and N:M sparsity pattern. In these cases, the pruned models maintain competitive performance.

You are right that model utility can deteriorate with one-shot pruning at very high sparsity levels — in such cases, we can recover model performance via retraining. Since ALPS-pruned models outperform models pruned by competing methods, they would require fewer retraining epochs to recover lost performance [1]. Consequently, ALPS remains valuable at high sparsity levels by reducing computational costs during the expensive retraining phase.

[1] Meng, X., Ibrahim, S., Behdin, K., Hazimeh, H., Ponomareva, N., & Mazumder, R. OSSCAR: One-Shot Structured Pruning in Vision and Language Models with Combinatorial Optimization.

Reply to W4: Thanks for your thoughtful comment. In our response to reviewer ZgXD, we report the running time of ALPS and other one-shot pruning methods. We emphasize that with retraining, the total runtimes of Wanda and DSnoT become much higher than ALPS. As reported in Wanda, fine-tuning LLaMa-7B with LoRA takes about 24 hours on a V100 GPU, and full parameter fine-tuning takes 24 days. In contrast, ALPS prunes the model in less than 1 hour. As a layer-wise pruning method, ALPS loads weights and activations onto the GPU layer by layer, enabling it to prune models as large as OPT-30B on a single V100 GPU, while retraining often requires more careful memory management. Given these considerations, we think that retraining methods would be much more computationally expensive than ALPS. Please recall that our focus here is on demonstrating the effectiveness of ALPS as a one-shot pruning method. As a side note, as we explained in our reply to ZgXD, the runtime of ALPS can be further decreased by loosening the convergence criteria of ALPS.

Reply to Q1 and Q4: Our proposed $\rho$ update scheme, theoretically supported by Theorem 1, ensures that ALPS converges rapidly while finding high-quality solutions. In contrast, ADMM with a fixed $\rho$ may fail to converge when applied to $\ell_0$ constrained least squares problems. To demonstrate this, we compared ALPS with ADMM using fixed $\rho$ values. We examined two key metrics: reconstruction loss (objective) and the rate of change of the support (of weights) between consecutive iterations (this measures the convergence speed of the algorithm). The results are presented in the general rebuttal (Tabl 3 and 4). Results reveal that ADMM with a large $\rho(=3)$ converges quickly but yields poor solutions, while a small $\rho(=0.3)$ fails to converge. ALPS, utilizing our $\rho$ update scheme, achieves both rapid convergence and high-quality solutions.

Reply to Q2: Thanks for your comment — please refer to our response to reviewer ZgXD.

Reply to Q3: Please refer to Table 5 in general rebuttal. In the table, we omit LLaMA results on LAMBADA due to its poor performance without modifications.

---

In light of the above discussions and results, we kindly ask the reviewer to consider increasing their evaluation.

We would like to thank the reviewer for their thoughtful feedback. Below, we provide some answers/clarifications.

Experiments could be more solid since many experiments are run on OPT, which is somewhat out-of-date. It would also be better to consider more challenging benchmarks, such as GSM8K or other questions that require generating a long answer.

Reply: We appreciate the reviewer's valuable suggestion. Based on your comment and that of reviewer ZgXD, we will consider comparing ALPS with other methods on the recent model LLaMA-3 and broaden our assessment by incorporating knowledge-intensive and reasoning-based tasks, specifically MMLU and GSM8K.

It would be better to compare with other representative pruning methods beyond the one-shot unstructured pruning, and to explore how this method could be integrated with others.

Reply: Thank you for this insightful suggestion! Beyond unstructured pruning, numerous representative pruning methods focus on structured sparsity patterns, such as block sparsity [1] and row sparsity [2]. Though direct numerical comparisons with these methods are potentially challenging due to their different settings, the idea of integrating ALPS with structured pruning methods is intriguing. Having a pruned model with both unstructured and structured sparsity could optimize both storage costs and inference time. It would be certainly interesting to explore such an integration as future work.

[1] Gray, S., Radford, A., & Kingma, D. P. GPU kernels for block-sparse weights.

[2] Meng, X., Ibrahim, S., Behdin, K., Hazimeh, H., Ponomareva, N., & Mazumder, R. OSSCAR: One-Shot Structured Pruning in Vision and Language Models with Combinatorial Optimization.

This work discusses engineering optimizations, such as incorporating PCG with GPU. Will you open-source the code or provide more efficient results from these engineering efforts?

Reply: If the paper is accepted, we will release the codes. Additionally, we will provide explanations of such implementations and a tutorial on their usage.

We would like to thank the reviewer for their thoughtful feedback. Below, we provide some answers/clarifications.

It would be beneficial to include evaluations on more recent and extremely large models to demonstrate the method's applicability. The current evaluation tasks are relatively limited. Expanding the evaluation to include knowledge-intensive tasks and reasoning-based tasks would provide a more comprehensive assessment.

Reply: We appreciate the reviewer's valuable suggestion. We will consider comparing ALPS with other methods on the recent model LLaMA-3 and broaden our assessment by incorporating knowledge-intensive and reasoning-based tasks, specifically MMLU and GSM8K. As a side remark, given our academic resource constraints (we conduct our experiments on a V100 GPU with 32GB memory), the largest model we’ve tried in the paper is OPT-30B.

The current focus is primarily on N:M sparsity patterns. It would be great to explore structured pruning. Can ALPS be extended to handle other types (structured sparsity) of structured sparsity patterns beyond N:M?

Reply: While we focus on unstructured sparsity and N:M sparsity, our approach can be generalized to handle other structured sparsity patterns, as discussed below.

Let $S$ denote the set of weight matrices with the desired structured sparsity pattern (in the case of unstructured sparsity, $S$ contains all matrices with no more than $k$ non-zero elements). Our method can be extended to find the optimal matrix in $S$ by modifying the $\mathbf{D}$-update step in Eq.(4) in our paper. We can set $\mathbf{D}^{(t+1)}$ to be the projection of $\mathbf{W}^{(t+1)}+\mathbf{V}^{(t+1)}/\rho$ onto $S$. Our approach remains efficient and can find high-quality solutions for sparsity patterns with low-cost projection operations, including:

• Block sparsity [1]
• Hierarchical sparsity [2]
• Row sparsity [3]

Importantly, our convergence results (Theorem 1) and proposed PCG step directly apply to these structured sparsity patterns as well.

We thank the reviewer for bringing up this point, and we will be happy to discuss these extensions in our revision.

[1] Gray, S., Radford, A., & Kingma, D. P. GPU kernels for block-sparse weights.

[2] Wu, Y. N., Tsai, P. A., Muralidharan, S., Parashar, A., Sze, V., & Emer, J. Highlight: Efficient and flexible DNN acceleration with hierarchical structured sparsity.

[3] Meng, X., Ibrahim, S., Behdin, K., Hazimeh, H., Ponomareva, N., & Mazumder, R. OSSCAR: One-Shot Structured Pruning in Vision and Language Models with Combinatorial Optimization.

How does the runtime of ALPS compare to other methods, especially for very large models?

Reply: We compare the runtime of ALPS with other methods in the following table. Here, the runtime includes the time for generating input activations $\mathbf{X}$.

ALPS employs an advanced optimization method to solve the layerwise reconstruction problem, which results in longer running times compared to other algorithms. However, it's important to note that ALPS's runtime is still negligible when compared to fine-tuning methods for LLMs. For context, Wanda [1] reports that fine-tuning LLaMa-7B with LoRA takes about 24 hours on a V100 GPU, while full parameter fine-tuning takes 24 days. In contrast, ALPS prunes the model in less than 1 hour. Furthermore, in our paper, ALPS is run with a tight convergence criterion. We can improve its runtime further by using a variant, ALPS-simple, which uses a looser convergence criterion. As shown in the above tables, ALPS-simple achieves shorter running times compared to SparseGPT while still maintaining better utility, as measured by single-layer reconstruction loss. We are happy to include the performance results of ALPS-simple in the revised paper if it is of interest to the reviewers.

[1] Sun, M., Liu, Z., Bair, A., & Kolter, J. Z. A simple and effective pruning approach for large language models.