Beyond Linear Approximations: A Novel Pruning Approach for Attention Matrix

W1.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).

Thanks for pointing out. Our work provides the foundation for future sparsity speedup. Specific sparsity patterns, e.g., structured N:M sparsity [1], and row/column/block-wise sparsity, are indeed important and widely used in modern GPU and CUDA implementation. However, the unstructured sparsity is a superset of specific sparsity patterns mentioned above, suggesting the potential algorithm design. To give a motivational example, we can modify our constraint to row-wise or column-wise pruning, which means we modify our pruning mask $M$ from a matrix of $d \times d$ to a vector of size $d$ denoting the index which row/column to prune, and then modify the corresponding loss function $\mathcal{L}$. This requires a re-analysis.

W1.2: The technique is applicable to a small fraction of the parameters in the full model, because in large models, the MLPs dominate the total parameter and FLOP counts.

Thanks for your valuable insights. Our method is applicable not only to attention layers but also to MLP layers, as we primarily employ gradient descent, which can be applied to the weights of MLP layers with tailored loss functions. Note that, MLP pruning is still in the linear approximation pruning regime, so it is not our main focus since their convergence is guaranteed.

Furthermore, our approach can be seamlessly integrated with other linear pruning methods, such as Wanda [1] and SparseGPT [2], to prune MLP layers, achieving greater overall acceleration.

W2: Experiments limited to synthetic data.

Thank you for pointing out your concerns. We provide some results on the real dataset in C4 and the model in Llama 3.2 1B, where our method outperforms Wanda and SparseGPT. See global response Experiments on real LLMs.

Q1: How does the convergence rate in Theorem 4.1 compare to the rates one can achieve with naive baselines?

Thanks for your valuable question. Since pruning is typically performed only once, the associated time cost can be relatively high. However, after pruning, inference time is accelerated and unaffected by the initial pruning process. Wanda, being a non-iterative method, requires time equivalent to a single forward pass. In contrast, SparseGPT performs pruning through $d$ iterations, where $d$ represents the dimensionality.

Our time cost of pruning mainly depends on the iterations which is $O(d \cdot \mathrm{poly}(n)/\epsilon)$. This is due to the relatively slow convergence of Gradient Descent. However, future work will explore second-order methods like Newton’s method or Adam optimizer and try to optimize the time and memory cost.

[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

We appreciate your raising score!

We use LLama 3.2 1B and C4 Dataset and we do not truncate the input sequences for new experiments. We will add these information and new empirical results in our camera ready. We will also add results of larger model and different datasets in camera ready.

We sincerely thank you for your valuable suggestions and time!

W1: Assumed conditions may not always hold in real-world applications, limiting the practical relevance of the results.

Thank you for pointing out your concerns. We provide some results to verify the assumptions on the real dataset in C4 and the model in Llama 3.2 1B. See global response Assumptions verification on real LLMs.

W2: 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.

Thank you for pointing out. We provide some results on the real dataset in C4 and model in Llama 3.2 1B for a variation of input context length $n$. See global response Experiments on real LLMs.

W1: 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 fusion, 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.

Thanks for your insightful comments. Our methods can easily extend to GQA setting.

In our new experiments in global response Experiments on real LLMs, we add two pruning masks to $W_K$ and $W_Q$ separately, where the two masks are the optimization variables, and the objective/loss function is calculated based on these two masks (see Line 479-502 in revision). Thus, when we consider GQA, we can use joint objective function over all attention heads in a group of GQA. For example, Llama 3.2 1B has "num_attention_heads": 32, "num_key_value_heads": 8, and “num_key_value_groups”: 4. Thus, we can use 4 pruning masks for 4 $W_Q$ and 1 pruning mask for 1 $W_K$ in this group as the optimization variables, and, correspondingly, the overall objective/loss function is the summation of 4 loss functions over 4 attention heads.

There is a small gap between our current analysis setup and the GQA extension. However, this small gap may not hurt the theoretical intuition of our non-linear approximation pruning. We are willing to discuss more per the reviewer’s request.

W2: Lacks of empirical results on pruning real-world LLM. The low error rates on synthetic data may not translate to real world dowstream performances.

Thank you for pointing out your concerns. We provide some results on the real dataset in C4 and the model in Llama 3.2 1B, where our method outperforms Wanda and SparseGPT. See global response Experiments on real LLMs for details.

Q1: How to revert the pruned fused matrix back to the query and key projection matrices respectively?

We can add two pruning masks to $W_K$ and $W_Q$, and let them be optimization variables as discussed in W1 above. In this way, we can get query and key projection matrices respectively.

Q2: How does the proposed attention pruning affect length extrapolation ability?

Thank you for your insightful question. The length extrapolation ability is a very important feature for LLMs. But we have not forsee an answer based on our current analysis directly. We will keep it in mind and leave this important topic as our future work.

W1 & W2 & Q3: Experiments with real LLMs and real data

Thank you for pointing out your concern. We provide some results on the real dataset in C4 and the model in Llama 3.2 1B, where our method outperforms Wanda and SparseGPT. See global response Experiments on real LLMs for all details.

W3: 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.

Thanks for your valuable suggestions. Our method is applicable not only to attention layers but also to MLP layers, as we primarily employ gradient descent, which can be applied to the weights of MLP layers with tailored loss functions. Note that, MLP pruning is still in the linear approximation pruning regime, so it is not our main focus since their convergence is guaranteed.

Furthermore, our approach can be seamlessly integrated with other linear pruning methods, such as Wanda [1] and SparseGPT [2], to prune MLP layers, achieving greater overall acceleration.

Q1: 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?

Thanks for your suggestions. The memory cost is not big since we have the close form of the gradient. The time cost of pruning mainly depends on the iterations which is $O(d \cdot \mathrm{poly}(n)/\epsilon)$. This is due to the relatively slow convergence of Gradient Descent. However, future work will explore second-order methods like Newton’s method or Adam optimizer and try to optimize the time and memory cost.

Q2: 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?

Thanks for pointing out. Yes, we check our assumptions on Llama 3.2 1B. See global response Assumptions verification on real LLMs.

[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