Parameter Averaging in Link Prediction

TL;DR

This paper introduces a parameter averaging method for knowledge graph embedding models that improves link prediction performance without the high computational costs of traditional ensemble methods.

Contribution

It proposes a weighted averaging approach for model merging in KGE models, reducing training overhead while maintaining or improving accuracy.

Findings

Weighted averaging consistently outperforms baseline ensemble methods.

Selective updates based on validation performance enhance model generalization.

Approach is effective across various KGE models and 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.