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

2025-03-03

AI Summary

This work investigates whether a minimalist, universal architecture can replace complex physics-informed models in molecular dynamics (MD) simulation. We propose MD-ET—a lightweight variant of the edge Transformer—that deliberately discards built-in rotational equivariance and energy conservation priors. To our knowledge, MD-ET is the first purely data-driven, non-equivariant, non-conservative “off-the-shelf” Transformer achieving state-of-the-art (SOTA) performance on MD tasks. It is pretrained on 30 million structures from the QCML database under supervised learning, followed by fine-tuning and NVE ensemble simulation. We further introduce a novel error attribution framework to disentangle errors arising from non-equivariance versus numerical integration. Experiments demonstrate: (1) SOTA accuracy across multiple MD benchmarks; (2) approximate energy conservation in small-molecule systems; and (3) the first empirical characterization of bounded energy drift in large-scale systems—revealing its asymptotic saturation behavior.

Technology Category

• Machine Learning: Deep Neural Architectures and Foundation Models
• Multiagent Systems: Other Foundations of Multi Agent Systems
• Search and Optimization: Sampling/Simulation-based Search

Application Category

• User Modeling, Personalization and Recommendation: Accountability, Transparency, and Ethics for personalization
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Security and Privacy: Data transparency and provenance

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, 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 pre-training scheme on $sim$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. Additionally, we examine the effects of being only approximately equivariant and energy conserving for MD simulations, proposing a novel method for distinguishing the errors resulting from non-equivariance from other sources of inaccuracies like numerical rounding errors. While our model exhibits runaway energy increases on larger structures, we show approximately energy-conserving NVE simulations for a range of small structures.

Problem

• Evaluates minimal architectural features in molecular dynamics models.
• Proposes an off-the-shelf transformer for molecular dynamics simulations.
• Assesses effects of non-equivariance and energy conservation in simulations.

Innovation

• Off-the-shelf transformer for molecular dynamics
• Minimally modified Edge Transformer architecture
• Supervised pre-training on QCML database