Parameter Averaging in Link Prediction

TL;DR

The paper tackles the high cost of ensemble methods for knowledge graph embeddings (KGE) by proposing parameter-averaging ensembles. It extends Stochastic Weight Averaging (SWA) to KGE and introduces ASWA, an adaptive, validation-driven parameter-averaging method that selectively updates the ensemble. Across link prediction, literal-augmented models, and multi-hop query answering, SWA and ASWA consistently improve generalization over single-model training and prior ensembles, with ASWA often outperforming SWA by mitigating overfitting. The approach maintains test-time and memory efficiency comparable to a single model, offering a practical path to robust KGEs in diverse tasks.

Abstract

Ensemble methods are widely employed to improve generalization in machine learning. This has also prompted the adoption of ensemble learning for the knowledge graph embedding (KGE) models in performing link prediction. Typical approaches to this end train multiple models as part of the ensemble, and the diverse predictions are then averaged. However, this approach has some significant drawbacks. For instance, the computational overhead of training multiple models increases latency and memory overhead. In contrast, model merging approaches offer a promising alternative that does not require training multiple models. In this work, we introduce model merging, specifically weighted averaging, in KGE models. Herein, a running average of model parameters from a training epoch onward is maintained and used for predictions. To address this, we additionally propose an approach that selectively updates the running average of the ensemble model parameters only when the generalization performance improves on a validation dataset. We evaluate these two different weighted averaging approaches on link prediction tasks, comparing the state-of-the-art benchmark ensemble approach. Additionally, we evaluate the weighted averaging approach considering literal-augmented KGE models and multi-hop query answering tasks as well. The results demonstrate that the proposed weighted averaging approach consistently improves performance across diverse evaluation settings.

Figures (1)

• Figure 1: Training and validation MRR scores of SWA and ASWA across epochs for different models on UMLS and KINSHIP datasets.