# On the Theoretical Limitations of Embedding-based Link Prediction

arXiv: https://arxiv.org/abs/2506.22271 · HTML: https://mlsanity.com/paper/2506.22271/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 38 of its 70 references that are in the map.
- First posted on arXiv 27 Jun 2025.

## Builds on (most cited first, 30 of 38)

- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Modeling Relational Data with Graph Convolutional Networks](https://mlsanity.com/paper/1703.06103/) (2017)
- [Embedding Entities and Relations for Learning and Inference in Knowledge Bases](https://mlsanity.com/paper/1412.6575/) (2014)
- [RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space](https://mlsanity.com/paper/1902.10197/) (2019)
- [Convolutional 2D Knowledge Graph Embeddings](https://mlsanity.com/paper/1707.01476/) (2017)
- [Composition-based Multi-Relational Graph Convolutional Networks](https://mlsanity.com/paper/1911.03082/) (2019)
- [Neural Bellman-Ford Networks: A General Graph Neural Network Framework for Link Prediction](https://mlsanity.com/paper/2106.06935/) (2021)
- [A Review of Relational Machine Learning for Knowledge Graphs](https://mlsanity.com/paper/1503.00759/) (2015)
- [SimplE Embedding for Link Prediction in Knowledge Graphs](https://mlsanity.com/paper/1802.04868/) (2018)
- [Go for a Walk and Arrive at the Answer: Reasoning Over Paths in Knowledge Bases using Reinforcement Learning](https://mlsanity.com/paper/1711.05851/) (2017)
- [Breaking the Softmax Bottleneck: A High-Rank RNN Language Model](https://mlsanity.com/paper/1711.03953/) (2017)
- [Canonical Tensor Decomposition for Knowledge Base Completion](https://mlsanity.com/paper/1806.07297/) (2018)
- [BoxE: A Box Embedding Model for Knowledge Base Completion](https://mlsanity.com/paper/2007.06267/) (2020)
- [Complex Query Answering with Neural Link Predictors](https://mlsanity.com/paper/2011.03459/) (2020)
- [SimKGC: Simple Contrastive Knowledge Graph Completion with Pre-trained Language Models](https://mlsanity.com/paper/2203.02167/) (2022)
- [Knowledge Graph Completion via Complex Tensor Factorization](https://mlsanity.com/paper/1702.06879/) (2017)
- [Message Passing for Hyper-Relational Knowledge Graphs](https://mlsanity.com/paper/2009.10847/) (2020)
- [Bringing Light Into the Dark: A Large-scale Evaluation of Knowledge Graph Embedding Models Under a Unified Framework](https://mlsanity.com/paper/2006.13365/) (2020)
- [Knowledge Base Completion: Baselines Strike Back](https://mlsanity.com/paper/1705.10744/) (2017)
- [Closing the Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/2310.01693/) (2023)
- [OpenBioLink: A benchmarking framework for large-scale biomedical link prediction](https://mlsanity.com/paper/1912.04616/) (2019)
- [Node Embeddings and Exact Low-Rank Representations of Complex Networks](https://mlsanity.com/paper/2006.05592/) (2020)
- [Sigsoftmax: Reanalysis of the Softmax Bottleneck](https://mlsanity.com/paper/1805.10829/) (2018)
- [How to Turn Your Knowledge Graph Embeddings into Generative Models](https://mlsanity.com/paper/2305.15944/) (2023)
- [Understanding the Performance of Knowledge Graph Embeddings in Drug Discovery](https://mlsanity.com/paper/2105.10488/) (2021)
- [On Multi-Relational Link Prediction with Bilinear Models](https://mlsanity.com/paper/1709.04808/) (2017)
- [Low-Rank Softmax Can Have Unargmaxable Classes in Theory but Rarely in Practice](https://mlsanity.com/paper/2203.06462/) (2022)
- [I Know What You Do Not Know: Knowledge Graph Embedding via Co-distillation Learning](https://mlsanity.com/paper/2208.09828/) (2022)
- [Breaking the Softmax Bottleneck via Learnable Monotonic Pointwise Non-linearities](https://mlsanity.com/paper/1902.08077/) (2019)

## Built on by (most cited first, 1 of 1)

- [On the Theoretical Limitations of Embedding-Based Retrieval](https://mlsanity.com/paper/2508.21038/) (2025)

## Abstract (arXiv)

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.