WILTing Trees: Interpreting the Distance Between MPNN Embeddings

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.

Message Passing Neural Networks (MPNNs) are remarkably good at making predictions for graph-structured data, such as molecules or social networks. But how exactly do these complex models achieve such high performance?We tackled this question by developing a new, human-understandable way to measure the "distance" or dissimilarity between graphs, called the Weisfeiler Leman Labeling Tree distance. We used this to approximate and analyze the internal "MPNN distance"—how the MPNN model learns to distinguish between different input graphs.Our experiments revealed that MPNNs perform well because they learn to capture the task-related functional distance between graphs (e.g., differences in a molecule's properties like oil-friendliness) rather than just their structural distance (e.g., differences in their physical shape). This result indicates that it can be more useful to interpret MPNNs from a functional distance perspective, rather than a structural distance perspective as in most previous studies. We also discovered that MPNNs focus on tiny, functionally important sub-regions within the graphs to capture the functional distance. Our method offers a new way to understand the "black-box" MPNNs, contributing to safer and more robust machine learning.