Enhancing the local expressivity of geometric graph neural networks

A central operation in geometric graph neural networks (GNNs) is the equivariant pairwise embedding, which encodes the local environment of each node as a learned representation. In this work, we examine the role of the pairwise embedding and consider a series of generalizations of its functional form beyond previous work. The new embeddings that we design considerably advance the state of the art in challenging distributions: as a highlight, when applied as an interatomic potential, we achieve a 29% relative reduction of force errors on diverse allotropes of lithium-intercalated carbon with a 4-fold reduction in parameter count. Furthermore, we demonstrate improved transferrability in molecular datasets by varying the locality of the network according to the depth of the representation.