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

Philipp Thölke; Gianni De Fabritiis

Source-linked AI summary

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

Philipp Thölke, Gianni De Fabritiis

arXiv:2202.02541v2Feb 5, 2022cs.LGcs.AIphysics.chem-ph

TL;DR

Quantum-mechanical property prediction requires balancing accuracy with computational efficiency. TorchMD-NET introduces an equivariant Transformer with distance-aware attention and evaluates it across QM9, MD17, and ANI-1. The model achieves strong benchmark performance while attention analysis distinguishes representations learned from static and dynamical data.

Problem
Quantum-mechanical property prediction is limited by the computational cost of accurate methods, motivating efficient molecular potentials.
Method
TorchMD-NET uses an equivariant Transformer with learned atomic and neighborhood features, distance-aware attention, and scalar–vector feature updates.
Results
TorchMD-NET achieves state-of-the-art accuracy across the reported benchmarks, including all MD17 targets except benzene force prediction.
Takeaways & Limitations
Attention analysis indicates that dynamical datasets emphasize hydrogen more strongly than static QM9 structures, supporting the importance of off-equilibrium conformations.
Takeaways & Limitations
The comparison is constrained by scope: NequIP later reported better MD17 accuracy, while using high-order spherical harmonics likely increases computational cost.

Abstract

from arXiv · show

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.

1 INTRODUCTION

TorchMD-NET addresses the tension between quantum-mechanical accuracy and computational efficiency with an equivariant Transformer based on learned atomic types and coordinates. It is evaluated across QM9, MD17, and ANI-1, while attention analysis examines learned representations across data regimes.

Quantum-mechanical methods provide accurate molecular-property estimates but remain computationally impractical at larger scales.
Existing molecular-potential approaches include Behler–Parrinello representations, ANI, graph convolutions, and rotationally equivariant networks.
Machine-learning molecular potentials offer a parametric alternative suited to large datasets, whereas kernel methods generally perform best with scarce data.
TorchMD-NET introduces an equivariant Transformer that uses attention and learned featurization of atomic types and coordinates for quantum-property prediction.
The paper analyzes attention weights and latent representations across energy-minimized, molecular-dynamics, and normal-mode-sampled data.

2 METHODS

TorchMD-NET represents molecules as graphs with learned atomic and distance-based features, then uses modified attention to update scalar and vector representations. Equivariant update layers exchange scalar, vector, and directional information before producing molecular predictions.

2 METHODS: Molecules are represented as graphs in which self-attention constructs fully connected node interactions augmented with interatomic-distance edge information.
2 METHODS: The architecture comprises an embedding layer, repeated update layers, and an output network that aggregates atomwise predictions.
2.2 EMBEDDING LAYER: Each atom receives intrinsic and neighborhood embeddings; the neighborhood component is distance-filtered to distinguish atomic information from pairwise interactions.
2.2 EMBEDDING LAYER: Interatomic distances are expanded with exponential radial basis functions and smoothly truncated at the 5 Å cutoff using a cosine cutoff.
2.3 MODIFIED ATTENTION MECHANISM: Modified multi-head attention projects distances into filters and combines queries, keys, and distance projections to compute attention weights.
2.3 MODIFIED ATTENTION MECHANISM: SiLU replaces softmax in attention, preserving the distance cutoff while omitting the traditional softmax scaling factor.
2.3 MODIFIED ATTENTION MECHANISM: A continuous-filter graph convolution in the value pathway injects interatomic distances directly into feature vectors, not only into attention weights.
2.4 UPDATE LAYER: Update layers compute interatomic interactions and exchange information between scalar and vector features using directional edge information and equivariant transformations.

3 EXPERIMENTS AND RESULTS

TorchMD-NET is evaluated across QM9, MD17, and ANI-1, including equilibrium, dynamical, and off-equilibrium molecular data. The experiments combine benchmark comparisons, attention analysis, ablations, and model-size checks to assess accuracy, interpretation, and efficiency.

QM9: QM9 evaluation compares TorchMD-NET with invariant, covariant, and equivariant architectures across all reported targets using MAE.LieTF denotes the best-performing LieTransformer-T3+SO3 Aug variant.
MD17: MD17 models are trained in a limited-data setting using 1000 samples, with a combined energy-and-force loss weighted toward forces.The remaining trajectory data supports evaluation against SchNet, PhysNet, DimeNet, PaiNN, and NequIP.
ANI-1: ANI-1 supplies 22,057,374 off-equilibrium configurations from 57,462 molecules, split into 80% training, 5% validation, and 15% testing data.The configurations are generated through exhaustive normal-mode sampling of energy-minimized molecules.
Attention weight analysis: Attention analysis extracts rollout matrices across all heads and layers, then compares element-pair attention with conditional bond probabilities.The analysis covers ANI-1, QM9, and MD17 test sets.
Attention weight analysis: Training data changes attention patterns: QM9 assigns little importance to hydrogens, whereas dynamical ANI-1 and MD17 data produce strong hydrogen attention.The authors relate this pattern to hydrogen-bond-type interactions and argue that dynamical datasets are important for learning chemical representations.
Ablation studies: Ablating neighbor embeddings raises QM9 U0 MAE roughly 6%, from 6.24 to 6.60 meV, while replacing them with an update layer raises it around 10%, from 6.24 to 6.85 meV.The update-layer comparison may partly reflect its approximately 4.6-times larger parameter count.
Model size: Smaller TorchMD-NET models remain competitive with PaiNN and NequIP in parameter count and outperform previous state-of-the-art results on MD17.The standard MD17 and ANI-1 configuration has 1.34M trainable parameters, versus 600k for PaiNN and 290k for NequIP.

4 DISCUSSION

TorchMD-NET combines an attention-based equivariant architecture with analysis of learned molecular representations. The discussion reports strong benchmark performance while noting a newer method with better MD17 accuracy and likely slower computation.

TorchMD-NET achieves high QM9 accuracy and is particularly effective for predicting energies and atomic forces in molecular dynamics.The architecture sets a new state-of-the-art on all MD17 targets except benzene force prediction and can operate in a low-data regime.
Attention-weight analysis reveals that learned molecular representations depend on whether training data contain equilibrium or off-equilibrium conformations.The model’s attention patterns differ between energy-minimized molecules and molecular-dynamics trajectories.
NequIP’s updated preprint reports better MD17 accuracy than TorchMD-NET, while likely requiring substantially slower high-order spherical-harmonic features.This comparison qualifies the paper’s state-of-the-art claim for MD17.

SOFTWARE AND DATA

The paper provides public datasets and source code alongside implementation details for the equivariant Transformer. Training uses distributed computation, while the listed hyperparameters distinguish QM9, MD17, and ANI-1 configurations.

TorchMD-NET is implemented in PyTorch with PyTorch Geometric and pytorch-lightning, and its datasets and source code are publicly available.The release covers training, running, and analyzing the presented models.
Training uses distributed data parallelism across two NVIDIA RTX 2080 Ti GPUs.
Training times are 16h for QM9, 10h for MD17, and 83h for ANI-1.
The ET differs from the SE(3)-Transformer by using interatomic distances instead of expensive Clebsch-Gordan coefficients and spherical harmonics, with separate scalar and equivariant pathways.

B IMPORTANCE OF HYDROGEN

Hydrogen matters more for models trained on molecular-dynamics data than for models trained on equilibrium-only QM9 data. Removing hydrogen therefore causes a substantially larger energy-accuracy loss on MD17 and an additional force-prediction decline.

Models trained on QM9 mostly attend to carbon, whereas models trained on molecular-dynamics trajectories strongly attend to carbon-hydrogen interactions.
The energy loss from excluding hydrogen is one order of magnitude larger for MD17 than for QM9 ground-state molecules.
Force-prediction accuracy drops by another 1.5x compared with MD17 energy predictions when hydrogen is excluded.
On QM9, removing hydrogen changes the reported total-energy MAE from 6.37 meV to 20.83 meV, a relative change of 227.0%.

C ATOM DISPLACEMENT

The atom-displacement analysis compares attention to equilibrium and displaced atoms in H, C, and O molecules. Models trained on off-equilibrium conformations attend more strongly to displaced hydrogen than models trained only on equilibrium structures.

Figure 5 averages absolute attention scores after displacing a single atom from its equilibrium structure.The comparison is restricted to molecules containing hydrogen, carbon, and oxygen.
QM9-trained models focus on carbon and largely neglect hydrogen when evaluated on displaced-atom structures.
Models trained on off-equilibrium conformations show higher attention to displaced hydrogen atoms.
Figure 5 uses blue bars for attention to equilibrium atoms, orange bars for attention involving the displaced atom, and black bars for standard deviations.Attention scores are normalized within each molecule.

D EQUIVARIANT TRANSFORMERS WITH REDUCED PARAMETER COUNT

TorchMD-NET remains competitive with state-of-the-art molecular-potential models even when its parameter count is reduced to match PaiNN and NequIP. These comparisons test whether its performance reflects architecture rather than simply model size.

Reduced-parameter ET models matched PaiNN and NequIP model sizes to test whether performance gains came from architecture rather than parameter count.The comparison used 600k- and 290k-parameter reference sizes.
The smaller ET models remained competitive and outperformed state-of-the-art models on most MD17 targets.Table 7 reports the corresponding energy and force MAE comparisons.
Table 6 compares the full ET with PaiNN-sized and NequIP-sized variants using the hyperparameters that changed.

E ABLATION OF EQUIVARIANT FEATURES

Removing equivariant features worsened TorchMD-NET’s errors, especially on the dynamical MD17 data. The reported increases are substantially larger for aspirin than for QM9 total energy.

146%: removing equivariance increased aspirin energy MAE from 5.37 to 13.23 meV.
174%: removing equivariance increased aspirin force errors from 11.05 to 30.27 meV/Å.
6%: removing equivariance increased QM9 total-energy U0 MAE from 6.24 to 6.64 meV.
The authors hypothesize that equivariant features are particularly useful for dynamical data with non-zero forces.

F MOLECULAR REPRESENTATION BY DATASET

Attention patterns differ across QM9, MD17, and ANI-1. QM9 models emphasize carbon-carbon interactions, whereas MD17 and ANI-1 models focus more on hydrogen-carbon and hydrogen-oxygen interactions; dataset composition helps interpret these patterns.

Dataset-dependent attention: QM9 attention focuses largely on carbon-carbon interactions, while MD17 and ANI-1 emphasize hydrogen-carbon and hydrogen-oxygen interactions.The visualizations show the 10 largest absolute attention weights for random molecules or conformations.
Dataset-dependent attention: The visualized examples use random test-set molecules or conformations, ensuring that the model had not seen those structures during training.
Dataset composition: Nitrogen and fluorine receive comparatively limited attention partly because they are underrepresented across the datasets.MD17 contains only one nitrogen-containing molecule, while fluorine occurs only in QM9.

H MD17 ATTENTION WEIGHTS

MD17 attention analysis compares learned attention with bond probabilities across target molecules. Hydrogen receives strong attention for most molecules, while aspirin, malondialdehyde, and uracil additionally show strong oxygen focus.

Attention patterns: Hydrogen atoms receive high attention across most MD17 target molecules.The analysis covers aspirin, benzene, ethanol, malondialdehyde, naphthalene, salicylic acid, toluene, and uracil.
Attention patterns: Aspirin, malondialdehyde, and uracil additionally exhibit strong attention to oxygen.
Bond-probability comparison: The bond-probability matrix gives the conditional probability of a bond between atom types zj and zi given zi, with row sums equal to one.Nbonded(zi, zj) denotes the total number of bonds between the atom types in the molecular collection.
Bond-probability comparison: Figure 8 compares rolled-out attention scores and bond probabilities for individual MD17 molecules, alongside atom-type counts.
Higher-accuracy trajectories: Additional ET results were evaluated on higher-accuracy CCSD and CCSD(T) trajectories using the original MD17 training protocol.The associated table reports energy and force MAE averaged over two random splits.