[Paper Review] Multi-task Pre-training Language Model for Semantic Network Completion

Da Li, Sen Yang+4 more|arXiv (Cornell University)|Jan 13, 2022

Advanced Graph Neural Networks4 citations

TL;DR

This paper proposes LP-BERT, a multi-task pre-training language model for knowledge graph completion that improves link prediction by jointly learning entity, relation, and context representations. By introducing masked entity and relation modeling tasks and a batch-wise triple-style negative sampling strategy, LP-BERT achieves state-of-the-art performance, improving Hits@10 by 5% on WN18RR and reaching 100% on UMLS.

ABSTRACT

Semantic networks, such as the knowledge graph, can represent the knowledge leveraging the graph structure. Although the knowledge graph shows promising values in natural language processing, it suffers from incompleteness. This paper focuses on knowledge graph completion by predicting linkage between entities, which is a fundamental yet critical task. Semantic matching is a potential solution as it can deal with unseen entities, which the translational distance based methods struggle with. However, to achieve competitive performance as translational distance based methods, semantic matching based methods require large-scale datasets for the training purpose, which are typically unavailable in practical settings. Therefore, we employ the language model and introduce a novel knowledge graph architecture named LP-BERT, which contains two main stages: multi-task pre-training and knowledge graph fine-tuning. In the pre-training phase, three tasks are taken to drive the model to learn the relationship from triples by predicting either entities or relations. While in the fine-tuning phase, inspired by contrastive learning, we design a triple-style negative sampling in a batch, which greatly increases the proportion of negative sampling while keeping the training time almost unchanged. Furthermore, we propose a new data augmentation method utilizing the inverse relationship of triples to improve the performance and robustness of the model. To demonstrate the effectiveness of our method, we conduct extensive experiments on three widely-used datasets, WN18RR, FB15k-237, and UMLS. The experimental results demonstrate the superiority of our methods, and our approach achieves state-of-the-art results on WN18RR and FB15k-237 datasets. Significantly, Hits@10 indicator is improved by 5% from previous state-of-the-art result on the WN18RR dataset while reaching 100% on the UMLS dataset.

Motivation & Objective

• Address the challenge of knowledge graph incompleteness by improving link prediction for unseen entities.
• Overcome the limitations of translational distance models that fail on out-of-vocabulary entities.
• Enhance semantic matching-based models, which suffer from low negative sampling ratios and poor generalization.
• Develop a pre-training framework that effectively captures relational and semantic knowledge from structured and unstructured data.
• Improve model robustness and performance through novel data augmentation using inverse triples.

Proposed method

• Introduce a multi-task pre-training phase with three tasks: masked language modeling (MLM), masked entity modeling (MEM), and masked relation modeling (MRM) to learn contextual, entity, and relational representations.
• Design a triple-style negative sampling strategy within a batch, inspired by contrastive learning, to increase negative sampling proportion without increasing training time.
• Propose a data augmentation method using inverse relationships of triples to increase training sample diversity and model robustness.
• Fine-tune the pre-trained model on knowledge graph completion using a contrastive learning-inspired objective to improve generalization.
• Leverage BERT-style transformer architecture to encode textual descriptions of entities and relations for semantic matching.
• Combine pre-training on large-scale knowledge graph triples with fine-tuning on downstream link prediction tasks.

Experimental results

Research questions

• RQ1 Can multi-task pre-training with entity and relation masking improve zero-shot and few-shot link prediction in knowledge graphs?
• RQ2 How does batch-wise triple-style negative sampling affect model convergence and performance compared to standard negative sampling?
• RQ3 To what extent does data augmentation using inverse triples improve model generalization and robustness?
• RQ4 Can a semantic matching-based model outperform translational distance models in link prediction, especially for unseen entities?
• RQ5 What is the impact of pre-training on downstream knowledge graph completion performance across diverse benchmarks?

Key findings

• LP-BERT achieves state-of-the-art performance on the WN18RR and FB15k-237 datasets, outperforming existing methods in link prediction.
• On WN18RR, LP-BERT improves Hits@10 by 5% compared to the previous state-of-the-art result, demonstrating significant performance gains.
• On the UMLS dataset, LP-BERT achieves 100% Hits@10, indicating perfect performance on this benchmark.
• The case study shows that LP-BERT consistently ranks correct entities higher (average rank 1.4) than MLM-only pre-training (rank 3.4) and no pre-training (rank 8.2).
• The proposed batch-wise negative sampling increases negative sampling ratio substantially while maintaining training efficiency.
• The inverse triple data augmentation method enhances model robustness and generalization, especially in low-resource settings.