kNN-Prompt: Nearest Neighbor Zero-Shot Inference (2205.13792v2)

Abstract: Retrieval-augmented LLMs (LMs) use non-parametric memory to substantially outperform their non-retrieval counterparts on perplexity-based evaluations, but it is an open question whether they achieve similar gains in few- and zero-shot end-task accuracy. We extensively study one such model, the k-nearest neighbor LM (kNN-LM), showing that the gains marginally transfer. The main challenge is to achieve coverage of the verbalizer tokens that define the different end-task class labels. To address this challenge, we also introduce kNN-Prompt, a simple and effective kNN-LM with automatically expanded fuzzy verbalizers (e.g. to expand terrible to also include silly and other task-specific synonyms for sentiment classification). Across nine diverse end-tasks, using kNN-Prompt with GPT-2 large yields significant performance boosts over strong zero-shot baselines (13.4% absolute improvement over the base LM on average). We also show that other advantages of non-parametric augmentation hold for end tasks; kNN-Prompt is effective for domain adaptation with no further training, and gains increase with the size of the retrieval model.

kNN-Prompt: Nearest Neighbor Zero-Shot Inference

The paper "kNN-Prompt: Nearest Neighbor Zero-Shot Inference" explores the enhancement of zero-shot inference in LLMs (LMs) through retrieval-augmented techniques. Traditionally, retrieval-augmented LMs leverage non-parametric memory, which has shown significant improvements in perplexity-based evaluations. This paper extends the exploration to zero-shot and few-shot task evaluations, demonstrating the potential of enhanced retrieval using a $k$-nearest neighbor LLM ($k$NN-LM).

Key Contributions

The primary challenge addressed in this work is achieving comprehensive coverage of verbalizer tokens that define end-task class labels. The paper introduces kNN-Prompt, an extension of the $k$NN-LM model, integrating fuzzy verbalizers which automatically expand the set of tokens corresponding to each output label. This process involves associating task-specific synonyms such as expanding "terrible" to include "silly" for sentiment classification. The method shows a substantial average performance boost of 13.4% over strong zero-shot baselines across a diverse set of nine tasks when applied with GPT-2 large models.

Performance and Findings

1. Improved Performance: kNN-Prompt was evaluated on a range of tasks including sentiment analysis, topic classification, and entailment. The methodology demonstrated significant improvements compared to both baseline LMs and standard $k$NN-LMs, showcasing its robustness and versatility across domains.
2. Fuzzy Verbalizers: The introduction of fuzzy verbalizers plays a crucial role in bridging the gap between sparse $k$NN distributions and the dense representations needed for effective task completion. This not only increases the support for output labels within the $k$NN retrieval but also mitigates issues related to verbalizer sensitivity.
3. Zero-Shot and Few-Shot Adaptability: The kNN-Prompt model is adaptable for both zero-shot and few-shot scenarios, providing consistent improvements without additional training. The paper also demonstrates the model's proficiency in domain adaptation by effectively utilizing domain-specific datastores.
4. Scalability: Performance gains were shown to scale with the size of the retrieval model. Larger retrieval models lead to better results, albeit with increased memory and computational demands.

Implications and Future Directions

The implications of this work lie in the potential of retrieval-augmented LMs to serve broader applications in AI by enhancing task-specific inference without extensive task-specific training data. The findings suggest that such models can be effectively adapted using domain-specific data, providing practical solutions for real-world applications where specific training data might be limited or costly to obtain.

Future research could investigate optimizing datastore sizes and efficiency in retrieval processes, perhaps through compression or more efficient indexing mechanisms. Additionally, the methodology could be extended to explore varying levels of context abstraction and retrieval at coarser granularities like sentences or paragraphs, which might enhance reasoning capabilities.

In conclusion, kNN-Prompt represents a significant stride in non-parametric LLM enhancement, providing a pathway for improved task versatility and zero-shot learning capabilities while maintaining computational feasibility.