Learning local equivariant representations for large-scale atomistic dynamics

Abstract

A simultaneously accurate and computationally efficient parametrization of the potential energy surface of molecules and materials is a long-standing goal in the natural sciences. While atom-centered message passing neural networks (MPNNs) have shown remarkable accuracy, their information propagation has limited the accessible length-scales. Local methods, conversely, scale to large simulations but have suffered from inferior accuracy. This work introduces Allegro, a strictly local equivariant deep neural network interatomic potential architecture that simultaneously exhibits excellent accuracy and scalability. Allegro represents a many-body potential using iterated tensor products of learned equivariant representations without atom-centered message passing. Allegro obtains improvements over state-of-the-art methods on QM9 and revMD17. A single tensor product layer outperforms existing deep MPNNs and transformers on QM9. Furthermore, Allegro displays remarkable generalization to out-of-distribution data. Molecular simulations using Allegro recover structural and kinetic properties of an amorphous electrolyte in excellent agreement with ab-initio simulations. Finally, we demonstrate parallelization with a simulation of 100 million atoms.

Drug Design in the Exascale Era: A Perspective from Massively Parallel QM/MM Simulations

“…44,[124][125][126][127] These ML models work natively on GPUs, and because they normally rely on local interactions alone, they can be exceptionally scalable on distributed architectures. 72,73 Furthermore, their training requires datasets generated through many single-point QM(/MM) calculations that are expensive but embarrassingly parallelizable. Finally, the recent introduction of ML-accelerated perturbative techniques provides an efficient and highly parallelizable way of recovering the accuracy of QM/MM potentials from simulations using cheaper methods (such as force fields or even ML/MM models) at the cost of only a few single-point energy and force QM/MM calculations.…”

“…Machine learning methods have already shown to make excellent utilization of GPU resources, and could be excellent candidates to push DFT QM/MM MD into the exascale regime. 72,73 Here, after summarizing some salient aspects of the MiMiC QM/MM interface and demonstrating its scalability, we present applications of the code to systems of pharmacological relevance, from enzymatic reactions for the prediction of the transition state, to inhibitor-enzyme binding towards the investigation of k off values. We close by giving our perspective about QM/MM MD simulations for drug design in the exascale era.…”

Representations of Materials for Machine Learning

“…Characterizing systems based off local environments allows models to extrapolate to cases where global representations may vary substantially (e.g. an extended supercell of a crystal structure) [14] and enables highly-scalable methods of computation that can extend the practical limit of simulations to much larger systems [159]. However, Unke et al notes that the required complexity of the representation can grow quickly when modeling systems with many distinct elements and the quality of ML predictions will be sensitive to the selected hyperparameters, such as the characteristics distances and angles in atom-centered symmetry functions [160].…”

EquiFold: Protein Structure Prediction with a Novel Coarse-Grained Structure Representation

“…SE(3)-equivariant neural networks have been used in various 3D object modeling tasks, including atomic potential prediction [22, 23, 24], molecular property prediction [25, 26, 27, 28, 29], protein structure prediction, [2] and docking [30, 31]. They incorporate the symmetry of 3D space and are substantially more data efficient than their non-symmetry-aware counterparts [32, 33]. Input to SE(3)-equivariant networks consists of geometric tensors, or the irreducible representations of the SO(3) group (the group of rotations in ), of various degrees l , such as scalars ( l = 0), vectors ( l = 1), and higher degree ( l ≥ 2) tensors that transform under a rotation R through multiplication with corresponding Wigner D-matrices D l ( R ).…”