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

TL;DR

This work introduces P^2 GNN, a plug-and-play technique leveraging prototypes to optimize message passing, enhancing the performance of the base GNN model and establishing it as a leading approach.

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 (6)

• 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. \ref{['eqn:fwd']}) are mixed to form $H_\mathcal{N}$ (Eq. \ref{['eqn:hnl_attn']}). $H_\mathcal{N}$ is then aligned with $K_\mathcal{A}$ learnable class-independent $P_\mathcal{A}$ prototypes to create $H_\mathcal{A}$ (Eq. \ref{['eqn:align_sf']}), followed by mixing with $H_\mathcal{N}$ for final output $H_{out}$ (Eq. \ref{['eqn:align_attn']}).
• 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.
• Figure 3: The plots illustrate the level of attention given by test nodes to two types of prototypes. The top row depicts the attention paid to the prototypes $P_{\mathcal{N}}$, while the bottom row shows the attention given to prototypes $P_{\mathcal{A}}$.
• Figure 4: Learned Prototype $P_{\mathcal{N}}$ representation for first GNN layer and raw node features.
• 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).