Expressivity and Generalization: Fragment-Biases for Molecular GNNs

TL;DR

This work addresses expressivity and generalization gaps in molecular GNNs by introducing Fragment-WL, a unified test that analyzes fragment-biased models, and FragNet, a message-passing architecture operating on both the molecular graph and a higher-level fragment graph. It establishes a Hierarchy NF-WL < FR-WL < HLG-WL and demonstrates that higher-level fragment abstractions enhance expressivity, aided by an infinite-vocabulary rings-paths fragmentation and an ordinal encoding scheme. Empirically, FragNet achieves state-of-the-art or near-state-of-the-art results on ZINC and Peptides among GNNs, with strong generalization including zero-shot QM9 and robustness to unseen fragments, while maintaining linear time complexity. The approach reduces over-squashing via shortcutting through the fragment graph, showing substantial practical impact for molecular property prediction and related tasks, and suggesting future extensions to orbit information and multi-task learning.

Abstract

Although recent advances in higher-order Graph Neural Networks (GNNs) improve the theoretical expressiveness and molecular property predictive performance, they often fall short of the empirical performance of models that explicitly use fragment information as inductive bias. However, for these approaches, there exists no theoretic expressivity study. In this work, we propose the Fragment-WL test, an extension to the well-known Weisfeiler & Leman (WL) test, which enables the theoretic analysis of these fragment-biased GNNs. Building on the insights gained from the Fragment-WL test, we develop a new GNN architecture and a fragmentation with infinite vocabulary that significantly boosts expressiveness. We show the effectiveness of our model on synthetic and real-world data where we outperform all GNNs on Peptides and have 12% lower error than all GNNs on ZINC and 34% lower error than other fragment-biased models. Furthermore, we show that our model exhibits superior generalization capabilities compared to the latest transformer-based architectures, positioning it as a robust solution for a range of molecular modeling tasks.

Figures (10)

• Figure 1: Example graph ${G}$ with corresponding augmented variants. ${\text{NF}}({G})$ includes node features, ${\text{FR}}({G})$ also includes a representation for each fragment and ${\text{HLG}}({G})$ also has connections between neighboring fragment represenations.
• Figure 2: Graphs ${G}^1$ and ${G}^2$ with their corresponding higher-level graph of fragments. The edges of the fragment representation to the vertices of ${G}^1$ and ${G}^2$ are omitted. ${G}^1$ and ${G}^2$ are indistinguishable by WL, NF-WL and FR-WL but distinguishable by HLG-WL as the higher-level graphs exhibit different connections from the 3-ring nodes to the 4-ring node.
• Figure 3: Overview of our model and our fragmentation. The molecular graph is fragmented with our rings-paths fragmentation into three cycles, three paths, and a junction node. The figure shows the messages $m^{t}_{{F} \rightarrow f}$, $m^{t}_{{\mathcal{V}}\rightarrow f}$ to one fragment $f$, and the messages $m^t_{{\mathcal{V}}\rightarrow v}$, $m^t_{{F} \rightarrow v}$ to one vertex $v$.
• Figure 4: Ordinal encoding applied to a 2-path, 4-path, 5-ring, and 6-ring. The encoding comprises two components: one learned embedding $e$ for every fragment class (i.e., path, cycle, or junction) and another learned embedding $s$ that is proportionally scaled based on the fragment size.
• Figure 5: Recovery rate of messages sent from the star node to all other nodes. A recovery rate of 0.1 corresponds to random guessing. The first graph has no fragmentation, the second one a rings fragmentation (like in CIN/CIN++), the third a rings and paths fragmentation (like in our model).

Theorems & Definitions (56)

• Definition 4.1
• Definition 4.2
• Definition 4.3
• Definition 4.4
• Theorem 4.5
• Theorem 4.6
• Theorem 4.7
• Theorem 4.8
• Theorem 4.9
• Theorem 5.1