Nearest Neighbor Zero-Shot Inference (2022)

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 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 zeroshot 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.

## 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 nonretrieval-based counterparts on language modeling tasks @xcite @xcite , but it is an open question whether they also achieve similar gains in fewshot and zero-shot end task evaluations @xcite . 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 different 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 verbalizer 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 @xcite , which interpolates the lm softmax distribution with a nearest-neighbor distribution. the nearest neighbours 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 improvements (section 4). the main challenge is that the support of the knn distribution is sparse (covering at most k tokens, often less), as it only assigns probability mass to nearest neighbors. this means it often entirely misses the tokens that are used to verbalize the output label in the standard application of lms to zero-shot classification: across the 1 3254 datasets we test, an output label receives nonzero probability under the knn distribution only 44. ## method to perform zero-shot prediction on a downstream task using a pretrained language model, we recast the task as language modeling @xcite by converting each input instance into a natural language prompt (section 2.1). we then augment the pretrained model with the knearest-neighbors language modeling technique from @xcite . to better benefit from the sparse knn distribution, we introduce fuzzy verbalizers for mapping from the lm's outputs to a distribution over task-specific labels (section 2.3). finally, we decode the output from this label distribution using the domain-conditional pmi scoring method of @xcite . ## experimental results results for zero-shot prediction are in interestingly, the knn-lm alone yields a fairly small improvement over the base lm (about 0.4% on average). this suggests that the fuzzy verbalizer and pmi calibration methods may help knn-prompt better leverage the information in the knearest neighbors distribution. we examine possible sources of knn-prompt's performance gains in section 6. few-shot inference for a subset of tasks, we additionally compare to a few-shot setting where we prepend four examples uniformly sampled from the training data to the input ( ## knn-prompt for domain adaptation one of the advantages of retrieval-based lms is that they can be adapted to new domains with no further training. to test this capability, we replace our heterogeneous datastore (section 3.2) with domain-specific ones for several tasks. to build these domainspecific datastores, we select amazon reviews for cr, cc-news for hyp and imdb for mr, and encode them separately. for comparison, we consider domain-adaptive pretraining @xcite , which further trains the lm on the domain-specific corpus. we train gpt-2 large on each domain corpus for a single pass, then apply it to downstream tasks using our prompting and verbalizer setup and domain-conditional pmi scoring. as shown in ## analysis we perform several experiments to understand the contribution of each component of knn-prompt and inform our choice of hyperparameters. interpolation (section 2.2), fuzzy verbalizers (section 2.3), and pmi scoring (section 2.4). ## conclusions we present knn-prompt, a new technique to augment lms with nearest neighbor retrieval for zeroshot inference on end tasks. knn-prompt substantially improves zero-shot performance on a wide range of multiple-choice and classification tasks. with a domain-or task-relevant datastore, knn-prompt enables efficient domain adaptation with no additional training, and its benefits scale with the size of the retrieval model.