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

• URL: http://arxiv.org/abs/2303.08169v1
• Date: Tue, 14 Mar 2023 18:36:44 GMT
• Title: Allegro-Legato: Scalable, Fast, and Robust Neural-Network Quantum Molecular Dynamics via Sharpness-Aware Minimization
• Authors: Hikaru Ibayashi, Taufeq Mohammed Razakh, Liqiu Yang, Thomas Linker, Marco Olguin, Shinnosuke Hattori, Ye Luo, Rajiv K. Kalia, Aiichiro Nakano, Ken-ichi Nomura, and Priya Vashishta
• Abstract summary: 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. 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. We
• Score: 1.8431330466822737
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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_\textrm{failure}$, without sacrificing computational speed or accuracy. Specifically, Allegro-Legato exhibits much weaker dependence of timei-to-failure on the problem size, $t_{\textrm{failure}} \propto N^{-0.14}$ ($N$ is the number of atoms) compared to the SOTA Allegro model $\left(t_{\textrm{failure}} \propto N^{-0.29}\right)$, 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.

Related papers

• Scalable Mechanistic Neural Networks [52.28945097811129]
  We propose an enhanced neural network framework designed for scientific machine learning applications involving long temporal sequences. By reformulating the original Mechanistic Neural Network (MNN) we reduce the computational time and space complexities from cubic and quadratic with respect to the sequence length, respectively, to linear. Extensive experiments demonstrate that S-MNN matches the original MNN in precision while substantially reducing computational resources.
  arXiv  Detail & Related papers  (2024-10-08T14:27:28Z)
• NeuralStagger: Accelerating Physics-constrained Neural PDE Solver with Spatial-temporal Decomposition [67.46012350241969]
  This paper proposes a general acceleration methodology called NeuralStagger. It decomposing the original learning tasks into several coarser-resolution subtasks. We demonstrate the successful application of NeuralStagger on 2D and 3D fluid dynamics simulations.
  arXiv  Detail & Related papers  (2023-02-20T19:36:52Z)
• Towards Neural Variational Monte Carlo That Scales Linearly with System Size [67.09349921751341]
  Quantum many-body problems are central to demystifying some exotic quantum phenomena, e.g., high-temperature superconductors. The combination of neural networks (NN) for representing quantum states, and the Variational Monte Carlo (VMC) algorithm, has been shown to be a promising method for solving such problems. We propose a NN architecture called Vector-Quantized Neural Quantum States (VQ-NQS) that utilizes vector-quantization techniques to leverage redundancies in the local-energy calculations of the VMC algorithm.
  arXiv  Detail & Related papers  (2022-12-21T19:00:04Z)
• MemSE: Fast MSE Prediction for Noisy Memristor-Based DNN Accelerators [5.553959304125023]
  We theoretically analyze the mean squared error of DNNs that use memristors to compute matrix-vector multiplications (MVM) We take into account both the quantization noise, due to the necessity of reducing the DNN model size, and the programming noise, stemming from the variability during the programming of the memristance value. The proposed method is almost two order of magnitude faster than Monte-Carlo simulation, thus making it possible to optimize the implementation parameters to achieve minimal error for a given power constraint.
  arXiv  Detail & Related papers  (2022-05-03T18:10:43Z)
• Interfacing Finite Elements with Deep Neural Operators for Fast Multiscale Modeling of Mechanics Problems [4.280301926296439]
  In this work, we explore the idea of multiscale modeling with machine learning and employ DeepONet, a neural operator, as an efficient surrogate of the expensive solver. DeepONet is trained offline using data acquired from the fine solver for learning the underlying and possibly unknown fine-scale dynamics. We present various benchmarks to assess accuracy and speedup, and in particular we develop a coupling algorithm for a time-dependent problem.
  arXiv  Detail & Related papers  (2022-02-25T20:46:08Z)
• Quantum dynamics simulations beyond the coherence time on NISQ hardware by variational Trotter compression [0.0]
  We demonstrate a post-quench dynamics simulation of a Heisenberg model on present-day IBM quantum hardware. We show how to measure the required cost function, the overlap between the time-evolved and variational states, on present-day hardware. In addition to carrying out simulations on real hardware, we investigate the performance and scaling behavior of the algorithm with noiseless and noisy classical simulations.
  arXiv  Detail & Related papers  (2021-12-23T15:44:47Z)
• Learning Large-Time-Step Molecular Dynamics with Graph Neural Networks [14.388196138756195]
  We introduce a graph neural network (GNN) based model, MDNet, to predict the evolution of coordinates and momentum with large time steps. We demonstrate the performance of MDNet on a 4000-atom system with large time steps, and show that MDNet can predict good equilibrium and transport properties.
  arXiv  Detail & Related papers  (2021-11-30T07:32:39Z)
• Fixed Depth Hamiltonian Simulation via Cartan Decomposition [59.20417091220753]
  We present a constructive algorithm for generating quantum circuits with time-independent depth. We highlight our algorithm for special classes of models, including Anderson localization in one dimensional transverse field XY model. In addition to providing exact circuits for a broad set of spin and fermionic models, our algorithm provides broad analytic and numerical insight into optimal Hamiltonian simulations.
  arXiv  Detail & Related papers  (2021-04-01T19:06:00Z)
• ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations [86.41674945012369]
  We develop a scalable and expressive Graph Neural Networks model, ForceNet, to approximate atomic forces. Our proposed ForceNet is able to predict atomic forces more accurately than state-of-the-art physics-based GNNs.
  arXiv  Detail & Related papers  (2021-03-02T03:09:06Z)
• Fast and differentiable simulation of driven quantum systems [58.720142291102135]
  We introduce a semi-analytic method based on the Dyson expansion that allows us to time-evolve driven quantum systems much faster than standard numerical methods. We show results of the optimization of a two-qubit gate using transmon qubits in the circuit QED architecture.
  arXiv  Detail & Related papers  (2020-12-16T21:43:38Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.