Parameter Averaging in Link Prediction

2025-10-29

AI Summary

To address the high computational and memory overhead of ensemble learning in knowledge graph embedding (KGE) models for link prediction, this paper proposes a parameter-level model merging approach. It maintains model parameters via sliding weighted averaging during training and introduces a dynamic update strategy triggered solely by validation-set performance improvement—eliminating the need for parallel multi-model training. This work is the first to effectively adapt lightweight model merging to KGE link prediction, simultaneously enhancing literal-awareness and multi-hop query generalization. Experiments across multiple benchmark datasets demonstrate that the method consistently outperforms state-of-the-art ensembles: inference speed improves by 1.8–3.2×, GPU memory consumption decreases by 40%–65%, and cross-task generalization capability is significantly strengthened.

Technology Category

• Machine Learning: Ensemble Methods
• Data Mining & Knowledge Management: Linked Open Data, Knowledge Graphs & KB Completion
• Natural Language Processing: Sentence-level Semantics, Textual Inference, etc.

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Reducing computational overhead in ensemble knowledge graph embedding models
• Improving link prediction performance through parameter averaging techniques
• Addressing latency and memory issues in ensemble learning methods

Innovation

• Parameter averaging replaces ensemble training for KGE models
• Selective updating based on validation performance improves generalization
• Weighted averaging enhances link prediction and multi-hop tasks