Two-Stage Regularization-Based Structured Pruning for LLMs

1. Problem Statement

Large Language Models (LLMs) face significant deployment challenges due to their massive parameter counts, which lead to high computational costs and memory requirements. While structured pruning (removing entire layers or channels) is a promising solution for compression, existing layer-wise pruning methods suffer from critical limitations:

• Knowledge Loss: Current methods typically calculate layer importance, directly remove "unimportant" layers, and then attempt to recover performance via extensive retraining. However, even layers deemed "unimportant" often contain valuable knowledge. Direct removal causes immediate, irreversible knowledge loss.
• Retraining Overhead: The performance degradation caused by direct pruning necessitates substantial retraining (fine-tuning) to recover, which incurs high computational costs and data requirements.
• Sequential Flaw: The standard "select-then-prune" approach fails to redistribute the knowledge contained in the layers scheduled for deletion to the remaining layers before the deletion occurs.

2. Methodology: TRSP

The authors propose TRSP (Two-Stage Regularization-Based Structured Pruning), a novel framework that performs regularization before pruning to facilitate knowledge transfer. The process consists of four stages:

A. Data Preparation

TRSP requires only a minimal amount of data (e.g., 128 random samples from standard benchmarks like WikiText-2) to guide the regularization process, significantly reducing computational overhead compared to full retraining datasets.

B. Stage 1: Iterative Layer Weight Learning

• Mechanism: Each transformer layer is multiplied by a learnable scalar weight $S[i]$. Initially, all weights are set to 1.
• Optimization: The model is optimized to minimize the language modeling loss ($L(W, X)$) plus an $\ell_1$-norm regularization term on the layer weights ($\lambda_1 \sum \|S[i]\|_1$).
• Process: This is performed iteratively (greedy strategy). In each step, the layer with the smallest weight is identified and masked (temporarily removed) for the next iteration. This prevents the "one-shot" removal of consecutive layers, which can destabilize the model.
• Mathematical Trick: Since the $\ell_1$-norm is non-differentiable, the authors transform the problem into an equivalent constrained optimization problem (using auxiliary variables) to allow standard backpropagation.

C. Stage 2: Knowledge Transfer Regularization

• Mechanism: Once the set of layers to be pruned ($P$) is identified, a second regularization step is applied.
• Objective: The loss function adds a term penalizing the difference between the input and output of the layers in set $P$: $\lambda_2 \sum \|X_{out}^i - X_{in}^i\|$.
• Goal: This forces the layers marked for removal to behave like identity functions ($X_{out} \approx X_{in}$). By minimizing the transformation performed by these layers, the model is compelled to redistribute the knowledge originally held in these layers into the remaining layers before the layers are physically deleted.
• Norm Choice: The authors use either $\ell_1$ or $\ell_2$ norms, finding that $\ell_2$ generally yields slightly better results.

D. Final Pruning

After the two-stage regularization, the layers in set $P$ are permanently removed. The resulting model is ready for inference without any further retraining.

3. Key Contributions

1. Knowledge Retention via Regularization: Unlike traditional methods that prune then retrain, TRSP uses regularization to migrate knowledge from pruned layers to preserved layers in advance, minimizing information loss.
2. Retraining-Free: The method achieves high performance without the need for costly fine-tuning or retraining on large downstream datasets after pruning.
3. Iterative Weight Learning: The authors demonstrate that iteratively learning layer weights (greedy approach) is superior to one-shot selection, as it avoids removing consecutive layers that might collectively be critical.
4. Minimal Data Dependency: The entire process relies on a tiny fraction of data (e.g., 128 samples), making it highly efficient.

4. Experimental Results

The authors evaluated TRSP on various models (Phi-2, OPT, LLaMA2, LLaMA3) across generation and zero-shot tasks.

• Performance Superiority:
  • TRSP consistently outperforms strong baselines (SLEB, ShortGPT, LaCo, Shortened LLaMA).
  • On LLaMA2-7B with 25% pruning, TRSP achieved a 20% lower perplexity than the second-best method (ShortGPT).
  • On OPT-13B, TRSP dropped only 1% in average accuracy compared to the dense model, whereas baselines dropped significantly more.
• Acceleration:
  • TRSP delivers notable end-to-end speedups. For OPT-13B at 50% pruning, it achieved a 1.75× throughput increase and a 1.85× latency reduction.
  • For LLaMA2-13B at 50% pruning, it achieved a 1.71× throughput increase.
• Robustness:
  • TRSP maintains low perplexity even at high pruning ratios (up to 60%), where other methods fail.
  • It is robust across different calibration datasets (Alpaca, WikiText-2, PTB, C4).
• Scalability: Experiments on LLaMA-2-70B confirmed that TRSP scales effectively, outperforming baselines in both perplexity and zero-shot accuracy without retraining.
• Comparison with Other Granularities: TRSP (layer pruning) outperformed unstructured pruning (Wanda) and channel pruning (SliceGPT) in terms of the trade-off between performance and hardware-friendly acceleration.

5. Significance

TRSP represents a paradigm shift in LLM pruning strategies. By decoupling the "knowledge redistribution" phase from the "physical removal" phase via a two-stage regularization process, it solves the fundamental issue of knowledge loss in structured pruning.

• Efficiency: It drastically reduces the computational cost of model compression by eliminating the need for retraining.
• Practicality: The minimal data requirement makes it feasible for organizations with limited resources to compress large models.
• Generalizability: The method is applicable to various LLM architectures and scales effectively to very large models (70B+ parameters), offering a viable path for deploying efficient LLMs in resource-constrained environments.

Code Availability: The implementation is open-sourced at https://github.com/fmk345/TRSP.