TorchMD-NET: Equivariant Transformers for Neural Network based Molecular Potentials

TL;DR

TorchMD-NET introduces a rotationally equivariant Transformer tailored for molecular potentials, leveraging edge-aware attention and update mechanisms to predict energies and forces with high accuracy. The architecture demonstrates state-of-the-art results on QM9, MD17, and ANI-1, and provides insights into learned representations via attention-weight analyses. A key finding is the critical role of off-equilibrium conformations in training and evaluation, particularly for dynamics-related properties. The work also analyzes model components through ablations and shows favorable computational efficiency, making it a practical alternative to existing SE(3)-transformer-based approaches.

Abstract

The prediction of quantum mechanical properties is historically plagued by a trade-off between accuracy and speed. Machine learning potentials have previously shown great success in this domain, reaching increasingly better accuracy while maintaining computational efficiency comparable with classical force fields. In this work we propose TorchMD-NET, a novel equivariant transformer (ET) architecture, outperforming state-of-the-art on MD17, ANI-1, and many QM9 targets in both accuracy and computational efficiency. Through an extensive attention weight analysis, we gain valuable insights into the black box predictor and show differences in the learned representation of conformers versus conformations sampled from molecular dynamics or normal modes. Furthermore, we highlight the importance of datasets including off-equilibrium conformations for the evaluation of molecular potentials.

Figures (10)

• Figure 1: Overview of the equivariant Transformer architecture. Thin lines: scalar features in $\mathbb{R}^{F}$, thick lines: vector features in $\mathbb{R}^{3 \times F}$, dashed lines: multiple feature vectors. (a) Transformer consisting of an embedding layer, update layers and an output network. (b) Residual update layer including attention based interatomic interactions and information exchange between scalar and vector features. (c) Modified dot-product attention mechanism, scaling values (blue) by the attention weights (red).
• Figure 2: Depiction of bond probabilities and attention scores extracted from the ET model of TorchMD-NET using QM9 (total energy $U_0$), MD17 (average over 8 discussed molecules) and ANI-1 testing data. Attention scores are given as $z_i$ attending $z_j$, bond probabilities follow the same idea, showing the conditional probability of a bond between $z_i$ and $z_j$, given $z_i$. Darker colors correspond to larger values, element pairs without data are grayed out. See Appendix \ref{['appendix:element-frequency']} for an overview of elemental composition in the respective datasets.
• Figure 3: Visualization of five molecules from the QM9 dataset with attention scores corresponding to models trained on ANI-1, MD17 (uracil) and QM9. Blue and red lines represent negative and positive attention scores respectively.
• Figure 4: Averaged attention weights extracted from the ET on the QM9 test set (molecules consisting of H, C, and O only) with a displacement of 0.4Å in single atoms. Blue bars show attention towards atoms in equilibrium locations, orange bars correspond to attention weights involving the displaced atom. Attention scores are normalized inside each molecule. The black bars show the attention weights' standard deviation.
• Figure 5: Visualization of 10 largest attention weights by absolute value on random molecules from QM9 (a), MD17-aspirin (b) and ANI-1 (c). Each column shows the same molecule, rows correspond to the same ET model trained on QM9, MD17-uracil and ANI-1 respectively.