Probe Pruning: Accelerating LLMs through Dynamic Pruning via Model-Probing

TL;DR

Probe Pruning (PP) tackles the heavy compute cost of large language model inference by introducing an online dynamic structured pruning framework that operates batch‑wise without additional modules or fine‑tuning. PP uses a small, batch‑dependent probe of hidden states to guide per‑batch channel pruning through a three‑stage process: probing, history‑informed pruning with fusion, and full inference on the remaining weights. The PPsp pruning metric preserves per‑weight importance and aggregates channel importance via $L^2$ norms to select channels, while an importance‑scaled fusion integrates probe information with historical states to stabilize decisions across batches. Empirically, PP achieves substantial speedups with minimal loss of performance across LLaMA‑2/3 and OPT models, often outperforming calibration‑dataset baselines and matching or exceeding fine‑tuned baselines in zero‑shot tasks such as WikiText2 perplexity and commonsense reasoning, with probing consuming only about $1.5\%$ of the dense FLOPs. system-level results show PP yields a lower PRR than prior methods, indicating more efficient performance degradation trade‑offs for a given runtime reduction, and place PP as a practical tool for accelerating LLM inference in real‑world deployments.

Abstract

We introduce Probe Pruning (PP), a novel framework for online, dynamic, structured pruning of Large Language Models (LLMs) applied in a batch-wise manner. PP leverages the insight that not all samples and tokens contribute equally to the model's output, and probing a small portion of each batch effectively identifies crucial weights, enabling tailored dynamic pruning for different batches. It comprises three main stages: probing, history-informed pruning, and full inference. In the probing stage, PP selects a small yet crucial set of hidden states, based on residual importance, to run a few model layers ahead. During the history-informed pruning stage, PP strategically integrates the probing states with historical states. Subsequently, it structurally prunes weights based on the integrated states and the PP importance score, a metric developed specifically to assess the importance of each weight channel in maintaining performance. In the final stage, full inference is conducted on the remaining weights. A major advantage of PP is its compatibility with existing models, as it operates without requiring additional neural network modules or fine-tuning. Comprehensive evaluations of PP on LLaMA-2/3 and OPT models reveal that even minimal probing-using just 1.5% of FLOPs-can substantially enhance the efficiency of structured pruning of LLMs. For instance, when evaluated on LLaMA-2-7B with WikiText2, PP achieves a 2.56 times lower ratio of performance degradation per unit of runtime reduction compared to the state-of-the-art method at a 40% pruning ratio. Our code is available at https://github.com/Qi-Le1/Probe_Pruning.

Figures (5)

• Figure 1: Probe Pruning (PP) is executed in four stages: (1) PP selects key samples and tokens from the layer-normalized hidden states, based on residual importance, to create a small yet crucial probe. (2) PP deploys this probe to run a few model layers ahead and obtains the probe's intermediate hidden states. (3) PP integrates the probing states with historical states and uses the integrated states to calculate the pruning metric and prune weight channels. (4) PP performs full inference on the remaining weights.
• Figure 2: Jaccard Index of common pruning channels: comparing PP and Full-Batch Probing, and comparing fix-pruned model (without PP) and Full-Batch Probing for each batch.
• Figure 3: Performance of different probe combinations at a 40% pruning ratio.
• Figure 4: Importance-scaled fusion studies.
• Figure 5: Visualization of the $L^2$ norm for each token position of the input hidden states at layers 10 and 20 across the batch and feature dimensions. Experiments are conducted on the LLaMA-2-7B model using the WikiText2 dataset.