Towards Lightweight Graph Neural Network Search with Curriculum Graph Sparsification

Beini Xie,Heng Chang,Ziwei Zhang,Zeyang Zhang,Simin Wu,Xin Wang,Yuan Meng,Wenwu Zhu

DOI: https://doi.org/10.1145/3637528.3671706

2024-06-24

Abstract:Graph Neural Architecture Search (GNAS) has achieved superior performance on various graph-structured tasks. However, existing GNAS studies overlook the applications of GNAS in resource-constraint scenarios. This paper proposes to design a joint graph data and architecture mechanism, which identifies important sub-architectures via the valuable graph data. To search for optimal lightweight Graph Neural Networks (GNNs), we propose a Lightweight Graph Neural Architecture Search with Graph SparsIfication and Network Pruning (GASSIP) method. In particular, GASSIP comprises an operation-pruned architecture search module to enable efficient lightweight GNN search. Meanwhile, we design a novel curriculum graph data sparsification module with an architecture-aware edge-removing difficulty measurement to help select optimal sub-architectures. With the aid of two differentiable masks, we iteratively optimize these two modules to efficiently search for the optimal lightweight architecture. Extensive experiments on five benchmarks demonstrate the effectiveness of GASSIP. Particularly, our method achieves on-par or even higher node classification performance with half or fewer model parameters of searched GNNs and a sparser graph.

Machine Learning,Artificial Intelligence,Social and Information Networks

What problem does this paper attempt to address?

The paper aims to address the design and optimization of Lightweight Graph Neural Networks (LGNNs), particularly in resource-constrained scenarios. Specifically, the paper proposes a method called GASSIP (Graph Architecture Search with Sparsification and Pruning), which tackles two core challenges through the following approaches:

1. **Efficiently searching for GNN sub-architectures**: Traditional Graph Neural Architecture Search (GNAS) methods become very time-consuming and complex when manually designing GNN architectures for diverse graph data. GASSIP introduces an operation-pruned architecture search module, which can efficiently identify important parts, thereby constructing lightweight GNNs with fewer parameters.

2. **Efficient implementation of lightweight GNAS**: Directly adopting a search-then-prune approach incurs significant computational costs. GASSIP simplifies the optimization process and reduces computational burden by jointly optimizing the graph structure and architecture search stages.

Additionally, the paper proposes a novel curriculum graph data sparsification algorithm, which helps select the optimal sub-architecture by measuring the difficulty of edge removal in a structure-aware manner. By iteratively optimizing these two modules, GASSIP can reduce redundant edges while retaining critical information, thus improving model efficiency and maintaining high classification performance.

Experimental results show that GASSIP performs excellently on multiple benchmark datasets, significantly reducing the number of parameters while achieving classification performance that is superior to or on par with traditional GNNs and existing GNAS baseline methods. For example, on the Cora dataset, GASSIP reduces the number of parameters to half or even less, while achieving a 2.42% improvement in node classification performance over traditional GNNs and a 2.11% improvement over GNAS baselines. Additionally, the search time is reduced from the original 16 minutes to less than 1 minute.