# Distributional Alignment as a Criterion for Designing Task Vectors in In-Context Learning

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 44 of its 54 references that are in the map.
- First posted on arXiv 20 May 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Qwen3 Technical Report](https://mlsanity.com/paper/2505.09388/) (2025)
- [Qwen2.5 Technical Report](https://mlsanity.com/paper/2412.15115/) (2024)
- [The Power of Scale for Parameter-Efficient Prompt Tuning](https://mlsanity.com/paper/2104.08691/) (2021)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [Character-level Convolutional Networks for Text Classification](https://mlsanity.com/paper/1509.01626/) (2015)
- [Editing Models with Task Arithmetic](https://mlsanity.com/paper/2212.04089/) (2022)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [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)
- [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)
- [Function Vectors in Large Language Models](https://mlsanity.com/paper/2310.15213/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [In-context Vectors: Making In Context Learning More Effective and Controllable Through Latent Space Steering](https://mlsanity.com/paper/2311.06668/) (2023)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [Task Arithmetic in the Tangent Space: Improved Editing of Pre-Trained Models](https://mlsanity.com/paper/2305.12827/) (2023)
- [Many-Shot In-Context Learning](https://mlsanity.com/paper/2404.11018/) (2024)
- [Hate Speech Dataset from a White Supremacy Forum](https://mlsanity.com/paper/1809.04444/) (2018)
- [What and How does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalization](https://mlsanity.com/paper/2305.19420/) (2023)
- [In-Context Learning through the Bayesian Prism](https://mlsanity.com/paper/2306.04891/) (2023)
- [Which Attention Heads Matter for In-Context Learning?](https://mlsanity.com/paper/2502.14010/) (2025)
- [When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformers](https://mlsanity.com/paper/2504.10957/) (2025)
- [Knowledge Composition using Task Vectors with Learned Anisotropic Scaling](https://mlsanity.com/paper/2407.02880/) (2024)

## Abstract (arXiv)

In-context learning (ICL) allows large language models (LLMs) to adapt to new tasks through demonstrations, yet it suffers from escalating inference costs as context length increases. While task vectors offer a promising alternative by compressing demonstrations into compact hidden-state representations, their quality has been evaluated only through downstream task accuracy. This indirect criterion provides limited insight into how to design more effective task vector extraction methods. In this paper, we posit that inference using task vectors should align their predictive distribution with that of ICL. To quantify this, we introduce dNTP, a metric that measures the discrepancy in next-token probabilities between task vector-based and ICL-based inference. Our empirical analysis reveals that dNTP serves as a performance proxy, exhibiting a strong negative correlation with downstream accuracy. Motivated by this, we develop Linear Task Vector (LTV), a method designed to minimize dNTP via a closed-form linear mapping that estimates demonstration effects through regression. Across eight classification benchmarks and five LLMs, LTV consistently outperforms existing task vector baselines, improving average accuracy by 9.2% while reducing inference latency. We further show that LTV outperforms the baselines on regression tasks. Moreover, we investigate the transferability of LTV across different model scales; an aspect that has remained nascent in task vector research. Specifically, we empirically show that task vectors from a larger model can enhance a smaller model's performance by 6.4%, suggesting a new utility for extracted task representations.