可扩展GNN网络结构搜索系统(PaSca论文解读)

作者：于璠

文章来源：https://zhuanlan.zhihu.com/p/525292195

基于消息传播机制的GNN网络，在消息聚合阶段需要获取邻居节点特征进行计算，这一过程随着数据规模的扩大，内存读写和数据通信增加导致性能变差，分布式并行的加速比远达不到线性比，随着模型聚合层数的增加，性能更差。关于GNN模型结构在数据规模和聚合层数的扩展性问题，有研究尝试将消息传播机制计算流程进行拆解，将基于图结构的聚合计算作为采样部分，消息更新预测作为训练部分并设计了单个GNN网络例如APPNP，并为了模型的准确性进行了精心的结构设计，但没有提出通用的解法，今年WWW的最佳学生论文尝试从GNN模型结构范式和系统的角度解决可扩展GNN网络结构设计问题，提出了针对可扩展GNN网络结构的搜索系统，包含可扩展GNN结构范式和多目标GNN结构搜索系统两个部分，在15w的搜索空间，搜索出的结构在工业级大数据集上相比于SOTA精度提升了0.4%，分布式并行训练加速了28倍。

论文：PaSca: a Graph Neural Architecture Search System under the Scalable Paradigm

链接：https://arxiv.org/pdf/2203.00638.pdf

可扩展GNN结构范式

本文将消息传播机制的流程拆解为前处理、模型训练、后处理三个过程。前处理进行基于图结构的数据聚合，获取邻域节点的特征进行聚合，这里将不再限于获取一跳邻居，而是可以通过randomwalk、pagerank、自适应等多种方式在大范围内选取邻域特征进行聚合，可以实现低阶邻居信息与高阶邻居信息的平衡。模型训练将对前处理后的邻域特征进行结合和转换生成每个节点特征向量，可以不进行多步邻域聚合特征结合直接取最后一步聚合特征，这里转换采用多层MLP。后处理的思路类似标签传播，将模型训练后的特征再次进行邻域聚合生成预测特征。不同于消息传播机制，这一范式特征聚合与参数训练相互独立，训练过程可以具有很好的扩展性和加速，此外不同于很多GNN网络仅取用最后一次聚合特征，邻域特征结合机制可以实现低阶和扩展邻域信息的结合。而后处理过程则被论文实验证明可以提升模型阶段。

结构搜索引擎

针对每一个过程，定义步数和算法的选择范围就构成了可扩展GNN搜索空间，整个搜索系统分成搜索引擎和验证引擎两部分。搜索部分采用了基于贝叶斯优化的多目标搜索算法，Suggestion Server基于历史观察数据建模模型结构与目标之间的关系，从搜索空间随机采样的样本中预测最大化增量值EHVI的模型结构提供给验证引擎。验证引擎对样本模型结构进行训练验证，反馈验证结果观察指标给搜索引擎更新历史观察数据。为了加速搜索过程，验证引擎分为图数据聚合和模型训练两部分并进行分布式存储和计算。前处理和后处理复用图数据聚合计算区别仅在于输入节点特征不同，图数据被切分为batch交给worker并行计算，每一步的计算结果分布式存储在worker上，下一步从其中拉取一跳结果作为输入，因为可以分布式并行，大大减少了计算时间。模型训练部分同其他深度学习基于参数服务器的分布式训练。

结果

从实验结果可以看到搜索出满足可扩展架构范式的GNN模型，在分布式训练上具有很好的扩展性，当worker数增加时，能获得近乎理想的线性加速，因而更适用于大规模图数据训练。

此外在精度上，搜索模型相比于最好的baseline略有提升。在模型扩展性上，搜索模型更优。随着模型加深即聚合次数增加，由于over-smoothing，超过两层时，GCN等模型精度呈下降趋势，而可扩展搜索模型随着聚合次数增加，一定范围精度上升并保持不降。

总结

优点：从GNN模型结构范式出发，将邻域节点特征获取聚合与特征转换等消息更新过程相互交替的依赖过程解耦，提出一种具有良好扩展性GNN结构范式。针对新的范式，设计了一个模型结构搜索系统，搜索出了满足多种目标的可扩展GNN结构，相比于现有人工设计结构模型精度略有提升，但是表现出了良好的分布式计算加速和模型扩展性。

不足：论文中模型搜索进行了2000次验证，虽然没有提及具体耗时但是这一过程资源消耗并不少。此外搜索出新范式的模型相比于SOTA模型层数增加了不少并没有显著的精度提升，将图进行切分以及将邻域特征聚合作为前处理不进行参数训练都将造成精度损失，新范式的模型设计也不容易。