Nearest Neighbor Zero-Shot Inference

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 3254 - 3265 December 7-11, 2022 ©2022 Association for Computational Linguistics kNN-Prompt: Nearest Neighbor Zero-Shot Inference Weijia Shi Julian Michael Suchin Gururangan Luke Zettlemoyer Paul G. Allen School of Computer Science & Engineering, University of Washington, Seattle, WA {swj0419, julianjm, sg01, lsz}@cs.washington.edu Abstract Retrieval-augmented language models (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 ex- tensively 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 sig- nificant 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 aug- mentation 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. 1 Introduction Retrieval-augmented language models (LMs) have access to a non-parametric memory, allowing them to directly access a large external text collection during inference. Previous work has shown that these models substantially outperform their non– retrieval-based counterparts on language modeling tasks (Khandelwal et al., 2020; He et al., 2021; Borgeaud et al., 2021), but it is an open question whether they also achieve similar gains in few- shot and zero-shot end task evaluations (Radford et al., 2019; Brown et al., 2020a). In this paper, we demonstrate that, with some extensions to improve coverage of the verbalizer tokens, the performance gains of retrieval-augmented LMs generalize well to a wide range of downstream tasks. Figure 1: kNN-Prompt incorporates information from a large, heterogeneous corpus (unlabeled texts from differ- ent domains) to facilitate few- and zero-shot inference. The datastore contains key-value pairs where the key is an encoding of a leftward context and the value is the next token following the context. Our fuzzy ver- balizer expands "terrible" to include "silly" and "great" to include "excellent". Because the encoded corpus is unlabeled plain text, some datastore entries contain next tokens not in the verbalizer tokens (e.g., "cinema"). We study the k-nearest neighbors language model (Khandelwal et al., 2020, kNN-LM), which interpolates the LM softmax distribution with a nearest-neighbor distribution. The nearest neigh- bours are computed based on the distance in LM output embeddings and can be drawn from any text corpus, in our case, a heterogeneous corpus that contains unlabeled data from different domains. We are the first to study the zero-shot application of kNN-LM to end tasks, and we find that applying the technique naïvely produces only marginal im- provements (Section 4). …

Content truncated to 500 words