Distilling Knowledge from Reader to Retriever for Question Answering

1 INTRODUCTION . Information retrieval is an important component for many natural language processing tasks , such as question answering ( Voorhees et al. , 1999 ) or fact checking ( Thorne et al. , 2018 ) . For example , many real world question answering systems start by retrieving a set of support documents from a large source of knowledge such as Wikipedia . Then , a finer-grained model processes these documents to extract the answer . Traditionally , information retrieval systems were based on hand-crafted sparse representations of text documents , such as TF-IDF or BM25 ( Jones , 1972 ; Robertson et al. , 1995 ) . Recently , methods based on dense vectors and machine learning have shown promising results ( Karpukhin et al. , 2020 ; Khattab et al. , 2020 ) . Deep neural networks based on pre-training , such as BERT ( Devlin et al. , 2019 ) , have been used to encode documents into fixed-size representations . These representations are then queried using approximate nearest neighbors ( Johnson et al. , 2019 ) . These techniques have lead to improved performance on various question answering tasks . A challenge of applying machine learning to information retrieval is to obtain training data for the retriever . To train such models , one needs pairs of queries and the corresponding list of documents that contains the information corresponding to the queries . Unfortunately , hand-labeling data to that end is time consuming , and many datasets and applications lack such annotations . An alternative approach is to resort to heuristics , or weakly supervised learning , for example by considering that all documents containing the answer are positive examples . However , these approaches suffer from the following limitations . First , frequent answers or entities might lead to false positive examples . As an example , consider the question “ where was Ada Lovelace born ? ” . The sentence “ Ada Lovelace died in 1852 in London ” would be considered as a positive example , because it contains the answer “ London ” . A second limitation is that for some tasks , such as fact checking or long form question answering , such heuristics might not be applicable directly . In this paper , we propose a procedure to learn retriever systems without strong supervision in the form of pairs of queries and documents . Following previous work ( Chen et al. , 2017 ) , our approach uses two models : the first one retrieves documents from a large source of knowledge ( the retriever ) , the second one processes the support documents to solve the task ( the reader ) . Our method is inspired by knowledge distillation ( Hinton et al. , 2015 ) , and uses the reader model to obtain synthetic labels to train the retriever model . More precisely , we use a sequence-to-sequence model as the reader , and use the attention activations over the input documents as synthetic labels to train the retriever . Said otherwise , we assume that attention activations are a good proxy for the relevance of documents . We then train the retriever to reproduce the ranking of documents corresponding to that metric . We make the following contributions : • First , we show that attention scores from a sequence-to-sequence reader model are a good measure of document relevance ( Sec . 3.2 ) ; • Second , inspired by knowledge distillation , we propose to iteratively train the retriever from these activations , and compare different loss functions ( Sec . 3.4 ) ; • Finally , we evaluate our method on three question-answering benchmarks , obtaining stateof-the-art results ( Sec . 4 ) . Our code is available at : github.com/facebookresearch/FiD . 2 RELATED WORK . We briefly review information retrieval based on machine learning . We refer the reader to Manning et al . ( 2008 ) and Mitra et al . ( 2018 ) for a more exhaustive introduction to the subject . Vector space models . In traditional information retrieval systems , documents and queries are represented as sparse vectors , each dimension corresponding to a different term . Different schemes have been considered to weigh the different term , the most well known being based on inverse document frequency , or term specificity ( Jones , 1972 ) . This technique was later extended , leading to the BM25 weighting scheme which is still widely used today ( Robertson et al. , 1995 ) . A limitation of sparse representations is that the terms of the query need to match the terms of the returned documents . To overcome this , Deerwester et al . ( 1990 ) proposed to use latent semantic analysis for indexing , leading to low-dimension dense representations of documents . Neural information retrieval . Following the success of deep learning for other natural processing tasks , neural networks were applied to the task of information retrieval . Huang et al . ( 2013 ) proposed a deep bag-of-words model , where queries and documents were embedded independently , a technique known as bi-encoder . Documents were then ranked by using the cosine similarity with the query , and the model was trained on clickthrough data from a search engine . This technique was later extended by using convolutional neural networks ( Shen et al. , 2014 ) and recurrent neural networks ( Palangi et al. , 2016 ) . A limitation of independently embedding documents and query is that it does not capture fine-grained interactions between the query and documents . This lead Nogueira & Cho ( 2019 ) and Yang et al . ( 2019 ) to use a BERT model to jointly embed documents and query , a technique known as cross-encoder . End-to-end retrieval . Most of the methods described in the previous paragraph were used to rerank a small number of documents , usually returned by a traditional IR systems . In the context of ad-hoc document retrieval , Gillick et al . ( 2018 ) showed that bi-encoder models could be competitive with traditional IR systems . For open domain question-answering , Karpukhin et al . ( 2020 ) introduced dense passage retrieval ( DPR ) , which uses dense embeddings and nearest neighbors search . More precisely , question and passage embeddings are obtained using a BERT-based biencoder model , which is trained on a small dataset of question and passage pairs . Then , the full knowledge source ( Wikipedia ) is encoded using this model , and passages are queried by computing the k-nearest neighbors of the embedding of the question . Jointly embedding the query and documents makes the application of cross-encoder models intractable to large database . To address this limitation , Humeau et al . ( 2019 ) introduced the poly-encoder architecture , in which each documents is represented by multiple vectors instead of one . Similarly , Khattab et al . ( 2020 ) proposed a scoring function where each term of the query and documents is represented by a single vector . To make the method tractable , their system retrieves documents with an approximate score , which are then re-ranked with the exact one . Finally , Luan et al . ( 2020 ) conducts a theoretical and empirical study of sparse , dense and cross-attention information retrieval systems . Unsupervised learning . Closest to our work , there is growing body of work trying to learn information retrieval systems from unsupervised data . Lee et al . ( 2019 ) introduced the inverse cloze task for pre-training retrievers , which can then be fine-tuned end-to-end on question-answering tasks . This pre-training scheme was later evaluated for ad-hoc document retrieval by Chang et al . ( 2020 ) . Guu et al . ( 2020 ) proposed to augment language model pre-training with a retriever module , which is trained using the masked language modeling objective . Similarly , Lewis et al . ( 2020a ) introduced a sequence-to-sequence model that is pre-trained by generating a target text , after retrieving a set of related texts . Lewis et al . ( 2020b ) further train the retriever obtained in Karpukhin et al . ( 2020 ) by backpropagating to the retriever the error between the generated output and the gold answer . Simultaneously to our work , Yang & Seo ( 2020 ) proposes to train a retriever with knowledge distillation . The main difference with our method is the nature of the synthetic labels that are used to train the retriever . Yang & Seo ( 2020 ) uses the DPR reader , which includes a classifier that predicts which passage contains the answer , and can be seen as a cross-encoder reranker . This technique thus performs the distillation of a cross-encoder retriever to a bi-encoder retriever . In contrast , our method uses the internal attention scores of the reader , which does not require additional supervision besides pairs of question and answer . 3 METHODOLOGY . Our system is composed of two modules , the retriever and the reader , following the standard pipeline for open-domain question answering . Given an input question these modules are used in a two-step process to generate an answer . First the retriever selects support passages in a large knowledge source . Then these passages are processed by the reader , along with the question , to generate an answer . For the reader module we use the Fusion-in-Decoder model ( Izacard & Grave , 2020 ) , which achieves state-of-the-art performance when combined with BM25 or DPR ( Karpukhin et al. , 2020 ) . It is based on a sequence-to-sequence architecture , and is initialized from pre-trained models such as T5 or BART ( Raffel et al. , 2019 ; Lewis et al. , 2019 ) . The focus of this work is to train the retriever without strong supervision or weakly supervised learning based on heuristics . For this we propose to train the retriever by learning to approximate the attention score of the reader . The training scheme outlined here can be seen as a student-teacher pipeline , where the teacher , the reader module , produces targets which are used to train a student network , the reader . By doing so , we hope to leverage the signal extracted from the question-answer pairs by the reader . Since the goal of the retriever is to retrieve the most relevant passages , by training the retriever to estimate the reader attention scores , we implicitly make the assumption that these scores are a good proxy for the usefulness of a passage to answer the question . In this section we will first describe the Fusion-in-Decoder architecture , before elaborating on the signal which is used to train the retriever , the design of the retriever , and how this module is trained . 3.1 CROSS-ATTENTION MECHANISM . First , let us briefly review the Fusion-in-Decoder model ( FiD , Izacard & Grave , 2020 ) . The underlying architecture is a sequence-to-sequence model , composed of an encoder and a decoder . The encoder independently processes np different text inputs ( sk ) 1≤k≤np . In the case of open-domain question answering based on Wikipedia , each input sk is the concatenation of the question q and a support passage , with special tokens question : , title : and context : added before the question , the title of the Wikipedia article and the text of each passage . The output representations of the encoder are then concatenated to form a global representation X of dimension ( ∑ k ` k ) × d , where ` k is the length of the k-th segment and d is the dimension of the embeddings and hidden representations of the model . Then , the decoder processes this representation as a regular autoregressive model , alternating self-attention , cross-attention and feed-forward modules . Only the cross-attention module explicitly takes as input the global output representation X of the encoder . If H ∈ Rd denotes the output of the previous self-attention layer of the decoder , the crossattention operation consists in the following operations . First , queries Q , keys K and values V are computed by applying linear transformations : Q = WQH , K = WKX , V = WVX . Then a similarity score between the query at position i , Qi , and the key at position j , Kj , is obtained by computing the dot-product between these two elements , and normalized over the dimension : αi , j = Q T i Kj , α̃i , j = exp ( αi , j ) ∑ m exp ( αi , m ) . A new representation is obtained as a sum of the values , weighted by the attention probabilities , before going through a final linear transformation Wo : Oi = WO ∑ j α̃i , jVi , j The operations described above are performed in parallel with different linear transformations in the case of multi-head attention . Finally a normalization layer is applied , and this pipeline is wrapped by a skip connection . See Vaswani et al . ( 2017 ) for more details on the structure of Transformers .

Summary: Many Information Retrieval systems rely on two components: a retriever that identifies a small set of "support" documents from a large corpus, followed by a reader that re-scores these support documents more finely. For retrievers, metrics like BM25 were once common but they are increasingly replaced by machine learned components. However, most datasets do not provide direct supervision information for the retriever.