MLPInit：MLP的权重可以直接迁移到GNN上！一种简单的GNN训练加速方法

科技

2023-04-14 20:04

©PaperWeekly 原创 · 作者 | Xiaotian Han

单位 | 德州农工大学

研究方向 | 图神经网络，语言大模型

我们发现多层感知机（MLP）和图神经网络（GNN）的可训练的权重是可以共享的。所以我们想能不能使用训练好的 MLP 的权重直接应用到 GNN 上，因为 MLP 的训练非常快，如果权重可以直接迁移，那么就可以大大加速 GNN 的训练。

当我们把 MLP 训练好权重直接用到 GNN 上去做推理时，我们发现效果非常好。结果如下：

基于这个有意思的现象，我们提出了一种非常简单的 GNN 训练加速的方法， MLPInit。顾名思义， MLPInit 就是用训练好的 MLP 的权重去初始化 GNN，可以使得 GNN 又快又好地收敛，如下图的结果：

一直以来，在大型图上训练图神经网络（GNN）非常复杂且耗时。这其中有稀疏矩阵乘法比较慢的原因。来自德州农工大学，SnapChat 和莱斯大学的研究者提出一个极简的加速方法，使用 MLP 的训练来加速 GNN 的训练。文章发表在 ICLR2023 上。

论文题目：

MLPInit: Embarrassingly Simple GNN Training Acceleration with MLP Initialization

发表会议：

ICLR 2023

论文链接：

https://arxiv.org/abs/2210.00102

代码链接：

https://github.com/snap-research/MLPInit-for-GNNs

PPT链接：

https://ahxt.github.io/files/mlpinit_slides.pdf

为了方便研究者试用研究我们发现的现象，我们在开源的代码库中提供了几个独立的 Jupyter Notebooks，展示了 MLP 和 GNN 权重的迁移对比和 MLPInit 在 ogb-products 数据集上的结果。欢迎大家试用和反馈！

简介

在大型图上训练图神经网络（GNN）非常复杂且耗时。这归因于稀疏矩阵乘法引起的开销，在仅使用节点特征训练多层感知器（MLP）时则避开了这些开销。通过忽略图上下文，MLP 在图数据上简单且更快，但通常牺牲了预测准确性，限制了它们在图数据的应用。

我们发现，对于大多数基于消息传递的 GNN，我们可以通过设置相同形状的可训练参数，轻易地推导出一个具有等效可训练权重空间的 MLP（称之为 PeerMLP），这让我们好奇，直接使用已经完全训练的 PeerMLP 权重的 GNN 推理表现如何？令人惊讶的是，我们发现使用这种权重初始化的 GNN 显著优于它们的 PeerMLP，激发我们将 PeerMLP 训练作为 GNN 训练的初始化步骤。

为此，我们提出了一种非常简单，但效果极佳的 GNN 训练加速初始化方法，称为 MLPInit。我们在多个大型图数据集上的广泛实验验证了 MLPInit 可以加速 GNN 的训练（在 OGBN-Products 上加速高达 33 倍）并通常提高预测性能（例如，在 7 个数据集上的 GraphSAGE 节点分类上提高了高达 7.97%，在 4 个数据集上的链路预测上的 Hits@10 提高了高达 17.81%）。

主要贡献

我们首次进行了实验研究，通过实验观察揭示了 MLP 和 GNN 在可训练权重空间方面的关系：（i）GNN 和 MLP 具有相同的可训练的权重空间。（ii）GNN 可以通过训练其 PeerMLP 的权重进行优化。（iii）在节点分类任务上，直接使用 PeerMLP 的收敛权重的 GNN 的推理表现令人惊讶地优于其收敛 PeerMLP 的表现。

基于上述观察，我们提出了一种及其简单且非常有效的初始化方法来加速 GNN 训练。我们提出的方法 MLPInit 用收敛的 PeerMLP 的权重初始化 GNN 的权重。初始化后，我们观察到，与随机初始化相比，GNN 训练需要不到一半的周期就可以收敛。由于训练 MLP 比训练 GNN 更便宜且更快，所以 MLPInit 能够加速 GNN 的训练。

在多个大型图上的广泛实验结果验证了 MLPInit 可以加速 GNN 的训练（在 OGB-products 上加速高达 33 倍），同时通常可以提高模型性能（例如，在 GraphSAGE 上的节点分类上提高了 7.97%，在 Hits@10 上的链路预测上提高了 17.81%）。

MLPInit 及其容易实现，与传统的 GNN 训练方法相比，几乎没有额外的计算开销。此外，它与其他 GNN 加速方法（例如权重量化和图粗化）正交，进一步增加了 GNN 训练加速的空间。

方法：MLPInit

MLPInit 的基本思想非常简单：我们采用收敛的 PeerMLP 的权重来初始化 GNN，然后对 GNN 进行微调。具体而言，

1. 构建对应的 PeerMLP：为待训练的 GNN 创建一个具有相同可训练权重的 MLP；

2. 训练 PeerMLP 至收敛；

3. 用 PeerMLP 的收敛权重来初始化 GNN；

4. 对 GNN 进行微调。

我们方法的 python 风格的算法流程图如下：

实验

我们做了大量的实验来验证MLPInit的有效性。

第一，我们对比了随机初始化的 GNN 和 MLPInit 的 GNN 的训练速度。我们计算了随机初始化的 GNN 达到最佳测试性能所需的训练周期。结果显示，MLPInit 可以显著减少 GNN 的训练时间。下表显示 MLPInit 通常可以使 GNN 的训练速度提高 2-5 倍，甚至在某些情况下超过 30 倍。

第二，我们进行了实验，以展示所提出方法在节点分类和链接预测任务上收敛的 GNN 模型性能方面的优越性。下面两个表展示了 MLPInit 与随机初始化相比在节点分类和链接预测方面的性能提升。大多数情况下，MLPInit 提高了节点分类和链接预测任务的预测性能。

第三：我们通过实验来分析经过微调的 GNN 模型的收敛性。MLPInit 在 Loss Landscape 中为 GNNs 找到了更大的低损失区域。Loss Landscape 显示，MLPInit 的低损失区域在相同损失水平上更大，这表明使用 MLPInit 训练的模型的损失景观比随机初始化的模型具有更大的低损失区域。总之，MLPInit 帮助在 Loss Landscape 中为 GNNs 找到了更大的低损失区域。