PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Alberto Bernardi, Luca Costabello, Christophe Gueret

arXiv:2609.36996v1cs.AIcs.LG

TL;DR

Degree imbalance is associated with substantially poorer predictions of low-degree entities, across a range of graph-learning models. IMBALANCE applies latent search at inference time to refine embeddings for imbalanced queries, and the authors report improvements on benchmark triples, with performance dependent on dataset and oracle quality.

  • Problem

    Prediction quality varies across the graph, and greater degree imbalance in test triples is associated with poorer predictions of the lower-degree entity.

    Source: §1 · p. 2, §3 · p. 4

  • Method

    IMBALANCE applies inference-time latent search to fine-tune the anchor and selected entity embeddings, using training context and oracle-generated triples.

    Source: §1 · p. 2, Training Context Term · p. 8

  • Results

    The authors report that IMBALANCE improves predictions on heavily imbalanced benchmark triples and can make simpler models perform on the same level as more involved ones.

    Source: §1 · p. 3, Baseline Comparison · p. 11

  • Takeaways & Limitations

    The method's reported efficacy and efficiency support its use in link-prediction settings where degree imbalance is a consistent concern.

    Source: §7 · p. 14

    IMBALANCE requires non-empty training context, and its benefit depends on the dataset and oracle-input quality.

    Source: §6 · pp. 13–14

Abstract

from arXiv · show

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.

1 Introduction

Knowledge graph embeddings predict missing links, but prediction quality varies across the graph and worsens for low-degree targets when test triples have strongly imbalanced entity degrees. The paper proposes IMBALANCE, an inference-time latent search method, and reports improvements on highly imbalanced benchmark triples.

1 Introduction

  • Link prediction infers a missing entity in an incomplete triple, using the known entity as the anchor and the missing one as the target.KGE models learn low-dimensional representations of entities and relations for this task.Source: §1 · p. 1
  • Degree imbalance correlates with prediction quality: higher-degree entities are predicted better while lower-degree entities are predicted much more poorly.The problem is reported across shallow KGE models and GNN-like models, including NBFNet.Source: §1 · p. 2
  • IMBALANCE fine-tunes the anchor and selected entity embeddings at inference time to improve predictions of low-degree entities.Its dual-term objective explores the embedding space using selected training facts and extends sparse long-tail knowledge with an oracle.Source: §1 · pp. 2–3
  • The authors report significant improvements on the most imbalanced triples in common benchmark datasets.They describe the method as a lightweight, efficient alternative to more computationally expensive models.Source: §1 · p. 3

2 Related Work

Prior work links prediction difficulty to high-degree bias, long-tail entity distributions, or relation-type heterogeneity. This paper frames degree imbalance as a broader problem and applies inference-time latent search to KGE models.

2 Related Work

  • NBFNet also suffers from degree imbalance, but its entangled entity embeddings prevent independent updates and make IMBALANCE inapplicable.The authors also characterize this approach as not scalable and prone to overfitting.Source: Knowledge Graph Embedding Models · p. 3
  • Prior work associated topology-related difficulty with high-degree bias and noted that aggregate metrics can distort model performance, but did not propose a countermeasure.Other work addressed entity-relation frequency through synthetic embeddings.Source: Topology-related Issues on KGs · pp. 3–4
  • Several methods address relation heterogeneity, especially 1-to-many and many-to-1 relations, by making relation representations more expressive.Source: Topology-related Issues on KGs · p. 4
  • The paper identifies inference-time latent search as previously used to explore latent program spaces, and says this is its first application to KGE models.The authors note that applying it to KGE requires non-trivial adaptations.Source: Latent Search Optimization · p. 4

3 What Aggregate Metrics Do not Show

Aggregate metrics obscure a widening performance gap as test-triple degree imbalance increases. The section connects low-degree prediction difficulty to overfitting or failed convergence during training.

3 What Aggregate Metrics Do not Show

  • The normalized degree difference is defined from the subject and object degrees, with node degree counting incoming and outgoing edges.The triples are binned by quartiles of the absolute normalized difference.Source: §3 · p. 4
  • As normalized degree difference increases, the prediction gap between high- and low-degree entities increases; aggregate metrics can hide this difficulty.Figure 2 bins test triples by normalized degree difference and separates subject from object corruptions.Source: §3 · p. 4, Figure 2 · p. 6
  • Entities with persistently large embedding updates show failed convergence rather than simply needing more training.The analysis plots distances between successive-epoch embeddings for RotatE on FB15k-237.Source: §3 · pp. 4–5
  • The convergent entity group has a much higher degree, while degree-matched low-degree entities show that greater training-triple imbalance favors convergence.For low-degree nodes, this apparent convergence is interpreted as overfitting; nodes connected to lower-degree entities can instead fail to converge.Source: §3 · p. 5
  • The low-degree prediction challenge has two forms: overfitting on one side and extremely poor learning on the other.IMBALANCE is proposed to improve predictions on the most imbalanced test triples involving low-degree entities.Source: §3 · p. 5

4 Inference Time Latent Search for Degree Imbalance

The section formalizes link prediction and explains how query-specific latent search addresses degree imbalance. It refines the anchor using training context and adjusts target embeddings with oracle triples.

4.1 Preliminary and Notation

  • Link prediction ranks candidate triples by a score proportional to their likelihood of being true.A knowledge graph is a set of subject–predicate–object triples.Source: §4.1 · p. 6

4.2 Intuition

  • A high-degree anchor may have a one-size-fits-all embedding poorly suited to a query, while low-degree entities may be underlearned or overfit.The method uses training triples sharing the query’s anchor and predicate as its training context, and an external oracle to suggest similar triples.Source: §4.2 · pp. 7–8

4.3 ImbalancE

  • IMBALANCE performs inference-time search on a pretrained KGE model, refining relevant embeddings for a query with a two-term objective.The approach can be applied to queries in either subject- or object-prediction form.Source: §4.3 · p. 8
  • The training-context term re-presents known query-context triples and updates only the anchor embedding, tailoring it to the query while avoiding further overfitting of learned triples.The other embeddings in those triples remain frozen.Source: Training Context Term · pp. 8–9
  • The oracle-enhanced term adds likely-answer triples from an out-of-band source and updates both anchor and target embeddings to bias targets toward the query.Benchmark experiments use an LLM to assess similarity from entity descriptions; other proposed sources include user interaction patterns or domain expertise.Source: Oracle-Enhanced Term · p. 9
  • Oracle triples may be false positives in the evaluation ranking, while freezing relation and context-target embeddings avoids negatives and acts as a regularizer.The method can be selectively applied to individual queries and is described as scalable.Source: Oracle-Enhanced Term · p. 9

5 Experiments

Experiments evaluate IMBALANCE on three benchmark knowledge graphs, including degree-imbalanced test splits. Results show substantial gains on FB15k-237, smaller or metric-specific gains elsewhere, and competitive performance against two baselines.

5.1 Experimental Settings

5.2 Results

  • On FB15k-237, metrics improve at least 2× on Low-High and over 6× on High-Low; WN18RR gains are smaller, YAGO3-10 improves only Hits@10, and Fig. 4 shows >500 ranks reaching the top 10.The authors attribute the smaller WN18RR gains to its less polarized degree distribution and the YAGO3-10 results to lower-quality oracle triples.Source: §5.2 · pp. 10–11, Table 3 · p. 11, Figure 4 · p. 12
  • The gains depend on oracle quality, making oracle selection a critical design choice.Source: Baseline Comparison · p. 12
  • On imbalanced triples, KG-Mixup provides little benefit compared with IMBALANCE; IMBALANCE outperforms CSProm-KG or falls behind slightly, despite being much more lightweight.On the entire FB15k-237 test set, CSProm-KG reports MRR = 0.36 and H@10 = 0.54, superior in aggregate metrics to ComplEx and RotatE.Source: Baseline Comparison · p. 11
  • The training-context and oracle loss terms each contribute positively on FB15k-237, and their joint contribution exceeds either term alone.The authors describe the terms as complementary.Source: Baseline Comparison · p. 11, Table 5 · p. 13
  • IMBALANCE is reported as more efficient than retraining the KGE model from scratch, with latent-search cost depending on query-context and oracle triples.The latency comparison covers 30 search epochs, though convergence is typically reached within the first 10.Source: Time-Complexity and Latency · p. 12, Table 6 · p. 13

6 Limitations and Future Directions

IMBALANCE requires non-empty training context, and its benefits also depend on dataset degree distributions and oracle quality.

6 Limitations and Future Directions

  • IMBALANCE applies only to queries with non-empty training context.The authors identify refining training-triple selection as a direction for overcoming this limitation.Source: §6 · p. 13
  • Less polarized degree distributions, such as WN18RR, may reduce IMBALANCE’s benefit.Source: §6 · p. 14
  • On YAGO3-10, oracle triple quality affects result quality, making performance dependent on the LLM or other oracle.The authors leave comparison of alternative oracles, including human feedback, for future work.Source: §6 · p. 14

7 Conclusions

The conclusion frames degree imbalance as a widespread KGE problem with overfitting and failed convergence among low-degree failure modes, and presents IMBALANCE as an inference-time mitigation. Supplemental materials provide reproducibility details and code.

7 Conclusions

  • Low-degree entities show two distinct failure modes: overfitting and failed convergence.Source: §7 · p. 14
  • IMBALANCE is presented as the first inference-time latent-search method for KGE models, with efficacy on highly imbalanced triples.The authors connect its efficacy and efficiency to reliable link prediction in use cases where degree imbalance is a consistent concern.Source: §7 · p. 14
  • Supplemental materials include dataset-split instructions, dataset statistics, hyperparameters, and code for reproducing the study.Source: §7 · p. 14

A Degree Imbalance for Additional Models

Additional-model results show that increasing degree imbalance widens the prediction gap between high- and low-degree entities, while NBFNet remains difficult to improve with IMBALANCE.

A Degree Imbalance for Additional Models

B Dataset Statistics

The appendix points readers to dataset statistics and describes how Figure 5 groups test triples by normalized degree difference.

B Dataset Statistics

  • Table 8 reports statistics for the three datasets considered in the experiments.The supplied passages do not state the table’s numerical values.Source: B Dataset Statistics · p. 18, Table 8 · p. 18
  • Figure 5 bins test triples by quartiles of normalized degree difference, separates high-degree-object/low-degree-subject from high-degree-subject/low-degree-object cases, and isolates subject and object corruption.Source: Figure 5 · p. 19

C Oracle Technical Details

The oracle constructs candidate triples by encoding entity labels and descriptions, then selecting entities near the barycenter of observed context entities. LLM triples are precomputed to avoid query overhead during experiments.

Oracle construction

  • For q = (s, p, ?), the oracle selects the top-m entities closest to the barycenter of observed context entities in the LLM embedding space.Entity labels and descriptions are encoded with NovaSearch/stella_en_400M_v51 at dimension 1,024; cosine similarity measures distance.Source: C Oracle Technical Details · p. 19
  • Yago3-10 uses entity labels alone, while the other datasets use labels and descriptions.Source: C Oracle Technical Details · p. 19

Preprocessing

Degree-binned performance

  • Figure 6 bins FB15k-237 test triples by quartiles of normalized degree difference and separates subject from object corruption.Its left side shows high-degree objects with low-degree subjects; its right side shows the reverse.Source: Figure 6 · p. 20

D Additional Experiments

Additional experiments examine loss-term contributions, rank changes, optimization stability, seed variability, and training-context size. Together, they report stronger results from combining loss terms, stable validation behavior, and rank improvements across datasets.

Loss Function Terms Contribution

  • The two loss terms together give the best results; single-term contributions are more limited on Yago3-10 and WN18RR than on FB15k-237.The authors attribute the difference to IMBALANCE having a more limited overall impact on Yago3-10 and WN18RR.Source: Loss Function Terms Contribution · p. 20, Table 9 · p. 21

Sankey Plots

Degenerate Behaviour Loss Function

  • The positive-fact objective avoids observed collapse because the first loss term updates only the anchor embedding, while target and relation embeddings remain fixed.The fixed relation embedding also constrains overfitting in both loss terms.Source: Degenerate Behaviour Loss Function · pp. 20–21
  • On FB15k-237, validation MRR and H@10 plateaued after the first few epochs and remained stable through 100 epochs.The authors recommend limited inference-time iterations; the extended run showed no collapse.Source: Degenerate Behaviour Loss Function · p. 21, Figure 7 · p. 22

Training Context Size v Performance

Execution with different seeds

  • Across five seeds, the method produced exactly the same averaged performance on all datasets, with only extremely limited GPU floating-point variability.The procedure starts from fixed pretrained embeddings and uses fixed observed triples without negative sampling, stochastic minibatching, or dropout.Source: Execution with different seeds · p. 22

Related papers

Loading 2609.36996v1…