Title: Breaking Rank Bottlenecks in Knowledge Graph Embeddings
Primary area: learning on graphs and other geometries & topologies
Keywords: Knowledge Graph Embeddings, Rank Bottlenecks, Representation Learning

Abstract: Many knowledge graph embedding (KGE) models for link prediction use powerful encoders. However, they often rely on a simple hidden vector-matrix multiplication to score subject-relation queries against candidate object entities. When the number of entities is larger than the model's embedding dimension, which is often the case in practice by several orders of magnitude, we have a linear output layer with a rank bottleneck. Such bottlenecked layers limit model expressivity.
We investigate both theoretically and empirically how rank bottlenecks affect KGEs. We find that, by limiting the set of feasible predictions, rank bottlenecks hurt the ranking accuracy and distribution fidelity of scores. Inspired by the language modelling literature, we propose KGE-MoS, a mixture-based output layer to break rank bottlenecks in many KGEs. Our experiments show that KGE-MoS improves ranking performance of KGE models on large-scale datasets at a low parameter cost.

TL;DR: Knowledge Graph Embeddings models for link prediction often suffer from “rank bottlenecks” in their output layers resulting in limited prediction accuracy and miscalibrated probabilities. We address this by proposing a mixture-based output layer.