Towards Infinite-Long Prefix in Transformer

TL;DR

This paper introduces a theoretically grounded method for efficient prefix tuning in transformers, enabling near-infinite-long prefix approximation with fewer parameters, and demonstrates competitive performance across multiple domains.

Contribution

It provides a convergence guarantee for ultra-long prefix training using NTK and proposes a practical algorithm for polynomial-small error approximation.

Findings

Achieves superior or competitive performance on vision, language, and math tasks.

Requires only a few extra trainable parameters instead of infinite-long prefix.

Demonstrates the effectiveness of the method through preliminary experiments.

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. Our code can be found at \url{https://github.com/ChristianYang37/chiwun/tree/main/src/NTK-Attention}.