GreaseLM: Graph REASoning Enhanced Language Models

1 INTRODUCTION . Question answering is a challenging task that requires complex reasoning over both explicit constraints described in the textual context of the question , as well as unstated , relevant knowledge about the world ( i.e. , knowledge about the domain of interest ) . Recently , large pretrained language models fine-tuned on QA datasets have become the dominant paradigm in NLP for question answering tasks ( Khashabi et al. , 2020 ) . After pretraining on an extreme-scale collection of general text corpora , these language models learn to implicitly encode broad knowledge about the world , which they are able to leverage when fine-tuned on a domain-specific downstream QA task . However , despite the strong performance of this two-stage learning procedure on common benchmarks , these models struggle when given examples that are distributionally different from examples seen during fine-tuning ( McCoy et al. , 2019 ) . Their learned behavior often relies on simple ( at times spurious ) patterns to offer shortcuts to an answer , rather than robust , structured reasoning that effectively fuses the explicit information provided by the context and implicit external knowledge ( Marcus , 2018 ) . On the other hand , massive knowledge graphs ( KG ) , such as Freebase ( Bollacker et al. , 2008 ) , Wikidata ( Vrandečić & Krötzsch , 2014 ) , ConceptNet ( Speer et al. , 2017 ) , and Yago ( Suchanek et al. , 2007 ) capture such external knowledge explicitly using triplets that capture relationships between entities . Previous research has demonstrated the significant role KGs can play in structured reasoning and query answering ( Ren et al. , 2020 ; 2021 ; Ren & Leskovec , 2020 ) . However , extending these reasoning advantages to general QA ( where questions and answers are expressed in natural language and not easily mapped to strict logical queries ) requires finding the right integration of knowledge from the KG with the information and constraints provided by the QA example . Prior 1All code , data and pretrained models are available at https : //github.com/snap-stanford/ GreaseLM . methods propose various ways to leverage both modalities ( i.e. , expressive large language models and structured KGs ) for improved reasoning ( Mihaylov & Frank , 2018 ; Lin et al. , 2019 ; Feng et al. , 2020 ) . However , these methods typically fuse the two modalities in a shallow and non-interactive manner , encoding both separately and fusing them at the output for a prediction , or using one to augment the input of the other . Consequently , previous methods demonstrate restricted capacity to exchange useful information between the two modalities . It remains an open question how to effectively fuse the KG and LM representations in a truly unified manner , where the two representations can interact in a non-shallow way to simulate structured , situational reasoning . In this work , we present GREASELM , a new model that enables fusion and exchange of information from both the LM and KG in multiple layers of its architecture ( see Figure 1 ) . Our proposed GREASELM consists of an LM that takes as input the natural language context , as well as a graph neural network ( GNN ) that reasons over the KG . After each layer of the LM and GNN , we design an interactive scheme to bidirectionally transfer the information from each modality to the other through specially initialized interaction representations ( i.e. , interaction token for the LM ; interaction node for the GNN ) . In such a way , all the tokens in the language context receive information from the KG entities through the interaction token and the KG entities indirectly interact with the tokens through the interaction node . By such a deep integration across all layers , GREASELM enables joint reasoning over both the language context and the KG entities under a unified framework agnostic to the specific language model or graph neural network , so that both modalities can be contextualized by the other . GREASELM demonstrates significant performance gains across different LM architectures . We perform experiments on several standard QA benchmarks : CommonsenseQA , OpenbookQA and MedQA-USMLE , which require external knowledge across different domains ( commonsense reasoning and medical reasoning ) and use different KGs ( ConceptNet and Disease Database ) . Across both domains , GREASELM outperforms comparably-sized prior QA models , including strong fine- tuned LM baselines ( by 5.5 % , 6.6 % , and 1.3 % , respectively ) and state-of-the-art KG+LM models ( by 0.9 % , 1.8 % , and 0.5 % , respectively ) on the three competitive benchmarks . Furthermore , with the deep fusion of both modalities , GREASELM exhibits strong performance over baselines on questions that exhibit textual nuance , such as resolving multiple constraints , negation , and hedges , and which require effective reasoning over both language context and KG . 2 RELATED WORK . Integrating KG information has become a popular research area for improving neural QA systems . Some works explore using two-tower models to answer questions , where a graph representation of knowledge and language representation are fused with no interaction between them ( Wang et al. , 2019 ) . Other works seek to use one modality to ground the other , such as using an encoded representation of a linked KG to augment the textual representation of a QA example ( e.g. , Knowledgeable Reader , Mihaylov & Frank , 2018 ; KagNet , Lin et al. , 2019 ; KT-NET , Yang et al. , 2019 ) . Others reverse the flow of information and use a representation of the text ( e.g. , final layer of LM ) to provide an augmentation to a graph reasoning model over an extracted KG for the example ( e.g. , MHGRN , Feng et al. , 2020 ; Lv et al. , 2020 ) . In all of these settings , however , the interaction between both modalities is limited as information between them only flows one way . More recent approaches explore deeper integrations of both modalities . Certain approaches learn to access implicit knowledge encoded in LMs ( Bosselut et al. , 2019 ; Petroni et al. , 2019 ; Hwang et al. , 2021 ) by training on structured KG data , and then use the LM to generate local KGs that can be used for QA ( Wang et al. , 2020 ; Bosselut et al. , 2021 ) . However , these approaches discard the static KG once they train the LM on its facts , losing important structure that can guide reasoning . More recently , QA-GNN ( Yasunaga et al. , 2021 ) proposed to jointly update the LM and GNN representations via message passing . However , they use a single pooled representation of the LM to seed the textual component of this joint structure , limiting the updates that can be made to the textual representation . In contrast to prior works , we propose to make individual token representations in the LM and node representations in the GNN mix for multiple layers , enabling representations of both modalities to reflect particularities of the other ( e.g. , knowledge grounds language ; language nuances specifies which knowledge is important ) . Simultaneously , we retain the individual structure of both modalities , which we demonstrate improves QA performance substantially ( §5 ) . Additionally , some works explore integrating knowledge graphs with language models in the pretraining stage . However , much like for QA , the modality interaction is typically limited to knowledge feeding language ( Zhang et al. , 2019 ; Shen et al. , 2020 ; Yu et al. , 2020 ) , rather than designing interactions across multiple layers . Sun et al . ( 2020 ) ’ s work is perhaps most similar , but they do not use the same interaction bottleneck , requiring high-precision entity mention spans for linking , and they limit expressivity through shared modality parameters for the LM and KG . 3 PROPOSED APPROACH : GREASELM . In this work , we augment large-scale language models ( Devlin et al. , 2019 ; Liu et al. , 2019 ; Lan et al. , 2020 ; Liu et al. , 2021 ) with graph reasoning modules over KGs . Our method , GREASELM ( depicted in Figure 1 ) , consists of two stacked components : ( 1 ) a set of unimodal LM layers which learn an initial representation of the input tokens , and ( 2 ) a set of upper cross-modal GREASELM layers which learn to jointly represent the language sequence and linked knowledge graph , allowing textual representations formed from the underlying LM layers and a graph representation of the KG to mix with one another . We denote the number of LM layers as N , and the number of GREASELM layers as M . The total number of layers in our model is N +M . Notation . In the task of multiple choice question answering ( MCQA ) , a generic MCQA-type dataset consists of examples with a context paragraph c , a question q and a candidate answer set A , all expressed in text . In this work , we also assume access to an external knowledge graph ( KG ) G that provides background knowledge that is relevant to the content of the multiple choice questions . Given a QA example ( c , q , A ) , and the KG G as input , our goal is to identify which answer a ∈ A is correct . Without loss of generality , when an operation is applied to an arbitrary answer , we refer to that answer as a . We denote a sequence of tokens in natural language as { w1 , . . . , wT } , where T is the total number of tokens , and the representation of a token wt from the ` -th layer of the model as h ( ` ) t . We denote a set of nodes from the KG as { e1 , . . . , eJ } , where J is the total number of nodes , and the representation of a node ej in the ` -th layer of the model as e ( ` ) j . 3.1 INPUT REPRESENTATION . We concatenate our context paragraph c , question q , and candidate answer a with separator tokens to get our model input [ c ; q ; a ] and tokenize the combined sequence into { w1 , . . . , wT } . Second , we use the input sequence to retrieve a subgraph of the KG G ( denoted Gsub ) , which provides knowledge from the KG that is relevant to this QA example . We denote the set of nodes in Gsub as { e1 , . . . , eJ } . KG Retrieval . Given each QA context , we follow the procedure from Yasunaga et al . ( 2021 ) to retrieve the subgraph Gsub from G. We describe this procedure in Appendix B.1 . Each node in Gsub is assigned a type based on whether its corresponding entity was linked from the context c , question q , answer a , or as a neighbor to these nodes . In the rest of the paper , we use “ KG ” to refer to Gsub . Interaction Bottlenecks . In the cross-modal GREASELM layers , information is fused between both modalities , for which we define a special interaction token wint and a special interaction node eint whose representations serve as the bottlenecks through which the two modalities interact ( §3.3 ) . We prepend wint to the token sequence and connect eint to all the linked nodes Vlinked in Gsub .

Summary: The paper presents a novel method of fusion of information from two modalities: text (context and question) and a Knowledge Base, for the task of question answering. Concretely, they propose GreaseLM layers that combine information from a Language Model (LM) and a Graph Neural Network (GNN) by mixing representations of embeddings from two special tokens (one in the LM and one in the GNN). The embedding in the LM passes information to the text modality via self attention; while the embedding in the GNN passes information to the KB modality via connections to other nodes. The modality interaction layer (dubbed MInt) is a simple MLP for combining the embeddings. The paper demonstrates superior performance compared to baseline methods on 3 different datasets. Additional ablations quantitatively show that the model performs better on questions with negations, hedging and constraints (modelled with a proxy of number of prepositional phrases).