论文笔记--kNN PROMPTING: BEYOND-CONTEXT LEARNING WITH CALIBRATION-FREE NEAREST NEIGHBOR INFERENCE

WITH CALIBRATION-FREE NEAREST NEIGHBOR INFERENCE)

1. 文章简介

• 标题：kNN PROMPTING: BEYOND-CONTEXT LEARNING WITH CALIBRATION-FREE NEAREST NEIGHBOR INFERENCE
• 作者：Benfeng Xu, Quan Wang, Zhendong Mao, Yajuan Lyu, Qiaoqiao She, Yongdong Zhang
• 日期：2023
• 期刊：ICLR

2. 文章概括

  文章提出了一种kNN prompting的方法，解决了传统In-Context Learning(ICL)d中Context长度受限的问题，且可以充分利用大量标注样本进行推理。在文章测试的10个文本分类任务中，kNN prompting相比于ICL和calibration-based方法性能有了显著提升。

3 文章重点技术

3.1 In-Context Learning(ICL)

  ICL的概念可参考博客[1]。这里只给出ICL的标准模型：给定标记训练数据集$\mathcal{T}=\{(x_i,y_i)\}$，其中$y_i\in \mathcal{Y}$表示样本$x_i$的标签，给定LLM为$\theta$，则ICL的任务是给定测试样本$x_{test}$时，根据训练样本构建其prompt：$$P=\pi (x_1,y_1)\oplus \pi (x_2,y_2)\oplus \cdots \oplus \pi (x_{|\mathcal{T}|},y_{|\mathcal{T}|})\oplus \pi (x_{test},\ast )$$，其中$\pi$表示verbalization操作，即将$x,y$映射到其对应的标准prompt格式。比如$1->positive,0->negative$或者$0->white,1->red,2->blue....$。下表是文章采用的一些prompt的样例：

模型学习到一个$y\in \mathcal{Y}$：$${\hat{y}}_{test}=\underset{y\in \mathcal{Y}}{\mathrm{argmax}}(v|P,\theta )=\underset{y\in \mathcal{Y}}{\mathrm{argmax}}(\pi (y)|P,\theta ).$$
  ICL满足大模型的幂律法则，也就是说，当标记数据$|\mathcal{T}|$增大时，模型的性能也会提升。但ICL的一个劣势是需要将标记数据放在context中输入到LLM，而LLM一般的输入长度限制为1024或2048等，这就造成了我们只能使用一小部分的标记数据作为上下文。

3.2 kNN prompting

  为了解决上述问题，文章提出了kNN prompting方法。给定训练集$\mathcal{T}$，我们将其分为示范集(demonstration set)$\mathcal{D}$和锚集(anchor set) $\mathcal{A}$，如下图所示，kNN prompting分为两个阶段

• Meta Test：首先我们利用$\mathcal{A}$中所有标记样本：对任意$(x_i^a,y_i^a)\in \mathcal{A}$，我们将$x_i^a$和所有$\mathcal{D}$中的标记样本组成prompt，得到prompt为 $$P_i=\pi (x_1^d,y_1^d)\oplus \pi (x_2^d,y_2^d)\oplus \cdots \oplus \pi (x_{|\mathcal{D}{|}^d},y_{|\mathcal{D}|}^d)\oplus \pi (x_i^a,\ast )$$，再将prompt放入LLM中得到一个分布$p(v|P_i,\theta )$。这里我们不将$y$转化成$v$，而是直接将$k_i=p(v|P_i,\theta )$缓存，记作key representation(KR)。从而我们构建出了一组 { k i , y i a } \{k_i, y_i^a\} {ki​,yia​}的数据库
• Formal Test：正式推理阶段，对每个测试样本$x_{test}$，我们类似上面的公式构建prompt: $$P_{test}=\pi (x_1^d,y_1^d)\oplus \pi (x_2^d,y_2^d)\oplus \cdots \oplus \pi (x_{|\mathcal{D}{|}^d},y_{|\mathcal{D}|}^d)\oplus \pi (x_{test},\ast )$$，并得到$p_{test}=p(v|P_{test},\theta )$。然后我们将此分布与数据库中所有分布进行KL散度比较：$$D_{KL}(p_{test}||k_i)=\sum _{v}p(v|P_{test},\theta )\log \frac{p(v|P_{test},\theta )}{p(v|P_i,\theta )}$$。然后得到$x_{test}$对应的预测输出为$${\hat{y}}_{pred}=\underset{y\in \mathcal{Y}}{\mathrm{argmax}}\sum _{t\in N{N}^k(p_{test},\mathcal{K})}1(y_i^1=y)$$，其中$N{N}^k(\ast ,\mathcal{K})$表示测试分布的$k$个近邻。翻译下上式，简单来说先计算当前分布和历史所有分布的KL散度，找到最近的$k$个分布，这$k$个分布对应的类别最多的即为预测值。
  

4. 数值实验

  文章将kNN prompting和ICL、Calibration-based方法进行了比较，发现kNN prompting显著提高了baselines。可能原因为

• kNN prompting信息使用更充分：使用了全部的标记样本
• 多个标签的单词之间彼此竞争可能会影响传统ICL判断。
    此外，数值实验表明，kNN prompting显著提高了prompt方法的鲁棒性:
  
    最后，文章通过t-SNE将得到的数据进行可视化。如下图所示，LLM学习到的分布并不总是按照标签聚类的，这使得ICL有一定的概率预测错误。如下表中的样本A，真实标签为artist，但ICL预测其为book。但通过kNN prompting得到它的三个近邻标签均为artist，从而可NN prompting给出了artist的正确标签。
  

4. 文章亮点

  文章提出了kNN prompting，将标记样本的LLM分布缓存，在推理阶段获得与测试样本分布的KL散度最接近的k个近邻，将近邻中最多的分类作为预测结果。kNN prompting解决了ICL受context长度限制的问题，充分利用了标记数据，且提升了ICL的鲁棒性。

5. 原文传送门

kNN PROMPTING: BEYOND-CONTEXT LEARNING WITH CALIBRATION-FREE NEAREST NEIGHBOR INFERENCE

6. References

[1] 论文笔记–Learning To Retrieve Prompts for In-Context Learning