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

TL;DR

This work tackles the instability and extrapolation limitations of machine learned force fields in long MD runs. It introduces SO3krates, a Euclidean transformer that decouples invariant and equivariant information through Euclidean self-attention, avoiding expensive tensor products while maintaining directional sensitivity. The approach delivers high stability, accuracy, and substantially faster inference than prior equivariant models, enabling nanosecond-scale MD and extensive PES exploration for peptides and supramolecular structures, including thousands of minima with near-DFT validation. The findings demonstrate robust extrapolation to larger systems and higher temperatures, offering a practical path toward scalable, reliable MD with MLFFs that can inform biochemistry and materials science at unprecedented scales.

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.

Figures (20)

• Figure 1: (a) Illustration of an invariant convolution, an SO(3) convolution and of the Euclidean attention mechanism that underlies the SO3krates transformer. We decompose the representation of molecular structure into high dimensional invariant features and equivariant Euclidean variables (EV), which interact via self-attention. (b) Computational efficiency of SO3krates allows the calculation of velocity-auto correlation functions from converged MD simulations for supra-molecular structures. (c)SO3krates enables to explore thousands of minima of the potential energy surface of small chainlike molecules such as Ac-Ala3-NHMe or DHA, where SO3krates can robustly extrapolate beyond the training data.
• Figure 2: SO(3) convolutions are constructed as triplet tensor products in the spherical harmonics basis, which is performed $F$ times along the feature dimension. We replace SO(3) convolutions by a parametrized filter function on the invariants (red blocks), which effectively reduces the tripled tensor product to taking the partial (per-degree) trace of a simple tensor product. Colored volumes correspond to the non-zero entries in the Clebsch-Gordan coefficients, which mask the tensor products.
• Figure 3: SO3krates architecture and building blocks. Taking the atomic types and positions as input they are embedded into invariant features $F$ and equivariant EV $X$ (methods section \ref{['sec:feature-and-sphc-initialization']}). They are then refined by $T$ Euclidean transformer blocks (ecTblock) (Eq. \ref{['eq:eqTblock']}) before the final invariant features are used to predict the potential energy (Eq. \ref{['eq:energy-pooling']}). After the Euclidean attention block, features and EV exchange per-atom information within the interaction block. Both blocks are enveloped by skip connection which allows to carry over information from prior layers. For an in detail description of the individual parts see methods section.
• Figure 4: (a) Number of frames per second (FPS) vs. the averaged stability coefficient (upper panel) and FPS vs. the averaged force MAE (lower panel) for four small organic molecules from the MD17 data set as reported in fu2022forces for different state-of-the-art MPNN architectures batzner20223schutt2018schnetschutt2021equivarianthu2021forcenetklicpera2021gemnetklicpera2020directionalliu2021sphericallu2019deeponet. (b) Since run times are sensitive to hardware and software implementation details, we re-implement two representative models along the trade-off lines under settings identical to the SO3krates MLFF (using jax), which yields framework-corrected FPS (dahshed vs. solid line). We observe speed-ups between 28 (for NequIP) and 15 (for SchNet) in our re-implementations. We find, that SO3krates enables reliable MD simulations and high accuracies without sacrificing computational performance. (c) MD step time vs. the number of atoms in the system. The smaller pre-factor in the computational complexity compared to SO(3) convolutions (Tab. \ref{['tab:theoretical-scaling']}) results in computational speed-ups that grow in system size. (d) Stability and speed of SO3krates enable nanosecond-long MD simulations for supra-molecular structures within a few hours. For the buckyball catcher, the ball stays in the catcher over the full simulation time of 20 ns, illustrating that the model successfully picks up on weak, non-covalent bonding.
• Figure 5: (a) Per-structure error distributions for an invariant and an equivariant SO3krates model with the same mean error on the test set. Spread and mean of the error distributions are given in SI Tab. \ref{['app:tab:error-distribution']}. (b) The MD stability observed at temperatures 300 K and 500 K. The transition to higher temperatures results in a drop of stability for the invariant model, hinting towards less robustness and weaker extrapolation behavior. Flexible molecules such as DHA pose a challenge for the invariant model at 300 K already.