[论文解读] $k$NN Prompting: Beyond-Context Learning with Calibration-Free Nearest Neighbor Inference

Benfeng Xu, Quan Wang+4 more|arXiv (Cornell University)|Mar 24, 2023

Domain Adaptation and Few-Shot Learning被引用 12

一句话总结

该论文介绍了 kNN Prompting，一种无需校准、无梯度的方法，使用对语言模型分布的最近邻搜索来利用超出上下文长度的无限训练数据，在多种大型语言模型上实现了强劲的少-shot 和全监督性能。

ABSTRACT

In-Context Learning (ICL), which formulates target tasks as prompt completion conditioned on in-context demonstrations, has become the prevailing utilization of LLMs. In this paper, we first disclose an actual predicament for this typical usage that it can not scale up with training data due to context length restriction. Besides, existing works have shown that ICL also suffers from various biases and requires delicate calibration treatment. To address both challenges, we advocate a simple and effective solution, $k$NN Prompting, which first queries LLM with training data for distributed representations, then predicts test instances by simply referring to nearest neighbors. We conduct comprehensive experiments to demonstrate its two-fold superiority: 1) Calibration-Free: $k$NN Prompting does not directly align LLM output distribution with task-specific label space, instead leverages such distribution to align test and training instances. It significantly outperforms state-of-the-art calibration-based methods under comparable few-shot scenario. 2) Beyond-Context: $k$NN Prompting can further scale up effectively with as many training data as are available, continually bringing substantial improvements. The scaling trend holds across 10 orders of magnitude ranging from 2 shots to 1024 shots as well as different LLMs scales ranging from 0.8B to 30B. It successfully bridges data scaling into model scaling, and brings new potentials for the gradient-free paradigm of LLM deployment. Code is publicly available.

研究动机与目标

• 通过重复使用提示长度之外的训练数据，解决上下文相关学习（ICL）的上下文长度限制。
• 通过利用 LM 分布进行匹配来代替直接的标签-标记对齐，消除对校准技巧的依赖。
• 通过将训练数据使用与提示长度解耦，促进可扩展的数据利用，从而实现对更大数据集的学习。
• 在不同数据集和模型规模上研究鲁棒性、划分策略，以及基于锚点的表示的有效性。

提出的方法

• 将训练数据划分为演示集 D 和锚点集 A。
• 用包含 D 的提示以及每个锚点 x^a 向 LLM 进行查询，以获得完整的 LM 分布 p(v|P,θ)，并将其缓存为键 k_i，标签 y^a 作为值存储在数据存储中。
• 推理时，计算测试实例的 p(v|P_test,θ)，并使用 KL 散度 D_KL(p_test || k_i) 在数据存储中检索最近的 k 个邻居。
• 通过对最近的 k 个邻居的标签进行聚合来预测 y：通过类似多数投票的方式取 argmax_y ∑_{i in NN^k(p_test, K)} 1(y_i^a = y)。
• 该框架是无校准的，因为它避免将 LM 输出与固定的任务特定标签空间对齐，并通过使用锚点集合来利用无限数据来支持超越上下文的学习。
• 不需要训练；该方法设计为可与现有的 LLM 和训练数据一起工作。

实验结果

研究问题

• RQ1 在少量样本设置下，无校准最近邻推断是否能优于基于校准的 ICL 方法？
• RQ2 在上下文长度之外可用的更多训练数据时，kNN Prompting 如何扩展，跨不同模型规模和任务？
• RQ3 演示集与锚点集的划分对性能的影响如何，方法对不同的训练分割有多鲁棒？
• RQ4 利用完整 LM 分布（相对于仅标签词）进行检索是否能提升或降低准确性与鲁棒性？

主要发现

• 在严格的少-shot 设置下领先于最先进的基于校准的方法（在某些任务上达到 +7.07）。
• 在超越上下文范围的扩展方面表现出色，随着训练数据从 2 次到 1024 次的增长以及从 0.8B 到 30B 参数的 LLM 的跨越，提升稳定。
• 仅需 32 次示例时，其高峰期实现平均增益 +13.58（0.8B 模型），在全监督设置下最高可达 +18.84。
• 相比于原生 ICL，表现出鲁棒性提升，输出不稳定性下降，跨种子的一致性标准差减小。
• 在各种数据规模下领先于 ICL Ensemble 和标准微调，并且在模型尺寸范围包括 0.8B–30B 时保持竞争力。