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

TL;DR

This paper introduces Function-Aware Neuron Grouping (FANG), a novel post-training pruning method for LLMs that improves generalization to downstream tasks by preserving functionally critical neurons, achieving state-of-the-art results.

Contribution

FANG is a new pruning framework that groups neurons by function, adaptively allocates sparsity, and enhances downstream accuracy, addressing calibration bias in LLM pruning.

Findings

FANG outperforms existing methods by 1.5%-8.5% in accuracy at 30-40% sparsity.

FANG improves downstream task performance while maintaining language modeling capabilities.

Combining FANG with FLAP and OBC yields state-of-the-art results.

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.