GPrune-LLM: Generalization-Aware Structured Pruning for Large Language Models

Abstract

Structured pruning is widely used to compress large language models (LLMs), yet its effectiveness depends heavily on neuron importance estimation. Most existing methods estimate neuron importance from activation statistics on a single calibration dataset, which introduces calibration bias and degrades downstream cross-task generalization. We observe that neurons exhibit heterogeneous distribution sensitivity, with distribution-robust neurons maintaining consistent rankings across datasets and distribution-sensitive neurons showing high cross-dataset ranking variance. Based on this, we identify two structural limitations in existing methods. First, ranking all neurons within a shared space causes distribution-sensitive neurons that strongly activate on calibration inputs to dominate, crowding out distribution-robust neurons critical for out-of-distribution tasks. Second, applying activation-based importance metrics uniformly can be unreliable. Distribution-sensitive neurons that infrequently activate on calibration data receive insufficient activation signal for accurate local ranking. To address these limitations, we propose GPrune-LLM, a generalization-aware structured pruning framework that explicitly accounts for neuron differences in cross-distribution behavior. We first partition neurons into behavior-consistent modules to localize ranking competition, then evaluate activation-based metric reliability per module according to distribution sensitivity and score magnitude. For modules where activation-based scoring is unreliable, we switch to an activation-independent metric. Finally, we adaptively learn module-wise sparsity. Extensive experiments across multiple downstream tasks demonstrate GPrune-LLM's consistent improvements in post-compression generalization, particularly at high sparsity, and reduced dependence on importance metric choice.

Figures (3)

• Figure 1: Two structural limitations of activation-based structured pruning. Middle: Cross-distribution behavioral heterogeneity. Per-module rank drift and Kendall correlation across datasets show that many modules maintain consistent importance rankings across distributions (distribution-robust), while some exhibit high rank drift and near-zero cross-dataset correlation (distribution-sensitive). Right: Cross-behavior ranking interference. When task-specific and general neurons compete within a shared ranking space, task-specific neurons that strongly activate on the calibration set dominate the ranking under that calibration, suppressing distribution-robust neurons that are broadly important across tasks. Left: Local ranking distortion. Neurons in distribution-sensitive modules that activate infrequently on calibration data receive insufficient activation signal for reliable local ranking.
• Figure 2: Overview of GPrune-LLM. FFN neurons are first partitioned into behavior-consistent modules by jointly modeling data-free parameter similarity and cross-dataset rank drift, with assignments refined via a differentiable objective. Each module's mean rank drift then serves as a sensitivity indicator of this module. Distribution-robust (low-drift) modules retain activation-based metrics, while distribution-sensitive (high-drift) modules with low activation scores switch to activation-independent alternatives, ensuring reliable within-module ranking. Finally, module-specific sparsity is adaptively learned, yielding a pruned LLM with improved generalization across diverse downstream tasks.
• Figure 3: Effect of threshold training epoch number in adaptive sparsity learning on WikiText2 perplexity at different retention ratios on LLaMA2-7B with Wanda-sp.