Source-linked AI summary
JAKET: Joint Pre-training of Knowledge Graph and Language Understanding
Donghan Yu, Chenguang Zhu, Yiming Yang, Michael Zeng
TL;DR
Pre-trained language models struggle to incorporate structured world knowledge, while existing approaches have difficulty aligning representations, avoiding over-parameterization, and adapting to unseen graphs. JAKET jointly pre-trains knowledge and language modules that exchange structural and contextual representations. It outperforms baseline methods across several knowledge-aware language-understanding tasks and is designed to adapt to unseen knowledge graphs.
Problem
Existing language models struggle with world knowledge, while prior knowledge-integration methods face alignment, over-parameterization, and unseen-graph adaptation challenges.
Method
JAKET jointly pre-trains a knowledge module and language module that provide structure-aware entity embeddings and context-aware initial embeddings for graph entities and relations.
Results
JAKET achieves the best performance against strong baselines across few-shot relation classification, question answering, and entity classification, including previously unseen knowledge graphs.
Takeaways & Limitations
The framework can quickly adapt to unseen knowledge graphs in new domains while leveraging knowledge for language understanding.
Takeaways & Limitations
The framework does not consider discontinuous entity mentions.
Abstract
from arXiv · showhide
Knowledge graphs (KGs) contain rich information about world knowledge, entities and relations. Thus, they can be great supplements to existing pre-trained language models. However, it remains a challenge to efficiently integrate information from KG into language modeling. And the understanding of a knowledge graph requires related context. We propose a novel joint pre-training framework, JAKET, to model both the knowledge graph and language. The knowledge module and language module provide essential information to mutually assist each other: the knowledge module produces embeddings for entities in text while the language module generates context-aware initial embeddings for entities and relations in the graph. Our design enables the pre-trained model to easily adapt to unseen knowledge graphs in new domains. Experimental results on several knowledge-aware NLP tasks show that our proposed framework achieves superior performance by effectively leveraging knowledge in language understanding.
1 INTRODUCTION
JAKET addresses the difficulty of integrating structured knowledge into language models while preserving contextual understanding. It jointly pre-trains knowledge and language modules, enabling adaptation to unseen knowledge graphs and strong performance across knowledge-aware tasks.
- Pre-trained language models capture linguistic patterns but struggle with world knowledge, concepts, and relations important for language understanding.
- Knowledge graphs represent entities and relations structurally, providing a precise alternative to learning facts only from repeated textual examples.
- Existing knowledge-integrated language models face representation-alignment, over-parameterization, and adaptation problems on unseen knowledge graphs.
- JAKET jointly pre-trains a knowledge module and language module that mutually provide structure-aware and context-aware representations.The knowledge module supplies entity embeddings for language modeling, while the language module initializes graph entities and relations from descriptive text.
- A two-step language module resolves cyclic dependence, and an entity context embedding memory speeds pre-training by 15x.
- JAKET achieves the best performance against strong baselines across few-shot relation classification, question answering, and entity classification, including unseen-graph settings.
2 RELATED WORK
Prior work integrates knowledge into language models through explicit entity representations, implicit knowledge modeling, or task-specific combinations of graph networks and language models. JAKET instead jointly pre-trains graph representation and language modeling.
- Knowledge integration approaches explicitly inject entity representations that are either externally pre-computed or learned as model parameters.
- Other approaches implicitly model knowledge through entity-level masking, entity replacement prediction, or knowledge-embedding losses.
- Some methods combine graph neural networks with pre-trained language models for specific tasks, but do not integrate knowledge into language modeling.
3 METHOD
JAKET jointly models knowledge graphs and text through mutually assisting knowledge and language modules. Its decomposed language architecture, entity-description initialization, and periodically updated memory address cyclic computation and pre-training cost.
- Framework: JAKET defines a framework with knowledge and language modules that mutually provide information for semantic analysis.The knowledge module supplies structure-aware entity embeddings, while the language module supplies contextual initial embeddings for graph entities and relations.
- Knowledge module: The knowledge module uses a multi-relational graph attention network to compute structure-aware entity representations from entity and relation embeddings.It composes entity and relation embeddings and uses neighborhood sampling because updating all entities in a large graph is intractable.
- Language module: The language module uses RoBERTa to learn context-aware representations and initializes graph entities and relations from their descriptive text.Entity descriptions provide contextual embeddings, while relation embeddings are generated similarly from relation descriptions.
- Language module: The language module uses RoBERTa-base, while this work does not consider discontinuous entity mentions.The limitation concerns the representation of entity mentions in the language module.
- Solving cyclic dependency: LM1 and LM2 resolve cyclic dependency by generating initial graph embeddings with LM1, integrating knowledge-module outputs, and producing final text representations with LM2.The computation proceeds from text through LM1, graph initialization and propagation, fusion at entity mentions, and final processing by LM2.
- Entity context embedding memory: The entity context embedding memory pre-computes initial entity embeddings and periodically refreshes them to reduce language-module computation and maintain consistency with updated parameters.The memory avoids computing description embeddings for all involved entities on the fly, and the design speeds pre-training by about 15x.
4 EXPERIMENT
JAKET is evaluated on few-shot relation classification, KGQA, and entity classification, including settings involving unseen or incomplete knowledge graphs. It outperforms the reported baselines across these evaluations, with gains attributed to joint use of language context, graph structure, and pre-training.
- Experimental Setup: The experiments compare JAKET with RoBERTa, RoBERTa+GNN, and RoBERTa+GNN+M across few-shot classification, KGQA, and unseen-graph entity classification.
- Few-shot Relation Classification: JAKET achieves the best results in all three few-shot relation-classification settings.
- KGQA: JAKET improves MetaQA accuracy by 0.6% to 2.5% under both full and 50% knowledge-graph settings.
- Entity Classification: JAKET achieves the best performance under all entity-classification settings, while its advantage over RoBERTa+GNN+M grows with less training data.
5 CONCLUSION
The paper concludes that JAKET jointly pre-trains knowledge-graph and language understanding, adapts to unseen knowledge graphs, and outperforms baselines on several knowledge-aware NLU tasks.
- JAKET jointly pre-trains models for knowledge-graph and language understanding, with its knowledge and language modules mutually assisting each other.
- JAKET can quickly adapt to unseen knowledge graphs in new domains after pre-training.
- The entity context embedding memory speeds pre-training by 15x.
- JAKET outperforms baseline methods on few-shot relation classification, KGQA, and entity classification.
A.1 IMPLEMENTATION DETAILS
The implementation specifies the hidden-state size, attention configuration, batch sizes, training duration, and AdamW optimization settings used during pre-training.
- The knowledge module uses 768-dimensional hidden states and 8 attention heads, matching RoBERTaBASE’s hidden-state dimension.
- Pre-training uses text batches of 1024 sequences of length 512 and knowledge-graph entity batches of 16,384.
- JAKET is trained for 8 epochs with AdamW, β1 = 0.9, β2 = 0.999, ϵ = 1e-8, and weight decay 0.01.