Provable Benefits of Task-Specific Prompts for In-context Learning
This work addresses the challenge of enhancing task adaptation in AI models, offering a theoretical framework for more efficient learning, though it appears incremental as it builds on prior mathematical analyses of linear attention models.
The paper tackles the problem of improving in-context learning in language models by using task-specific prompts and prediction heads, showing that this approach facilitates a decoupling of covariance and mean in the loss landscape, with prompt-tuning explaining the conditional mean and in-context learning handling the variance.
The in-context learning capabilities of modern language models have motivated a deeper mathematical understanding of sequence models. A line of recent work has shown that linear attention models can emulate projected gradient descent iterations to implicitly learn the task vector from the data provided in the context window. In this work, we consider a novel setting where the global task distribution can be partitioned into a union of conditional task distributions. We then examine the use of task-specific prompts and prediction heads for learning the prior information associated with the conditional task distribution using a one-layer attention model. Our results on loss landscape show that task-specific prompts facilitate a covariance-mean decoupling where prompt-tuning explains the conditional mean of the distribution whereas the variance is learned/explained through in-context learning. Incorporating task-specific head further aids this process by entirely decoupling estimation of mean and variance components. This covariance-mean perspective similarly explains how jointly training prompt and attention weights can provably help over fine-tuning after pretraining.