Fine-Tuning Without Forgetting In-Context Learning: A Theoretical Analysis of Linear Attention Models
| Authors | Chungpa Lee et al. |
| Year | 2026 |
| Field | NLP |
| arXiv | 2602.23197 |
| Download | |
| Categories | cs.CL, cs.LG, stat.ML |
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 empirically validate our theoretical results.
Engineering Breakdown
Plain English
This paper analyzes why fine-tuning large language models degrades their ability to learn from in-context examples (few-shot prompting), a critical capability of modern LLMs. The authors use linear attention models as a theoretical lens to characterize exactly how fine-tuning modifies attention parameters and under what conditions this hurts few-shot performance. They discover that fine-tuning all attention parameters causes the most damage, but selectively freezing certain parameters during fine-tuning can preserve in-context learning while still improving zero-shot performance. This is a theoretical contribution that explains a real production tension: you want models that work without examples to reduce inference cost, but not at the expense of flexibility on novel tasks.
Core Technical Contribution
The core novelty is a formal theoretical analysis of the trade-off between fine-tuning for zero-shot performance and preserving in-context learning capability in transformer models. Rather than empirically observing that fine-tuning hurts few-shot performance, the authors provide mathematical characterization of how different fine-tuning objectives modify attention parameters in ways that degrade in-context adaptation. They specifically identify that unrestricted parameter updates harm few-shot learning, while selective parameter freezing during fine-tuning can maintain in-context capability. The key insight is using linear attention models as a tractable theoretical framework to derive conditions and guarantees about when and why this degradation occurs, moving beyond black-box empirical observations.
How It Works
The approach models transformer attention as linear transformations and analyzes how standard fine-tuning objectives (like supervised loss on downstream task demonstrations) modify the attention weight matrices Q (query), K (key), and V (value). During pre-training, the model learns to use attention for in-context learning—storing and retrieving information from prompts within the context window. When fine-tuning occurs, the objective shifts to optimize for zero-shot performance on a specific task, which can reorient the attention parameters away from the general in-context learning behavior. The authors characterize this mathematically by tracking how gradient updates during fine-tuning change the learned attention subspaces, and identify critical parameter dimensions that, when frozen, preserve the pre-trained in-context mechanism while still allowing adaptation. The theoretical framework shows trade-offs: fully fine-tuning attention matrices optimizes for the target task but catastrophically interferes with few-shot adaptation, while selective updates (e.g., freezing specific parameter groups) maintain a Pareto frontier between zero-shot and few-shot performance.
Production Impact
For engineers deploying LLMs at scale, this provides actionable guidance on fine-tuning strategies that don't sacrifice model flexibility for single-task optimization. In production, you often need models that work both zero-shot (to reduce prompt engineering and inference cost) and few-shot (to adapt to novel variants not seen during training), and this work explains how to achieve both. The practical recommendation is to use selective parameter freezing or LoRA-like approaches that only update a subset of attention parameters during fine-tuning, rather than full model fine-tuning. This reduces compute cost during fine-tuning (fewer parameters to update), maintains model versatility across tasks, and provides theoretical justification for why parameter-efficient methods are superior to full fine-tuning—not just empirically, but by architectural design. The trade-off is slightly reduced performance on the target task (since you're not fully optimizing), but the flexibility gain usually outweighs this in multi-task production environments.
Limitations and When Not to Use This
The analysis relies on the linear attention approximation, which is a simplification of actual softmax attention in production transformers—the theoretical guarantees may not transfer perfectly to real nonlinear attention mechanisms used in GPT-4 style models. The paper focuses on attention parameters and doesn't deeply analyze how fine-tuning other components (embeddings, feed-forward layers, layer normalization) affect in-context learning, so the complete picture of full model fine-tuning is incomplete. The theoretical framework assumes a specific structure of in-context learning that may not hold across all model scales or training regimes; very large models or those trained with different objectives might have different mechanisms underlying few-shot adaptation. Additionally, the paper doesn't provide empirical validation on large-scale production models or comprehensive benchmarks comparing the proposed selective freezing approach against other parameter-efficient methods (LoRA, prefix tuning, adapters) on real downstream tasks, leaving a gap between theory and practice.
Research Context
This work builds on the growing literature observing that fine-tuning hurts in-context learning (empirical observations from scaling studies), and contributes the missing theoretical framework to explain why this occurs. It relates to broader research on parameter-efficient fine-tuning (LoRA, adapters, prompt tuning), providing theoretical justification for why these methods preserve pre-training capability better than full fine-tuning. The paper fits into the larger research direction on understanding and preserving foundation model capabilities during adaptation, alongside work on catastrophic forgetting, continual learning, and multitask learning. This opens research directions on: (1) extending the theory to nonlinear attention, (2) characterizing trade-offs across more than two objectives, and (3) developing fine-tuning algorithms that automatically identify which parameters to freeze based on the analysis.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::