BERT-kNN: Adding a kNN Search Component to Pretrained Language Models for Better QA

Findings of the Association for Computational Linguistics: EMNLP 2020, pages 3424–3430 November 16 - 20, 2020. c⃝2020 Association for Computational Linguistics 3424 BERT-kNN: Adding a kNN Search Component to Pretrained Language Models for Better QA Nora Kassner, Hinrich Sch¨utze Center for Information and Language Processing (CIS) LMU Munich, Germany kassner@cis.lmu.de Abstract Khandelwal et al. (2020) use a k-nearest- neighbor (kNN) component to improve lan- guage model performance. We show that this idea is beneficial for open-domain ques- tion answering (QA). To improve the recall of facts encountered during training, we com- bine BERT (Devlin et al., 2019) with a tra- ditional information retrieval step (IR) and a kNN search over a large datastore of an embed- ded 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 factu- ally correct answer (e.g., “Miami”). iii) Com- pared to BERT, BERT-kNN excels for rare facts. iv) BERT-kNN can easily handle facts not covered by BERT’s training set, e.g., re- cent events. 1 Introduction Pretrained language models (PLMs) like BERT (Devlin et al., 2019), GPT-2 (Radford et al., 2019) and RoBERTa (Liu et al., 2019) have emerged as universal tools that not only capture a diverse range of linguistic, but also (as recent evidence seems to suggest) factual knowledge. Petroni et al. (2019) introduced LAMA (LAn- guage Model Analysis) to test BERT’s perfor- mance on open-domain QA and therefore inves- tigate PLMs’ capacity to recall factual knowledge without the use of finetuning. Since the PLM train- ing objective is to predict masked tokens, ques- tion answering tasks can be reformulated as cloze questions; e.g., “Who wrote ‘Ulysses’?” is refor- mulated as “[MASK] wrote ‘Ulysses’.” In this setup, Petroni et al. (2019) show that, on QA, PLMs outperform baselines trained on automatically ex- tracted knowledge bases (KBs). Figure 1: BERT-kNN interpolates BERT’s prediction for question q with a kNN-search. The kNN search runs in BERT’s embedding space, comparing the em- bedding of q with the embeddings of a retrieved sub- set of a large text collection: Pairs of a word w in the text collection and the BERT embedding of w’s con- text (BERT(s)) are stored in a key-value datastore. An IR step is used to define a relevant subset of the full datastore (yellow). BERT(q) (red) is BERT’s embedding of the question. The kNN search runs be- tween BERT(q) and BERT(s) and the correspond- ing distance d and word w is returned (orange). Fi- nally, BERT’s predictions (blue) are interpolated with this kNN search result. 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 (Poerner et al., 2019). To address PLMs’ poor performance on facts and choosing BERT as our PLM, we introduce …

Content truncated to 500 words