🤖 AI

On the Theoretical Limitations of Embedding-based Link Prediction

This paper demonstrates that linear output layers in knowledge graph embedding models create rank bottlenecks that limit expressivity as graph size and connectivity increase, and proposes a parameter-efficient non-linear mixture-based output layer that theoretically and empirically overcomes these limitations to improve performance on large, dense datasets.

Original authors: Samy Badreddine, Emile van Krieken, Luciano Serafini

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Samy Badreddine, Emile van Krieken, Luciano Serafini

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Picture: The "Too Small Table" Problem

Imagine you are trying to organize a massive library of books (the Knowledge Graph). You want a computer to predict which books go together. To do this, the computer gives every book and every relationship a "ID card" (an embedding) that is a short list of numbers.

Usually, these ID cards are short and simple (low-dimensional), like a 10-digit phone number. But the library has millions of books (high-dimensional output space).

The paper argues that most current computer models try to match these short ID cards to the massive library using a linear output layer. Think of this layer as a small, flat table where you try to lay out all the books.

The Problem: If you have a million books but your table is only big enough for 100 items, you physically cannot lay them all out in the right order. No matter how smart your ID cards are, the table is too small to hold all the possible arrangements. The paper calls this a "Rank Bottleneck." It's like trying to fit a 3D sculpture onto a 2D piece of paper; you lose information and can't represent the full shape.

The Three Ways We Try to "Read" the Library

The authors look at three different ways we try to use these models, and show that the "small table" breaks all of them when the library gets big:

Ranking (Who is #1?): We want to know which book is the best match.
- The Limit: If the library is huge, the small table can't create enough unique "heights" to rank every single book correctly. Some books will always be stuck in the wrong order.
Sign Reconstruction (Yes or No?): We want to know if a book belongs in a category (True/False).
- The Limit: The small table can't draw enough distinct "Yes" and "No" zones. It's like trying to draw a complex map with only two colors; you can't show the details.
Distribution (How likely is it?): We want to know the exact probability of a match.
- The Limit: The small table forces the probabilities into a rigid, straight line. Real life is curved and complex. The model can't bend the probabilities to fit the truth.

The Theory: The authors did some math to prove that to fix this with the old "small table" method, you would need to make the ID cards as long as the number of books in the library. For a library with a million books, your ID cards would need to be a million numbers long. This is impossible to train and use in practice.

The Solution: The "Pop-Up Book" (KGE-MOS)

Since we can't make the ID cards huge (it's too expensive), the authors propose a new way to use the table. They introduce KGE-MOS (Mixture of Softmaxes).

The Analogy:
Instead of one small, flat table, imagine a Pop-Up Book.

The Old Way: You have one flat page. You can only show one arrangement of books.
The New Way (KGE-MOS): You have a book with several layers (mixtures). Depending on which book you are looking at, the page "pops up" into a different 3D shape.

By mixing several different "views" (softmaxes) together, the model can create a complex, curved shape that fits the data perfectly, even though the underlying ID cards are still short.

Efficiency: It's like having a small set of Lego bricks (the ID cards) but using a clever instruction manual (the mixture) to build a huge, complex castle. You don't need more bricks; you just need a better way to assemble them.
Cost: This new method adds very few extra parameters (memory cost) compared to just making the ID cards longer.

What the Experiments Showed

The authors tested this on several real-world knowledge graphs (like drug discovery networks and biological databases).

Small Libraries: On small datasets (like FB15k-237), the "small table" worked fine. The new "Pop-Up Book" didn't help much and sometimes even confused the model.
Big, Dense Libraries: On large, complex datasets (like openbiolink or ogbl-biokg), the old models struggled. The "Pop-Up Book" (KGE-MOS) significantly improved the ability to rank items correctly and predict probabilities.
The Trade-off: The new method is slightly slower to train (about 2x slower), but it is much faster than trying to make the ID cards huge, which would require massive amounts of memory and often fail to fit on computers.

Summary

The Problem: Current AI models for linking data use a "bottleneck" that limits how complex their predictions can be. They try to squeeze a huge amount of information into a tiny space, losing accuracy.
The Proof: The paper mathematically proves that you can't fix this by just making the data slightly bigger; you'd need impossibly large data sizes.
The Fix: They built a new output layer (KGE-MOS) that acts like a multi-layered pop-up book. It allows the model to represent complex patterns without needing massive amounts of memory.
The Result: For large, real-world data, this new method makes predictions more accurate and reliable, offering a practical way to scale up AI for big knowledge graphs.

Technical Summary: On the Theoretical Limitations of Embedding-based Link Prediction

1. Problem Statement

The paper addresses a fundamental limitation in modern Knowledge Graph Embedding (KGE) models: the rank bottleneck introduced by linear output layers. In standard KGE architectures, entities and relations are mapped to low-dimensional embeddings (typically $d \in [10^2, 10^4]$ ), which are then projected to a high-dimensional output space of entities (typically $|E| \in [10^5, 10^9]$ ) via a linear layer.

The authors argue that this linear projection confines the model's output scores to a $d$ -dimensional linear subspace within the $|E|$ -dimensional space. When $d \ll |E|$ , this creates a "rank bottleneck" that severely restricts the model's expressivity. Specifically, the model cannot represent arbitrary ranking orders, sign patterns (positive/negative triples), or probability distributions required for complex downstream reasoning tasks, regardless of the power of the upstream neural encoder (e.g., Graph Neural Networks or Transformers).

2. Methodology

Theoretical Analysis

The authors derive necessary bounds for the embedding dimension $d$ required for Linear-Decoder KGEs (LD-KGEs) to satisfy three distinct prediction tasks:

Ranking Reconstruction (RR): Correctly ordering objects for a query.
Sign Reconstruction (SR): Correctly classifying triples as true or false.
Distributional Reconstruction (DR): Assigning calibrated probabilities to triples.

They analyze these tasks under two conditions:

Universality: The ability to represent all possible patterns for graphs of a given size. The paper proves that for universality, $d$ must be at least $|E|-1$ (for RR and DR) or $|E|$ (for SR). These bounds are practically impossible for large-scale graphs.
Realisability: The ability to represent patterns consistent with a specific observed graph. The authors show that the necessary dimension depends on the sign rank of the graph's adjacency matrix (or its sign matrix). While sign rank is generally smaller than the full rank, it is NP-hard to compute.

To provide a practical sufficient condition, the authors connect sign rank to graph connectivity. They demonstrate that the required dimension is bounded by the maximum number of sign changes in the rows of the adjacency matrix. Using a Reverse Cuthill-McKee (RCM) heuristic to reorder entities, they show that the required dimension scales with the graph's connectivity (specifically, the number of contiguous blocks of positive entries) rather than just the total number of entities. However, even with optimal ordering, dense graphs require dimensions far exceeding standard practice.

Proposed Solution: KGE-MOS

To break this bottleneck without incurring the massive parameter cost of increasing $d$ , the authors propose KGE-MOS (Knowledge Graph Embedding with Mixture of Softmaxes).

Architecture: Instead of a single linear layer followed by a softmax, KGE-MOS replaces the output layer with a mixture of $K$ softmaxes.
Mechanism: For a query $(s, r)$ , the model computes a mixture of $K$ distributions:
$P(O|s, r) = \sum_{k=1}^K \pi_k(h_{s,r}) \cdot \text{softmax}(f_k(h_{s,r})E^\top)$
where $h_{s,r}$ is the hidden state, $E$ is the shared entity embedding matrix, $f_k$ are component-specific non-linear projections (MLPs), and $\pi_k$ are mixture weights.
Efficiency: The entity embeddings $E$ are shared across all components. The additional parameters are limited to the projection weights $f_k$ and mixture weights $\pi_k$ , resulting in a cost of $O(Kd^2)$ , which is negligible compared to the $O(d|E|)$ cost of increasing the embedding dimension $d$ .

3. Key Contributions

Theoretical Bounds: The paper provides the first necessary bounds for all LD-KGEs (both bilinear and neural) regarding universality and realisability. It establishes that linear output layers fundamentally limit expressivity in large, dense graphs.
Connectivity Analysis: It derives a sufficient dimension for realisability based on graph connectivity (sign changes), showing that dense graphs are theoretically harder to fit than sparse ones.
KGE-MOS: It introduces a drop-in non-linear output layer that breaks the rank bottleneck efficiently.
Empirical Validation: It demonstrates that KGE-MOS improves both ranking accuracy (MRR) and probabilistic fit (NLL) on large, dense datasets, while simple increases in embedding dimension yield diminishing returns due to training instability (e.g., batch size constraints).

4. Experimental Results

The authors evaluated KGE-MOS on five datasets (FB15k-237, Hetionet, ogbl-biokg, openbiolink, WN18RR) using four baseline models (DISTMULT, COMPLEX, RESCAL, CONVE) and one neural encoder (COMPGCN).

Performance on Large/Dense Graphs: On the three largest and most densely connected datasets (Hetionet, ogbl-biokg, openbiolink), KGE-MOS consistently outperformed the baseline LD-KGEs in both MRR and NLL. The gains were most pronounced in the high-connectivity regime where the theoretical bottleneck is most severe.
Performance on Small Graphs: On the smaller FB15k-237 dataset, KGE-MOS showed negligible improvement or slight degradation, confirming that the bottleneck is less critical for smaller, less dense graphs.
Dimension Scaling: Increasing the embedding dimension $d$ from 200 to 1000 improved performance for baselines but required reducing the batch size due to GPU memory constraints, which ultimately hurt performance. In contrast, KGE-MOS at $d=200$ achieved performance comparable to baselines at $d=1000$ with significantly fewer parameters.
Ablation: Increasing the number of softmax components ( $K$ ) consistently improved performance. The non-linear projections $f_k$ also contributed to gains, but the mixture itself was the primary driver of breaking the bottleneck.
Rank Breaking: Empirical rank analysis of the log-probability matrix showed that while standard models were limited to rank $d+1$ , KGE-MOS achieved ranks approaching the full dimension of the entity set, effectively breaking the bottleneck.

5. Significance and Claims

The paper claims that the field of link prediction has focused heavily on improving encoders (e.g., using GNNs or Transformers) while neglecting the limitations of the output layer. The authors argue that exploring more expressive non-linear output layers is a necessary and promising avenue for scaling KGEs to large and dense graphs.

They emphasize that KGE-MOS offers a "drop-in" solution that improves performance and probabilistic fidelity at a low parameter cost, challenging the assumption that simply increasing embedding dimensions is the only way to handle larger graphs. The work suggests that future research should move beyond encoder-centric improvements to address the structural limitations of the decoding process in embedding-based models.