[论文解读] DOTIN: Dropping Task-Irrelevant Nodes for GNNs

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

Advanced Graph Neural Networks被引用 4

一句话总结

DOTIN 提出了一种新颖的、任务感知的图剪枝方法，通过利用 K 个可学习的虚拟节点来捕捉特定任务的全局表征，自适应地丢弃高达 90% 的与任务无关的节点。通过使用基于 Transformer 的注意力机制对节点重要性进行排序，DOTIN 在 D&D 数据集上实现了约 50% 的训练加速和约 60% 的内存减少，且性能无损失，优于基于池化的基线方法，在单任务和多任务设置下均表现更优。

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.

研究动机与目标

• 通过降低计算和内存成本，解决 GNN 在图级任务中的可扩展性限制。
• 克服现有池化方法优先关注局部信息且与任务无关的缺点。
• 通过虚拟节点学习全局表征，实现自适应的、任务特定的节点剪枝。
• 在不牺牲性能的前提下，提升多任务 GNN 设置下的效率和可扩展性。
• 提供一种与 GNN 无关、参数高效的分层图规模缩减方法。

提出的方法

• 引入 K 个可学习的虚拟节点，每个节点代表 K 个不同图级任务的特定任务全局嵌入。
• 使用基于 Transformer 的注意力机制，计算所有任务中节点的注意力得分。
• 识别并丢弃在所有任务中平均注意力得分较低的节点（即与任务无关的节点）。
• 在每个 GNN 层后分层应用节点丢弃，减少图规模，且无需额外的 GNN 或聚类层。
• 通过保留具有高任务特定注意力得分的节点来维持模型性能，尤其在多任务场景中表现更优。
• 通过排序虚拟节点注意力得分实现 O(N log N) 的时间复杂度，避免了昂贵的聚类或额外的 GNN 层。

实验结果

研究问题

• RQ1 是否可以在不降低图级任务中 GNN 性能的前提下，识别并丢弃与任务无关的节点？
• RQ2 在图分类和图编辑距离等不同图级任务中，节点的重要性如何变化？
• RQ3 一种与 GNN 无关的方法是否能比基于池化的方案实现更高的效率和可扩展性？
• RQ4 基于多任务注意力的自适应节点丢弃是否能提升性能并减少内存使用？
• RQ5 该方法是否能在未见过的任务和数据集上保持效率和准确率？

主要发现

• 在 D&D 数据集上，DOTIN 丢弃高达 90% 的节点，与基线相比性能无显著下降。
• 在 D&D 数据集上，DOTIN 使用 GAT 主干网络时，训练速度提升约 50%，内存使用减少约 60%。
• 对于隐藏维度为 512 和 4096 的情况，DOTIN 在 90% 丢弃率下分别仅使用基线 42% 和 57% 的内存。
• 可视化结果表明，不同任务中节点的重要性存在显著差异——例如，对图分类重要的节点通常对 GED 无关，反之亦然。
• DOTIN 优于基于池化的方法（如 gPool），后者因设计上与任务无关，可能错误丢弃与任务相关的节点。
• 该方法在多个数据集（D&D、PROTEINS、NCI1、NCI109、FRANKENSEIN）上泛化良好，无论主干网络或丢弃率如何，均保持高准确率。