分布对齐作为设计上下文学习中任务向量的准则

# 分布对齐作为上下文学习中任务向量设计的准则  
来源：https://arxiv.org/html/2605.20730  
Jihoon Kwon  
首尔大学  
kog0712@snu\.ac\.kr  
&Jiwon Choi∗  
延世大学  
jiii111@yonsei\.ac\.kr  
&Jy\-yong Sohn†  
延世大学  
jysohn1108@yonsei\.ac\.kr  

###### 摘要  
上下文学习（ICL）使得大语言模型（LLM）能够通过演示适应新任务，但随着上下文长度增加，推理成本不断攀升。任务向量通过将演示压缩为紧凑的隐藏状态表示提供了一种有前景的替代方案，但其质量仅通过下游任务准确率来评估。这一间接标准对如何设计更有效的任务向量提取方法提供的洞察有限。在本文中，我们提出，使用任务向量的推理应使其预测分布与ICL的预测分布对齐。为了量化这一点，我们引入了\(d_{\text{NTP}}\)，该指标衡量基于任务向量推理与基于ICL推理在下一个词元概率上的差异。我们的实证分析表明，\(d_{\text{NTP}}\)可作为性能代理，与下游准确率呈强烈负相关。受此启发，我们开发了线性任务向量（LTV），该方法通过闭式线性映射最小化\(d_{\text{NTP}}\)，并通过回归估计演示效果。在八个分类基准和五个LLM上，LTV始终优于现有任务向量基线，平均准确率提高9.2%，同时降低推理延迟。我们还展示了LTV在回归任务上优于基线。此外，我们研究了LTV在不同模型规模间的可迁移性——这在任务向量研究中仍处于初期阶段。具体来说，我们通过实验表明，来自较大模型的任务向量可使较小模型性能提升6.4%，这揭示了提取任务表示的一种新用途。

## 1 引言  
上下文学习（ICL）已成为一种强大的范式，只需在查询前添加带标签的演示即可使大语言模型（LLM）适应新任务[6 (https://arxiv.org/html/2605.20730#bib.bib46),47 (https://arxiv.org/html/2605.20730#bib.bib68)]。ICL已被证明在多种任务上无需任何模型参数更新即可实现显著的性能提升，且性能通常随着提供更多演示而改善[2 (https://arxiv.org/html/2605.20730#bib.bib65),9 (https://arxiv.org/html/2605.20730#bib.bib73)]。然而，这种提升是有代价的：更长的演示要么增加推理时的计算量（因为输入长度增长），要么在缓存激活时带来内存开销[32 (https://arxiv.org/html/2605.20730#bib.bib81),28 (https://arxiv.org/html/2605.20730#bib.bib31),12 (https://arxiv.org/html/2605.20730#bib.bib62),14 (https://arxiv.org/html/2605.20730#bib.bib82)]。这些计算和内存限制阻碍了ICL在资源受限条件下的实际应用。  

为解决这些限制，近期工作提出了*上下文学习中的任务向量*，作为一种免训练方法，能够在推理时无需直接使用演示即可实现任务适应[16 (https://arxiv.org/html/2605.20730#bib.bib29),41 (https://arxiv.org/html/2605.20730#bib.bib32),29 (https://arxiv.org/html/2605.20730#bib.bib33),24 (https://arxiv.org/html/2605.20730#bib.bib41)]。在该方法中，任务向量（TV）是从LLM执行ICL时的内部激活中提取的浓缩向量，封装了LLM从演示中隐式推断的任务信息[15 (https://arxiv.org/html/2605.20730#bib.bib30),50 (https://arxiv.org/html/2605.20730#bib.bib34),10 (https://arxiv.org/html/2605.20730#bib.bib67)]。将该向量应用于零样本推理（即不提供任何演示）时，模型在该任务上可获得显著的性能提升[16 (https://arxiv.org/html/2605.20730#bib.bib29)]。  

尽管在过去几年中提出了各种TV提取方法，但下游任务性能仍然是比较它们的唯一既定标准，这限制了我们对*为什么*一种方法优于另一种方法以及*如何*改进现有提取方法的洞察。  

在本文中，我们提出，与ICL的分布对齐是任务向量的一个理想属性，也是设计TV提取方法的有用准则。这一观点源于以下想法：由于任务向量的作用是浓缩演示的效果，基于TV的推理自然应产生与ICL紧密对齐的预测分布。基于这一视角，我们做出以下贡献：  

- • 我们提出了\(d_{\text{NTP}}\)指标，通过衡量基于TV推理与基于ICL推理在*下一个词元概率（NTP）*上的*差异*来量化TV方法的质量。我们通过实验表明，\(d_{\text{NTP}}\)与下游性能呈强烈负相关，可作为TV质量的指标。  
- • 我们开发了线性任务向量（LTV）方法，该方法旨在最小化\(d_{\text{NTP}}\)。具体来说，LTV采用线性映射来估计演示效果，并使用回归问题的闭式解来提取任务向量。  
- • 在实验中，LTV在八个分类基准和五个LLM上始终优于现有TV方法，平均准确率提高9.2%，同时降低推理延迟。此外，LTV在回归任务上也优于基线。最后，我们将研究扩展到任务向量的可迁移性，这是现有研究中很大程度上尚未探索的维度。通过将较大模型提取的任务向量应用于较小模型（这些模型可能容量或上下文长度有限），我们在分类准确率上获得了6.4%的提升。

## 2 相关工作  
##### 任务向量。  
开创性工作Ilharco等人[19 (https://arxiv.org/html/2605.20730#bib.bib6)]引入了任务向量的概念，将其定义为预训练模型与针对特定任务微调后的模型在参数空间中的差异。核心思想是，沿任务向量方向移动模型权重可提升在该任务上的性能[34 (https://arxiv.org/html/2605.20730#bib.bib13),52 (https://arxiv.org/html/2605.20730#bib.bib27),25 (https://arxiv.org/html/2605.20730#bib.bib19)]。近期工作表明，任务向量也可以从表示空间（如激活空间[16 (https://arxiv.org/html/2605.20730#bib.bib29),50 (https://arxiv.org/html/2605.20730#bib.bib34)]或软提示空间[5 (https://arxiv.org/html/2605.20730#bib.bib72)]）中提取。  

##### 上下文学习。  
ICL使LLM能够通过简单地将查询-标签对作为演示添加到模型输入来适应新任务[6 (https://arxiv.org/html/2605.20730#bib.bib46)]。ICL的成功引发了多种理论解释[55 (https://arxiv.org/html/2605.20730#bib.bib74),44 (https://arxiv.org/html/2605.20730#bib.bib44),3 (https://arxiv.org/html/2605.20730#bib.bib66),33 (https://arxiv.org/html/2605.20730#bib.bib16),51 (https://arxiv.org/html/2605.20730#bib.bib15),4 (https://arxiv.org/html/2605.20730#bib.bib25),27 (https://arxiv.org/html/2605.20730#bib.bib24)]。一个值得注意的工作方向将ICL解释为隐式贝叶斯推理[48 (https://arxiv.org/html/2605.20730#bib.bib38),37 (https://arxiv.org/html/2605.20730#bib.bib23),54 (https://arxiv.org/html/2605.20730#bib.bib9)]：当模型处理演示时，它隐式地从演示中推断出*潜在任务概念*，并根据由此产生的后验分布进行预测。这一观点为上下文学习中的任务向量提供了理论基础[31 (https://arxiv.org/html/2605.20730#bib.bib77)]，我们将详细描述这一研究方向。  

##### 上下文学习中的任务向量。  
ICL的一个关键限制是，随着输入长度增加，推理时的计算和内存开销会显著增大[28 (https://arxiv.org/html/2605.20730#bib.bib31),2 (https://arxiv.org/html/2605.20730#bib.bib65)]。为减少这些开销，近期工作旨在将ICL引起的任务适应内化，通过调整模型参数或激活，使演示的效果编码到模型自身中。一个研究方向是通过少量样本的参数高效微调（PEFT）来实现这一点[28 (https://arxiv.org/html/2605.20730#bib.bib31),20 (https://arxiv.org/html/2605.20730#bib.bib80),12 (https://arxiv.org/html/2605.20730#bib.bib62),26 (https://arxiv.org/html/2605.20730#bib.bib69)]。另一个研究方向探索上下文学习中的任务向量，提供了一种免训练的替代方案，实现任务适应。已有多种方法被提出用于在ICL中提取任务向量，并展示了相对于零样本推理的性能提升[16 (https://arxiv.org/html/2605.20730#bib.bib29),41 (https://arxiv.org/html/2605.20730#bib.bib32),29 (https://arxiv.org/html/2605.20730#bib.bib33),24 (https://arxiv.org/html/2605.20730#bib.bib41),28 (https://arxiv.org/html/2605.20730#bib.bib31),21 (https://arxiv.org/html/2605.20730#bib.bib4),46 (https://arxiv.org/html/2605.20730#bib.bib3)]。由于模型激活跨越多个模块，现有方法在任务向量的提取位置和方式上差异很大。这种多样性突显了需要一个直接标准来评估提取任务向量的质量。

## 3 背景  
在本节中，我们首先回顾论文中使用的相关概念和符号。第3.1节 (https://arxiv.org/html/2605.20730#S3.SS1)描述了LLM如何预测下一个词元，第3.2节 (https://arxiv.org/html/2605.20730#S3.SS2)定义了我们目标分类任务，第3.3节 (https://arxiv.org/html/2605.20730#S3.SS3)介绍了LLM的三种推理模式——零样本、ICL以及使用任务向量。  

### 3.1 模型：大语言模型（LLM）  
我们考虑预训练的自回归LLM，它们根据输入提示\(p\)（一个词元序列）预测下一个词元\(u\)。该模型由三个组件组成：  
- • 嵌入层：将提示\(p\)中的每个词元转换为嵌入向量，  
- • 由\(L\)层组成的Transformer [43 (https://arxiv.org/html/2605.20730#bib.bib49)] (TF) 解码器：将嵌入向量序列转换为隐藏状态序列，  
- • 语言建模（LM）头：基于TF的输出预测下一个词元\(u\)的概率。  

令\([\boldsymbol{h}_1(p), \boldsymbol{h}_2(p), \dots, \boldsymbol{h}_l(p)]\)表示TF解码器最后一层的隐藏状态，其中\(l\)是序列长度，\(\boldsymbol{h}_l(p) \in \mathbb{R}^d\)是\(d\)维向量。LM头通过对最后一个隐藏状态\(\boldsymbol{h}_l(p)\)应用线性投影来预测下一个词元\(u \in \mathcal{U}\)，其中\(\mathcal{U}=\{1,2,\ldots,N_{\mathcal{U}}\}\)是词表集合；每个词元由其索引表示。具体来说，下一个词元的概率计算如下：
\[
P(u \mid p) = \sigma(\boldsymbol{W}_{\mathrm{lm}} \boldsymbol{h}_l(p))[u], \quad u \in \mathcal{U},
\tag{1}
\]
其中\(\boldsymbol{W}_{\mathrm{lm}} \in \mathbb{R}^{N_{\mathcal{U}} \times d}\)表示LM头的权重矩阵，\(\sigma(\cdot)\)表示softmax函数。为简化符号，由于我们只使用最后一个词元的隐藏状态，后文将\(\boldsymbol{h}_l(p)\)简写为\(\boldsymbol{h}(p)\)。我们也用\(\text{TF}(p)\)表示通过嵌入提示\(p\)并经过TF后得到的隐藏状态\(\boldsymbol{h}(p)\)。

### 3.2 任务：分类  
尽管大语言模型（LLM）可应用于广泛的下游任务，但上下文学习中任务向量的前期工作[16 (https://arxiv.org/html/2605.20730#bib.bib29),28 (https://arxiv.org/html/2605.20730#bib.bib31),39 (https://arxiv.org/html/2605.20730#bib.bib59)]主要集中在分类设置上。沿用这一研究方向，我们也将注意力限制在分类任务上。我们将分类任务定义为给定查询-标签对\((x, y)\)的分布\(\mathcal{D}\)，其中查询\(x\)是文本序列，标签\(y\)属于任务特定的标签集\(\mathcal{C} \subseteq \mathcal{U}\)，且\(|\mathcal{C}| = K\)类。给定查询\(x\)，目标是预测其对应标签\(y\)。我们考虑*限制*在标签集\(\mathcal{C}\)上的下一个词元分布：
\[
P(c \mid p;\mathcal{C}) = \frac{P(c \mid p)}{\sum_{c' \in \mathcal{C}} P(c' \mid p)}, \quad c \in \mathcal{C}.
\tag{2}
\]
为简化符号，后文将\(P(c \mid p;\mathcal{C})\)简写为\(P(c \mid p)\)。在贪心解码中，预测标签\(\hat{y}\)由选择概率最高的类别决定：
\[
\hat{y} = \mathrm{argmax}_{c \in \mathcal{C}} \, P(c \mid p).
\tag{3}
\]

参照标题  
图1：三种推理模式的比较。在*零样本*推理模式（左）中，模型仅基于测试查询\(x_{\text{test}}\)预测下一个词元\(\hat{y}_{\text{zs}}\)。在*上下文学习*模式（中）中，模型基于演示\(Z\)与测试查询\(x_{\text{test}}\)的拼接预测下一个词元\(\hat{y}_{\text{icl}}\)。在*任务向量*模式（右）中，模型不仅基于测试查询\(x_{\text{test}}\)，还基于注入的任务向量\(\boldsymbol{v}=f(Z)\)（添加到模型激活中）预测下一个词元\(\hat{y}_{\text{tv}}\)。这里，任务向量\(\boldsymbol{v}\)由函数\(f\)利用演示\(Z\)构建。

### 3.3 方法：LLM的推理模式  
我们介绍LLM的三种推理模式：第3.3.1节 (https://arxiv.org/html/2605.20730#S3.SS3.SSS1)中的零样本推理、第3.3.2节 (https://arxiv.org/html/2605.20730#S3.SS3.SSS2)中的上下文学习（ICL）以及第3.3.3节 (https://arxiv.org/html/2605.20730#S3.SS3.SSS3)中的任务向量。

#### 3.3.1 零样本推理模式  
在零样本推理模式中，LLM为测试查询\(x_{\text{test}}\)预测\(y_{\text{test}}\)，而不提供任何目标任务的带标签示例\((x, y)\)。图1 (https://arxiv.org/html/2605.20730#S3.F1)最左侧部分展示了详细过程。首先，测试查询\(x_{\text{test}}\)通过TF获得隐藏状态\(\boldsymbol{h}_{\text{zs}} = \text{TF}(x_{\text{test}})\)。然后，LM头根据\(\boldsymbol{h}_{\text{zs}}\)计算每个类别\(c\)的概率\(P(c \mid x_{\text{test}})\)。最后，通过贪心解码选择预测标签：
\[
\hat{y}_{\text{zs}} = \mathrm{argmax}_{c \in \mathcal{C}} \, P(c \mid x_{\text{test}}).
\tag{4}
\]
在本文中，我们使用下标‘zs’表示该量适用于零样本推理。

#### 3.3.2 上下文学习（ICL）模式  
假设我们从任务分布\(\mathcal{D}\)中采样得到\(k\)个演示\(Z = \{(x_i, y_i)\}_{i=1}^k\)。如图1 (https://arxiv.org/html/2605.20730#S3.F1)中间部分所示，ICL将\(Z\)添加到测试查询\(x_{\text{test}}\)之前，并将它们传入模型。