Abstract
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.
Abstract (translated)
现代语言模型的上下文学习能力促使了对序列模型进行更深层次的数学理解。近期的研究表明，线性注意力模型可以通过模拟投影梯度下降迭代来隐式地从提供的上下文中学习任务向量。在本研究中，我们考虑了一种新颖的情况：全球任务分布可以被划分为条件任务分布的并集。然后，我们探讨了使用特定于任务的提示和预测头部来利用单层注意力模型学习与条件任务分布相关的先验信息的方法。我们的损失景观结果表明，特定于任务的提示促进了协方差-均值解耦，在此过程中提示调整解释了分布的条件均值，而方差则是通过上下文学习被学习/解释出来的。进一步引入特定于任务的头部有助于这个过程，完全分离了对均值和方差分量的估计。这种从协方差-均值的角度同样可以解释为什么联合训练提示与注意力权重能够确凿地帮助在预训练后进行微调的效果优于单独微调。
URL
https://arxiv.org/abs/2503.02102
