A fast neural network approach for direct covariant forces prediction in complex multi-element extended systems

Neural network force field (NNFF) is a method for performing regression on atomic structure-force relationships, bypassing the expensive quantum mechanics calculations that prevent the execution of long ab initio quality molecular dynamics (MD) simulations. However, most NNFF methods for complex multi-element atomic systems indirectly predict atomic force vectors by exploiting just atomic structure rotation-invariant features and network-feature spatial derivatives, which are computationally expensive. Here, we show a staggered NNFF architecture that exploits both rotation-invariant and -covariant features to directly predict atomic force vectors without using spatial derivatives, and we demonstrate 2.2x NNFF-MD acceleration over a state-of-the-art C++ engine using a Python engine. This fast architecture enables us to develop NNFF for complex ternary- and quaternary-element extended systems composed of long polymer chains, amorphous oxide and surface chemical reactions. The rotation-invariant-covariant architecture described here can also directly predict complex covariant vector outputs from local environments, in other domains beyond computational material science. Neural network force fields promise to bypass the computationally expensive quantum mechanical calculations typically required to investigate complex materials, such as lithium-ion batteries. Mailoa et al. accelerate these approaches with an architecture that exploits both rotation-invariant and -covariant features separately.