arXiv:2604.03815cs.LGcs.AI2026-04中稿 · ICLR

提出k-MIP注意力机制,让图Transformer高效处理超大规模图。

k-Maximum Inner Product Attention for Graph Transformers and the Expressive Power of GraphGPS

  • 用top-k选关键节点,实现稀疏灵活的注意力模式
  • 内存线性增长,处理超50万节点图仅需单张A100 GPU
  • 理论证明可逼近全注意力模型,适合大规模图任务

图Transformer在克服传统图神经网络的过度压缩和长程依赖建模困难方面展现出潜力,但其在大规模图上的应用受限于全连接注意力机制带来的二次方内存与计算开销。尽管已有线性注意力和受限注意力等替代方案,但往往导致性能下降或表达能力受限。为此,本文提出k-最大内积(k-MIP)注意力机制,通过top-k操作为每个查询节点选择最相关的键节点,形成稀疏且灵活的注意力模式。结合符号矩阵计算注意力得分,实现线性内存复杂度,并在实际中达到比全连接注意力快一个数量级的加速效果,使单张A100 GPU即可处理超过50万节点的图。我们还提供了理论分析,证明k-MIP Transformer可任意精度逼近全注意力变压器。此外,我们分析了集成该机制的GraphGPS框架的表达能力,建立了其在S-SEG-WL测试下的图区分能力上界。最后,在长程图基准、城市网络基准以及两个自定义的大规模归纳点云数据集上验证,本方法始终位列可扩展图变压器的顶尖表现之列。

原文摘要 · Abstract (English)

Graph transformers have shown promise in overcoming limitations of traditional graph neural networks, such as oversquashing and difficulties in modeling long-range dependencies. However, their application to large-scale graphs is hindered by the quadratic memory and computational complexity of the all-to-all attention mechanism. Although alternatives such as linearized attention and restricted attention patterns have been proposed, these often degrade performance or limit expressive power. To better balance efficiency and effectiveness, we introduce k-Maximum Inner Product (k-MIP) attention for graph transformers. k-MIP attention selects the most relevant key nodes per query via a top-k operation, yielding a sparse yet flexible attention pattern. Combined with an attention score computation based on symbolic matrices, this results in linear memory complexity and practical speedups of up to an order of magnitude compared to all-to-all attention, enabling the processing of graphs with over 500k nodes on a single A100 GPU. We provide a theoretical analysis of expressive power, showing that k-MIP attention does not compromise the expressiveness of graph transformers: specifically, we prove that k-MIP transformers can approximate any full-attention transformer to arbitrary precision. In addition, we analyze the expressive power of the GraphGPS framework, in which we integrate our attention mechanism, and establish an upper bound on its graph distinguishing capability in terms of the S-SEG-WL test. Finally, we validate our approach on the Long Range Graph Benchmark, the City-Networks benchmark, and two custom large-scale inductive point cloud datasets, consistently ranking among the top-performing scalable graph transformers.

图Transformer注意力机制大规模图表达能力

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