Direct Molecular Polarizability Prediction with SO(3) Equivariant Local Frame GNNs

TL;DR

The paper presents an $SO(3)$-equivariant GNN that directly regresses molecular polarizability tensors by leveraging per-atom local reference frames and separate scalar, vector, and tensor channels. It introduces a charge-weighted PCA frame construction and edge-wise tensorial message passing to preserve rotational symmetry and improve geometric information exchange. On QM7-X, the tensorial model consistently outperforms a scalar baseline across multiple tensor-related metrics, with equivariance demonstrated under rotations and end-to-end frame recomputation. The work advances geometry-aware molecular property prediction and points to robust frame definitions and broader benchmarks as avenues for future impact.

Abstract

We introduce a novel equivariant graph neural network (GNN) architecture designed to predict the tensorial response properties of molecules. Unlike traditional frameworks that focus on regressing scalar quantities and derive tensorial properties from their derivatives, our approach maintains $SO(3)$-equivariance through the use of local coordinate frames. Our GNN effectively captures geometric information by integrating scalar, vector, and tensor channels within a local message-passing framework. To assess the accuracy of our model, we apply it to predict the polarizabilities of molecules in the QM7-X dataset and show that tensorial message passing outperforms scalar message passing models. This work marks an advancement towards developing structured, geometry-aware neural models for molecular property prediction.

Figures (1)

• Figure 1: Equivariant Architecture for Rank-2 Tensors