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

TL;DR

This paper adapts an SE(3)-equivariant transformer neural network to predict molecular energies and forces on the ANI-1x dataset, aiming to improve molecular dynamics simulations with higher accuracy.

Contribution

It introduces a novel application of the SE(3)-Transformer to molecular energy prediction and demonstrates potential for improved accuracy with deeper networks.

Findings

Deeper SE(3)-Transformer networks show promise for accurate energy and force predictions.

Preliminary results suggest potential for integration into molecular dynamics simulations.

Faster implementations are needed for practical use.

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.