Deep Bidirectional Language-Knowledge Pretraining
Anonymous Author(s)
AfﬁliationAddress 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.
Herewe propose D RAGON ( D eep Bidi r ectional L a n g uage-Kn
o wledge Pretrai n ing), aself-supervised method to pretrain a deeply joint language-knowledge model fromraw text and KG at scale.
Speciﬁcally, our model takes pairs of text segments andrelevant 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 ﬁnetuning on speciﬁc 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 Pretrai n ing), an approach that performs deeplybidirectional, 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, and
[Table]
(cid:17)(cid:45)(cid:32)(cid:47)(cid:45)(cid:28)(cid:36)(cid:41)(cid:1) (cid:5) (cid:19)(cid:2)(cid:8)(cid:16)(cid:15) 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-uniﬁed language-knowledgepretrained model where information ﬂows 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 producessigniﬁcantly 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],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 [28, 16, 15, 29, 30].
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 ﬁnd 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 [31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41].
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 ﬁnetuning models on speciﬁc 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 speciﬁc QA tasks.
We ﬁnd that our proposed pretraining approach (D RAGON ) offerssigniﬁcant 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 [42, 43], and variousworks study methods to learn KG entity and relation embeddings for link prediction, such as TransE[44], DistMult [45] and RotatE [46].
Several works additionally use textual data or pretrained LMsto aid learning KG embeddings and link prediction [47, 48, 49, 50, 51, 52].
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 Pretraining (D RAGON )
We propose D RAGON , an approach that performs deeply bidirectional, self-supervised pretraining ofa language-knowledge model from text and KG.
Speciﬁcally, 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 uniﬁes 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 ﬁnetune 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 ﬁrst to bring these pieces together, present how to unifythem effectively and show that this produces a signiﬁcantly performant pretrained model (§3, §4).
Deﬁnitions.
We deﬁne a text corpus W as a set of text segments W = { W } , and each text segment 118 W as a sequence of tokens (words), W = { w 1 , ..., w I } .
We deﬁne 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 ⇥ V 120 is 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 deﬁne a language-knowledge model to 124 be 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 ﬁnetuning 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.
Speciﬁcally, 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 ﬁnal 138 local KG, G =
( V, E ) , as well as our ﬁnal 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).
Speciﬁcally, 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 ﬁnetuning for classiﬁcation 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 [53] and graph neural networks (GNNs) to fuse text-KG inputs.
Speciﬁcally, GreaseLM ﬁrst 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 ﬁnal token/node representations,
[Equation]
where each of the fusion layers ( ` =1 , ..., M ) performs the following:
[Equation]
1) int , H ( `  1),..., H ( `  1) I } ) , (5) { e V ( ` ) int , V ( ` ) 1 ,..., V ( ` ) J } = GNN-Layer ( { V ( `  1) int , V ( `  1),..., V ( `  1) J } ) , (6) [ H ( ` ) int ; V ( ` ) int ] =
MInt ([ e H ( ` ) int ; e 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 ﬂow,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[46] 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 deﬁnes ascoring function  r ( h , t ) to model positive/negative triplets.
Speciﬁcally, 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 [45]: h h , r , t i ,
TransE [44]:  k h + r  t k , RotatE [46]:  k [Equation]
[Equation]
)where h · , · , · i denotes the trilinear dot product and  the Hadamard product.
A higher  indicates a 199 higher 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 0 , r, t 0 ) 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 uniﬁes the effects of MLM and LinkPred, which 207 encourage the model to simultaneously ground text with KG structure and contextualize KG withtext, facilitating bidirectional information ﬂow 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 ﬁnetune D RAGON for downstream tasks such as text classiﬁcationand 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 2 { v 1 , ..., v J }} using 217 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 ﬁnetuning as described in thissection (hence, it is an LM ﬁnetuned with KGs), D RAGON performs self-supervised pretraining asdescribed in §2.3 (hence, it can be viewed as an LM pretrained + ﬁnetuned with KGs).
Experiments: General domain
We experiment with the proposed approach in a general domain ﬁrst.
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 signiﬁcantly 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 [54], 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 [45] 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 ﬁnetune and evaluate D RAGON on nine diverse commonsense reasoning benchmarks: CommonsenseQA ( CSQA ) [55], OpenbookQA ( OBQA ) [56], RiddleSense ( Riddle ) [57], AI2 Reasoning
Challenge–Challenge Set ( ARC ) [58], CosmosQA [59], HellaSwag [60], Physical Interaction QA( PIQA ) [61], Social Interaction QA ( SIQA ) [62], and Abductive Natural Language Inference ( aNLI )[63].
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 ﬁnetuning 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 ﬁnetuning of RoBERTa on downstream tasks.
LM ﬁnetuned with KG.
We also compare with existing KG-augmented QA models, QAGNN [8]and GreaseLM [9], which ﬁnetune 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 signiﬁcant 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 ﬁrst key contribution of D RAGON (w.r.t.
existing LM pretraining methods) is that we incorporate
KGs.
We ﬁnd that this signiﬁcantly 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 signiﬁcantly outperformsthe baseline LM (RoBERTa) across all these categories (e.g., +14% accuracy for negation), whichconﬁrms that our joint language-knowledge pretraining boosts reasoning performance.
D RAGONalso consistently outperforms the existing KG-augmented QA models (QAGNN, GreaseLM).
Weﬁnd 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 ﬁnetuning-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 ﬁnd that D RAGON exhibits abilities to extrapolateand perform robust reasoning.
For instance, D RAGON adjusts the entity attention weights and ﬁnalpredictions 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 ﬁnetuning (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 ﬁnd that D RAGON is especially helpful for the following three scenarios.
Downstream tasks with limited data.
In Table 1, we ﬁnd that D RAGON provides signiﬁcant boostsover GreaseLM on downstream tasks with limited ﬁnetuning data available, such as ARC (3K training
[Table] [Table]
(cid:335)(cid:313)(cid:446)(cid:274)
[Table]
(cid:274)(cid:346)(cid:415)(cid:274)(cid:393)(cid:415)(cid:232)(cid:313)(cid:346)(cid:344)(cid:274)(cid:346)(cid:415) (cid:403)(cid:415)(cid:232)(cid:415)(cid:313)(cid:355)(cid:346)
[Table]
(cid:3) (cid:544) (cid:17951) (cid:545) (cid:573)(cid:573)(cid:573)(cid:573) (cid:116)(cid:374)(cid:280)(cid:288)(cid:352)(cid:807)(cid:164)(cid:412)(cid:288)(cid:280)(cid:329)(cid:273)(cid:432)(cid:329)(cid:374)(cid:363)
[Table]
[Figure]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Table]
fusion layers—is increased (“-Ex”).
Increased capacity does not help for the ﬁnetuning-only model (GreaseLM), but helps when pretrained (D RAGON ), suggesting the promise of D RAGON to be further scaled up.
[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 ﬁnetuningdata is used (Table 3).
Here we also see that D RAGON attains signiﬁcant gains over GreaseLM (+5%accuracy on PIQA ), suggesting the improved data-efﬁciency of D RAGON .
Complex downstream tasks.
In Table 1, we ﬁnd 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 ﬁnd that increased capacity does not help for the ﬁnetuning-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 beneﬁcial 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 ﬁrst 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 ﬁnd 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 ﬁnd 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 ﬁnal 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 [32], append them to the main text input, and perform vanilla
MLM pretraining.
We ﬁnd 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 [64, 65, 66, 1], and experts curatevarious knowledge bases for it [67, 17, 68, 69].
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 [70], a widely-used corpus in biomedial LMtraining (e.g., BioBERT [71], PubmedBERT [72]).
It contains the abstracts of biomedical paperson PubMed and has 21GB of text.
For the KG data, we use the Uniﬁed 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 ﬁnetune and evaluate D RAGON on three popular biomedical
NLP and reasoning benchmarks: MedQA-USMLE ( MedQA ) [73], PubMedQA [74], and BioASQ[75].
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 ﬁnetuned 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 signiﬁcant efﬁcacy 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 reﬂect 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] See Appendix A
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] See Appendix A2.
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]3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental 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.3(c)
Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes](d) Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Section 3.14.
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 andSection
(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 identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
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
spent on participant compensation?
[N/A]