kNN Prompting：超越上下文学习的无校准最近邻推理
kNN Prompting: Beyond-Context Learning with Calibration-Free Nearest Neighbor Inference
0 背景
当大规模语言模型（LLMs）的参数规模扩展至数十亿时，已在各种自然语言处理（NLP）任务中展现出非凡的能力。然而，这些模型的训练成本极其高昂，超出了大多数科研级或消费级设备的承受范围。LLM的一种新兴使用方式是将其托管在远程数据中心，用户或应用程序通过简单的API请求来访问。
在这种模式下，LLM的典型使用方式为上下文学习（In-Context Learning，ICL），即LLM读取并完成一个提示序列，就像其在海量文本语料上进行预训练时的方式一样。提示（prompt）通常是由若干训练示例与一个测试实例连接而成的，预测结果则通过将LLM生成的单词序列映射回标签空间来获得。
1 现有研究局限性
2.1 上下文长度的限制
尽管增加提示示例可以带来性能提升，但上下文长度的限制是一个主要瓶颈。现有的大规模语言模型通常只能处理固定长度的序列（例如1024或2048个标记），这意味着在超过此长度后，额外的训练数据无法被模型处理。因此，即使存在更多的训练数据，ICL的提升效果也受到限制。
2.2 输出分布的偏差问题
ICL的另一个脆弱性是LLMs输出分布中存在的严重偏差，这导致了显著的性能下降和不稳定性，这些现有研究已有所表明。因此，许多研究者提出了各种校准输出分布的方法。然而，由于LLMs是在通用领域的自然语言上进行预训练的，其完成一个虚构提示的能力与下游任务特定标签空间并不一致。因此，这种基于校准的方法只能在有限程度上缓解偏差。
2 研究目标
提出一种称为k近邻提示（kNN Prompting）的新方法，来解决上下文学习（ICL）中存在的两个主要挑战：上下文长度限制和输出分布偏差问题。
3 方法
3.1 基本思路
kNN提示的核心思想是把大量的训练数据保存下来，以便在需要预测的时候能找到和当前问题最相似的“邻居”来辅助模型做出决策。这种方式避免了传统方法中的复杂校准过程，也能利用更多的数据。
3.2 数据集划分
我们把训练数据分成两个部分：
- 演示集 (D)：用于帮助语言模型理解任务，并构建一个示例序列，作为模型输入的前半部分。
- 锚定集 (A)：用于保存每个示例的“特征分布”（即模型对每个锚定例子的预测概率分布），之后预测时可以用来比较。
3.3 执行步骤
1 构建数据存储库 (Meta Test)
首先，我们对每个锚定集的例子，创建一个包含演示集和锚定例子的“提示序列”（Prompt），然后让语言模型对这个提示序列进行预测，得到一个分布（也就是它认为每个词的概率）。
我们把这个分布存储下来作为锚定例子的“特征”，并且记录它的标签。这个存储库会在推理阶段用来查找。
2 预测阶段 (Formal Test)
当我们要预测一个新例子时：
- 先构建一个类似的提示序列，让模型输出一个概率分布。
- 然后，把这个分布和存储库里的分布进行对比，看哪个锚定例子的分布最接近。这个过程使用的是KL散度（即一种度量概率分布差异的方式），找到最接近的几个锚定例子。
- 最后，根据找到的锚定例子的标签，综合判断出新例子的标签。
3.4 优点
无需校准：不像传统方法需要复杂的校准，kNN提示方法只需要找到相似的邻居，就能做出预测。
突破数据限制：传统方法只能利用有限的上下文（即提示序列长度），而kNN提示可以利用所有存储下来的数据，适合大数据量任务。
4 创新点
1 打破上下文长度的限制
传统的上下文学习（In-Context Learning, ICL）方法受限于语言模型的上下文长度，通常只能在有限的文本长度内利用少量训练数据。这会限制模型的表现，尤其是在需要大量示例的复杂任务中。
kNN提示通过将锚定集的数据存储为分布特征，并在推理阶段直接查找匹配，避免了上下文长度的限制。这一创新让模型能够随意使用大量训练数据来提升预测性能，不再依赖于有限的提示序列。
2 无校准优化
传统的ICL方法往往需要复杂的“校准”步骤来减小模型输出的偏差，以确保模型的输出与特定任务的标签空间相符。校准过程耗时且难以处理大规模任务中的偏差问题。
kNN提示则不需要这种校准过程。通过将模型的预测结果存储为概率分布，再在推理阶段通过相似度匹配找到“邻居”，直接用邻居的标签来做预测，从而规避了复杂的校准问题。这大幅简化了模型的使用流程，提升了实用性。
5 实验
数据集：10个广泛应用的文本分类数据集，分别为SST2、SUBJ、MPQA、AGNews、CB、CR、DBPedia、MR、RTE、TREC
LLM：GPT-2（0.8B和1.5B参数）、GPT-2 XL、OPT系列（3B、6.7B、13B和30B参数）
表格2展示了在少样本场景下，kNN提示方法与其他基于校准的方法的对比结果。校准基线方法使用公开的代码进行复现，从而确保结果的公平性和可比性。
标记为“y”的地方表示在特定情况下需要进行截断处理。这通常是因为LLM的上下文长度限制，在处理少样本任务时可能无法加载完整的输入序列。因此，某些提示被截断以适应模型的上下文窗口。
表格中的“Partial”表示仅利用了标签词（label-words）的分布信息，而非完整的分布特征。这意味着在这种方法下，模型只考虑了标签词的概率分布，而没有参考整个预测分布，可能会影响其预测精度。
kNN提示方法的一个重要优势是可以扩展到更多的训练样本，这些样本在标准ICL方法中可能无法被有效利用。实验中，将训练样本的数量增加到128，进行比较：
- ICL Ensemble：这是一种直观的替代方案，用于扩展ICL，之前的研究也采用了这种方法。该方法通过将训练数据集拆分成多个不重叠的演示集，构建不同的提示并进行集成预测。
- 微调标准预训练语言模型（PLM）：例如BERT或GPT大型模型，使用如此数量的数据可以产生有意义的结果。
在表格3中，kNN提示方法的表现继续提升，分别比ICL和其集成方法提高了16.96和16.08（0.8B模型，平均分数）。这一结果表明，kNN提示方法在处理更多训练数据时表现更加优越。
图4：少样本场景下的数据扩展。与基于校准的方法进行比较。
图5：完全监督场景下的数据扩展。 在不同规模的语言模型（LLM）上进行的实验。
图6：与演示选择进行比较。
表4：与演示选择的上限比较。括号内的1‰表示该结果可以在1,000次搜索中实现1次。
实验对创新点支撑
1 打破上下文长度的限制
实验中，kNN提示方法在多个数据集上被验证，尤其是在需要大量示例的复杂任务中表现出色。通过将锚定集中的数据存储为概率分布特征，kNN提示能够在推理阶段直接查找与测试实例相似的邻居。这种方法使得模型在实际应用中能够使用更多的训练数据，而不受上下文长度的限制。
在实验结果中，kNN提示在使用32个训练示例时，显著提高了模型的性能，表现出更高的得分（如在图5中展示），这表明kNN提示能够有效利用更多的数据来增强模型的预测能力。
2 无校准优化
在对比实验中，研究人员将kNN提示与传统的ICL方法以及校准基线方法进行比较。实验结果显示，kNN提示在没有复杂校准步骤的情况下，依然能达到更好的性能。这是通过在推理阶段仅依赖邻居的标签进行预测，而不是尝试将模型输出与特定标签空间对齐来实现的。
数据表明，kNN提示在进行大规模任务时，其性能优势尤为明显，尤其是在模型规模超过6B参数时，这种优势更加突出（如表3和图6中所示）。这表明kNN提示不仅简化了模型使用流程，同时也提高了实用性。
6 结论
论文提出了kNN提示作为一种简单有效的解决方案，旨在推进大语言模型（LLM）推理的无梯度部署。kNN提示通过避免复杂的校准过程，显著超越了现有的基于校准的方法，尤其在少样本场景下表现出色。
kNN提示的一个主要优点是在训练数据增多且无法容纳在上下文中的情况下，能够有效扩展使用的训练数据。这一方法成功地将数据扩展的效用与模型扩展结合起来，从而提升了模型的性能。
论文强调，kNN提示在实际应用中旨在实现更有效、更高效且更适用的大语言模型利用。相比于传统的方法，kNN提示不仅能够在少样本场景下降低计算成本，还能减少在实际使用中的反复查询所带来的经济负担。
7 局限性
尽管kNN提示方法能够有效利用大量训练数据，但作者承认，在理想情况下，专门微调的LLM可能仍然在数据利用方面优于kNN提示。这意味着，如果有足够的计算资源和理想条件，微调模型可能会实现更好的性能。
虽然kNN提示在处理相对小规模的训练数据时表现良好，但当涉及更大规模的数据集时，检索模型的效率可能成为一个潜在问题。在这种情况下，数据存储和检索的开销可能会影响kNN提示的实际应用。
kNN提示的性能很大程度上依赖于训练数据的质量。如果训练数据包含噪声或偏差，这些问题可能会在推理过程中被放大，从而影响最终的预测结果。