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 study, 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 × 106 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.