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

cs.LGApr 4, 2026

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

Authors: Jonas De Schouwer, Haitz Sáez de Ocáriz Borde, Xiaowen Dong

Organizations: Stanford University · University of Oxford

Abstract

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.

Figures & tables

Graph Subsampling	Engineered Patterns	Linearized Attention	Learnt Patterns
Gophormer ( Zhao et al., 2021 )	Exphormer ( Shirzad et al., 2023 )	Nodeformer ( Wu et al., 2022 )	GOAT ( Kong et al., 2023 )
NAGphormer ( Chen et al., 2022b )	GPS+BigBird ( Rampášek et al., 2022 )	Difformer ( Wu et al., 2023 )	GPS+k-MIP (ours)
	SpExphormer ( Shirzad et al., 2024 )	SGFormer ( Wu et al., 2024 )
		GPS+Performer ( Rampášek et al., 2022 )

Table 1: Categorization of scalable graph transformer methods.

Figure 1: One layer in the GraphGPS framework. The dashed lines indicate residual connections. Based on Rampášek et al. (2022) .

Panel of Figure 2 — Figure 2: Comparison of full attention and k-MIP attention (with/without symbolic matrices). Shown is the mean $\pm$ std over 5 runs, measured on a single 40GB A100 GPU. Full attention gave OOM errors in the training setting for $N\geq 10^{5}$ . The data behind this figure can be found in Appendix J .

Figure 3: Runtime breakdown of full attention and k-MIP attention in the training setting at

N=10^{4.5}

, measured on a single 40GB A100 GPU. The left and middle bars compare full and k-MIP attention at the same scale; the right panel zooms into k-MIP attention.

	PascalVOC-SP (F1 $\uparrow$ )	COCO-SP (F1 $\uparrow$ )	Pept-Func (AP $\uparrow$ )	Pept-Struct (MAE $\downarrow$ )
GCN †	20.78 $\pm$ 0.31	13.38 $\pm$ 0.07	68.60 $\pm$ 0.50	0.2460 $\pm$ 0.0007
GINE †	27.18 $\pm$ 0.54	21.25 $\pm$ 0.09	66.21 $\pm$ 0.67	0.2473 $\pm$ 0.0017
GAT	27.15 $\pm$ 0.49	18.86 $\pm$ 0.11	67.87 $\pm$ 0.56	0.2488 $\pm$ 0.0018
GatedGCN †	38.80 $\pm$ 0.40	29.22 $\pm$ 0.18	67.65 $\pm$ 0.47	0.2477 $\pm$ 0.0009
GPS + BigBird	38.75 $\pm$ 0.64	35.25 $\pm$ 0.27	64.83 $\pm$ 0.73	0.2566 $\pm$ 0.0019
GPS + Performer	37.92 $\pm$ 1.13	27.70 $\pm$ 0.27	66.06 $\pm$ 0.64	0.2643 $\pm$ 0.0008

Table 2: Test performance of GPS+k-MIP and baselines on LRGB. Shown is the mean

\pm

std over four runs. The first , second , and third best results are highlighted. † Taken from ( Tönshoff et al., 2023 ) . ‡ Taken from ( Shirzad et al., 2023 ) .

Figure 4: Tradeoffs between training time and accuracy across different datasets in City-Networks ( Liang et al., 2025 ) . Shown is the mean

\pm

std over four runs, except for the London dataset where we only ran one run for the graph transformer models. GPS+BigBird was not evaluated on LA due to long training times. GPS+Transformer, Exphormer, and GPS+BigBird returned OOM on London.

	ShapeNet-Part (F1 $\uparrow$ )	S3DIS (mIoU $\uparrow$ )
GCN	60.18 $\pm$ 0.04	39.98 $\pm$ 1.47
GINE	64.57 $\pm$ 0.35	44.16 $\pm$ 0.62
GAT	63.01 $\pm$ 0.17	44.24 $\pm$ 1.14
GatedGCN	76.20 $\pm$ 0.32	63.71 $\pm$ 1.28
GPS + BigBird	79.65 $\pm$ 0.98	67.92 $\pm$ 0.91
GPS + Performer	77.36 $\pm$ 1.23	60.83 $\pm$ 0.56

Table 3: Test performance of GPS+k-MIP and baselines on ShapeNet-Part and S3DIS. Shown is the mean

\pm

std over three runs. The first , second , and third best results are highlighted.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 5 — Figure 5: Illustration of a single iteration of the Weisfeiler-Lehman test. The node labels (in $\mathcal{C}=\mathbb{N}$ ) are written inside the nodes. The mapping $\tau$ is written next to the nodes.

Panel of Figure 6 — Figure 6: Two graphs that are indistinguishable by the 1-WL test.

Panel of Figure 7 — Figure 7: A comparison of dense, sparse, and symbolic matrices. Based on Feydy et al. ( Feydy et al., 2020 ) .

Hyperparameter	Default	Range
Dropout	0.1	[0, 0.1, 0.2]
Depth	8	[6, 8, 10]
Learning rate	0.001	[0.001, 0.0005, 0.0001]
Head depth	2	[1, 2, 3]
Encoding	none	[none, LapPE, RWSE]
$k$ (only GPS + k-MIP)	15	[5, 15, 45]

Appendix

Table 5: Hyperparameter search ranges and default values for the LRGB experiment.

Hyperparameter	Default	Range
Dropout	0.1	[0, 0.1, 0.2]
Depth	8	[6, 8, 12, 16]
Learning rate	0.001	[0.005, 0.001, 0.0005]
Head depth	2	[1, 2, 3]
Encoding	none
$k$ (only GPS + k-MIP)	15	[5, 15, 45]

Appendix

Table 6: Hyperparameter search ranges and default values for the City-Networks experiment.

HP/Dataset	ShapeNet-Part		S3DIS
	MPNNs	GTs	MPNNs	GTs
Metric	F1 score	F1 score	mIoU	mIoU
PE/SE type	none	none	none	none
Optimizer	adamW	adamW	adamW	adamW
Loss function	weighted CE	weighted CE	weighted CE	weighted CE
# Epochs	150	200	120	200

Appendix

Table 7: The hyperparameters used on ShapeNet-Part and S3DIS. GT stands for graph transformer.

Panel of Figure 8 — Figure 8: Analysis of k-MIP attention’s relationship to full attention. a. The L2 distance between k-MIP attention and full attention outputs. b. The cumulative attention weight for the top- $k$ keys.

	Cifar10	MalNet-Tiny	PascalVOC-SP	Peptides-Func
$\bm{k}$	Accuracy $\uparrow$	Accuracy $\uparrow$	F1 score $\uparrow$	AP $\uparrow$
$\bm{1}$	74.37 $\pm$ 0.56	93.00 $\pm$ 0.58	35.57 $\pm$ 1.06	60.64 $\pm$ 1.13
$\bm{2}$	74.84 $\pm$ 0.41	92.92 $\pm$ 0.69	37.65 $\pm$ 0.70	62.97 $\pm$ 0.27
$\bm{3}$	74.75 $\pm$ 0.33	92.88 $\pm$ 0.41	37.47 $\pm$ 0.48	63.13 $\pm$ 0.65
$\bm{5}$	75.08 $\pm$ 0.27	93.54 $\pm$ 0.52	37.30 $\pm$ 0.99	64.65 $\pm$ 0.37
$\bm{10}$	75.05 $\pm$ 0.32	93.52 $\pm$ 0.49	36.58 $\pm$ 1.50	65.10 $\pm$ 0.42

Appendix

Table 8: Comparison of the performance of GraphGPS with k-MIP attention for varying

k

, on four small-graph datasets. Shown is the mean

\pm

std over 5 runs. The first and second best results are highlighted.

	Cifar10	MalNet-Tiny	PascalVOC-SP	Peptides-Func
$\bm{k}$	Time (s) $\uparrow$	Time (s) $\uparrow$	Time (s) $\uparrow$	Time (s) $\uparrow$
$\bm{1}$	135.34 $\pm$ 8.72	22.09 $\pm$ 1.19	16.14 $\pm$ 0.95	9.29 $\pm$ 0.29
$\bm{2}$	146.64 $\pm$ 11.51	23.96 $\pm$ 0.89	15.61 $\pm$ 1.00	9.70 $\pm$ 0.18
$\bm{3}$	146.50 $\pm$ 20.20	24.25 $\pm$ 0.91	16.81 $\pm$ 1.11	10.05 $\pm$ 0.19
$\bm{5}$	147.94 $\pm$ 14.23	26.44 $\pm$ 0.90	17.12 $\pm$ 0.86	10.55 $\pm$ 0.25
$\bm{10}$	146.60 $\pm$ 19.97	30.14 $\pm$ 0.94	17.70 $\pm$ 0.89	12.79 $\pm$ 0.26

Appendix

Table 9: Comparison of the average training epoch duration in seconds of GraphGPS with k-MIP attention for varying

k

, on four small-graph datasets. Shown is the mean

\pm

std over 5 runs. The first and second best results are highlighted.

Figure 9: Scatter plots of the test performance against the average training epoch duration for varying values of $k$ , for the datasets Cifar10, MalNet-Tiny, PascalVOC-SP, and Peptides-Func. The Pareto frontier is drawn for Cifar10 and Peptides-Func.

	Inference Setting	Training Setting
$\bm{N}$	Total runtime	Forward	Backward	Total	GPU memory (MB)
$\bm{10^{2.0}}$	1.67 ms	1.68 ms	0.96 ms	2.64 ms	0.19
$\bm{10^{2.5}}$	1.69 ms	1.79 ms	0.98 ms	2.76 ms	0.58
$\bm{10^{3.0}}$	1.89 ms	1.93 ms	1.01 ms	2.94 ms	1.84
$\bm{10^{3.5}}$	2.14 ms	2.23 ms	1.04 ms	3.27 ms	5.80
$\bm{10^{4.0}}$	3.14 ms	3.18 ms	1.41 ms	4.59 ms	18.31

Appendix

Table 10: k-MIP attention with symbolic matrices: Detailed runtimes and peak GPU memory usage across different graph sizes, in the controlled experiment of section 5.1 . All experiments were conducted on a single A100 GPU with 40GB memory. No gradient checkpointing was used. Averages are reported over 5 runs.

	Inference Setting	Training Setting
$\bm{N}$	Total runtime	Forward	Backward	Total	GPU memory (MB)
$\bm{10^{2.0}}$	2.06 ms	2.01 ms	1.12 ms	3.13 ms	8.31
$\bm{10^{2.5}}$	2.00 ms	2.00 ms	1.00 ms	3.00 ms	8.71
$\bm{10^{3.0}}$	2.09 ms	2.09 ms	1.07 ms	3.16 ms	12.81
$\bm{10^{3.5}}$	2.19 ms	2.20 ms	1.12 ms	3.33 ms	49.00
$\bm{10^{4.0}}$	5.00 ms	4.96 ms	1.31 ms	6.27 ms	398.22

Appendix

Table 11: k-MIP attention without symbolic matrices: Detailed runtimes and peak GPU memory usage across different graph sizes, in the controlled experiment of section 5.1 . All experiments were conducted on a single A100 GPU with 40GB memory. No gradient checkpointing was used. Averages are reported over 5 runs.

	Inference Setting	Training Setting
$\bm{N}$	Total runtime	Forward	Backward	Total	GPU memory (MB)
$\bm{10^{2.0}}$	0.64 ms	0.75 ms	1.04 ms	1.79 ms	16.42
$\bm{10^{2.5}}$	0.83 ms	0.75 ms	0.94 ms	1.69 ms	17.84
$\bm{10^{3.0}}$	0.67 ms	0.76 ms	1.04 ms	1.80 ms	31.70
$\bm{10^{3.5}}$	0.70 ms	0.79 ms	1.42 ms	2.21 ms	169.42
$\bm{10^{4.0}}$	2.61 ms	2.69 ms	7.41 ms	10.10 ms	1544.04

Appendix

Table 12: Full attention: Detailed runtimes and peak GPU memory usage across different graph sizes, in the controlled experiment of section 5.1 . All experiments were conducted on a single A100 GPU with 40GB memory. No gradient checkpointing was used. Averages are reported over 5 runs.

	Paris	Shanghai	LA	London
# Nodes	114k	184k	241k	569k
# Edges	183k	263k	343k	759k
GCN	52.93 $\pm$ 0.06	57.75 $\pm$ 0.24	56.65 $\pm$ 0.04	55.25 $\pm$ 0.06
GINE	53.36 $\pm$ 0.23	63.35 $\pm$ 0.20	58.21 $\pm$ 0.56	57.60 $\pm$ 0.20
GAT	55.83 $\pm$ 0.42	72.53 $\pm$ 0.23	65.53 $\pm$ 0.65	57.19 $\pm$ 1.09
GatedGCN	53.27 $\pm$ 0.10	68.80 $\pm$ 0.21	63.42 $\pm$ 0.12	61.47 $\pm$ 0.14

Appendix

Table 13: Test accuracies of GPS + k-MIP and baselines on City-Networks ( Liang et al., 2025 ) . Shown is the mean

\pm

std over four runs, except for the London dataset where we only ran one run for the graph transformer models. GPS + BigBird was not evaluated on LA due to long training times.

	Paris	Shanghai	LA	London
# Nodes	114k	184k	241k	569k
# Edges	183k	263k	343k	759k
	Time (h)	Time (h)	Time (h)	Time (h)
GCN	0.051	0.066	0.089	0.193
GINE	0.048	0.065	0.087	0.183
GAT	0.071	0.093	0.125	0.278

Appendix

Table 14: Training times of GPS + k-MIP and baselines on City-Networks ( Liang et al., 2025 ) . Shown is the mean over four runs, except for the London dataset where we only ran one run for the graph transformer models. All times were measured on a single 40GB A100 GPU for Paris, and a single 80GB A100 GPU for the other datasets.

$N$	k-MIP (fp32) runtime	k-MIP (fp32) peak memory (MB)	FlashAttention (fp16) runtime	FlashAttention (fp16) peak memory
$10^{2.0}$	1.67 ms	0.14	0.20 ms	0.02
$10^{2.5}$	1.69 ms	0.45	0.17 ms	0.13
$10^{3.0}$	1.89 ms	1.41	0.17 ms	0.69
$10^{3.5}$	2.14 ms	4.47	0.18 ms	2.17
$10^{4.0}$	3.14 ms	14.18	0.30 ms	6.87
$10^{4.5}$	6.38 ms	44.63	1.80 ms	6.87

Appendix

Table 15: Inference runtime and memory comparison between k-MIP (fp32) and FlashAttention (fp16).

$N$	k-MIP (fp32) total runtime	k-MIP (fp32) peak memory (MB)	FlashAttention (fp16) runtime	FlashAttention (fp16) peak memory
$10^{2.0}$	2.64 ms	0.19	0.98 ms	1.73
$10^{2.5}$	2.76 ms	0.58	1.01 ms	5.17
$10^{3.0}$	2.94 ms	1.84	1.14 ms	13.83
$10^{3.5}$	3.27 ms	5.80	1.29 ms	43.23
$10^{4.0}$	4.59 ms	18.31	1.86 ms	136.60
$10^{4.5}$	8.46 ms	57.90	7.89 ms	428.88

Appendix

Table 16: Training runtime and memory comparison between k-MIP (fp32) and FlashAttention (fp16).

Explore similar work

Gaussian Process Limit Reveals Structural Benefits of Graph Transformers
Mar 18, 2026Nil Ayday, Lingchu Yang, Debarghya GhoshdastidarGraph Neural Networks Graph Transformer
School of Computation, Information and Technology, Technical University of Munich · Graduate School of Information Science and Technology, The University of Tokyo
Capacity-Controlled Global Attention for Graph Transformers
Apr 19, 2026Yang Liu, Dongxin Guo, Tom Zheng +3Gated Attention Transformer Attention
Brain Investing Limited Hong Kong, China · The University of Hong Kong Hong Kong, China · Stellaris AI Limited Hong Kong, China
Attention Graphons: A Graph Limit Perspective on Graph Transformers
Sep 28, 2026Caio F. Deberaldini Netto, Moshe Eliasof, Luana RuizTransformer Attention Graph Transformer
Johns Hopkins University Baltimore, MD · University of Cambridge Cambridge, UK

Mar 18, 2026J/K move · Enter open · S save

Gaussian Process Limit Reveals Structural Benefits of Graph Transformers

Nil Ayday, Lingchu Yang, Debarghya Ghoshdastidar

School of Computation, Information and Technology, Technical University of Munich · Graduate School of Information Science and Technology, The University of Tokyo

Graph Neural Networks Graph Transformer Node Classification

Open paper AskPDF arXiv

Graph transformers are the state-of-the-art for learning from graph-structured data and are empirically known to avoid several pitfalls of message-passing architectures. However, there is limited theoretical analysis on why these models perform well in practice. In this work, we prove that attention-based architectures have structural benefits over graph convolutional networks in the context of node-level prediction tasks. Specifically, we study the neural network gaussian process limits of graph transformers (GAT, Graphormer, Specformer) with infinite width and infinite heads, and derive the node-level and edge-level kernels across the layers. Our results characterise how the node features and the graph structure propagate through the graph attention layers. As a specific example, we prove that graph transformers structurally preserve community information and maintain discriminative node representations even in deep layers, thereby preventing oversmoothing. We provide empirical evidence on synthetic and real-world graphs that validate our theoretical insights, such as integrating informative priors and positional encoding can improve performance of deep graph transformers.

Abstract

Explore similar work

Gaussian Process Limit Reveals Structural Benefits of Graph Transformers

Capacity-Controlled Global Attention for Graph Transformers

Attention Graphons: A Graph Limit Perspective on Graph Transformers