# Allegro-Legato: Scalable, Fast, and Robust Neural-Network Quantum Molecular Dynamics via Sharpness-Aware Minimization

arXiv: https://arxiv.org/abs/2303.08169 · HTML: https://mlsanity.com/paper/2303.08169/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 9 of its 26 references that are in the map.
- First posted on arXiv 14 Mar 2023.

## Builds on (most cited first, 9 of 9)

- [Explaining and Harnessing Adversarial Examples](https://mlsanity.com/paper/1412.6572/) (2014)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials](https://mlsanity.com/paper/2101.03164/) (2021)
- [Adversarial Training for Free!](https://mlsanity.com/paper/1904.12843/) (2019)
- [Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics](https://mlsanity.com/paper/2204.05249/) (2022)
- [Descending through a Crowded Valley - Benchmarking Deep Learning Optimizers](https://mlsanity.com/paper/2007.01547/) (2020)
- [Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations](https://mlsanity.com/paper/2210.07237/) (2022)
- [Differentiable sampling of molecular geometries with uncertainty-based adversarial attacks](https://mlsanity.com/paper/2101.11588/) (2021)

## Built on by (most cited first, 3 of 3)

- [Uncertainty-biased molecular dynamics for learning uniformly accurate interatomic potentials](https://mlsanity.com/paper/2312.01416/) (2023)
- [Learning Interatomic Potentials at Multiple Scales](https://mlsanity.com/paper/2310.13756/) (2023)
- [nabla2DFT: A Universal Quantum Chemistry Dataset of Drug-Like Molecules and a Benchmark for Neural Network Potentials](https://mlsanity.com/paper/2406.14347/) (2024)

## Abstract (arXiv)

Neural-network quantum molecular dynamics (NNQMD) simulations based on machine learning are revolutionizing atomistic simulations of materials by providing quantum-mechanical accuracy but orders-of-magnitude faster, illustrated by ACM Gordon Bell prize (2020) and finalist (2021). State-of-the-art (SOTA) NNQMD model founded on group theory featuring rotational equivariance and local descriptors has provided much higher accuracy and speed than those models, thus named Allegro (meaning fast). On massively parallel supercomputers, however, it suffers a fidelity-scaling problem, where growing number of unphysical predictions of interatomic forces prohibits simulations involving larger numbers of atoms for longer times. Here, we solve this problem by combining the Allegro model with sharpness aware minimization (SAM) for enhancing the robustness of model through improved smoothness of the loss landscape. The resulting Allegro-Legato (meaning fast and "smooth") model was shown to elongate the time-to-failure t_textrmfailure, without sacrificing computational speed or accuracy. Specifically, Allegro-Legato exhibits much weaker dependence of timei-to-failure on the problem size, t_textrmfailure propto N-0.14 (N is the number of atoms) compared to the SOTA Allegro model left(t_textrmfailure propto N-0.29right), i.e., systematically delayed time-to-failure, thus allowing much larger and longer NNQMD simulations without failure. The model also exhibits excellent computational scalability and GPU acceleration on the Polaris supercomputer at Argonne Leadership Computing Facility. Such scalable, accurate, fast and robust NNQMD models will likely find broad applications in NNQMD simulations on emerging exaflop/s computers, with a specific example of accounting for nuclear quantum effects in the dynamics of ammonia.