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

📄 arXiv: 2605.20730v1 📥 PDF

作者: Jihoon Kwon, Jiwon Choi, Jy-yong Sohn

分类: cs.CL, cs.AI

发布日期: 2026-05-20

备注: 9 pages, preprint

💡 一句话要点

提出线性任务向量（LTV）方法，通过分布对齐提升上下文学习的任务向量性能。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 上下文学习 任务向量 分布对齐 线性回归 模型迁移

📋 核心要点

  1. 现有任务向量方法仅通过下游任务准确率评估，缺乏对任务向量设计方向的直接指导。
  2. 论文提出分布对齐的思想，即任务向量推理的预测分布应与上下文学习的预测分布对齐。
  3. 实验表明，提出的线性任务向量（LTV）方法在多个任务和模型上优于现有方法，平均准确率提升9.2%。

📝 摘要（中文）

上下文学习(ICL)允许大型语言模型(LLM)通过示例适应新任务，但随着上下文长度的增加，其推理成本也随之增加。任务向量通过将示例压缩成紧凑的隐藏状态表示提供了一种有希望的替代方案，但其质量仅通过下游任务的准确性来评估。这种间接的标准对如何设计更有效的任务向量提取方法提供的见解有限。本文提出，使用任务向量的推理应该使其预测分布与ICL的预测分布对齐。为了量化这一点，我们引入了$d_{ ext{NTP}}$，它衡量了基于任务向量的推理和基于ICL的推理之间的下一个token概率的差异。我们的经验分析表明，$d_{ ext{NTP}}$可以作为性能代理，与下游准确性呈强烈的负相关。受此启发，我们开发了线性任务向量(LTV)，这是一种旨在通过闭式线性映射最小化$d_{ ext{NTP}}$的方法，该映射通过回归估计示例效果。在八个分类基准和五个LLM上，LTV始终优于现有的任务向量基线，平均准确率提高了9.2%，同时降低了推理延迟。我们进一步表明，LTV在回归任务上优于基线。此外，我们研究了LTV在不同模型规模之间的可迁移性；这是任务向量研究中尚未成熟的一个方面。具体来说，我们通过实验表明，来自较大模型的任务向量可以通过6.4%提高较小模型的性能，这表明提取的任务表示具有新的效用。

🔬 方法详解

问题定义：现有任务向量方法的评估标准是下游任务的准确率，这种间接的评估方式无法直接指导任务向量的设计，导致任务向量的性能提升受限。此外，现有方法忽略了任务向量推理的预测分布与上下文学习的预测分布之间的差异，这可能导致任务向量无法有效捕捉上下文学习的本质。

核心思路：论文的核心思路是，一个好的任务向量应该能够尽可能地模拟上下文学习的预测行为，即任务向量推理的预测分布应该与上下文学习的预测分布尽可能地接近。通过最小化这两种分布之间的差异，可以有效地提升任务向量的性能。

技术框架：论文提出的线性任务向量（LTV）方法主要包含以下几个阶段：1) 使用上下文学习进行推理，得到每个token的预测概率分布；2) 使用任务向量进行推理，得到每个token的预测概率分布；3) 计算两个概率分布之间的差异，使用$d_{ ext{NTP}}$指标进行量化；4) 使用线性回归方法，学习一个线性映射，将上下文学习的隐藏状态映射到任务向量，从而最小化$d_{ ext{NTP}}$。

关键创新：论文最重要的技术创新点是提出了分布对齐的思想，并将其作为设计任务向量的标准。与现有方法仅关注下游任务准确率不同，论文直接关注任务向量推理的预测分布与上下文学习的预测分布之间的差异，从而更有效地指导任务向量的设计。此外，论文还提出了$d_{ ext{NTP}}$指标，用于量化两个概率分布之间的差异。

关键设计：LTV的关键设计在于使用闭式线性映射来估计示例效果，并通过回归最小化$d_{ ext{NTP}}$。具体来说，LTV通过线性回归学习一个矩阵，将上下文学习的隐藏状态映射到任务向量。损失函数是$d_{ ext{NTP}}$，优化目标是最小化任务向量推理的预测分布与上下文学习的预测分布之间的差异。这种线性映射的设计使得LTV具有计算效率高、易于实现的优点。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LTV在八个分类基准和五个LLM上始终优于现有的任务向量基线，平均准确率提高了9.2%，同时降低了推理延迟。此外，LTV在回归任务上也优于基线。更重要的是，实验证明了LTV在不同模型规模之间的可迁移性，来自较大模型的任务向量可以通过6.4%提高较小模型的性能。

🎯 应用场景

该研究成果可应用于各种需要快速适应新任务的场景，例如智能客服、对话系统、机器翻译等。通过使用任务向量，可以显著降低推理成本，提高响应速度，同时保持较高的准确率。此外，该研究还为任务向量的跨模型迁移提供了新的思路，可以利用大型模型的知识来提升小型模型的性能。

📄 摘要（原文）

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.