Improving Generalization in LLM Structured Pruning via Function-Aware Neuron Grouping

TL;DR

This work addresses the generalization gap in post-training structured pruning of large language models caused by calibration data distribution mismatch. It introduces Function-Aware Neuron Grouping (FANG), which segments FFN neurons into functionally coherent groups via context clustering, assigns group-specific importance with token-aware reweighting, retains a shared cross-context neuron group, and allocates sparsity adaptively based on functional complexity. Across multiple models (eg, LLaMA and Qwen) and sparsity levels, FANG improves downstream accuracy (1.5%–8.5% on average) while preserving perplexity, and achieves state-of-the-art results when combined with FLAP and OBC (denoted F-FANG and O-FANG). The approach leverages explicit functional specialization to maintain both specialized and general capacities, offering practical pruning efficiency with broad applicability to post-training pruning without fine-tuning. The framework is complemented by extensive ablations and extensions (Appendix) that demonstrate robustness and generality across architectures and downstream benchmarks like ARC, WinoGrande, BoolQ, HellaSwag, OpenBookQA, PIQA, and MMLU.

Abstract

Large Language Models (LLMs) demonstrate impressive performance across natural language tasks but incur substantial computational and storage costs due to their scale. Post-training structured pruning offers an efficient solution. However, when few-shot calibration sets fail to adequately reflect the pretraining data distribution, existing methods exhibit limited generalization to downstream tasks. To address this issue, we propose Function-Aware Neuron Grouping (FANG), a post-training pruning framework that alleviates calibration bias by identifying and preserving neurons critical to specific function. FANG groups neurons with similar function based on the type of semantic context they process and prunes each group independently. During importance estimation within each group, tokens that strongly correlate with the functional role of the neuron group are given higher weighting. Additionally, FANG also preserves neurons that contribute across multiple context types. To achieve a better trade-off between sparsity and performance, it allocates sparsity to each block adaptively based on its functional complexity. Experiments show that FANG improves downstream accuracy while preserving language modeling performance. It achieves the state-of-the-art (SOTA) results when combined with FLAP and OBC, two representative pruning methods. Specifically, FANG outperforms FLAP and OBC by 1.5%--8.5% in average accuracy under 30% and 40% sparsity.

Figures (3)

• Figure 1: Compared to traditional post-training pruning, our method mitigates calibration set distribution mismatch by pruning neurons based on functional grouping, reducing the risk of misestimating important neurons and improving generalization.
• Figure 2: Overview of our Function-Aware Neuron Grouping (FANG). For attention heads, we adopt baseline pruning methods such as OBC and FLAP. For FFN layers, neurons are grouped based on their functional roles, and a function-aware pruning strategy is applied. A shared neuron group is retained to preserve neurons contributing to multiple contexts. Additionally, block-wise sparsity is allocated according to functional complexity, with more complex blocks assigned lower sparsity.
• Figure 3: Execution Efficiency Comparison of Different Algorithms (LLaMA2-7B, 40% Sparsity).