Deep Bidirectional Language-Knowledge Graph Pretraining
Anonymous Author(s)
AffiliationAddress email
Abstract
Pretraining a language model (LM) on text helps various downstream NLP tasks.
Recent works show that a knowledge graph (KG) can complement text data, offering structured background knowledge and scaffold useful for reasoning.
However,these works are not pretrained to learn deep fusion of the two modalities at scale,limiting the potential to acquire fully joint representations of text and KG.
Here wepropose D RAGON ( D eep Bidi r ectional L a n g uage-
Kn o wledge Graph Pretrai n ing),a self-supervised method to pretrain a deeply joint language-knowledge modelfrom raw text and KG at scale.
Specifically, our model takes pairs of text segmentsand relevant KG subgraphs as input and bidirectionally fuses information from bothmodalities.
We pretrain this model by unifying two self-supervised reasoning objectives, masked language modeling and KG link prediction.
D RAGON outperformsexisting LMs and LM+KG models on diverse downstream tasks including questionanswering across general and biomedical domains, with +5% absolute gain onaverage across the board.
In particular, D RAGON achieves notable performance oncomplex reasoning about language and knowledge (+10% on questions involvinglong context or multi-step reasoning) and low-resource QA (+8% on OBQA and
RiddleSense), and new state-of-the-art results on various BioNLP tasks.
Introduction
Pretraining learns self-supervised representations from massive raw data to help various downstreamtasks [1].
Language models (LMs) pretrained on large text data, such as BERT [2] and GPTs [3],have shown strong performance on many natural language processing (NLP) tasks.
The success ofthese models comes from deeply interactive (contextualized) representations of inputs learned atscale via self-supervision [2, 4].
On the other hand, large knowledge graphs (KGs), like Freebase[5], Wikidata [6] and ConceptNet [7], provide complementary information to text data.
KGs offerstructured background knowledge by representing entities as nodes and relations between them asedges, and also offer scaffolds for structured reasoning such as multi-step reasoning about entities[8, 9, 10, 11] (§3.4.1).
This motivates research in combining the strengths of text data and KGs andpretraining deeply interactive representations of the two modalities at scale.
How to effectively combine text and KGs for pretraining is an open problem and presents challenges.
Given text and KG, we need both (i) a deeply bidirectional model for the two modalities to interact,and (ii) a self-supervised objective to learn joint reasoning over text and KG at scale.
Several existingworks [12, 13, 14, 15, 16] propose methods for self-supervised pretraining, but they only fuse textand KG in a shallow or uni-directional manner.
Another line of work [8, 9] proposes bidirectionalmodels for text and KG, but these models only perform finetuning on specific labeled downstreamtasks and do not perform self-supervised learning.
Consequently, existing methods may have limitedcapabilities to model and learn deep interactions over text and KG.
To address both of the above challenges and fully unify the strengths of text and KG, we propose
D RAGON ( D eep Bidi r ectional L a n g uage-Kn
o wledge Graph Pretrai n ing), an approach that performsdeeply bidirectional, self-supervised pretraining of a language-knowledge model from raw text and KG.
D RAGON has two core components: a cross-modal model that bidirectionally fuses text and KG,
[Table]
(cid:29)(cid:321)(cid:316)(cid:972)(cid:312)(cid:311)(cid:3)(cid:21)(cid:312)(cid:311)(cid:316)(cid:310)(cid:308)(cid:319)(cid:3)(cid:20)(cid:308)(cid:321)(cid:314)(cid:328)(cid:308)(cid:314)(cid:312)(cid:3)(cid:27)(cid:332)(cid:326)(cid:327)(cid:312)(cid:320)(cid:3) [Table] and a bidirectional self-supervised objective that learns joint reasoning over text and KG.
Concretely,as in Figure we take a text corpus and a KG as raw data and create inputs for the model bysampling a text segment from the corpus and extracting a relevant subgraph from the KG via entitylinking, obtaining a ( text , local KG ) pair.
We use a cross-modal model to encode this input intofused representations, where each layer of the model encodes the text with an LM and the KG witha graph neural network (GNN) and fuses the two with a bidirectional modality interaction module(GreaseLM [9]).
We pretrain this model by unifying two self-supervised reasoning tasks: (1) maskedlanguage modeling (MLM), which masks and predicts tokens in the input text, and (2) link prediction,which drops and predicts edges in the input KG.
The intuition is that by combining the two tasks,
MLM makes the model use the text jointly with structured knowledge in the KG to reason aboutmasked tokens in the text (e.g., in Figure 1, using the “round brush”–“art supply” multi-hop pathfrom the KG helps), and link prediction makes the model use the KG structure jointly with the textualcontext to reason about missing links in the KG (e.g., recognizing that “round brush could be usedfor hair” from the text helps).
This joint objective thus enables text to be grounded by KG structureand KG to be contextualized by text simultaneously, producing a deeply-unified language-knowledgepretrained model where information flows bidirectionally between text and KG for reasoning.
We pretrain D RAGON in two domains: a general domain, using the Book corpus and ConceptNet
KG [7] (§3), and a biomedical domain, using the PubMed corpus and UMLS KG [17] (§4).
Weshow that D RAGON improves on existing LMs and LM+KG models on diverse downstream tasksacross domains.
For the general domain, D RAGON outperforms RoBERTa [18], our base LM without
KGs, on various commonsense reasoning tasks such as CSQA, OBQA, RiddleSense and HellaSwag,with +8% absolute accuracy gain on average.
For the biomedical domain, D RAGON improves on theprevious best LM, BioLinkBERT [19], and sets new states of the art on BioNLP tasks such as MedQAand PubMedQA, with +3% accuracy gain.
In particular, D RAGON exhibits notable improvements on
QA tasks involving complex reasoning about language and knowledge (+10% gain on multi-step,negation, hedge, or long context reasoning) and on downstream tasks with limited training data (+8%gain).
These results show that our deep bidirectional self-supervision over text and KG producessignificantly superior language-knowledge representations compared to existing models.
All code,data and pretrained D RAGON models will be released upon publication.
1.1 Related work
Knowledge-augmented LM pretraining.
Knowledge integration is an active research for improving
LMs.
One line of works is retrieval-augmented LMs [20, 21, 22], which retrieve relevant text from acorpus and integrate it into LMs as additional knowledge.
Orthogonal to these works, we focus onusing knowledge bases as background knowledge, to ground reasoning about entities and facts.
Closest to our work are works that integrate knowledge base information in LM pretraining.
One lineof research adds entity features to LMs [12, 23, 24], but they do not incorporate KG structure.
Someworks use KG entity information or structure to create additional training signals [13, 25, 14, 26, 27,28], but they do not fuse text and KG in the same input context to model their interactions for reasoning.
Several works add KG information (e.g. triplets) directly to the LM input [29, 16, 15, 30, 31].
However, these methods propagate information from one side to the other (e.g., KG to text), limitingthe capacity to jointly reason over the two modalities.
Different from the above works, we propose tobidirectionally interact text and KG via a cross-modal model and joint self-supervised objective, sothat text and KG are grounded and contextualized by each other.
We find that this improves modelperformance on various reasoning tasks (§3).
We also note that while above works often focus addingentity- or triplet-level knowledge in KGs to LMs, we use larger KG subgraphs (200 nodes) as inputto enable richer contextualization and exploit the further potential of KGs as reasoning scaffolds.
KG-augmented question answering.
Various works designed KG-augmented reasoning modelsfor question answering [32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42].
In particular, recent works such as
QAGNN [8] and GreaseLM [9] suggest that a KG can scaffold reasoning about entities with its graphstructure and help for complex question answering (e.g., negation, multi-hop reasoning).
However,these works focus on training or finetuning models on specific QA datasets.
In contrast, we proposeto generalize this and integrate KG-augmented reasoning into general-purpose pretraining.
Ourmotivation is that self-supervised pretraining allows the model to learn from larger and more diversedata, helping to learn richer interactions between text and KGs and to acquire more diverse reasoningabilities beyond specific QA tasks.
We find that our proposed pretraining approach (D RAGON ) offerssignificant boosts over the baseline QA models (e.g. GreaseLM) on diverse downstream tasks (§3).
This opens a new research avenue in scaling up various carefully-designed QA models to pretraining.
KG representation learning.
Our link prediction task used in pretraining is motivated by researchin KG representation learning.
Link prediction is a fundamental task in KGs [43, 44], and variousworks study methods to learn KG entity and relation embeddings for link prediction, such as TransE[45], DistMult [46] and RotatE [47].
Several works additionally use textual data or pretrained LMsto aid learning KG embeddings and link prediction [48, 49, 50, 51, 52, 53].
While these works focuson the KG-side representations, we extend the scope and use the KG-side objective (link prediction)jointly with a text-side objective (language modeling) to train a mutually-interactive text-KG model.
Deep Bidirectional Language-Knowledge Graph Pretraining (D RAGON )
We propose D RAGON , an approach that performs deeply bidirectional, self-supervised pretraining ofa language-knowledge model from text and KG.
Specifically, as illustrated in Figure 1, we take a textcorpus and a large knowledge graph as raw data, and create input instances for the model by samplingcoarsely-aligned (text segment, local KG) pairs (§2.1).
To learn mutual interactions over text and
KG, D RAGON consists of a cross-modal encoder (GreaseLM) that fuses the input text-KG pairbidirectionally (§2.2), and a pretraining objective that performs bidirectional self-supervision on thetext-KG input (§2.3).
Our pretraining objective unifies masked language modeling (MLM) and KGlink prediction (LinkPred) to make text and KG mutually inform each other and learn joint reasoningover them.
Finally, we describe how we finetune the pretrained D RAGON model for downstreamtasks (§2.4).
While the individual piece of our approach (GreaseLM, MLM, LinkPred) is not newin itself, our contribution is that we are the first to bring these pieces together, present how to unifythem effectively and show that this produces a significantly performant pretrained model (§3, §4).
Definitions.
We define a text corpus W as a set of text segments W = { W } , and each text segment W as a sequence of tokens (words), W = { w 1 , ..., w I } .
We define a knowledge graph (KG) as amulti-relational graph G =
( V , E ) , where V is the set of entity nodes in the KG and E ⊆ V ×
R × Vis the set of edges (triplets) that connect nodes in V , with R being the set of relation types { r } .
Each triplet ( h, r, t ) in a KG can represent a knowledge fact such as ( Paris , in , France ) .
As araw KG is often large, with millions of nodes, a subgraph of the raw KG ( local KG ) is considered:
G =
( V, E ) where V = { v 1 , ..., v J } ⊆ V and E ⊆ E .
We define a language-knowledge model tobe a composition of two functions, f head ( f enc ( X )) , where the encoder f enc takes in an input X =( text segment W, local KG G ) , and produces a contextualized vector representation for each texttoken, { H 1 , ..., H I } , and for each KG node, { V 1 , ..., V J } .
A language model is a special case of alanguage-knowledge model with no KG ( J = 0 ).
The head f head uses these representations to performself-supervised tasks in the pretraining step and to perform downstream tasks in the finetuning step.
2.1 Input representation
Given a text corpus W and a large knowledge graph G , we create input instances for the model bypreparing (text segment W , local KG G ) pairs.
We want each pair’s text and KG to be (roughly)semantically aligned so that the text and KG can mutually inform each other and facilitate the modelto learn interactive reasoning between the two modalities.
Specifically, for each text segment W from
W , we extract a relevant local KG G for it from G via the following KG retrieval process.
KG retrieval.
Given a text segment W , we link entity mentions in W to entity nodes in G to get aninitial set of nodes V el .
We then add their 2-hop bridge nodes from G to get the total retrieved nodes
V ⊆ V .
Lastly, we add all edges that span these nodes in G to get E ⊆ E , which yields the finallocal KG, G =
( V, E ) , as well as our final input instance X
=
( W, G ) .
Appendix B.1 provides moredetails on KG retrieval.
Henceforth, we use “KG” to refer to this local KG G unless noted otherwise.
Modality interaction token/node.
For each obtained (text, KG) pair, we further add a special token(interaction token) w int to the text and a special node (interaction node) v int to the KG, which willserve as an information pooling point for each modality as well as an interface for modality interactionin our cross-modal encoder (§2.2).
Specifically, we prepend w int to the original text W = { w 1 , ..., w I } ,and connect v int to the entity-linked nodes in the original KG, V el ⊆ V = { v 1 , ..., v J } , using a newrelation type r el .
The interaction token and node can also be used to produce a pooled representationof the whole input, e.g., when finetuning for classification tasks (§2.4).
2.2 Cross-modal encoder
To model mutual interactions over the text and KG, we use a bidirectional sequence-graph encoderfor f enc which takes in the text tokens and KG nodes and exchanges information across them formultiple layers to produce a fused representation of each token and node (Figure 1 right):
[Equation]
While we may use any deep bidirectional sequence-graph encoder for f enc , for controlled comparisonwith existing works, we adopt the existing top-performing sequence-graph architecture, GreaseLM[9], which combines Transformers [54] and graph neural networks (GNNs) to fuse text-KG inputs.
Specifically, GreaseLM first uses N layers of Transformer language model (LM) layers to map theinput text into initial token representations, and uses KG node embeddings to map the input KGnodes into initial node representations,
[Equation]
Then it uses M layers of text-KG fusion layers to encode these token/node representations jointlyinto the final token/node representations,
[Equation]
where each of the fusion layers ( ℓ =1 , ..., M ) performs the following:
[Equation]
H ( ℓ − 1) I } ) , (5) { (cid:101) V ( ℓ ) int , V ( ℓ ) 1 ,..., V ( ℓ ) J } = GNN-Layer ( { V ( ℓ − 1) int , V ( ℓ − 1),..., V ( ℓ − 1) J } ) , (6) [ H ( ℓ ) int ; V ( ℓ ) int ] = MInt ([ (cid:101) H ( ℓ ) int ; (cid:101) V ( ℓ ) int ]) .
Here GNN induces graph structure-aware representations of KG nodes, [ · ; · ] does concatenation, and
MInt (modality interaction module) exchanges information between the interaction token (text side)and interaction node (KG side) via an MLP.
For more details on GreaseLM, we refer readers to [9].
2.3 Pretraining objective
We aim to pretrain the D RAGON model so that it learns joint reasoning over text and a KG.
To ensurethat the text and KG mutually inform each other and the model learns bidirectional information flow,we unify two self-supervised reasoning tasks: masked language modeling and KG link prediction.
Masked language modeling (MLM).
MLM is a common pretraining task used for language models(e.g., BERT [2], RoBERTa [18]), which masks some tokens in the input text and predicts them.
Thistask makes the model use non-masked context to reason about masked tokens, and in particular, as our approach takes a joint text-KG pair as input, we expect that MLM can encourage the model tolearn to use the text jointly with structured knowledge in the KG to reason about masks in the text(e.g.
, in the example of Figure 1, besides the textual context, recognizing the “round brush”–“artsupply” path from the KG can help together to predict the masked tokens “art supplies”).
Concretely, to perform the MLM task, we mask a subset of tokens in the input text, M ⊆ W , with aspecial token [MASK] , and let the task head f head be a linear layer that takes the contextualized tokenvectors { H i } from the encoder to predict the original tokens.
The objective is a cross-entropy loss:
[Equation]
Link prediction (LinkPred).
While the MLM task predicts for the text side, link prediction holdsout some edges from the input KG and predicts them.
Link prediction is a fundamental task in KGs[47] and makes the model use the structure of KGs to perform reasoning (e.g., using a compositionalpath “X’s mother’s husband is Y” to deduce a missing link “X’s father is Y”).
In particular, as ourapproach takes a joint text-KG pair as input, we expect that link prediction can encourage the modelto learn to use the KG structure jointly with the textual context to reason about missing links in the
KG (e.g., in Figure 1, besides the KG structure, recognizing that “round brush could be used for hair”from the text can help together to predict the held-out edge (round_brush, at, hair) ).
Concretely, to perform the link prediction task, we hold out a subset of edge triplets from the input
KG, S =
{ ( h, r, t ) } ⊆ E .
For the task head f head , we adopt a KG representation learning framework,which maps each entity node ( h or t ) and relation ( r ) in the KG to a vector, h , t , r , and defines ascoring function ϕ r ( h , t ) to model positive/negative triplets.
Specifically, we let h
= V h , t =
V t ,r = R r , with { V j } being the contextualized node vectors from the encoder, and R = { r 1 , ..., r |R| }learnable relation embeddings.
We consider a KG triplet scoring function ϕ r ( h , t ) such as
DistMult [46]: ⟨ h , r , t ⟩ ,
TransE
[Equation]
where ⟨· , · , ·⟩ denotes the trilinear dot product and ⊙ the Hadamard product.
A higher ϕ indicates ahigher chance of ( h, r, t ) being a positive triplet (edge) instead of negative (no edge).
We analyze thechoices of scoring functions in §3.4.3.
For training, we optimize the objective:
[Equation]
where ( h ′ , r, t ′ ) are n negative samples corresponding to the positive triplet ( h, r, t ) , γ is the margin,and σ is the sigmoid function.
The intuition of this objective is to make the model predict triplets ofthe held-out edges S as positive and other random triplets as negative.
Joint training.
To pretrain D RAGON , we optimize the MLM and LinkPred objectives jointly:
L = L MLM + L LinkPred .
This joint objective unifies the effects of MLM and LinkPred, whichencourage the model to simultaneously ground text with KG structure and contextualize KG withtext, facilitating bidirectional information flow between text and KGs for reasoning.
We show in§3.4.3 that the joint objective yields a more performant model than using one of the objectives alone.
2.4 Finetuning
Lastly, we describe how we finetune D RAGON for downstream tasks such as text classificationand multiple-choice QA (MCQA).
Given an input text W (e.g., concatenation of a question andan answer choice in the case of MCQA), we follow the same steps as §2.1 and §2.2 to retrieve arelevant local KG G and encode them jointly into contextualized token/node vectors, { H int , H 1 , ...,
H I } , { V int , V 1 , ..., V J } .
We then compute a pooled representation of the whole input as X =
MLP ( H int , V int , G ) , where G denotes attention-based pooling of { V j | v j ∈ { v 1 , ..., v J }} using H int as a query.
Finally, the pooled representation X is used to perform the downstream task, in thesame way as how the [CLS] representation is used in LMs such as BERT and RoBERTa.
The difference from GreaseLM is that while GreaseLM only performs finetuning as described in thissection (hence, it is an LM finetuned with KGs), D RAGON performs self-supervised pretraining asdescribed in §2.3 (hence, it can be viewed as an LM pretrained + finetuned with KGs).
Experiments: General domain
We experiment with the proposed approach in a general domain first.
We pretrain D RAGON using the
Book corpus and ConceptNet KG (§3.1), and evaluate on diverse downstream tasks (§3.2).
We showthat D RAGON significantly improves on existing models (§3.4).
We extensively analyze the effect of
D RAGON ’s key design choices such as self-supervision and use of KGs (§3.4.1, 3.4.2, 3.4.3).
3.1 Pretraining setup
Data.
For the text data, we use BookCorpus [55], a general-domain corpus widely used in LMpretraining (e.g., BERT, RoBERTa).
It has 6GB of text from online books.
For the KG data, we use
ConceptNet [7], a general-domain knowledge graph designed to capture background commonsenseknowledge.
It has 800K nodes and 2M edges in total.
To create a training instance, we sample a textsegment of length up to 512 tokens from the text corpus, then retrieve a relevant KG subgraph of sizeup to 200 nodes (details in Appendix B.1), by which we obtain an aligned (text, local KG) pair.
Implementation.
For our encoder (§2.2), we use the exact same architecture as GreaseLM [9](19 LM layers followed by 5 text-KG fusion layers; 360M parameters in total).
As done by [9], weinitialize parameters in the LM component with the RoBERTa-Large release [18] and initialize the
KG node embeddings with pre-computed ConceptNet entity embeddings (details in Appendix B.2).
For the link prediction objective (§2.3, Equation 10), we use DistMult [46] for KG triplet scoring,with a negative exampling of 128 triplets and a margin of γ = 0 .
To pretrain the model, we perform
MLM with a token masking rate of 15% and link prediction with an edge drop rate of 15%.
Wepretrain for 20,000 steps with a batch size of 8,192 and a learning rate of 2e-5 for parameters inthe LM component and 3e-4 for the others.
Training took 7 days on eight A100 GPUs using FP16.
Additional details on the hyperparameters can be found in Appendix B.3.
3.2 Downstream evaluation tasks
We finetune and evaluate D RAGON on nine diverse commonsense reasoning benchmarks: CommonsenseQA ( CSQA ) [56], OpenbookQA ( OBQA ) [57], RiddleSense ( Riddle ) [58], AI2 Reasoning
Challenge–Challenge Set ( ARC ) [59], CosmosQA [60], HellaSwag [61], Physical Interaction QA( PIQA ) [62], Social Interaction QA ( SIQA ) [63], and Abductive Natural Language Inference ( aNLI )[64].
Appendix B.4 provides details on these tasks and data splits.
For each task, we repeat three runsand report the average accuracy.
Hyperparameters used for finetuning can be found in Appendix B.3.
3.3 Baselines
LM.
To study the effect of using KGs, we compare D RAGON with the vanilla language model,
RoBERTa
As we initialize D RAGON ’s parameters using the RoBERTa-Large release (§3.1), forfair comparison, we let the baseline be such that we take the RoBERTa-Large release and continuepretraining it with the vanilla MLM objective on the same text data for the same number of steps as
D RAGON .
Hence, the only difference is that D RAGON uses KGs during pretraining while RoBERTadoes not.
We then perform standard LM finetuning of RoBERTa on downstream tasks.
LM finetuned with KG.
We also compare with existing KG-augmented QA models, QAGNN [8]and GreaseLM [9], which finetune a vanilla LM (i.e. RoBERTa-Large) with a KG on downstreamtasks, but do not pretrain with a KG.
GreaseLM is the existing top-performing model in this paradigm.
As we use the same encoder architecture as GreaseLM for D RAGON , the only difference from
GreaseLM is that D RAGON performs self-supervised pretraining while GreaseLM does not.
3.4 Results
Table 1 shows performance on the 9 downstream commonsense reasoning tasks.
Across all tasks,
D RAGON consistently outperforms the existing LM (RoBERTa) and KG-augmented QA models(QAGNN, GreaseLM), e.g., +7% absolute accuracy boost over RoBERTa and +
% over GreaseLM on
OBQA .
D RAGON also performs competitively to all other top systems on the leaderboad (see Appendix
C).
These accuracy boosts indicate the advantage of D RAGON over RoBERTa (KG reasoning)
andover GreaseLM (pretraining).
The gain is especially significant on datasets that have small training
[Table] [Table]
data such as ARC , Riddle and OBQA , and datasets that require complex reasoning such as CosmosQAand HellaSwag , which we analyze in more detail in the following sections.
3.4.1 Analysis: Effect of knowledge graph
The first key contribution of D RAGON (w.r.t.
existing LM pretraining methods) is that we incorporate
KGs.
We find that this significantly improves the model’s performance for robust and complexreasoning, such as resolving multi-step reasoning and negation, as we discuss below.
Quantitative analysis.
In Table 2, we study downstream task performance of D RAGON on questionsinvolving complex reasoning.
Building on [8, 9], we consider several proxies to categorize complexquestions: (i) presence of negation (e.g. no , never ), (ii) presence of conjunction (e.g. and , but ), (iii)presence of hedge (e.g. sometimes , maybe ), (iv) number of prepositional phrases, and (v) number ofentity mentions.
Having negation or conjunction indicates logical multi-step reasoning, having moreprepositional phrases or entity mentions indicates involving more reasoning steps or constraints, andhaving hedge terms indicates involving complex textual nuance.
D RAGON significantly outperformsthe baseline LM (RoBERTa) across all these categories (e.g., +14% accuracy for negation), whichconfirms that our joint language-knowledge pretraining boosts reasoning performance.
D RAGONalso consistently outperforms the existing KG-augmented QA models (QAGNN, GreaseLM).
Wefind that QAGNN and GreaseLM only improve moderately on RoBERTa for some categories likeconjunction or many prepositional phrases (=2, 3), but D RAGON provides substantial boosts.
Thissuggests that through self-supervised pretraining with larger and diverse data, D RAGON has learnedmore general-purpose reasoning abilities than the finetuning-only models like GreaseLM.
Qualitative analysis.
Using the CSQA dataset, we further conducted case studies on the behaviorof D RAGON ’s KG reasoning component, where we visualize how graph attention weights changegiven different question variations (Figure 2).
We find that D RAGON exhibits abilities to extrapolateand perform robust reasoning.
For instance, D RAGON adjusts the entity attention weights and finalpredictions accordingly when we add conjunction or negation about entities (A1, A2) or when weadd extra context to an original question (B1 → B2), but existing models, RoBERTa and GreaseLM,struggle to predict the correct answers.
As these questions are more complex than ones typically seenin the CSQA training set, our insight is that while vanilla LMs (RoBERTa) and finetuning (GreaseLM)have limitation in learning complex reasoning, KG-augmented pretraining (D RAGON ) helps acquiregeneralizable reasoning abilities that extrapolate to harder test examples.
3.4.2 Analysis: Effect of pretraining
Another key contribution of D RAGON (w.r.t.
existing QA models like GreaseLM) is pretraining.
Here we discuss when and why our pretraining is useful.
Considering the three core factors inmachine learning (data, task complexity, and model capacity), pretraining helps when the availabledownstream task data is smaller compared to the downstream task complexity or model capacity.
Concretely, we find that D RAGON is especially helpful for the following three scenarios.
Downstream tasks with limited data.
In Table 1, we find that D RAGON provides significant boostsover GreaseLM on downstream tasks with limited finetuning data available, such as ARC (3K training [Table] [Table] Glamping GNN Final Layer c amp
Int t r ip beach school beach
Int RoBERTa: A. camp ( ✗ ) GreaseLM: C. camp ( ✗ ) Glamping: B. school ( ✓ ) (A1)
Conjunction sc h
o o l Where would you use a folding chair but not store one?
A. garage B. school C .
beach t r ip
Glamping GNN 1st Layer
Glamping GNN Final Layer g a ra g e
Int t r ip beach school beach
Int (A2) Negation + Conjunction g a ra g e m ovi
e You will buy a ticket for entering what building for entertainment ?
A. station B .
movie theater [Figure] [Table] [Figure]
e n t e rtain m e nt
Int station (B1)
Single context m ovi
e You don't enjoy watching pre-recorded performance .
You will buy a ticket for entering what building for entertainment ?
A. station B. movie theater C .
concert hall
Glamping GNN 1st Layer
Glamping GNN Final Layer
Int l i ve movie Int (B2) Multi context m ovi e e n t e rtain m e nt station Int l i ve
Model Prediction
RoBERTa: B. school ( ✗ ) GreaseLM:
[Table]
✗
[Table]
✓ )
[Table]
RoBERTa: B. movie theater ( ✓ ) GreaseLM: B. movie theater ( ✓ ) Glamping: B. movie theater ( ✓ ) Model Prediction RoBERTa: B. movie theater ( ✗ ) GreaseLM: B. movie theater ( ✗ ) Glamping: C. concert hall ( ✓ ) Model Prediction
[Figure]
amp
[Figure]
r e cord r e cord
[Figure]
[Table] [Figure] ool [Figure]
t r i p
[Figure]
c amp
[Figure] t r i p [Figure] s c h ool
[Figure]
t r i p
[Figure] g a r a g e [Figure] t r i p
[Figure]
g a r a g e m ovie
[Figure] ent [Figure] l i v e
[Figure]
m ovie e n tertainm ent [Figure] l i v e [Figure]
c amp
, the chair is not used for a trip. A2: D RAGON shows strong attention to “trip” and “beach”, likely because the question now states “ but not store one”—hence, the chair is used for a trip. B1 → B2: D RAGON ’s final GNN layer shows strong attention to “movie” in the original question (B1), but after adding the extra context “don’t enjoy pre-record” (B2), D RAGON shows strong attention to “live” and “concert”, leading to making the correctly adjusted prediction “concert hall”. One interpretation of these findings is that D RAGON leverages the KG’s graph structure as a scaffold for performing complex reasoning.
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table]
[Table] [Table] [Table] instances; +4% accuracy gain), Riddle (3K instances; +4% accuracy) and OBQA (5K instances; +5%accuracy).
For other tasks, we also experimented with a low-resource setting where 10% of finetuningdata is used (Table 3).
Here we also see that D RAGON attains significant gains over GreaseLM (+5%accuracy on PIQA ), suggesting the improved data-efficiency of D RAGON .
Complex downstream tasks.
In Table 1, we find that D RAGON provides substantial gainsover GreaseLM on downstream tasks involving more complex reasoning, such as CosmosQA and
HellaSwag , where the inputs have longer context and more entities (thus bigger local KGs).
Forthese tasks, improvements of GreaesLM over RoBERTa were small (+0.1% on CosmosQA ), but
D RAGON provides substantial boosts (+1.8%).
Our insight is that through self-supervised pretrainingwith larger and more diverse data, D RAGON has learned richer text-KG interactions than GreaseLM,enabling solving more complex downstream tasks.
Similarly, as seen in §3.4.1, D RAGON also attainslarge gains over GreaseLM on complex questions containing negation, conjunction and prepositionalphrases (Table 2), and extrapolates to questions more complex than seen in training sets (Figure 2).
Increased model capacity.
In Table 4, we study downstream performance when the model capacityis increased—the number of text-KG fusion layers is increased from 5 to 7—for both GreaseLM and
D RAGON .
We find that increased capacity does not help for the finetuning-only model (GreaseLM) aswas also reported in the original GreaseLM paper, but it helps when pretrained (D RAGON ).
This resultreveals that increased model capacity can actually be beneficial when combined with pretraining, andsuggests the promise of D RAGON to be further scaled up.
3.4.3 Analysis: Design choices of D RAGON
Pretraining objective (Table 5 top).
The first important design choice of D RAGON is the jointpretraining objective: MLM + LinkPred (§2.3).
Using the joint objective outperforms using MLM or
LinkPred alone (+5% accuracy on OBQA ).
This suggests that having the bidirectional self-supervisedtasks on text and KG facilitates the model to fuse the two modalities for reasoning.
Link prediction head choice (Table 5 middle 1).
KG representation learning is an active area ofresearch, and various KG triplet scoring models are proposed (Equation 9).
We hence experimentedwith using different scoring models for D RAGON ’s link prediction head (§2.3).
We find that while
DistMult has a slight edge, all variants we tried (DistMult, TransE, RotatE) are effective, outperforming the baseline without LinkPred (“MLM only”).
This result suggests the generality of D RAGONand its promise to be combined with various KG representation learning techniques.
Cross-modal model (Table 5 middle 2).
Another core component of D RAGON is the cross-modalencoder with bidirectional text-KG fusion layers (§2.2).
We find that if we ablate them and simplyconcatenate text and KG representations at the end, the performance drops substantially.
This resultsuggests that deep bidirectional fusion is crucial to model interactions over text and KG for reasoning.
KG structure (Table 5 bottom).
The final key design of D RAGON is that we leverage the graphstructure of KGs via a sequence-graph encoder and link prediction objective.
Here we experimentedwith an alternative pretraining method that drops the graph structure: we convert triplets in the local
KG into sentences using a template [33], append them to the main text input, and perform vanilla
MLM pretraining.
We find that D RAGON substantially outperforms this variant (+2% accuracy on
OBQA ), which suggests that the graph structure of KGs helps the model perform reasoning.
Experiments: Biomedical domain
Biomedicine is a domain with extensive background knowledge [65, 66, 67, 1], and experts curatevarious knowledge bases for it [68, 17, 69, 70].
We hypothesize that these biomedical KGs canenable deeper understanding and reasoning about biomedical text.
With this motivation, we pretrain
D RAGON on a biomedical corpus and KG, and evaluate on biomedical downstream tasks.
Pretraining setup.
For the text data, we use PubMed [71], a widely-used corpus in biomedial LMtraining (e.g., BioBERT [72], PubmedBERT [73]).
It contains the abstracts of biomedical paperson PubMed and has 21GB of text.
For the KG data, we use the Unified Medical Language System(UMLS)
[17], a widely-used knowledge graph in biomedicine.
It has 300K nodes and 1M edges intotal.
For training, we follow the same procedure as the experiment in the general domain (§3.1),except that we initialize D RAGON ’s LM component with BioLinkBERT-Large [19], the state-of-theart biomedical LM, instead of RoBERTa-Large.
Note that while “BioLinkBERT” has “Link” in itsname, it is not about KG links but about citation links that the model was originally pretrained with.
Downstream evaluation tasks.
We finetune and evaluate D RAGON on three popular biomedical
NLP and reasoning benchmarks: MedQA-USMLE ( MedQA ) [74], PubMedQA [75], and BioASQ[76].
Appendix B.4 provides details on these tasks and data splits.
[Table] [Table]
Baselines.
We compare D RAGON with the vanilla
LM (BioLinkBERT) and LMs finetuned with the KG(QAGNN and GreaseLM seeded with BioLinkBERT).
Results.
Table 6 summarizes model performance onthe downstream tasks.
Across tasks, D RAGON outperforms all the existing biomedical LMs and KGaugmented QA models, e.g., +3% absolute accuracyboost over BioLinkBERT and +2% over GreaseLM on
MedQA , achieving new state-of-the-art performance onthese tasks.
This result suggests significant efficacy of KG-augmented pretraining for improvingbiomedical reasoning tasks.
Combined with the results in the general commonsense domain (§3.4),our experiments also suggest the domain-generality of D RAGON , serving as an effective pretrainingmethod across domains with different combinations of text, KGs and seed LMs.
Conclusion
We presented D RAGON , a self-supervised pretraining method to learn a deeply bidirectional languageknowledge model from text and knowledge graphs (KGs) at scale.
In both general and biomedicaldomains, D RAGON outperforms existing language models and KG-augmented models on various NLPtasks, and exhibits notably strong performance on complex reasoning such as answering questionsinvolving long context or multi-step reasoning.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes] See Section
(b) Did you describe the limitations of your work?
[Yes] See Appendix A
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] SeeAppendix A(d)
Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes] See Appendix A
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A]
(b) Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Section 3.1 and Appendix B.
(c) Did you report error bars (e g., with respect to the random seed after running experi- ments multiple times)?
[Yes] (d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[Yes] See Section 3.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] See Section 3 and Section
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
[Table] (a) Did you include the full text of instructions given to participants and screenshots, if [Table]
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A]
A Ethics, limitations and risks
We outline potential ethical issues with our work below.
First, D RAGON is a method to fuse languagerepresentations and knowledge graph representations for joint reasoning.
Consequently, D RAGONcould reflect the same biases and toxic behaviors exhibited by language models and knowledge graphsthat are used to initialize it.
For example, language models have been shown to encode biases aboutrace, gender, and other demographic attributes [77, 78] and generate toxic outputs [79].
Because
D RAGON is seeded with pretrained language models that often learn these patterns, it is possible toreflect them in open-world settings.
Second, the ConceptNet knowledge graph [7] used in this workhas been shown to encode stereotypes [80], rather than completely clean commonsense knowledge.
If D RAGON were used outside these standard benchmarks in conjunction with ConceptNet as a
KG, it might rely on unethical relationships in its knowledge resource to arrive at conclusions.
Consequently, while D RAGON could be used for applications outside these standard benchmarks,we would encourage implementers to use the same precautions they would apply to other languagemodels and methods that use noisy knowledge sources.
Another source of ethical concern is the use of the MedQA-USMLE evaluation.
While we find thisclinical reasoning task to be an interesting testbed for D RAGON and for joint language and knowledgereasoning in general, we do not encourage users to use these models for real world clinical prediction.
Finally, another limitation is that the current D RAGON is an encoder language model and does notperform language generation by itself.
Extending D RAGON to language generation and advancing
KG-enhanced language generation [28, 81] will be an important future research direction.
Acknowledgment: We referred to [9] for writing this statement.
B Experimental Setup Details
B.1 KG retrieval
Given each input text segment W , we follow the procedure from Yasunaga et al.
[8] to retrieve arelevant local KG G from the raw KG G =
( V , E ) .
First, we use the entity linker from the spaCylibrary 1 to link entity mentions in W to entity nodes in G , obtaining an initial set of nodes V el .
Second,we add any bridge entities in G that are in a 2-hop path between any pair of linked entities in V el to getthe total retrieved nodes V ⊆ V .
If the number of nodes in V exceeds 200, we prune V by randomlysampling 200 nodes from it to be the final retrieved nodes V .
Lastly, we retrieve all the edges in Gthat connect any two nodes in V to obtain E ⊆ E , forming the final local KG, G =
( V, E ) .
B.2 Graph initialization
For the ConceptNet knowledge graph used in the general commonsense domain (§3), we follow themethod of MHGRN [33] to prepare the initial KG node embeddings.
Specifically, we convert tripletsin the KG into sentences using pre-defined templates for each relation.
Then, these sentences are fedinto BERT-Large [2] to compute embeddings for each sentence.
Finally, for each entity, we collectall sentences containing the entity, extract all token representations of the entity’s mention spans inthese sentences, and return the mean pooling of these representations.
For the UMLS knowledge graph used in the biomedical domain (§4), node embeddings are initializedsimilarly using the pooled token output embeddings of the entity name from BioLinkBERT [19].
While extremely rare (< 1%), in case when the input text does not yield any linked entity, we representthe graph using a dummy node initialized with 0, i.e., D RAGON backs off to only using the text siderepresentations because the graph propagates no information.
B.3 Hyperparameters
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
B.4 Downstream evaluation tasks
We use the following nine commonsense reasoning benchmarks for the experiments in the generaldomain (§3).
CommonsenseQA (CSQA) [56] is a 5-way multiple-choice QA task testing commonsense reasoning.
The dataset has 12,102 questions.
We use the in-house data splits by [32].
OpenbookQA (OBQA)
[57] is a 4-way multiple-choice QA task containing elementary sciencequestions.
It has 5,957 questions.
We use the original data splits in [36].
RiddleSense (Riddle) [58] is a 5-way multiple-choice task testing complex riddle-style commonsensereasoning.
It has 5,715 questions.
We split the dev set in half to make in-house dev/test sets.
AI2 Reasoning Challenge, Challenge Set (ARC)
[59] is a 4-way multiple-choice QA task containingscience exam questions.
It has 2,590 questions.
We use the original data splits in [59].
CosmosQA [60] is a 4-way multiple-choice QA task testing commonsense reasoning with longnarratives.
It has 35.6K questions.
We split the dev set in half to make in-house dev/test sets.
HellaSwag [61] is a 4-way multiple-choice task testing grounded commonsense reasoning aboutevents.
It has 70K questions.
We split the dev set in half to make in-house dev/test sets.
Physical Interaction QA (PIQA) [62] is a 3-way multiple-choice QA task testing physics reasoningabout objects.
It has 20K questions.
We split the dev set in half to make in-house dev/test sets.
Social Interaction QA (SIQA) [63] is a 3-way multiple-choice QA task testing social commonsensereasoning.
It has 37K questions.
We use the original data splits in [63].
Abductive Natural Language Inference (aNLI) [64] is a 2-way multiple-choice task testingabductive commonsense reasoning.
It has 170K questions.
We use the original data splits in [64].
For the experiments in the biomedical domain (§4), we use the following three biomedical NLP andreasoning benchmarks.
MedQA-USMLE
(MedQA) [74] is a 4-way multiple-choice task containing United States Medical
License Exam questions.
The dataset has 12,723 questions.
We use the original data splits in [74].
PubMedQA [75] is a 3-way multiple-choice task testing biomedical language understanding andreasoning.
The dataset has 1,000 questions.
We use the original data splits in [75].
BioASQ [76] is a 2-way multiple-choice task testing biomedical language understanding and reason-
[Table]
ing.
The dataset has 885 questions.
We use the original data splits in [76]
[Table]
the school play, Robin played a hero in the struggle to the death with the angry villain.
How would others feel as a result?
(A) sorry for
[Table]
villain
(B) hopeful
[Table]
Robin will succeed
(C)
[Table]
Robin should lose the fight
[Table]
Obs1:
It was a gorgeous day outside.
Obs2:
She asked her neighbor for a jump-start.
Hyp1:
Mary decided to drive to the beach, but her car would not start due to a dead battery.
Hyp2:
It made a weird sound upon starting.
[Table]
A 57-year-old man presents to his primary care physician with a 2-month history of right upper and lower extremity weakness.
He noticed the weakness when he started falling far more frequently while running errands.
Since then, he has had increasing difficulty with walking and lifting objects.
His past medical history is significant only for well-controlled hypertension, but he says that some members of his family have had musculoskeletal problems.
His right upper extremity shows forearm atrophy and depressed reflexes while his right lower extremity is hypertonic with a positive Babinski sign.
Which of the following is most likely associated with the cause of this patients symptoms?
(A) HLA-B8 haplotype (B) HLA-DR2 haplotype (C) Mutation in SOD (D) Mutation in SMN
[Table]
Recent studies have demonstrated that statins have pleiotropic effects, including anti-inflammatory effects and atrial fibrillation (AF) preventive effects [...] 221 patients underwent CABG in our hospital from 2004 to 2007.
14 patients with preoperative AF and 4 patients with concomitant valve surgery [...]
The overall incidence of postoperative AF was 26%.
Postoperative AF was significantly lower in the Statin group compared with the Non-statin group (16%versus 33%, p=0.005).
Multivariate analysis demonstrated that independent predictors of AF [...]
Do preoperative statins reduce atrial fibrillation after coronary artery bypass grafting?
(A) yes (B) no (C)
maybe
BioASQ
LT4 absorption is unchanged by concomitant metformin ingestion.
It
[Table]
C Additional Experimental Results
[Table]