WILTing Trees: Interpreting the Distance Between MPNN Embeddings

TL;DR

This paper introduces a method to interpret the learned distance function of message passing neural networks (MPNNs) by distilling it into an interpretable graph distance using optimal transport on Weisfeiler Leman Labeling Trees, revealing key subgraphs influencing embeddings.

Contribution

We propose a novel approach that captures task-specific distances in MPNN embeddings through an interpretable graph distance based on optimal transport on WILT, generalizing existing graph kernels.

Findings

MPNNs focus on a small set of functionally important subgraphs.

The proposed method generalizes two well-known graph kernels.

The approach is computationally efficient, running in linear time.

Abstract

We investigate the distance function learned by message passing neural networks (MPNNs) in specific tasks, aiming to capture the functional distance between prediction targets that MPNNs implicitly learn. This contrasts with previous work, which links MPNN distances on arbitrary tasks to structural distances on graphs that ignore task-specific information. To address this gap, we distill the distance between MPNN embeddings into an interpretable graph distance. Our method uses optimal transport on the Weisfeiler Leman Labeling Tree (WILT), where the edge weights reveal subgraphs that strongly influence the distance between embeddings. This approach generalizes two well-known graph kernels and can be computed in linear time. Through extensive experiments, we demonstrate that MPNNs define the relative position of embeddings by focusing on a small set of subgraphs that are known to be functionally important in the domain.