# Provable Benefits of Task-Specific Prompts for In-context Learning

arXiv: https://arxiv.org/abs/2503.02102 · HTML: https://mlsanity.com/paper/2503.02102/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 14 of its 26 references that are in the map.
- First posted on arXiv 3 Mar 2025.

## Builds on (most cited first, 14 of 14)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing](https://mlsanity.com/paper/2107.13586/) (2021)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?](https://mlsanity.com/paper/2310.08391/) (2023)
- [Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond](https://mlsanity.com/paper/2407.10005/) (2024)
- [Provable Pathways: Learning Multiple Tasks over Multiple Paths](https://mlsanity.com/paper/2303.04338/) (2023)

## Built on by (most cited first, 4 of 4)

- [Out-of-Distribution Generalization of In-Context Learning: A Low-Dimensional Subspace Perspective](https://mlsanity.com/paper/2505.14808/) (2025)
- [Stability of In-Context Learning: A Spectral Coverage Perspective](https://mlsanity.com/paper/2509.20677/) (2025)
- [The Effect of Training Task Diversity on In-Context Learning through the Lens of Low-Dimensional Subspaces](https://mlsanity.com/paper/2606.06814/) (2026)
- [Understanding Generalization and Forgetting in In-Context Continual Learning](https://mlsanity.com/paper/2605.28705/) (2026)

## Abstract (arXiv)

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.