Parameterized Explainer for Graph Neural Network

In plain words

A graph neural network makes a prediction for a node or a whole graph, but it does not say which connections in the graph drove that prediction. This paper asks for an explanation in the form of a small set of edges that keeps the prediction stable.

The older method, GNNExplainer, solves this separately for every instance, so each explanation is a new optimisation run. PGExplainer instead trains one small network that looks at a graph and says how important each edge is. Once that network is trained, it can explain new graphs or nodes in a single forward pass, and it can be trained on a group of instances together so that it learns patterns shared across them.

The network is trained by sampling random subsets of edges, weighted by its scores, and rewarding it when the frozen model still gives the same answer on those subsets. The authors report that this is both more accurate on synthetic benchmarks with known answers and much faster at explaining new instances.

TL;DR

• Explains a trained GNN by picking edges (not node features) whose presence keeps the GNN's prediction.
• A shared MLP outputs one logit per edge, so one trained explainer covers many instances and unseen graphs §4.3.
• Trained with binary concrete relaxation of edge selection and a cross-entropy term against the GNN's original prediction §4.2 eq. 4 eq. 9.
• Author-reported: best explanation AUC in all six settings, and 12x-108x faster inference than GNNExplainer tab. 2.
• Use it when you need explanations for many instances or for new graphs. The per-instance speed-up does not include the explainer's own training time app. C.1.

Problem

Explaining a GNN means finding a subgraph whose information preserves the prediction §1. Exact search over subgraphs is exponential in the number of edges, so the problem is relaxed to a continuous optimisation §4.1.

The existing per-instance method, GNNExplainer, has three limitations, as the authors state them. Each explanation is built independently and does not give a global picture of the model. Each explanation needs its own optimisation, which is slow for large datasets. And the explanations do not transfer to unseen nodes in the inductive setting §1 §4.3.

The paper assumes the GNN is trained and fixed, and that the task is node or graph classification. Link prediction is mentioned as similar but not evaluated §4.3. Feature-level explanation is out of scope and only discussed briefly app. D.

Method

The GNN is split into two parts. A node-embedding function produces vectors Z from the graph and features X, and a classifier maps Z to the prediction §4.3 eq. 7. The explainer does not change the GNN's weights.

Objective §4.1. The explanation is the subgraph G_s that maximises the mutual information with the GNN prediction Y_o eq. 1. Since the GNN's entropy term is fixed, this reduces to minimising the conditional entropy of Y_o given G_s §4.1 eq. 3. Each edge is treated as an independent Bernoulli variable with parameter theta_ij, so the graph distribution factorises over edges eq. 2.

Relaxation §4.2. Each binary edge indicator is replaced by a sample from a binary concrete distribution with temperature tau. The sampled weight lies in (0, 1), is differentiable in omega_ij, and approaches the Bernoulli sample as tau goes to 0 when omega_ij = log(theta_ij / (1 - theta_ij)) §4.2 eq. 4. The authors use a Monte Carlo estimate of the expected cross-entropy between the original prediction and the prediction on the sampled subgraph eq. 5 eq. 6.

Amortised explanation network §4.3. Instead of treating omega_ij as a free parameter, omega_ij is produced by a shared MLP Psi from node embeddings eq. 8:

• node classification (explaining node v): the input is the concatenation of z_i, z_j and z_v for the edge (i, j) eq. 10;
• graph classification: the input is the concatenation of z_i and z_j eq. 11.

The training objective sums the cross-entropy over all instances in the set being explained eq. 9. Parameters Psi are the only trained parameters in this stage.

Training procedure alg. 1 alg. 2 app. A. Node tasks use the L-hop computation graph of each node as its instance; graph tasks use the whole graph app. A. The GNN is run once per instance to get Z and Y_o. For each epoch, the explainer produces omega for each instance, samples K edge masks, runs the GNN on each masked graph, and updates Psi by backpropagation on the loss of eq. 9 alg. 1 alg. 2.

Regularisers §4.4. These are optional and were not used in the reported experiments §5.2:

• size: an L1 penalty on the edge weights, with an element-wise entropy term to push weights towards 0 or 1 §4.4 app. B;
• budget: a ReLU penalty on the total sampled weight above a budget B eq. 12;
• connectivity: a cross-entropy between weights of edges that share a node, which encourages connected explanations eq. 13 app. C.2.

# Trained explainer for a frozen GNN  [alg. 1] [alg. 2] [§4.3]
# Input:  instances (graphs or node computation graphs), features X,
#         frozen GNN, explanation network Psi
# Output: trained Psi

def train_explainer(instances, gnn, psi, epochs, temperatures):
    # Run the frozen GNN once per instance to get embeddings and labels
    # [alg. 1 lines 2-5] [alg. 2 lines 2-4]
    for each instance:
        embeddings = gnn.embed(instance.graph, instance.features)
        original_pred = gnn.classify(embeddings)

    for epoch in range(epochs):
        # Temperature for this epoch  [app. B]
        tau = temperatures(epoch)
        for each instance:
            # One logit per edge from the shared MLP  [eq. 8] [eq. 10-11]
            logits = psi(instance.graph, embeddings)
            for k in range(K):  # K is not stated  [§4.2]
                # Relaxed edge weights in (0, 1)  [eq. 4]
                weights = binary_concrete(logits, tau)
                # Prediction on the graph masked by the weights  [§4.2]
                masked_pred = gnn.classify_masked(instance, weights)
                # Cross-entropy with the original prediction  [eq. 9]
                loss += cross_entropy(original_pred, masked_pred)
            # Optional size and entropy terms  [§4.4] [app. B]
            loss += size_and_entropy(weights)
        # Adam update on Psi only  [app. B]
        psi.update(loss)

# One relaxed edge weight  [eq. 4]
def binary_concrete(logit, tau):
    # Uniform noise in (0, 1), one per edge  [eq. 4]
    noise = uniform(0, 1)
    # Logistic noise added to the logit, scaled by temperature  [eq. 4]
    return sigmoid((log(noise) - log(1 - noise) + logit) / tau)

Key equations, copied from the paper:

\hat{e}_{ij} = \sigma\big((\log \epsilon - \log(1-\epsilon) + \omega_{ij})/\tau\big),\quad \epsilon\sim U(0,1) eq. 4

\min_{\Psi}\; -\sum_{i\in\mathcal{I}}\sum_{k=1}^{K}\sum_{c=1}^{C} P_\Phi(Y=c\mid G^{(i)}_o)\,\log P_\Phi(Y=c\mid \hat{G}^{(i,k)}_s) eq. 9

Hyperparameters

Values are as stated in the paper. Where the paper does not state a value, it is marked "not stated". All values come from the paper's text and appendix, not from code.

Results

Author-reported. Setup: five synthetic node and graph datasets (BA-Shapes, BA-Community, Tree-Cycles, Tree-Grid, BA-2motifs) and MUTAG §5.1 tab. 1. Metric: AUC of edge weights against the ground-truth motif edges, with only mutagen graphs used for MUTAG §5.3. Each PGExplainer result is the mean over 10 runs ± std §5.3. Baselines: GRAD, ATT, Gradient and GNNExplainer §5.2. Hardware and training details are in app. B. The GNN accuracy table is in tab. 3.

The printed "Improve" values match PGExplainer over GNNExplainer for the first four datasets and BA-2motifs, and over GRAD for MUTAG (0.783); see Pitfalls for BA-2motifs. Inference time in milliseconds, for explaining one new instance once the explainer is trained §5.3 tab. 2:

Inductive results (explaining unseen instances after training on alpha of them) are shown only as curves in fig. 3, not in a table. The text reports that test AUC grows with alpha and that variance falls as alpha grows app. C.1. Connectivity regularisation on a noisy BA-Shapes variant is shown in fig. 5.

Paper vs code

No comparison is recorded on this card. The repository linked from §5 (https://github.com/flyingdoog/PGExplainer, commit recorded in the frontmatter) has no LICENSE file, so its code license is not stated. The compiler looked at the source while checking the card, but under the licensing rules no content from an unlicensed repository is used here, so no discrepancies are recorded. Anyone who needs a paper-to-code comparison should review the repository directly, subject to its owners' terms.

Pitfalls

• Prediction target is ambiguous. The loss in eq. 9 uses P(Y = c | G_o), but the text does not say whether this is the GNN's softmax output or a one-hot label (predicted or ground truth) §4.3. The two give different gradients whenever the GNN is not confident.
• K is not stated. The number of sampled subgraphs per instance per step is not given anywhere in the paper §4.2.
• Edge independence and symmetry. The explanation treats each edge as an independent Bernoulli variable §4.1. The paper does not say how an undirected edge is represented (one weight, or separate weights for (i, j) and (j, i)).
• Which quantity the regularisers act on. The size term is written as an L1 norm on Omega §4.4, but the coefficient is given only in app. B, and the budget term is written on the sampled weights eq. 12. The entropy term is named but its formula is not given §4.4.
• Connectivity constraint details. eq. 13 is written for an arbitrary adjacent pair. Whether the loss sums over all adjacent pairs or a sample is not stated. The constraint is also not used in the reported experiments §4.4 app. C.2.
• Explanation-network dimensions do not chain. app. B gives FC(#input, 64, ReLU) followed by FC(20, 1, linear). The second layer should take 64 inputs if the first output feeds it. The 60/40 input sizes match 3x and 2x the embedding size 20 eq. 10-11 app. B. Treat the hidden-layer width as unclear.
• Temperature is high and annealed. tau_0 = 5 and tau_T = 2 app. B. The paper says small temperatures hinder optimisation and that relatively high temperatures worked well in practice. No other schedule is reported.
• Hidden training cost. The speed-up in the table compares only inference: GNNExplainer's cost includes retraining per instance, but PGExplainer's own training time is excluded §5.3 app. C.1. The paper says it dismisses that time because of the inductive use case.
• The 24.7% headline is relative to GNNExplainer. On BA-2motifs, Gradient scores 0.773 and GNNExplainer 0.742 tab. 2. Against the best printed baseline, PGExplainer's 0.926 is about 19.8% higher (computed from the table), not 24.7%.
• Dataset numbers conflict. The text describes 800 BA-2motifs graphs §5.1, while tab. 1 lists 1,000. The Tree-Cycles node count in tab. 1 (871) does not follow from the text (an 8-level tree plus 80 six-node cycles); the tree convention is not stated.
• Ground truth is constructed. Synthetic labels follow the motifs by construction, and the MUTAG mutagen motifs (NH2, NO2) are taken from prior work as known explanations §5.1. The AUC therefore measures recovery of planted structure, not explanation quality in general.
• Validation-tuned inductive numbers. The inductive experiments tune hyperparameters on validation data, so their numbers are not directly comparable to the main table app. C.1.
• Feature explanations are out of scope. Node-feature masks are only discussed, not evaluated app. D.

Lineage

• Builds on: GNNExplainer (Ying et al., 2019), a per-instance method. PGExplainer keeps its mutual-information objective, its size and entropy regularisers and its random-graph view, but replaces per-instance parameters with a shared network §1 §4.1 §4.4. It also uses the binary concrete relaxation of Maddison et al. and Jang et al. §4.2.
• Compared against: GRAD, ATT and Gradient §5.2.
• Superseded by: not stated in this card's sources.