Graph Neural Networks at a Fraction

TL;DR

The paper addresses the high parameter and energy costs of Graph Neural Networks by introducing Quaternion Message Passing Neural Networks (QMPNNs), which represent features and weights in quaternion space to reduce trainable parameters by $1/4$ while preserving or improving accuracy. It provides a generalized framework to convert any GNN to a quaternion counterpart and redefines Graph Lottery Tickets (GLTs) for graph-based tasks, enabling substantial sparsification without loss of performance. Empirical results across node classification, link prediction, and graph classification demonstrate that QMPNNs can match or surpass real-valued GNNs with far fewer parameters, and GLTs exist at small fractions of the original size, particularly for large-scale graphs. The work contributes a practical, scalable approach to efficient graph representation learning and lays groundwork for future extensions to dynamic graphs, multi-modal data, and interpretability of quaternion-based GNNs.

Abstract

Graph Neural Networks (GNNs) have emerged as powerful tools for learning representations of graph-structured data. In addition to real-valued GNNs, quaternion GNNs also perform well on tasks on graph-structured data. With the aim of reducing the energy footprint, we reduce the model size while maintaining accuracy comparable to that of the original-sized GNNs. This paper introduces Quaternion Message Passing Neural Networks (QMPNNs), a framework that leverages quaternion space to compute node representations. Our approach offers a generalizable method for incorporating quaternion representations into GNN architectures at one-fourth of the original parameter count. Furthermore, we present a novel perspective on Graph Lottery Tickets, redefining their applicability within the context of GNNs and QMPNNs. We specifically aim to find the initialization lottery from the subnetwork of the GNNs that can achieve comparable performance to the original GNN upon training. Thereby reducing the trainable model parameters even further. To validate the effectiveness of our proposed QMPNN framework and LTH for both GNNs and QMPNNs, we evaluate their performance on real-world datasets across three fundamental graph-based tasks: node classification, link prediction, and graph classification.

Figures (2)

• Figure 1: Pruning a quaternion message passing neural network. Graph's features are first transformed into quaternion features. We train QMPNN, which has quaternion weights, for the given task. Furthermore, to find our proposed winning lottery ticket, we prune and train the QMPNN until we get a model with pruned quaternion weights that gives a comparable accuracy. The pruned quaternion network can be trained with only a small fraction of the parameters in the real GNN.
• Figure 2: Performance After Pruning: The plots show GCN, GAT, and GraphSAGE performance on OGBN-ARXIV (node classification), OGBL-COLLAB (link prediction), and OGBG-MOHLIV (graph classification) at a pruning weight fraction of 0.3. GLTs are marked by red ($\star$) and green ($\star$) stars, indicating comparable performance despite sparsity, while dashed lines represent pre-pruning baselines. The plot sections correspond to node classification (top), link prediction (middle), and graph classification (bottom).