Scaling Neural-Network-Based Molecular Dynamics with Long-Range Electrostatic Interactions to 51 Nanoseconds per Day

2025-04-22

AI Summary

Neural network molecular dynamics (NNMD) suffers from severe performance bottlenecks in systems with long-range electrostatic interactions, primarily due to computational overhead from neural network inference and Ewald summation. To address this, we propose a co-optimization framework tailored for exascale supercomputing platforms. Our approach introduces: (i) hardware-accelerated FFT offloading for efficient long-range force computation; (ii) fine-grained intra-core overlap between neural network inference and long-range force evaluation; (iii) a ring-based atom-level load-balancing scheme to minimize inter-node communication; and (iv) deep integration with the DPLR framework and Fugaku’s heterogeneous architecture. Evaluated on the Fugaku supercomputer, our method achieves up to 37× speedup over baseline NNMD implementations, attaining simulation throughput of 51 ns/day—surpassing prior state-of-the-art for NNMD-based long-range electrostatics. This work establishes a scalable paradigm for high-fidelity, long-timescale simulations of large-scale ionic and polar systems.

Technology Category

• Machine Learning: Hardware-aware ML
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems
• Search and Optimization: Distributed Search

Application Category

• Economics, Online Markets and Human Computation: Incentives in network design for Web infrastructures and ecosystems
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization

Abstract

Neural network-based molecular dynamics (NNMD) simulations incorporating long-range electrostatic interactions have significantly extended the applicability to heterogeneous and ionic systems, enabling effective modeling critical physical phenomena such as protein folding and dipolar surface and maintaining ab initio accuracy. However, neural network inference and long-range force computation remain the major bottlenecks, severely limiting simulation speed. In this paper, we target DPLR, a state-of-the-art NNMD package that supports long-range electrostatics, and propose a set of comprehensive optimizations to enhance computational efficiency. We introduce (1) a hardware-offloaded FFT method to reduce the communication overhead; (2) an overlapping strategy that hides long-range force computations using a single core per node, and (3) a ring-based load balancing method that enables atom-level task evenly redistribution with minimal communication overhead. Experimental results on the Fugaku supercomputer show that our work achieves a 37x performance improvement, reaching a maximum simulation speed of 51 ns/day.

Problem

• Optimizing NNMD for faster long-range electrostatic simulations
• Reducing bottlenecks in neural network inference and force computation
• Enhancing DPLR package efficiency for large-scale molecular dynamics

Innovation

• Hardware-offloaded FFT reduces communication overhead
• Overlapping strategy hides long-range force computations
• Ring-based load balancing enables atom-level redistribution