[论文解读] MAG-GNN: Reinforcement Learning Boosted Graph Neural Network

Lecheng Kong, Jiarui Feng+4 more|arXiv (Cornell University)|Oct 29, 2023

Advanced Graph Neural Networks被引用 5

一句话总结

MAG-GNN 提出了一种基于强化学习（RL）的方法，用于为图神经网络（GNNs）识别最小且具有表现力的子图子集，从而替代计算成本高昂的全部子图枚举。通过使用强化学习智能体迭代选择高区分性的子图，MAG-GNN 在多个图基准测试中实现了最先进性能，同时相比子图 GNN 显著降低了推理时间。

ABSTRACT

While Graph Neural Networks (GNNs) recently became powerful tools in graph learning tasks, considerable efforts have been spent on improving GNNs' structural encoding ability. A particular line of work proposed subgraph GNNs that use subgraph information to improve GNNs' expressivity and achieved great success. However, such effectivity sacrifices the efficiency of GNNs by enumerating all possible subgraphs. In this paper, we analyze the necessity of complete subgraph enumeration and show that a model can achieve a comparable level of expressivity by considering a small subset of the subgraphs. We then formulate the identification of the optimal subset as a combinatorial optimization problem and propose Magnetic Graph Neural Network (MAG-GNN), a reinforcement learning (RL) boosted GNN, to solve the problem. Starting with a candidate subgraph set, MAG-GNN employs an RL agent to iteratively update the subgraphs to locate the most expressive set for prediction. This reduces the exponential complexity of subgraph enumeration to the constant complexity of a subgraph search algorithm while keeping good expressivity. We conduct extensive experiments on many datasets, showing that MAG-GNN achieves competitive performance to state-of-the-art methods and even outperforms many subgraph GNNs. We also demonstrate that MAG-GNN effectively reduces the running time of subgraph GNNs.

研究动机与目标

• 解决子图 GNN 因枚举所有可能子图以提升表达能力而带来的高计算成本问题。
• 探究一个小型、最优的子图子集是否可实现与完整子图枚举相当的表达能力。
• 开发一种高效且可扩展的方法，在无需全面搜索的情况下保持高 GNN 表达能力。
• 在保持或提升预测性能的同时，降低子图 GNN 的推理时间。

提出的方法

• MAG-GNN 从输入图中随机采样一组根节点子图作为候选集。
• 训练一个强化学习智能体，通过迭代替换候选集中的子图以最大化基于模型性能的奖励信号。
• 智能体使用 Q 表来估计用新候选子图替换目标子图的预期奖励。
• 子图替换过程由一个策略指导，该策略在每一步选择预期奖励最高的子图更新。
• 最终选定的子图集合被用作下游 GNN 的输入以执行预测任务。
• 该方法将指数级子图枚举复杂度降低为常数步长的强化学习搜索，从而实现高效且具有表现力的 GNN。

实验结果

研究问题

• RQ1 一个小型、经过优化选择的子图子集是否可实现与完整子图枚举在 GNN 中相当的表达能力？
• RQ2 强化学习能否有效识别最具区分性的子图以提升 GNN 性能？
• RQ3 将子图枚举简化为常数步长的强化学习过程，是否能在降低运行时间的同时保持或提升 GNN 性能？
• RQ4 MAG-GNN 在准确率和效率方面与最先进子图 GNN 相比表现如何？

主要发现

• MAG-GNN 实现了与最先进方法相当的性能，在包括 ZINC 和 OGBG-MOLHIV 在内的多个数据集上优于多个子图 GNN。
• 在 ZINC 数据集上，MAG-GNN 实现了 0.106 ± 0.014 的测试损失，显著优于 GIN（0.163 ± 0.004），并接近更复杂模型 CIN（0.079 ± 0.006）的性能。
• 通过预训练（MAG-GNN-PRE），模型性能进一步提升，在 ZINC 上实现 0.096 ± 0.009 的测试损失，在 OGBG-MOLHIV 上实现 78.30 ± 1.08 的准确率。
• MAG-GNN 显著降低了推理时间，比基于边根的子图 GNN（如 I²GNN）更快，并且在准确率上优于 NGNN，尽管其效率相近。
• 该模型对不同基础 GNN 具有适应性，表明通过架构整合可进一步提升性能。
• 消融实验确认，MAG-GNN 基于强化学习的子图选择方法能有效平衡表达能力与效率，避免了固定结构子图 GNN 中常见的过拟合问题。