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: This paper presents a KG augmented LM which fuses graph and contextual representations. The model, when evaluated on multi-choice QA, outperforms strong baselines.  The context, question, answer pairs are encoded with a standard LM and the nodes of a relevant subgraph which is retrieved from a KG are encoded with a GNN. The representations of (i) a special interaction token (same as CLS), and (ii) a special interaction node, are input to custom layers that mix/fuse them which potentially makes them more informative to the downstream MLP head.  The authors experiment on three MCQA datasets: commonsenseQA, openbookQA, and medQA-USMLE (from the biomedical domain). ConceptNet is used as the KG for the first two datasets and a self-constructed KG is used for the third. The evaluation and analysis show that this method outperforms strong baselines (sometimes with a much higher parameter count). especially when the questions are complex and have prepositional phrases, negation terms, etc.