BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Authors: Lian Liu
Xiandong Zhao
Guanchen Li
Dong Li
Mengdi Wang
Yinhe Han
Xiaowei Li
ying wang
Venue: ICML 2025 poster (2025)
Keywords: LLM Pruning; Automaic Framework
Primary area: general_machine_learning->hardware_and_software
TL;DR: We identify the existing bias in existing pruning metric and propose a novel automatic optimization framework to get the appropriate pruning metric for better pruning results.

Abstract: One-shot post-training pruning enhances the deployment of billion-scale large language models (LLMs), with the pruning metric playing a pivotal role in determining which weights to remove. However, existing metrics underperform due to their reliance on a simple symbolic combination of weights and activations, overlooking imbalanced weight magnitudes and the disproportionate influence of activation outliers.
To overcome these limitations, we introduce \textbf{BaWA}, a novel pruning metric that systematically \underline{Ba}lances \underline{W}eight and \underline{A}ctivation distributions for more effective pruning.
BaWA introduces two key innovations: \textbf{magnitude normalization}, which mitigates weight imbalance across channels for fairer pruning decisions, and \textbf{outlier regularization}, which reduces the impact of activation outliers, ensuring more appropriate channel prioritization. 
To further enhance its effectiveness, BaWA incorporates an efficient and automatic framework for optimizing normalization and regularization hyperparameters. Extensive experiments validate BaWA as a state-of-the-art (SOTA) pruning metric. For instance, applying BaWA to induce 2:4 sparsity in Mistral-7B reduces perplexity in language comprehension by 2.49 and improves average downstream task accuracy by 3.08\%, outperforming the previous SOTA method Wanda.

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Lian Liu 1234
Xiandong Zhao Guanchen Li"
Dong Li Mengdi Wang Yinhe Han 23 Xiaowei Li Ying Wang23 23
Abstract
2023a;b) has revolutionized AI research, demonstrating exceptional capabilities across diverse and complex tasks. However, their vast parameter sizes, often in the billions, pose significant hardware constraints, limiting their practi- cal deployment, especially in resource- constraine environ- ments (Liu al., 2025b). To mitigate this challenge. model compression techniques such quantization(Deetmerss et al., 2022; Xiao et al., 2023; Lin et al., 2023; Yuan et al., 2023; Frantar et al., 2023; Shao et al., 2023; Liu et al., 2024; 2025a) and pruning (Li et al., 2023a; Kwon et al., 2022; Frantar & Alistarh, 2022; Xia et al., 2023) have been ex- plored to reduce resource demands while maintaining perfor- mance. This paper focuses on unstructured pruning, which compresses models by zeroing out weights and benefits from sufficient hardware acceleration supports.
One-shot post-training pruning enhances the de- ployment of billion-scale large language mod- els (LLMs), with the pruning metric playing a pivotal role in determining which weights to re- move. However, existing metrics underperform due to their reliance on a simple symbolic com- bination of weights and activations, overlooking imbalanced weight magnitudes and the dispropor- tionate influence of activation outliers. To over- come these limitations, we introduce BaWA, a novel pruning metric that systematically Balances Weight and Activation distributions for more ef- fective pruning. BaWA introduces two key inno- vations: magnitude normalization, which miti- gates weight imbalance across channels for fairer pruning decisions, and outlier regularization, which reduces the impact of activation outliers, ensuring more appropriate channel prioritization. To further enhance its effectiveness, BaWA in- corporates an efficient and automatic framework for optimizing normalization and regularization hyperparameters. Extensive experiments validate BaWA as state-of-the-art (SOTA) pruning met- ric. For instance, applying BaWA to induce 2:4 sparsity in Mistral-7B reduces perplexity in lan- guage comprehension by 2.49 and improves aver- age downstream task accuracy by 3.08%, outper- forming the previous SOTA method Wanda.
Due to the scale of LLMs, one-shot post-training pruning has gained attention for its ability to efficiently remove weights without relying on gradient backpropagation or fine-tuning (Frantar & Alistarh, 2023; Sun al., 2023; Zhang al., 2024; Dong et al.). While effective, existing methods remain constrained by suboptimal pruming metrics. At the cost of additional mask training. MaskILLM (Fang et al.) demonstrated that well-designed pruning metric enables weight sparsity with minimal performance loss. However, existing one-shot LLM pruning methods, such as Wanda (Sun al., 2023) and Pruner-Zers (Dong et al.), rely on simple symbolic combination of weights and ac- tivations determine pruning decisions. While intuitive, they overlook the heterogeneity in weights and activations, leading to suboptimal pruning masks and significant perfor- mance degradation in pruned models. We summarize the key factors that influence pruning effectiveness as follows:
1. Introduction
Imbalanced Weight Magnitude Distribution. LLMs' weight magnitudes distribution is imbalanced, with certain input output channels exhibiting abnormally large small magnitudes. This imbalance leads to sub-optimal pruning decisions, as weights within a channel are either predominantly preserved or pruned, regardless of their ac- tual importance to the model's 's performance.
The remarkable performance of large language models (LL.Ms)qwee 2024; Dettmers et al., 2022; Touvron et al.,
'Advanced Micro Devices, Inc "Institute of Computing Tech- nology, CAS, China University of Chinese Academy of Sciences, Beijing, China "Zhongguaneun Laboratory. Correspondence to: Ying Wang @wangying2009@icc..ee ict.ac.en>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Disproportionate Influence of Outliers. few activation outliers can disproportionately inflate channel's norm,

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
leading to biased pruning. For instance, fewer than I % of outliers can increase channel's norm by over 5x, causing excessive pruning of outlier-free channels.
model pruning methods (Hassibi et al., 1993; Han al., 2015b). These methods be broadly categorized into structured and unstructured pruning.
To address these limitations, we propose BaWA, a novel pruning mask selection method that systematically Balances the contributions of Weight and Activation distributions. BaWA introduces two key innovations. Firstly, to address the imbalanced weight magnitude distribution, BaWA nor- malizes the weight magnitudes across both input and output channels, contributing to a fairer pruning mask selection. Furthermore, to mitigate the disproportionate influence of outliers, BaWA introduces learnable power factors that con- trol the impact of outliers on the pruning metric, ensuring that channels without outliers are not unevenly penalized. Moreover, BaWA employs a zeroth-order gradient-based optimization strategy to efficiently and automatically search for the optimal hyper-paramettes of normalization and reg- ularization, enabling the method to identify better pruning masks in just a few minutes. Notably, the reliable pruning metric provided by BaWA is orthogonal to conventional weight adjustment methods, and the combination of the two can bring better pruning outcomes. Through extensive experiments, we demonstrate that BaWA significantly out- performs existing state-of-the-ar pruning methods across a variety of LLMs and language benchmarks. For instance, applying BaWA to induce 2:4 sparsity in Mistral-7B reduces language comprehension perplexity by 2.49 and improves average downstream task accuracy by 3.08% compared to the previous SOTA method Wanda.
Structured pruning (An et al.. 2024: Xia et al.; Ma et al., 2023) removes entire substructures or weight groups-such as layers (Ling et al., 2024), FFN neurons (Ma et al., 2023), MHA heads, or embedding dimensions (Sreenivas al., 2024) enabiig hardware-agnostic efficiency gains. How- ever, due to its coarse gaarse- nature, structured pruning often leads to significant accuracy degradation, typically limiting the applicable sparsity ratio to 15%-30% Post- pruning fine-tuning can help recover performance at higher sparsity levels.
In contrast, unstructured pruning (Frantar & Alistarh, 2023; Sun et al., 2023) eliminates individual weight elements and stores them in compressed format. When combined with decompression techniques (for memory optimization) or hardware acceleration (e.g., 2:4 sparse tensor cores), un- structured sparsity can also deliver substantial efficiency improvements. Thanks to its fine-grained approach, unstruc- tured pruning generally preserves model accuracy more ef- fectively, allowing sparsity ratios exceeding 50%-or even stricter patterns like 2:4 sparsity. Further training techniques, such as PEFT (Lu et al., 2024) and STE (Ma et al., 2024), can also enhance pruned models.
Given these advantages, this paper focuses on optimizing unstructured pruning metrics.
2.2. Pruning Metric
In summary, the key contributions of this paper are:
The one-shot post-training pruning process primarily com- prises two stages: construction. Previous studies (Frantar & Alistarh, 2023; Li et al., 2023b; Liu et al., 2021) mainly focus on efficient weight reconstruction methodologies, while simply select- ing the pruning mask based on weight magnitude. However, due to the presence of outliers, a larger magnitude does not necessarily indicate greater importance for a weight. As a result, directly pruning weights based solely on weight magnitude can result in the removal of important weights and lead a significant performance drop for pruned LLMss In fact, the effectiveness of pruning in large language mod- els (LLMs) critically depends on the pruning metric (Li et al., 2023a; Frantar & Alistarh, 2023; Fang et al.). which quantifies the importance of each weight and guides the de- cision which weights to remove. well-designed pruning metric is essential for maintaining model performance after pruning, as it directly influences the quality of the pruned model.
A comprehensive analysis of the limitations of current pruning metrics, highlighting the bias introduced by the symbolic combination of weight and activation.
pruning mask selection and weight re-
The introduction of BaWA, novel pruning metric that balances weight and activation distributions through magnitude normalization and outlier regularization.
An efficient optimization strategy for identifying op- timal hyper-parameters for normalization and regular- ization, enabling BaWA to achieve superior pruning performance with minimal computational overhead.
Experimental results demonstrate that BaWA success- fully improves the performance of pruned LLMs com- pared to the statc-of-the-art (SOTA) pruning methods.
2. Background & Motivation
2.1. LLM Pruning
As presented in Table 1, current pruning metrics utilize the symbolic combination of weight (Han et al., 2015a), activation (Sun et al., 2023; Zhang et al., 2024) and even
The growing complexity of Transformer- based language models, which now scale to hundreds of billions of parame- ters, has intensified the demand for effective and efficient
2

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Table I:
The existing pruning metrics tailored for LLMs. SC denotes the symbolic combination operation, G denotes the gradient, and a denotes the min-max scaling operation.
either abnormally large or small, creating a skewed distri- bution. This imbalance results in scenario where most weights within a channel are either entirely preserved or pruned during compression. For unstructured and N:M sparsity, this concentration of pruning decisions on spe- cific channels leads to sub-optimal sparsity patterns, as the pruning process fails to account for the heterogeneity in weight importance. This observation underscores the need for a pruning metric that explicitly considers the imbalanced nature of weight distributions, rather than relying on over- simplified assumptions.
Method
Pruning Metric S
SC
Mangnitude
Wel W
Wanda
GBL.M-rrnner
W ||Gylla
Pruner-Zero
W sy x W&ll x ((G))
BaWA
W |X,11
X
Disproportionate Impact of the Small Set of Outliers. When constructing pruning metrics, existing methods such as Wanda (Equation 1) compute channel-wise norms (e.g.. the f--norm) to evaluate the importance of weights. How- ever, our analysis reveals that small subset of outliers in activations can disproportionately influence these norm val- ues. As shown in Figure 1c, even when outliers constitute less than 1% of a channel's activations, they can inflate the channel's norm value by up to 5x compared to the average. This disproportionate impact introduces significant bias into the pruning metric, as channels with outliers are erroneously prioritized for preservation, while others are undervalued. Consequently, pruning decisions based such metrics be- come skewed, leading to sub-optimal sparsity patterns and degraded model performance.
gradient (Dong et al.; Das et al., 2023) to determine the importance of weights. For example, Wanda defines the pruning metric as:
Su = Wol -1x,112,
where W E CCnn is the weight of a linear layer and X E B.Q is the input activation with batch size B and sequence length Q. | is the absolute value operator and ||,l represents the lo norm of jth features aggre- gated across B Q different tokens. Wanda prunes those weights that have lower final scores S. GBLM--Prune (Das et al., 2023) further utilizes gradient to adjust the pruning metric; however, the improved performance is limited. To improve the pruning performance, Pruner-Zero (Dong et al.) proposes an automatic framework for searching symbolic pruning metrics using genetic programming.
It is necessary to mention that the two phenomena are widely present in various LLMs (qwe, 2024; Jiang et al., 2023; Tou- vron et al., 2023b). Consequently, we argue that balanced pruning metric must consider both the weight magnitude distribution and the impact of outliers in activations. Such metric should normalize the weight magnitudes across channels to mitigate the effects of imbalanced distributions and regulate the influence of outliers to ensure a fair eval- uation of each weight's importance. By balancing these factors, a more accurate and effective pruning metric can be developed, enabling better pruning decisions and improved model performance.
2.3. Motivation
However, existing pruning metrics rely on simplistic, sym- bolic combinations of weight, activation, and gradient val- ues, failing to account for the intricate and heterogeneous distributions of weights and activations in LLMss This over- sight often leads to sub-optimal pruning decisions, resulting in significant performance degradation and inefficiencies in model compression.
To address this gap, we conduct comprehensive analysis and visualization of weight and activation distributions in LLaMA-7B, as illustrated in Figure 1, Our investigation reveals two critical phenomena that challenge conventional assumptions and motivate the development of a more robust pruning methodology:
3. Methods
3.1. Magnitude Normalization
According to our observation in Section 2.3, the weight val- ues of different channels (including both input and output channels) are quite different. Consequently, when establish- ing the pruning metric, it is imperative to consider not only the per-channel norm values of activation but also the norm values within the weight. This comprehensive considera- tion ensures a more balanced approach to determining the pruning mask, which is critical for maintaining the pruning performance. We present the details of the input and output channel normalization as follows.
Imbalanced Distribution of Weight Magnitude in LLM. Recent studies on LLM compression often assume that weights follow a normal distribution, leading many to over- look the nuanced analysis of weight distributions (Xiao et al., 2023; Yin et al., 2023; Shao et al., 2023). However, as demonstrated in Figure la and 1b, weight magnitudes exhibit significant imbalances across channels. Specifically, certain input or output channels contain weights that are
3

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Weight Distribution
Weight Distribution
Activation Distribution
0.5
0033
622
2.11
Negnitude
613
5.00
13144
1.00
0034
nr
1312
So Inpor Charneel
on Chatine
1330 Output Chonnee
Output Chonine
Channel bidex
(a) layes, sett.ummo..gn
(h) lagers 10.sel_.tttn xppeo
(c) layens s..elf..ttt Raproj
Figure 1: The visualization of the weight and activation magnitude distribution in the 10th layer of LLaMA-7B. The z-axis of the 3-D figure represents the magnitude of the weights and activations. Observations: (I) The distributions vary greatly not only across the channel dimension of activations but also across the input and output channel dimensions of weights. The three figures below demonstrate that the norms of different channels vary significantly. (2) The norm of one channel is often heavily influenced by outliers, despite their small proportion. Figure (c) is particularly notable, demonstrating that several outliers lead to a significant norm. * denotes outliers.
Input Channel Normalization. In contrast to the activa- tion values, the magnitude of weight exhibits relatively more gradual variation. Consequently, as presented in Equa- tion (1), compared to the weight magnitude Wiyl the pro- posed pruning metric in Wanda mainly focuses on capturing the fluctuations in activation values |XX, This strategy emphasizes that the dynamic range of activation is more susceptible outliers, which can significantly influence the pruning performance. However, observations indicate that the per-channel norm value within the weights also ex- hibits considerable variation. Therefore, we argue that the pruning metric should not be exclusively skewed towards outliers in activation. It is equally important to take into account the magnitude of weights themselves. Accordingly, we introduce the no-m weight value for each input chan- nel mitigate the excessive focus on activation values. We incorporate input channel normalization and redefine the pruning metric as outlined below.
its focus more towards the magnitude of the weight val- ues themselves. The pruning metric with output channel normalization is as follows:
((OCN)
(Wol ||x,122
(3)
W
It is important to note that during unstructured and N:M pruning, we often prune each output channel by the same sparsity ratio, rendering the consideration of output channel normalization alone unproductive. Therefore, we need to jointly consider input normalization and output channel normalization form the final magnitude normalization metric S(MN) = s +S(CN)+((OO)) As illustrated in Figure 2a, adopting magnitude normalization helps to make the distribution of pruning mask more balanced on different channels, improving the performance of pruned LLMs.
3.2. Outlier Regularization
(ICN) s S
||,
Wll
Our analysis in Section 2.3 has revealed that a small set of outliers significantly influence the norm value of channel. According to our evaluation, when using Wanda as the pruning metric for LLaMA-7B (Touvron et al., 2023a), considerable number of channels were eliminated across various layers (even exceeding 10% in specific layers). That is to say, when adopting the typical norm value of each channel as a part of the pruning metric, channels devoid of outliers are susceptible to being inadvertently compromised during the pruning process (Zhang et al., 2024), as illustrated in the left part of Figure 2b.
Output Channel Normalization. Recent studies have revealed that outliers tend to occur in specific channels of activation (Yuan et al., 2023; Hooper et al.. 2024; Yao et al., 2022). We notice that the presence of outliers in specific channels is closely related to the magnitude of weights across different output channels. Therefore, for these channels, it is critical pay attention the inher- magnitude of the original weights. Building these insights, we propose output channel normalization, tech- nique that leverages the fn-mrmm of weight values for each output channel to alleviate the undue emphasis on activation values. Especially when the fo-m of an output channel substantially large, our proposed pruning metric shifts
The outlier values have been unduly emphasized in the prun- ing process (Wei et al., 2023). It is necessary to reduce the proportion of weights pruned from channels that do not contain outliers. To address this, we propose novel

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
0.8
-6
-3.2
3.56
0.53
040
0.85
1.29
46
1:1
1.1
2.12
2.2
2.0
1.8
6.9
7.73
1.12
1.08
0.99
1.74
11.2
18.8
8.1
2044
4.97
6.13
3.82
8:45
Wi
|wyl - | ll,
0:57
0.71
0:70
143
0.99
-2.8
3.35
051
5.6
2.46
3.43
||W
||w]l,
2.63
2.56
2:35
8.14
[W]]2
1Wy1 + [Wyl
||
| 11g
5.1
9.4
4.5
1.5
|
11
2.26
3:07
2.12
1.22
8
0.5
(a) Magnitude Normalization
(b) Outlier Regularization
Figure 2: Illustration of our proposed magnitude normalization and outlier regularization methods. The leftmost figure represents the naïve pruning approach, which prunes weights with smaller magnitudes to satisfy the 2:4 constraints. The changed pruned positions are highlighted with red boxes. (a) magnitude normalization re-evaluates the importance of weight magnitude by considering the effect of norm values across input and output channel dimensions. (b) Outlier regularization suppresses the effect of larger activation norms and leads to different weight pruning masks.
strategy, outlier regularization, to reduce the impact of out- liers. Outlier regularization aims diminish the variability in magnitude across different channels, thereby ensuring more equitable and balanced pruning strategy. Specifically, when defining the pruning metric, we introduce a novel pruning parameter 0, to serve as the power factor for the norm value. Specifically, the pruning metric with outlier regularization can be modified as:
parameter optimization algorithm. According to our evalua- tion, the proposed efficient search process can be completed in about 16 minutes for LLaMA2-70B within a single GPU,
Search Target. A straightforward way to prevent the search space from growing exponentially is to replace the global search target with a local search target. Previous works (Frantar Alistarh, 2023; Nagel et al., 2020; Fran- tar & Alistarh, 2022; Hassibi et al., 1993) optimize the pruning loss layer-by-layer, reducing the search space from C² to C- L. C represents the search complexity for one layer and L is the number of linear layers in one model. However, this approach may lack consideration of the in- teraction between layers. Other works (Shao et al., 2023; Li et al., 2020) employ the idea of block-wise optimization to incorporate interaction between layers within a single block. Each transformer block is structurally identical and relatively independent, making it well-suited for holistic optimization. Suppose Transformer block contains lin- ear layer, the search space is C L/m. We formulate the search target as follows:
I i
||2
(4)
The proportion of unreasonable norm values generated by outliers not solely present in activation. As revealed in previous work (Nrusimha et al., 2024), outliers also exist in the distribution of weight values. Therefore, we apply out- lier regularization not only to the norm of activation but also to the per-channel norm values of the weights. Incorporating the two methods previously discussed, we can formulate the final pruning metric, BaWA, as follows:
S = W
W
, i
IW5l
W
input channel normalization
output channel
(5)
O* arg min L(O; X),
(6)
where 01, 62. and 0g represent the power factors for the input channel, output channel, and activation norm, respectively.
L(O; X) |RMSNorm (F(W; X))
-RMSNorm (F(W MM; X)) (7)
3.3. Efficient Power Factor Search
= S > top,
(8)
Based on the proposed pruming metric in Equation (5), each linear layer has three power factors, 0,, 02. and 0z. to be determined. As different layers often exhibit various distributions, the optimal values for the three factors may vary accordingly. As the number of layers increases, the search space expands exponentially, resulting in an ineffi- cient search process. We introduce the design of an efficient search method from two perspectives: search target and
where F represents the computation of a single transformer block, which contains self self-attention and feed-forwart network. W is the weights for linear layers and X is the input activation of the block. @ consists of m sets of (0,, 02 0g). The pruning mask M is determined by our prun- ing metric S in Equation (5), where weights with smaller pruning metric are pruned by tops function. The symbol
5

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
represents the element-wise multiplication, and || ll2 is the Ez norm function. Different blocks may exhibit varying sensitivities to pruning, resulting in significant fluctuations in pruning error across blocks. To address this issue, we adopt the R.MSNrrm (Touvron et al., 2023a) after the out- put feature O, containing n elements, to balance the error disparities among different transformer blocks:
4. Experiments
4.1. Experimental Setup
Power Factor Search. Due to the presence of numerous outliers in the activation distribution and only a minimal number of outlier values in the weights, we empirically ini- tialize the power factors for the weight norm as 1, while 0.5 for the activation norm. To optimize these pruning param- eters, we utilize a typical SGD optimizer without weight decay. The learning rate is set as 0.2. Following the setting in Omniquant (Shao et al., 2023), we employ a calibration dataset consisting of 128 randomly selected 2048-token seg- ments from C4 (Raffel et al., 2020). We set the batch size as 16 with only epochs for the optimization process. We further explore the different optimization settings in our eval- uation. The whole optimization process can be completed in about 16 minutes for LLaMA2-70B (Touvron et al., 2023b) within a single GPU, which is less than that of SParseGPT.
RMSNorm(OO = RMS(O)
Oi
(9)
1 T
where RMS(O)
of
We demonstrate the effectiveness of this error normalization strategy in Appendix 4.5,
Models and Evaluation. We evaluate BaWA on widely adopted LLMs, including LLaMA (7B-65B) (Touvron et al., 2023a), LLaMA2 (7B-70B), Mistral-7B (Jiang al., 2023) and Qwen2-72B (qwe, 2024). We measure the perfor- mance of pruned models on seven zero-shot tasks. We further evaluate the five-shot performance of MMLU in the Appendix. For zero-shot evaluation, we use seven tasks from EleutherAl LM Harness (Gao al., 2021), including HellaSwag (Clark et al., 2018), PIQA (Bisk et al., 2020), ARC (Clark al., 2018), BoolQ (Clark al., 2019), RTE and Winogrande. Following previous works on LLM com- pression (Frantar & Alistarh, 2023; Xiao et al., 2023), we also evaluate the perplexity on the WikiText-2 (Merity et al., 2016).
Parameter Optimization Algorithm. The reliance on em- pirical or grid search strategies (Zhang et al., 2024; Lin et al., 2023), while common, can become inefficient as the search space expands. Gradient-basee optimization another typi- cal strategy for parameter tuning, especially in the studies of neural network compression (Li et al., 2020; Nagel et al., 2020; Shao et al., 2023). Unfortunately, obtaining M in- volves non -i-differentiabee topk function, which poses an obstacle to computing the gradient of power factors.
Recent studies utilize zeroth-order optimizer to enable non- differentiable parameter tuning, while also conserving GPU memory and improving efficiency (Malladi et al., 2024; Spall, 1992). Motivated by this research, we use zeroth- order optimizer to estimate the corresponding gradient for each power factor. Specifically, the gradients of + can be described as:
Baselines. We compare BaWA with various methods in- cluding (1) Magnitude- pruning (Han al., 2015b) that discards weights based on their magnitudes. (2) SparseGPT (Frantar & Alistarh, 2023) that utilizes second- order Hessian inverses to ascertain unimportant weights. (3) Wanda (Sun et al., 2023) that removes weights with the smallest magnitudes multiplied by the corresponding input activation norms. Moreover, we also compare BaWA with diverse existing pruning metrics, as depicted in Table 1, that use the symbolic combination of weight, activation and gradients to demonstrate the necessity of a balanced distribution of weight and activation.
X) = [(0 +z; X)= L(0 EZ; X)
2F
(10)
R 22 V)
Here, L denotes the error loss in Equation 6, 2 is a sampled data that satisfies 2 ~ (0, Ia), and is the perturbation scale. la an identity matrix with the dimension d = 3m, as each layer has three power factors to be determined. Fol- lowing setting in MeZO (Malladi et al., 2024), e in our experiments is set 0.01. We perturb the power factors at each step and obtain the corresponding gradients. We then update them based on these gradients. It is often nec- essary to repeat this iteration process times to efficiently converge to the optimal solution. The details on updating these parameters are shown in Algorithm 2 of Appendix E.
4.2. Language Modeling
4.2.1. BAWA ON VARYING SPARSITY.
We investigate the efficacy of BaWA when pruning LLMs with varying pruning rates. Table 2 shows that BaWA per- forms better than Wanda at different sparsity levels. Particu- larly, this improvement becomes increasingly evident as the
6

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Table 2: WikiText-2 perplexity performance of BaWA and Wanda for different LLMs at varying sparsity rates.
LLaMA-7B
LLaMA-13B
LLaMA2-70B
Qwen2-72B
Sparsity 60%
70%
80%
60%
70%
80%
60%
70%
80%
60%
70%
80%
Wanda
10.57
74.79
4.80e3
8.69
51.94
4.95e3
4.97
10.23
149.76
6.26
9.00
40.50
BaWA
10.00
57.84
3,95c3
7.67
33.83
4.10c3
4.56
8.71
125.71
6.03
8.17
31.89
sparsity level grows. For instance, we can achieve a redue- tion of 16.95 and 18.11 for LLaMA1-7B and 13B models under 70% sparsity. For large models such as LLaMA2- 70B and Qwen2-72B, BaWA can still achieve a reduction of 24.05 and 8.61 under 80% sparsity compared to Wanda.
4.2.2. PERPLEXITY WITH N:M SPARSITY.
As presented in Table 3, we compare the perplexity of var- ious pruned LLMs with existing LLM pruning methods. As can notice, BaWA consistently achieves better per- formance when compared with the other strong baselines without introducing weight reconstruction. Specifically, em- ploying BaWA as the pruning metric achieves up to a 1.23 reduction in perplexity in the 2:4 sparsity (LLaMA2-13B) relative to Wanda. This result suggests that exact and ef- fective sparse sub-networks exist for LLMs, and finds the appropriate metric is non-trivial but important for LLM prun- ing. Furthermore, BaWA also combine with existing weight reconstruction pruning methods. As our experimen- tal results depict, BaWA + ADMM can achieve the best performance. We also present more evaluations on diverse LLMs and sparse patterns in Appendix B.
Table 3: WikiText-2 perplexity of various pruned LLMs for N:M sparsity. We compare BaWA with existing prune-only methods, including Magnitude (Han et al., 2015a), Wanda (Sun et al., 2023), GBLM (Das et al., 2023), RIA (Zhang et al., 2024) and Pruner-Zero (Dong et al.), as well as weight reconstruction methods, includ- ing SparseGPT (Frantar Alistarh, 2023), DSnoT (Zhang al., 2023) and ADMM-Iter (Boža, 2024). We further com- bine BaWA with the SOTA weight reconstruction method ADMM-Iter, denoted as BaWA+ADMM. In our evaluation, bold denotes the best performance and underline denotes the second best performance.
4.3. Zero-Shot Tasks
We further compare BaWA with baselines on typical zero- shot tasks. As shown Table 10, we present the mean zero-shot accuracy of the pruned models within the LLaMA, LLaMA2 families, Mistral-7B and Qwen2-72B (the detailed performance of each task can be found in Table 11 and Table 12). We implement the Im-eval- harness (Gao et al., 2021) for all zero-shot tasks, with the report including both the accuracy results on each henchmark and the overall average accuracy. As it illustrates, BaWA consistently outperforms Wanda. For LLaMA2-70RE the pruned model under 50% sparsity achieves better accuracy than the original dense model. In certain cases, SparseGPT achieves higher accu- racy than our proposed BaWA, suggesting that weight re- construction can further enhance the performance of pruned L.MMs. Fortunately, our proposed pruning metric method can seamlessly integrate with novel weight reconstruction techniques (Zhang et al., 2023; Kwon al., 2022), which will be detailed in Appendix
ILaMA2
Mistral
Qwen2
Method
Sparsity
13B
70B
7B
72B
Dense
0%
4.57
3.12
5.25
4.94
Magnitude
4:8
6.76
5.54
9.21
8.14
SparseGPT
4:8
6.60
4.59
8.07
5.97
Wanda
4:8
6.55
4.47
8.41
5.86
GBLM
4:8
6.54
4.49
8.31
5.85
RIA
4:8
6.29
4.37
8.27
5.81
Pruner-Zero
4:8
6.75
4.45
8.11
5.85
DSnoT
4:8
6.43
4.41
7.93
5.79
ADMM-
4:8
6.37
4.35
7.79
5.77
BaWA
4:8
6.16
4.32
7.54
5.74
BaWA ADMM
4:8
6.07
4.24
7.36
5.65
Magnitude
2:4
8.33
6.33
13.57
8.15
SparseGPT
2:4
8.32
5.40
10.52
6.51
Wanda
2:4
8.27
5.16
12.37
6.31
GBLM
2:4
8.80
5.47
10.97
6.39
RIA
2:4
7.77
5.11
10.51
6.28
Pruner-Zero
2:4
7.41
4.81
10.49
6.33
4.4. Ablation Study
DSnoT
2:4
8.09
5.11
10.24
6.26
In this paper, we introduce multiple strategies to achieve a balanced orchestration of weight and activation. To verify the efficacy of each strategy, we conducted a comprehensive ablation study, shown in Table 5. We evaluate the im- pact of magnitude normalization and outlier regularization
ADMM-Iter
2:4
7.78
5.19
10.29
6.21
BaWA
2:4
7.13
4.84
9.88
6.14
BaWA ADMM
2:4
7.04
4.71
9.53
6.01
7

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Table 4: Mean zero-shot accuracies (%) of pruned LLaMA, LLaMA2, Mistral-7B and Qwen2-72B models on 7 zero-shot tasks. BaWA performs competitively against the prior best methods SparseGPT and Wanda.
LLaMA
LLaMA2
Method
Weight Update
Sparsity
7B
13B
30B
65B
7B
13B
70B
Mistral-7B
Qwen2-72B
Dense
0%
59.99
62.59
65.38
66.97
59.71
63.03
67.08
64,30
69.82
Magnitude
X
50%
46.94
47.61
53.83
62.74
51.14
52.77
60.93
55.87
60.66
SparseGPT
V
50%
54.94
58.61
63.09
66.30
56.24
60.57
67.28
59.34
68.11
Wanda
X 50%
55.13
59.33
63.60
66.67
56.24
60.04
67.03
58.93
66.41
BaWA
X 50%
55,27
59.97
64.12
67.21
57.02
60.67
67.81
60.17
69.11
Table 5: Ablation study proposed methods. We propose several strategies to improve the pruning mask metric for LLMs. We evaluate the performance improvements provided by these strategies, respectively. Only adopting Output Channel Normalization does not alter the pruning position, so we don't present its experimental results.
LLaMA2 & Qwen2 (50%)
LLaMA2 & Qwen2 (4 8)
I.LaM22 & Qwen2 (2 : 4)
Method
13B
70B
72B
13B
70B
72B
13B
70B
72B
Wanda
5.56
3.98
5.48
6.55
4.47
5.86
8.27
5.16
6.31
Input Channel Normalization
5.47
3.89
5.48
6.38
4.42
5.84
7.93
5.13
6.30
Magnitude Normalization
5.45
3.88
5.44
6.27
4.41
5.81
7.74
5.04
6.27
Outlier Regularization (0.5)
5.46
3.90
5.46
6.20
4.39
5.77
7.54
4.95
6.21
BaWA w/o Automatic Search
5.45
3.88
5,43
6,27
4.41
5.80
7.74
5.05
6,23
BaWA w/ Automatic Search
5.42
3.84
5,41
6.16
4.32
5.74
7.13
4.84
6.14
on LLaMA2-15 13B, LLaMA2-70B and Qwen2-72B, respec- tively. The power factor is set as 0.5 for the outlier regular- ization. Compared to Wanda, both magnitude normalization and outlier regularization reduce perplexity some extent. Specifically, magnitude normalization has a greater impact on performance improvement in unstructured pruning, while outlier regularization has a more significant effect on per- plexity under N:M sparsity conditions. is noteworthy that even without careful search, our BaWA method can easily outperform Wanda. Furthermore, automating the search for the power factors θ consistently reduces model perplexity. As presented in Appendix C, different transformer blocks, different linear layers, and different sparse patterns prefer quite different power factors. It is necessary to find the appropriate pruning parameters for every linear layer.
SparseGP1
Ir-0.01
Wanda
1n-0.05
Premlosity
1m0.2
1r=10
Calbration Sample
Epochs
Figure 3: analysis of pruning methods learning rate (Ir) and epochs on calibration samples.
The robustness Figure 4: Effect analysis on
for parameter optimization.
only one sample, satisfactory results could be achieved. The remarkable robustness exhibited by BaWA primarily stems from the outlier regularization strategy proposed, which effectively suppresses the impact of anomalies in the sampling process on the pruning results. We further analyze the performance variance when adopting different random seeds, as illustrated in Appendix D.
Power Factor Analysis. further analyze how different power factor settings affect the pruning performance. As shown in Table 6, different power factor settings exhibit significant sensitivity (#24.1% PPL variance). Additionally, the optimized power factor setting effectively reduces per- plexity by 12.3% compared to the best-fixed scale (0=0.5).
Syyee-PPrammetee Exploration We also compare results under different hyper-parameter conditions to investigate the impact of varying typer-parameten (learning rate and epochs) on the search outcomes during the parameter op- timization process. As shown in Figure 4, we notice that when the learning rate small (< 0.2), perplexity decreases as the number of epochs increases. However, a minimal learning rate leads slow convergence, significantly im-
4.5. Effectiveness Analysis
Roubustness Analysis. We analyze the impact of sample size on the performance of the pruned model for LLaMA2- 13B. As depicted in Figure 3, it is evident that BaWA is almost insensitive to the number of samples. Even with
8

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Table 6: Power Factor Analysis.
Scaling Strategy
01 (Input)
02 (Output)
Og (Activation)
PPL
A vs. Best Fixed (Fixed at 0.5)
Fixed Scales 0 = 0.1
0.1
0.1
0.1
8.92
+24. 1%
Fixed Scales 0 0.5
0.5
0.5
0.5
7.18
+0% % (baseline)
Fixed Scales 0 = 1.0
1.0
1.0
1.0
7.53
+4.9%
BaWA Optimized
0.42
0.51
0.38
6.30
-12.3 %
proving the search cost. Conversely, a larger learning rate makes maintaining a sustained model perplexity reduction with increased epochs difficult. Therefore, in this paper, we set the learning rate 0.2 and use only epoch 2 for pa- rameter search. While larger epochs could further enhance the performance of the pruned LLM, the additional search cost makes such an attempt unnecessary.
0.40
BoWA
0.35
Loss
0.20
Table 7: Time overhead with BaWA for pruning LLaMA-2 model family. The performance is evaluated on single NVIDIA A100 40GB GPU.
0.15
Method
7B
13B
70B
0.05
Wanda
0.4
0.8s
2.2
25
30
Layers
SparseGPT
3.1 min
5.2 min
22.3 min
BaWA w/o Search
0.5 s
0.8 s
2
Figure 5: The loss value for different layers
BaWA w/ Search
1.9 min
3.7 min
16.2 min
Pruning Overhead. We further analyze the computing cost of BaWA. Following Wanda (Sun et al., 2023), we report the time of total parameter optimization and pruning. Table 7 shows the quantitative wall-clock overhead evaluated on single NVIDIA 1100 GPU. We have evaluated the opti- mization overhead of BaWA on different models, both with and without parameter search. As shown in Table 7, one can obtain a bigh-performance pruned LLM within seconds. As depicted in Section 4.4, BaWA can easily outperform SOTA methods even without automatic optimization. Cer- tainly, searching for suitable pruning parameters requires several minutes, which costs less optimization time than SparseGPT. For instance, approximately 16 minutes are needed on single A100 GPU search for appropriate pa- rameters for the LLaMA-2-70B model. Therefore, users can adaptively adjust the strategy for parameter search based on their computational resources with limited searching costs.
initial layers can further reduce the loss in BaWA. Therefore, by finely adjusting the hyper-parameters during the search process, we hope achieve better performance than the results reported in this paper.
5. Conclusion
In this work, we propose an improved pruning mask selec- tion strategy for pruning LLMs. Our insight is that existing works ignore the imbalanced distribution in LLMs. To this end, we propose novel algorithm, BaWA, which introduces magnitude normalization and outlier regularization to alle- viate the impact of imbalanced distributions in both weight and activation. By introducing several learnable pruning pa- rameters, BaWA typically discovers more effective pruning masks and enhances the existing pruning method. As the introduced parameters are non oon-differentiable, we carefully employ a forward-forward optimization strategy find the solution efficiently. It is worth noting that the optimiza- tion process is highly efficient, requiring only 2 minutes for LLaMA2-7B and 16 minutes for LLaMA2-70B, for exam- ple. Experimental results demonstrate that BaWA efficient and superior to existing methods across various language benchmarks under the same sparse ratio.
Sensitivity Analysis on Loss Function. Figure compares the loss across different layers for Wanda and BaWA on LLaMA2-13B. The loss value evaluated according to Equation (6). It reveals that the initial layers exhibit smaller losses, indicating lower sensitivity and reduced need for meticulous selection of pruning metrics. Conversely, the later layers demonstrate significant loss differentiation based on the choice of pruning strategies. Similarly, this suggests that using lower learning rates for parameter search in the

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Acknowledgments
from seratch for large language models. In Forty-first International Conference on Machine Learning.
We sincerely thank the anonymous reviewers for their in- sightful suggestions. Furthermore, we thank Lu Tian and Emad Barsoum from AMD for valuable discussions on the proposed metric design. This work was partially supported by the National Key R&D Program of China (Grant No. 2023YFB4404400) and the National Natural Science Foun- dation of China (Grant No. 62222411), I). Ying Wang is the corresponding author (wangying2009 (et.00.c9eitt..ccc........
Fang, G., Yin, H., Muralidharan, S., Heinrich, G., Pool, J., Kautz, J., Molchanov, P., and Wang, X. Maskllm: Learn- able semi-structured sparsity for large language models. In The Thirty-eighth Annual Conference on Neural Infor- mation Processing Systems.
Frantar, E. and Alistarh, D. Optimal brain compression: A framework for accurate post-training quantization and pruning. Advances in Neural Information Processing Systems, 35:4475 44888, 2022.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none of which we feel must be specifically highlighted here.
Frantar, E. and Alistarh, D. Sparsegpt: Massive language models can be accurately pruned in one-shot. In Inter- national Conference on Machine Learning, pp. 10323- 10337. PMLR, 2023.
Frantar, E., Ashkboos, S., Hoefler, T., and Alistarh, D. Gptq: Accurate post-training quantization for generative pre- trained transformers. In The Eleventh International Con- ference on Learning Representations, 2023.
References
Qwen2 technical report. 2024.
An, Y., Zhao, X., Yu. T., Tang, M., and Wang, J. Fluctuation- based adaptive structured pruning for large language mod- els. In Proceedings of the AAAI Conference Artificial Intelligence, volume 38, pp. 10865-10873, 2024.
Gao, L., Tow, J., Biderman, S., Black, S., DiPofi, A., Foster, C., Golding, L., Hsu, J., McDonell, K., Muennighoff, N., et al. A framework for few-shot language model evaluation. Version vO. 0. I. Sept, pp. 8, 2021.
Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense in natural language. In Pro- ceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Han, S., Mao, H., and Dally, W. J. Deep compres- sion: Compressing deep neural networks with pruning. trained quantization and huffman coding, arXiv preprint 2015a.
Boža, V. Fast and optimal weight update for pruned large language models. arXiv preprint arXiv:2401.02933 2024.
Han, S., Pool, J.. Tran, J., and Dally, W. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015b.
Clark, C., Lee, K., Chang, M. M., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions. arXiv preprint a9Xiv: 10044, 2019.
Hassibi, B., Stork, D. G., and Wolff, G. J. Optimal brain surgeon and general network pruning. In IEEE interna- tional conference on neural networks, pp. 293-299. IEEE, 1993.
Clark, P., Cowhey, I., Etzioni, O., Khot, T., Sabharwal, A., Schoenick, C., and Tafjord, O. Think you have solved question answering? try are, the ai2 reasoning challenge. arXiv preprint arXiv: 054557 2018.
Hendrycks, D., Burns, C., Basart, S., Critch, A., Li, J., Song, D., and Steinhardt, Aligning ai with shared human values. Proceedings of the International Conference on Learning Representations (ICLR), 2021a.
Das, J., Ma, L.,, and Shen, Z. Beyond size: How gradients shape pruning decisions in large language models. arXiv preprint arXiv:2311.04900 049002, 2023.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M.. Song, D., and Steinhardt, J. Measuring massive multitask language understanding. Proceedings of the International Conference on Learning Representations (ICLR), 2021b.
Dettmers, T. Lewis, M., Belkada, Y., and Zettlemoyer, L. Gpt3. int8 (): 8-bit matrix multiplication for transformers at scale. Advances in Neural Information Processing Systems, 35:30318 -033322 2022.
Hooper, C., Kim, S., Mohammadzadeh, H., Mahoney, M. W., Shao, Y. S., Keutzer, K., and Gholami, A. Kvquant: Towards 10 million context length Ilm in- ference with kv cache quantization. arXiv preprint arXiv:2401.1807 2024.
Dong, P., Li, L.. Tang, Z., Liu, X., Pan, X., Wang, Q., and Chu, X. Pruner-zero: Evolving symbolic pruning metric
10

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Jiang, A. Q.. Sablayrolles, A., Mensch, A., Bamford, C., Chaplot, D. S., Casas, D. d. I., Bressand, F,, Lengyel, G. Lample, G., Saulnier, L., et al. Mistral 7b, arXiv preprint arXiv:2310.06825, 2023.
Liu, S., Chen, T., Chen, X., Shen, L.., Mocanu, D. C., Wang. Z.., and Pechenizkiy, M. The unreasonable effectiveness of random pruning: Return of the most naive baseline for sparse training. In International Conference on Learning Representations, 2021.
Kwon, W., Kim, S., Mahoney, M. W., Hassoun, J., Keutzer, K., and Gholami, A. A fast post-training pruning frame- work for transformers. Advances in Neural Information Processing Systems, 35:24101- 2210001-21166 16, 2022.
Lu, X., Zhou, A., Xu, Y., Zhang, R., Gao, P., and Li, H. Spp: sparsity- sparsity-preservent preserved parameter-effiiint fine-tuning for large language models. In Proceedings of the 41st In- teratiinnl Conference on Machine Learning, pp. 33254 33269, 2024.
Li, Y., Gong, R., Tan, X., Yang. Y. Hu, P., Zhang. Q. Yu, F, Wang, W., and Gu, S. Breeq: Pushing the limit of post-training quantization by block reconstruction. In International Conference on Learning Representations, 2020.
Ma, D., Chen, L., Wang, P., Xu, H., Li, H., Sun, L., Zhu, S., Fan, S., and Yu, K. Sparsity- accelerated training for large language models. Findings of the Association for Computational Linguistics ACL 2024. pp. 14696-14700 2024.
Li, Y., Niu, L., Zhang, X., Liu, K., Zhu, J., and Kang, Z. E-sparse: Boosting the large language model inference through entropy-based n: M sparsity. arXiv preprint 1xXiv2230.11922 2023a.
Ma. X., Fang. G., and Wang, X. LIm-pruner: On the struc- tural pruning of large language models. Advances in neural information processing systems, 36:21702-21720, 2023.
Li, Y., Yu, Y., Zhang, Q.. Liang, C., He, P., Chen, W., and Zhao, T. Losparse: Structured compression of large lan- guage models based on low-rank and sparse approxima- tion. In International Conference Machine Learning, pp. 20336 20350. PMLR, 2023b.
Malladi, S., Gao, T., Nichani, E., Damian, A., Lee, J. D., Chen, D., and Arora, S. Fine-tuning language models with just forward passes. Advances in Neural Information Processing Systems, 36, 2024.
Lin, J. Tang, J., Tang, H., Yang, S., Dang, X., and Han, S. Awq: Activation-awer aware weight quantization for I1m compression and acceleration. arXiv preprint arXiv:2306.00097, 2023.
Merity, S., Xiong, C., Bradbury, J., and Socher, Pointer sentinel mixture models. In International Conference on Learning Representations, 2016.
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can a suit of armor conduct electricity? new dataset for open book question answering. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 2381-2391, 2018.
Ling, G., Wang, Z., and Liu, Q. Slimgpt: Layer-wise structured pruning for large language models. Advances in Neural Information Processing Systems. 37:1071 122 107137, 2024.
Nagel, M., Amjad, R. A., Van Baalen, M., Louizos, C., and Blankevoort, T. Up or down? adaptive rounding for post-training quantization. In International Conference on Machine Learning, pp. 7197-7206. PMLR, 2020.
Liu, L., Xu, Z., He, Y., Wang, Y. Li, H., Li, X., and Han, Y. Drift: Leveraging distribution-bast based dynamic precision quantization for efficient deep neural network acceler- ation. In Proceedings of the 61st ACM/IEEE Design Automation Conference, pp. 1-6, 2024.
Nrusimha, A., Mishra, M., Wang, N., Alistarh, D., Panda, R., and Kim, Y. Mitigating the impact of outlier channels for language model quantization with activation regular- ization. arXiv preprint vrXiv:2404.03605 2024.
Liu, L.,, Cheng, L., Ren, H., Xu, Z., Pan, Y., Wang, M., Li, X., Han, Y., and Wang, Y. Comet: Towards practical w4a4kv4 Ilms serving. In Proceedings of the 30th ACM International Conference on Architectural Support for Programming anguages and Operating Systems, Volume 2, pp. 131-146, 2025a.
Raffel, C., Shazeer, N. Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21 (I (140):1-67,200 2020.
Liu, L., Zhao, S., Li, B., Ren, H., Xu, Z., Wang, M., Li. X., Han, Y., and Wang, Y. Make I1m inference affordable to everyone: Augmenting gpu memory with nim-dimm. In 2025 IEEE International Symposium on High Perfor- mance Computer Architecture (HPCA), pp. 1751-1765.
Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9)999-106 2021.
IEEE, 2025b.

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Shao, W., Chen, M., Zhang, Z., Xu, P., Zhao, L., Li, Z.,, Zhang, K., Gao, P., Qiao, Y., and Luo, P. Omniquant: Omnidirectionally calibrated quantization for large lan- guage models. In The Twelfth International Conference on Learning Representations, 2023.
Xiao, G., Lin, J., Seznec, M., Wu, H., Demouth, J., and Han, S. Smoothquant: Accurate and efficient post-training quantization for large language models. In International Conference on Machine Learning, pp. 38087-38099, PMLR, 2023.
Spall, J. C. Multivariate stochastic approximation using simultaneous perturbation gradient approximation. IEEE transactions on automatic control, 37(3)333-331, 1992.
Yao, Z., Yazdani Aminabadi, R., Zhang, M., Wu, X., Li, C., and He, Y. Zeroquant: Efficient and affordable post- training quantization for large-scale transformers. Ad- vances in Neural Information Processing Systems, 35: 27168-277 183,, 2022.
Sreenivas, S. T., Muralidharan, S., Joshi, R., Chochowski, M., Wahabaleshwarkar, A. S., Shen, G., Zeng. J., Chen, Z., Suhara, Y., Diao, S., et al. Llm pruning and distilla- tion in practice: The minitron approach. arXiv preprint arXiv:2408. 1796, 2024.
Yin, L., Wu, Y., Zhang, Z., Hsich, C.-Y., Wang, Y., Jia, Y., Pechenizkiy. M., Liang, Y., Wang, Z., and Liu, S. Out- lier weighed layerwise sparsity (owl): A missing secret sauce for pruning Ilms to high sparsity. In Conference on Parsimony and Learning (Recent Spotlight Track), 2023.
Sun, M., Liu, Z., Bair, A., and Kolter, J. Z. A simple and effective pruning approach for large language models. In The Twelfth International Conference on Learning Representations, 2023.
Yuan, Z., Niu, I.., Liu, J., Liu, W., Wang, X., Shang, Y., Sun, G., Wu, Q., Wu, J., and Wu, B. Rptq: Reorder-based post- training quantization for large language models. arXiv preprint arXiv:2304.01089 2023.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., al. Llama: Open and efficient foundation lan- guage models. arXiv preprint arXiv:2302.139711 2023a.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can machine really finish your sentence? Proceedings of the 57th Annual Meeting of the Asso- ciation for Computational Linguistics, pp. 4791 4800,, 2019.
Touvron, H., Martin, L.., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXiv:2307.09288, 2023b.
Zhang, Y.. Zhao, L., Lin, M.. Yunyun, S., Yao. Y., Han, X., Tanner, J., Liu, S., and Ji, R. Dynamic sparse no training: Training-fre free fine fine-tuning for sparse llms. In The Twelfth International Conference on Learning Representations, 2023.
Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bowman, S. Glue: A multi-task benchmark and analysis platform for natural language understanding. In Pro- ceedings of the 2018 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp- 353-355, 2018.
Zhang, Y., Bai, H., Lin, H., Zhao, J., Hou, L.,, and Can- nistraci, C. V. Plug-and-play An efficient post-training pruning method for large language models. In The Twelfih International Conference on Learning Representations, 2024.
Wei, X., Zhang, Y., Y., Zhang, X., Gong, R., Guo, J., and Liu, X. Outlier suppressiont: Accurate quantization of large language models by equivalent and effective shifting and scaling. In Proceedings of the 2023 Conference on Empirical Methods Natural Language Processing, pp. 1648-1665 2023.
Xia, H., Zheng, Z., Li, Y., Zhuang, D., Zhou, Z. Qiu, X., Li, Y., Lin, W., and Song, S. L. Flash-llm: Enabling cost-effective and highly-efficient large generative model inference with unstructured sparsity. Proceedings of the VLDB Endowment, 7(2):211-224 2023.
Xia, M., Gao, T., Zeng. Z., and Chen, D. Sheared Ilama: Accelerating language model pre-training via structured pruning. In The Twelfth International Conference on Learning Representations.
12

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
A. Hardware Efficiency
Table 8: GEMM Speedup with 2:4 sparsity for one transformer block of LLaMA-65B, The performance is evaluated on single NVIDIA A100 40GB GPU.
Layer
Dense GEMM
Sparse GEMM
Speedup
Q/K/V
0.270 ms
0.163 ms
1.66x
Att_Out
0.097 ms
0,072 ms
1.35
Up/Gate
0.487 ms
0.285 ms
1.71x
Down
0.247 ms
0.176 ms
1.40x
To illustrate the enhanced efficiency of pruned L.Ms,, we present the actual speedup on different operators for LLaMA-65B
on a single A100 40GB GPU. Due to the lack of support for the 4:8 sparsity pattern in NVIDIA Ampere architecture, we only measure the latency of GEMM with 2:4 sparsity. As shown Table 8, the model pruned by BaWA achieves 1.58x speedup over dense FP16 GEMM baseline for the token generation stage. It noteworthy that pruned LLMs can achieve acceleration not only GPUs but also modern CPUs such AMD Ryzen 7 PRO 5850U.
B. Detailed Results
Table 9: WikiText perplexity of various pruned LLMs. BaWA performs competitively against prior best methods SparseGPT and Wanda, specifically for N:M sparsity.
ILLMA
LLaMA2
Method
Weight Update
Sparsity
7B
13B
30B
65B
7B
13B
70B
Mistral-7B
Qwen2-72B
Dense
0%
5.68
5.09
4.77
3.56
5.12
4.57
3.12
5.25
4.94
Magnitude
50%
17.29
20.21
7.54
5,90
14.89
6.37
4.98
7.87
7.84
SparseGPT
50%
7.22
6.21
5.31
4.57
6.51
5.63
6.46
5.51
Wanda
50%
7.26
6.15
5.24
4.57
6.42
5,56
6.54
5.48
BaWA
50%
7.08
6.01
5.06
4.40
6.23
5.37
3.84
6.39
5.41
Magnitude
4:8
16.84
13.84
7.62
6.36
16.48
6.76
5.54
9.21
8.14
SparseGPT
4:8
8.61
7.40
6.17
5.38
8.12
6.60
4.59
8.07
5.97
Wanda
x
4:8
8.57
7.40
5.97
5.30
7.97
6.55
4.47
8.41
5.86
BaWA
X
4:8
8,08
6.89
5.62
4,94
7.36
6,16
4.32
7.54
5.74
Magnitude
X
2:4
42.13
18.37
9.10
7.11
54.59
8.33
6.33
13.57
8.15
SparseGPT
2:4
11.00
9.11
7.16
6.28
10.17
8.32
5.40
10.52
6.51
Wanda
2:4
11.53
9.58
6.90
6.25
11.02
8.27
5.16
12.37
6.31
BaWA
2:4
10.32
7.94
6.37
5.61
7.13
6.14
Results WikiText2. Table 9 shows the Wiki'eex22 perplexity various LLMs with different sparsity configurations, including 50% sparsity, 2:4 and 4:8 sparse patterns. As one notice, BaWA always achieves the best performance on diverse LLMs and sparse patterns.
Results on Zero-Shot Tasks. For zero-shot results in Section 4.3, the seven evaluated zero-shot tasks are: BoolQ (Clark et al., 2019), RTE (Wang et al., 2018). HellaSwag (Zellers et al., 2019). WinoGrande (Sakaguchi et al., 2021), ARC Easy and Challenge (Clark al., 2018), and OpenbookQA (Mihaylov al., 2018). this section, we also present the mean zero-shot accuracies of pruned LLMs under N:M sparsity, as shown in Table 10. To offer a more nuanced demonstration of the advantages of BaWA across various datasets, we present the performance evaluations for each zero-shot task of LLaMA2-133 13B and Mistral-7B, as shown in Table 11 and 12. Compared to Wanda and SparseGPT, BaWA demonstrates more consistent performance across various tasks, indicating its robustness. Additionally, BaWA shows high accuracy in the ARC tests, suggesting its strong adaptability to the task. Overall, utilizing BaWA can effectively enhance the precision of pruned LLMs in zero-shot tasks.
Results on MMLU Validation, To further validate the adaptability of BaWA across diverse datasets, we conduct additional
13

---

BaWA:
Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Table 10: Mean zero-shot accuracies (%) of pruned LLaMA, LLaMA2, Mistral-7B and Qwen2-72B models on 7 zero-shot tasks. BaWA performs competitively against the prior best methods SparseGPT and Wanda.
LLaMA
LLaMA2
Method
Weight Update
Sparsity
13B
30B
7B
13B
70B
Mistral-7B
Qwen2-72B
Dense
0%
59.99
62.59
65.38
66.97
59.71
63.03
67.08
64.30
69.82
Magnitude
X
50%
46.94
47.61
53.83
62.74
51.14
52.77
60.93
55.87
60.66
SparseGPT
V
50%
54.94
58.61
63.09
66.30
56.24
60.57
67.28
59.34
68.11
Wanda
X
50%
55.13
59.33
63.60
66.67
56.24
60.04
67.03
58.93
66.41
BaWA
X
50%
55.27
59.97
64.12
67.21
57.02
60.67
67.81
60.17
69.11
Magnitude
X
4:8
46.03
50.53
53,53
62.17
50.64
$2.89
60.28
54,25
60.15
SparseGPT
/
4:8
52.80
55.99
60.79
64.87
53.80
59.18
65.84
56.51
66.42
Wanda
x
4:8
52.76
56.09
61.00
64.97
52.49
58.00
66.06
55.31
66.76
BaWA
X
53.12
56.33
61.62
65.64
53.65
58.38
66.31
57.24
67.34
Magnitude
X
44.73
53.16
61.28
45.58
49.20
59.95
49.68
60.39
SparseGPT
/
2:4
50.60
53.22
58.91
62.57
50.94
54.36
63.89
50.88
66.57
Wanda
X
2:4
48.53
52.30
59.21
62.84
48.75
53.19
64.14
50.15
66.16
BaWA
X
2:4
49.91
53.33
59.87
62.90
49.98
54.61
65.03
53.23
67.29
Table 11: Detailed analysis on zero-shot tasks for LLaMA2-133.
Method
Sparsity
BoolQ
RTE
HellaSwag
WinoGrande
ARC-e
ARC-c
OBQA
Avg
Dense
0%
80.61
65.34
60.00
72.37
79.41
48.37
35.20
63.04
Magnitude
50%
57.71
55.6
54.43
65.35
70.58
38.31
27.40
52.77
SparseGPT
50%
81.38
66.06
56.09
71.51
75.20
41.72
32.00
60.57
Wanda
50%
81.28
60.65
57.05
70.64
75.76
42.75
32.20
60.04
BaWA
50%
80.76
$7.76
56.87
70.72
76.56
44.62
32.00
60.67
Magnitude
4:8
63.39
57.76
53.97
64.72
68.52
35.84
26.00
52.89
SparseGPT
4:8
80.58
64.62
51.94
71.98
73.70
40.87
30.60
59.18
Wanda
4:8
79.60
60.06
52.32
69.53
73.95
40.36
30.20
58.00
BaWA
4:8
80.12
62.45
53.58
68.35
74.07
38.91
31,20
58.38
Magnitude
2:4
65.81
53.79
50.12
62.12
57.53
31.83
23.20
49.20
SparseGPT
2:4
77.89
55.26
46.97
69.40
69.37
35.03
26.60
54.36
Wanda
2:4
75.26
56.68
46.43
66.77
68.35
34.47
24.4
53.19
BaWA
2:4
78.26
56.32
48.5
66.93
70.79
35.49
26.00
54.61
tests on the MMLU dataset (Hendrycks et al., 2021a;b). As shown in Table 13, we compared the performance of various methods on the MMLU dataset. Specifically, have evaluated the 5-shot performance on the LLaMA-30B model, with unstructured 50%, 4:8 and 2:4 sparsity. The experimental results indicate that, average, BaWA achieves higher accuracy than existing SOTA methods. This further confirms the effectiveness of BaWA, and effectively demonstrates that pruned models are task-agnostic and generalizable to any downstream task.
C. Visualization
We have selected several different layers from LLaMA2-7B and analyzed the power factor parameters obtained through BaWA, as detailed in Table 14, 15 and 16. As one can notice Table 14, the parameters for the weights were consistently close 1, which is attributed the absence of outliers in the weight distribution, thus negating the need for lower power factor to suppress outliers. Additionally, we can note that the parameters for transformer block 2 are closer to 0.5 compared to blocks 15 and 28, which is due to the lower loss in the carlier layers requiring less significant parameter adjustments, which satisfies our observation in Figure 5. Similarly, we have also observed that different operators exhibit distinct tendencies in the selection of activation power factors. For instance, the parameter for Mlp MOW. is notably higher than other values, which substantiates that the proportion of activation outliers in Mlp DOWN lower.
14

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Table 12: Detailed analysis on zero-shot tasks for Mistral-7B.
Method
Sparsity
BoolQ
RTE
HellaSwag
WinoGrande
ARC-e
ARC-c
OBQA
Avg
Dense
0%
83.43
67.51
61.25
73.80
80.83
50.43
32.80
64.30
Magnitude
50%
71.13
55.96
56.63
66.30
72.60
41.30
27.20
55.87
SparseGPT
50%
83.06
58.12
55.82
72.06
75.59
42.75
28.00
59.34
Wanda
50%
82.91
58.12
55.18
71.27
74.66
43.17
27.20
58.93
BaWA
50%
83.21
63.54
56.17
70.56
73.42
43.26
29.00
60.17
Magnitude
4:8
74.37
56.68
53.72
65.11
68.69
36.95
24.20
54.25
SparseGPT
4:8
79.11
60.65
51.64
68.67
37.97
25.60
56.51
Wanda
4:8
75.32
63.90
50.32
67.80
70.41
35.84
23.60
55.31
BaWA
4:8
81.25
65.70
52.36
68.27
71.80
36.69
24.60
57.24
Magnitude
2:4
62.65
$5.60
48.74
61.88
66.75
31.14
21.00
49.68
SparseGPT
2:4
71.41
54.15
46.61
64.48
66.62
32.08
20.80
50.88
Wanda
2:4
67.80
57.76
44.61
63.93
64.98
30.97
21.00
50.15
BaWA
2:4
74.16
60.65
46.78
66.61
67.00
33.02
24.40
53.23
Table 13: MMLU performance of LLaMA-30B with unstructured 50%, 4:8 and 2:4 sparsity.
Method
Sparsity
Humanities
Other
Stem
Avg
Dense
0%
56.00
64.31
67.47
47.00
58,33
Magnitude
50%
38.02
47.12
47.22
35,49
41.48
Wanda
50%
48.65
59.16
59.57
42.63
52.02
SparseGPT
50%
50.69
57.68
59.18
41.58
52.05
BaWA
50%
49.27
58.77
61.16
43.64
52.71
Magnitude
4:8
27.61
32.89
33.77
31.27
30.95
Wanda
4:8
43.57
54.20
53.33
40.72
47.42
SparseGPT
4:8
44.82
54.29
54.18
39.15
47.70
BaWA
4:8
44.06
54.23
55.05
39.77
47.76
Magnitude
2:4
28.84
34,86
35.72
31.11
32.19
Wanda
2:4
38.30
43.39
45,95
36.28
40.65
SparseGPT
2:4
39.68
44.60
46.52
36.23
41.49
BaWA
2:4
41.82
44.35
47.25
36.54
41.82
Table 14: Visualization of parameters acquired for different transformer blocks on LLaMA2-7B with unstructured 50%
sparsity.
Transformer Block-2
Transformer Block-15
Transformer Block-28
Operator
01
02
0g
01
02
03
01
02
0s
Attm_Q
1.00
1.00
0.43
1.00
1.00
0.32
0.99
1.00
0.38
Attn K
0.98
0.97
0.44
0.98
0.99
0.38
0.98
1.00
0.34
Attn_V
0.97
1.12
0.37
1.02
1.06
0.35
1.02
1.08
0.33
Attn.O
1.01
0.97
0.43
1.02
0.99
0.37
1.00
0.98
0.35
Mlp GATE
0.98
1.01
0.43
1.02
0.99
0.37
1.00
0.98
0.35
MIp_UP
0,95
0.32
1.00
1.00
1.01
0.96
0.38
Mlp DOWN
1.04
1.06
0.54
1.00
0.99
0.48
0.99
0.98
0.45
15

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Table 15: Visualization of the parameters acquired for different transformer blocks on LLaMA2-7B with semi-structured 4:8
sparsity.
Transformer Block-2
Transformer Block-15
Transformer Block-28
Operator
01
0g
01
02
01
02
Attn_Q
0.52
0.48
0.41
0.54
0.43
0.30
0.52
0.48
0.36
Attn K
0.49
0.41
0.45
0.56
0.61
0.37
0.48
0.49
0.32
Attn_V
0.53
0.59
0.36
0.52
0.49
0.38
0.54
0.58
0.32
Attn.O
0.59
0.51
0.41
0.53
0.52
0.36
0.52
0.72
0.34
MIP.GATE
0,55
0.53
0,42
0.43
0,52
0.36
0.48
0.62
0.36
MIp_UP
0.51
0.49
0.29
0.49
0.56
0.38
0.49
0.55
0.39
DOWN
0.41
0.51
0.52
0.58
0.56
0.61
0.42
Table 16: Visualization of the parameters acquired for different transformer blocks on LLaMA2-7B with semi-structured 2:4 sparsity.
Transformer Block-2
Transformer Block-15
Transformer Block-28
Operator
0,
02
0g
01
02
0.
01
0₂
0₂
Attn_Q
0.42
0.45
0.41
0.34
0.58
0.25
0.48
0.41
0.34
Attn K
0.41
0.46
0.40
0.51
0.56
0.32
0.45
0.44
0.32
Attn_V
0.51
0.51
0.36
0.43
0,62
0.31
0.51
0.50
0.31
Attn_O
0.53
0.58
0.40
0.42
0.52
0.34
0.47
0.51
0.32
MIp_GATE
0.51
0.51
0.38
0.46
0.53
0.36
0.42
0.51
0.25
Mlp UP
0.42
0.28
0.44
0.54
0.43
0.51
MIp.DOWN
0.39
0.51
0.50
0.48
0.51
0.40
0.51
0.52
0.38
We then compare the power factors acquired from different sparsity patterns, as indicated in Table 14, 15, and 16. noteworthy that when we adopt structrrue N:M sparsity, the power factors for weight can be greatly reduced. This due to the similar distributions among adjacent channels, as revealed in Figure 1. Furthermore, the power factor for activation should be lower when adopting BaWA on N:M sparsity.
D. Robustness Analysis
In this part, we perform a robustness analysis different pruning methods. Here, our focus primarily on the pruning process of BaWA, rather than the parameter search process. As shown in Table we report the mean and standard deviation under random seeds for sparse LLaMA1 and LLaMA2 family models 50% and 60% sparsity. As one can
Table 17: WikiText-2 perplexity performance of BaWA for sparse LLaMA-1 and LLaMA2 family models at 50% and 60% sparsity. We report the mean and standard deviation under random seeds. Note that, when the standard deviation is less than 0.01, we denote it as 0.00, with the sign * However, the actual standard deviation is not equal to zero.
LLaMA-1
LLaMA2
Method
Sparsity
7B
13B
30B
7B
13B
Dense
0%
5.68 (#-0.00)
5.09 (4-0.00)
4.77 (#:0.00)
5.12 (+:0.00)
4.57 (4:0.00)
SparseGPT
50%
7.26 (+:0.04)
6.24 (+0.004
5.33 (#:0.08)
6.52 (#0.044
5.63 (+00001)
Wanda
50%
7.25 (1:0033)
6.15 (#00001)
5.24 (#:0.3)
6.44 (#00001)
5.57 (1.0.03
BaWA
50%
(1:0.03)
5.95 (1:0.01)
5.01 (4:0.01)
6.23 (4:0.00)
5.37 (4:0.00)
SparseGPT
60%
10.34 (100.14)
8.44 (10.15)
6.69 (10.033
0.58 (10.11)
7.71 (10.11))
Wanda
60%
10.62 (#:0.05)
8.72 (#:0.03
6.54 (#0.01)
(#0.01)
(1:0.7)
BaWA
60%
9.80 (+:0.02)
7.60(+10001
6.15 (+0000)
8.78 (40.01)
6.80(=0.01)
16

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Algorithm 1 PyTorch code for the pruning process of BaWA Algorithm 2 Pytorch code for the proposed zeroth-order optimization
weight matrix ou_ c_in);
X: input matrix (N c_in);
layer:
the specific transformer block that needs to be optimized;
9: learned parameters for BaWA;
8: desired sparsity level, between and 1;
inputs: the sampled input data for optimization;
def prune (W, X, θ, s):
81, 91,02.03 0/3 = =
fouts:
the output data of dense transformer block;
w w_norm weightt_ooteee wow.s/oooomaaiiinnoooeooo (W, 02)
act metric xo (p=2, dim=0) ** 03
@: learned parameters for BaNA;
metric w_norm.abs () act_metric
def po_.optimieerreeee (layer, inputs, fouts, θ): corch.norman (mean=0, std=1)
sorted idx torch. (metric, dim=1) pruned_ide idx w.scatter [dimte, (dim=l, indexepruned idx, sre=0) return
[iddd_ixx[;; (C_ =
é zo_ ( 1, f, loss1 = lossfunc (layer, inputs, fouts) @ po_pdatee (θ, -2, inputs, fouts)
z)
loss2 = losefunc
def weight normalization (W, 81, 0z):
@ = zo ro_update 1, c, z)
W_row_now _ow_ n.norm dim=-1)
grad ((loss] loss2) / zo update 1r, grad, z)
_.co_norm N.norm dim=0)
row_mette metric (1/W_row ['W_row_norm norm)
01
col_mett metrietic (1/W_col norm) 02
def update (θ, ir, grad, z): for 8 in G: 8 - ir grad return θ
return trow_metric + col_mmmet mel_merrc)
notice, compared with previous methods, the standard deviation is lower, indicating its robustness under different random seeds.
E. Code Details
In this section, we further present the code details for implementing BaWA, enabling readers to easily replicate the experimental results reported. The code for BaWA's pruning process and parameter optimization is shown in Algorithm I and 2, respectively.
The pruning process of BaWA can be implemented and integrated seamlessly within a single forward pass of the LLM model, can presented in Algorithm 1. Given a pre-trained LLM, compute our pruning metric from the initial the final layers of the network. After pruning preceding layer, the subsequent layer receives updated input activations, based on which its pruning metrics will be computed. As one can notice, the complexity of the pruning process of BaWA is (dPiidden) which is the same as Wanda. dhidden denotes the channel dimension of the hidden layer.
However, we also introduce a parameter optimization stage in BaWA, which uses a zeroth-order optimizer to adjust the corresponding power factors in Equation (6). As Algorithm shows, the entire parameter optimization process requires two times loss computations and four times parameter updates. Each loss computation requires one forward pass. Therefore, the complexity of the whole parameter optimization process is consistent with the complexity of the model inference process, which is O(BQ d. Here, B denotes the batch size and Q denotes the sequence length.
F. Weight Reconstruction with Efficient Pruning Mask
As mentioned in Section the pruning process can be divided into two stages, including pruning mask selection and weight reconstruction. Our evaluation demonstrates that BaWA presents nigh-performance pruning mask metric, which can preserve the performance of original LLMs without introducing any weight update. It is trivial combine BaWA with other weight reconstruction methods to achieve better pruning performance. Consequently, reconstruction methods, DS2T (Zhang et al., 2023) and ADMM-Iter (Boža, 2024). DSOT introduces novel weight connection reconstruction strategy that iteratively grows and prunes weight connection based on their proposed metric. ADMM-Iter views the weight update process as an ADMM problem and solves it by iteratively adjusting the weight data that are not masked. We use two different pruning metrics, Wanda and BaWA, and compare their pruning results, as shown Table 18.
adopt two novel weight
17

---

BaWA: Automatic Optimizing Pruning Metric for Large Language Models with Balanced Weight and Activation
Table 18: Weight reconstruction strategies with efficient pruning mask achieve better performance.
LLaMA
LLaMA2
Method
Sparsity
30B
65B
7B
13B
70B
Wanda w/ DS2T
50%
7.12
6.08
5.12
4.54
6.31
5.48
3.95
BaWA DS2T
50%
6.97
5.94
5.01
4.37
6.28
5.42
3.87
Wanda ADMM-Iter
50%
7.06
6.07
5.18
4.51
6.33
5.52
3.95
BaWA w/ ADMM-liC Iter
50%
6.91
5.93
4.98
4.37
6.22
5.33
3.81
Wanda w/ DS2T
4:8
8,45
7.25
5.91
5.26
7.83
6.47
4.43
BaWA w/ DS2T
4:8
8.01
6,83
5.65
4.93
7.36
6.19
4.35
Wanda w/ ADMM-Iter
4:8
8.13
6.97
5.73
5.17
7.71
6.41
4.40
BaWA w/ ADMM-Iter
4:8
7.96
6.76
5.63
4.95
7.34
6.11
4.31
Wanda DS2T
2:4
10.89
9.05
6.76
6.14
10.46
8.09
5.11
BaWA w/ DSOT
2:4
10.21
7.91
6.42
5.69
9.84
7.08
4.86
Wanda ADMM-Iter
2:4
9.90
8.47
6.72
6.23
9.89
7.94
5.19
BaWA w/ ADMM-Iter
2:4
9.71
7.86
6.39
5.60
9.75
7.04
4.71
is worth noting that using BaWA as the pruning metric achieves superior LLM performance compared to the adoption of Wanda. Furthermore, the performance pruned LLMs enhanced with the addition of a weight reconstruction strategy beyond what is achieved by pruning mask alone most cases. This underscores the complementary nature of the two phases:
an effective pruning mask can significantly improve pruning performance, while appropriate weight reconstruction strategy further refine the results in post-training pruning.
18


Reviews:
[{"review_id": "DdqxzhJiCZ", "reviewer": "Reviewer_mDNk", "summary": "Existing pruning metrics are limited by their reliance on simple symbolic combinations of weights and activations, failing to account for imbalanced weight magnitudes and the disproportionate impact of activation outliers. To address these shortcomings, this paper introduces BaWA, a pruning metric that balances Weight and Activation distributions for more effective pruning. BaWA incorporates two key innovations:\n\n1. Magnitude Normalization, which mitigates weight imbalances across channels, enabling fairer pruning decisions.\n2. Outlier Regularization, which reduces the influence of activation outliers, ensuring more appropriate channel prioritization.\n\nTo further improve its effectiveness, BaWA includes an efficient, automated framework for optimizing normalization and regularization hyperparameters. Extensive experiments demonstrate that BaWA outperforms existing pruning metrics. For instance, applying BaWA to induce 2:4 sparsity in Mistral-7B reduces perplexity by 2.49 and increases average downstream task accuracy by 3.08%, surpassing the previous method Wanda.\n\n## update after rebuttal\n\nI revisited the paper and would like to keep my original rating for two reasons:\n1. The experimental comparisons are somewhat outdated. Except for Table 3, most of the comparisons are Wanda (proposed in mid-2023) and SparseGPT, which is even older. In Table 4, the method is only slightly better than Wanda. Even in Table 3, the proposed method only shows good improvement when combined with the weight reconstruction method ADMM. Without it, the improvement seems marginal. So, I feel the proposed method may not offer a significant advancement.\n2. Techniques like *Magnitude Normalization* and *Outlier Regularization* have already been extensively studied in previous work. This paper doesn't introduce anything particularly new or exciting for me.\n\nI think this work is a slight extension of Wanda, with an additional normalization step applied to the weights. The method is reasonable but the contribution feels moderate to me, I'd be fine with the paper being either accepted or rejected.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear mDNk:\n\nWe sincerely appreciate your thoughtful feedback regarding BaWA’s novelty and computational overhead. Below, we address your concerns in detail.\n\n### **Novelty of BaWA**:\n\nWe respectfully disagree with the provided novelty concern for three key reasons:\n\n(a) Problem Characterization in LLMs\n+ Structural Outliers: LLMs exhibit extremely sparse activation outliers (>50% outliers in <5% channels, Fig 1c), unlike CNNs with uniform noise patterns [1].\n+ Cross-Layer Imbalance: Weight magnitudes vary by 100× within layers (Fig 1a), violating CNN pruning assumptions.\n\n(b) Methodological Advancements\n\n+ Dual-Channel Normalization: Joint input/output scaling (Eq. 5) handles asymmetric LLM distributions, unlike standard normalization.\n\n+ Dynamic Outlier Suppression: Learnable threshold $\\theta_3$ adapts to layer-wise outlier density, improving upon fixed-threshold methods [2].\n\n(c) Empirical Superiority\n\n+ BaWA reduces perplexity by 2.49 over Wanda (Table 3), which is a great improvement.\n\n+ Direct application of robust CNN pruning degrades accuracy by 4.1% (Table 8 in Appendix).\n\n### **Computational Overhead**:\n\nFurthermore, the evaluation results in our paper have shown that the additional overhead of BaWA is negligible in two aspects. \n\n(a) One-Time Search Cost:\nSearching a 70B model takes only 16 minutes (Table 6), <0.01% of typical training costs (thousands of GPU hours).\n\n(b) Search-Free Mode:\nEven without search (BaWA w/o search), BaWA outperforms Wanda while maintaining the same pruning efficiency (Table 5).\n\n**References**\n\n[1] Li et al., Pruning Filters for Efficient ConvNets, ICLR 2017.\n\n[2] Wei et al., Outlier Suppression+, EMNLP 2023.", "claims_and_evidence": "The claims in the paper are well-supported by clear empirical evidence. However, the concepts of magnitude normalization and outlier regularization are not entirely novel, and the overall contribution may not appear significantly innovative.", "methods_and_evaluation_criteria": "The proposed methods and evaluation criteria are appropriate for the problem.", "theoretical_claims": "This is an empirical paper, no theoretical claims are presented in the paper.", "experimental_designs_or_analyses": "Yes, all parts.", "relation_to_broader_scientific_literature": "The paper builds on prior work in LLM pruning that uses weight and activation magnitudes to guide pruning decisions (e.g., Wanda). While magnitude normalization and outlier regularization have been widely explored in the context of model adaptive sparsity and robust pruning, BaWA refines these ideas by introducing a balancing mechanism for weight and activation distributions. The contribution is a little bit incremental and limited for the community.", "essential_references_not_discussed": "No.", "other_strengths_and_weaknesses": "Strengths:\n\nThe motivation of the paper is clear, the proposed method is simple and easy to follow.\n\nWeaknesses:\n\n1. The proposed two methods *magnitude normalization* and *outlier regularization* in LLM pruning are a little bit trivial and not significant enough.\n\n2. The novelty of the paper is somehow limited. The contributions appear somewhat incremental, making the overall presentation less engaging.\n\n3. The proposed method will involve additional complexity and computation than baseline Wanda, especially w/ the search strategy.", "other_comments_or_suggestions": "Please see the comments above.", "questions_for_authors": "Please see above.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vENgWQfcX7", "reviewer": "Reviewer_xLm6", "summary": "This work proposes a weight pruning method based on Wanda by performing normalization through input and output channels and scaling normalization factors. Wanda is a simple weight pruning method which uses scores measured by L1 of weight and L2 of input, but it suffers issues of imbalance in weight magnitude and influence from outliers. This work alleviate the issues by normalizing weights according to the L2 norm of input and output channels and by scaling the normalization terms. Scaling terms are optimized by forward-only methods for faster computation. Experiments are carried out on several zero-shot benchmarks and shows that the proposed method achieves better results when compared with other SOTAs, including ADMM-Iter.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear xLm6:\n\nWe greatly appreciate your insightful comments. Below we provide a point-by-point response to your concerns.\n\n### **Scaling Factor Analysis**：\n\nWe agree that analyzing scaling factors is critical. In the revised manuscript, we will add:\n\n+ A new comparative table (below) demonstrating the superiority of optimized scales over fixed values on LLaMA2-7B\n\n+ A novel discussion in the revised Section 4.4 will highlight how adaptive scaling addresses LLM-specific distribution challenges. Specifically, different models and task settings will be explored to illustrate the effectiveness of BaWA's scaling strategy.\n\nThe evaluation results can be shown as follows. \n\n| Scaling Strategy       | $\\theta_1$ (Input) | $\\theta_2$ (Output) | $\\theta_3$ (Activation) | PPL  | Δ vs. Best Fixed (Fixed at 0.5) |\n|------------------------|------------|-------------|------------------|------|------------------|\n| **Fixed Scales**       |            |             |                  |      |                  |\n| - $\\theta$=0.1               | 0.1        | 0.1         | 0.1              | 8.92 | +24.1%           |\n| - $\\theta$=0.5               | 0.5        | 0.5         | 0.5              | 7.18 | +0% (baseline)   |\n| - $\\theta$=1.0               | 1.0        | 1.0         | 1.0              | 7.53 | +4.9%            |\n| **BaWA Optimized**     | 0.42       | 0.51        | 0.38             | 6.30 | **-12.3%**       |\n\nThe key findings include: \n\n+ Optimized scales reduce perplexity by 12.3% compared to the best-fixed scale ($\\theta$=0.5)\n\n+ Fixed scales exhibit significant sensitivity (±24.1% PPL variance)\n\n### **Relationship with ADMM-Iter/DSnoT**：\n\nWe apologize for the lack of explanation for the relationship between BaWA and ADMM-Iter/DSnoT. In fact, the LLM pruning procedure can be divided into two stages, pruning mask selection and weight reconstruction. Different from BaWA which optimizes pruning metric (Stage 1), these methods (both ADMM and DSnoT) focus on reconstructing post-pruning weights (Stage 2). To illustrate the orthogonality, we validate when adding weight reconstruction methods (both DSnoT and ADMM) with BaWA, as well as comparing them with SparseGPT, ADMM-Iter and DSnoT without BaWA pruning metric on various LLaMA models with 2:4 sparsity. The results (Table below) clearly depict that using BaWA pruning metric with weight reconstruction methods achieves the best pruning performance. Furthermore, our evaluation in Table 17 (Appendix G) also demonstrates the effectiveness of adding weight reconstruction to BaWA pruning metric. We will add this pipeline diagram to Appendix G.\n\n| PPL        | 1-7B  | 1-13B | 1-30B | 1-65B | 2-7B  | 2-13B | 2-70B |\n|------------|-------|-------|-------|-------|-------|-------|-------|\n| SparseGPT  | 11.00 | 9.11  | 7.16  | 6.28  | 10.17 | 8.32  | 5.40  |\n| Admm-Iter  | 9.90  | 8.60  | 6.89  | 6.02  | 9.74  | 7.78  | 5.19  |\n| DSnoT      | 10.89 | 9.05  | 6.76  | 6.15  | 10.46 | 8.09  | 5.11  |\n| BaWA       | 10.32 | 7.94  | 6.37  | 5.61  | 9.93  | 7.13  | 4.84  |\n| BaWA+DSnoT | 10.21 | 7.91  | 6.42  | 5.69  | 9.84  | 7.08  | 4.86  |\n| BaWA+Admm  | 9.71  | 7.86  | 6.39  | 5.60  | 9.75  | 7.04  | 4.71  |", "claims_and_evidence": "The major claim regarding the weight magnitude normalization for input and output channels is supported by the discussion and analysis in section 2 and 3, and by ablation studies in section 4.4. The use of scaling is clearly described in section 2 and 3, and the proposed optimization algorithm is a nice contribution. However, further analysis on the more controlled scaling might be helpful to clearly understand their impact. For instance, we could intentionally introduce scales in the range of 0.1, 0.25, 0.5, 1.0 etc. and measured the impact in perplexity. The combination of optimization algorithm together with normalization is not throughly measured in the ablation studies if my understanding is correct. For instance, it is possible to optimize a single scale parameter $\\theta$ in Equation 4 when showing ablations in Table 5 to measure the impact of optimization. Similarly, outlier regularization in Table 5 uses a fixed scaling of 0.5, but it could be optimized to show the effectiveness of the proposed algorithm.", "methods_and_evaluation_criteria": "Experiments cover diverse benchmarks for zero-shot settings as well as WikiText-2 for perplexity to compare prior baselines.", "theoretical_claims": "No theoretical proofs for the proposed method, since this work focuses on empirical studies by carefully analyzing the behavior of weight magnitudes for pruning.", "experimental_designs_or_analyses": "Experimental design sounds good to me, but I'd suggest additional ablation studies as noted in the \"Claims and Evidence\" section, i.e., running experiments by optimizing scales when ablating normalization factors.\n\n* Ablations in the rebuttal, that look promising.", "relation_to_broader_scientific_literature": "This work is an extension of prior work in unstructured weight pruning such as Wanda as noted in this manuscript.", "essential_references_not_discussed": "This work is missing discussion on ADMM-Iter and DSOT, given that this work shows experiments by combining them with the proposed method. It is not clear whether the proposed method is orthogonal to those method, and it is even not clear what kind of findings or conclusion could be drawn by showing the results.\n\n* Additional details in the rebuttal.", "other_strengths_and_weaknesses": "Strengths\n\n- It is an extension of prior work on unstructured pruning by adding normalizations and scaling. The design is well motivated and optimization for scaling factor is a yet another contribution to this field.\n\nWeaknesses\n\n- Further ablations are necessary to justify the claim regarding the scaling, since it is not well supported by the experiments.", "other_comments_or_suggestions": "None.", "questions_for_authors": "It is not clear whether scales are optimized in Table 5 when showing the proposed method with, e.g., input channel normalization.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "h6KPkdmrwN", "reviewer": "Reviewer_zn1Z", "summary": "This paper focuses on unstructured pruning of LLMs and introduces a new pruning metric. Unlike previous methods that estimate parameter importance based solely on magnitude, activations, or gradients, the proposed approach also considers the impact of outliers in model parameters.\n\nThe authors first demonstrate that a small number of outlier parameters with large magnitudes can significantly affect existing pruning metrics. To address this issue, the proposed method normalizes each model parameter using the $\\ell_2$ norms of the corresponding input and output channels. Additionally, to handle outliers in the input, the authors introduce a power factor after computing the input norm.\n\nTo optimize the newly introduced hyperparameters, the authors propose using a zeroth-order gradient approach, which allows optimization without backpropagation. Experimental results show that the proposed method outperforms baseline approaches.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Dear zn1Z:\n\nWe sincerely appreciate the valuable suggestions provided by the reviewer. We note the two main concerns you raised, which we address below.\n\nFirstly, we thank the reviewer for emphasizing the importance of comparing with structured pruning methods. We would like to compare structured and unstructured pruning in terms of pruning granularity, sparsity, accuracy, efficiency, and training after pruning. Structured pruning [1, 2, 3] removes complete substructures or weight groups (layers [4], FFN neurons [1], MHA heads [1], embedding dim [5], they are coarse-grained) from LLMs, enabling hardware-independent efficiency gains. However, under the constraints of this coarse-grained pruning, the post-pruned accuracy of LLM is prone to be drastically reduced, and therefore the generally applicable sparsity ratio is between 15% and 30%. Post-pruning fine-tuning/training/knowledge distillation may be used to restore the performance of the pruned model when faced with high sparsity ratios [2, 5]. Unstructured pruning removes elements from the weights (fine-grained) and stores or loads the pruned weights in a compressed assortment. Combined with decompression (memory bottleneck optimization) or hardware advantage (2:4 sparse tensor core), unstructured sparse LLM can also obtain considerable efficiency improvement. Unstructured pruning is less likely to harm the accuracy of the model due to the fine-grained constraints, so the sparsity can generally exceed 50%, and the adoption of a stricter 2:4 sparsity is also acceptable. It is possible to do unstructured pruning without loss of accuracy for large LLMs such as Llama-2-70B, and BaWA has done it. Unstructured pruned LLM can also be further trained, and some techniques such as PEFT [6] and STE [7] have emerged to provide support. In the revised manuscript, we will include a dedicated discussion (Section 2) comparing structured and unstructured pruning for LLMs to demonstrate the necessity of high-performance unstructured/semi-structured LLM pruning. The following table compares a portion of the experimental results of structured pruning and unstructured pruning as a reference.\n\n| Method | Type | Sparsity | reference speedup | BoolQ  | RTE | HellaSwag | WinoGrande  | ARC-e  | ARC-c  | OBQA | AVG   |\n|--|--|--|--|--|--|--|--|--|--|--|--|\n| Llama-2-13B | Dense | 0% | 1.0x | 83.43  | 67.51 | 61.25 | 73.8 | 80.83  | 50.43  | 32.8 | 64.3  |\n| LLM-Pruner  | Structured Pruning | 25%  | 1.25x | 68.35  | 50.54 | 46.81 | 61.56  | 70.2   | 37.63  | 28.8 | 51.98 |\n| ShortGPT    | Layer Pruning | 25%  | 1.25x | 62.54  | 59.57 | 47.7 | 70.96 | 61.24  | 37.88  | 27   | 52.41 |\n| Wanda       | 2:4 Semi | 50% | 1.4x | 75.26  | 56.68 | 46.43 | 66.77 | 68.35  | 34.47  | 24.4 | 53.19 |\n| BaWA        | 2:4 Semi | 50% | 1.4x | 78.26  | 56.32 | 48.5  | 66.93 | 70.79  | 35.49  | 26   | 54.61 |\n\nSecondly, although unstructured pruning has traditionally lacked hardware support, recent advancements such as Flash-LLM [8] demonstrate its practical speedup through specialized kernels. Moreover, our method’s compatibility with N:M sparsity (natively supported by Ampere GPUs) further ensures deployability while maintaining higher accuracy than structured alternatives (Table 9). Additionally, we would like to clarify that the slight performance gap between BaWA and SparseGPT in semi-structured settings (Table 8) stems from SparseGPT’s weight reconstruction—a complementary technique orthogonal to pruning metrics, which is explained in our rebuttal to the reviewer xLm6. As shown in \"BaWA+ADMM\" (Table 3), combining our metric with weight reconstruction outperforms all baselines universally. Furthermore, larger models (e.g., LLaMA2-70B) exhibit greater robustness to N:M constraints due to inherent redundancy, reducing accuracy drops to <1% in 4:8 sparsity (Table 12).\n\n**Reference**\n\n[1] Llm-pruner: On the structural pruning of large language models, NeurIPS'23.\n\n[2] Sheared llama: Accelerating language model pre-training via structured pruning, ICLR'24. \n\n[3] Instruction-Following Pruning for Large Language Models, ArXiv'25. \n\n[4] SlimGPT: Layer-wise Structured Pruning for Large Language Models, NeurIPS'24. \n\n[5] Llm pruning and distillation in practice: The minitron approach, ArXiv'24. \n\n[6] SPP: Sparsity-preserved parameter-efficient fine-tuning for large language models, ICML'24. \n\n[7] Sparsity-accelerated training for large language models, ACL'24. \n\n[8] Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity, VLDB'23.", "claims_and_evidence": "The claims made in the submission are well-supported. The authors provide clear evidence for their core arguments regarding the influence of outliers in model parameters and input activations. Additionally, the proposed method is well-reasoned and justified.", "methods_and_evaluation_criteria": "The proposed method and empirical evaluation follow common practices in this domain. The reviewer generally agrees with the evaluation settings and does not find any significant issues with them.", "theoretical_claims": "This paper does not contain theoretical claims.", "experimental_designs_or_analyses": "The experimental design is sound and valid. It follows previous methods (e.g., Wanda, SparseGPT) and uses widely accepted benchmarks.", "relation_to_broader_scientific_literature": "This paper is closely related to previous literature and methods in this domain. It builds upon existing pruning metrics and further enhances the effectiveness of unstructured model pruning.", "essential_references_not_discussed": "I have only one main comment here. The paper does not fully discuss past research on model pruning for LLMs, focusing mainly on the baselines used in the experiments. While this may be due to space constraints caused by the extensive evaluation and analysis, I still recommend that the authors provide a broader discussion of related work. In particular, discussing structured pruning methods [1-4] for LLMs would be valuable, as structured pruning is generally more hardware-friendly compared to unstructured pruning, which is the focus of this paper.\n\n[1] Xia, Mengzhou, et al. \"Sheared llama: Accelerating language model pre-training via structured pruning.\" *arXiv preprint arXiv:2310.06694* (2023).\n\n[2] Sreenivas, Sharath Turuvekere, et al. \"Llm pruning and distillation in practice: The minitron approach.\" arXiv preprint arXiv:2408.11796 (2024).\n\n[3] Ling, Gui, Ziyang Wang, and Qingwen Liu. \"SlimGPT: Layer-wise Structured Pruning for Large Language Models.\" Advances in Neural Information Processing Systems 37 (2024): 107112-107137.\n\n[4] Hou, Bairu, et al. \"Instruction-Following Pruning for Large Language Models.\" *arXiv preprint arXiv:2501.02086* (2025).", "other_strengths_and_weaknesses": "This paper is well-motivated, with clear writing and a coherent logical flow. The reviewer enjoyed reading it. Additionally, the proposed method is reasonable and well-structured, and the evaluation is rigorous and comprehensive.\n\nThe only concern is the performance under semi-structured sparsity. I assume that the unstructured sparsity results in Table 2 and Table 4 do not lead to efficiency improvements. As shown in Table 5, Table 8, and Table 10, the performance degradation remains significant compared to the original dense model. Furthermore, SparseGPT sometimes outperforms the proposed method. However, the results improve in Table 11 and Table 12.\n\nOverall, the reviewer considers this a strong paper.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper introduces a new metric for pruning language models, focused on addressing imbalanced weight magnitudes and activation outliers.  While there is some disagreement about the novelty of the approach as it extends prior work, the proposed method improves over prior work and the reviewers agree it is clearly explained and experiments are sound. Other concerns raised by reviewers (including speed of pruning, scaling factor sensitivity, comparison to structured pruning methods) are largely addressed by the authors in the rebuttal. Based on this, we recommend this paper for acceptance.
Decision: Accept (poster)