Towards Infinite-Long Prefix in Transformer
This work addresses the problem of parameter-efficient fine-tuning for researchers and practitioners in machine learning, offering a theoretically grounded approach with practical benefits.
The paper tackles the theoretical limitation of prefix learning in transformers by providing a convergence guarantee for training ultra-long prefixes using the Neural Tangent Kernel framework, and demonstrates that their method achieves superior or competitive performance on vision, natural language, and math data compared to existing fine-tuning methods.
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. Our code can be found at \url{https://github.com/ChristianYang37/chiwun/tree/main/src/NTK-Attention}.