Graph Pooling via Dropping Task-Irrelevant Nodes

Abstract

Graph neural networks (GNNs) face scalability challenges. While recent approaches have adopted pooling strategies inspired by convolutional neural networks (CNNs) to reduce graph size and improve efficiency, these methods often focus on local information and are optimized for single graph-level tasks. This limitation hinders their effectiveness in multi-task scenarios that require task-specific global information. We present DOTIN (Dropping Out Task-Irrelevant Nodes), an approach to graph size reductio. DOTIN utilizes K learnable virtual nodes to represent graph embeddings for K distinct graph-level tasks. By employing a transformer-based attention model, it adaptively removes up to 90% of low-attentiveness raw nodes without notable performance degradation. Our method achieves comparable accuracy to state-of-the-art techniques while offering substantial benefits in efficiency. Specifically, DOTIN accelerates Graph Attention Networks (GAT) by approximately 50% on graph-level tasks such as graph classification and graph edit distance (GED). Additionally, it reduces memory usage by about 60% on the D&D dataset. These results show DOTIN's potential to enhance the scalability and efficiency of deep GNNs across multiple graph-level tasks while maintaining high performance.