Frustratingly Easy Task-aware Pruning for Large Language Models

TL;DR

This work addresses the challenge of pruning large language models without sacrificing task-specific capabilities. It introduces a simple, plug-in task-aware pruning framework that uses two calibration datasets, general-domain $\mathcal{D}_G$ and task-specific $\mathcal{D}_T$, to compute activation-weighted importance scores and fuse them into a mixed score for pruning. The method partitions parameters into shared, general-only, and task-only groups via activation-norm differences with a threshold $\alpha$, enabling selective pruning that preserves task-relevant representations. Experiments on Qwen-3 (32B) show consistent improvements over a baseline Wanda pruning across both unstructured and structured sparsity settings, validating the approach's practicality and robustness.

Abstract

Pruning provides a practical solution to reduce the resources required to run large language models (LLMs) to benefit from their effective capabilities as well as control their cost for training and inference. Research on LLM pruning often ranks the importance of LLM parameters using their magnitudes and calibration-data activations and removes (or masks) the less important ones, accordingly reducing LLMs' size. However, these approaches primarily focus on preserving the LLM's ability to generate fluent sentences, while neglecting performance on specific domains and tasks. In this paper, we propose a simple yet effective pruning approach for LLMs that preserves task-specific capabilities while shrinking their parameter space. We first analyze how conventional pruning minimizes loss perturbation under general-domain calibration and extend this formulation by incorporating task-specific feature distributions into the importance computation of existing pruning algorithms. Thus, our framework computes separate importance scores using both general and task-specific calibration data, partitions parameters into shared and exclusive groups based on activation-norm differences, and then fuses their scores to guide the pruning process. This design enables our method to integrate seamlessly with various foundation pruning techniques and preserve the LLM's specialized abilities under compression. Experiments on widely used benchmarks demonstrate that our approach is effective and consistently outperforms the baselines with identical pruning ratios and different settings.

Figures (3)

• Figure 1: The overall process of our framework for LLM pruning. The bottom part presents the LLM. The top presents the pruning process for a particular parameter matrix in the LLM, where the general and task-specific calibration data (marked by the blue and purple colors, respectively) are used to compute the pruning scores of each parameter.
• Figure 2: Performance of the pruned Qwen-3 (32B) on MedQA at 50% sparsity under different $\alpha$ values (blue curve). The dashed red line denotes the Wanda baseline. The parameter $\alpha$ controls how strictly the model separates shared and task-specific channels during pruning.
• Figure 3: Distribution of parameter groups across Transformer layers in Qwen-3 (32B). The x-axis represents the layer index, and the y-axis shows the proportion of parameters categorized as shared (blue), general-only (orange), and task-only (red).