Beyond Linear Approximations: A Novel Pruning Approach for Attention Matrix

Authors: Yingyu Liang, Jiangxuan Long, Zhenmei Shi, Zhao Song, Yufa Zhou
Primary area: foundation or frontier models, including LLMs
Keywords: Weights Pruning, Attention Approximation, Gradient Descent Optimization

Abstract:
Large Language Models (LLMs) have shown immense potential in enhancing various aspects of our daily lives, from conversational AI to search and AI assistants. However, their growing capabilities come at the cost of extremely large model sizes, making deployment on edge devices challenging due to memory and computational constraints. This paper introduces a novel approach to LLM weight pruning that directly optimizes for approximating the attention matrix, a core component of transformer architectures. Unlike existing methods that focus on linear approximations, our approach accounts for the non-linear nature of the Softmax attention mechanism. We provide theoretical guarantees for the convergence of our Gradient Descent-based optimization method to a near-optimal pruning mask solution. Our empirical results demonstrate the effectiveness of our non-linear pruning approach in maintaining model performance while significantly reducing computational costs, which is beyond the current state-of-the-art methods, i.e., SparseGPT and Wanda, by a large margin. This work establishes a new theoretical foundation for pruning algorithm design in LLMs, potentially paving the way for more efficient LLM inference on resource-constrained devices.

Decision: accept (poster)
Reviewer scores: 5, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730697535726

thread_type:
official_review

<thread_object>
creator: reviewer_qlv4

<content>
summary:
This paper proposes a method for unstructured pruning of the query and key projection matrices in attention. The paper provides theoretical guarantees for the convergence of the pruning method. It also provides some experiments on synthetic data demonstrating that the pruning method outperforms other pruning methods.

strengths:
*Original problem.* The paper highlights an interesting, niche problem in the pruning literature, which, to my knowledge, has not been analyzed theoretically. 

*Clear presentation.* The paper is well-written and the structure is easy to follow. 

*Extensive theoretical analysis.* The authors provide an extensive theoretical analysis on the convergence of their pruning algorithm.

weaknesses:
*Limited potential for speedup on large scale models.* I have two concerns around the potential for this method to provide significant speedups for large language models. (1) It is typically difficult to leverage fine-grained, unstructured sparsity for speedups on GPUs because standard vector and matrix instructions are not compatible with sparsity (though, recent changes to Nvidia GPUs could begin to mitigate this problem [Nvidia Blog](https://developer.nvidia.com/blog/accelerating-inference-with-sparsity-using-ampere-and-tensorrt/)). (2) The technique is applicable to a small fraction of the parameters in the full model. In large models, the MLPs dominate the total parameter and FLOP counts, but this method is only applicable to the Q and K projections, which account for only 9.4% of parameters in Llama-405B (excluding embeddings, [source](https://huggingface.co/meta-llama/Llama-3.1-405B/tree/main)). So even if the method could achieve 100% sparsity, the relative reduction in FLOPs/memory consumption would be limited. 

*Experiments limited to synthetic data.* Because experiments are only performed on synthetic data and a single layer, it remains unclear how the performance of real language models would degrade with this method. Without experiments on real language models it is hard to assess the improvement of this method relative to prior work (Wanda and SparseGPT).

questions:
How does the convergence rate in Theorem 4.1 compare to the rates one can achieve with naive baselines?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730795951004

thread_type:
official_review

<thread_object>
creator: reviewer_n1ar

<content>
summary:
The paper introduces a novel approach to pruning weights in large language models (LLMs) to improve computational efficiency, particularly for deployment on resource-constrained devices. This method focuses on optimizing the approximation of the attention matrix, a critical component in transformer architectures, by accounting for the non-linear properties of the Softmax attention mechanism. The authors present a Gradient Descent-based optimization technique, offering theoretical guarantees for its convergence toward a near-optimal pruning solution. Preliminary results indicate that this approach maintains model performance while significantly reducing computational costs. This work provides a new theoretical basis for designing pruning algorithms for LLMs, potentially enabling more efficient inference on devices with limited resources.

strengths:
1. The presented technique enhances convergence guarantees for non-convex functions, crucial for optimization challenges in areas like deep learning, through the application of the Polyak-Łojasiewicz inequality. It showcases mathematical rigor with detailed derivations and validations of key conditions, supported by established theorems, notably Theorem 5.2 from Frei & Gu (2021). Practical implications include improved computational efficiency via attention weight pruning without notable performance trade-offs, assured by the method’s convergence mechanisms. Additionally, its framework offers potential applicability across various non-convex optimization scenarios, expanding its usefulness in the optimization toolkit.
2.  The experimental results suggest that the proposed algorithm outperforms existing methods under certain conditions, the above weaknesses highlight areas where the study could be improved.

weaknesses:
1. The convergence guarantees depend on strict conditions like the boundedness and positive definiteness of matrices (e.g., $\( X^T X \succeq \beta I \))$. These conditions may not always hold in real-world applications, limiting the practical relevance of the results. Additionally, the analysis does not consider the stochasticity and noise common in real-world data and optimization, which could affect gradient descent’s performance.

2. The experimental section lacks clarity on whether the tested input sequence lengths cover the full range required for real-world corpus. Given that modern llm often deals with lengthy texts like entire documents or extended conversations, the study does not explore the algorithm’s performance in such scenarios, raising concerns about its practical utility.

questions:
Conside the weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730811368790

thread_type:
official_review

<thread_object>
creator: reviewer_wrvg

<content>
summary:
The paper proposes to prune the projection weight matrices of the Softmax attention conditioned on the similarity between the pruned and the original attention matrices. The proposed approach leverages gradient decent to solve the constraint opimitzation problem and has theoretical gurantees for near-optimality with polynormal time.

strengths:
1. Considering the similarity between attention matrices during pruning of the projection weight is novel and interesting.
2. The assumptions in theoretical analyses are mild and sensible.

weaknesses:
1. The applicability of the approach is limited. The proposed pruning approach needs to fuse the query and the key projection matrices into a single matrix through matrix multiplication. This means that it will eliminate the inference efficiency benefits of popular LLMs, such as Llama 3, which use Grouped Query Attention that shares the key-value pairs between queries. This is because after the proposed fushion, the original input X needs to be saved during decoding, while for GQA, only much smaller key-value vectors are needed which has much smaller memory I/O footprint.
2. Lacks of empirical results on pruning real-world LLM. The low error rates on synthetic data may not translate to real world dowstream performances.

questions:
1. How to revert the pruned fused matrix back to the query and key projection matrices respectively?
2. How does the proposed attention pruning affect length extrapolation ability?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730823379826

thread_type:
official_review

<thread_object>
creator: reviewer_opkw

<content>
summary:
This paper proposes a novel approach to pruning QK projection matrices by directly optimizing for an approximation of the attention matrix, bypassing traditional linear approximations applied before the Softmax operation. By addressing the non-linear nature of the attention mechanism, the authors introduce a gradient descent-based optimization method with theoretical guarantees on convergence, aiming to produce a near-optimal pruning mask. Preliminary empirical results, although limited, demonstrate that this approach can maintain model performance while significantly reducing computational costs.

strengths:
- The paper offers a robust theoretical analysis, including convergence guarantees for the gradient descent algorithm, which enhances confidence in the method’s foundational soundness.
- The proposed pruning method consistently outperforms baseline approaches in terms of relative error, effectively recovering the attention matrix activations. This suggests its potential for maintaining model performance while reducing computation.

weaknesses:
- The empirical analysis relies heavily on synthetic data, specifically a full-rank Gaussian random matrix, with matrix dimensions (64x128) that are small relative to the scale of practical LLMs. This somewhat limits confidence in the method’s real-world applicability.
- I understand the main contributions are theoretical, but I think the paper would benefit from experiments with actual LLM activations and weights, even small-scale LLMs. Results from real models would give a stronger sense of how effective this method might be in practice. 
- One big limitation is that the approach only targets $W_Q$ and $W_K$ matrices in the attention mechanism, while the majority of parameters in transformers come from the MLP layers. It would be interesting to see if this technique could be extended to cover the MLP parameters as well, which could have a more substantial impact on reducing model size.

questions:
- I assume the cost of pruning itself is very small given the closed-form gradient, but it would be helpful to see an analysis of this cost. Is this approach lightweight enough for large models?
- The assumptions regarding the Lipschitz continuity and PL inequality of the gradient are central to the theoretical guarantees. Are there empirical checks or practical insights into how these assumptions hold up in real LLMs?
- The paper uses synthetic data by generating $W_Q$ and $W_K$ through SVD on Gaussian random matrices. A bit more explanation for this choice—or testing on real model weights—would clarify whether the setup could translate to practical applications.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732132998800

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
Thanks for the reviews' valuable comments!

We appreciate reviewers opkW, N1Ar, WRvG, and QLv4 for recognizing the novelty and theoretical rigor of our analysis. They collectively acknowledge our convergence guarantees for the gradient descent method, with opkW, N1Ar, and QLv4 highlighting the lower relative error achieved compared to existing pruning methods. WRvG commends the mild and sensible assumptions in our theoretical analysis, while N1Ar notes our method’s potential applicability to various non-convex optimization scenarios, enhancing versatility. QLv4 appreciates the clarity and organization of our paper. Additionally, WRvG and QLv4 find our optimization approach to attention matrix pruning particularly intriguing. 

We have updated a **revision** for our draft. We also updated the code about experiments in the supplemental material under the folder named *7932_supp_rebuttal*. We summarize all the updates (in brown color) we made in the revision. All line numbers in the rebuttal correspond to the revised version.
- Line 472-530: We add our new experimental results performed on Llama 3.2 1B and C4 dataset.

Here, we address some common questions. 

## Experiments on real LLMs
We update the below experiments in our revision Line 472-530.

Due to the time constraint, we implement our method on [LLama 3.2 1B](https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct/blob/main/config.json). We extract the real attention weights and input hidden states for one specific attention layer using a customized hook function and use the calibration data from [C4 Dataset](https://huggingface.co/datasets/legacy-datasets/c4) as Wanda [1] and SparseGPT [2]. We also submit our code in the supplementary material. We show our results below, where the number in each cell is the output **relative error** (Line 418-421, Equation (1) in the draft) between the pruned and original attention layers: 
 
All three methods under different input context length $n$: 
| **Method \ $n$** | **64**       | **128**      | **256**      | **512**      | **1024**     | **2048**     | **4096**     |
|---------------|--------------|--------------|--------------|--------------|--------------|--------------|--------------|
| **Wanda**     | 0.1146       | 0.1204       | 0.1271       | 0.1229       | 0.1533       | 0.1851       | 0.1804       |
| **Sparse GPT**| 0.0595       | 0.0692       | 0.0873       | 0.0749       | 0.0968       | 0.1526       | 0.1472       |
| **Ours** | 0.0576       | 0.0399       | 0.0542       | 0.0471       | 0.0847       | 0.1262       | 0.1346       |

All three methods under different pruning ratios $\rho$: 
| Method \ $\rho$         | 0.4      | 0.5      | 0.6      | 0.7      | 0.8      | 
|-----------------|--------------|--------------|--------------|--------------|--------------|
| **Wanda**| 0.0337       | 0.0491       | 0.0750       | 0.1159       | 0.1730       | 
| **SparseGPT**| 0.0331      | 0.0329       | 0.0476       | 0.0763       | 0.0989       | 
| **Ours**   | 0.0333       | 0.0333       | 0.0390       | 0.0423       | 0.0485       | 


As we can see that our result is quite promising, where our method outperforms Wanda and SparseGPT in all cases. These results show that our method may inspire new pruning algorithms, and the current pruning techniques used for LLMs have potential improvement space. Lastly, we emphasize that our work is theoretical focus, introducing a new approach to pruning the non-linear attention matrix. We are actively researching ways to extend our pruning method to larger LLMs to show more promising empirical results. 


## Assumptions verification on real LLMs

Our work made two assumptions:
- We assume the positive semi-definiteness of input data, i.e., $XX^\top \succeq \beta I$ in Line 256.
- We assume the pruned attention matrix is lower bounded, i.e. $\min_{i,j \in [n]} (\widetilde{D}^{-1}\widetilde{A})_{i,j} \ge \delta > 0$ in Line 256.

We verify the two assumptions on Llama 3.2 1B described in the table below:
| **Assumption**         | **Value**           |
|-------------------------|---------------------|
| PSD input data, $\beta$ assumption  | 0.034 |
| Attention matrix lower bound, $\delta$ assumption | 0.0025 |


We see our assumptions hold for Llama 3.2 1b on the C4 dataset, suggesting the practicality of our assumptions. 

[1] Sun, M., Liu, Z., Bair, A., & Kolter, J. Z. A Simple and Effective Pruning Approach for Large Language Models. ICLR 2024

[2] Frantar, E., & Alistarh, D. Sparsegpt: Massive language models can be accurately pruned in one-shot. ICML 2023
</content>
</thread_object>

</thread>
<thread>
creation date:
1732902099707

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
End-to-end Perplexity Evaluation

comment:
Per reviewers QLv4 and opkW's valuable suggestions, we evaluate our methods Perplexity (PPL) on a real model and real dataset from end to end, comparing Wanda [1] and SparseGPT [2]. 

*Setup:* We use LLama 3.2 1B as the target model to prune, in which the hidden state dimension is 2048. We set the pruning ratio as 0.5 for MLP and 0.5 for Attention Layer. If we prune the MLP and Attention Layer, we have a pruning ratio of 0.5 for the whole model. We use 8 sentences in the C4 training dataset [3] as calibration data, we use 128 sentences in the C4 validation dataset [3] to evaluate perplexity. We set 0.005 as the learning rate and 0.05 as the regularization coefficient. 

The results are in the following table, where Dense means non-pruning/original weights. 

| Method | PPL |
|------|----|
| Dense MLP + Dense Attn | 12.487  |
| Dense MLP + SparseGPT Attn | 14.269 |
| Dense MLP + Wanda Attn | 14.912 |
| Dense MLP + Our Attn | **13.885**  |
| Wanda MLP + Wanda Attn | 30.426 |
| Wanda MLP + SparseGPT Attn |  26.074 |
| Wanda MLP + Our Attn | **24.427** |
| SparseGPT MLP + Wanda Attn |45.435 |
| SparseGPT MLP + SparseGPT Attn | 36.641  |
| SparseGPT MLP + Our Attn | **34.946** |

As we can see in the table, Our Attention pruning methods always outperform Wanda Attention pruning and SparseGPT Attention pruning by a large margin when combined with Dense/Wanda/SparseGPT MLP pruning methods. 
- These empirical results support our theoretical analysis that our pruning method can converge. 
- The empirical results support that our method is practical in the real world case.
- Our method is broadly applicable and can be combined with many other pruning methods.

[1] Sun, M., Liu, Z., Bair, A., & Kolter, J. Z. A Simple and Effective Pruning Approach for Large Language Models. ICLR 2024

[2] Frantar, E., & Alistarh, D. Sparsegpt: Massive language models can be accurately pruned in one-shot. ICML 2023

[3] Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., ... & Liu, P. J. Exploring the limits of transfer learning with a unified text-to-text transformer. JMLR’20.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper introduces a method for pruning query and key weight matrices in attention modules. The paper proposes a theoretical justification for the proposed scheme, with convergence guarantees. The method was initially evaluated on synthetic data, and as requested by the reviewers, also on real-world language models. The proposed approach seems to perform relatively well, but does not provide a viable strategy for pruning in real-world applications. As mentioned by reviewer opkW, most of the computation and parameters are in the feedforward layers. Tu summarize, this paper proposes an interesting finding, for which the empirical implications are yet a bit unclear. While this is a borderline paper, I recommend this paper for acceptance.