Source-linked AI summary
CoLAKE: Contextualized Language and Knowledge Embedding
Tianxiang Sun, Yunfan Shao, Xipeng Qiu, Qipeng Guo, Yaru Hu, Xuanjing Huang, Zheng Zhang
TL;DR
Existing knowledge-enhanced language models mainly rely on shallow, static, separately pre-trained entity embeddings, leaving deep contextualized knowledge representation underexplored. CoLAKE jointly learns contextualized language and knowledge representations through an extended MLM objective over unified WK graphs, outperforming prior counterparts on most tasks and achieving strong WK graph-completion performance.
Problem
Existing models commonly use separately pre-trained entity embeddings and do not fully capture entities’ rich contextual knowledge-graph information.
Method
CoLAKE jointly pre-trains contextualized language and knowledge representations on heterogeneous word-knowledge graphs using an extended MLM objective.
Results
CoLAKE outperforms previous counterparts on most evaluated tasks and achieves high performance on word-knowledge graph completion in transductive and inductive settings.
Takeaways & Limitations
The results support simultaneously contextualizing language and knowledge representation through the WK graph.
Takeaways & Limitations
The study does not compare with KEPLER because its authors had not released the data split and model.
Abstract
from arXiv · showhide
With the emerging branch of incorporating factual knowledge into pre-trained language models such as BERT, most existing models consider shallow, static, and separately pre-trained entity embeddings, which limits the performance gains of these models. Few works explore the potential of deep contextualized knowledge representation when injecting knowledge. In this paper, we propose the Contextualized Language and Knowledge Embedding (CoLAKE), which jointly learns contextualized representation for both language and knowledge with the extended MLM objective. Instead of injecting only entity embeddings, CoLAKE extracts the knowledge context of an entity from large-scale knowledge bases. To handle the heterogeneity of knowledge context and language context, we integrate them in a unified data structure, word-knowledge graph (WK graph). CoLAKE is pre-trained on large-scale WK graphs with the modified Transformer encoder. We conduct experiments on knowledge-driven tasks, knowledge probing tasks, and language understanding tasks. Experimental results show that CoLAKE outperforms previous counterparts on most of the tasks. Besides, CoLAKE achieves surprisingly high performance on our synthetic task called word-knowledge graph completion, which shows the superiority of simultaneously contextualizing language and knowledge representation.
1 Introduction
CoLAKE jointly contextualizes language and knowledge representations by integrating sentence words with entity-centered knowledge subgraphs in a unified WK graph. It outperforms previous counterparts on most evaluated tasks and shows strong performance on WK graph completion.
- Motivation: Existing knowledge-enhanced language models use separately pre-trained, fixed entity embeddings that incompletely capture entities’ rich knowledge-graph context.These limitations restrict their ability to jointly learn language and knowledge representations.
- Approach: CoLAKE jointly learns contextualized language and knowledge representations in a common representation space.Its extended MLM objective learns both modalities simultaneously.
- Approach: CoLAKE dynamically represents each entity using its surrounding knowledge subgraph together with the current language context.This lets the model access different entity facts as background knowledge for different sentences.
- Approach: The WK graph unifies a fully connected word graph with entity-centered knowledge subgraphs extracted from large knowledge graphs.This heterogeneous structure integrates language context and knowledge context for the model.
- Results: CoLAKE outperforms previous semi-contextualized counterparts on most evaluated tasks and exceeds several knowledge-embedding models on WK graph completion.The graph-completion evaluation covers both transductive and inductive settings.
- Contribution: CoLAKE is characterized as a structure-aware, extensible pre-trained graph neural network.Its three stated properties are simultaneous contextualization, heterogeneous WK-graph integration, and a GNN formulation.
2 Related Work
Prior work developed contextualized language and knowledge representations separately and explored joint models that align words and entities. Knowledge-enhanced language models nevertheless remained limited by static or indirectly generated entity representations.
- Language Representation Learning: Pre-trained language representations evolved from static word vectors to contextualized representations learned with deep neural networks.Contextualized language models have become successful across many NLP tasks.
- Knowledge Representation Learning: Knowledge representation learning maps entities and relations into low-dimensional vectors, with traditional methods learning static embeddings from triplets.More recent methods use subgraphs or paths and report improved knowledge-graph task performance.
- Joint Language and Knowledge Models: Joint language-and-knowledge models map words and entities into compatible spaces because language and knowledge graphs share mutual information.Entity linking is one motivating application requiring compatible word and entity embeddings.
- Joint Language and Knowledge Models: Injecting pre-trained entity embeddings into language models has succeeded, but its knowledge gains are limited by static, inflexible embeddings.KEPLER jointly trains language modeling and knowledge embedding but generates entity embeddings from descriptions rather than directly learning each entity embedding.
3 CoLAKE
CoLAKE integrates language and factual knowledge in unified word-knowledge graphs and jointly contextualizes both through a modified Transformer and extended MLM objective. Its graph construction, masked attention, node-type prediction tasks, and training strategies adapt pre-training to heterogeneous knowledge inputs.
- Overview: CoLAKE pre-trains jointly contextualized language and knowledge representations on structured, unlabeled word-knowledge graphs.The model uses an extended masked language modeling objective over graph nodes.
- Graph Construction: WK graphs unify fully connected word graphs with knowledge sub-graphs centered on entities mentioned in the sentence.Mentions are linked to entities, and surrounding knowledge triplets are extracted and combined with the word graph.
- Model Architecture: The embedding layer combines token, type, and position embeddings while maintaining separate lookup tables for words, entities, and relations.Soft position indices allow repeated positions and keep tokens in the same triplet continuous.
- Model Architecture: Masked multi-head self-attention restricts each node to information from connected neighbors, making the Transformer reflect WK-graph structure.Each layer gathers information from 1-hop neighbors, similarly to a graph attention network.
- Pre-Training Objective: Extended MLM randomly masks 15% of graph nodes and predicts words, entities, or relations according to node type.Word masking uses linguistic and knowledge context; entity masking aligns words with entities and contextualizes entities; relation masking supports relation prediction.
- Training: Anchor-neighbor discarding mitigates trivial entity prediction, while mixed CPU-GPU training and negative sampling address the large entity vocabulary.Neighbors of anchor nodes are discarded half the time, entity embeddings are updated in shared CPU memory, and prediction uses one positive plus k negative entities.
4 Experiments
CoLAKE is evaluated on knowledge-driven, knowledge-probing, language-understanding, and word-knowledge graph completion tasks. It generally improves factual-knowledge performance while retaining comparable language-understanding results.
- Evaluation setup: CoLAKE is evaluated on knowledge-driven tasks, knowledge probing, GLUE language-understanding tasks, and word-knowledge graph completion.The experiments include entity typing, relation extraction, LAMA and LAMA-UHN probing, GLUE, and transductive or inductive graph completion.
- Pre-training setup: CoLAKE uses Wikipedia aligned to Wikidata5M to construct WK-graph training samples and pre-trains a Transformer encoder with entity and relation embeddings.The pre-training data contain 21M fact triplets, 26M training samples, 3,085,345 entity embeddings, and 822 relation embeddings.
- Knowledge probing: CoLAKE outperforms RoBERTaBASE by a significant margin on LAMA and improves over K-Adapter by 2.5% on LAMA-Google-RE and 1.2% on LAMA-UHN-Google-RE.K-Adapter uses RoBERTaLARGE and a subset of T-REx, which may contribute to its superiority on LAMA-T-REx and LAMA-UHN-T-REx.
- Language understanding: On GLUE, CoLAKE is slightly degraded from RoBERTa but improves over KEPLER by 1.4% on average.The authors conclude that CoLAKE achieves comparable language-understanding results while improving knowledge-required tasks.
- Word-knowledge graph completion: The transductive setting exposes entities during training, whereas the inductive setting evaluates an unknown entity whose neighboring entities were seen during training.This distinction tests whether representations can be formed for unseen entities from their graph neighborhoods.
- Word-knowledge graph completion: CoLAKE outperforms other models by a large margin on word-knowledge graph completion in both transductive and inductive settings.Unlike description-based DKRL and KEPLER, CoLAKE generates entity embeddings from neighboring entities.
5 Conclusion
CoLAKE jointly contextualizes language and knowledge through a unified word-knowledge graph. Its results support knowledge-task effectiveness and suggest applications in graph-aware relation extraction and graph-to-text evaluation.
- Conclusion: CoLAKE jointly learns contextualized language and knowledge representations in a unified word-knowledge graph.The paper characterizes CoLAKE as a structure-aware and inductive graph neural network.
- Conclusion: The word-knowledge graph completion results motivate potential uses for denoising distantly annotated relation-extraction samples and measuring graph-to-text template quality.The paper leaves these applications for future work.