[Paper Review] LoRAPrune: Structured Pruning Meets Low-Rank Parameter-Efficient Fine-Tuning

Mingyang Zhang, Hao-Zhen+5 more|arXiv (Cornell University)|May 28, 2023

Sparse and Compressive Sensing Techniques4 citations

TL;DR

LoRAPrune proposes a unified framework that jointly performs structured pruning and LoRA-based parameter-efficient fine-tuning on large pre-trained models, using LoRA gradients for importance estimation to avoid memory-intensive pre-trained weight gradients. It achieves state-of-the-art compression with 50% model size reduction, reducing perplexity by 8.0 on WikiText2 and 16.05 on PTB, while cutting memory usage by 52.6% compared to prior methods.

ABSTRACT

Large Language Models (LLMs), such as LLaMA and T5, have shown exceptional performance across various tasks through fine-tuning. Although low-rank adaption (LoRA) has emerged to cheaply fine-tune these LLMs on downstream tasks, their deployment is still hindered by the vast model scale and computational costs. Post-training model pruning offers a way to compress LLMs. However, the current pruning methods designed for LLMs are not compatible with LoRA. This is due to their utilization of unstructured pruning on LLMs, impeding the merging of LoRA weights, or their dependence on the gradients of pre-trained weights to guide pruning, which can impose significant memory overhead. To this end, we propose LoRAPrune, a new framework that delivers an accurate structured pruned model in a highly memory-efficient manner. Specifically, we first design a LoRA-guided pruning criterion, which uses the weights and gradients of LoRA, rather than the gradients of pre-trained weights for importance estimation. We subsequently integrate this criterion into an iterative pruning process, effectively removing redundant channels and heads. Extensive experimental results demonstrate the superior performance of our LoRAPrune over existing approaches on the LLaMA series models. At a 50\% compression rate, LoRAPrune demonstrates superior performance over LLM-Pruner, achieving a reduction in perplexity by 4.81 on WikiText2 and 3.46 on PTB, while also decreasing memory usage by 52.6%. Besides, LoRAPrune also matches semi-structural pruning across multiple LLMs, proving its wide applicability. The code is available at https://github.com/aim-uofa/LoRAPrune.

Motivation & Objective

• Address the incompatibility between unstructured pruning and LoRA, which prevents efficient merging of LoRA weights and increases inference latency.
• Overcome the high memory overhead of existing pruning methods that require gradients of pre-trained weights during pruning.
• Develop a memory-efficient, structured pruning method compatible with LoRA that enables seamless merging of LoRA weights post-pruning.
• Enable joint pruning and fine-tuning to maintain model performance under high compression ratios without separate fine-tuning after pruning.
• Propose a novel LoRA-guided pruning criterion that estimates parameter importance using only LoRA gradients, avoiding computation of pre-trained weight gradients.

Proposed method

• Introduce a LoRA-guided pruning criterion that computes parameter importance based solely on gradients of LoRA weights (A), eliminating the need for gradients of the pre-trained weights (W₀).
• Design a structured iterative pruning procedure that removes redundant channels and attention heads from both the pre-trained model and LoRA adapters in a joint manner.
• Use the Frobenius norm of LoRA gradients as a proxy for parameter importance, enabling efficient and accurate pruning without storing large intermediate tensors.
• Implement a joint pruning and fine-tuning pipeline where pruning and LoRA fine-tuning are alternated in a single training loop to maintain performance.
• Ensure that after pruning, LoRA weights (BA) can be seamlessly merged into the pruned pre-trained weights, enabling inference without additional computational cost.
• Optimize pruning frequency to balance model recovery and overfitting, with an optimal frequency of 10 fine-tuning steps per pruning iteration.

Experimental results

Research questions

• RQ1 Can a pruning criterion based solely on LoRA gradients achieve comparable or better performance than criteria relying on pre-trained weight gradients, while drastically reducing memory usage?
• RQ2 How does joint pruning and fine-tuning compare to separate pruning and fine-tuning in terms of model performance and efficiency on large pre-trained models?
• RQ3 What is the optimal pruning frequency that balances model recovery and overfitting during iterative pruning with LoRA fine-tuning?
• RQ4 To what extent can structured pruning with LoRA-guided importance estimation reduce memory overhead while maintaining model accuracy on downstream tasks?
• RQ5 Can the proposed method enable efficient, memory-effective, and hardware-compatible inference by allowing direct merging of LoRA weights into pruned models?

Key findings

• At a 50% compression rate, LoRAPrune reduces perplexity by 8.0 on WikiText2 and 16.05 on PTB compared to LLM-Pruner, demonstrating superior accuracy.
• LoRAPrune reduces memory usage by 52.6% compared to LLM-Pruner when pruning LLaMA-65b, requiring only one A100 (80GB) GPU instead of three.
• The LoRA-guided criterion reduces GPU memory usage by 52.6% and increases throughput by 64.6% compared to the vanilla gradient-based criterion.
• Joint pruning and fine-tuning outperforms separate pruning and fine-tuning, especially at high compression ratios, with a 2.87-point improvement in average accuracy on classification tasks at 50% ratio.
• The optimal pruning frequency is 10 fine-tuning steps per pruning iteration, as lower frequencies (e.g., 1) limit recovery and higher frequencies (e.g., 20) lead to overfitting.
• Mask similarity between the LoRA-guided and vanilla criteria remains high (especially at low compression rates), confirming that the LoRA-guided criterion is a strong approximation that enables accurate, low-cost pruning.