[论文解读] MLPInit: Embarrassingly Simple GNN Training Acceleration with MLP Initialization

Xiaotian Han, Tong Zhao+3 more|arXiv (Cornell University)|Sep 30, 2022

Advanced Graph Neural Networks被引用 12

一句话总结

MLPInit 提出了一种简单而有效的方法，通过使用具有相同架构和权重空间的预训练 PeerMLP 权重来初始化图神经网络（GNN），从而加速 GNN 的训练。该方法在多个基准测试中，将训练时间最多缩短 33 倍，并在节点分类任务上提升预测准确率最多 7.97%，在链接预测任务上提升最多 17.81%。

ABSTRACT

Training graph neural networks (GNNs) on large graphs is complex and extremely time consuming. This is attributed to overheads caused by sparse matrix multiplication, which are sidestepped when training multi-layer perceptrons (MLPs) with only node features. MLPs, by ignoring graph context, are simple and faster for graph data, however they usually sacrifice prediction accuracy, limiting their applications for graph data. We observe that for most message passing-based GNNs, we can trivially derive an analog MLP (we call this a PeerMLP) with an equivalent weight space, by setting the trainable parameters with the same shapes, making us curious about extbf{\emph{how do GNNs using weights from a fully trained PeerMLP perform?}} Surprisingly, we find that GNNs initialized with such weights significantly outperform their PeerMLPs, motivating us to use PeerMLP training as a precursor, initialization step to GNN training. To this end, we propose an embarrassingly simple, yet hugely effective initialization method for GNN training acceleration, called MLPInit. Our extensive experiments on multiple large-scale graph datasets with diverse GNN architectures validate that MLPInit can accelerate the training of GNNs (up to 33X speedup on OGB-Products) and often improve prediction performance (e.g., up to $7.97\%$ improvement for GraphSAGE across $7$ datasets for node classification, and up to $17.81\%$ improvement across $4$ datasets for link prediction on metric Hits@10). The code is available at \href{https://github.com/snap-research/MLPInit-for-GNNs}.

研究动机与目标

• 解决大规模图上 GNN 训练计算成本高和耗时长的问题。
• 探究消息传递 GNN 与 MLP 的权重空间是否可以实现等价。
• 探索是否可以先训练 PeerMLP，然后使用其收敛后的权重来初始化 GNN，以加速训练并提升性能。
• 开发一种简单、即插即用的初始化方法，无需架构修改或超参数调优，仅需标准 GNN 训练流程。
• 证明使用 PeerMLP 权重初始化的 GNN 在性能上优于随机初始化的 GNN 及其 PeerMLP 对应模型。

提出的方法

• 将 PeerMLP 定义为与给定 GNN 具有相同架构和权重维度的 MLP，共享相同的可学习参数空间。
• 仅使用节点特征在相同图数据上训练 PeerMLP，绕过昂贵的稀疏消息传递，实现更快收敛。
• 将 PeerMLP 收敛后的权重用作对应 GNN 的初始权重。
• 从该初始化状态开始微调 GNN，利用预学习的特征表示加速收敛。
• 在多种 GNN 架构（如 GraphSAGE、GCN、ClusterGCN）和大规模数据集（如 OGB-products、Reddit、Flickr）上应用该方法。
• 采用标准训练协议，结合早停和基于验证的模型选择，确保与随机初始化的公平比较。

实验结果

研究问题

• RQ1 GNN 及其对应 PeerMLP 的权重空间在形状和维度上是否可以完全一致？
• RQ2 仅使用节点特征训练 PeerMLP 是否能学习到一种表示，当迁移至 GNN 后，其下游性能优于随机初始化？
• RQ3 PeerMLP 的收敛权重是否可作为 GNN 的有效初始化，从而减少达到收敛所需的训练轮次？
• RQ4 该初始化方法是否在多种 GNN 架构和大规模图数据集上均能持续加速训练？
• RQ5 使用 PeerMLP 权重初始化的 GNN 在节点分类和链接预测任务上的性能是否优于 PeerMLP 本身？

主要发现

• 与随机初始化相比，MLPInit 在 OGB-products 数据集上将 GNN 训练时间最多缩短 33 倍。
• 使用 MLPInit 的 GraphSAGE 在七个数据集上的节点分类准确率最高提升 7.97%，相比随机初始化。
• 在链接预测任务中，MLPInit 在四个数据集上的 Hits@10 性能最高提升 17.81%。
• 使用 PeerMLP 权重初始化的 GNN 所需训练轮次少于随机初始化 GNN 的一半。
• PeerMLP 本身在使用其收敛权重初始化后，性能反而低于 GNN，表明 GNN 的消息传递机制能从迁移的表示中获益。
• 该方法在多种 GNN 架构（如 GraphSAGE、GCN、ClusterGCN、GraphSAINT）和数据集（如 Reddit、Flickr、OGB-arXiv、OGB-products）上均表现出鲁棒性。