ImbalancE: Inference-Time Latent Search against Degree Imbalance in Link Prediction

ICLR 2026 Conference Submission20041 Authors

Keywords: Knowledge Graph, Link Prediction, Degree Imbalance, Knowledge Graph Completion, Latent Search, Knowledge Graph Embeddings

TL;DR: We scope the degree imbalance problem that affects knowledge graph embedding models for knowledge graph completion and we propose ImbalancE, an inference-time latent search method to address this widespread issue

Abstract:

Knowledge Graph Embedding models have been extensively used to learn representations of entities and relations in Knowledge Graphs for predicting missing links. However, the quality of the learned representations varies a lot across different areas of the same Knowledge Graph. If previous research efforts have loosely linked the problem to relation types or degree bias, we show that it is much more widespread, and it more precisely lies in the degree imbalance of the entities in test triples. In particular, we show that the prediction of a target entity that has a degree much smaller than the degree of the anchor entity is extremely problematic. This is critical in use cases like \textit{drug target discovery}, where these triples are predominant, or \textit{recommender systems}, where they represent important corner cases.
To address this issue, we propose an inference-time latent search optimization method capable of significantly improving model predictions on the most imbalanced triples. Built on top of a pre-trained model, it explores the embedding space at evaluation time, blending known and out-of-band information to mitigate the degree imbalance bias. We show the value of our approach on imbalanced triples from common benchmark datasets, where we outperform conventional methods, opening the door to the successful adoption of Knowledge Graph Embedding models on these critical corner cases.

Primary Area: learning on graphs and other geometries & topologies

Submission Number: 20041