On the Theoretical Limitations of Embedding-based Link Prediction

Neural networks often map low-dimensional embeddings to high-dimensional output spaces. Usually, the output layer is linear, which can create a rank bottleneck that limits the functions a model can represent. Such bottlenecks are ubiquitous in link prediction models, such as knowledge graph embeddings (KGEs), as the output space of entities can be orders of magnitude larger than the embedding dimension. We investigate how rank bottlenecks limit model expressivity for fitting the training data. While previous work focused on sufficient bounds on the embedding dimension required for specific KGEs, we show necessary bounds for all KGEs with a linear output layer, which grow with graph size and connectivity. We also consider a non-linear output layer using mixtures to break the bottleneck without significant parameter overhead. Empirically, we show that models using this non-linear layer improve in ranking performance and probabilistic fit for large and dense datasets at a low parameter cost, as predicted by our theory. Our work reveals how linear output layers limit KGEs and motivates non-linear alternatives for scaling to large and dense graphs.

Artificial Intelligence is increasingly used to infer missing facts from known ones, for example, suggesting which drugs might treat a disease, or which products a customer might want. Many of these "link prediction" systems work by turning every entity (a drug, a person, a concept) into an embedding, a short list of numbers that captures its properties, and then scoring possible facts using these numbers. To stay fast and scalable even when there are millions of candidate answers, the final scoring step is kept deliberately simple. In this paper, we show that this simplicity carries a hidden cost: it imposes a hard mathematical limitation on what the model can represent, one that no amount of training can overcome, and this limitation grows more severe as the dataset gets larger. We calculate how large the embeddings would need to be to capture every pattern, and find these sizes are impractical for realistic datasets. We then propose a small, inexpensive change to the final step that lifts the limitation and improves predictions on large datasets. Our results suggest that link prediction research should keep exploring such directions for their final scoring steps.