DistShap: Scalable GNN Explanations with Distributed Shapley Values

TL;DR

DistShap introduces a scalable, GPU-based distributed algorithm for explaining GNN predictions using Shapley values, enabling efficient analysis of large-scale graphs with millions of features.

Contribution

It presents the first scalable distributed method for GNN explanations using Shapley values, leveraging multiple GPUs for large graph analysis.

Findings

Outperforms existing explanation methods in accuracy

Scales to GNNs with millions of features using 128 GPUs

Operates efficiently on supercomputers like NERSC Perlmutter

Abstract

With the growing adoption of graph neural networks (GNNs), explaining their predictions has become increasingly important. However, attributing predictions to specific edges or features remains computationally expensive. For example, classifying a node with 100 neighbors using a 3-layer GNN may involve identifying important edges from millions of candidates contributing to the prediction. To address this challenge, we propose DistShap, a parallel algorithm that distributes Shapley value-based explanations across multiple GPUs. DistShap operates by sampling subgraphs in a distributed setting, executing GNN inference in parallel across GPUs, and solving a distributed least squares problem to compute edge importance scores. DistShap outperforms most existing GNN explanation methods in accuracy and is the first to scale to GNN models with millions of features by using up to 128 GPUs on the NERSC Perlmutter supercomputer.