Separate-and-Aggregate: A Transformer-based Patch Refinement Model for Knowledge Graph Completion

TL;DR

This paper introduces PatReFormer, a Transformer-based model for knowledge graph completion that segments embeddings into patches and uses cross-attention to improve understanding and prediction of missing facts.

Contribution

The paper proposes a novel Patch Refinement Model using Transformer-based cross-attention for enhanced knowledge graph completion.

Findings

Significant performance improvements on WN18RR, FB15k-237, YAGO37, and DB100K benchmarks.

Better capture of KG information with larger relation embedding dimensions.

Enhanced performance on complex relation types.

Abstract

Knowledge graph completion (KGC) is the task of inferencing missing facts from any given knowledge graphs (KG). Previous KGC methods typically represent knowledge graph entities and relations as trainable continuous embeddings and fuse the embeddings of the entity $h$ (or $t$) and relation $r$ into hidden representations of query $(h, r, ?)$ (or $(?, r, t$)) to approximate the missing entities. To achieve this, they either use shallow linear transformations or deep convolutional modules. However, the linear transformations suffer from the expressiveness issue while the deep convolutional modules introduce unnecessary inductive bias, which could potentially degrade the model performance. Thus, we propose a novel Transformer-based Patch Refinement Model (PatReFormer) for KGC. PatReFormer first segments the embedding into a sequence of patches and then employs cross-attention modules to allow bi-directional embedding feature interaction between the entities and relations, leading to a better understanding of the underlying KG. We conduct experiments on four popular KGC benchmarks, WN18RR, FB15k-237, YAGO37 and DB100K. The experimental results show significant performance improvement from existing KGC methods on standard KGC evaluation metrics, e.g., MRR and H@n. Our analysis first verifies the effectiveness of our model design choices in PatReFormer. We then find that PatReFormer can better capture KG information from a large relation embedding dimension. Finally, we demonstrate that the strength of PatReFormer is at complex relation types, compared to other KGC models