$P^2$GNN: Two Prototype Sets to boost GNN Performance

Arihant Jain; Gundeep Arora; Anoop Saladi; Chaosheng Dong

cs.LGMar 10, 2026

$P^2$ GNN: Two Prototype Sets to boost GNN Performance

Authors: Arihant Jain, Gundeep Arora, Anoop Saladi, Chaosheng Dong

Organizations: Amazon, India

Abstract

Message Passing Graph Neural Networks (MP-GNNs) have garnered attention for addressing various industry challenges, such as user recommendation and fraud detection. However, they face two major hurdles: (1) heavy reliance on local context, often lacking information about the global context or graph-level features, and (2) assumption of strong homophily among connected nodes, struggling with noisy local neighborhoods. To tackle these, we introduce $P^2$ GNN, a plug-and-play technique leveraging prototypes to optimize message passing, enhancing the performance of the base GNN model. Our approach views the prototypes in two ways: (1) as universally accessible neighbors for all nodes, enriching global context, and (2) aligning messages to clustered prototypes, offering a denoising effect. We demonstrate the extensibility of our proposed method to all message-passing GNNs and conduct extensive experiments across 18 datasets, including proprietary e-commerce datasets and open-source datasets, on node recommendation and node classification tasks. Results show that $P^2$ GNN outperforms production models in e-commerce and achieves the top average rank on open-source datasets, establishing it as a leading approach. Qualitative analysis supports the value of global context and noise mitigation in the local neighborhood in enhancing performance.

Figures & tables

Figure 1: Schematic of $P^{2}$ GNN . Two computational graphs for each target node: one with neighbouring nodes from source graph ( $CG_{1}$ ) and another with learnable $K_{\mathcal{N}}$ class-independent prototypes $P_{\mathcal{N}}$ as neighbouring nodes ( $CG_{2}$ ). GNN embeddings of target node obtained via message passing from both ( $H_{base}$ from $CG_{1}$ , $H_{P_{\mathcal{N}}}$ from $CG_{2}$ as in Eq. ( 1 )) are mixed to form $H_{\mathcal{N}}$ (Eq. ( 2 )). $H_{\mathcal{N}}$ is then aligned with $K_{\mathcal{A}}$ learnable class-independent $P_{\mathcal{A}}$ prototypes to create $H_{\mathcal{A}}$ (Eq. ( 3 )), followed by mixing with $H_{\mathcal{N}}$ for final output $H_{out}$ (Eq. ( 4 )).

Criterion	DPGNN	HP-GMN	ProtoGNN	$P^{2}$ GNN
GNN-based	✓	✗	✓	✓
No Pre-training	✓	✗	✗	✓
Class independent	✗	✓	✗	✓
Task-agnostic	✗	✗	✗	✓

Table 1:

P^{2}

GNN vs. Other Prototype-based Methods

Notation	Description
$\mathcal{G}=(\mathcal{V},\mathcal{E})$	undirected graph without self-loops
$\mathcal{V}=\{v_{i}\}_{i=1}^{V}$ , $V$	set of nodes, number of nodes
$\mathcal{E}\subseteq\mathcal{V}\times\mathcal{V}$	set of edges
$A\subseteq\{0,1\}^{V\times V}$	adjacency matrix
$X\in\mathbb{R}^{V\times d_{0}}$	initial node feature matrix
$K_{\mathcal{N}}$	Number of learnable $P_{\mathcal{N}}$ prototypes

Table 2: List of Notations

Section	Question	Answer (Yes/No)
RQ1	Does the $P^{2}$ GNN model outperform baseline methods on node recommendation tasks?	Yes
	Does the $P^{2}$ GNN model achieve superior performance on node classification tasks?	Yes
	Does the $P^{2}$ GNN model scale effectively to large-scale graphs?	Yes
	Does $P^{2}$ GNN drastically increase training time?	No
	Are the performance improvements of $P^{2}$ GNN statistically significant over the backbone GNN?	Yes
RQ2	Does incorporating neighborhood prototypes ( $P_{\mathcal{N}}$ ) enhance the model’s performance?	Yes

Table 3: Summary of Research Questions and Findings in the Experimental Evaluation of the Proposed

P^{2}

GNN Model

Dataset	Model	Hits@1	Hits@3	Hits@5	Hits@10	MRR@10
E-comm1	GCN	3.77x	3.81x	3.81x	3.81x	3.79x
E-comm1	$P^{2}$ GNN	3.93x	3.95x	3.96x	3.96x	3.94x
E-comm2	GCN	0.285x	0.753x	1.056x	1.531x	0.613x
E-comm2	$P^{2}$ GNN	0.333x	0.782x	1.090x	1.568x	0.657x

Table 4: Results on proprietary e-commerce datasets for the node recommendation task. Results are relative to a production model at an e-commerce store with GCN backbone.

	Cornell	Wisconsin	Texas	Film	Chameleon	Squirrel	Cora	CiteSeer	Avg. Rank
$\mathcal{H}_{node}$	0.30	0.21	0.11	0.22	0.23	0.22	0.81	0.74
MLP	91.30 $\pm$ 0.70	93.87 $\pm$ 3.33	92.26 $\pm$ 0.71	38.58 $\pm$ 0.25	46.72 $\pm$ 0.46	31.28 $\pm$ 0.27	76.44 $\pm$ 0.30	76.25 $\pm$ 0.28	11.0
GAT [ 50 ]	76.00 $\pm$ 1.01	71.01 $\pm$ 4.66	78.87 $\pm$ 0.86	35.98 $\pm$ 0.23	63.9 $\pm$ 0.46	42.72 $\pm$ 0.33	76.70 $\pm$ 0.42	67.20 $\pm$ 0.46	13.6
GPRGNN [ 8 ]	91.36 $\pm$ 0.70	93.75 $\pm$ 2.37	92.92 $\pm$ 0.61	39.30 $\pm$ 0.27	67.48 $\pm$ 0.40	49.93 $\pm$ 0.53	79.51 $\pm$ 0.36	67.63 $\pm$ 0.38	8.4
H2GCN [ 64 ]	86.23 $\pm$ 4.71	87.5 $\pm$ 1.77	85.90 $\pm$ 3.53	38.85 $\pm$ 1.17	52.30 $\pm$ 0.48	30.39 $\pm$ 1.22	87.52 $\pm$ 0.61	79.97 $\pm$ 0.69	10.8
GCN+JK [ 20 ]	66.56 $\pm$ 13.82	62.50 $\pm$ 15.75	80.66 $\pm$ 1.91	32.72 $\pm$ 2.62	64.68 $\pm$ 2.85	53.40 $\pm$ 1.90	86.90 $\pm$ 1.51	73.77 $\pm$ 1.85	13.1

Table 5: Average classification accuracy and standard deviation on 8 real-world benchmark datasets with a random 60%/20%/20% data split. The best results are highlighted in bold , and the runner-up score with an underline. NA means that neither the reported results nor the code were available for evaluation. ACM-GCN is used as the backbone for

P^{2}

GNN across datasets for a fair comparison. In principle, any high-performing GNN can be used as the

P^{2}

GNN backbone.

Panel of Figure 2 — Figure 2: Test accuracy performance of $P^{2}$ GNN compared against its backbone ACM-GCN, segmented by the level of heterophily in test-nodes. $P^{2}$ GNN consistently boosts the performance on strong (SHet) and weak (WHet) heterophilous node segments.

	Penn94	pokec	arXiv-year	snap-patents	genius	twitch-gamers	Avg. Rank
$\mathcal{H}_{node}$	0.47	0.44	0.22	0.07	0.61	0.54
MLP	$73.61\pm 0.40$	$62.37\pm 0.02$	$36.70\pm 0.21$	$31.34\pm 0.05$	$86.68\pm 0.09$	$60.92\pm 0.07$	10.5
GCN	$82.47\pm 0.27$	$75.45\pm 0.17$	$46.02\pm 0.26$	$45.65\pm 0.04$	$87.42\pm 0.37$	$62.18\pm 0.26$	7.8
GAT	$81.53\pm 0.55$	$71.77\pm 6.18$	$46.05\pm 0.51$	$45.37\pm 0.44$	$55.80\pm 0.87$	$59.89\pm 4.12$	9.3
MixHop	$83.47\pm 0.71$	$81.07\pm 0.16$	$51.81\pm 0.17$	$52.16\pm 0.09$	$90.58\pm 0.16$	$65.64\pm 0.27$	5.2
GCNII	$82.92\pm 0.59$	$78.94\pm 0.11$	$47.21\pm 0.28$	$37.88\pm 0.69$	$90.24\pm 0.09$	$63.39\pm 0.61$	6.8

Table 7: Results (%) on 6 large-scale datasets released in [ 23 ] . We compare the AUC score on genius as in [ 23 ] . For other datasets, we show the classification accuracy. The error bar (

\pm

) denotes the standard deviation score of results over 5 trials. We highlight the best score on each dataset in bold and the runner-up score with an underline. Note that approaches such as [ 35 , 64 , 45 ] suffers out-of-memory issue; We choose not to include them on large datasets.

Model	Penn94	pokec	arXiv-year	snap-patents	genius	twitch-gamers
ACM-GNN	1.52	2.51	1.02	3.43	1.19	1.23
$P^{2}$ GNN	2.02	2.92	1.32	3.72	1.46	1.84

Table 8: Comparison of average running time per epoch (in seconds) of

P^{2}

GNN with ACM-GNN on large datasets

	Model Components		Cornell	Wisconsin	Texas	Film	Chameleon	Squirrel	Cora	CiteSeer
Models	$P_{\mathcal{N}}$	$P_{\mathcal{A}}$	Acc $\pm$ Std	Acc $\pm$ Std	Acc $\pm$ Std	Acc $\pm$ Std	Acc $\pm$ Std	Acc $\pm$ Std	Acc $\pm$ Std	Acc $\pm$ Std
GCN w/	✗	✗	82.30 $\pm$ 2.82	72.50 $\pm$ 3.71	82.62 $\pm$ 4.82	35.52 $\pm$ 0.77	64.38 $\pm$ 2.44	45.25 $\pm$ 2.07	87.88 $\pm$ 0.81	81.39 $\pm$ 1.21
	✓	✗	84.26 $\pm$ 3.45	78.25 $\pm$ 4.75	83.61 $\pm$ 3.20	36.45 $\pm$ 1.29	64.49 $\pm$ 2.56	45.37 $\pm$ 1.76	89.16 $\pm$ 1.01	80.46 $\pm$ 2.12
	✓	✓	84.10 $\pm$ 2.84	78.38 $\pm$ 4.61	83.44 $\pm$ 3.47	36.81 $\pm$ 0.94	64.22 $\pm$ 2.37	45.37 $\pm$ 1.37	89.97 $\pm$ 0.89	80.72 $\pm$ 1.57
SGC w/	✗	✗	70.98 $\pm$ 8.39	70.38 $\pm$ 2.85	83.28 $\pm$ 5.43	25.26 $\pm$ 1.18	51.53 $\pm$ 1.52	31.89 $\pm$ 1.85	76.60 $\pm$ 2.62	79.02 $\pm$ 1.74
	✓	✗	78.20 $\pm$ 4.02	72.75 $\pm$ 4.99	83.93 $\pm$ 2.72	32.63 $\pm$ 2.24	53.81 $\pm$ 1.70	37.69 $\pm$ 1.58	80.11 $\pm$ 2.87	78.99 $\pm$ 1.60

Table 9: Ablation study of

P^{2}

GNN on 8 real-world datasets using random 60%/20%/20% data split. The "Models" refers to backbone chosen for

P^{2}

GNN addition with ✓meaning the component is added to backbone. The best test results are bolded.

Panel of Figure 5 — Figure 5: tSNE plots of the penultimate layer embeddings of $P^{2}$ GNN and the backbone ACM-GCN model. We progressively add the two different prototypes $P_{\mathcal{N}}$ and $P_{\mathcal{A}}$ and observed embeddings become well separated among the different classes (colors).

Panel of Figure 6 — Figure 6: Sensitivity of incremental performance of $P^{2}$ GNN over ACM-GCN backbone towards different hyper-parameters. We experiment with different number of prototypes for $P_{\mathcal{N}}$ and $P_{\mathcal{A}}$ and relative weights of each of the 3 loss functions. While varying each hyper-parameter, the values of others are kept constant.

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Type	Dataset	# Nodes	# Edges	# Feat.	# C	$\mathcal{H}_{node}$
Small	Texas	183	309	1,703	5	.11
	Wisconsin	251	499	1,703	5	.21
	Cornell	183	295	1,703	5	.30
	Film	7,600	29,926	931	5	.22
	Chameleon	2,277	36,101	2,325	5	.23
	Squirrel	5,201	216,933	2,089	5	.22

Appendix

Table 10: Statistics of different datasets used for experimentation. # C is the number of distinct node classes.

\mathcal{H}_{node}

measures the average proportion of label consistency between a node and its neighbors.

	Texas	Wisconsin	Cornell	Film	Squirrel	Chameleon	Cora	Citeseer	Avg. Rank
$\mathcal{H}_{node}$	0.11	0.21	0.30	0.22	0.22	0.23	0.81	0.74
MLP	$80.81\pm 4.75$	$85.29\pm 3.31$	$81.89\pm 6.40$	$36.53\pm 0.70$	$28.77\pm 1.56$	$46.21\pm 2.99$	$75.69\pm 2.00$	$74.02\pm 1.90$	12.4
GCN [ 20 ]	$55.14\pm 5.16$	$51.76\pm 3.06$	$60.54\pm 5.30$	$27.32\pm 1.10$	$53.43\pm 2.01$	$64.82\pm 2.24$	$86.98\pm 1.27$	$76.50\pm 1.36$	13.0
GAT [ 50 ]	$52.16\pm 6.63$	$49.41\pm 4.09$	$61.89\pm 5.05$	$27.44\pm 0.89$	$40.72\pm 1.55$	$60.26\pm 2.50$	$87.30\pm 1.10$	$76.55\pm 1.23$	13.4
H 2 GCN [ 64 ]	${84.86\pm 7.23}$	${87.65\pm 4.98}$	$82.70\pm 5.28$	$35.70\pm 1.00$	$36.48\pm 1.86$	$60.11\pm 2.15$	$87.87\pm 1.20$	$77.11\pm 1.57$	9.0
WRGAT [ 45 ]	$83.62\pm 5.50$	${86.98\pm 3.78}$	$81.62\pm 3.90$	$36.53\pm 0.77$	$48.85\pm 0.78$	$65.24\pm 0.87$	$88.20\pm 2.26$	$76.81\pm 1.89$	8.3

Appendix

Table 11: The classification accuracy

\pm

standard deviation (in %) over the methods on 8 small-scale datasets is reported on the fixed 48%/32%/20% data split. We highlight the best score on each dataset in bold and the runner-up score with underline .

Explore similar work

Enhancing Graph Representations with Neighborhood-Contextualized Message-Passing
Nov 14, 2025Brian Godwin Lim, Galvin Brice Lim, Renzo Roel Tan +2Graph Neural Networks Graph Representation Learning
Nara Institute of Science and Technology, Kyoto University, and Ateneo de Manila University · UNI-President Information Philippines Corporation · Nara Institute of Science and Technology and Ateneo de Manila University +2
Revisiting Pre-Propagation GNNs: Robust Diffusion Operators and Hidden-State Re-Propagation
May 24, 2026Zichao Yue, Zhiru ZhangGraph Neural Networks Graph Diffusion
School of Electrical and Computer Engineering, Cornell University, Ithaca, New York, USA.
How Much Does Message Passing Matter? A Drop-In Study of GNN Layers for Neural Network Graph Regression
Jul 29, 2026Keith G. Mills, Aedan J. DeFrates, Joong Ho KimGraph Attention Networks Graph Neural Networks
LSU ATHENA Lab

Nov 14, 2025J/K move · Enter open · S save

Enhancing Graph Representations with Neighborhood-Contextualized Message-Passing

Brian Godwin Lim, Galvin Brice Lim, Renzo Roel Tan, Irwin King+1

Nara Institute of Science and Technology, Kyoto University, and Ateneo de Manila University · UNI-President Information Philippines Corporation · Nara Institute of Science and Technology and Ateneo de Manila University+2

Graph Neural Networks Graph Representation Learning GNN Expressivity Message Passing Neural Networks

Open paper AskPDF arXiv

Graph neural networks (GNNs) have become an indispensable tool for analyzing relational data. Classical GNNs are broadly classified into three variants: convolutional, attentional, and message-passing. While the standard message-passing variant is expressive, its typical pair-wise messages only consider the features of the center node and each neighboring node individually. This design fails to incorporate contextual information contained within the broader local neighborhood, potentially hindering its ability to learn meaningful relationships within the entire set of neighboring nodes. To address this, the paper first refines the concept of neighborhood-contextualization within GNNs, leveraging ideas from set-based aggregation methods and a key property of the attentional variant. This then serves as the basis for generalizing the message-passing variant to the proposed neighborhood-contextualized message-passing (NCMP) framework. To demonstrate its utility, a simple, mathematically grounded method to parametrize and operationalize NCMP is presented, leading to the development of the proposed Soft-Isomorphic Neighborhood-Contextualized Graph Convolution Network (SINC-GCN). Across a diverse set of synthetic and benchmark datasets, SINC-GCN strikes a highly favorable balance between expressivity and efficiency. Notably, while more complex models incur significant computational overhead, SINC-GCN delivers substantial performance gains with considerable effect sizes over baseline GNN models while maintaining a highly efficient asymptotic runtime complexity, further underscoring the distinctive utility of neighborhood-contextualization. Overall, by integrating multiset neighborhood context, the proposed NCMP framework serves as a practical and scalable path toward enhancing the graph representational power of classical GNNs.

Abstract

Explore similar work

Enhancing Graph Representations with Neighborhood-Contextualized Message-Passing

Revisiting Pre-Propagation GNNs: Robust Diffusion Operators and Hidden-State Re-Propagation

How Much Does Message Passing Matter? A Drop-In Study of GNN Layers for Neural Network Graph Regression