GradMAP: Faster Layer Pruning with Gradient Metric and Projection Compensation

TL;DR

GradMAP tackles the redundancy in LLM layers by introducing a two-stage pruning framework that first ranks layers using a global gradient-magnitude metric and then recovers performance with a lightweight projection compensation matrix. The Stage 1 gradient-based importance scores allow pruning decisions with a single backward pass, while Stage 2 targets the largest activation drift with a trainable down-projection compensation, reparameterized as a corrected downstream mapping. Empirically, GradMAP achieves about a $4\times$ speedup in pruning time and consistently outperforms baselines on zero-shot tasks and perplexity across multiple models, with the Stage 2 compensation providing additional gains and generality when integrated with other pruning methods. This approach offers a practical, hardware-friendly pathway to deploy compressed LLMs, balancing pruning efficiency and task performance with minimal extra calibration data and computation.

Abstract

Large Language Models (LLMs) exhibit strong reasoning abilities, but their high computational costs limit their practical deployment. Recent studies reveal significant redundancy in LLMs layers, making layer pruning an active research topic. Layer pruning research primarily focuses on two aspects: measuring layer importance and recovering performance after pruning. Unfortunately, the present works fail to simultaneously maintain pruning performance and efficiency. In this study, we propose GradMAP, a faster layer pruning method with \textbf{Grad}ient \textbf{M}etric \textbf{A}nd \textbf{P}rojection compensation, which consists of two stages. In the first stage, we introduce a novel metric based on gradient magnitudes, enabling a global assessment of layer importance. Note that, it requires only a single backward propagation step per pruning decision, substantially enhancing pruning efficiency. In the second stage, we first analyze the layers with the largest mean shift resulting from pruning, and then incorporate a simple yet effective projection compensation matrix to correct this drift in one step. In this way, the degradation of model performance caused by layer pruning is effectively alleviated. Extensive experiments show that GradMAP outperforms previous layer pruning methods in both pruning speed (achieving an average $4\times$ speedup) and performance.

Figures (8)

• Figure 1: Comparison of pruning time and accuracy across different methods on the Vicuna-7B.
• Figure 2: The Stage 1 of GradMAP introduces a novel gradient-based metric for estimating layer importance in LLMs, which quantifies each layer's contribution by analyzing gradient magnitudes. This metric enables the iterative identification and pruning of unimportant layers.
• Figure 3: The Stage 2 of GradMAP. To mitigate capacity loss, we introduce a learnable projection compensation matrix that optimizes reconstruction on a small calibration dataset.
• Figure 4: Perplexity of LLaMA2-7B under varying compression ratios.
• Figure 5: Distribution of first-order moment offsets before and after pruning for OPT-6.7B and OPT-2.7B.