Enhancing Contrastive Link Prediction With Edge Balancing Augmentation

2025-08-20

AI Summary

Existing contrastive learning-based link prediction methods suffer from two major limitations: (i) a lack of theoretical grounding, and (ii) neglect of degree distribution imbalance and its adverse impact on contrastive optimization. This paper establishes the first theoretical analysis framework for contrastive learning in link prediction. We propose Edge-Balanced Augmentation (EBA), a degree-aware graph structural reweighting scheme that mitigates gradient bias in the contrastive loss between high- and low-degree nodes. Furthermore, we design a tailored contrastive loss function compatible with EBA and integrate it into an autoencoder architecture to enable end-to-end training. Extensive experiments on eight benchmark datasets demonstrate that our method consistently outperforms state-of-the-art models, validating its effectiveness, robustness, and generalizability.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Search and Optimization: Learning to Search
• Data Mining & Knowledge Management: Linked Open Data, Knowledge Graphs & KB Completion

Application Category

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

Abstract

Link prediction is one of the most fundamental tasks in graph mining, which motivates the recent studies of leveraging contrastive learning to enhance the performance. However, we observe two major weaknesses of these studies: i) the lack of theoretical analysis for contrastive learning on link prediction, and ii) inadequate consideration of node degrees in contrastive learning. To address the above weaknesses, we provide the first formal theoretical analysis for contrastive learning on link prediction, where our analysis results can generalize to the autoencoder-based link prediction models with contrastive learning. Motivated by our analysis results, we propose a new graph augmentation approach, Edge Balancing Augmentation (EBA), which adjusts the node degrees in the graph as the augmentation. We then propose a new approach, named Contrastive Link Prediction with Edge Balancing Augmentation (CoEBA), that integrates the proposed EBA and the proposed new contrastive losses to improve the model performance. We conduct experiments on 8 benchmark datasets. The results demonstrate that our proposed CoEBA significantly outperforms the other state-of-the-art link prediction models.

Problem

• Addresses lack of theoretical analysis for contrastive link prediction
• Resolves inadequate consideration of node degrees in contrastive learning
• Proposes edge balancing augmentation to adjust node degree distribution

Innovation

• Edge Balancing Augmentation adjusts node degrees
• Proposes new contrastive losses for improved performance
• Integrates theoretical analysis with practical augmentation