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

Alberto Bernardi; Luca Costabello; Christophe Gueret

cs.AISep 29, 2026

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

Authors: Alberto Bernardi, Luca Costabello, Christophe Gueret

Organizations: Accenture Labs, Dublin, Ireland

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 graph. If previous research has loosely linked the problem to relation types or degree bias, we show that it is more widespread and it correlates with the degree imbalance of the entities in test triples. In particular, 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 recommender systems and other use cases, where these triples 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.

Figures & tables

Figure 1 : We consider a query $q$ with high degree imbalance. In this example taken from FB15k-237, Jackie Chan has degree 49 and J-Pop has degree 8 (limited in the figure for clarity of visualization). ImbalancE takes the pre-trained embeddings of the KGE model for the anchor entity and performs an inference-time latent search optimization on a subset of the original neighborhood enhanced with the triples generated by the oracle. It then outputs a representation for the anchor node that better aligns with the selected query, improving the ranking of the target answer from 67 down to 3.

	FB15k-237		WN18RR		YAGO3-10
	MRR	H@10	MRR	H@10	MRR	H@10
ComplEx	0.31	0.49	0.51	0.58	0.36	0.56
RotatE	0.31	0.51	0.52	0.61	0.37	0.57

Table 1 : Aggregate performance of KGE models on FB15k-237, WN18RR, and YAGO3-10.

Panel of Figure 2 — Figure 2 : Performance of conventional KGE models on test triples binned by normalized degree difference $\hat{\Delta}$ , using quartiles of $\lvert\hat{\Delta}\rvert$ . On the left hand side of each plot we have triples with high-degree object and low-degree subject, on the right hand side, high-degree subjects and low-degree objects. As the degree imbalance increases, the performance gap between high-degree and low-degree prediction increases.

Panel of Figure 3 — Figure 3 : (Top) Euclidean Norm of the difference between embeddings of converging and non-converging entities while training RotatE on FB15k-237. (Bottom) Comparison of degree statistics for converging vs. non-converging entities.

	FB15k-237	WN18RR	Yago3-10
#H Nodes (Min degree)	$3536$ ( $41$ )	$2296$ (6)	$28\,913$ (5)
#L Nodes (Max degree)	$4015$ ( $11$ )	$32\,697$ (10)	$31\,487$ (16)
#Valid H-L / L-H	$338$ / $775$	$295$ (689)	$25$ / $298$
#Test H-L / L-H	$396$ / $942$	$277$ / $753$	$53$ / $270$

Table 2 : Statistics of the datasets.

Dataset	Model	High-Low				Low-High
Dataset	Model	MRR	H@1	H@3	H@10	MRR	H@1	H@3	H@10
FB15k-237	ComplEx-N3	0.03	0.01	0.01	0.06	0.03	0.01	0.02	0.08
	+ ImbalancE	0.13	0.05	0.14	0.28	0.08	0.04	0.08	0.17
	RotatE	0.02	0.01	0.01	0.05	0.04	0.01	0.04	0.09
	+ ImbalancE	0.13	0.06	0.11	0.36	0.08	0.04	0.09	0.17
Yago3-10	ComplEx-N3	0.06	0.00	0.03	0.23	0.07	0.02	0.07	0.17

Table 3 : Results on the three benchmark datasets on the High-Low and Low-High triples splits. The best value for each metric on each dataset is reported in bold except in case of a tie, when they are underlined .

Figure 4 : Sankey plots showing the distribution of test triples ranks assigned by RotatE (before) and how the distribution changed applying ImbalancE (after). ImbalancE greatly improves ranks even from far below the ranking ( $>$ 500) bringing them in the top 10.

Dataset	Model	High-Low				Low-High
Dataset	Model	MRR	H@1	H@3	H@10	MRR	H@1	H@3	H@10
FB15k-237	ComplEx-N3 + ImbalancE	0.13	0.05	0.14	0.28	0.08	0.04	0.08	0.17
	RotatE + ImbalancE	0.13	0.06	0.11	0.36	0.08	0.04	0.09	0.17
	TuckER [ 1 ]	0.08	0.03	0.08	0.21	0.08	0.04	0.07	0.16
	+KG-Mixup [ 21 ]	0.09	0.03	0.08	0.24	0.08	0.03	0.08	0.17
	CSProm-KG [ 5 ]	0.09	0.02	0.07	0.29	0.09	0.05	0.09	0.17

Table 4 : Results on FB15k-237 for ImbalancE and two baselines. The best value for each metric on each dataset is reported in bold except in case of a tie, when they are underlined . ImbalancE outperforms previous work to address “degree bias”, and it also brings simpler and less capable KGE models to the same levels of a much more computationally intensive option.

Dataset	Model	High-Low		Low-High
Dataset	Model	MRR	H@10	MRR	H@10
FB15k-237	ComplEx-N3	0.03	0.06	0.03	0.08
	+Context Only	0.05	0.14	0.05	0.12
	+Oracle Only	0.10	0.22	0.08	0.17
	+ ImbalancE	0.13	0.28	0.08	0.17
	RotatE	0.02	0.05	0.04	0.09

Table 5 : Results on FB15k-237 isolating the two terms of the loss function. Both terms provide complementary contributions.

	Model	FB15k-237	Yago3-10
Single query	ImbalancE	0.91	1.39
Full test set	ImbalancE	440.7	400.1
Pre-training	RotatE	935.9	$24\,131.3$

Table 6 : Latency of ImbalancE compared to re-training the KGE model from scratch. Runtime is reported in seconds and 30 epochs of latent search are considered, though convergence is typically reached within the first 10.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

	FB15k-237		YAGO3-10
	MRR	H@10	MRR	H@10
TransE	0.31	0.49	0.35	0.55
DistMult	0.30	0.48	0.34	0.53
NBFNet	0.41	0.60	-	-

Appendix

Table 7 : Aggregate performance of KGE models on FB15k-237 and YAGO3-10

Panel of Figure 5 — Figure 5 : Performance of conventional KGE models on test triples binned by normalized degree difference. Bins are obtained using quartiles of $\lvert\hat{\Delta}(t)\rvert$ . On the left hand side we have triples with high-degree object and low-degree subject, on the right hand side, high-degree subjects and low-degree objects. We isolate subject and object corruption.

Figure 6 : Performance of NBFNet on FB15k-237 test triples binned by normalized degree difference. Bins are obtained using quartiles of

\lvert\hat{\Delta}\rvert

. On the left hand side we have triples with high-degree object and low-degree subject, on the right hand side, high-degree subjects and low-degree objects. We isolate subject and object corruption.

Dataset	#Entities	#Relations	#Train	#Valid	#Test
FB15k-237	$15\,541$	$237$	$272\,115$	$17\,535$	$20\,466$
WN18RR	$40\,943$	$18$	$86\,835$	$3034$	$3134$
Yago3-10	$123\,182$	$37$	$1\,079\,040$	$5000$	$5000$

Appendix

Table 8 : Statistics of the datasets.

Dataset	Model	High-Low				Low-High
Dataset	Model	MRR	H@1	H@3	H@10	MRR	H@1	H@3	H@10
Yago3-10	ComplEx-N3	0.06	0.00	0.03	0.23	0.07	0.02	0.07	0.17
	+Context Only	0.06	0.00	0.03	0.23	0.06	0.02	0.07	0.17
	+Oracle Only	0.06	0.00	0.03	0.23	0.07	0.00	0.10	0.20
	+ ImbalancE	0.06	0.00	0.03	0.23	0.08	0.02	0.11	0.20
	RotatE	0.18	0.14	0.20	0.26	0.03	0.01	0.03	0.05

Appendix

Table 9 : Results on WN18RR, and YAGO3-10 isolating the contribution of the two terms of the loss function.

Panel of Figure 7 — Figure 7 : Validation metrics for FB15k-237 extending the latent search to 100 epochs. Metrics plateau after reaching convergence without any collapse, showing how the optimization process guarantees no degenerate collapse of the embeddings.

Figure 8 : Correlation between the size of the training context and the rank improvement in FB15k-237 and Yago3-10.

Panel of Figure 9 — Figure 9 : Sankey plots detailing how ImbalancE altered RotatE and ComplEx ranks on FB15k-237.

Panel of Figure 10 — Figure 10 : Sankey plots detailing how ImbalancE altered RotatE and ComplEx ranks on Yago3-10.

Panel of Figure 11 — Figure 11 : Sankey plots detailing how ImbalancE altered RotatE and ComplEx ranks on WN18RR.

Explore similar work

Jun 2, 2026cs.LG

Link Prediction or Perdition: the Seeds of Instability in Knowledge Graph Embeddings

Embedding models (KGEMs) constitute the main link prediction approach to complete knowledge graphs. Standard evaluation protocols emphasize rank-based metrics such as MRR or Hits@

K

, but usually overlook the influence of random seeds on result stability. Moreover, these metrics conceal potential instabilities in individual predictions and in the organization of embedding spaces. In this work, we conduct a systematic stability analysis of multiple KGEMs across several datasets. We find that high-performance models actually produce divergent predictions at the triple level and highly variable embedding spaces. By isolating stochastic factors (i.e., initialization, triple ordering, negative sampling, dropout, hardware), we show that each independently induces instability of comparable magnitude. Furthermore, for a given model, hyperparameter configurations with better MRR are not guaranteed to be more stable. Moreover, voting, albeit a known remediation mechanism, only provides a limited enhancement of stability. These findings highlight critical limitations of current benchmarking protocols, and raise concerns about the reliability of KGEMs for knowledge graph completion.

Guillaume Méroué, Fabien Gandon, Pierre Monnin

Université Côte d’Azur, Inria, CNRS, I3S, France

Knowledge Graph Embeddings Zero-Shot Link Prediction Knowledge Graphs Instability

Ask PDF

Jun 15, 2026cs.AI

Model Graph Inductive Learning for Knowledge Graph Completion

Link prediction in knowledge graphs fundamentally depends on the quality of learned embeddings for entities and relations. However, most existing methods derive these embeddings by aggregating only the local neighborhood of each entity, neglecting the global structure of the knowledge graph. This limited view prevents models from capturing higher-level structural patterns that are essential for accurate and generalizable link prediction. To address these limitations, we introduce Model Graph Inductive Learning (\textbf{MGIL}), a framework that constructs a model graph by clustering entities based on the similarity of their incoming and outgoing relational structures or their entity types. A GNN is then applied to this model graph to produce embeddings that capture the global view of the knowledge graph. These embeddings subsequently serve as high-quality initial features %embeddings for the original knowledge graph, replacing random initialization and leading to more stable and expressive representations. Extensive experiments on standard and recently proposed inductive benchmarks demonstrate that MGIL achieves state-of-the-art or highly competitive performance in inductive link prediction, highlighting its effectiveness across diverse graph settings.

Mohommad Esmaei Khani, Mahdieh Hasheminejad, Ali Taherkhani +1

Faculty of Mathematical Science Department of Computer Science, Yazd Unversity,, 100190, Iran. · Department of Mathematics, Institute for Advanced Studies in Basic Sciences (IASBS), 45137-66731, Iran. · Diagnostic and Research Institute of Human Genetics, Medizinische Universität Graz, Graz, 8010, Austria.

Zero-Shot Link Prediction Graph Representation Learning Knowledge Graphs Graph Representations Relational

Ask PDF

Jun 29, 2026cs.AI

Beyond Triplet Plausibility: Relation Set Completion in Knowledge Graphs

Knowledge graphs (KGs) organize real-world knowledge as triplets and underpin many downstream applications. Due to their inherent incompleteness, knowledge graph completion (KGC) is widely studied and is typically formulated as triplet prediction, with link prediction as the dominant paradigm. However, this formulation focuses on the incompleteness of triplet-wise information and overlooks the incompleteness of entity-relation compatibility information. To address this limitation, we introduce a relation set completion task (RSC), which complements the link prediction task and aims to reason about missing relations that are semantically compatible with a given entity. We further propose a Relation Set Embedding model (RelSetE), which models latent patterns among the observed relations of entities to infer missing ones. To evaluate RelSetE, we derive three benchmark datasets from standard KG benchmarks. Extensive experiments demonstrate that RelSetE effectively captures entity-relation compatibility patterns and performs favorably in inferring missing relations of entities. Code and data are publicly available.

Zihao Zheng, Borui Cai, Yao Zhao +2

School of Information Technology, Deakin University, VIC 3125, Australia · Hangzhou International Innovation Institute, Beihang University, Hangzhou 311115, China · School of Information Technology, Victoria University, VIC 3011, Australia +1

Knowledge Graph Completion Knowledge Graphs Zero-Shot Link Prediction

Ask PDF