$D^2Prune$: Sparsifying Large Language Models via Dual Taylor Expansion and Attention Distribution Awareness

TL;DR

D^2Prune tackles the dual challenges of pruning large language models under distribution shift and preserving long-tail attention. It introduces a dual Taylor expansion to jointly model activation and weight perturbations for accurate mask selection and updates, and an attention distribution-aware dynamic weight update strategy that preserves critical Q/K/V patterns via a KL-divergence regularization term. Across diverse LLMs and even ViT models, it achieves state-of-the-art perplexity and zero-shot accuracy at high sparsity and demonstrates practical speedups, with LoRA fine-tuning further enhancing performance. This approach offers a principled, deployment-friendly path to compressing LLMs without sacrificing crucial attention dynamics, enabling broader accessibility of large-scale models. The method’s hyperparameters for activation perturbations ($\lambda_1$, $\lambda_2$) and the scaling factor $s$ provide tunable knobs to adapt pruning to task distribution gaps, underscoring its practical relevance for real-world systems.

Abstract

Large language models (LLMs) face significant deployment challenges due to their massive computational demands. % While pruning offers a promising compression solution, existing methods suffer from two critical limitations: (1) They neglect activation distribution shifts between calibration data and test data, resulting in inaccurate error estimations; (2) They overlook the long-tail distribution characteristics of activations in the attention module. To address these limitations, this paper proposes a novel pruning method, $D^2Prune$. First, we propose a dual Taylor expansion-based method that jointly models weight and activation perturbations for precise error estimation, leading to precise pruning mask selection and weight updating and facilitating error minimization during pruning. % Second, we propose an attention-aware dynamic update strategy that preserves the long-tail attention pattern by jointly minimizing the KL divergence of attention distributions and the reconstruction error. Extensive experiments show that $D^2Prune$ consistently outperforms SOTA methods across various LLMs (e.g., OPT-125M, LLaMA2/3, and Qwen3). Moreover, the dynamic attention update mechanism also generalizes well to ViT-based vision models like DeiT, achieving superior accuracy on ImageNet-1K.

Figures (10)

• Figure 1: Activation distributions and their shift between upstream and downstream data in LLaMA-2-7B (layer 10). During pruning, a small subset of C4 is used as calibration data (upstream), while downstream tasks such as WinoGrande are evaluated in a zero-shot setting. All inputs are formatted as 128-sample sequences with maximum embedding length. The $\mathrm{L}_2$ norms represent the average activation magnitude across samples.
• Figure 2: Visualizing of uniformized multi-head attention in LLaMA-2-7B (80% Sparsity, 128-token sequences from C4 as Calibraion input). We compare the 3D attention scores of the surfaces in the final Transformer layer for dense models against pruned counterparts (Wanda, SparseGPT, and our $D^2Prune$). $D^2Prune$ selectively updates projections to optimize for distribution consistency, successfully preserving these patterns with minimal distortion (lowest KL/RMSE).
• Figure 3: Illustration of the $D^2Prune$ framework. Using the update status of q, k, and v weights in each attention layer as the search space, we take the minimization of ppl as the objective to dynamically adapt the update strategy of the q, k, and v weights.
• Figure 4: Perplexity change of different pruning models with varying scaling factor $s$ on WikiText2
• Figure 5: Comparison of model perplexity on WikiText2 before and after introducing the scaling factor $s$ ("w/o s" denotes the setting without the scaling factor, while "s" indicates the setting with the scaling factor applied.