Towards Infinite-Long Prefix in Transformer

We sincerely thank you for your reviews and helpful suggestions. We address your concerns as follows:

W1: This analysis does not cover various prompt tuning methods but is instead focused solely on P-Tuning V2.

Thank you for pointing out these insightful comments. In our method, we can have different trainable parameters for different layers to approximate P-Tuning V2. Also, we can (1) only introduce trainable parameters in the first layer or (2) share the same trainable parameters among different layers to approximate other prompt tuning methods. However, on the theoretical side, these methods may require new analysis and we leave these interesting directions as our future works.

W2: NTK-attention is not parameter-efficient.

Thank you for your comments. We used low-rank adaptation to achieve an efficient parameter representation of $Z$. In detailed, we decompose $Z(0) = Z_A(0) \cdot Z_B(0)$, where $Z_A(0) \in \mathbb{R}^{r \times s}$ and $Z_B \in \mathbb{R}^{s \times d}$. We choose $s \leq \lfloor d/2 \rfloor$ as an appropriately small integer, in Line 361-370 of the revision.

On the other hand, in updated Table 1, 2, and 3, we disclose the number of parameters each fine-tuning uses. Since the number of trainable parameters of NTK-Attention is close to LoRA but indicates better performance, this confirms the efficiency and effectiveness of NTK-Attention. We refer the reviewer to the Global Response section Supplementary Implementation Detail for more details.

W3: Computational efficiency also does not surpass that of prefix-tuning.

We acknowledge this computational efficiency limitation of NTK-Attention compared to prefix attention with $m \leq d$. However, in the scope of this paper, we focus on the comparison with prefix attention with large prefix length $m$. Since we confirm the strong learning ability of an ultra-long prefix learning in Section 3, the proposition of NTK-Attention is how to catch this strong ability efficiently. NTK-Attention provides a probability for the longer extension of prefix length in prefix-based fine-tuning methods, while only costing a few computational complexity (see Figure 3) and few trainable parameters (see Table 1, 2, 3, 4). Moreover, the superior performance of NTK-Attention we showed in our experiments is worth using such complexity in exchange for.

On the other hand, the assumption of a big $m$ is common, and it could also be supported by Reclaim the Practicality of Assumption towards $m \gg L$ and $m \gg d$ section in Global Response.

W4: The advantage of prefix-tuning lies in achieving fine-tuning level performance by adjusting only the input, enabling parallel execution across various tasks. However, since NTK-attention modifies the model itself, this advantage is lost.

Thank you for pointing this concern out. We argue this is not our issue by providing a very easy Python code (only 10 lines) in Appendix C, which is a naive implementation of NTK-Attention. This substantiates the simplicity of our method of implementation. As we access the Flash Attention function [1], we only need to add such a shortcode to implement our method.

W5: The author should conduct additional experiments across various foundation models and tasks.

Thank you for suggesting this. We add new experiments to extend the experimental scope, which is regarding scalability with large language models in real-world scenarios. We refer the reviewer to Supplementary Experiments in Global Response to see the details and results of our new experiments.

On the other hand, a prompt-based fine-tuning method processes to fit on various tasks at once fine-tuning, while our NTK-Attention can also qualify it since we can approximate a prefix attention computation with arbitrary prefix length by following Eq. (7).

References

[1] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. NeurIPS’22

Thank the authors for the response.

Regarding the first question in the "Weakness" section, compared to the previous NTK analysis, what distinguishes your approach from the original NTK analysis? It seems the only addition is an extra layer, which makes the analysis appear less novel.

Additionally, the main theme of your paper is that infinitely long prefixes can be reduced to your NTK attention. However, the final form appears quite similar to LoRA (The authors did not even mention the use of low-rank approximation in $Z$ in the initial version. ).

Regarding the practicality of long prefixes, I believe it mostly relates to hard prompting and in-context learning. However, as demonstrated in recent RAG papers (which pre-compress the context), the concept of an infinitely long prefix seems unnatural. Moreover, prior research on soft prompting suggests that longer prefixes don’t necessarily lead to better performance and may degrade it at certain points.

I also noticed that the computational complexity in Figure 1 appears to be incorrect, and this hasn't been addressed in the current version.

I feel that my concerns have not been adequately addressed in the paper.

We appreciate your time and further reply. We hope that our answer to the following can relieve your concerns.

Q1: Justification for applying a low-rank approximation to Z.

Thank you so much for your questions! In the revision Table 4 (we also highlight part of them below) on page 24 (previously used to reply to Q1 of Reviewer LAin), we can see that the low-rank approximation version, e.g, $(r,s)=(128,4)$, and large-rank version $(r,s)=(128,64)$ have similar performance, as shown in below table. Thus, we claim that $Z$ has a low-dimensional structure. We hope that this ablation study can be a justification for addressing your concerns. We will conduct more justification experiments in the next version. Furthermore, theoretically studying why $Z$ can be approximated by low rank is an interesting future direction. We leave it as a future work.

Q2: Could you clarify how your analysis can be extended to the case where P is updated or modified across layers in the forward pass?

To analyze this situation, we would like to reclaim that we are able to consider the case where the prefix matrix parameters $P$ are only in the first layer. Although the prefix matrix is involved in the computations of further layers in the model and eventually input to the final linear classifier, during the back-propagation, this could only introduce some additional terms on gradients that compare to P-Tuning V2.

In detail, assume we have $h$ layers of transformers, denoted as $\mathsf{TF}\_1, \dots, \mathsf{TF}\_h$. Then, the model is $f(x) = \mathsf{TF}\_h \circ \mathsf{TF}\_{h-1} \circ \dots \circ \mathsf{TF}\_1(x)$. Given a data point $(x,y)$, for one layer setting, we have the gradient be $\frac{d L(f(x), y)}{ d P} = \frac{d L(f(x), y)}{ d \mathsf{TF}\_1(x)} \cdot \frac{d \mathsf{TF}\_1(x) }{d P}$ by chain rule. For $h$ layer, by chain rule we still have $\frac{d L(f(x), y)}{ d P} = \frac{d L(f(x), y)}{ d \mathsf{TF}\_1(x)} \cdot \frac{d \mathsf{TF}\_1(x) }{ d P}$. Here, we can view the loss function as $L( (\mathsf{TF}\_h \circ \mathsf{TF}\_{h-1} \circ \dots \circ \mathsf{TF}\_2) \circ f’(x), y)$, where $f’ = \mathsf{TF}\_1$, rather than $L(f(x), y)$. In other words, we can absorb the latter layers into the loss function as they are deterministic.

Now we focus on these additional terms on gradients, e.g., $\frac{d L(f(x), y)}{ d \mathsf{TF}\_1(x)}$, it can also be trivially handled by a considerably large $m$ as we proved in Appendix I. Therefore, we can conclude this situation within some extra equations in our proof. We will add this additional analysis in the next version since we are now not allowed to modify the manuscript.

Thank you!

Moreover, we would like to sincerely thank you again since you provided lots of constructive suggestions and insightful questions, which helped us greatly enhance the contribution of our work and further polish our presentation. This work luckily gains good progress during the revision.

We sincerely thank you for your reviews and helpful suggestions. We address your concerns as follows:

W1: Lack of Practical Applicability of Theoretical Analysis

Thank you for your question. We acknowledge that a very large setting of prefix length $m$ is impractical in prompt-based parameter-efficient-fine-tuning methods. However, in in-context learning (ICL), chain-of-thought (CoT), and LLMs as an agent, an ultra-long prompt as the prefix for input is quite common. We refer the reviewer to Reclaim the Practicality of Assumption towards $m \gg L$ and $m \gg d$ section in the Global Response to see more practical examples that support the rationality of our assumption in this paper.

W2: Implementation Complexity of NTK-Attention

Thank you for your question. We would like to argue that the implementation is not complex, as seen in our pseudo code Line 390-404, eg, compared to prefix attention, and also the PyTorch code Line 1291-1311.

On the other hand, we would like to clarify that a main contribution of our NTK-Attention is improving the complexity from a prompt-based fine-tuning method, e.g. P-Tuning V2. Currently, P-Tuning V2 and Prefix Tuning are gaining popularity in industries, as these methods also introduce extra parameters and computational complexity, but may not affect themselves to become popular.

W3: Limited Experimental Scope

Thank you for suggesting this. We add new experiments to extend the experimental scope to relieve your issue, which is regarding scalability with large language models in real-world scenarios. It shows that our methods are effective under a wide range of model sizes. We refer the reviewer to the Supplementary Experiments section in Global Response to see the details and results of our new experiments.

We sincerely thank you for your reviews and helpful suggestions. We address your concerns as follows:

W1: Incomplete Evaluation Across Different Model Architectures

Thank you for pointing this out. We state additional experiments across LLAMA and OPT architecture to relieve this issue, which shows that our methods are effective on these model architectures. We refer the reviewer to the Supplementary Experiments section in the Global Response to see the details and results of our new experiments.

W2: Limited Applicability of NTK-Attention’s Efficiency Claims

Since we confirm the strong learning ability of an ultra-long prefix learning in Section 3, the proposition of NTK-Attention is how to catch this strong ability efficiently. NTK-Attention provides a probability for the longer extension of prefix length in prefix-based fine-tuning methods, while only costing few computational complexities (see Figure 3) and fewer trainable parameters than P-Tuning V2 when $m$ is large (see Table 1, 2, 3, 4 in revision). In short, we have compatible or better performance than LoRA, when we share the same number of trainable parameters. We refer the reviewer to the Global Response section Supplementary Implementation Detail for more details.

W3 and Q3: Inappropriate Comparison with Prefix Attention

Thank you for pointing out your concern. Carefully following the derivation in Section 4.1, the NTK-Attention method comes up from introducing a polynomial attention trick to modify the prefix attention computation. It needs to be clarified that by following Section 4.1, Section 4.2, and Section 4.3, our NTK-Attention is more similar to prefix attention both in mathematics and architecture instead of LoRA.

Moreover, prefix attention introduces extra parameters and computational complexity as we state in Figure 1, while our NTK-Attention enhances it from $mL+L^2$ to $L^2$ in complexity, from $md$ to $rd+r$ in the number of trainable parameters.

Q1: Sensitivity to Hyperparameters

Thanks for your suggestion; we add an ablation study on the hyper-parameters of NTK-Attention in the revision. We refer the reviewer to the Supplementary Experiments section of Global Response to see the details and results of the ablation study.

Q2: Applicability of the m >> L Assumption

In fact, in the next-token-prediction mechanism, in order to improve the models’ outputs, all previous tokens can be considered as a prefix of the language model. Hence, in this situation, how to find an optimal prefix for the best next-token prediction becomes a considerable problem that is interesting to study. In particular, the above is only one case that Prefix Learning can stand for, we recommend the reviewer check Reclaim the Practicality of Assumption towards $m \gg L$ and $m \gg d$ section in Global Response to see more practical examples that support the rationality of our assumption in this paper.

Q4: Additional Inference Costs Compared to LoRA

Thank you for pointing this out. Since our NTK-Attention is more similar to prefix attention, the born disadvantage of prefix attention - extra parameters and computational complexity - is unavoidable. However, our NTK-Attention still makes a great improvement compared to traditional prompt-based fine-tuning methods (see rebuttal to W3 and Q3: Inappropriate Comparison with Prefix Attention) and demonstrates high performance compared to LoRA, full fine-tuning, and P-Tuning V2 (see Table 1, 2, 3). We believe that it is worth using such complexity in exchange for.

References

[1] Xiao Liu, Kaixuan Ji, Yicheng Fu, Weng Lam Tam, Zhengxiao Du, Zhilin Yang, and Jie Tang. P-tuning v2: Prompt tuning can be comparable to fine-tuning universally across scales and tasks. ACL’22.

[2] Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficient prompt tuning. ACL’21.

[3] Rishabh Agarwal, Avi Singh, Lei M Zhang, Bernd Bohnet, Stephanie Chan, Ankesh Anand, Zaheer Abbas, Azade Nova, John D Co-Reyes, Eric Chu, et al. Many-shot in-context learning. NeurIPS’24.

[4] Yao Fu, Hao Peng, Ashish Sabharwal, Peter Clark, and Tushar Khot. Complexity-based prompting for multi-step reasoning. ICLR’22.

[5] Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. OpenAI’20.

We sincerely thank you for your reviews and helpful suggestions. We address your concerns as follows:

W1: Is the "infinite-long" or "sufficiently long" prefix assumption practical (i.e., beneficial in terms of the performance)?

Thank you for your question. Our paper considers the learning setting when $m \to \infty$, and then approximates the $m = \infty$ learning setting with NTK-attention. Thus, our NTK-attention method is not approximating small $m$.

As we claimed in Section 2, Line 174, prior works [1, 2, 3, 4] practically discovered the benefit of scaling the prefix length up in prompt-based methods, where all of them could be abstract as a learning problem of prefix matrix as we stated in Eq. (2). We clarify that the larger $m$ would provide the transformer-based model with strong learning ability, whereas it doesn’t guarantee the superior performance on generalization, the model can be over-fitted on training set to perform not well on evaluation set. Furthermore, scaling law requires comprehensive high-quality datasets since models can master complicated skills from training on them due to their strong learning capabilities under scaling law, but it may not mean that we can enhance performance by only improving the model size [5].

In short, large $m$ means better capacity, but this capacity may not be fully exploited due to different learning algorithms. On the other hand, our method does have convergence guarantees which means it can exploit the capacity and typically shows better performance.

W2: Experimental results (Tables and Figures) should indicate the number of trainable parameters for each configuration.

Thank you for pointing this out! We have updated Table 1 and Table 2 in the pdf file of our revision indicating the number of parameters of each method in our experiments, where we can change the number of parameters by adjusting $r$ and $s$, which is defined in Line 361-370. For the new experiments we add during revision, we also clearly claim the number of parameters of each method for fair comparison.

In short, we have compatible or better performance than LoRA, when we share the same number of trainable parameters. We refer the reviewer to the Global Response section Supplementary Implementation Detail for more details.

Q1.1: Does NTK-Attention still theoretically supported for the small m (~=100) case? How about the performance of NTK-Attention for small m, in terms of the number of training parameters?

As we claimed in W1, our paper considers the learning setting when $m \to \infty$, and then approximates the $m = \infty$ learning setting with NTK-attention. Thus, our NTK-attention method is not approximating small $m$. Also, we would like to clarify that $m$ is not a parameter of NTK-Attention but prefix attention. Only the choices of $r$ and $s$ (see Global Response section Supplementary Implementation Detail) will affect the number of trainable parameters in NTK-Attention.

Q1.2: I am not sure if the theoretical derivation also applies to L >> 1 case and causal attention setup, which is the most common case for Transformer-based generative models.

Thank you for pointing this concern out. In fact, assuming $L=1$ in Section 3.1 is for simplifying the setting of our analysis, it helps us to consider only the learning of one token query state (Q) in attention. We are able to extend our theoretical results to the $L \gg 1$ case easily by the NTK analysis framework. We are willing to provide more discussion per the reviewer’s request.

Q2: How should we determine the size "r" of NTK-Attention?

It depends on the performance of NTK-Attention we aim to achieve and the computational resource we have for parameter-efficient-fine-tuning. Theorem 4.1 literally says that the larger choice of the value of $r$ will lead to NTK-Attention approximating prefix attention with ultra-long prefix length with small errors. However, larger $r$ would also bring requirements of larger computational complexity, which may be difficult to implement in some computational resource-limited situations.

The choice of $r$ is based on the degree we choose for the Taylor expansion of the exponential function. When choosing $r \in [d, d^2)$, function $\phi(Q_i)^\top \phi(K_j)$ is first-order approximation to $\exp(Q_i^\top K_j / \sqrt{d})$ (see Eq. (6)), when choosing $r \in [d^2, d^3)$, function $\phi(Q_i)^\top \phi(K_j)$ is second-order approximation to $\exp(Q_i^\top K_j / \sqrt{d})$, etc. We provide the detailed implementation of the function $\phi()$ in our supplementary code as well.