Picture a large company with multiple departments. If each employee is likened to a node, they would be interconnected through relationships such as collaborations or reporting lines.
Graph neural networks, or GNNs, constitute a class of machine learning methods capable of leveraging relationships between nodes to make predictions, such as predicting employee performance or identifying key influencers.
Traditional methods for distilling GNNs tend to focus on preserving local structures, potentially overlooking preserving global interactions and latent relationships. For instance, the distilled graph might only consider day-to-day interactions among employees who work closely yet miss the broader perspective when predicting the performance of employees.
A*STAR researchers Fayao Liu, Xun Xu and Chuan-Sheng Foo from the Institute for Infocomm Research (I2R) explored a new approach to preserving global topology in knowledge distillation of GNNs to enhance their performance.
The team introduced Graph Contrastive Representation Distillation (G-CRD) to assist the 'student' (smaller, less complex GNN models) in learning from a 'teacher', a large, pre-trained GNN capturing intricate node relationships. The researchers explained that instead of merely matching features one-to-one, G-CRD accentuates similarities and differences across the entire network to bolster learning efficiency.
“Teacher and student embeddings belonging to the same node were encouraged to be pulled closer, while those from different nodes were pushed apart,” the team commented.
The researchers assessed G-CRD by applying it across various GNN architectures and datasets, including large-scale networks, batches of small graphs and 3-D point clouds, focusing on real-world applications where data might be noisy and constantly evolving. They compared G-CRD with existing distillation techniques for node classification and molecular graph property prediction.
Their findings revealed that G-CRD surpassed existing methods of preserving local structure by effectively capturing both immediate and broader network connections, thereby significantly enhancing robustness and performance. G-CRD emerged as a faster, more accurate means of preserving complex relationships, even with noisy data or in resource-constrained environments.
“Lightweight and efficient GNN models trained with our distillation techniques may be useful for deployment on edge devices or in scenarios where resources are limited,” the researchers explained. “For instance, the low GPU memory consumption and fast inference speed of smaller models can be very useful for resource-constrained robotics applications.”
The introduction of G-CRD signifies a significant advancement in resource-efficient artificial intelligence research, offering a more nuanced and effective approach for smaller models to learn intricate patterns, with far-reaching implications for technology advancement in areas such as drug discovery, autonomous vehicle navigation and beyond.
The A*STAR-affiliated researchers contributing to this research are from the Institute for Infocomm Research (I2R).