[Paper Review] DOTIN: Dropping Task-Irrelevant Nodes for GNNs

Shaofeng Zhang, Feng Zhu+3 more|arXiv (Cornell University)|Apr 28, 2022

Advanced Graph Neural Networks4 citations

TL;DR

DOTIN proposes a novel, task-aware graph pruning method that adaptively drops up to 90% of task-irrelevant nodes in GNNs by leveraging K learnable virtual nodes to capture task-specific global representations. By using a transformer-based attention mechanism to rank node importance, DOTIN achieves ~50% faster training and ~60% less memory usage on D&D with no performance drop, outperforming pooling-based baselines in both single- and multi-task settings.

ABSTRACT

Scalability is an important consideration for deep graph neural networks. Inspired by the conventional pooling layers in CNNs, many recent graph learning approaches have introduced the pooling strategy to reduce the size of graphs for learning, such that the scalability and efficiency can be improved. However, these pooling-based methods are mainly tailored to a single graph-level task and pay more attention to local information, limiting their performance in multi-task settings which often require task-specific global information. In this paper, departure from these pooling-based efforts, we design a new approach called DOTIN (\underline{D}r\underline{o}pping \underline{T}ask-\underline{I}rrelevant \underline{N}odes) to reduce the size of graphs. Specifically, by introducing $K$ learnable virtual nodes to represent the graph embeddings targeted to $K$ different graph-level tasks, respectively, up to 90\% raw nodes with low attentiveness with an attention model -- a transformer in this paper, can be adaptively dropped without notable performance decreasing. Achieving almost the same accuracy, our method speeds up GAT by about 50\% on graph-level tasks including graph classification and graph edit distance (GED) with about 60\% less memory, on D\&D dataset. Code will be made publicly available in https://github.com/Sherrylone/DOTIN.

Motivation & Objective

• To address the scalability limitations of GNNs in graph-level tasks by reducing computational and memory costs.
• To overcome the shortcomings of existing pooling methods that prioritize local information and are task-agnostic.
• To enable adaptive, task-specific node pruning by learning global representations through virtual nodes.
• To improve efficiency and scalability in multi-task GNN settings without sacrificing performance.
• To provide a GNN-agnostic, parameter-efficient method for hierarchical graph size reduction.

Proposed method

• Introduces K learnable virtual nodes, each representing a task-specific global embedding for K different graph-level tasks.
• Uses a transformer-based attention mechanism to compute node attentiveness scores across all tasks.
• Identifies and drops nodes with low average attentiveness across tasks, defined as task-irrelevant.
• Applies node dropping hierarchically after each GNN layer, reducing graph size without additional GNN or clustering layers.
• Maintains model performance by preserving nodes with high task-specific attentiveness, especially in multi-task scenarios.
• Achieves O(N log N) complexity via sorting virtual node attentiveness scores, avoiding expensive clustering or additional GNN layers.

Experimental results

Research questions

• RQ1 Can task-irrelevant nodes be identified and dropped without degrading GNN performance in graph-level tasks?
• RQ2 How does task-specific node importance vary across different graph-level tasks like classification and graph edit distance?
• RQ3 Can a GNN-agnostic method achieve better efficiency and scalability than pooling-based approaches?
• RQ4 Does adaptive node dropping based on multi-task attention improve performance and reduce memory usage?
• RQ5 Can the method generalize to unseen tasks and datasets while maintaining efficiency and accuracy?

Key findings

• DOTIN reduces the number of nodes by up to 90% on the D&D dataset with no significant performance drop compared to the baseline.
• On the D&D dataset, DOTIN achieves approximately 50% faster training speed and 60% less memory usage with a GAT backbone.
• For hidden dimensions of 512 and 4096, DOTIN with 90% drop ratio uses only 42% and 57% of the memory, respectively, compared to the baseline.
• Visualizations show that node importance varies significantly between tasks—e.g., nodes important for graph classification are often irrelevant for GED, and vice versa.
• DOTIN outperforms pooling-based methods like gPool, which may drop task-relevant nodes due to their task-agnostic design.
• The method generalizes well across multiple datasets (D&D, PROTEINS, NCI1, NCI109, FRANKENSEIN), maintaining high accuracy across different backbones and drop ratios.