From Peptides to Nanostructures: A Euclidean Transformer for Fast and Stable Machine Learned Force Fields

TL;DR

This paper introduces SO3krates, a transformer-based machine learned force field that combines efficiency, accuracy, and stability, enabling extended molecular dynamics simulations and exploration of complex molecular conformations.

Contribution

The paper presents SO3krates, a novel transformer architecture that uses sparse equivariant representations and self-attention to improve stability and speed in MLFFs for molecular simulations.

Findings

Achieves stable MD trajectories for peptides and large structures

Balances stability with exploration of new conformations

Demonstrates efficiency and accuracy in quantum property analysis

Abstract

Recent years have seen vast progress in the development of machine learned force fields (MLFFs) based on ab-initio reference calculations. Despite achieving low test errors, the reliability of MLFFs in molecular dynamics (MD) simulations is facing growing scrutiny due to concerns about instability over extended simulation timescales. Our findings suggest a potential connection between robustness to cumulative inaccuracies and the use of equivariant representations in MLFFs, but the computational cost associated with these representations can limit this advantage in practice. To address this, we propose a transformer architecture called SO3krates that combines sparse equivariant representations (Euclidean variables) with a self-attention mechanism that separates invariant and equivariant information, eliminating the need for expensive tensor products. SO3krates achieves a unique combination of accuracy, stability, and speed that enables insightful analysis of quantum properties of matter on extended time and system size scales. To showcase this capability, we generate stable MD trajectories for flexible peptides and supra-molecular structures with hundreds of atoms. Furthermore, we investigate the PES topology for medium-sized chainlike molecules (e.g., small peptides) by exploring thousands of minima. Remarkably, SO3krates demonstrates the ability to strike a balance between the conflicting demands of stability and the emergence of new minimum-energy conformations beyond the training data, which is crucial for realistic exploration tasks in the field of biochemistry.