Enhancing In-Context Learning Performance with just SVD-Based Weight Pruning: A Theoretical Perspective

We appreciate your response on finding our work thorough and informative. Below we address specific questions.

W1: It would be good to define deep and shallow, as these are subjective terms depending on the reference frame.

A: Indeed, there is no universally accepted definition for the terms "deep" and "shallow" as their interpretation can be subjective and dependent on the reference frame (eg. model size). Intuitively, in this paper, our definition is similar to that of Work [1]: “shallow” layers refer to those closer to the input, while “deep” layers are closer to the output. In Figure 1 and Figure 2 (Section 2), "shallow" typically denotes the first few layers, and "deep" denotes the last few layers of the network.

[1] Lean Wang et al. Label Words are Anchors: An Information Flow Perspective for Understanding In-Context Learning. EMNLP 2023.

W2: Figure 1 cpation says: "We operate on the whole of MLP or ATTN." What does this mean?

A: It refers to pruning all the weight matrices in these modules. The striking model architecture Transformer block consists of an attention (ATTN) layer and a feed-forward (MLP) layer. As mentioned in Section 2, the attention (ATTN) layers consist of key, query, value, and output matrices in both GPT-J-6B and LLAMA2-7B. The MLP layers in GPT-J-6B include input and output matrices, while in LLAMA2-7B, they are composed of up, gate, and down matrices.

W3: If as figure 1 states, you can clip 99.5% of the original weights, what happens if you just drop that layer entirely? Recent work has shown that the deeper layers can be completely dropped without much effect. [1]

• I cannot see much benefit gained from pruning part of the weights with SVD when it seems that the in nearly all cases, the benefit can be had by dropping the layer entirely.

A: Thank you for your good question. To provide a thorough and accurate answer, could you please share the reference for [1] that you mentioned? Understanding the specifics of this recent work would help us address your inquiry more effectively.

Regarding your question, I have two key points to consider:

(i) Even though we can "clip 99.5% of the original weights," the remaining weights are still significant and play a crucial role. This is because the remaining 0.5% corresponds to the larger singular values after performing SVD, which are vital for maintaining the model's performance.

(ii) As stated in Theorem 1, from the perspective of ICL (In-Context Learning) and gradient descent, each attention layer corresponds to an implicit gradient descent step (L layers with Implicit Gradient Descent Trajectories $[\Delta W_t]_1^L$/$[G_t]_t^L$).

Therefore, if we just drop the entire layer, it will result in two main issues: (a) It will become challenging to adjust the weights for another downstream task. (b) The implicit gradient update will be reduced by one step, generally leading to a decline in model performance.

Thank you again for your understanding.

W4: Is the mask on L138 supposed to represent a causal mask? If so, I do not think the notation is correct, as the Identity matrix would only have $N$ binary values which is much less than is needed for a causal mask.

W5: How can equation 1 and 2 use the same mask?

A: Our notation here follows the related work [1], which explains: Note that the prompt is asymmetric since the label for $x_{N+1}$ is excluded from the input. To reflect this asymmetric structure, the mask matrix $M$ is included in the attention.

More specifically, if you pay attention to the $(N+1)$-th item, L138 is supposed to represent a causal mask.

(For $H\in \mathbb{R}^{(dout+din)\times(N+1)}$, $HM=(h_1,..,h_N,h_{N+1})M=(h_1,..,h_N,0)=H_s$)

Besides, this mask method is used in GLM training [2]: Part A tokens can attend to each other, but cannot attend to any tokens in B. Part B tokens can attend to Part A and antecedents in B, but cannot attend to any subsequent tokens in B. So it is reasonable for equation 1 and 2 use the same mask.

[1] wangjun Ahn et al. Transformers learn to implement preconditioned gradient descent for in-context learning. NeurIPS 2023.

[2] Zhengxiao Du et al. GLM: General Language Model Pretraining with Autoregressive Blank Infilling. ACL 2022.

W6: Example 1 appears to be incorrect:

• There is no parentheses around $W^k_{V_r}+ \delta_Vh_i^{k−1}$ in the first line.
• The triangle inequality seems to say that line 2 ≥ line 1

A: Thank you for pointing out the typo (parentheses). We appreciate your attention to detail and will have made the necessary corrections. However, this minor typo does not affect the final conclusion. A detailed analysis is as follows:

In triangle inequality, it is ture that $||\mathbf{a}||_2+||\mathbf{b}||_2\geq ||\mathbf{a}+\mathbf{b}||_2$ for any vector $\mathbf{a},\mathbf{b}$ (same dimension).

But in this Example, we did not use the triangle inequality:

$||W^k\_{V\_r}h\_i^{k−1}+\delta\_Vh\_i^{k−1}||\_2^2=||W^k\_{V\_r}h\_i^{k−1}||\_2^2+2(W^k\_{V\_r}h\_i^{k−1})^T(\delta\_Vh\_i^{k−1})+||\delta\_Vh\_i^{k−1}||\_2^2=||W^k\_{V\_r}h\_i^{k−1}||\_2^2+||\delta\_Vh\_i^{k−1}||\_2^2$.

Note that $W^k_{V_r}=U_{:r}\Sigma_{:r}V_{:r}^T$ and $\delta_V=U_{r:}\Sigma_{r:}V_{r:}^T$, and $U_{:r},U_{r:}$ are orthometric (properties of SVD).

Therefore, $(W^k_{V_r}h_i^{k−1})^T(\delta_Vh_i^{k−1})=(h_i^{k−1})^TV_{:r}\Sigma_{:r}(U_{:r}^TU_{r:})\Sigma_{r:}V_{r:}^Th_i^{k−1}=0$, line 2 = line 1.

Q1: I don't believe the clipping process was adequately explained. Once the SVD operation is done, do you clip starting from the largest singular value? or starting from the smallest?

A: As mentioned in Section 2, "The optimal rank-r approximation and SVD," we retain the components corresponding to the largest r singular values. Therefore, the clipping process begins with the smallest singular values.

Thanks so much for your time and insightful comments. Please find our point-by-point response below.

W1: The theoretical analysis primarily focuses on linear attention, which may not fully capture the complexities of standard Softmax attention used in most transformer models

A: Firstly, to facilitate qualitative analysis, our main text primarily focuses on the theoretical aspects of linear attention. This is our first step on this issue, and our goal is to provide insightful conclusions under simplified conditions.

Secondly, a considerable portion of the works also consider from the linear attention to explore the ICL [1,2].

However, we also discuss the standard Softmax attention setting in Appendix B.2. Specifically, it can be considered that the mapping function $\phi$, or rather the effect of the Softmax function, is to project the original features into a higher-dimensional space to capture more profound features. Subsequently, learning and meta-optimization are conducted under this new feature space.

$\hat{h}\_{N+1}=h\_{N+1} + \Delta W\_{icl}^{'}\phi(W\_Q h\_{N+1}),\ \Delta W\_{icl}^{'}=\frac{1}{D^{'}} \left[\sum\_{i=1}^N (W\_VH\_s)\_i \otimes \phi(W\_KH\_s)\_i \right]$

[1] Ekin Akyiurek et al. What learning algorithm is in-context learning? investigations with linear models. ICLR 2023.

[2] Johannes von Oswald et al. Transformers learn in-context by gradient descent. ICML 2023.

W2: The proposed algorithm is derivative-free, but the search for optimal clipping rates may still be computationally expensive for very large models or datasets

A: (1) Our primary objective in this work is to provide a general theoretical framework that reveals the underlying mechanism behind the phenomenon that SVD-based weight pruning can enhance ICL performance. Based on this theory, a simple algorithm is provided to demonstrate the effectiveness of theoretical analysis in guiding experimental procedures.

In Theorem 2, the key insight is that managing the expected generalization error by controlling the norm of $\Delta W_t/G_t$, indicating that pruning methods are not unique. Therefore, some computationally efficient pruning methods could be considered: (a) Random SVD, (b) Magnitude-based pruning [1], and so on.

[1] Wen, W et al. Learning structured sparsity in deep neural networks. NeurIPS 2016.

W3: There is a substantial lack of comparison with other pruning methods: the study focuses on SVD-based pruning but doesn't compare it with other pruning techniques, which could provide context for the method's effectiveness

A: Thanks for your suggestion. Similar to the answer to the above W2, we primarily focus on theoretical analysis. And the simple algorithm is just used to validate the effectiveness of the theory, so we have not provided comparative methods. There is a possibility that better pruning methods exist.

For instance, if we use magnitude-based pruning on matrix $A$, the $||A||_F$ will be decrease.

W4: Poor language, frequent typos, and grammatical errors are significant issues in this paper. This does not help readability, and would likely be a barrier to publication in its current form.

A: Thank you for your feedback. We sincerely apologize for the language issues, typos, and grammatical errors present in the paper. Rest assured, we are fully committed to addressing these concerns and will do everything possible to correct the issues before resubmission. Your patience and understanding are greatly appreciated, and we are dedicated to improving the paper to meet the required standards.

W5: An essential part of the paper, which is the discussion of related works is not part of the main text. Furthermore, this discussion is prone to criticism. For example, quoting the seminal paper by Frankle and Carbin as an example of low-rank properties of neural networks is clearly misleading. I think that this discussion should be an essential part of the main text, and should also be substantially revised in order to avoid conceptual confusions.

A: (1) Due to page limitations, we previously omitted the detailed discussion of related works in the main text, but we will include it in future versions. (2) We acknowledge your viewpoint and will make substantial revisions accordingly. Our findings show that pruning does not necessarily worsen performance (Note that performance depends not only on generalization but also on optimization). In some datasets, it may even improve results. This aligns with previous work, which we will discuss in detail in the Related Works section, such as [1]. In future research, we will study the effects of pruning different layers and the performance variations across different tasks.

[1] Pratyusha Sharma et al. The truth is in there: Improving reasoning in language models with layer-selective rank reduction. ICLR 2024.

Q1: I suggest the Authors to address 1) the concerns regarding a better framing of the work in the current literature. In particular, there is a large body of evidence that is growing on the resilience of LLMs to pruning (see for example "The Unreasonable Ineffectiveness of the Deeper Layers", Gromov et al, https://arxiv.org/pdf/2403.17887, and references therein), and 2) the quality of writings by proofreading it together with a native speaker

A: Thanks. We will definitely revise the literature review (Gromov et al,and references therein) & our writings.

However, please let us demonstrate the advantages of our theoretical analysis (pruning weight) compared to the example (remove layer) you provided. If we just remove the entire layer, it will result in two main issues: (a) It will become challenging to adjust the weights for another downstream task. (b) The implicit gradient update will be reduced by one step (Theorem 1), perhaps leading to a decline in model performance.

Dear Reviewer Nxqr,

We sincerely appreciate your time and effort in reviewing our manuscript and providing valuable feedback.

As the author-reviewer discussion phase nears completion, we wish to confirm whether our responses have effectively addressed your concerns. We provided detailed responses to your concerns a few days ago and hope they have adequately resolved any issues. If you require further clarification or have any additional concerns, please do not hesitate to contact us. We are more than willing to continue our communication with you.

Best regards.

Thank you for your detailed review and valuable comments. Below we address specific questions.

W1: More details of algorithms is not shared. e.g. the range / number of clipping rate candidates set.

A: Firstly, the details of the algorithm can be reviewed in the code provided. Specifically, the set of clipping rate candidates is predefined. In our study, the clipping rate candidates are set as shown in Figures 1 and 2: [0, 0.1, 0.5, 0.75, 0.9, 0.95, 0.99, 0.995]. Besides, we analyze the impact of different hyperparameters through comparative experiments (details in Section 2 and Section 3.4), as detailed below.

• Clipping rate $\xi$：We search for the optimal $\xi$ in the predefined clipping rate candidates.

• Predefined module $\mathcal{O}$ : The module containing the target pruning weights, which can be chosen from ['k_proj', 'q_proj', 'v_proj', 'out_proj', 'fc_in', 'fc_out', 'all', 'mlp', 'attn']

• Selected layer $l$ : The layer containing the target pruning weights. For examlpe: In Section 2, we mainly focuse on comparing the impact of weight pruning to the first two and the last two layers of the model.

• ICL shot number : The demonstration number in ICL, we analyze the effect of different ICL shot numbers in Section 2.2.

W2: In experiments result of C.5, the optimal $\xi$ varies a lot across different tasks and different modules. However, this phenomenon is not touched in the theoretical part.

A: That's a great question! The impact of different tasks and various module prunings on performance is indeed significant.

However, our aim here is to provide a general analytical framework, so we have not delineated the impact of different layers/modules on performance in detail.

In our theory, the simplification of the MLP and the handling method for the ATTN layer （mentioned in the main text） ：

• Feed-forward (MLP) layer：In Appendix B.3, it can be seen as a dimensional adaptation to correspond to residual connection. Then we can get the Implicit Gradient Descent with MLP: $\Delta W_{icl}^{''}=W_{MLP}\left( \sum_{i=1}^N W_V{h_i} \otimes W_K{h_i} \right) W_Q$

• Attention (ATTN) layer:

  (i) In main text Section 3.2, our theoretical analysis is mainly focuses on linear attention setting.

  $\hat{h}\_{N+1} = h\_{N+1}+ \Delta W\_{icl} h\_{N+1}, \Delta W\_{icl}=\left( \sum\_{i=1}^N W_V{h_i} \otimes W\_K{h\_i} \right) W\_Q$

  (ii) In Section Appendix B.2, it can be considered that the mapping function $\phi$, or rather the effect of the Softmax function, is to project the original features into a higher-dimensional space to capture more profound features. Subsequently, learning and meta-optimization are conducted under this new feature space.

  $\hat{h}\_{N+1}=h\_{N+1} + \Delta W_{icl}^{'}\phi(W\_Q h\_{N+1}), \Delta W_{icl}^{'}=\frac{1}{D^{'}} \left[\sum_{i=1}^N (W_VH_s)_i \otimes \phi(W_KH_s)_i \right]$

Then if one consider more than one layer, a more detailed analysis can be provided within the current framework to depict the interaction relationship between the input feature and these different modules.

Q1: Matrix condition number is an option for the indicator. But could there be more options, such as compute the decreasing rate of eigenvalues? Because when p=2, conditional number only leverages two values among all the eigenvalues.

A: Good question! Let's first review our theory in Theorem 2: Modulating the norm of $[G\_t]_1^L$ or $[\Delta W\_{t}]_1^L$ could enhance performance when utilizing ICL.

Our primary objective in this work is to provide a general analytical framework. The simple algorithm is employed to demonstrate the effectiveness of theoretical analysis in guiding experimental procedures. So, Any indicator that can guide the control of norms is a potential option. For instance:

(i) Consider the behavior of the first few and last few singular values.

(ii) Assessing the intrinsic rank of the matrix can be valuable, as it offers another dimension of understanding beyond just the condition number. [1] [2]

[1] Armen Aghajanyan et al. Intrinsic Dimensionality Explains the Effectiveness of Language Model Fine-Tuning. URL https://arxiv.org/abs/2012.13255.

[2] Edward Hu et al. Lora: Low-rank adaptation of large language models. ICLR 2022.

Q2: Could authors provide further more clarification why optimal $\xi$ varies, and is there a way to explain this phenomenon under current theoretical framework provided in this paper?

A: (i) As we mentioned in Section 3.3 and Section C.4, the implicit gradients produced by Transformers in practical applications are noisy due to factors such as the extent of model pre-training and data characteristics (eg. ICL shot number/task difficulty). Therefore, $[\Delta W_t]_1^L$/$[G_t]_t^L$ in Theorem 1 have varies noise. That is why optimal $\xi$ varies.

(ii) Besides, expected generalization error (Theorem 2) = population risk - empirical risk ($L_{\mu}$-$L_{H_s}$). we search the optimal $\xi$ on the validation set and subsequently evaluate it on the test set (detailed in Section 3.4).

On the one hand, in Theorem 2, clipping weight controls the F-norm of the implicit gradient ( $[\Delta W\_t]\_{1}^{L}/[G\_t]\_{1}^{L}$), which reduces the expected generalization error, on the other hand, we can evaluate the empirical risk ($L_{H_s}$) by observing the model's performance on the validation set and estimate the population risk ($L_{\mu}$). (Note that different types of tasks affect $L_{H_s}$ differently, and different modules have varying impacts on the expected generalization error. See A to W2.)

Thank you for your review, and for your thoughtful comments and questions. They will certainly help improve the revised paper. Please see our response below.

A to Weaknesses: Thanks to Reviewer Zh6a for raising the issue, which gives us the opportunity to clarify this matter.

(1) Our goal is to provide insightful conclusions under simplified conditions. However, we also discuss the standard Softmax attention setting in Appendix B.2 and feed-forward (MLP) layers in Appendix B.3.

(2) The goal of this paper is to study ICL from a generalization perspective. According to the traditional statistical learning viewpoint, performance can be defined by the sum of optimization error and generalization error. This paper reveals that SVD decomposition is beneficial for generalization error, but does not analyze its impact on optimization. Given that theoretical analysis of ICL is rapidly developing, the impact of weight pruning on optimization has not yet been fully explored. We will consider this aspect in our future work.

A to Q1 (How should Theorem 2 be interpreted? Can this also be applied to empirical risk ?):

Expected generalization error (Theorem 2) = population risk ($L\_{\mu}$) - empirical risk ($L\_{H\_s}$).

More specifically, on the one hand, in Theorem 2, clipping weight controls the F-norm of the implicit gradient $([\Delta W\_t]\_1^L/[G_t]\_1^L)$), which reduces the expected generalization error, on the other hand, we can evaluate the empirical risk ($L\_{H\_s}$) by observing the model's performance on the validation set. If the generalization error is determined, it is at least possible to estimate the population risk ($L_{\mu}$). Thus, the most challenging aspect is addressing the generalization error. For the direct theoretical analysis of empirical risk, we will consider this aspect in our future work.

A to Q2 (Could this mean that fewer context examples are more robust for SVD weight pruning?):

That's a pretty worthy question! Yes, fewer context examples can indeed lead to more robust results in SVD weight pruning.

(1) Experimentally, in Section 2.2, we analyze the effect of different ICL shot numbers. As shown in Figure 2, with a decrease in the number of shots, the rate of performance collapse in the model slows down after a sharp reduction at the shallow layer.

(2) Theoretically, $\Delta {W}(N)-\Delta {W}(N^{'}) = \left( \sum\_{i=N^{'}+1}^N{W}\_V{{h}\_i} \otimes {W}\_K{{h}\_i} \right) {W}\_Q,$ so it is supposed that the implicit gradient is more sensitive for SVD weight pruning.

(3) The robustness mainly concerns the operation "SVD weight pruning", rather than the model's performance.

A to Q3 (What would be the effect of pruning only a single module?):

Good question! In our theory, the example (pruning only K or V) is provided in Appendix B.6, demonstrating how weight pruning can affect the norm of $[\Delta W_t]_1^L$/$[G_t]_t^L$. Thus, it may confer advantages on the performance of Transformers in ICL inference. However, in Theorem 2 (Remark 5), it also suggests that modifications to the shallow layers have a less steady impact.

Additionally, please review the supplementary experimental results provided below. (We choose key projection matrix and select the 3-layer (large matrix condition number) of GPT-J-6B, others are the same to Section 2.1).

Task	Optimal $\xi$	Accuracy/F1 Improve
SST-2	0.0	0.7828 - 0. 7828
RTE	0.5	0.5413 - 0.5413
COPA	0.995	0.53 - 0.54

A to Q4 (What would happen if we apply the same clipping rate to other datasets?):

As we mentioned in Section 3.3 and Section C.4, the implicit gradients produced by Transformers in practical applications are noisy due to factors such as the extent of model pre-training and data characteristics (eg. ICL shot number/task difficulty). Therefore, $[\Delta W_t]_1^L$/$[G_t]_t^L$ in Theorem 1 exhibit varying levels of noise, causing the optimal clipping rate to vary among different tasks, as it is dependent on the specific task and data. So if we apply the clipping rate (0.95) as used in SST-2 to other datasets, the model performance can either improve or deteriorate. Additionally, It is possible to conduct a certain number of experiments to find a range of optimal clipping rate that is broadly applicable.

A to Limitations:

• For the equation between line 228 and 229.

$||W^k\_{V\_r}h\_i^{k−1}+\delta\_Vh\_i^{k−1}||\_2^2=||W^k\_{V\_r}h\_i^{k−1}||\_2^2+2(W^k\_{V\_r}h\_i^{k−1})^T(\delta\_Vh\_i^{k−1})+||\delta\_Vh\_i^{k−1}||\_2^2=||W^k\_{V\_r}h\_i^{k−1}||\_2^2+||\delta\_Vh\_i^{k−1}||\_2^2$.

Note that $W^k_{V_r}=U_{:r}\Sigma_{:r}V_{:r}^T$ and $\delta_V=U_{r:}\Sigma_{r:}V_{r:}^T$, and $U_{:r},U_{r:}$ are orthometric (properties of SVD).

Therefore, $(W^k_{V_r}h_i^{k−1})^T(\delta_Vh_i^{k−1})=(h_i^{k−1})^TV_{:r}\Sigma_{:r}(U_{:r}^TU_{r:})\Sigma_{r:}V_{r:}^Th_i^{k−1}=0$.

• For the equation for about F-norm in lines 230,231,232. Thank you for pointing out the typo (231). We appreciate your attention to detail and will have made the necessary corrections. However, this minor typo does not affect the final conclusion. A detailed analysis is as follows:

  For any vector $\mathbf{a}$ and $\mathbf{b}$, $\text{rank}(\mathbf{a}\otimes\mathbf{b})=1$, so the matrix ($\mathbf{a}\otimes\mathbf{b}$) only has one non-zero singular value.

  Combined with $||A||\_F=\sqrt{\sum\_{i}\sigma\_i^2(A)}$ and singular value is nonnegative, we can get

  $||\mathbf{a}\otimes\mathbf{b}||\_F = \sqrt{\sum\_{i}\sigma\_i^2(\mathbf{a}\otimes\mathbf{b})}=\max\_i[\sigma\_i(\mathbf{a}\otimes\mathbf{b})]$. Therefore, the unique non-zero singular value will decrease after performing SVD on $\mathbf{a}$ and $\mathbf{b}$.