Impact of Local Descriptors Derived from Machine Learning Potentials in Graph Neural Networks for Molecular Property Prediction

In this study, we present a framework aimed at enhancing molecular property prediction through the integration of local descriptors obtained from large-scale pretrained machine learning potentials into three-dimensional graph neural networks (3D GNNs). As an illustration, we developed an EGNN-PFP model by integrating descriptors derived from the preferred potential (PFP) features, acquired through Matlantis, into an equivariant graph neural network (EGNN), and evaluated its effectiveness. When tested on the QM9 dataset, comprising small organic molecules, the proposed model demonstrated superior accuracy compared to both the original EGNN models and the baseline models without PFP-derived descriptors for 11 out of the 12 molecular properties. Furthermore, when evaluated on the tmQM dataset, which encompasses transition metal complexes, notable enhancements in performance were observed across all five target properties, indicating the significance of the local atomic environment surrounding transition metals. In essence, the proposed methodology is adaptable to any 3D GNN architecture, and further enhancements in prediction accuracy are anticipated when integrated with continually evolving GNN architectures.

💡 Research Summary

In this work the authors introduce a general framework that augments three‑dimensional graph neural networks (3D‑GNNs) with local atomic descriptors extracted from a large‑scale pretrained machine‑learning potential, the Preferred Potential (PFP). The central hypothesis is that conventional 3D‑GNNs, which rely solely on atomic numbers and Cartesian coordinates, implicitly assume that all relevant electronic information can be inferred from geometry alone. While this may hold for simple organic molecules, it becomes questionable for chemically diverse systems, especially transition‑metal complexes where d‑orbital occupancy, oxidation state, and ligand field effects play a decisive role.

To test the hypothesis, the authors first obtain a 256‑dimensional embedding for each atom from the final hidden layer of the PFP neural network, just before the energy prediction head. This embedding encodes the local electronic environment learned from millions of quantum‑chemical calculations. They then concatenate the embedding with the atomic number and four simple geometric features (distance to the molecular centre, number of neighbours within a cutoff, normalized local atomic density, and a normalized coordinate norm) to form the initial node feature vector. Edge features are constructed from the Euclidean distance and the cosine similarity of the two PFP embeddings, yielding a four‑dimensional edge descriptor that simultaneously captures spatial proximity and electronic similarity. An interaction weight (\gamma_{ij}= \cos_{ij}\exp(-d_{ij})) is defined to modulate the message passing.

The backbone GNN is an Equivariant Graph Neural Network (EGNN). In each EGNN layer, edge messages are computed by a two‑layer MLP, multiplied by (\gamma_{ij}), and aggregated (sum or mean) at each node. Node updates use a residual connection and another two‑layer MLP, while optional coordinate updates refine the geometry during training. An additional attention MLP learns a scalar importance for each edge, allowing the network to emphasize strong covalent bonds or, conversely, weak non‑bonded interactions as needed. After a stack of EGNN layers, node features are decoded, pooled (sum or mean) to a graph‑level representation, and finally passed through a graph‑level MLP to predict the target molecular property.

The methodology is evaluated on two benchmark datasets. QM9 contains ~133 k small organic molecules with 13 quantum‑chemical targets (dipole moment, polarizability, HOMO/LUMO energies, zero‑point energy, thermochemical quantities, etc.). tmQM comprises ~86 k transition‑metal complexes with five targets, including metal‑center partial charge and frontier orbital energies. For QM9, the EGNN‑PFP model outperforms the vanilla EGNN on 11 of the 12 properties, achieving an average mean absolute error (MAE) reduction of about 7 % and up to 10 % for dipole moments and HOMO‑LUMO gaps. For tmQM, the gains are even more pronounced: all five properties see MAE reductions of roughly 10 %–15 %, highlighting the importance of explicit electronic descriptors for metal‑containing systems.

Importantly, the authors demonstrate that the same PFP descriptors can be plugged into other 3D‑GNN architectures (e.g., DimeNet, PaiNN, SE(3)‑Transformer) with comparable improvements, confirming the architecture‑agnostic nature of the approach. The study therefore provides three key insights: (1) providing pretrained local electronic embeddings simplifies the learning task for GNNs and improves data efficiency; (2) such embeddings are especially valuable for chemically complex domains where geometry alone is insufficient; and (3) the framework is modular and can be combined with future, more expressive potentials such as MACE or Allegro, as well as with advanced GNN designs like multi‑scale attention or graph transformers.

The paper also discusses limitations. The PFP embedding dimension is fixed at 256, and exploring higher‑dimensional or multi‑layer representations could capture richer physics. Enabling coordinate updates sometimes destabilizes training, suggesting a need for more physically informed constraints. Finally, scalability to millions of molecules and memory consumption of large embeddings remain open challenges.

In conclusion, by integrating high‑quality, pretrained local descriptors into 3D‑GNNs, the authors achieve a notable leap in molecular property prediction accuracy, particularly for transition‑metal chemistry. The proposed strategy is broadly applicable, paving the way for more reliable, data‑efficient models in drug discovery, catalyst design, and materials science.