[Paper Review] Deep Bidirectional Language-Knowledge Graph Pretraining

Michihiro Yasunaga, Antoine Bosselut+5 more|arXiv (Cornell University)|Oct 17, 2022

Topic Modeling87 citations

TL;DR

Dragon pretrains a deeply joint language-knowledge model from text and KG, combining MLM and KG link prediction to achieve strong gains across general and biomedical QA tasks, including +10% on complex reasoning and +8% in low-resource settings.

ABSTRACT

Pretraining a language model (LM) on text has been shown to help various downstream NLP tasks. Recent works show that a knowledge graph (KG) can complement text data, offering structured background knowledge that provides a useful scaffold for reasoning. However, these works are not pretrained to learn a deep fusion of the two modalities at scale, limiting the potential to acquire fully joint representations of text and KG. Here we propose DRAGON (Deep Bidirectional Language-Knowledge Graph Pretraining), a self-supervised approach to pretraining a deeply joint language-knowledge foundation model from text and KG at scale. Specifically, our model takes pairs of text segments and relevant KG subgraphs as input and bidirectionally fuses information from both modalities. We pretrain this model by unifying two self-supervised reasoning tasks, masked language modeling and KG link prediction. DRAGON outperforms existing LM and LM+KG models on diverse downstream tasks including question answering across general and biomedical domains, with +5% absolute gain on average. In particular, DRAGON achieves notable performance on complex reasoning about language and knowledge (+10% on questions involving long contexts or multi-step reasoning) and low-resource QA (+8% on OBQA and RiddleSense), and new state-of-the-art results on various BioNLP tasks. Our code and trained models are available at https://github.com/michiyasunaga/dragon.

Motivation & Objective

• Motivate learning deeply interactive representations by jointly modeling text and knowledge graphs (KGs).
• Develop a scalable, self-supervised pretraining objective that fuses text and KG representations bidirectionally.
• Evaluate the approach across general-domain and biomedical tasks to demonstrate improved reasoning and data efficiency.
• Show that joint pretraining with MLM and KG link prediction yields stronger performance than single-objective or non-pretrained baselines.

Proposed method

• Construct aligned (text segment, local KG) pairs by sampling text and retrieving relevant KG subgraphs via entity linking and 2-hop expansion.
• Use a cross-modal encoder (GreaseLM) to bidirectionally fuse text token representations and KG node representations across multiple fusion layers.
• Pretrain with a joint objective combining masked language modeling (MLM) and KG link prediction (LinkPred).
• Evaluate different KG scoring functions (DistMult, TransE, RotatE) and modality interaction strategies through ablations.
• Finetune on downstream tasks by pooling fused text-KG representations for classification or QA tasks.

Experimental results

Research questions

• RQ1 Can a deeply bidirectional, self-supervised pretraining regime effectively fuse text and KG information at scale?
• RQ2 Do joint MLM and KG link prediction objectives yield better joint representations than single-objective pretraining?
• RQ3 How does Dragon perform on commonsense reasoning and biomedically oriented QA tasks compared to LM-only and KG-augmented baselines?
• RQ4 Is the approach more data-efficient and scalable when KG structures are leveraged during pretraining?

Key findings

• Dragon consistently outperforms RoBERTa and KG-augmented baselines (QAGNN, GreaseLM) across nine commonsense reasoning tasks, e.g., on OBQA, CosmosQA, and HellaSwag.
• The model shows substantial gains on tasks requiring complex reasoning (e.g., +10% on multi-step/negation/long-context questions) and low-resource QA (+8% in OBQA and RiddleSense).
• In the general-domain setup, Dragon achieves about +5% absolute improvement on average over baselines.
• In the biomedical domain, Dragon sets new state-of-the-art on BioNLP tasks with notable improvements in MedQA and PubMedQA.
• Ablations show joint MLM + LinkPred outperforming either objective alone, and bidirectional cross-modal fusion outperforming end-concatenation or KG-as-sentences baselines.
• Increasing model capacity helps Dragon but not the finetuning-only baselines, indicating data-efficient scaling with pretraining.