2SSP: A Two-Stage Framework for Structured Pruning of LLMs

2025-01-29

AI Summary

To address the low inference efficiency and high computational overhead of large language models (LLMs), this paper proposes a two-stage collaborative structured pruning framework. In the first stage, width pruning removes entire neurons based on an output-impact importance score; in the second stage, depth pruning dynamically eliminates attention submodules guided by perplexity. Our key contributions include: (i) the first dual-path pruning mechanism jointly optimizing width and depth, (ii) novel constraints preserving intermediate-layer connectivity and enforcing global sparsity, and (iii) a dynamic sparsity-ratio balancing strategy. Evaluated across four major LLM families under three sparsity levels, our method consistently outperforms five state-of-the-art pruning approaches on both language modeling and downstream tasks, while reducing pruning time by two orders of magnitude.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

We propose a novel Two-Stage framework for Structured Pruning (2SSP) for pruning Large Language Models (LLMs), which combines two different strategies of pruning, namely Width and Depth Pruning. The first stage (Width Pruning) removes entire neurons, hence their corresponding rows and columns, aiming to preserve the connectivity among the pruned structures in the intermediate state of the Feed-Forward Networks in each Transformer block. This is done based on an importance score measuring the impact of each neuron over the output magnitude. The second stage (Depth Pruning), instead, removes entire Attention submodules. This is done by applying an iterative process that removes the Attention submodules with the minimum impact on a given metric of interest (in our case, perplexity). We also propose a novel mechanism to balance the sparsity rate of the two stages w.r.t. to the desired global sparsity. We test 2SSP on four LLM families and three sparsity rates (25%, 37.5%, and 50%), measuring the resulting perplexity over three language modeling datasets as well as the performance over six downstream tasks. Our method consistently outperforms five state-of-the-art competitors over three language modeling and six downstream tasks, with an up to two-order-of-magnitude gain in terms of pruning time. The code is available at available at url{https://github.com/FabrizioSandri/2SSP}.

Problem

• Simplification
• Large Language Models
• Efficiency

Innovation

• 2SSP Framework
• Two-step Pruning Strategy
• Efficient Language Model Simplification