So3krates - Self-attention for higher-order geometric interactions on arbitrary length-scales

The application of machine learning (ML) methods in quantum chemistry has enabled the study of numerous chemical phenomena, which are computationally intractable with traditional ab initio methods. However, some quantum mechanical properties of molecules and materials depend on non-local electronic effects, which are often neglected due to the difficulty of modelling them efficiently. This work proposes a modified attention mechanism adapted to the underlying physics, which allows to recover the relevant non-local effects. Namely, we introduce spherical harmonic coordinates (SPHCs) to reflect higher order geometric information for each atom in a molecule, enabling a non-local formulation of attention in the SPHC space. Our proposed model So3krates -- a self-attention based message passing neural network (MPNN) -- uncouples geometric information from atomic features, making them independently amenable to attention mechanisms. We show that in contrast to other published methods, So3krates is able to describe quantum mechanical effects due to orbital overlap over arbitrary length scales. Further, So3krates is shown to match or exceed state-of-the-art performance on the popular MD-17 and QM-7X benchmarks, notably, requiring a significantly lower number of parameters while at the same time giving a substantial speedup compared to other models.

• General Science

• Machine Intelligence