Fine-Tuning Without Forgetting In-Context Learning: A Theoretical Analysis of Linear Attention Models
Abstract
Transformer-based large language models exhibit in-context learning, enabling adaptation to downstream tasks via few-shot prompting with demonstrations. In practice, such models are often fine-tuned to improve zero-shot performance on downstream tasks, allowing them to solve tasks without examples and thereby reducing inference costs. However, fine-tuning can degrade in-context learning, limiting the performance of fine-tuned models on tasks not seen during fine-tuning. Using linear attention models, we provide a theoretical analysis that characterizes how fine-tuning objectives modify attention parameters and identifies conditions under which this leads to degraded few-shot performance. We show that fine-tuning all attention parameters can harm in-context learning, whereas restricting updates to the value matrix improves zero-shot performance while preserving in-context learning. We further show that incorporating an auxiliary few-shot loss enhances in-context learning primarily on the target task, at the expense of degraded in-context learning ability on tasks not seen during fine-tuning. We provide empirical evidence from synthetic and real-world datasets consistent with the qualitative predictions of our theory.
Lay Summary
Large language models can learn new tasks from examples given in the input, but they are often further trained to perform a specific task without such examples. This can improve performance on that task, but may reduce the model’s ability to learn from examples later. This paper studies why this happens and shows that updating only a specific part of the model can improve task performance while preserving much of its flexible learning ability.