Balancing Graph Embedding Smoothness in Self-Supervised Learning via Information-Theoretic Decomposition

TL;DR

This paper tackles the polarized performance of graph self-supervised learning methods by reframing SSL through an information-theoretic lens that includes a neighbor representation. It introduces BSG, a framework that adds three loss terms—neighbor loss, minimal loss, and divergence loss—to balance the terms arising from incorporating neighbor information into the SSL objective, alongside a standard graph masking-based SSL loss. The authors provide theoretical analyses linking each loss to graph smoothing and mutual information with downstream tasks, and demonstrate state-of-the-art results on node classification and link prediction across multiple real-world datasets, with robust improvements when integrating BSG into other SSL objectives. Overall, BSG offers a principled, tunable approach to achieving robust, balanced graph representations suitable for a range of downstream tasks. The work suggests that carefully balancing local neighborhood smoothness and task-relevant information is key to generalizable graph SSL.

Abstract

Self-supervised learning (SSL) in graphs has garnered significant attention, particularly in employing Graph Neural Networks (GNNs) with pretext tasks initially designed for other domains, such as contrastive learning and feature reconstruction. However, it remains uncertain whether these methods effectively reflect essential graph properties, precisely representation similarity with its neighbors. We observe that existing methods position opposite ends of a spectrum driven by the graph embedding smoothness, with each end corresponding to outperformance on specific downstream tasks. Decomposing the SSL objective into three terms via an information-theoretic framework with a neighbor representation variable reveals that this polarization stems from an imbalance among the terms, which existing methods may not effectively maintain. Further insights suggest that balancing between the extremes can lead to improved performance across a wider range of downstream tasks. A framework, BSG (Balancing Smoothness in Graph SSL), introduces novel loss functions designed to supplement the representation quality in graph-based SSL by balancing the derived three terms: neighbor loss, minimal loss, and divergence loss. We present a theoretical analysis of the effects of these loss functions, highlighting their significance from both the SSL and graph smoothness perspectives. Extensive experiments on multiple real-world datasets across node classification and link prediction consistently demonstrate that BSG achieves state-of-the-art performance, outperforming existing methods. Our implementation code is available at https://github.com/steve30572/BSG.

Figures (6)

• Figure 1: A comparison of existing graph SSL baselines, with colors indicating the smoothness of representations on a log-scale ((a)) and the height of each bar corresponds to the performance of downstream tasks ((b)).
• Figure 2: Illustration of the representations and loss functions of BSG. The figure first shows the process of obtaining three representations $\mathbf{z}_\mathbf{x}$, $\mathbf{z}_\mathbf{s}$, and $\mathbf{z}_\text{neigh}$. The right section defines three loss functions related to Eq \ref{['eq:1']}.
• Figure 3: The effect of $\mathcal{L}_\text{nei}$ and $\mathcal{L}_\text{div}$ respect to graph smoothness. The y-axis denotes the normalized graph embedding smoothness score, and the low values indicate oversmoothing.
• Figure 4: Ablation study of the proposed loss functions. We compared the performance without applying each loss function.
• Figure 5: Extended experiment of BSG in the Cora dataset, with values indicating the node classification results. (a) compares the performance of BSG and MaskGAE with different mask ratios. (b) analyzes the sensitivity of loss functions.

Theorems & Definitions (4)

• Definition 1
• Theorem 1
• Theorem 2
• Theorem 3