Towards Infinite-Long Prefix in Transformer

Authors: Yingyu Liang, Zhenmei Shi, Zhao Song, Chiwun Yang
Primary area: learning theory
Keywords: large language model, prefix learning, neural tangent kernel

Abstract:
Prompting and context-based fine-tuning methods, which we call Prefix Learning, have been proposed to enhance the performance of language models on various downstream tasks. They are empirically efficient and effective, matching the performance of full parameter fine-tuning, but the theoretical understandings are limited. In this paper, we aim to address this limitation by studying their ability from the perspective of prefix length. 
In particular, we provide a convergence guarantee for training an ultra-long prefix in a stylized setting using the Neural Tangent Kernel (NTK) framework. Based on this strong theoretical guarantee, we design and implement an algorithm that only needs to introduce and fine-tune a few extra trainable parameters instead of an infinite-long prefix in each layer of a transformer, and can approximate the prefix attention to a guaranteed polynomial-small error.
Preliminary experimental results on vision, natural language, and math data show that our method achieves superior or competitive performance compared to existing methods like full parameters fine-tuning, P-Tuning V2, and LoRA. This demonstrates our method is promising for parameter-efficient fine-tuning.

Decision: reject
Reviewer scores: 6, 6, 5, 3

Reviews and discussion:
<thread>
creation date:
1730361362176

thread_type:
official_review

<thread_object>
creator: reviewer_ofkk

<content>
summary:
This paper presents an analysis of ultra-long prefix vectors (P-Tuning V2), focusing exclusively on this version of prefix-tuning. The study utilizes the Neural Tangent Kernel (NTK) to examine ultra-long prefix vector convergence and subsequently proposes NTK-attention, derived from observations in the analysis.

strengths:
The paper analyzes P-Tuning V2 through the lens of the Neural Tangent Kernel (NTK), offering a fresh perspective on prefix-tuning. The analysis is precise, highly detailed, and appears to have involved considerable effort.

weaknesses:
1. This analysis does not cover various prompt tuning methods but is instead focused solely on P-Tuning V2, which employs different embeddings per layer. In this case, it is evident that the model is equivalent to using two linear layers with tied weights, making this insight neither new nor surprising. This raises questions about the novelty of the analysis itself. The difficulty in analyzing prefix-tuning lies in the fact that the prefixes passed through the input layer interact with each other in every layer via attention before moving to the next layer. However, this paper lacks analysis in that regard.

2. In my view, contrary to the author’s claim, NTK-attention is not parameter-efficient. The number of parameters is quite large:  as shown on Line 345, $r = d$, resulting in $d^2+d$ parameters. This count is significantly higher than the $md$ parameters used in prefix-tuning (experiments were conducted with prefix counts $m$ ranging from 1 to 200). This discrepancy calls into question whether the comparative experiments with P-Tuning V2 and LoRA are indeed fair. The author notes using LLama 3.1 1B for the experiments, and given that this model has an embedding dimension of 2048, a considerable number of parameters are indeed used (more than x10). If the author intends to argue that their method is more parameter-efficient than prefix-tuning, they should disclose the parameter counts used in Tables 1 and 2.

3. Computational efficiency also does not surpass that of prefix-tuning. The $\text{O}(\text{L}^2)$ notation in Figure 1 is incorrect; due to the $\Phi(Q)Z$ operation in Line 5 of Algorithm 2, the complexity of NTK-attention is $\text{O}(\text{(L+d+1)L})$ ($d=r$ as mentioned above). Although Appendix Figure 3 provides a running time comparison, it is not a fair comparison with $d=32$ as the value of $m$ is ranging from $1$~$2^{16}$ (The complexity of prefix-tuning $\text{O}(\text{(L+m)L})$ is quite large in this setting. And nobody use more than 1K tokens for the prefix-tuning.). For a fair experiment, the author should test with common embedding size $d$ such as 2048 or 4096, frequently used in various foundation models.  In this experimental setting, since $m\leq d$ in most cases, NTK-attention is likely to incur an even greater computational load.

4. 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.

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

questions:
See weaknesses.

+) If the mentioned weaknesses are adequately addressed in the rebuttal, I would be willing to raise the score.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730644037497

thread_type:
official_review

<thread_object>
creator: reviewer_ud7d

<content>
summary:
This paper aims to enhance the theoretical understanding of "Prefix Learning" in Transformer-based models. The authors use the Neural Tangent Kernel (NTK) framework to guarantee convergence when training ultra-long prefixes, essentially addressing the theoretical limits of prefix length. They propose a new NTK-Attention algorithm that approximates infinite-long prefix learning using a few extra trainable parameters per layer, thus reducing computational complexity and memory usage. Experimental results show that NTK-Attention achieves competitive performance compared to methods like P-Tuning V2 and LoRA on datasets involving vision, natural language understanding, and mathematical inference.

strengths:
Theoretical Contribution: The paper provides a solid theoretical analysis based on NTK to understand the efficiency and convergence of Prefix Learning, which deepens the theoretical foundation of this area.
Efficient Training: The NTK-Attention algorithm significantly reduces computational complexity and memory requirements by replacing ultra-long prefixes with a small number of additional trainable parameters.
Experimental Validation: The authors validate the performance of NTK-Attention on diverse datasets (vision, natural language understanding, and math inference) and demonstrate superior or competitive results compared to existing methods.

weaknesses:
Lack of Practical Applicability of Theoretical Analysis: While this paper presents an NTK-based analysis to establish the theoretical foundation of Prefix Learning, there is a lack of empirical evaluation regarding its applicability to large-scale models. In particular, there is little discussion on how the approach involving ultra-long prefixes can be practically applied in real-world industrial settings.

Implementation Complexity of NTK-Attention: The NTK-Attention algorithm introduces a small number of additional trainable parameters to replace ultra-long prefixes, but this also brings about implementation complexity and the need for additional hyperparameter tuning. This could limit the method's applicability and make it challenging to ensure model stability in practical applications.

Limited Experimental Scope: The experiments validate the performance of NTK-Attention on vision, natural language, and mathematical inference datasets, but the scope of the experiments is limited. In particular, there is a lack of validation regarding scalability with very large language models and performance in diverse real-world scenarios, leaving questions about the generalizability of the proposed method. Additional evaluations on a wider range of datasets and real-world applications are needed.

questions:
see weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730657977667

thread_type:
official_review

<thread_object>
creator: reviewer_lain

<content>
summary:
This paper explores “Prefix Learning,” a method that enhances large language models (LLMs) by fine-tuning an added prefix matrix in each layer of a transformer. The authors investigate the effect of infinitely long prefixes on model performance through theoretical analysis using the Neural Tangent Kernel (NTK) framework. Based on this theoretical insight, they propose “NTK-Attention,” an efficient approximation algorithm that replaces the infinite prefix matrix with a finite number of trainable parameters. This method achieves polynomially small errors while significantly reducing computational and memory costs.

strengths:
1. Strong Theoretical Foundation: The paper provides a rigorous theoretical analysis of prefix learning with ultra-long prefixes, leveraging the Neural Tangent Kernel (NTK) framework. This contributes a new perspective to understanding prefix-based methods in transformers, especially regarding convergence and scaling laws.
2. Efficient Fine-Tuning Method: The proposed NTK-Attention significantly reduces the number of trainable parameters and computational complexity by replacing the large prefix matrix with two smaller trainable parameters,  Z  and  k . This approach is practical and resource-efficient, addressing scalability issues that arise with traditional long prefix methods.
3. High Performance Across Multiple Domains: Experimental results demonstrate NTK-Attention’s competitive or superior performance on tasks in natural language understanding, vision, and mathematical inference. This versatility and effectiveness across diverse datasets and tasks underscore its broad applicability.
4. Reduced Memory and Computational Costs: NTK-Attention achieves considerable savings in memory and computation compared to standard prefix learning methods. This efficiency is particularly beneficial for large-scale models where memory and computational constraints are a concern.
5. Reproducibility: Although the code is not included in the OpenReview submission, the arXiv version of the paper provides the accompanying code, enhancing reproducibility.

weaknesses:
1.Incomplete Evaluation Across Different Model Architectures: The experiments presented focus primarily on a few specific architectures, such as pretrained ViT and ChatGLM3-6B. The generalizability of NTK-Attention across a broader range of transformer architectures and model sizes remains unexplored, which limits the applicability of the findings to other models commonly used in different domains​.
2. Limited Applicability of NTK-Attention’s Efficiency Claims: The reduced computational complexity of NTK-Attention compared to Prefix Attention relies on the condition  m \gg L , where  m  is the prefix length and  L  is the input length. However, such conditions may not be common in many real-world scenarios, which restricts the practical applications of NTK-Attention to specific cases.
3. Inappropriate Comparison with Prefix Attention: The paper’s comparison with Prefix Attention may be less relevant, as NTK-Attention does not introduce additional prefill tokens and is structurally closer to LoRA. From Table 1, NTK-Attention achieves comparable performance to LoRA, but it introduces additional inference costs that LoRA does not have. Given these trade-offs, LoRA may provide better practical value, limiting NTK-Attention’s advantage over existing methods.

questions:
1. Sensitivity to Hyperparameters: The paper does not extensively cover the sensitivity of NTK-Attention to key hyperparameters, such as the dimension of the feature mapping  r  and the learning rates for  Z  and  k . Could the authors provide an ablation study or guidelines for tuning these parameters? This would assist researchers in optimizing NTK-Attention across different model architectures and tasks.
2. Applicability of the  m >> L  Assumption: The computational efficiency of NTK-Attention compared to Prefix Attention assumes  m >> L , which may not hold in many real-world scenarios. Could the authors provide more context or examples where this assumption is practical? Alternatively, could they discuss how NTK-Attention performs in settings where m and L are closer in scale?
3. Justification for Comparison with Prefix Attention: Given that NTK-Attention does not introduce additional prefill tokens, it is more similar to LoRA than Prefix Attention. Could the authors clarify why they chose to emphasize the comparison with Prefix Attention rather than focusing on comparisons with methods like LoRA? This would help align the experimental evaluations with the method’s structural similarities.
4. Additional Inference Costs Compared to LoRA: NTK-Attention introduces additional inference costs, as highlighted in Table 1, where its performance is comparable to LoRA. Could the authors elaborate on scenarios where the increased inference cost of NTK-Attention would be justified over LoRA? This could clarify the practical advantages of NTK-Attention in resource-constrained environments.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730704659856

thread_type:
official_review

<thread_object>
creator: reviewer_jxua

<content>
summary:
The paper proposes NTK-Attention, an alternative to Prefix-Tuning method for Transformer fine-tuning. Specifically, NTK-Attention is derived from theoretical understanding that Prefix-Tuning can be approximated using Neural Tangent Kernel (NTK) when the prefix length is sufficiently large. The key difference of NTK-Attention is to utilize a learnable weight matrix multiplied to Query, instead of concatenating a learnable prefix to Key and Value. Due to the architecture of NTK-Attention, it can avoid quadratic memory increase as the prefix length increases. Experimental results show that NTK-Attention achieves comparable performance to competitors, P-Tuning-v2 and LoRA.

strengths:
* The proposed NTK-based analysis is novel and opens a new direction to understand prefix-tuning. Although the extremely-long prefix setup is not widely used yet, it seems to be an interesting direction for Transformer-based models.

weaknesses:
* Is the "infinite-long" or "sufficiently long" prefix assumption practical (i.e., beneficial in terms of the performance)? In Table 1, m=200 is worse than m=100 case. The "Many-shot In-context Learning" paper (Agarwal et al., 2024) empirically showed that using lots of few-shot examples help, but the paper increases the length of input prompt rather than the number of training parameters; the result would not directly mapped one-to-one.
* Experimental results (Tables and Figures) should indicate the number of trainable parameters for each configuration. To be fair, it is important to compare the performance under similar resource usage constraint.

questions:
* Theoretical questions
  * 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?
  * 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.
* How should we determine the size "r" of NTK-Attention?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732235625099

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response (Part 2)

comment:
### Reclaim the Practicality of Assumption towards $m \gg L$ and $m \gg d$

Since Reviewer LAin, ud7D, and oFkK mentioned their concerns about the practicality of prefix length $m \rightarrow +\infty$ assumption, we here explain and address as follows:

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. In this paper, we examine the implications of using a significantly large prefix length, denoted as $m \gg L$ and $m \gg d$, which is prevalent across various prompt-based methods. The primary objective of Prefix Learning is to enhance the LLMs' outputs by identifying an advanced prefix during the generation process. For instance, the search for optimal example pairs to improve ICL and the development of prompt engineering tailored for agent frameworks to address specific task requirements often necessitate the use of exceptionally long prefixes. Moreover, given the modern application demands related to long-context scenarios, optimizing previous tokens to improve next-token prediction can be framed as a prefix optimization problem. Thus, a thorough investigation into the optimization of infinitely long prefixes is essential for understanding the theoretical significance of the prefix matrix in LLMs.

### Clarification of notation

We found that the $r$ in LoRA is reused with the $r$ notation in our NTK Attention. We re-denote the $r$ in LoRA as $r’$ in the revision.


### References

[1] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. Meta’22.

[2] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. ICLR’17.

[3] Mitch Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building a large annotated corpus of english: The penn treebank. Computational linguistics’93.

[4] Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. ACL’16.

[5] Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. ICLR’22.

[6] Xiao Liu, Yanan Zheng, Zhengxiao Du, Ming Ding, Yujie Qian, Zhilin Yang, and Jie Tang. Gpt understands, too. AI Open’23.

[7] 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.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732235732004

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response (Part 1)

comment:
We thank all constructive and valuable reviews from all reviewers.

We are grateful that reviewers all agree that our NTK analysis in Transformers is novel and theoretically solid, contributing considerably to understanding prefix learning in pretrained transformers-based architectures. Besides, reviewers (LAin, ud7D) admit the potential of our NTK-Attention as a parameter-efficient-fine-tuning method that approximates that ultra-long prefix attention computation and demonstrates superior empirical results.


We have updated a **revision** for our draft. We summarize all the major updates (in **brown** color) we made in the revision. All line numbers in the rebuttal correspond to the revised version.
- Line 102-107: we supplement to introduce new experiments we add in the revision.
- Line 188-197: we add an explanation for the practicality of infinite-long prefix length assumption $m \gg L$ and $m \gg d$
- Line 367-394: we reclaim the implementation detail of our method, and add an analysis of the number of trainable parameters compared to other popular methods.
- Line 381-388: we clarify the number of trainable parameters of each fine-tuning method in Table 1
- Line 440-441: we clarify the setting of hyper-parameters we use for fine-tuning ChatGLM3-6B with NTK-Attention
- Line 459-467 & 486-505: we add a detailed instruction for new experiments and we state the results of new experiments in Table 3
- Line 474: we clarify the setting of hyper-parameters we use for fine-tuning LLAMA with NTK-Attention on math inference tasks
- Line 489-495: we clarify the number of trainable parameters of each fine-tuning method in Table 2
- Line 520-523 & 1253-1288: we add detailed instructions for the ablation study and we provide the results of the ablation study in Table 4
- Line 1291-1311: we provide a 10-line naive implementation of NTK-Attention by Python code to support the implementation simplicity of our method.

Here, we address some common questions. 

### Supplementary Implementation Detail
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. 

Since given $r$ and $s$ as two hyper-parameters in NTK-Attention, for each attention layer in transformer-based architecture, we denote $\beta:= \frac{r}{d}$, then the number of trainable parameters could be computed by $(\beta s+ \beta + s)d$ where integer $\beta \ge 1$ and $s \leq \lfloor d/2 \rfloor$. This is more flexible when adjusting the practical efficiency needs. For LoRA with its hyper-parameter $r' \leq \lfloor d/2 \rfloor$, where $r'$ is the rank number used for approximation, its number of trainable parameters is $4r'd$ and for prefix attention with its hyper-parameter $m \ge 1$, its number of trainable parameters is $md$ in each attention layer. By choosing $(\beta s+ \beta + s) \leq 4r'$, the higher efficiency of NTK-Attention compared to LoRA will be satisfied. The above comparison is also in Line 371-377 of the revision. 

In our first submitted manuscript, for the convenience of readers' understanding and to make the paper focus on our main message of theoretical analysis, we did not disclose the specific implementation of a trainable matrix $Z$, but the details we claim below can be found in the code of the supplementary materials in our first submitted version. 

### Supplementary Experiments

Since Reviewer LAin, ud7D, and oFkK mentioned our limitation on the experimental scope, we add two more experiments to support our work.

**Evaluation on Language Modeling Tasks.** In this experiment, we focus on the scalability of NTK-Attention on a family of language models of different sizes, the OPT family with the model sizes 125M, 350M, 1.3B, 2.7B and 6.7B [1]. We introduce three text datasets, which are WikiText-103 [2], Penn TreeBank [3], and LAMBADA [4], to compare the scalability of NTK-Attention with LoRA [5] and P-Tuning V2 [6, 7]. As we choose $r'=8$ for LoRA, $m=32$ for P-Tuning V2, and $r=2d$ and $s = 10$ for our NTK-Attention, the numbers of trainable parameters are aligned to the same as $32d$ for each attention layer. The results are stated in Table 3, which shows the improvement of NTK-Attention compared to baselines when scaling the model size.

**Ablation Study.** We validate the sensitivity of hyper-parameters $r$ and $s$ and give the results in Appendix B.3. The results firstly indicate that choosing $r=d$ and $s=4$ is enough for high-performance fine-tuning on LLAMA-3.1-8B. Also, we follow Table 4 to suggest choosing a larger value of $r$ primarily instead of $s$ to achieve supernal accuracy.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732789303588

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Second version revision update

comment:
Per reviewer oFkK valuable suggestions, we made a second revision, where the update is in the color **purple** (we also keep our old updates). We summarize all major updates below. Note that all line numbers in the rebuttal correspond to the second revision version.

- Line 54-70: we modify Figure 1 and its caption to clarify the complexity and number of parameters of our NTK-Attention as reviewer oFkK requires.
- Line 94-95: we modify the introduction for low-rank approximation to matrix $Z$.
- Line 1140-1182: we update the detailed version of NTK-Attention with low-rank approximation to matrix $Z$ and update the newest complexity and the memory usage of trainable parameters.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper addresses the theoretical underpinnings and practical applications of Prefix Learning in enhancing the performance of large language models. By utilizing the Neural Tangent Kernel (NTK) framework, the study provides a theoretical guarantee for the convergence when training an ultra-long prefix, which was previously not well-understood. Some reviewers questioned the practicality of the assumption of ultra-long prefixes and the complexity of implementing the proposed NTK-Attention, suggesting that these aspects may limit the method’s real-world applicability.  There are concerns about the fairness and adequacy of the comparative analysis with existing methods. Specifically, details about parameter counts and computational efficiency need clearer exposition to validate the claimed advantages. AC agrees with reviewers, while this work presents significant theoretical contributions and innovative methods, requires more rigorous empirical support and a clearer exposition of its comparative advantages over existing methods.