AlphaPruning: Using Heavy-Tailed Self Regularization Theory for Improved Layer-wise Pruning of Large Language Models

Weakness

We present the differences between AlphaPruning and [1], as detailed below:

• Different research focus. Our study investigates post-training LLM pruning, whereas [1] studies model training.

• The underlying principles of the two works are different. [1] aims to balance layer quality (or make layers equally well-trained) by tuning learning rates, aiming to improve generalization performance. This work focuses on minimizing pruning damage to the high-quality layers by making them less sparse, thereby reducing performance loss. While both studies use heavy-tailed metrics from HT-SR theory to estimate layer quality, this shared aspect does not make our work a trivial extension or incremental idea of [1]. This work is the first to explore how HT-SR theory can be applied to allocate sparsity, a concept previously unexplored in the literature.

• Important technical difference: transformer-block-wise measurement instead of matrix-wise measurement. [1] measures the PL_Alpha_Hill metric for each weight matrix of the model individually. However, as we demonstrate in Appendix F, this approach provides suboptimal results for sparsity allocation. We improved upon this by averaging the metric scores across matrices within a transformer block and using the block-wise average score for sparsity allocation, which yielded significantly better results.

Question 1

We thank the reviewer for suggesting a new mapping function method. We implemented the proposed method and compared it with the linear mapping function used in our current approach, as shown in Figure 12 of the rebuttal PDF. The results show that both methods perform similarly when combined with Wanda, but linear mapping slightly outperforms the proposed logarithmic method when combined with SparseGPT.

Here are the experimental setups for Figure 12. We pruned the LLaMA-V1-7B model to different sparsity levels using two allocation methods and reported the perplexity on the WikiText validation set. The hyperparameter search and setup are consistent with the original settings in Appendix G.

Question 2

We thank the reviewer for the new reference. We integrated AlphaPruning with OSSCAR and provided the updated results in Figure 13 of the rebuttal PDF. OSSCAR prunes only the linear sublayer of multi-head attention and the second sublayer of the feed-forward network, applying uniform pruning across each transformer block. By incorporating AlphaPruning's layer-wise sparsity allocation, we achieved non-uniform block-wise pruning ratios while keeping the global pruning ratio the same. The results show that integrating AlphaPruning with OSSCAR can reduce perplexity at different sparsities. We will include this experiment in the updated draft and cite the corresponding work.

Weakness 1 and 2

AlphaPruning is grounded in heavy-tail self-regularization (HT-SR) theory, which we use to quantify the training quality of each layer and determine layer-wise sparsity. Here we provide a detailed overview of this theory and explain how AlphaPruning is built based on it. We will include these explanations in the updated draft.

• HR-SR theory originated as a semi-empirical theory, with early empirical work [1-2] examining the empirical spectral density (ESD) of weight matrices, specifically the eigenspectrum of the correlation matrix $W^\top W$. This research found that the heavy-tailed structure of the ESD strongly correlates with training quality. These findings are rooted in heavy-tailed random matrix theory and statistical physics, as detailed in Table 1 of [2].
• Recent theoretical work studies how heavy-tails in ESD emerge and why it correlates with training quality. It is well-known [3-4] that spikes in ESD represent "signals", while the bulk represents noise, which follows the Marchenko-Pastur law. In the theoretical setting of [3], the signal or the spike aligns with ground-truth features from the teacher model, and that corresponds to increased correlations in weight elements [1-2]. Furthermore, [5] shows that heavy tails in ESD originate from the interaction between spikes and bulk, which can be quantified precisely using recent advances in the free-probability theory [6], and that is the "bulk-decay" phase in the five-plus-one phase model in [2]. These studies indicate that layers with more heavy-tailed ESDs have extracted more useful signals during training, indicating better training quality.
• This insight motivates our sparsity assignment method: layers with more heavy-tailed ESD contain more learned signals and are assigned lower sparsity by our method, while layers with less heavy-tailed ESD retain fewer signals and are assigned higher sparsity. In practice, the heavy-tailed structure is measured by fitting a power-law distribution to the ESD, and extracting the power-law exponent $\alpha$ as the indicator. This is why our method is named "AlphaPruning".

References

[1] Martin et al. Predicting trends in the quality of state-of-the-art neural networks without access to training or testing data.

[2] Martin et al. Traditional and Heavy-Tailed Self Regularization in Neural Network Models

[3] Wang et al. Spectral Evolution and Invariance in Linear-width Neural Networks

[4] Couillet and Liao et al. Random Matrix Methods for Machine Learning

[5] Kothapalli et al. Crafting Heavy-Tails in Weight Matrix Spectrum without Gradient Noise

[6] Landau et al. Singular vectors of sums of rectangular random matrices and optimal estimation of high-rank signals: The extensive spike model

Weakness 3

We thank the reviewers for providing the new reference. The method used in [Gale+19] that performs best for Transformer models is uniform magnitude pruning with a gradual pruning schedule. The gradual pruning schedule refers to gradually increasing the sparsity of the network while training the semi-pruned model. This schedule is computationally intensive for LLMs due to the iterative model training and has not been adopted in LLM pruning literature. Our study focuses on pruning the LLM in a one-shot way without training, which aligns with previous studies [7-9] to ensure a fair comparison. We have included one-shot uniform magnitude pruning in the submitted paper, as shown in Tables 2 and 3 of the submitted paper. Our results demonstrate that our approach outperforms this baseline.

We clarify that in the submitted paper, we have compared with the most relevant and competitive baseline, OWL [7], which is the current SOTA non-uniform sparsity allocation method. We have also compared with other recent LLM pruning baselines, such as Wanda [8], and SparseGPT [9], as well as other baselines from CV pruning literature such as global, ER [10], rank-selection [11], and layer-wise error thresholding [12].

References

[7] Yin et al. 2024

[8] Sun et al. 2024

[9] Frantar et al. 2023

[10] Mocanu et al 2018

[11] Kuzmin et al. 2019

[12] Ye et al. 2020

Question 1

We thank the reviewer for pointing out this writing issue. We will remove the $\mu_{\mathbf{x}_i}$ in the updated draft.

Question 2

Both unstructured pruning and quantization are effective methods for improving inference speed and reducing memory footprints. The comparison between these two approaches in terms of efficiency can be nuanced depending on hardware deployment and the algorithms used. For example, unstructured sparsity has limited acceleration support in GPUs, compared to quantization, but it holds the great value of speedups on other hardware such as CPU, cerebras' chip, IPU, etc.

Recent studies indicate that unstructured pruning can slightly outperform quantization in inference speedup. A SOTA approach SqueezeLLM ([13]) demonstrates that quantizing LLMs to 3 bits can achieve a 2.1× speedup. Meanwhile, recent advances in sparse CPU kernels (DeepSparse) have better support in accelerating unstructured pruning, leading to 3.35$\times$ speedup in CPU runtimes, as shown in Table 3 of [14], as well as Table 4 of our submitted paper. This is partly because pruning can better maintain and recover performance through fine-tuning compared to quantization, as noted in [14].

Pruning and quantization are compatible and complementary, and combining both approaches further enhances efficiency. Table 3 of [14] shows that combining unstructured pruning with quantization (INT8) can achieve up to 9.08× speedup, significantly higher than using either method alone.

References

[13] Kim et al. 2024

[14] Kurtic et al. Sparse Fine-tuning for Inference Acceleration of Large Language Models, 2023

Weakness 1

We provide a detailed explanation of the parts the reviewer suggested, including HT-SR theory, terms in the method, and Figure 1a. We will include these in the updated draft.

• More details of HT-SR theory: HT-SR theory [1-2] examines the empirical spectral density (ESD) of weight matrices, specifically the eigenspectrum of the correlation matrix $W^\top W$. This research found that the heavy-tailed structure in ESD is strongly correlated with training quality. Recent theoretical work [3-4] finds that such a structure is a result of feature learning, a process of extracting various useful correlations (or features) from data during optimization. These studies indicate that layers with more heavy-tailed ESDs have extracted more useful signals during training, indicating better training quality. In practice, the heavy-tailed structure is measured by fitting a power-law distribution to the ESD, and extracting the power-law exponent $\alpha$ as the indicator. These studies motivate our work, and we allocate the layer-wise sparsity based on the $\alpha$ metric. This is why our method is named "AlphaPruning".

• Terms used in method: In our method section, $\lambda$ represents the eigenvalues of the weight matrices' correlation matrix. The interval ($\lambda_{\min }$, $\lambda_{\max }$) defines the range of eigenvalues considered as the tail part of ESDs.

• Figure 1a: The blue histograms depict the empirical spectral density. The $x$-axis represents eigenvalue magnitudes, and the $y$-axis represents the density, both on a logarithmic scale. The solid red curves depict the empirical distribution of the ESD tail, while the dashed red curves represent the fitted PL distribution. The PL_Alpha_Hill metric in the title is the fitted PL exponent.

Weakness 2

We thank the reviewer for providing the new baseline method LAMP, we provide new experiments on comparing our method with LAMP. The results are shown in Table 25 of the rebuttal PDF. We implement the original LAMP method, which allocates different sparsity for each matrix, and a variant LAMP called LAMP (per block), which allocates the same sparsity for all matrices within a transformer block. This adaption is based on our ablation study comparing per-matrix and per-block strategies in Appendix F. It shows that LAMP (per-block) outperforms the original LAMP, and shows better performance than Uniform pruning baseline when both are combined with SparseGPT. However, our method AlphaPruning outperforms this baseline method.

Here are the experimental settings. We pruned the LLaMA-V1-7B model to three sparsity levels (60%, 70%, 80%) using different methods and reported the perplexity on the WikiText validation set. The hyperparameter search and setup are consistent with the original settings in Appendix G.

Question 1

In Table 26 of the rebuttal PDF, we present the PL_Alpha_Hill and performance metrics (perplexity, zero-shot task accuracy) before and after fine-tuning. We can see that the PL_Alpha_Hill metric is reduced as well as the two performance metrics improved. Here is the experimental setup. We fine-tune the pruned LLaMA-V1-7B on 30K tokens from the C4 dataset. The model is pruned to 70% sparsity. The perplexity is evaluated on WiKiText, and the accuracy is evaluated on 7 zero-shot datasets as listed in Section 4.1.

We select the initial stage of fine-tuning to demonstrate how the metric value evolves. Figure 14 of the rebuttal PDF shows that the PL_Alpha_Hill metric indeed continues decreasing during the fine-tuning.

Question 2

The phenomenon that layers/blocks are not equally well-trained is well-documented in [5-7]. For example, [5] shows that layers or blocks of a model can have imbalanced training quality during training, and balancing the training speed of layers/blocks improves performance. As another example, [6-7] shows that some LLM layers are less useful than others, and removing these LLM layers has negligible impact on performance. However, the underlying reason why some layers are less well-trained than others remains an open question, and we will include that in future studies.

Incorporating PL_Alpha_Hill into the training process is possible. Recent work [5] has used this metric to dynamically adjust the learning rate of each layer during the training and demonstrated their method makes the layers more equally well-trained and improves the generalization performance.

Question 3

The non-linear relationship is due to the low-level hardware operations of sparsity structures rather than the maintenance of early layers. As shown in Figure 4 of [8], for low sparsity, the computational performance (red lines) increases slowly due to overheads in storing sparse structures and controlling sparse computations. As sparsity increases to moderate and high levels, we see sustained growth of performance until it usually levels off at extremely high sparsities where storage and control overheads dominate. This explains the slow speedup in our low-sparsity regime (e.g., less than 70% sparsity).

A new ablation study shows that early layer maintenance is irrelevant to the non-linear relationship. In Figure 11 of the rebuttal PDF, we compare the inference speedup of non-uniform pruning (AlphaPruning) and uniform pruning. The negligible speedup differences between the two methods indicate that non-uniform distribution is unrelated to the observed non-linearity. This is because our method allocates sparsity at the transformer block level, not at the individual layer level, and in transformer architectures, all blocks use identical computational resources.

References

[1] Ref [10] in submission

[2] Ref [13] in submission

[3] Ref [58] in submission

[4] Kothapalli et al. 2024

[5] Zhou et al. 2023 Temperature Balancing

[6] Gromov et al. 2024 The Unreasonable Ineffectiveness of the Deeper Layers

[7] Men et al. 2024 ShortGPT

[8] Hoefler et al. 2021 Sparsity in Deep Learning

Weakness 1 and Question 1

Motivation of our method and why property of $W^\top W$ can decide layer sparsity. Our method is grounded in heavy-tail self-regularization (HT-SR) theory, which we use to quantify the training quality of each layer and determine layer-wise sparsity. The rationale is as follows:

• HR-SR theory originated as a semi-empirical theory, with early empirical work [1-2] examining the empirical spectral density (ESD) of weight matrices, specifically the eigenspectrum of the correlation matrix $W^\top W$. This research found that the heavy-tailed structure of the ESD strongly correlates with training quality. These findings are rooted in heavy-tailed random matrix theory and statistical physics, as detailed in Table 1 of [2].
• Recent theoretical work studies how heavy-tails in ESD emerge and why it correlates with training quality. It is well-known [3,5] that spikes in ESD represent "signals", while the bulk represents noise, which follows the Marchenko-Pastur law. In the theoretical setting of [3], the signal or the spike aligns with ground-truth features from the teacher model, and that corresponds to increased correlations in weight elements [1-2]. Furthermore, [4] shows that heavy tails in ESD originate from the interaction between spikes and bulk, which can be quantified precisely using recent advances in the free-probability theory [6], and that is the "bulk-decay" phase in the five-plus-one phase model in [2]. These studies indicate that layers with more heavy-tailed ESDs have extracted more useful signals during training, indicating better training quality.
• This insight motivates our sparsity assignment method: layers with more heavy-tailed ESD contain more learned signals and are assigned lower sparsity by our method, while layers with less heavy-tailed ESD retain fewer signals and are assigned higher sparsity. In practice, the heavy-tailed structure is measured by fitting a power-law distribution to the ESD, and extracting the power-law exponent $\alpha$ as the indicator. This is why our method is named "AlphaPruning".

Why our method can be combined with other pruning metrics? Our layer-wise sparsity assignment is complementary to other pruning metrics, such as Wanda and SparseGPT. Our method determines the training quality of layers, while other pruning metrics identify the importance of components within each layer (or weight matrix). Thus, our approach decides how much to prune for each layer, while other metrics determine which components to prune within each layer.

References

[1] Martin et al. Predicting trends in the quality of state-of-the-art neural networks without access to training or testing data.

[2] Martin et al. Traditional and Heavy-Tailed Self Regularization in Neural Network Models

[3] Wang et al. Spectral Evolution and Invariance in Linear-width Neural Networks

[4] Kothapalli et al. Crafting Heavy-Tails in Weight Matrix Spectrum without Gradient Noise

[5] Couillet and Liao et al. Random Matrix Methods for Machine Learning

[6] Landau et al. Singular vectors of sums of rectangular random matrices and optimal estimation of high-rank signals: The extensive spike model

Question 2

In addition to PL_Alpha_Hill proposed in our work, [7-8] are other studies that investigated methods that measure whether a layer is well-trained or not, demonstrating LLMs layers are not equally well-trained. [7] developed a method that assesses the similarity between the representations at different layers, defined as the angular distance between feature vectors. They found that deeper layers are more similar to their neighboring layers than shallow layers, suggesting that LLMs may not fully utilize the parameters in these deeper layers, indicating these layers are not well-trained. Similarly, [8] introduced a metric called Block Influence, which measures the impact of each transformer block on hidden states to gauge layer significance. Their findings showed varying degrees of ineffectiveness/redundancy across layers, suggesting that these layers are not well-trained.

References

[7] Gromov et al. The Unreasonable Ineffectiveness of the Deeper Layers

[8] Men et al. ShortGPT: Layers in Large Language Models are More Redundant Than You Expect

Question 3

In Figure 10 of the rebuttal PDF, we compare the sparsity allocation of the two methods. We show that the general trends of sparsity distribution generated by the two methods are similar, with lower sparsities allocated to earlier layers and higher sparsities allocated to deeper layers. However, our method produces a more granular distribution with clearer distinctions between consecutive deep layers, resulting in improved pruning performance.