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

2026-05-20

AI Summary

This work addresses the limitation of existing task vector designs, which lack direct evaluation of alignment with the in-context learning (ICL) prediction distribution and rely solely on downstream accuracy, thereby hindering effective extraction. To overcome this, the study proposes minimizing the discrepancy between task vectors and ICL in next-token probability distributions as a core principle, introducing for the first time a distribution alignment metric, $d_{\text{NTP}}$, as a direct measure of task vector quality. Based on this criterion, the authors develop LTV, a closed-form linear mapping method. Experiments demonstrate that LTV improves average accuracy by 9.2% across eight classification benchmarks and five large language models, while significantly reducing inference latency. It also excels in regression tasks and cross-model transfer, enabling smaller models to achieve a 6.4% performance gain by leveraging task vectors from larger models.

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.

Problem

• task vectors
• in-context learning
• distributional alignment
• next-token probability
• model compression

Innovation

• Distributional Alignment
• Task Vectors
• In-Context Learning
• Linear Task Vector
• Cross-model Transfer