NN Prompting: Beyond-Context Learning with Calibration-Free Nearest Neighbor Inference

Original Paper: https://arxiv.org/abs/2303.13824

By: Benfeng Xu, Quan Wang, Zhendong Mao, Yajuan Lyu, Qiaoqiao She, Yongdong Zhang

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, kNN 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: kNN 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: kNN 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.

Summary Notes

Beyond-Context Learning with kNN Prompting: A New Way to Boost Language Model Performance

Large language models (LLMs) are at the forefront of AI, offering remarkable capabilities in processing and generating text that resembles human writing.

However, they face significant challenges, such as limited context lengths and biases in their output, which hinder their full potential. This has led to the exploration of new methods to enhance their performance effectively.

Challenges with Traditional Learning Methods

LLMs traditionally rely on in-context learning (ICL), where prompts are designed with a mix of training and test examples. This approach, though useful, has its limitations:

• It doesn't scale well; adding more examples eventually doesn't improve performance due to the models' maximum token limits.
• The output distributions are biased, affecting both performance and stability.

Introducing kNN Prompting

kNN Prompting emerges as a novel solution to these challenges, moving away from directly predicting from LLM distributions. Instead, it uses a two-step process:

• Meta Test: Creating a datastore by querying the LLM with training examples and saving the language modeling distributions.
• Formal Test: For each test instance, the closest matches from the datastore are retrieved based on KL divergence, and predictions are made based on these nearest neighbors.

Advantages of kNN Prompting

• No Need for Calibration: It simplifies the prediction process by eliminating the need to align output distributions with specific labels.
• Scales Beyond Context Limits: Unlike traditional ICL, kNN Prompting can handle increasing amounts of training data without being constrained by fixed context lengths.

Testing and Results

Through experiments across 10 text classification tasks and various LLM sizes, kNN Prompting showed:

• Superior performance compared to calibration-based methods and traditional ICL, especially when dealing with large volumes of training data.
• Consistent effectiveness across different model sizes and datasets, proving its versatility.

How kNN Prompting Stands Out

While there are several methods to improve ICL, such as prompt engineering and retrieval-based techniques, kNN Prompting shines due to its simplicity and effectiveness. Its success in other machine learning areas further supports its utility in enhancing LLMs.

Advantages and Challenges

kNN Prompting leverages the scalability of LLMs and vast datasets without the need for retraining or fine-tuning.

However, it's essential to be aware of potential limitations, like performance in highly imbalanced data scenarios or specific biases in the data.

Looking Ahead in LLM Deployment

kNN Prompting marks a significant advance in overcoming the challenges of data scalability and model utilization.

By avoiding the computational demands of traditional methods, it provides a promising avenue for practical LLM applications. As AI evolves, such methodologies will be crucial in maximizing LLM utility.

Ethical Considerations

It's imperative to address biases in LLM outputs carefully. Future efforts should focus on integrating efficient methods for assessing and mitigating these biases, ensuring AI advancements are both responsible and beneficial.

Innovation in LLM enhancement, like kNN Prompting, equips us better to address the ongoing challenges in AI.

For AI engineers in enterprise settings, adopting such cutting-edge techniques is vital for unlocking the full capabilities of language models.