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

• URL: http://arxiv.org/abs/2605.20730v1
• Date: Wed, 20 May 2026 05:26:38 GMT
• Title: Distributional Alignment as a Criterion for Designing Task Vectors in In-Context Learning
• Abstract summary: In this paper, we argue that inference using task vectors should align their predictive distribution with that of ICL.<n>We introduce $d_textNTP$, a metric that measures the discrepancy in next-token probabilities between task vector-based and ICL-based inference.<n>We develop Linear Task Vector (LTV), a method designed to minimize $d_textNTP$ via a closed-form linear mapping.
• Score: 6.840854574584369
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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 $d_{\text{NTP}}$, a metric that measures the discrepancy in next-token probabilities between task vector-based and ICL-based inference. Our empirical analysis reveals that $d_{\text{NTP}}$ 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 $d_{\text{NTP}}$ 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.

Related papers

• DeCoVec: Building Decoding Space based Task Vector for Large Language Models via In-Context Learning [6.461940812531588]
  textscDeCoVec captures the task essence as the difference between the output logit distributions of few-shot and zero-shot prompts.<n>textscDeCoVec consistently outperforms standard few-shot baselines, with gains up to +5.50 average accuracy.
  arXiv  Detail & Related papers  (2026-04-13T07:42:54Z)
• Model Specific Task Similarity for Vision Language Model Selection via Layer Conductance [92.72779885657373]
  We propose a framework that grounds model selection in the internal functional dynamics of the visual encoder.<n>Our approach represents each task via layer wise conductance and derives a target-conditioned block importance distribution through entropy regularized alignment.<n>Building on this, we introduce Directional Conductance Divergence (DCD), an asymmetric metric that quantifies how effectively a source task covers the target's salient functional blocks.
  arXiv  Detail & Related papers  (2026-02-01T17:29:43Z)
• From Brute Force to Semantic Insight: Performance-Guided Data Transformation Design with LLMs [48.83701310501069]
  Large language models (LLMs) have achieved notable performance in code synthesis.<n>We introduce a performance-aware, closed-loop solution that enables LLMs to autonomously engineer optimal transformations.<n>We fine-tune LLMs with Low-Rank Adaptation on a novel repository of more than 6,000 empirically evaluated PyTorch augmentation functions.
  arXiv  Detail & Related papers  (2026-01-07T11:13:02Z)
• Test time training enhances in-context learning of nonlinear functions [51.56484100374058]
  Test-time training (TTT) enhances model performance by explicitly updating designated parameters prior to each prediction.<n>We investigate the combination of TTT with in-context learning (ICL), where the model is given a few examples from the target distribution at inference time.
  arXiv  Detail & Related papers  (2025-09-30T03:56:44Z)
• Stable Diffusion Models are Secretly Good at Visual In-Context Learning [9.829303881652548]
  We show that off-the-shelf Stable Diffusion models can be repurposed for visual in-context learning (V-ICL)<n>We formulate an in-place attention re-computation within the self-attention layers of the Stable Diffusion architecture.<n>We show that this repurposed Stable Diffusion model is able to adapt to six different tasks.
  arXiv  Detail & Related papers  (2025-08-13T17:08:22Z)
• Adaptive Task Vectors for Large Language Models [14.108866468832623]
  Adaptive Task Vectors (ATV) is a simple and effective framework that dynamically generates task vectors conditioned on each input query.<n>ATV demonstrates strong performance and generalization capabilities, even for unseen tasks.
  arXiv  Detail & Related papers  (2025-06-03T22:12:28Z)
• FTP: A Fine-grained Token-wise Pruner for Large Language Models via Token Routing [17.01412432658081]
  Large language models (LLMs) have demonstrated superior performance across various tasks by adhering to scaling laws.<n>We propose a fine-grained token-wise pruning approach for the LLMs, which presents a learnable router to adaptively identify the less important tokens.<n>Our approach achieves state-of-the-art (SOTA) pruning results, surpassing other existing pruning methods.
  arXiv  Detail & Related papers  (2024-12-16T07:09:46Z)
• Multi-Task Model Merging via Adaptive Weight Disentanglement [69.7292615212444]
  We introduce an Adaptive Weight Disentanglement method for model merging.<n>We successfully extract redundant vectors, and after their subtraction, the task vectors retain robust performance.
  arXiv  Detail & Related papers  (2024-11-27T20:08:55Z)
• LoFiT: Localized Fine-tuning on LLM Representations [60.99814930367597]
  We introduce a framework called Localized Fine-Tuning on LLM Representations (LoFiT) LoFiT identifies a subset of attention heads that are most important for learning a specific task, then trains offset vectors to add to the model's hidden representations at those selected heads. For truthfulness and reasoning tasks, we find that LoFiT's intervention vectors are more effective for LLM adaptation than vectors from representation intervention methods such as Inference-time Intervention.
  arXiv  Detail & Related papers  (2024-06-03T17:45:41Z)
• Meta-Learned Confidence for Few-shot Learning [60.6086305523402]
  A popular transductive inference technique for few-shot metric-based approaches, is to update the prototype of each class with the mean of the most confident query examples. We propose to meta-learn the confidence for each query sample, to assign optimal weights to unlabeled queries. We validate our few-shot learning model with meta-learned confidence on four benchmark datasets.
  arXiv  Detail & Related papers  (2020-02-27T10:22:17Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.