How simple can you go? An off-the-shelf transformer approach to molecular dynamics

Max Eissler, Tim Korjakow, Stefan Ganscha, Oliver T.\ Unke, Klaus-Robert M\"uller, Stefan Gugler

Main figure for How simple can you go? An off-the-shelf transformer approach to molecular dynamics

Abstract

Most current neural networks for molecular dynamics (MD) include physical inductive biases, resulting in specialized and complex architectures. This is in contrast to most other machine learning domains, where specialist approaches are increasingly replaced by general-purpose architectures trained on vast datasets. In line with this trend, several recent studies have questioned the necessity of architectural features commonly found in MD models, such as built-in rotational equivariance or energy conservation. In this work, we contribute to the ongoing discussion by evaluating the performance of an MD model with as few specialized architectural features as possible. We present a recipe for MD using an edge transformer (ET), an ``off-the-shelf'' transformer architecture that has been minimally modified for the MD domain, termed MD-ET. Our model implements neither built-in equivariance nor energy conservation. We use a simple supervised pretraining scheme on $ $30 million molecular structures from the QCML database. Using this ``off-the-shelf'' approach, we show state-of-the-art results on several benchmarks after fine-tuning for a small number of steps. Using MD-ET as a simple and expressive testbed, we examine the effects of being only approximately equivariant and energy conserving for MD simulations and thereby try to evaluate the practical usefulness of unconstrained MD models. While our model exhibits runaway energy increases on larger structures, we show approximately energy-conserving NVE simulations for a range of small structures.

Benchmark Results

Tested on 1 dataset 5 baseline values reported too low

Tables show presented results in global context (best 3 papers + GCN + MLP).

Quantum Chemistry

#ModelMAE ↓Paper
1
QHNet
0.0000
2
QHNetV2
0.0000
3
PhiSNet
0.0000
7
NequIP
actual #7 vs paper #114 107 places worse
0.0009
best reported
actual #7 vs paper #114 107 places worse
paper 1.1000
22
PaiNN
actual #22 vs paper #119 97 places worse
0.0100
best reported
actual #22 vs paper #119 97 places worse
paper 3.8000
46
SchNet
actual #46 vs paper #125 79 places worse
0.0500
best reported
actual #46 vs paper #125 79 places worse
paper 16.8000
52
SphereNet
actual #52 vs paper #114 62 places worse
0.0520
best reported
actual #52 vs paper #114 62 places worse
paper 1.5000
64
DimeNet++
actual #64 vs paper #119 55 places worse
0.0640
best reported
actual #64 vs paper #119 55 places worse
paper 4.2000
paper 4.2000 Δ at pub 4.1360 Δ today 4.1360
113
MD
1.0000
This paper
114
Orb
1.0000
This paper
116
GemNet-dT
1.7000
This paper
122
DeepPot-SE
8.9000
This paper
124
ForceNet
9.9000
This paper
SOTA performance (not significantly different from the best) Baseline / comparison (this paper) Global context (other papers) Best at pub — reported result worse than the best known when published (best reported in green)