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

Jonas De Schouwer, Haitz Sáez de Ocáriz Borde, Xiaowen Dong
Proceedings of GRaM: the Second Edition of the Workshop on Geometry-grounded Representation Learning and Generative Modeling, PMLR 326:145-183, 2026.

Abstract

Graph transformers have shown promise in overcoming limitations of traditional graph neural networks, such as oversquashing and difficulties in modelling longrange 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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v326-de-schouwer26a, title = {k-{M}aximum {I}nner {P}roduct {A}ttention for {G}raph {T}ransformers and the {E}xpressive {P}ower of {G}raphGPS}, author = {De Schouwer, Jonas and S\'{a}ez de Oc\'{a}riz Borde, Haitz and Dong, Xiaowen}, booktitle = {Proceedings of GRaM: the Second Edition of the Workshop on Geometry-grounded Representation Learning and Generative Modeling}, pages = {145--183}, year = {2026}, editor = {Pouplin, Alison and Vadgama, Sharvaree and Bekkers, Erik and Kaba, Sékou-Oumar and Lawrence, Hannah and Lecha, Manuel and Baker, Elizabeth and Suk, Julian and Walters, Robin and Tomczak, Jakub and Jegelka, Stefanie}, volume = {326}, series = {Proceedings of Machine Learning Research}, month = {26 Apr}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v326/main/assets/de-schouwer26a/de-schouwer26a.pdf}, url = {https://proceedings.mlr.press/v326/de-schouwer26a.html}, abstract = {Graph transformers have shown promise in overcoming limitations of traditional graph neural networks, such as oversquashing and difficulties in modelling longrange 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.} }
Endnote
%0 Conference Paper %T k-Maximum Inner Product Attention for Graph Transformers and the Expressive Power of GraphGPS %A Jonas De Schouwer %A Haitz Sáez de Ocáriz Borde %A Xiaowen Dong %B Proceedings of GRaM: the Second Edition of the Workshop on Geometry-grounded Representation Learning and Generative Modeling %C Proceedings of Machine Learning Research %D 2026 %E Alison Pouplin %E Sharvaree Vadgama %E Erik Bekkers %E Sékou-Oumar Kaba %E Hannah Lawrence %E Manuel Lecha %E Elizabeth Baker %E Julian Suk %E Robin Walters %E Jakub Tomczak %E Stefanie Jegelka %F pmlr-v326-de-schouwer26a %I PMLR %P 145--183 %U https://proceedings.mlr.press/v326/de-schouwer26a.html %V 326 %X Graph transformers have shown promise in overcoming limitations of traditional graph neural networks, such as oversquashing and difficulties in modelling longrange 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.
APA
De Schouwer, J., Sáez de Ocáriz Borde, H. & Dong, X.. (2026). k-Maximum Inner Product Attention for Graph Transformers and the Expressive Power of GraphGPS. Proceedings of GRaM: the Second Edition of the Workshop on Geometry-grounded Representation Learning and Generative Modeling, in Proceedings of Machine Learning Research 326:145-183 Available from https://proceedings.mlr.press/v326/de-schouwer26a.html.

Related Material