{"id":"256b6c6f-0301-4da0-a672-c9c1c55ad317","arxiv_id":"2505.24642","paper_version":1,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":8.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":1,"one_line_summary":"The authors propose WILT, a weighted Weisfeiler-Leman tree distance that can be learned to approximate MPNN embedding distances, and show that MPNNs align with task-relevant functional distances rather than structural ones.","lead":"This paper studies the distance between graph embeddings learned by message passing neural networks and proposes a trainable tree-based distance, WILT, that approximates it and reveals which subgraphs drive the embedding geometry. A smart generalist might read it because it offers a new way to interpret what graph neural networks learn on molecular data, potentially identifying functionally important substructures.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Held-out generalization of the distilled WILT distance is never tested; in-sample RMSE does not establish that large edge weights identify subgraphs that drive dMPNN on new graphs.","rationale":"The theoretical core of the paper appears sound: Proposition 5.3 is a standard tree-Wasserstein identity, and the expressiveness results in Section 5.5 and Appendix A.4 give a principled reason why size-normalized WILT should track mean-pooled MPNNs and dummy-node-normalized WILT should track sum-pooled MPNNs. I give credit for the efficient linear-time computation and the explicit generalization of the WWL and WLOA distances. The central empirical claim, however, is that fitting WILT to dMPNN reveals which WL subgraphs drive the embedding geometry. That claim depends on generalization of the fitted distance and on the stability of the resulting color-ranking. As described above, all reported RMSE values and all interpretability outputs come from the same full-dataset fit, so the current evidence does not rule out memorization or arbitrary assignment of weight to idiosyncratic colors. The proposed cross-validated fit with a held-out RMSE comparison and top-color overlap across folds would settle this directly: if held-out RMSE collapses toward the fixed WWL/WLOA baselines, or if the top-color set changes across folds, the interpretation claims should be weakened to in-sample descriptive statements. I agree with the reader's weakest_assumption; this is essentially the same point. The CONDITIONAL verdict remains appropriate: the concern is empirical and addressable, not a defect in the theoretical construction.","tokens_in":26517,"tokens_out":8135,"duration_ms":109197,"concrete_test":"Run a 5-fold cross-validation on each dataset. For each fold, fit the WILT edge weights using Algorithm 1 with D equal to the training folds only, then compute RMSE(d_MPNN, alpha * d_WILT) on pairs whose two graphs are both in the held-out fold, optimizing alpha on those held-out pairs. Compare this held-out RMSE against the fixed-weight dWWL/dWLOA baselines and against weights fit to a permuted version of d_MPNN. Also compute, for each fold, the set of top-10 weight colors for the Mutagenicity model and report their overlap (e.g., Jaccard) with the full-data top-10 list. If the held-out RMSE is not substantially lower than the fixed baselines, or if the top-color sets are unstable across folds, the paper's central interpretability claim is not supported.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The load-bearing step is the inference in Section 5.4: a large learned edge weight identifies a WL color whose presence or absence significantly affects dMPNN. This inference requires dWILT to be a faithful proxy for dMPNN as a function, not merely on the pairs used to fit the weights. The paper never tests that. The distillation objective L is minimized over all pairs in D (Appendix E: 'We used the entire data set for D in L'), and the RMSE in Table 2 is computed from 1000 pairs of the same D^2. With thousands of free edge weights, the in-sample fit can memorize training pairs, so the low RMSE does not establish that dWILT captures dMPNN on unseen graphs. The sparsity analysis and the top-weight subgraphs (Figures 4-6, Tables 7-8) are produced from that same full-data fit, so they do not independently confirm the abstract's claim that MPNNs focus on a small set of functionally important subgraphs. The conclusion's limitation paragraph also narrows the evidence to two architectures with fixed hyperparameters. If the approximation fails out-of-sample, the 'large weight => large dWILT => large dMPNN' chain in Section 5.4 breaks, and the interpretation result is unsupported.","agreement_with_reader":"agree"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper studies the Euclidean distance between MPNN graph embeddings (d_MPNN) and asks whether this learned distance is task-aligned and interpretable. It first measures alignment between d_MPNN and a task-defined functional pseudometric d_func, reporting that training increases this alignment and that alignment correlates with predictive performance. The main methodological contribution is the Weighted Weisfeiler-Leman Labeling Tree (WILT): a weighted tree of WL colors whose shortest-path/optimal-transport distance d_WILT is shown to generalize the Wasserstein WL distance and the WL optimal assignment distance, to be computable in linear time via an explicit embedding, and to sit in a proved expressiveness hierarchy relative to d_MPNN. Edge weights are fitted by minimizing squared error between d_WILT and d_MPNN on a graph dataset, after which large learned weights are interpreted as identifying WL subgraphs that strongly influence the MPNN embedding distance. Experiments on five datasets with GCN and GIN report good in-sample RMSE for the distilled distance, a skewed weight distribution, and top-weight subgraphs that match known toxicophores on Mutagenicity.","tokens_in":26750,"tokens_out":5210,"duration_ms":67333,"significance":"If its empirical claims hold, the paper offers a concrete, computationally affordable surrogate for MPNN embedding geometry and a new way to produce global, subgraph-level interpretations of trained MPNNs. The theoretical parts are a genuine contribution: Proposition 5.3 gives an exact tree-Wasserstein/Manhattan equivalence, Theorems A.7 and A.8 identify WWL and WLOA as special cases, and Theorems A.9--A.11 give a clean expressiveness ordering that correctly predicts the mean-pooling versus sum-pooling experimental trends. The paper ships code and the derivations are self-contained. The main weakness is that the central interpretability claim rests on an in-sample fit of d_WILT to d_MPNN; without held-out validation, the chain from learned edge weights to statements about which subgraphs 'determine' d_MPNN is not fully established.","major_comments":[{"comment":"The approximation quality of d_WILT is evaluated only in-sample. Appendix E states that \"We used the entire data set for D in L\" when minimizing the distillation loss, and Section 6 reports RMSE computed from 1000 pairs drawn from the same D^2. Because d_WILT is linear in the edge weights (Proposition 5.3) and the number of weights can be large, a low in-sample RMSE does not establish that d_WILT approximates d_MPNN as a function on unseen graphs. This is load-bearing: the Section 5.4 inference \"large w => large d_WILT => large d_MPNN\" and the sparsity and top-subgraph analyses in Figures 4--6 and Tables 7--8 all use the same full-data fit. I ask the authors to split the dataset (or at least the graph pairs) into distillation and evaluation parts, report RMSE on held-out pairs, and recompute the top-weight subgraphs from a model fitted only on the training split.","section":"Section 5.4"},{"comment":"The qualitative claim that the largest edge weights identify subgraphs known to be functionally important is not compared with any baseline. Since the method can only output WL colors that survive a frequency filter, it is unclear whether the 4/10 or 3/10 toxicophore matches are significantly better than what a random or frequency-matched set of WL colors would achieve. A simple enrichment analysis, such as precision@k against the Kazius toxicophore list with a random-color control, would make the central interpretability claim much stronger.","section":"Section 6"},{"comment":"The claim that alignment between d_MPNN and d_func is \"key to the high predictive performance of MPNNs\" is supported by Spearman correlations computed across models that simultaneously vary in architecture, depth, width, and pooling. This design confounds alignment with model capacity and other hyperparameters. The evidence is consistent with correlation but does not establish that alignment is the causal driver; an ablation that varies training signal or regularizes d_MPNN toward/away from d_func while holding architecture fixed would be needed for the stronger causal wording used in the abstract and Section 1.","section":"Section 4"}],"minor_comments":[{"comment":"The word \"psudometrics\" should be \"pseudometrics\".","section":"Section 7"},{"comment":"In the paragraph discussing partial matches, \"limiation\" should be \"limitation\".","section":"Appendix E"},{"comment":"The label \"alphatic halide\" should be \"aliphatic halide\".","section":"Tables 7 and 8"},{"comment":"In the text as provided, the colors in the figure are rendered as blank boxes; if this is not a text-extraction artifact, the example in Figure 3 is very hard to follow and should be fixed in the camera-ready version.","section":"Figure 3"},{"comment":"The projection step \"w <- max(w, 0)\" is applied after every Adam update; since Adam state depends on past gradients, this is a valid projected-gradient approach but the paper should state that the projection is applied per iteration rather than only at the end.","section":"Algorithm 1"}],"recommendation":"major_revision","confidential_remarks":"The main risk is the in-sample-only evaluation of the distilled distance. If the authors add a held-out split for the distillation and rerun the sparsity/top-subgraph analysis, the paper would be much stronger. The theoretical contribution appears sound and the paper is within scope. I recommend major revision rather than rejection because the missing evidence is fixable within the manuscript's scope."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"Colleague,\n\nThe thing to know about this paper: the WILT distance is a real and useful idea, but the paper's main interpretability claim is not backed by the experiments as written. The authors fit edge weights of a WILT to MPNN embeddings on the full dataset, report RMSE on the same pairs, and then read off which subgraphs have large weights. Without a held-out set, the low RMSE only shows the model can fit the training pairs; it does not establish that dWILT captures dMPNN as a function, and the 'large weight => large dWILT => large dMPNN' chain in Section 5.4 has no out-of-sample support.\n\nWhat is genuinely new: the WILT distance with learned edge weights is a clean trainable generalization of WWL and WLOA, and the equivalence to weighted Manhattan distance on the tree (Proposition 5.3) is a nice, useful observation. The expressiveness analysis in Section 5.5 is also a genuine addition; the comparison between size normalization and dummy-node normalization, and the matching with mean vs sum pooling, is a solid theoretical contribution that explains a real empirical pattern. The results in Section 4, on functional alignment of MPNN distances, are based on a reasonable metric (ALI) and a sizable sweep of 48 models, and they support the claim that training improves alignment with task distance and that alignment correlates with performance. That part is credible.\n\nThe soft spots are concentrated in the distillation experiments. The stress test is right: the entire data set is used in L (Appendix E), and the RMSE is computed from 1000 pairs of the same D^2. With thousands of free edge weights, in-sample RMSE is not informative about the generalization of the approximation. The sparsity analysis and the top-weight subgraph visualizations come from that same full-data fit, so they cannot independently confirm the abstract's claim that MPNNs focus on a small set of functionally important subgraphs. The authors do include an L1 regularization experiment, but that is also in-sample. The paper's own limitation paragraph notes only two architectures with fixed hyperparameters, which is honest but does not address the held-out gap.\n\nThe theoretical results appear correct; the generalization of the two kernel distances is cleanly established, and the expressiveness theorems are consistent with known WL bounds. The citation pattern is reasonable, and the connection to previous structural-distance work is well drawn.\n\nWho this is for: someone working on GNN interpretability or graph kernels will get value from the WILT construction and the expressiveness results. The empirical claims should be read with caution. The paper deserves a serious referee; the main fix is to evaluate the distilled distance on a held-out split, or at least report cross-validated RMSE, and to soften the causal language about edge weights until out-of-sample agreement is shown.\n\nMy recommendation: send it to peer review, probably conditional acceptance with the held-out experiments required. The core idea is worth keeping, and the Section 4 result is a useful contribution on its own.","headline":"Solid theoretical core and a genuinely new distance, but the interpretability claims rest on an in-sample fit that is never tested on held-out graphs.","tokens_in":27307,"tokens_out":2734,"would_cite":true,"duration_ms":30020,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":["68T07","68R10"],"pacs":[],"model":"deepseek-v4-flash","headline":"Trained MPNNs position graphs by task-relevant functional distance, not structural similarity, and a distilled Weisfeiler-Leman tree reveals the few subgraphs that set this geometry.","keywords":["graph neural networks","message passing","Weisfeiler-Leman labeling tree","graph distance","optimal transport","interpretability","graph kernels","functional distance"],"falsifier":"Fit WILT edge weights on a training split of a dataset, then compute the RMSE between $d_{\\mathrm{WILT}}$ and $d_{\\mathrm{MPNN}}$ on held-out graphs and compare the highest-weight WL colors to those fitted on the full dataset: if the held-out fit degrades sharply relative to the in-sample fit, or the top-weight colors change substantially between splits, the claim that the learned weights identify the causal drivers of the MPNN distance fails its central test.","tokens_in":2388,"feed_emoji":"🌳","tokens_out":6393,"duration_ms":188814,"temperature":0.7,"pith_summary":"The paper asks what geometric property of a trained message-passing neural network's embedding space explains its predictive power, and whether that geometry can be translated into human-readable form. It argues that after training, the Euclidean distance between graph embeddings tracks the task's functional distance, meaning how far apart the prediction targets are, and that this alignment, measured by the $\\mathrm{ALI}_k$ criterion, is highly correlated with predictive performance across datasets (higher alignment means higher accuracy or lower RMSE), while alignment with task-agnostic structural distances is neither improved by training nor correlated with performance. To expose the mechanism, the authors distill the MPNN distance into a trainable optimal-transport distance on a weighted Weisfeiler-Leman labeling tree (WILT), which is computable in linear time. The learned edge weights are sharply concentrated on a small set of WL colors, and on the Mutagenicity dataset the corresponding subgraphs match known toxicophores. If correct, this yields a global, task-aware reading of what MPNNs learn, rather than instance-level explanations of single predictions.","feed_headline":"GNN embedding distances boil down to a few task-relevant subgraphs","feed_subtitle":"A weighted WL tree distills the MPNN geometry and flags the subgraphs that matter, matching known toxicophores.","key_machinery":"The central object is the Weisfeiler-Leman Labeling Tree (WILT): a rooted tree whose non-root vertices are the WL colors that appear during the WL test on the dataset, with the root connected to the initial colors and edges linking a color to the color of the same node in the previous iteration; injectivity of WL hashing makes this a tree. With non-negative edge weights $w$, the ground metric is shortest-path length, and the WILTing distance is the optimal transport distance between the color multisets of two graphs. Proposition 5.3 reduces it to the weighted Manhattan distance $\\sum_{c \\in V(T_D)\\setminus\\{r\\}} w(e_{\\{c,p(c)\\}}) |\\nu^G_c - \\nu^H_c|$ on WILT embeddings $\\nu^G$, which is what makes it linear-time computable and gives each edge weight a direct reading: a color with a large weight contributes proportionally to the distance whenever its abundance differs between two graphs. Two normalizations adapt the distance to unequal graph sizes, namely size normalization ($\\dot{\\nu}^G = \\nu^G/|V_G|$), which generalizes $d_{\\mathrm{WWL}}$, and dummy-node normalization ($\\bar{\\nu}$), which generalizes $d_{\\mathrm{WLOA}}$, and the expressiveness theorem (Theorem 5.4) links the former to mean-pooled and the latter to sum-pooled MPNNs. The learning procedure (Algorithm 1) fits the edge weights by minimizing the squared error between $d_{\\mathrm{WILT}}$ and $d_{\\mathrm{MPNN}}$ over graph pairs, after which the largest weights identify the WL colors whose presence or absence most affects the MPNN distance.","core_discovery":"The central claim is that the distance an MPNN places between graphs is not generic structural similarity but a task-specific pseudometric encoding the functional distance between prediction targets. Training consistently increases $\\mathrm{ALI}_k(d_{\\mathrm{MPNN}}, d_{\\mathrm{func}})$, and the Spearman correlation between this alignment and performance is positive and consistent across datasets (accuracy for classification, RMSE for regression), whereas the same check against structural distances shows neither improvement nor correlation. Because MPNN embeddings are functions of multisets of WL colors, the paper defines $d_{\\mathrm{WILT}}$, the optimal transport distance on a weighted Weisfeiler-Leman labeling tree, and proves it equals a weighted $\\ell^1$ difference of WILT embeddings, giving linear-time computation; size and dummy-node normalizations recover the Wasserstein WL distance ($d_{\\mathrm{WWL}}$) and the WL optimal assignment distance ($d_{\\mathrm{WLOA}}$) as special cases, and expressiveness theorems match the size-normalized variant to mean-pooled MPNNs and the dummy-node variant to sum-pooled MPNNs. Distilled onto trained GCNs and GINs, the fitted $d_{\\mathrm{WILT}}$ approximates $d_{\\mathrm{MPNN}}$ much better than the fixed-weight baselines, and most edge weights collapse toward zero: with $\\ell^1$ regularization, about 95% of edges can be zeroed while keeping a better fit than $d_{\\mathrm{WLOA}}$. On Mutagenicity, the WL colors with the largest weights correspond to subgraphs with known mutagenic function, supporting the claim that MPNNs build their embedding geometry from a small set of task-relevant substructures.","pith_inferences":["The ALI-to-performance correlation is measured across a grid of model architectures and hyperparameters, so it is a between-model correlation; the paper does not test whether alignment rises within a single training run before accuracy does, which would be needed to argue that the alignment causes the performance rather than merely co-occurs with it.","The WILT fit is evaluated on the same graphs used to fit the edge weights, so an out-of-sample check, fitting on a training split and measuring RMSE and top-weight colors on a held-out split, would show whether the identified subgraphs are stable or an artifact of in-sample fitting.","The authors note the tree construction extends to higher-order WL hierarchies and to node-embedding distances; a natural next experiment is training WILT edge weights from scratch as a standalone interpretable graph classifier, a use they mention as a possibility.","In domains without established toxicophore lists, the distillation could serve as an automated hypothesis generator: the highest-weight WL colors propose candidate functional substructures that domain experts could then validate."],"forward_implications":["The property that tracks MPNN predictive performance is alignment of embedding distance with the task's functional distance, not alignment with task-agnostic structural distances; the paper's data show that training increases this alignment and that higher alignment correlates with higher accuracy or lower RMSE.","The WILTing distance with learned edge weights approximates $d_{\\mathrm{MPNN}}$ substantially better than the fixed-weight WL-based distances $d_{\\mathrm{WWL}}$ and $d_{\\mathrm{WLOA}}$, which are its special cases.","Size-normalized WILT is the right surrogate for mean-pooled MPNNs and dummy-node-normalized WILT for sum-pooled MPNNs, matching the expressiveness analysis; sum-pooled MPNNs are in fact not bounded by the size-normalized distance.","MPNN embedding geometry is sparse: only a small fraction of WL colors carry it, since distilling with $\\ell^1$ regularization keeps a good fit with roughly 95% of the edge weights zeroed.","The subgraphs behind the largest weights align with domain-known functional substructures (toxicophores in mutagenicity prediction), giving a global interpretation of where graphs sit relative to each other in MPNN space."],"supporting_citations":[{"why":"Defines the WL optimal assignment kernel whose distance $d_{\\mathrm{WLOA}}$ becomes a special case of dummy-node-normalized WILT and serves as a baseline.","marker":"Kriege et al., 2016"},{"why":"Defines the Wasserstein WL distance $d_{\\mathrm{WWL}}$ that size-normalized WILT generalizes and uses as a baseline.","marker":"Togninalli et al., 2019"},{"why":"Establishes that MPNNs are bounded in expressiveness by the WL test, the fact behind the pseudometric expressiveness theorems.","marker":"Xu et al., 2019"},{"why":"Supplies the tree-Wasserstein equivalence that turns $d_{\\mathrm{WILT}}$ into a weighted $\\ell^1$ difference computable in linear time.","marker":"Le et al., 2019"},{"why":"Provides the toxicophore substructure list used to judge whether the highest-weight WL colors are functionally meaningful.","marker":"Kazius et al., 2005"},{"why":"Proposes the tree mover's distance, the structural pseudometric whose alignment with $d_{\\mathrm{MPNN}}$ the paper tests and finds non-predictive.","marker":"Chuang & Jegelka, 2022"},{"why":"Supplies the TUDataset benchmarks (Mutagenicity, ENZYMES, IMDB-BINARY, COLLAB) used in the experiments.","marker":"Morris et al., 2020"},{"why":"Supplies the Lipophilicity benchmark used in the regression experiments.","marker":"Wu et al., 2018"}],"fun_headline_variants":["MPNN distance equals a weighted WL tree that spots task subgraphs","Linear-time WILT distills MPNN embeddings into functional subgraphs","Optimal transport on WL labeling trees explains MPNN distance","MPNNs learn task distances, WILT extracts the relevant subgraphs","WILT: an interpretable linear-time distance for MPNN embeddings"],"cache_read_input_tokens":29440,"weakest_assumption_plain":"The load-bearing premise is that the distilled WILT distance faithfully reproduces the MPNN's embedding distance, so that the largest learned edge weights mark the subgraphs that actually steer where graphs sit in embedding space, and the paper only checks this fit on the same graphs used to fit the weights.","fun_headline_variants_meta":{"raw":{"variants":["MPNN distance equals a weighted WL tree that spots task subgraphs","Linear-time WILT distills MPNN embeddings into functional subgraphs","Optimal transport on WL labeling trees explains MPNN distance","MPNNs learn task distances, WILT extracts the relevant subgraphs","WILT: an interpretable linear-time distance for MPNN embeddings"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000344,"raw_usage":{"total_tokens":1937,"prompt_tokens":1043,"completion_tokens":894,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":659,"completion_tokens_details":{"reasoning_tokens":801}},"tokens_in":659,"tokens_out":894,"duration_ms":11553,"temperature":1.0,"reasoning_tokens":801,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-07T12:18:10.453923+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Fit WILT edge weights on a training split of a dataset, then compute the RMSE between $d_{\\mathrm{WILT}}$ and $d_{\\mathrm{MPNN}}$ on held-out graphs and compare the highest-weight WL colors to those fitted on the full dataset: if the held-out fit degrades sharply relative to the in-sample fit, or the top-weight colors change substantially between splits, the claim that the learned weights identify the causal drivers of the MPNN distance fails its central test.","supporting_citations":[{"cited_title":"M., Giscard, P., and Wilson, R","cited_arxiv_id":null,"evidence_quote":"Defines the WL optimal assignment kernel whose distance $d_{\\mathrm{WLOA}}$ becomes a special case of dummy-node-normalized WILT and serves as a baseline."},{"cited_title":"Wasserstein Weisfeiler-Lehman graph kernels","cited_arxiv_id":null,"evidence_quote":"Defines the Wasserstein WL distance $d_{\\mathrm{WWL}}$ that size-normalized WILT generalizes and uses as a baseline."},{"cited_title":"Derivation and validation of toxicophores for mutagenicity prediction","cited_arxiv_id":null,"evidence_quote":"Provides the toxicophore substructure list used to judge whether the highest-weight WL colors are functionally meaningful."},{"cited_title":"and Jegelka, S","cited_arxiv_id":null,"evidence_quote":"Proposes the tree mover's distance, the structural pseudometric whose alignment with $d_{\\mathrm{MPNN}}$ the paper tests and finds non-predictive."},{"cited_title":"M., Bause, F., Kersting, K., Mutzel, P., and Neumann, M","cited_arxiv_id":null,"evidence_quote":"Supplies the TUDataset benchmarks (Mutagenicity, ENZYMES, IMDB-BINARY, COLLAB) used in the experiments."}],"review_version":1}