# Learning Task Representations from In-Context Learning

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 78 references that are in the map.
- First posted on arXiv 8 Feb 2025.

## Builds on (most cited first, 30 of 49)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [CommonsenseQA: A Question Answering Challenge Targeting Commonsense Knowledge](https://mlsanity.com/paper/1811.00937/) (2018)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [Inference-Time Intervention: Eliciting Truthful Answers from a Language Model](https://mlsanity.com/paper/2306.03341/) (2023)
- [Editing Models with Task Arithmetic](https://mlsanity.com/paper/2212.04089/) (2022)
- [Steering Llama 2 via Contrastive Activation Addition](https://mlsanity.com/paper/2312.06681/) (2023)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Attention is not Explanation](https://mlsanity.com/paper/1902.10186/) (2019)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Attention is not not Explanation](https://mlsanity.com/paper/1908.04626/) (2019)
- [Dissecting Recall of Factual Associations in Auto-Regressive Language Models](https://mlsanity.com/paper/2304.14767/) (2023)
- [Word Translation Without Parallel Data](https://mlsanity.com/paper/1710.04087/) (2017)
- [Larger language models do in-context learning differently](https://mlsanity.com/paper/2303.03846/) (2023)
- [Learning to Compress Prompts with Gist Tokens](https://mlsanity.com/paper/2304.08467/) (2023)
- [In-Context Learning Creates Task Vectors](https://mlsanity.com/paper/2310.15916/) (2023)
- [Revealing the Dark Secrets of BERT](https://mlsanity.com/paper/1908.08593/) (2019)
- [Direct and Indirect Effects](https://mlsanity.com/paper/1301.2300/) (2013)
- [In-context Vectors: Making In Context Learning More Effective and Controllable Through Latent Space Steering](https://mlsanity.com/paper/2311.06668/) (2023)
- [Label Words are Anchors: An Information Flow Perspective for Understanding In-Context Learning](https://mlsanity.com/paper/2305.14160/) (2023)
- [Extracting Latent Steering Vectors from Pretrained Language Models](https://mlsanity.com/paper/2205.05124/) (2022)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [Visualizing and Measuring the Geometry of BERT](https://mlsanity.com/paper/1906.02715/) (2019)

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

- [Large Language Models Encode Semantics and Alignment in Linearly Separable Representations](https://mlsanity.com/paper/2507.09709/) (2025)
- [Do different prompting methods yield a common task representation in language models?](https://mlsanity.com/paper/2505.12075/) (2025)
- [Internal Chain-of-Thought: Empirical Evidence for Layer-wise Subtask Scheduling in LLMs](https://mlsanity.com/paper/2505.14530/) (2025)
- [Just-in-time and distributed task representations in language models](https://mlsanity.com/paper/2509.04466/) (2025)
- [Schema for In-Context Learning](https://mlsanity.com/paper/2510.13905/) (2025)
- [Transmuting prompts into weights](https://mlsanity.com/paper/2510.08734/) (2025)
- [Beyond Plain Demos: A Demo-centric Anchoring Paradigm for In-Context Learning in Alzheimer's Disease Detection](https://mlsanity.com/paper/2511.06826/) (2025)
- [Distributional Alignment as a Criterion for Designing Task Vectors in In-Context Learning](https://mlsanity.com/paper/2605.20730/) (2026)
- [Self-Improving In-Context Learning](https://mlsanity.com/paper/2605.23180/) (2026)
- [Efficient Task Adaptation in Large Language Models: A Survey of Weight-Based, Prompt-Based, and Embedding-Based Adaptations](https://mlsanity.com/paper/2610.00928/) (2026)
- [Capturing In-Context Learning Dynamics with Task Operators](https://mlsanity.com/paper/2610.01054/) (2026)

## Abstract (arXiv)

Large language models (LLMs) have demonstrated remarkable proficiency in in-context learning (ICL), where models adapt to new tasks through example-based prompts without requiring parameter updates. However, understanding how tasks are internally encoded and generalized remains a challenge. To address some of the empirical and technical gaps in the literature, we introduce an automated formulation for encoding task information in ICL prompts as a function of attention heads within the transformer architecture. This approach computes a single task vector as a weighted sum of attention heads, with the weights optimized causally via gradient descent. Our findings show that existing methods fail to generalize effectively to modalities beyond text. In response, we also design a benchmark to evaluate whether a task vector can preserve task fidelity in functional regression tasks. The proposed method successfully extracts task-specific information from in-context demonstrations and excels in both text and regression tasks, demonstrating its generalizability across modalities.