Enhancing Graph Neural Networks: A Mutual Learning Approach

TL;DR

The paper tackles the challenge of deploying accurate graph neural networks on resource-constrained devices by moving from traditional unidirectional knowledge distillation to mutual learning among multiple shallow GNNs. It proposes Graph Mutual Learning (GML), where peers exchange knowledge via posterior distributions $p_j$ with temperature scaling $T_v$ and are guided by a mutual KL-divergence loss, augmented by adaptive logit weighting and an entropy-based uncertainty penalty. The framework demonstrates improvements on node and graph classification benchmarks and shows that learned representations can be distilled into a graph-less MLP for faster inference, using architectures such as $GCN$, $GAT$, and $GraphSage$. Across experiments, GML scales with cohort diversity and maintains gains over baselines, while recognizing increased training costs that could be mitigated through parallelization; future work includes robustness to noisy data and scalability to larger graphs.

Abstract

Knowledge distillation (KD) techniques have emerged as a powerful tool for transferring expertise from complex teacher models to lightweight student models, particularly beneficial for deploying high-performance models in resource-constrained devices. This approach has been successfully applied to graph neural networks (GNNs), harnessing their expressive capabilities to generate node embeddings that capture structural and feature-related information. In this study, we depart from the conventional KD approach by exploring the potential of collaborative learning among GNNs. In the absence of a pre-trained teacher model, we show that relatively simple and shallow GNN architectures can synergetically learn efficient models capable of performing better during inference, particularly in tackling multiple tasks. We propose a collaborative learning framework where ensembles of student GNNs mutually teach each other throughout the training process. We introduce an adaptive logit weighting unit to facilitate efficient knowledge exchange among models and an entropy enhancement technique to improve mutual learning. These components dynamically empower the models to adapt their learning strategies during training, optimizing their performance for downstream tasks. Extensive experiments conducted on three datasets each for node and graph classification demonstrate the effectiveness of our approach.

Figures (8)

• Figure 1: The GML architecture with two GNN models, $GNN_{tar}$ and $GNN_{col}$, training together. The adaptive logit weight unit prioritizes crucial knowledge and the uncertainty enhancement unit penalizes models with low entropy. The process yields two improved models $GNN_{tar}^{\prime}$ and $GNN_{col}^{\prime}$.
• Figure 2: Centered kernel alignment similarity between model layers 1, 2, and 3. (a) and (b) show the similarity between models of the same architecture but different initializations. (c) and (d) show the similarity between models of different architectures and different initializations.
• Figure 3: Results of KD with the best performing GML technique. We show the difference in the performance of KD with the MLP without KD as influenced by GML.
• Figure 4: Sensitivity analysis of the balancing parameters $\gamma$, $\beta$ and the temperature $T_v$.
• Figure 5: Comparison of GML ensemble method with deep ensembles.