BERT -k NN : Adding a k NN Search Component to Pretrained Language Models for Better QA (2020)

Abstract: Khandelwal et al. (2020) use a k-nearest-neighbor (kNN) component to improve language model performance. We show that this idea is beneficial for open-domain question answering (QA). To improve the recall of facts encountered during training, we combine BERT (Devlin et al., 2019) with a traditional information retrieval step (IR) and a kNN search over a large datastore of an embedded text collection. Our contributions are as follows: i) BERT-kNN outperforms BERT on cloze-style QA by large margins without any further training. ii) We show that BERT often identifies the correct response category (e.g., US city), but only kNN recovers the factually correct answer (e.g.,“Miami”). iii) Compared to BERT, BERT-kNN excels for rare facts. iv) BERT-kNN can easily handle facts not covered by BERT’s training set, e.g., recent events.

## introduction pretrained language models (plms) like bert @xcite , gpt-2 @xcite and roberta @xcite have emerged as universal tools that not only capture a diverse range of linguistic, but also (as recent evidence seems to suggest) factual knowledge. @xcite introduced lama (language model analysis) to test bert's performance on open-domain qa and therefore investigate plms' capacity to recall factual knowledge without the use of finetuning. since the plm training objective is to predict masked tokens, question answering tasks can be reformulated as cloze questions; e.g., "who wrote 'ulysses'?" is reformulated as " @xcite wrote 'ulysses'." in this setup, @xcite show that, on qa, plms outperform baselines trained on automatically extracted knowledge bases (kbs). still, given that plms have seen more text than humans read in a lifetime, their performance on open-domain qa seems poor. also, many lama facts that plms do get right are not "recalled" from training, but are guesses instead @xcite . to address plms' poor performance on facts and choosing bert as our plm, we introduce bert-knn. bert-knn combines bert's predictions with a knn search. the knn search runs in bert's embedding space, comparing the embedding of the question with the embeddings of a retrieved subset of a large text collection. the text collection can be bert's training set or any other suitable text corpus. due qa by large margins. a schematic depiction of the model is shown in figure 1 . specifically, we use bert to embed each token's masked context s in the text collection (bert (s)). each pair of context embedding and token is stored as a key-value pair in a datastore. testing for a cloze question q, the embedding of q (bert (q)) serves as query to find the k contexttarget pairs in the subset of the datastore that are closest. the final prediction is an interpolation of the knn search and the plm predictions. we find that the knn search over the full datstore alone does not obtain good results. therefore, we first query a separate information retrieval (ir) index with the original question q and only search over the most relevant subset of the full datastore when finding the k-nearest-neighbors of bert (q) in embedding space. we find that the plm often correctly predicts the answer category and therefore the correct answer is often among the top k-nearest-neighbors. a typical example is "albert einstein was born in [mask]": the plm knows that a city is likely to follow and maybe even that it is a german city, but it fails to pick the correct city. on the other hand, the topranked answer in the knn search is "ulm" and so the correct filler for the mask can be identified. bert-knn sets a new state-of-the-art on the lama cloze-style qa dataset without any further training. even though bert-knn is based on bert-base, it also outperforms bert-large. the performance gap between bert and bert-knn is most pronounced on hard-to-guess facts. our method can also make recent events available to bert without any need of retraining: we can simply add embedded text collections covering recent events to bert-knn's datastore. the source code of our experiments is available under: https://github.com/norakassner/ bert-knn . ## data the lama dataset is a cloze-style qa dataset that allows to query plms for facts in a way analogous to kb queries. a cloze question is generated using a subject-relation-object triple from a kb and a templatic statement for the relation that contains variables x and y for subject and object; e.g, "x was born in y". the subject is substituted for x and [mask] for y. in all lama triples, y is a single-token answer. lama covers different sources: the googlere foot_0 set covers the relations "place of birth", "date of birth" and "place of death". trex @xcite consists of a subset of wikidata triples covering 41 relations. conceptnet @xcite combines 16 commonsense relations among words and phrases. the underlying open mind common sense corpus provides matching statements to query the language model. squad @xcite is a standard question answering dataset. lama contains a subset of 305 context-insensitive questions. unlike kb queries, squad uses manually reformulated cloze-style questions which are not based on a template. we use squad and an additional 305 concept-net queries for hyperparamter search. @xcite introduce lama-uhn, a subset of lama's trex and googlere questions from which easy-to-guess facts have been removed. to test bert-knn's performance on unseen facts, we collect wikidata triples containing trex relations from wikipedia pages created january-may 2020 and add them to the datastore. ## method bert-knn combines bert with a knn search component. our method is generally applicable to plms. here, we use bert-base-uncased @xcite . bert is pretrained on the bookcorpus @xcite and the english wikipedia. datastore. our text collection c is the 2016-12-21 english wikipedia. 2 for each single-token word occurrence w in a sentence in c, we com- information retrieval. we find that just using the datastore d does not give good results (see result section). we therefore use @xcite 's ir system to first select a small subset of d using a keyword search. the ir index contains all wikipedia articles. an article is represented as a bag of words and word bigrams. we find the top 3 relevant wikipedia articles using tf-idf search. for kb queries, we use the subject to query the ir index. if the subject has its dedicated wikipedia page, we simply use this. for non-knowledge base queries, we use the cloze-style question q ([mask] is removed). inference. during testing, we first run the ir search to identify the subset d of d that corresponds to the relevant wikipedia articles. for the knn search, q is embedded in the same way as the context representations c in d: we set bert (q) to the embedding computed by bert for [mask]. we then retrieve the @xmath0. we convert the distances (euclidean) between bert (q) and the knns to a probability distribution using softmax. since a word w can occur several times in knn, we compute its final output probability as the sum over all occurrences. in the final step, we interpolate knn's (weight 0.3) and bert's original predictions (weight 0.7). we optimize hyperparameters on dev. see supplementary for details. evaluation. following @xcite we report mean precision at rank r (p@r). p@r is 1 if the top r predictions contain the correct answer, otherwise it returns 0. to compute mean precision, we first average within each relation and then across relations. ## results and discussion it also outperforms entity-enhanced versions of bert (see related work) -ernie @xcite , knowbert @xcite and e-bert @xcite -on lama. bert-knn answers facts unseen during pretraining better than bert, see ## related work plms are top performers for many tasks, including qa @xcite @xcite . @xcite introduced the lama qa task to probe plms' knowledge of facts typically modeled by kbs. the basic idea of bert-knn is similar to khandelwal et al. (2020)'s interpolation of a plm and knn for language modeling. in contrast, we address qa. we introduce an ir step into the model that is essential for good performance. also, our context representations differ as we use embeddings of the masked token. @xcite and @xcite , inter alia, also make use of memory to store hidden states. they focus on recent history, making it easier to copy rare vocabulary items. drqa @xcite is an open-domain qa model that combines an ir step with a neural reading comprehension model. we use the same ir module, but our model differs significantly. drqa does not predict masked tokens, but extracts answers from text. it does not use plms nor a knn module. most importantly, bert-knn is fully unsupervised and does not require any extra training. some work on knowledge in plms focuses on injecting knowledge into bert's encoder. ernie @xcite and knowbert @xcite are entity-enhanced versions of bert. they introduce additional encoder layers that are integrated into bert's original encoder by expensive additional pretraining. @xcite injects factual entity knowledge into bert's embeddings without pretraining by aligning wikipedia2vec entity vectors @xcite ) with bert's wordpiece vocabulary. this approach is also limited to labeled entities. our approach on the other hand is not limited to labeled entities nor does it require any pretraining. our approach is conceptually different from entity-enhanced versions of bert and could potentially be combined with them for even better performance. also, these models address language modeling, not qa. the combination of plms with an ir step/knn search has attracted a lot of recent research interest. the following paragraph lists concurrent work: petroni et al. (2020) also combine bert with an ir step to improve cloze-style qa. they do not use a knn search nor an interpolation step but feed the retrieved contexts into bert's encoder. @xcite augment plms with a latent knowledge retriever. in contrast to our work they continue the pretraining stage. they jointly optimize the masked language modeling objective and backpropagate through the retrieval step. @xcite ; izacard and grave (2020) leverage retrieved contexts for better qa using finetuned generative models. they differ in that the latter fuse evidence of multiple contexts in the decoder. @xcite integrate retrieved contexts into pmls for better reading comprehension. ## conclusion this work introduced bert-knn, an interpolation of bert predictions with a knn search for unsupervised cloze-style qa. bert-knn sets a stateof-the-art on lama without any further training. bert-knn can be easily enhanced with knowledge about new events that are not covered in the training text used for pretraining bert. in future work, we want to exploit the utility of the knn component for explainability: knn predictions are based on retrieved contexts, which can be shown to users to justify an answer.