Skip to main content
Back to timeline
arXivSource publication:

Molecular graph networks stop reading only the last layer: PRISM lets each atom pick its message-passing depth, improving four of six endpoints on a pre-trained GIN

Synopsis

Molecular properties depend on chemical structure at different ranges, yet graph neural networks usually fix one propagation depth for every atom. PRISM turns each atom's own encoded features into a bounded scalar that positions a unimodal distribution over retained message-passing states, so the allocated depth is analytically tractable and moves monotonically with that scalar. Under matched scaffold partitions and tuning budgets, PRISM trained from scratch reaches the lowest or tied-lowest mean error on ESOL and FreeSolv, and on a pre-trained GIN it improves on last-layer readout on four of six endpoints.

Source-provided article image: Atom-Wise Allocation of Message-Passing Scale for Molecular Property Prediction
Figure 1 · arXiv · Page 4

Interpretation

A molecular graph network can now allocate message-passing depth per atom instead of sharing one propagation scale across the molecule. The prior practice was to read only the final layer, or to apply one shared layer-weight profile to every atom, so a heteroatom could not favor its immediate environment while an aromatic carbon favored a more extended substructure. On five shared Bemis-Murcko scaffold partitions, PRISM trained from scratch has a mean RMSE of 0.741 on ESOL versus 0.746 for GIN, and 0.459 on FreeSolv versus 0.459 for EquiformerV2, a margin of 2×10⁻⁴ that amounts to a tie.

The allocation rule is a plug-in layer-combination module that does not depend on a particular message-passing operator. The module changes only how the retained layer states are weighted and summed; the encoder, state bank, readout and training schedule stay fixed. Attached alone to GCN, GIN and GAT backbones, the rule gives the best of the three values on 7 of 12 backbone-endpoint pairs; on FreeSolv, relative to last-layer readout, GCN RMSE falls from 0.951 to 0.673 and GIN from 0.676 to 0.596.

Each atom's routing scalar is produced from its input encoding through a tanh bound into (−1,1), then clipped into a center over the state bank, and a Gaussian layer-weight profile forms a unimodal distribution. One coordinate translates the whole profile, keeping neighboring layer weights coupled rather than setting each layer weight freely. At the five-layer initialization, routing scalars of −0.8, 0 and +0.8 give allocated indices of 0.409, 2.000 and 3.591; in trained models the within-model standard deviation of the unclipped center is 0.70 to 1.18 index units across the six datasets, with 83% to 94% of that variation lying within molecules.

Permuting the learned layer-weight profiles among a molecule's atoms raises error, so which atom receives which profile carries predictive information. Permutation preserves the molecule's multiset of profiles and its mean allocated depth and breaks only the atom-to-profile correspondence. In the final model, permutation raises Lipophilicity RMSE by 0.107 and lowers SIDER ROC-AUC by 0.024, both significant over five models; on the pre-trained bank, permutation raises RMSE by 0.140, 0.109 and 0.219 on ESOL, FreeSolv and Lipophilicity.

Perspective

The demonstrated scope is a comparison on six MoleculeNet endpoints, ESOL, FreeSolv, Lipophilicity, BACE, BBBP and SIDER, under Bemis-Murcko scaffold partitions, plus a depth sweep from 5 to 15 layers on Peptides-struct; all graph models use the molecular graph only, without three-dimensional coordinates.

The standard deviation across the five scaffold partitions is comparable to the gaps among leading models, and differences smaller than one partition standard deviation are not resolved by five partitions; the Peptides-struct depth sweep is a single run per configuration, and in that sweep the allocation parameters are trained at a lower learning rate and held fixed early, so the error advantage over GIN cannot be credited to the allocation alone. Default routing reads only the atom's own descriptors, so atoms with identical input descriptors receive identical profiles in any molecular environment; whether the learned type-to-scale assignment corresponds to a physical length scale remains open.