arXiv:2412.01519cs.LG2024-12

用动态重分配机制让图注意力模型实现线性复杂度,同时保持高性能。

ReHub: Linear Complexity Graph Transformers with Adaptive Hub-Spoke Reassignment

  • 通过虚拟枢纽动态重分配节点,减少计算冗余。
  • 在LRGB数据集上超越基线方法,性能接近全连接模型。
  • 适合大规模图数据的高效建模,尤其对资源受限场景友好。

我们提出ReHub,一种新型图变换器架构,通过节点与虚拟节点间的高效重分配技术实现线性复杂度。图变换器因其能显式处理长程节点通信,在图学习中日益重要,可缓解消息传递图网络中的过度平滑和过度压缩问题。然而,其密集注意力机制随节点数呈二次增长,限制了在大规模图上的应用。ReHub借鉴航空业的枢纽-航线模式,将图节点(支线)在每层动态重分配至固定数量的虚拟节点(枢纽)。近期工作Neural Atoms(Li et al., 2024)已证明使用虚拟节点可显著提升性能;其结果表明枢纽数量对性能影响显著。但增加枢纽数通常会提高复杂度,需权衡以维持线性复杂度。我们的关键洞察是:每个节点只需与少量枢纽交互即可达到线性复杂度,即使枢纽总数较大。为此,我们提出基于枢纽间相似度的自适应重分配策略,无需额外节点-枢纽计算开销。实验表明,在LRGB数据集上,ReHub持续优于基线Neural Atoms,且保持线性复杂度;其稀疏模型性能与非稀疏版本相当。此外,ReHub在多个基准测试中表现优异,始终位列前列。

原文摘要 · Abstract (English)

We present ReHub, a novel graph transformer architecture that achieves linear complexity through an efficient reassignment technique between nodes and virtual nodes. Graph transformers have become increasingly important in graph learning for their ability to utilize long-range node communication explicitly, addressing limitations such as oversmoothing and oversquashing found in message-passing graph networks. However, their dense attention mechanism scales quadratically with the number of nodes, limiting their applicability to large-scale graphs. ReHub draws inspiration from the airline industry's hub-and-spoke model, where flights are assigned to optimize operational efficiency. In our approach, graph nodes (spokes) are dynamically reassigned to a fixed number of virtual nodes (hubs) at each model layer. Recent work, Neural Atoms (Li et al., 2024), has demonstrated impressive and consistent improvements over GNN baselines by utilizing such virtual nodes; their findings suggest that the number of hubs strongly influences performance. However, increasing the number of hubs typically raises complexity, requiring a trade-off to maintain linear complexity. Our key insight is that each node only needs to interact with a small subset of hubs to achieve linear complexity, even when the total number of hubs is large. To leverage all hubs without incurring additional computational costs, we propose a simple yet effective adaptive reassignment technique based on hub-hub similarity scores, eliminating the need for expensive node-hub computations. Our experiments on LRGB indicate a consistent improvement in results over the base method, Neural Atoms, while maintaining a linear complexity. Remarkably, our sparse model achieves performance on par with its non-sparse counterpart. Furthermore, ReHub outperforms competitive baselines and consistently ranks among top performers across various benchmarks.

图神经网络线性复杂度注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。