Bypass Back-propagation: Optimization-based Structural Pruning for Large Language Models via Policy Gradient

Authors: Yuan Gao, Zujing Liu, WEIZHONG ZHANG, Bo Du, Gui-Song Xia
Primary area: optimization
Keywords: Optimization-based Pruning, Back-Propagation-Free, Structural Pruning, Large Language Models

Abstract:
In contrast to moderate-size neural network pruning, structural weight pruning on the Large-Language Models (LLMs) imposes a novel challenge on the efficiency of the pruning algorithms, due to the heavy computation/memory demands of the LLMs. Recent efficient LLM pruning methods typically operate at the post-training phase without the expensive weight finetuning, however, their pruning criteria often rely on heuristically hand-crafted metrics, potentially leading to suboptimal performance. We instead propose a novel optimization-based structural pruning that learns the pruning masks in a probabilistic space directly by optimizing the loss of the pruned model. To preserve the efficiency, our method eliminates the back-propagation through the LLM per se during the optimization, requiring only the forward pass of the LLM. We achieve this by learning an underlying Bernoulli distribution to sample binary pruning masks, where we decouple the Bernoulli parameters from the LLM loss, thus facilitating an efficient optimization via a policy gradient estimator without back-propagation. As a result, our method is able to 1) operate at structural granularities of channels, heads, and layers, 2) support global and heterogeneous pruning (i.e., our method automatically determines different redundancy for different layers), and 3) optionally initialize with a metric-based method (for our Bernoulli distributions). Extensive experiments on LLaMA, LLaMA-2, LLaMA-3, Vicuna, and Mistral using the C4 and WikiText2 datasets demonstrate that our method operates for 2.7 hours with around 35GB memory for the 13B models on a single A100 GPU, and our pruned models outperform the state-of-the-arts w.r.t. both perplexity and the majority of various zero-shot tasks. Codes will be released.

Decision: reject
Reviewer scores: 5, 6, 5, 6, 3

Reviews and discussion:
<thread>
creation date:
1730214544075

thread_type:
official_review

<thread_object>
creator: reviewer_pajx

<content>
summary:
This paper proposes a novel optimization-based structural pruning method for Large Language Models (LLMs) that avoids expensive back-propagation through the LLM. The key idea is to formulate pruning as learning binary masks sampled from underlying Bernoulli distributions. By decoupling the Bernoulli parameters from the LLM loss, the method enables efficient optimization using policy gradient estimation that only requires forward passes. The approach supports different structural granularities (channels, heads, layers) and enables global/heterogeneous pruning across the model. Extensive experiments on various LLMs demonstrate strong performance while maintaining efficiency.

strengths:
- Clear writing and good organization of the paper
- The paper presents a novel methods for structural pruning that leverages policy gradient estimation, bypassing the need for back-propagation through large models.
- The method supports multiple pruning granularities (channels, heads, layers) and can be initialized using metric-based methods, showcasing its adaptability and potential for widespread application.

weaknesses:
1. **Lack of Comparison with Recent Baselines**:  
   The paper does not compare the proposed method with recent pruning techniques such as Shortened LLaMA ([arXiv:2402.02834](https://arxiv.org/abs/2402.02834)) and SLEB ([arXiv:2402.09025](https://arxiv.org/abs/2402.09025)), which were published prior to ShortGPT ([arXiv:2403.03853](https://arxiv.org/abs/2403.03853)) and Gromov et al.'s work ([arXiv:2403.17887](https://arxiv.org/abs/2403.17887)). Additionally, MKA ([arXiv:2406.16330](https://arxiv.org/abs/2406.16330)) is also not included. Including comparisons with these methods on benchmarks like MMLU and GSM8K, using metrics like accuracy and perplexity, would provide a more comprehensive evaluation of the method's performance relative to the latest advancements in the field.

questions:
1. Does the paper include accuracy (ACC) results on more advanced benchmarks like MMLU or GSM8K? If not, could the authors provide such evaluations to demonstrate the method's effectiveness on tasks that require higher reasoning capabilities?

2. The proposed method exhibits longer training times compared to metric-based pruning methods. Are there potential optimizations or strategies that the authors are considering to reduce the training duration without compromising performance?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730309785973

thread_type:
official_review

<thread_object>
creator: reviewer_ixef

<content>
summary:
This paper presents a novel approach to structural pruning of LLMs that bridges the gap between optimization-based and metric-based pruning methods. The key innovation lies in the use of Bernoulli distributions to sample binary pruning masks, coupled with a policy gradient estimator that eliminates the need for back-propagation through the LLM. This is a clever solution that maintains the benefits of optimization-based approaches while achieving the efficiency typically associated with metric-based methods.

strengths:
- The proposed method is both theoretically sound and practically implementable, requiring only forward passes through the LLM during optimization.
(1) The efficiency claims are impressive (2.7 hours, 35GB memory for 13B models on a single A100)
(2) The approach is flexible, supporting multiple structural granularities (channels, heads, and layers)
(3) The method automatically handles heterogeneous pruning across layers, addressing a key limitation of existing approaches
(4) Extensive experimental validation across multiple LLM architectures and datasets

weaknesses:
The paper could benefit from more detailed ablation studies on the impact of different structural granularities.

The main weaknesses of this LLM pruning work include: 
(1) limited evaluation beyond perplexity and basic zero-shot tasks, particularly lacking analysis of inference speed improvements and downstream task performance, 
(2) methodological constraints of using basic REINFORCE rather than more advanced policy gradient methods, 
(3) heavy reliance on the C4 dataset with some cross-dataset performance issues on specific zero-shot tasks like WinoGrande and Hellaswag. 
(4) While the method is more efficient than traditional back-propagation approaches, it still requires longer training time compared to metric-based methods (2.7 hours for LLaMA-2-13B) and significant memory usage (35GB).

questions:
- While the method avoids back-propagation through the LLM, how does the efficiency-accuracy trade-off compare when scaling to much larger models beyond 13B parameters? This is crucial for understanding the method's practical applicability.

- Is there any patterns among the structures pruned through using this method ?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730386086617

thread_type:
official_review

<thread_object>
creator: reviewer_scbn

<content>
summary:
The paper introduces an optimization-based structural pruning method for LLMs that eliminates the need for back-propagation. By employing Bernoulli distributions to learn pruning masks, the authors optimize these masks using a policy gradient estimator, enabling gradient estimation solely through the forward pass. This approach enhances efficiency in memory and computation compared with methods requiring back-propagation, and it achieves performance superior to heuristic-based pruning metric methods.

strengths:
1. The paper presents a novel pruning method that leverages policy gradient estimators instead of back-propagation, addressing key computational challenges in gradient-based LLM pruning methods.
2. The method supports multiple structural granularities (channels, heads, layers), providing flexibility in how the model is pruned. It also allows for global and heterogeneous pruning, which is more aligned with the varying redundancy across layers in LLMs.

weaknesses:
1. While the paper suggests using a policy gradient estimator to bypass back-propagation, policy gradient methods can suffer from high variance, which may lead to unstable training. The paper does propose a variance-reduction technique, but the effectiveness of this could be further elaborated or validated with more ablation studies. For example, how is the performance of the proposed methods compared with the results of using back-propagation? 
2. Following up on Weakness 1, could you clarify the exact speedup over direct back-propagation? A runtime comparison on a specific model and hardware or a theoretical analysis of computational complexity would be helpful.
3. The proposed method requires 120K samples with a sequence length of 128, whereas baseline methods like LLM-Pruner, SliceGPT, and Wanda-SP use significantly smaller calibration data, such as 128 samples with a sequence length of 2048. Are the results for the baseline methods obtained with the same calibration data as yours? If not, this may lead to unfair comparisons. Please clarify if all methods used the same calibration data in the reported results. If not, provide results with all methods using the same calibration data, or explain why this is not feasible
4. The sparsity ratios used in the experiments exceed 30%, which may be excessive for structured pruning, as the high PPL results may not be meaningful in practice and could significantly degrade real inference performance, such as in question answering tasks. Therefore, it is important to report performance under lower sparsity levels, such as 10% and 20%.
5. The experiments primarily compare the proposed method with other non-weight-updating pruning techniques. While this makes sense in terms of efficiency, it would be interesting to see how the method stacks up against pruning methods that do involve weight updates, such as [1], especially in terms of final model performance.

[1] Xuan Shen, Pu Zhao, Yifan Gong, Zhenglun Kong, Zheng Zhan, Yushu Wu, Ming Lin, Chao Wu, Xue Lin, and Yanzhi Wang. Search for efficient large language models. arXiv preprint arXiv:2409.17372, 2024.

questions:
In addition to the questions regarding weaknesses, I recommend including the results of dense models in your tables to highlight the performance degradation resulting from pruning.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730691102582

thread_type:
official_review

<thread_object>
creator: reviewer_1xw7

<content>
summary:
This paper proposes an optimization-based structural pruning method for Large Language Models (LLMs) which notably does not require any gradient back-propogation. The method works by casting the binary pruning mask over structural components as a bernoulli variable. This reparametrization allows us to solve the pruning problem via policy gradient. The authors conduct thorough experimental evaluation on open-sourced LLMs and demonstrate the effectiveness of the proposed pruning method.

strengths:
- This paper proposes a novel method, which casts the optimization problem of selecting optimal pruning mask as a reinforcement learning problem. This allows us to avoid the inefficiency for performing computationally intensive back-propogation.  
- The authors have conducted extensive experimental evaluation and compare with many existing baseline methods for structural pruning. The results show that the proposed method is a promising approach for structural pruning.

weaknesses:
- It might be good to evaluate a larger LLM model, e.g., LLaMA-3-70B.  
- The authors stress the memory efficiency of optimizating without back-propogation. It would be good to dedicate one section comparing the resource consumptions of different pruning approaches, especially to those with gradient computation.

questions:
- Could the proposed approach be used for width pruning [1]?  
- In principle, could we solve problem 4 via back-propogation? For this setting, is it expected that the solver with gradient information be better?  
- Could iterative pruning be beneficial for this approach? By iterative pruning, I mean doing the pruning procedure multiple times until the target prune ratio is reached. 

[1] LLM Pruning and Distillation in Practice: The Minitron Approach. NVIDIA 2024.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730692790997

thread_type:
official_review

<thread_object>
creator: reviewer_22xp

<content>
summary:
In this work, the authors leverage policy gradient estimation to optimize pruning masks without relying on backpropagation. They validate the efficacy of their approach on various pre-trained models and compare it with several heuristic-based baselines.

strengths:
- The performance gain is significant.
- The idea of leveraging policy gradient for pruning is novel and insightful.

weaknesses:
- The tables are hard to read. Using a figure instead will help the visualization.
- How does the method compare in performance with Gumbel-Softmax approaches? The authors primarily focus on heuristic-based comparisons, leaving out Gumbel-Softmax methods. For example, [1] employs Gumbel-Softmax but focus on semi-structured sparsity. Although the objectives differ slightly, including a performance and cost comparison with Gumbel-based methods would strengthen the study. 
- Additionally, the comparison between Gumbel-Softmax and the proposed method in Section 3.2 is unclear. Adding a formulation-based comparison would clarify the differences. Will leveraging gumbel-softmax leads to more accurate gradient calculation? More detailed discuss on the trade-off will be beneficial.
- The probabilistic modeling of pruning masks is not entirely novel, as it has been applied in various Gumbel-based pruning methods [1].
- A comparison and analysis of the pruning masks generated by different methods would be beneficial, such as evaluating the Hamming distance between the masks.

[1] MaskLLM: Learnable Semi-Structured Sparsity for Large Language Models

questions:
N/A
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper titled "Bypass Back-propagation: Optimization-based Structural Pruning for Large Language Models via Policy Gradient" proposes a novel method for pruning large language models (LLMs) that bypasses traditional back-propagation techniques. The authors introduce an optimization-based approach that utilizes policy gradient methods to identify and remove less significant model parameters, thereby reducing model size and improving efficiency without significantly sacrificing performance. The key claims include that this method can lead to effective pruning while maintaining the integrity of model performance, and it is particularly beneficial for deployment in resource-constrained environments.

The findings suggest that the proposed pruning technique not only achieves comparable results to existing methods but also offers advantages in terms of computational efficiency during the pruning process. However, reviewers raised concerns regarding the empirical validation of these claims and the overall effectiveness of the proposed approach. Given the weaknesses, I recommend rejecting this paper. Further work is required to address the issues proposed by reviewers for the next version of this work.