Differentiable Reasoning over a Virtual Knowledge Base

1 INTRODUCTION . Large knowledge bases ( KBs ) , such as Freebase and WikiData , organize information around entities , which makes it easy to reason over their contents . For example , given a query like “ When was the Grateful Dead ’ s lead singer born ? ” , one can identify the entity Grateful Dead and the path of relations LeadSinger , BirthDate to efficiently extract the answer—provided that this information is present in the KB . Unfortunately , KBs are often incomplete ( Min et al. , 2013 ) . While relation extraction methods can be used to populate KBs , this process is inherently error-prone , expensive and slow . Advances in open-domain QA ( Moldovan et al. , 2002 ; Yang et al. , 2019 ) suggest an alternative— instead of performing relation extraction , one could treat a large corpus as a virtual KB by answering queries with spans from the corpus . This ensures facts are not lost in the relation extraction process , but also poses challenges . One challenge is that it is relatively expensive to answer questions using QA models which encode each document in a query-dependent fashion ( Chen et al. , 2017 ; Devlin et al. , 2019 ) —even with modern hardware ( Strubell et al. , 2019 ; Schwartz et al. , 2019 ) . The cost of QA is especially problematic for certain complex questions , such as the example question above . If the passages stating that “ Jerry Garcia was the lead singer of the Grateful Dead ” and “ Jerry Garcia was born in 1942 ” are far apart in the corpus , it is difficult for systems that retrieve and read a single passage to find an answer—even though in this example , it might be easy to answer the question after the relations were explicitly extracted into a KB . More generally , complex questions involving sets of entities or paths of relations may require aggregating information from multiple documents , which is expensive . One step towards efficient QA is the recent work of Seo et al . ( 2018 ; 2019 ) on phrase-indexed question answering ( PIQA ) , in which spans in the text corpus are associated with question-independent ∗Part of this work was done during an internship at Google . 1Code available at http : //www.cs.cmu.edu/˜bdhingra/pages/drkit.html contextual representations and then indexed for fast retrieval . Natural language questions are then answered by converting them into vectors that are used to perform maximum inner product search ( MIPS ) against the index . This can be done efficiently using approximate algorithms ( Shrivastava & Li , 2014 ) . However , this approach can not be directly used to answer complex queries , since by construction , the information stored in the index is about the local context around a span—it can only be used for questions where the answer can be derived by reading a single passage . This paper addresses this limitation of phrase-indexed question answering . We introduce an efficient , end-to-end differentiable framework for doing complex QA over a large text corpus that has been encoded in a query-independent manner . Specifically , we consider “ multi-hop ” complex queries which can be answered by repeatedly executing a “ soft ” version of the operation below , defined over a set of entities X and a relation R : Y = X.follow ( R ) = { x′ : ∃x ∈ X s.t . R ( x , x′ ) holds } In past work soft , differentiable versions of this operation were used to answer multi-hop questions against an explicit KB ( Cohen et al. , 2019 ) . Here we propose a more powerful neural module which approximates this operation against an indexed corpus ( a virtual KB ) . In our module , the input X is a sparse-vector representing a weighted set of entities , and the relation R is a dense feature vector , e.g . a vector derived from a neural network over a natural language query . X and R are used to construct a MIPS query used for retrieving the top-K spans from the index . The output Y is another sparse-vector representing the weighted set of entities , aggregated over entity mentions in the top-K spans . We discuss pretraining schemes for the index in §2.3 . For multi-hop queries , the output entities Y can be recursively passed as input to the next iteration of the same module . The weights of the entities in Y are differentiable w.r.t the MIPS queries , which allows end-to-end learning without any intermediate supervision . We discuss an implementation based on sparse-matrix-vector products , whose runtime and memory depend only on the number of spans K retrieved from the index . This is crucial for scaling up to large corpora , providing up to 15x faster inference than existing state-of-the-art multi-hop and open-domain QA systems . The system we introduce is called DrKIT ( for Differentiable Reasoning over a Knowledge base of Indexed Text ) . We test DrKIT on the MetaQA benchmark for complex question answering , and show that it improves on prior text-based systems by 5 points on 2-hop and 9 points on 3-hop questions , reducing the gap between text-based and KB-based systems by 30 % and 70 % , respectively . We also test DrKIT on a new dataset of multi-hop slot-filling over Wikipedia articles , and show that it outperforms DrQA ( Chen et al. , 2017 ) and PIQA ( Seo et al. , 2019 ) adapted to this task . Finally , we apply DrKIT to multi-hop information retrieval on the HotpotQA dataset ( Yang et al. , 2018 ) , and show that it significantly improves over a BERT-based reranking approach , while being 10x faster . 2 DIFFERENTIABLE REASONING OVER A KB OF INDEXED TEXT . We want to answer a question q using a text corpus as if it were a KB . We start with the set of entities z in the question q , and would ideally want to follow relevant outgoing relation edges in the KB to arrive at the answer . To simulate this behaviour on text , we first expand z to set of cooccurring mentions m ( say using TFIDF ) . Not all of these co-occurring mentions are relevant for the question q , so we train a neural network which filters the mentions based on a relevance score of q to m. Then we can aggregate the resulting set of mentions m to the entities they refer to , ending up with an ordered set z′ of entities which are answer candidates , very similar to traversing the KB . Furthermore , if the question requires more than one hop to answer , we can repeat the above procedure starting with z′ . This is depicted pictorially in Figure 1 . We begin by first formalizing this idea in a probabilistic framework in §2.1 . In §2.2 , we describe how the expansion of entities to mentions and the filtering of mentions can be performed efficiently , using sparse-matrix products and MIPS algorithms ( Johnson et al. , 2017 ) . Lastly we discuss a pretraining scheme for constructing the mention representations in §2.3 . Notation : We denote the given corpus as D = { d1 , d2 , . . . } , where each dk = ( d1k , . . . , dLkk ) is a sequence of tokens . We start by running an entity linker over the corpus to identify mentions of a fixed set of entities E . Each mention m is a tuple ( em , km , im , jm ) denoting that the text span dimkm , . . . , d jm km in document km mentions the entity em ∈ E , and the collection of all mentions in the corpus is denoted asM . Note that typically |M| |E| . 2.1 DIFFERENTIABLE MULTI-HOP REASONING . We assume a weakly supervised setting where during training we only know the final answer entities a ∈ E for a T -hop question . We denote the latent sequence of entities which answer each of the intermediate hops as z0 , z1 , . . . , zT ∈ E , where z0 is mentioned in the question , and zT = a . We can recursively write the probability of an intermediate answer as : Pr ( zt|q ) = ∑ zt−1∈E Pr ( zt|q , zt−1 ) Pr ( zt−1|q ) ( 1 ) Here Pr ( z0|q ) is the output of an entity linking system over the question , and Pr ( zt|q , zt−1 ) corresponds to a single-hop model which answers the t-th hop , given the entity from the previous hop zt−1 , by following the appropriate relation . Eq . 1 models reasoning over a chain of latent entities , but when answering questions over a text corpus , we must reason over entity mentions , rather than entities themselves . Hence Pr ( zt|q , zt−1 ) needs to be aggregated over all mentions of zt , which yields Pr ( zt|q ) = ∑ m∈M ∑ zt−1∈E Pr ( zt|m ) Pr ( m|q , zt−1 ) Pr ( zt−1|q ) ( 2 ) The interesting term to model in the above equation isPr ( m|q , zt−1 ) , which represents the relevance of mentionm given the question and entity zt−1 . Following the analogy of a KB , we first expand the entity zt−1 to co-occuring mentions m and use a learned scoring function to find the relevance of these mentions . Formally , let F ( m ) denote a TFIDF vector for the document containingm , G ( zt−1 ) be the TFIDF vector of the surface form of the entity from the previous hop , and st ( m , z , q ) be a learnt scoring function ( different for each hop ) . Thus , we model Pr ( m|q , zt−1 ) as Pr ( m|q , zt−1 ) ∝ 1 { G ( zt−1 ) · F ( m ) > } ︸ ︷︷ ︸ expansion to co-occurring mentions × st ( m , zt−1 , q ) ︸ ︷︷ ︸ relevance filtering ( 3 ) Another equivalent way to look at our model in Eq . 3 is that the second term retrieves mentions of the correct type requested by the question in the t-th hop , and the first term filters these based on co-occurrence with zt−1 . When dealing with a large set of mentions m , we will typically retain only the top-K relevant mentions . We will show that this joint modelling of co-occurrence and relevance is important for good performance , as was also observed by Seo et al . ( 2019 ) . The other term left in Eq . 2 is Pr ( z|m ) , which is 1 if mention m refers to the entity z else 0 , based on the entity linking system . In general , to compute Eq . 2 the mention scoring of Eq . 3 needs to be evaluated for all latent entity and mention pairs , which is prohibitively expensive . However , by restricting st to be an inner product we can implement this efficiently ( §2.2 ) . To highlight the differentiability of the proposed overall scheme , we can represent the computation in Eq . 2 as matrix operations . We pre-compute the TFIDF term for all entities and mentions into a sparse matrix , which we denote as AE→M [ e , m ] = 1 ( G ( e ) · F ( m ) > ) . Then entity expansion to co-occuring mentions can be done using a sparse-matrix by sparse-vector multiplication between AE→M and zt−1 . For the relevance scores , let TK ( st ( m , zt−1 , q ) ) denote the top-K relevant mentions encoded as a sparse vector in R|M| . Finally , the aggregation of mentions to entities can be formulated as multiplication with another sparse-matrix BM→E , which encodes coreference , i.e . mentions corresponding to the same entity . Putting all these together , using to denote elementwise product , and defining Zt = [ Pr ( zt = e1|q ) ; . . . ; Pr ( zt = e|E||q ) ] , we can observe that for large K ( i.e. , as K → |M| ) , Eq . 2 becomes equivalent to : Zt = softmax ( [ ZTt−1AE→M TK ( st ( m , zt−1 , q ) ) ] BM→E ) . ( 4 ) Note that every operation in above equation is differentiable and between sparse matrices and vectors : we will discuss efficient implementations in §2.2 . Further , the number of non-zero entries in Zt is bounded by K , since we filtered ( the element-wise product in Eq . 4 ) to top-K relevant mentions among TFIDF based expansion and since each mention can only point to a single entity in BM→E . This is important , as it prevents the number of entries in Zt from exploding across hops ( which might happen if , for instance , we added the relevance and TFIDF scores instead ) . We can view Zt−1 , Zt as weighted multisets of entities , and st ( m , z , q ) as implicitly selecting mentions which correspond to a relation R. Then Eq . 4 becomes a differentiable implementation of Zt = Zt−1.follow ( R ) , i.e . mimicking the graph traversal in a traditional KB . We thus call Eq . 4 a textual follow operation . Training and Inference . The model is trained end-to-end by optimizing the cross-entropy loss between ZT , the weighted set of entities after T hops , and the ground truth answer set A . We use a temperature coefficient λ when computing the softmax in Eq , 4 since the inner product scores of the top-K retrieved mentions are typically high values , which would otherwise result in very peaked distributions of Zt . We also found that taking a maximum over the mention set of an entity Mzt in Eq . 2 works better than taking a sum . This corresponds to optimizing only over the most confident mention of each entity , which works for corpora like Wikipedia that do not have much redundancy . A similar observation was made by Min et al . ( 2019 ) in weakly supervised settings .

Summary: The paper proposes a model that can perform multi-hop question-answering based on a textual knowledge base. Results show that the proposed model -- DrKIT -- performed close to or better than the state of the art results over MetaQA and WikiData datasets. Ablation study is offered to show that a few tricks in the model are necessary to make it work, and comparisons with baseline models such as DrQA and PIQA are presented. The paper also provides additional means to speed up the DrKIT model using the hashing trick and approximated top-k methods.