[论文解读] Efficient Graph Neural Network Inference at Large Scale

Xinyi Gao, Wentao Zhang+4 more|arXiv (Cornell University)|Nov 1, 2022

Advanced Graph Neural Networks被引用 4

一句话总结

本文提出节点自适应推理（NAI），一种即插即用的框架，通过基于拓扑结构动态调整每个节点的传播顺序，实现大规模图神经网络推理的加速，减少冗余计算。结合膨胀蒸馏（Inception Distillation）实现多尺度特征学习，NAI 在 ogbn-products 数据集上实现了高达 75 倍的加速，同时在不同延迟约束下保持高精度。

ABSTRACT

Graph neural networks (GNNs) have demonstrated excellent performance in a wide range of applications. However, the enormous size of large-scale graphs hinders their applications under real-time inference scenarios. Although existing scalable GNNs leverage linear propagation to preprocess the features and accelerate the training and inference procedure, these methods still suffer from scalability issues when making inferences on unseen nodes, as the feature preprocessing requires the graph is known and fixed. To speed up the inference in the inductive setting, we propose a novel adaptive propagation order approach that generates the personalized propagation order for each node based on its topological information. This could successfully avoid the redundant computation of feature propagation. Moreover, the trade-off between accuracy and inference latency can be flexibly controlled by simple hyper-parameters to match different latency constraints of application scenarios. To compensate for the potential inference accuracy loss, we further propose Inception Distillation to exploit the multi scale reception information and improve the inference performance. Extensive experiments are conducted on four public datasets with different scales and characteristics, and the experimental results show that our proposed inference acceleration framework outperforms the SOTA graph inference acceleration baselines in terms of both accuracy and efficiency. In particular, the advantage of our proposed method is more significant on larger-scale datasets, and our framework achieves $75 imes$ inference speedup on the largest Ogbn-products dataset.

研究动机与目标

• 解决图神经网络（GNN）在实时、延迟敏感应用场景（如推荐系统和欺诈检测）中推理的可扩展性挑战。
• 克服现有线性传播 GNN 的局限性，即需要固定预计算图结构，无法高效处理推理阶段的未见节点。
• 通过基于节点拓扑上下文动态确定最优传播深度，减少特征传播中的冗余计算。
• 通过可调超参数实现推理延迟与精度之间的灵活权衡。
• 通过一种新颖的多尺度知识蒸馏技术，补偿因传播提前终止导致的精度下降。

提出的方法

• 提出节点自适应推理（NAI），一种即插即用的方法，基于节点到特征平滑稳定状态的距离，为每个节点计算个性化的传播顺序。
• 使用一个简单的全局超参数，通过确定每个节点何时停止特征传播，控制推理速度与精度之间的权衡。
• 引入膨胀蒸馏（Inception Distillation），一种多尺度知识蒸馏模块，利用来自多个传播深度的特征，扩展学生模型的感受野。
• 采用在线蒸馏，使用多个教师模型（如 f^(1), f^(2), f^(3)）的软标签来提升学生模型性能，而无需重新训练整个模型。
• 设计一种无参数、与图无关的特征传播机制，避免需要联合训练多个退出头，从而提升训练稳定性。
• 利用拓扑信息（如节点度数、邻域结构）指导提前终止，避免计算机视觉/NLP 中常见的重复预测问题。

实验结果

研究问题

• RQ1 自适应、节点特异的传播深度是否能在不牺牲模型精度的前提下减少 GNN 推理中的冗余计算？
• RQ2 基于拓扑结构的动态传播顺序与固定顺序传播相比，在推理效率和精度方面表现如何？
• RQ3 多尺度知识蒸馏在多大程度上可以缓解因特征传播提前终止导致的性能下降？
• RQ4 所提方法是否能在大规模真实图上实现显著加速，同时在不同延迟约束下保持高精度？
• RQ5 该框架在具有不同规模和拓扑特性的多样化图数据集上是否具备良好的泛化能力？

主要发现

• NAI 在 ogbn-products 数据集上实现了高达 75 倍的推理加速，证明了其在最大规模图上的有效性。
• 该方法在所有数据集上均保持高分类精度，通过膨胀蒸馏显著最小化性能下降，尤其在使用软标签和 3 个教师模型的集成时效果更优。
• 膨胀蒸馏在不同超参数设置下稳定了性能表现，当 λ ∈ [0.5, 1] 时标签软化效果最佳。
• 与 SOTA 基线相比，该框架在精度和效率方面均表现出更优性能，且在更大图上优势更加明显。
• 自适应传播机制通过避免不必要的传播步骤，减少了冗余计算，尤其使高阶数节点（易出现过平滑）受益。
• 该方法灵活且可部署于真实世界的归纳设置中，可在无需固定图结构的前提下，实现对未见节点的高效在线推理。