On the Theoretical Limitations of Embedding-based Link Prediction

Full paper: arxiv.org/abs/2506.22271

We show a structural limitation which prevents most knowledge graph embeddings from scaling to large graphs. Link prediction is how we fill gaps in knowledge graphs. Most approaches encode queries into low-dimensional embeddings and decode linearly to (much) larger answer spaces, which introduces a bottleneck.

This matters regardless of the encoder you use (GNNs, LLMs, …). We show how it behaves across different task objectives and then break the limitation by adapting a Mixture of Softmaxes from the LLM literature. 🔨 We derive an output layer that can be equiped on most KGE models and find it improves their performance on large graphs at a low parameter cost.

TL;DR: Let’s focus on decoders, not just encoders! 💡

For more visualisations of bottlenecks in output layers (classification and rankings which we also discuss in the paper), check also the website of Andreas Grivas !