Learning Small Molecule Energies and Interatomic Forces with an Equivariant Transformer on the ANI-1x Dataset

Abstract: Accurate predictions of interatomic energies and forces are essential for high quality molecular dynamic simulations (MD). Machine learning algorithms can be used to overcome limitations of classical MD by predicting ab initio quality energies and forces. SE(3)-equivariant neural network allow reasoning over spatial relationships and exploiting the rotational and translational symmetries. One such algorithm is the SE(3)-Transformer, which we adapt for the ANI-1x dataset. Our early experimental results indicate through ablation studies that deeper networks - with additional SE(3)-Transformer layers - could reach necessary accuracies to allow effective integration with MD. However, faster implementations of the SE(3)-Transformer will be required, such as the recently published accelerated version by Milesi.