Atom-Wise Allocation of Message-Passing Scale for Molecular Property Prediction
Regina Qianru Zhang, Honggang Wen, Zongru Li, Jinsong Ouyang, Zhiqin Chu, Pietro Liò, Xiaofeng Liu, Siu-Ming Yiu
Abstract
Molecular properties depend on chemical structure at different ranges, from local polar groups governing solvation to extended substructures shaping bioactivity, yet graph neural networks usually fix one propagation scale for every atom. We present PRISM (Per-atom Routing over Intermediate States of Message Passing), in which each atom's encoded features produce a bounded scalar that positions a unimodal distribution over retained message-passing states and matching circular-fingerprint radii. Because one coordinate translates the whole profile, allocated depth responds to the routing scalar in closed form: at fixed temperature it is nondecreasing, its sensitivity scales with the profile variance, and it saturates in the clipped regions. Beyond providing closed-form analyzability, this constrained one-parameter profile matches or significantly outperforms parameter-matched, unconstrained node-wise gates and Jumping Knowledge attention. Under scaffold partitions with equal tuning budgets against eight graph, fingerprint, and geometric baselines, PRISM trained from scratch without three-dimensional coordinates achieves the lowest mean error on aqueous solubility and hydration free energy and strong accuracy on BACE and BBBP. As a general layer-combination rule, PRISM transfers directly to GCN, GIN, and GAT backbones, reducing hydration free energy error by up to 29% relative to last-layer readout. Applied to a pre-trained GIN, it improves on last-layer readout on four of six endpoints and gives the lowest lipophilicity error overall. In trained models, allocated depth differs systematically between atom types in a task-dependent way, and permuting profiles among a molecule's atoms raises prediction error on several endpoints, so the rule acts as a learned counterpart of atom typing.