[论文解读] A Generalization of ViT/MLP-Mixer to Graphs

Xiaoxin He, Bryan Hooi+4 more|arXiv (Cornell University)|Dec 27, 2022

Advanced Graph Neural Networks被引用 7

一句话总结

本文提出 Graph ViT/MLP-Mixer，一种新颖的图神经网络架构，通过将图划分为图块并应用 token 混合与通道混合机制，将视觉 Transformer 和 MLP-Mixer 通用化至图结构数据。该方法在节点与边数上实现线性时间复杂度，缓解了过度压缩问题，并在长距离依赖与表达能力基准测试中优于消息传递 GNN 和图 Transformer。

ABSTRACT

Graph Neural Networks (GNNs) have shown great potential in the field of graph representation learning. Standard GNNs define a local message-passing mechanism which propagates information over the whole graph domain by stacking multiple layers. This paradigm suffers from two major limitations, over-squashing and poor long-range dependencies, that can be solved using global attention but significantly increases the computational cost to quadratic complexity. In this work, we propose an alternative approach to overcome these structural limitations by leveraging the ViT/MLP-Mixer architectures introduced in computer vision. We introduce a new class of GNNs, called Graph ViT/MLP-Mixer, that holds three key properties. First, they capture long-range dependency and mitigate the issue of over-squashing as demonstrated on Long Range Graph Benchmark and TreeNeighbourMatch datasets. Second, they offer better speed and memory efficiency with a complexity linear to the number of nodes and edges, surpassing the related Graph Transformer and expressive GNN models. Third, they show high expressivity in terms of graph isomorphism as they can distinguish at least 3-WL non-isomorphic graphs. We test our architecture on 4 simulated datasets and 7 real-world benchmarks, and show highly competitive results on all of them. The source code is available for reproducibility at: \url{https://github.com/XiaoxinHe/Graph-ViT-MLPMixer}.

研究动机与目标

• 解决消息传递 GNN 的局限性，特别是过度压缩与长距离依赖能力差的问题。
• 在保持或提升表达能力的同时，克服图 Transformer 的二次方时间复杂度。
• 开发一种图神经网络架构，能够高效捕捉长距离依赖，而无需依赖堆叠的消息传递层。
• 在 1-WL 基础上增强图同构表达能力，以区分更多复杂的非同构图。
• 在多种图基准测试（包括大规模与稀疏图）中实现高性能，同时保持线性时间与内存复杂度。

提出的方法

• 使用 METIS 图划分算法将图划分为 P 个图块，通过一跳重叠以保持连通性。
• 对每个图块应用基础 GNN（如 GCN、GIN）计算图块嵌入，相比标准 GNN 仅引入常数级额外开销。
• 应用 token 混合层，通过跨图块注意力或前馈操作建模图块间的长距离依赖。
• 应用通道混合层以细化每个图块内的节点表示，实现特征层面的交互。
• 使用全局分类头聚合图块表示以完成下游任务。
• 借鉴 MLP-Mixer 架构，解耦空间（图块）与通道级特征学习，实现高效的全局建模。

实验结果

研究问题

• RQ1 能否有效将受 ViT/MLP-Mixer 启发的架构推广至图结构数据，以克服消息传递 GNN 的局限性？
• RQ2 所提出的 Graph ViT/MLP-Mixer 在长距离依赖建模方面是否优于标准 MP-GNN 与图 Transformer？
• RQ3 该模型是否能在保持高于 1-WL 与 2-WL GNN 的表达能力的同时，维持线性的时间与内存复杂度？
• RQ4 在包含长距离依赖与结构复杂性的基准测试（如 TreeNeighbourMatch 与 Long Range Graph Benchmark）中，模型表现如何？
• RQ5 模型性能在多大程度上依赖于图块数量与图划分策略？

主要发现

• Graph ViT/MLP-Mixer 在 Peptides-func 数据集上平均准确率提升 0.056，在 Peptides-struct 上 MAE 降低 0.028，表明其在长距离交互建模方面表现更优。
• 在 TreeNeighbourMatch 数据集中，模型可泛化至深度 r=7，而标准 MP-GNN 在 r=4 时即失效，表明其有效缓解了过度压缩问题。
• 模型性能与 SUN 相当（Peptides-func 上为 0.6730，Peptides-struct 上为 0.2498），但内存使用减少 44 倍，训练速度提升 19 倍，验证了其高效性。
• Graph MLP-Mixer 展现出高表达能力，可区分至少 3-WL 非同构图，优于标准 MP-GNN。
• 模型在时间与内存复杂度上均保持线性 O(N + E)，显著优于图 Transformer 与 SUN 等表达能力强的 GNN，在可扩展性方面表现更优。
• 在 4 个模拟与 7 个真实世界基准测试中的实证结果表明，模型在所有任务中均表现出高度竞争力，验证了该架构的鲁棒性与泛化能力。