Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Characterizing Machine Learning Force Fields as Emerging Molecular Dynamics Workloads on Graphics Processing Units

Udari De Alwis, Benjamin E. Mayer, Tom J. Ashby, Maria Barrera, Timon Evenblij, Joyjit Kundu

TL;DR

The analysis identifies key bottlenecks in descriptor and force computation, memory handling, highlighting the opportunities for improvements in the emerging area of MLFF based MD in drug discovery, that has received limited attention from a computer architecture perspective.

Abstract

Molecular dynamics (MD) simulates the time evolution of atomic systems governed by interatomic forces, and the fidelity of these simulations depends critically on the underlying force model. Classical force fields (CFFs) rely on fixed functional forms fitted to experimental or theoretical data, offering computational efficiency and broad applicability but limited accuracy in chemically diverse or reactive environments. In contrast, machine learning force fields (MLFFs) deliver near quantum chemical accuracy at molecular-mechanics cost by learning interatomic interactions directly from high level electronic structure data. While MLFFs offer improved accuracy at a fraction of the cost of quantum methods, they introduce significant computational overhead, particularly in descriptor evaluation and neural network inference. These operations pose challenges for parallel hardware due to irregular memory access, minimum data reuse and inefficient kernel execution. This work investigates the hardware performance of such models using poly alanine chains, a novel benchmark molecule system(s) with controllable input size, which used as performance evaluation test cases highlighting the computational bottlenecks of the graphical processor units when scaling out MLFF simulations. The analysis identifies key bottlenecks in descriptor and force computation, memory handling, highlighting the opportunities for improvements in the emerging area of MLFF based MD in drug discovery, that has received limited attention from a computer architecture perspective.

Characterizing Machine Learning Force Fields as Emerging Molecular Dynamics Workloads on Graphics Processing Units

TL;DR

The analysis identifies key bottlenecks in descriptor and force computation, memory handling, highlighting the opportunities for improvements in the emerging area of MLFF based MD in drug discovery, that has received limited attention from a computer architecture perspective.

Abstract

Molecular dynamics (MD) simulates the time evolution of atomic systems governed by interatomic forces, and the fidelity of these simulations depends critically on the underlying force model. Classical force fields (CFFs) rely on fixed functional forms fitted to experimental or theoretical data, offering computational efficiency and broad applicability but limited accuracy in chemically diverse or reactive environments. In contrast, machine learning force fields (MLFFs) deliver near quantum chemical accuracy at molecular-mechanics cost by learning interatomic interactions directly from high level electronic structure data. While MLFFs offer improved accuracy at a fraction of the cost of quantum methods, they introduce significant computational overhead, particularly in descriptor evaluation and neural network inference. These operations pose challenges for parallel hardware due to irregular memory access, minimum data reuse and inefficient kernel execution. This work investigates the hardware performance of such models using poly alanine chains, a novel benchmark molecule system(s) with controllable input size, which used as performance evaluation test cases highlighting the computational bottlenecks of the graphical processor units when scaling out MLFF simulations. The analysis identifies key bottlenecks in descriptor and force computation, memory handling, highlighting the opportunities for improvements in the emerging area of MLFF based MD in drug discovery, that has received limited attention from a computer architecture perspective.
Paper Structure (27 sections, 17 equations, 14 figures)

This paper contains 27 sections, 17 equations, 14 figures.

Table of Contents

  1. Introduction
  2. Background
  3. Molecular Dynamics in Drug Discovery and Machine Learning Force Fields (MLFFs)
  4. Molecular Systems for MLFF Simulations and GPU Profiling
  5. Setup of Simulation Workload
  6. GPU Profiler Setup and Data Acquisition
  7. MLFF Architecture Analysis
  8. ANI2x
  9. Atomic Environment Vector (AEV)
  10. Neural Network Potentials (NNP)
  11. Energy and Force Evaluation
  12. Ani2x Performance analysis
  13. Pytorch Function Analysis
  14. CUDA Kernel Analysis
  15. TorchMD-Net : Equivariant Transformer(ET) Architecture
  16. ...and 12 more sections

Figures (14)

  • Figure 1: Trajectory calculation in Molecular Dynamics and Potential Energy function approximates for force calculation step in the trajectory (considering a fixed cutoff system excluding long range interactions).
  • Figure 2: Visualization of benchmark system creation by incrementally extracting segments from a 1000-residue poly-alanine protein. Both the sequence and structure of the full-length protein are shown, with red highlights indicating cut points used to generate systems of varying sizes. Each resulting system is equilibrated in a water box.
  • Figure 3: ANI2x potential/force compute time comparison with AEV compute time with varying system size.
  • Figure 4: ANI2x AEV/potential/force memory transfers with varying neighborhood sizes.
  • Figure 5: Combined Operator Timing: AEV vs Energy vs Force
  • ...and 9 more figures