Source-linked AI summary
Deep Bidirectional Language-Knowledge Graph Pretraining
Michihiro Yasunaga, Antoine Bosselut, Hongyu Ren, Xikun Zhang, Christopher D Manning, Percy Liang, Jure Leskovec
TL;DR
Existing methods do not deeply fuse text and knowledge graphs during large-scale pretraining, despite their complementary information for reasoning. DRAGON jointly pretrains a bidirectionally fused model on aligned text–KG pairs using masked language modeling and link prediction, outperforming existing models across general and biomedical tasks, especially complex and low-resource QA.
Problem
Existing text-and-KG methods lack deep fusion during large-scale pretraining, limiting fully joint representations and reasoning over the two modalities.
Method
DRAGON retrieves local KG subgraphs for text segments, bidirectionally fuses both modalities, and jointly trains masked language modeling with KG link prediction.
Results
DRAGON outperforms existing LM and LM+KG models across general and biomedical downstream tasks, including complex and low-resource question answering.
Takeaways & Limitations
Deep bidirectional self-supervision over text and KG yields improved language-knowledge representations for complex reasoning across domains.
Takeaways & Limitations
DRAGON is currently an encoder model and does not perform language generation.
Abstract
from arXiv · showhide
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.
1 Introduction
DRAGON addresses the open problem of deeply combining text and knowledge graphs during pretraining. It uses bidirectional fusion and joint self-supervision to improve language-knowledge representations across general and biomedical NLP tasks.
- 1 Introduction: Existing approaches typically combine text and knowledge graphs shallowly or unidirectionally, limiting fully joint reasoning over both modalities.Other bidirectional models focus on downstream finetuning rather than large-scale self-supervised pretraining.
- 1 Introduction: DRAGON performs deeply bidirectional, self-supervised pretraining of a language-knowledge model from text and knowledge graphs.Its core components are a cross-modal model that fuses text and KG representations and a bidirectional objective for joint reasoning.
- 1 Introduction: DRAGON creates training inputs by sampling text segments and extracting relevant local KG subgraphs through entity linking.The broader approach uses knowledge bases as background knowledge to ground reasoning about entities and facts.
- 1 Introduction: +8% absolute accuracy on average is reported over RoBERTa on general-domain commonsense reasoning tasks.The evaluated tasks include CSQA, OBQA, RiddleSense, and HellaSwag.
- 1 Introduction: +3% accuracy is reported over BioLinkBERT on biomedical tasks, alongside new state-of-the-art results on BioNLP tasks.The biomedical setup uses PubMed and UMLS, while the general setup uses BookCorpus and ConceptNet.
2 Deep Bidirectional Language-Knowledge Graph Pretraining (DRAGON)
DRAGON pretrains a language-knowledge model by deeply fusing text with relevant local KG subgraphs and jointly optimizing self-supervised objectives for both modalities.
- 2.1 Input representation: DRAGON pairs sampled text segments with relevant local KG subgraphs and uses a cross-modal encoder to fuse them bidirectionally.The encoder exchanges information across text tokens and KG nodes through multiple sequence-graph layers.
- 2.1 Input representation: KG retrieval links text entity mentions to KG nodes, adds their 2-hop bridge nodes, and includes edges spanning the retrieved nodes.The resulting local KG forms the KG component of each text-KG input instance.
- 2.2 Cross-modal encoder: GreaseLM combines Transformer language-model layers with graph neural networks and modality-interaction layers to produce joint token and node representations.An interaction token and interaction node provide pooling points and an interface between modalities.
- 2.3 Pretraining objective: DRAGON unifies masked language modeling and KG link prediction so text and KG mutually inform one another during pretraining.MLM predicts masked text tokens using contextualized token vectors, while link prediction scores held-out KG triplets against negative samples.
- 2.3 Pretraining objective: The joint objective is L = LMLM + LLinkPred, and the paper reports that joint training outperforms using either objective alone.MLM encourages reasoning with textual context and structured KG knowledge; link prediction uses KG structure for reasoning about missing edges.
3 Experiments: General domain
DRAGON combines text and relevant knowledge-graph subgraphs through bidirectional pretraining and improves commonsense reasoning across general-domain benchmarks. Its gains are strongest for complex reasoning, limited-data settings, and configurations that benefit from increased capacity.
- Experimental setup: DRAGON uses BookCorpus and ConceptNet, pairing sampled text segments with relevant KG subgraphs for general-domain pretraining.The implementation uses text segments up to 512 tokens and relevant subgraphs up to 200 nodes.
- Overall results: DRAGON consistently outperforms RoBERTa, QAGNN, and GreaseLM across nine commonsense reasoning tasks.On OBQA, it achieves a +7% absolute accuracy boost over RoBERTa and +5% over GreaseLM.
- Complex reasoning: +14% accuracy for negation is one example of DRAGON’s gains over RoBERTa across complex-reasoning categories.The categories include negation, conjunction, hedges, prepositional phrases, and entity mentions; DRAGON also consistently outperforms QAGNN and GreaseLM.
- Limited-data settings: DRAGON’s gains are especially substantial on limited-data tasks, including +4% on ARC and Riddle, +5% on OBQA, and +5% on PIQA with 10% of finetuning data.These results suggest improved downstream data efficiency.
- Complex tasks: +1.8% on CosmosQA contrasts with GreaseLM’s +0.1% over RoBERTa, while DRAGON also extrapolates to harder questions than those seen during training.The complex tasks involve longer contexts, more entities, and larger local KGs.
- Capacity and design: Increasing text-KG fusion layers helps DRAGON but not finetuning-only GreaseLM, indicating that pretraining makes added capacity useful.The compared capacity increases the number of fusion layers from 5 to 7.
- Capacity and design: +5% accuracy on OBQA results from jointly using MLM and LinkPred instead of either objective alone.Ablations also show benefits from bidirectional cross-modal fusion and preserving KG graph structure rather than converting triples into sentences.
4 Experiments: Biomedical domain
DRAGON is pretrained on biomedical text and the UMLS knowledge graph, then evaluated on three biomedical NLP and reasoning benchmarks. It outperforms existing biomedical language and KG-augmented models across tasks, achieving new state-of-the-art performance.
- Pretraining setup: DRAGON uses PubMed text and the UMLS knowledge graph for biomedical pretraining.PubMed contains 21GB of text, while UMLS has 300K nodes and 1M edges.
- Downstream evaluation tasks: The evaluation covers MedQA-USMLE, PubMedQA, and BioASQ.These are presented as popular biomedical NLP and reasoning benchmarks.
- Baselines: DRAGON is compared with BioLinkBERT, QAGNN, and GreaseLM.BioLinkBERT is the vanilla language-model baseline; QAGNN and GreaseLM are KG-finetuned models seeded with BioLinkBERT.
- Results: +3% absolute accuracy over BioLinkBERT and +2% over GreaseLM on MedQA.Across biomedical tasks, DRAGON outperforms existing biomedical language models and KG-augmented QA models, achieving new state-of-the-art performance.
- Results: DRAGON’s biomedical gains, combined with general-domain results, support its use across domains with different text, KGs, and seed language models.The authors describe this as evidence of domain-generality and effective pretraining across domains.
5 Conclusion
The paper presents DRAGON as a deeply bidirectional, self-supervised language-knowledge pretraining method. It reports improvements across general and biomedical NLP tasks, including complex reasoning, while identifying language generation as an unresolved limitation.
- Conclusion: DRAGON learns a deeply bidirectional language-knowledge model from text and knowledge graphs at scale.It is presented as a self-supervised pretraining method.
- Conclusion: DRAGON outperforms existing language models and KG-augmented models across general and biomedical NLP tasks.The reported gains extend across domains and task types.
- Conclusion: DRAGON performs strongly on complex reasoning involving long contexts or multi-step reasoning.The conclusion highlights these reasoning settings specifically.
- Limitations: DRAGON is currently an encoder model and does not perform language generation.The paper identifies extending DRAGON to generation as future research.
A Ethics, limitations and risks
The paper identifies risks from inherited biases and toxic behaviors in the pretrained language models and knowledge graphs used by DRAGON. It also cautions against applying the MedQA-USMLE evaluation model to real-world clinical prediction.
- Ethical risks: DRAGON may reflect biases and toxic behaviors inherited from its language models and knowledge graphs.The paper specifically mentions demographic biases, toxic outputs, and stereotypes in ConceptNet.
- Ethical risks: ConceptNet may encode stereotypes that influence conclusions outside standard benchmarks.The authors recommend applying precautions used for other language models and noisy knowledge resources.
- Clinical use: The authors do not encourage using DRAGON for real-world clinical prediction.This caution is tied specifically to the MedQA-USMLE evaluation.
B.1 KG retrieval
For each text segment, DRAGON retrieves a relevant local knowledge-graph subgraph through entity linking and two-hop bridge expansion, with pruning when necessary.
- KG retrieval: Entity linking maps mentions in each text segment to initial knowledge-graph nodes.The procedure uses the spaCy entity linker.
- KG retrieval: Two-hop bridge entities between linked entities are added to form the retrieved node set.This expands the local graph beyond directly linked entities.
- KG retrieval: The retrieved node set is randomly pruned when it exceeds 200 nodes.This bounds the size of the local KG input.
B.2 Graph initialization
Graph initialization differs by domain while retaining entity-specific representations. When no entity is linked, DRAGON backs off to text-only representations.
- Graph initialization: ConceptNet entity embeddings are initialized by converting KG triplets into templated sentences and encoding them with BERT-Large.Entity mention representations are collected from sentences containing each entity.
- Graph initialization: UMLS entity embeddings are initialized from pooled entity-name token outputs produced by BioLinkBERT.
- Graph initialization: For fewer than 1% of inputs without linked entities, DRAGON uses a zero-initialized dummy node and only text-side representations.The graph propagates no information in this fallback case.
B.3 Hyperparameters
Table 7 presents the hyperparameter settings used for models and experiments.
- Hyperparameters: Table 7 reports hyperparameter settings for the models and experiments.
B.4 Downstream evaluation tasks
The downstream evaluation covers nine general-domain commonsense benchmarks and three biomedical NLP and reasoning benchmarks. These datasets span multiple-choice question answering, commonsense reasoning, science, biomedical understanding, long narratives, and complex reasoning settings.
- General-domain tasks: The general-domain evaluation uses nine commonsense reasoning benchmarks.
- General-domain tasks: CommonsenseQA, OpenbookQA, RiddleSense, ARC, CosmosQA, HellaSwag, PIQA, and SIQA provide varied multiple-choice reasoning settings.The cited datasets cover commonsense, elementary science, riddles, science exams, narratives, events, physics, and social interaction.
- General-domain tasks: PIQA tests physics reasoning about objects with 20K questions.
- General-domain tasks: aNLI is a 2-way multiple-choice abductive commonsense reasoning task with 170K questions.The experiments use its original data splits.
- Biomedical tasks: The biomedical evaluation uses three biomedical NLP and reasoning benchmarks.
- Biomedical tasks: MedQA, PubMedQA, and BioASQ contain 12,723, 1,000, and 885 questions, respectively, with multiple-choice biomedical reasoning formats.MedQA is 4-way, PubMedQA is 3-way, and BioASQ is 2-way multiple choice.
C Additional Experimental Results
Additional experiments evaluate DRAGON on ConceptNet KG link prediction. DRAGON improves over DistMult while using retrieved text in addition to the KG.
- Additional results: Table 9 reports KG link prediction performance on ConceptNet.
- Additional results: DRAGON performs ConceptNet link prediction using retrieved text alongside the knowledge graph.
- Additional results: DRAGON achieves improved ConceptNet link prediction performance compared with DistMult, which does not use text.