Empowering Language Models with Knowledge Graph Reasoning for Open-Domain Question Answering (2022)

Abstract: Answering open-domain questions requires world knowledge about in-context entities. As pre-trained Language Models (LMs) lack the power to store all required knowledge, external knowledge sources, such as knowledge graphs, are often used to augment LMs. In this work, we propose knOwledge REasOning empowered Language Model(OREO-LM), which consists of a novel Knowledge Interaction Layer that can be flexibly plugged into existing Transformer-based LMs to interact with a differentiable Knowledge Graph Reasoning module collaboratively. In this way, LM guides KG to walk towards the desired answer, while the retrieved knowledge improves LM.By adopting OREO-LM to RoBERTa and T5, we show significant performance gain, achieving state-of-art results in the Closed-Book setting. The performance enhancement is mainly from the KG reasoning’s capacity to infer missing relational facts. In addition, OREO-LM provides reasoning paths as rationales to interpret the model’s decision.

## introduction open-domain question answering (odqa), one of the most knowledge-intensive nlp tasks, requires qa models to infer out-of-context knowledge to the given single question. following the pioneering work by @xcite , odqa systems often assume to access an external text corpus (e.g., wikipedia) as an external knowledge source. due to the large scale of such textual knowledge sources (e.g., 20gb for wikipedia), it cannot be encoded in the model parameters. therefore, most works retrieve relevant passages as knowledge and thus named open-book models @xcite , with an analogy of referring to textbooks during an exam. another line of closedbook models @xcite assume knowledge could be stored implicitly in parameters of language models (lm, e.g. bert @xcite , 2019) and t5 @xcite ). these lms directly generate answers without retrieving from an external corpus and thus benefit from faster inference speed and simpler training. however, current lms still miss a large portion of factual knowledge @xcite , and are not competitive with open-book models. to improve the knowledge coverage of lm, one natural choice is to leverage knowledge stored in knowledge graph (kg, e.g. freebase @xcite and wikidata @xcite ), which explicitly encodes world knowledge via relational triplets between entities. there are several good properties of kg: 1) a kg triplet is a more abstract and compressed representation of knowledge than text, and thus kg could be stored in memory and directly enhance lm without using an additional retrieval model; 2) the structural nature of kg could support logical reasoning @xcite and infer missing knowledge through high-order paths @xcite . taking the question "what cheese is used to make the desert cannoli?" as an example, even if this relational fact is missing in kg, we could still leverage high-order relationships, e.g., both ricotta cheese and cannoli are specialties in italy, to infer the answer "ricotta cheese." in light of the good properties of kg, there are several efforts to build knowledge base question answering (kbqa) systems. as is illustrated in figure 1 (a), most kbqa models use lm as a parser to map textual questions into a structured form (e.g., sql query or subgraph), and then based on kg, the queries could be executed by symbolic reasoning @xcite or neural reasoning (e.g. graph neural networks) @xcite to get the answer. another recent line of research @xcite tries to encode the knowledge graph as the memory into lm parameters. however, for most methods discussed above, lm is not interacting with kg to correctly understand the question, and the answer is usually restricted to a node or edge in kg. in this paper, we propose knowledge reasoning empowered language model (oreolm), a model architecture that can be applied to transformer-based lms to improve closed-book odqa. as is illustrated in figure 1(b), the key component is the knowledge interaction layers (kil) inserted amid lm layers, which is like cream filling within two waffles, leading to our model's name oreo. kil interacts with a kg reasoning module, in which we maintain different reasoning paths for each entity in the question. we formulate the retrieval and reasoning process as a contextualized random walk over the kg, starting from the in-context entities. each kil is responsible for one reasoning step. it first predicts a relation distribution for every in-context entity, and then the kg reasoning module traverses the graph following the predicted relation distribution. the reasoning result in each step is summarized as a weighted averaged embedding over the retrieved entities from the traversal. by stacking t layers of kil, oreolm can retrieve entities that are t -hop away from in-context entities and help lm to answer open questions that require out-of-context knowledge or multi-hop reasoning. the whole procedure is fully differentiable, and thus oreolm learns and infers in an end-toend manner. we further introduce how to pre-train oreolm over unlabelled wikipedia corpus. in addition to the salient entity span masking objective, we introduce two self-supervised objectives to guide oreolm to learn better entity and relation representations and how to reason over them. we test oreolm with roberta and t5 as our base lms. by evaluating on several single-hop odqa datasets in closed-book setting, we show that oreolm outperforms existing baselines with fewer model parameters. specifically, oreolm helps more for questions with missing relations in kg, and questions that require multi-hop reasoning. we further show that oreolm can serve as a backbone for open-book setting and achieves comparable performance compared with the stateof-the-art qa systems with dedicated design. in addition, oreolm has better interpretability as it can generate reasoning paths for the answered question and summarize general relational rules to infer missing relations. this key contributions are as follows: ## methodology preliminary we denote a knowledge graph @xmath0, r, @xmath1, where each @xmath2. a @xmath3, 1} |e|×|e| is a sparse adjacency matrix indicating whether relation r holds between a pair of entities. the task of knowledge graph reasoning aims at answering a factoid query (s, r, ?), i.e., which target entity has relation r with the source entity s. if kg is complete, we could simply get answers by checking the adjacency matrix, i.e., {∀t : a r [s, t] = 1}. for incomplete kg where many relational facts are missing, path-based reasoning approaches @xcite @xcite have been proposed to answer the one-hop query via finding multi-hop paths. for example, to answer the query (s, mother, ?), a path s father ---→ j wife --→ t could reach the target answer t. in this paper we try to integrate symbolic kg reasoning into neural lms and help it deal with odqa problems. overview of oreolm we illustrate the overall architecture of oreolm in figure 2 . all the light blue blocks are our added components to support kg reasoning, while the dark blue transformer layers are knowledge-injected lm. the key component of oreolm for conducting kg reasoning is the knowledge interaction layers (kil), which are added amid lm layers to enable deeper interaction with the kg. given a question @xmath4, qa model needs to extract knowledge about all n in-context entity mentions @xmath5, e.g., the history of "germany" at the time when "bauhaus" is founded, to get the answer @xmath6. such open-domain q&a can be abstracted as p (a|q, m ). starting from each mentioned entity m i , we desire the model to learn to walk over the graph to retrieve relevant knowledge and form a t -length reasoning path for answering this question, where t is a hyper-parameter denote the longest reasoning path required to answer the questions. we define each reasoning path starting from the entity mention m i as a chain of entities (states) random variables ρ @xmath7, where each mentioned entity is the initial state, i.e., e 0 @xmath8. the union of all paths for this question is defined as ϱ = {ρ i }, which contains the reasoning paths from each mentioned entity to answer the question. oreolm factorizes p a|q, m by incorporat-ing possible paths ϱ as a latent variable, yielding: p a|q, @xmath9, {mi} n @xmath10, m, ϱ = ϱ n @xmath11, mi • p a|q, {mi, ρi} n @xmath12, @xmath13.1) p a|q, {e 0:t i } n @xmath14.2) we assume (1) reasoning paths starting from different entities are generated independently; and (2) reasoning paths can be generated autoregressively. in this way, the qa problem can be decomposed into two entangled steps: 1) kg reasoning, which autoregressively walks through the graph to get a path ρ i starting from each entity mention m i ; and 2) knowledge-injected lm, which benefits from the reasoning paths to obtain the out-context knowledge for answer prediction. the relational path ρ i in kg reasoning requires the selection of next entity e t i at each step t. we further decompose it into two steps: 1.a) relation prediction, in which lm is involved to predict the next-hop relation based on the current state and context; and 1.b) the non-parametric state transition, which is to predict the next-hop entity based on the kg and the predicted relation. formally: p e t i |q, @xmath15.1) = r p rel r t i |q, @xmath16.1.1) • p walk e t i |r t i , @xmath17.1.2) we keep track of the entity distribution at each step t via the probability vector 1 π (t) @xmath18, with π (t) i [e] being the probability of staying at entity e, i.e., p e t @xmath19, @xmath20. we highlight the three procedures in red dotted box in figure 2 . we take the first reasoning step starting from entity mention "bauhaus" as an example. in the first red box within kil, we predict which relation action should be taken for entity "bauhaus", and send the prediction (e.g. "founded") to kg. in the second red box, kg reweights the graph and conducts contextualized random walk to update entity distribution, where "walter" has the highest probability. finally, weighted by the entity distribution, an aggregated entity embedding is sent back to kil and added into a placeholder token as the knowledge, so the later lm layer knows to focus on the retrieved "walter". we introduce these steps in the following. input initially, we first identify all n entity mentions {m i } n @xmath21.. for each mention m i we add three special tokens as the interface for knowledge interaction layers (kil) to send instruction and receive knowledge: we add a [s-ent] token before, and [rel], [t-ent] tokens after each entity mention m i . kil can be flexibly inserted into arbitrary lm intermediate layer. by default, we just insert each kil every n transformer-based lm layers, thus the input to the t-th kil are contextualized embeddings of each token k as lm (t) k , including added special tokens. ## experiments the proposed kil layers can be pugged into most transformer-based language models without hurting its original structure. in this paper, we experiment with both encoder-based lm, i.e. robertabase (@xmath22, @xmath23, and encoder-decoder lm, i.e. t5-base (@xmath24, @xmath25, @xmath26. for all lms, add 1 kil layer or 2 kil layers to the encoder layers. the statistics of kg are shown in ## related work open-domain question answering (odqa) gives qa model a single question without any context and asks the model to infer out-of-context knowledge. following the pioneering work by @xcite , most odqa systems assume the model can access an external text corpus (e.g. wikipedia). due to the large scale of web corpus (20gb for wikipedia), it could not be simply encoded in the qa model parameters, and thus most works propose a retrieval-reader pipeline, by firstly index the whole corpus and use a retriever model to identify which passage is relevant to the question; then the retrieved text passage concatenate with question is re-encoded by a seperate reader model (e.g., lm) to predict answer. as the knowledge is outside of model parameter, roberts et al. (2020) defines these methods as open-book, with an analogy to referring textbooks during exam. closed-book qa models (mostly a single lm) try to answer open questions without accessing external knowledge. this setting is much harder as it requires lm to memorize all pertinent knowledge in its parameters, and even recent lms with much larger model parameters is still not competitive to state-of-the-art open-book models. knowledge-augmented language models explicitly incorporate external knowledge (e.g. knowledge graph) into lm @xcite . overall, these approaches can be grouped into two categories: the first one is to explicitly inject knowledge representation into language model pre-training, where the representations are precomputed from external sources @xcite @xcite . for example, ernie @xcite encodes the pre-trained transe @xcite embeddings as input. the second one is to implicitly model knowledge information into language model by performing knowledge-related tasks, such as entity category prediction @xcite and graph-text alignment @xcite . for example, jaket @xcite jointly pre-trained both the kg representation and language representation by adding entity category and relation type prediction self-supervised tasks. there also exists several qa works using kg to help odqa. for example, @xcite and @xcite expand the entity graph following wikipedia hyperlinks or triplets in knowledge base. @xcite extract entities from current context via entity-linking and turn them into a cognitive graph, and a graph neural network is applied on top of it to extract answer. @xcite and @xcite construct an entity-mention bipartite graph and then model the qa reasoning as graph traversal by filtering only the contexts that are relevant to the question. lin et al. (2019), feng et al. (2020) and yasunaga et al. (2021) parse the question into a sub-graph of knowledge base, and apply graph neural networks as reasoner for extracting one of the entities as the answer. to encode knowledge (significantly smaller than the web corpus) as memory into lm parameter, a line of works try compressed knowledge including qa pairs @xcite @xcite , entity embedding (févry et al., 2020) and reasoning cases @xcite . there's also several works utilizing knowledge graph (kg) to augment lm. film @xcite turns kg triplets into memory. given a question, lm retrieves most relevant triplet as answer. greaselm @xcite propose to interact lm with kg via a interaction node. ## conclusion we presented oreolm, a novel model that incorporates symbolic kg reasoning with existing lms. we showed that oreolm can bring significant performance gain to open-domain qa benchmarks, both for closed-book and open-book settings, as well as encoder-only and encoder-decoder models. additionally, oreolm produces reasoning paths that helps interpret the model prediction. in future, we'd like to improve oreolm by training to conduct more reasoning steps, supporting locial reasoning, and apply oreolm to a broader range of knowledge-intensive nlp tasks.