Research map / Graph learning / GNN architectures
WILTing Trees: Interpreting the Distance Between MPNN Embeddings
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that trained MPNN embedding distances align with task-relevant functional distances (unlike structural distances), and distills MPNN distances into a trainable weighted Weisfeiler-Leman Labeling Tree (WILT) optimal-transport distance whose edge weights reveal functionally important subgraphs.
Problem
Existing analyses of MPNN embedding distances link them to task-agnostic structural graph distances, leaving unclear what distance MPNNs learn for a specific task and how they learn it.
Approach
Define a functional pseudometric from target labels, measure alignment of MPNN distances with it, and distill MPNN embedding distances into a weighted WL Labeling Tree optimal-transport distance (d_WILT) whose learned edge weights identify influential WL colors/subgraphs; the approach generalizes WWL and WLOA kernel distances and is computable in linear time.
Key results
- WILT distances approximate MPNN distances far better than WWL and WLOA (e.g., Mutagenicity sum-pooling GCN: RMSE 0.82e-2 for bar-d_WILT vs 5.98e-2 for WLOA).
- About 95% of WILT edge weights can be set to zero with L1 regularization while still aligning better than WLOA.
- Alignment ALI_k correlates positively with accuracy on Mutagenicity and ENZYMES (SRC up to 0.89 on train) and negatively with RMSE on Lipophilicity.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Distill the GCN embedding distance with size-normalized WILT instead of the strongest fixed-weight baseline | GCN with mean pooling · Mutagenicity | RMSE (× 10^-2): 9.25$\pm$0.87 -> 1.74$\pm$0.52 | Helped |
| Distill the GCN embedding distance with dummy-normalized WILT instead of the strongest fixed-weight baseline | GCN with sum pooling · Mutagenicity | RMSE (× 10^-2): 5.98$\pm$1.60 -> 0.82$\pm$0.17 | Helped |
| Distill the GCN embedding distance with size-normalized WILT instead of the strongest fixed-weight baseline | GCN with mean pooling · ENZYMES | RMSE (× 10^-2): 12.18$\pm$0.23 -> 2.71$\pm$0.38 | Helped |
| Distill the GCN embedding distance with dummy-normalized WILT instead of the strongest fixed-weight baseline | GCN with sum pooling · ENZYMES | RMSE (× 10^-2): 6.83$\pm$0.41 -> 1.43$\pm$0.10 | Helped |
| Distill the GCN embedding distance with size-normalized WILT instead of the strongest fixed-weight baseline | GCN with mean pooling · Lipophilicity | RMSE (× 10^-2): 10.92$\pm$0.42 -> 3.11$\pm$0.54 | Helped |
| Distill the GCN embedding distance with size-normalized WILT instead of the strongest fixed-weight baseline | GCN with sum pooling · Lipophilicity | RMSE (× 10^-2): 10.00$\pm$1.34 -> 2.50$\pm$0.67 | Helped |
| Add L1 regularization to the WILT edge-weight objective (ablation) | Dummy-normalized WILT distilled from a 3-layer GCN with sum pooling · Mutagenicity | RMSE | Hurt |
| Train WILT edge weights with L1 regularization at λ=1.0, leaving about 5% of edges non-zero (ablation) | Dummy-normalized WILT distilled from a 3-layer GCN with sum pooling · Mutagenicity | RMSE | Helped |
| Train the MPNN rather than use an untrained MPNN | MPNN · Multiple datasets | ALI_k alignment with functional distance | Helped |
| Increase alignment between the MPNN embedding distance and functional distance | Trained MPNNs · Mutagenicity, ENZYMES, and Lipophilicity | Predictive performance: 0.66, 0.70; 0.89, 0.49; -0.65, -0.57 | Mixed |
Surprising findings
- Structural alignment with MPNN distances does not improve with training and does not correlate with performance.
- On IMDB-BINARY test, ALI_1 correlates negatively with accuracy (SRC -0.43).
Limitations
- Only GCN and GIN with fixed hyperparameters were distilled.
- Analysis limited to final-layer embeddings.
- Restricted to 1-WL colors; higher-order WL extensions not explored.
Experimental setup
Datasets: Mutagenicity ENZYMES Lipophilicity IMDB-BINARY COLLAB
Benchmarks: Mutagenicity ENZYMES Lipophilicity IMDB-BINARY COLLAB
Metrics: Spearman rank correlation RMSE between distances alignment ALI_k
Baselines: WWL distance WLOA distance
Builds on: Weisfeiler Leman test Wasserstein Weisfeiler Leman graph kernel Weisfeiler Leman optimal assignment kernel graph convolutional networks graph isomorphism networks
Closest papers by meaning
- Fine-grained Expressivity of Graph Neural Networks (NeurIPS 2023 · similarity 0.88)
- Distance Encoding: Design Provably More Powerful Neural Networks for Graph Representation Learning (NeurIPS 2020 · similarity 0.87)
- Covered Forest: Fine-grained generalization analysis of graph neural networks (ICML 2025 · similarity 0.87)
- Probabilistic Graph Rewiring via Virtual Nodes (NeurIPS 2024 · similarity 0.86)
- Probabilistically Rewired Message-Passing Neural Networks (ICLR 2024 · similarity 0.85)
- Improving the Effective Receptive Field of Message-Passing Neural Networks (ICML 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.