Learning the energy relaxation manifold from unrelaxed structures with RelaxNet

Emily Lin, Evelyn N. Wang, Sili Deng

Keywords: neural ODEs, energy minimization, trajectory, relaxation, forcefield, optimization

TL;DR: We developed RelaxNet, a dynamics-aware, equivariant model that leverages neural ordinary differential equations and message passing neural networks for predicting energy relaxation landscapes between initial unrelaxed and final relaxed structures.

Abstract: In an effort to bypass computationally expensive density functional theory (DFT) calculations for energy minimization and structure relaxation, rapid progress in the development of machine learning force fields/interatomic potentials (MLFF/MLIPs) and more robust models that adhere to quantum chemistry/physical paradigms and constraints have been realized. However, most research to date involves static-frame energy predictions only (i.e., given a specific atomic configuration, predict the energy of the current or final instance), neglecting intermediary physical insight-providing contexts. In this work, we developed RelaxNet, the first end-to-end, dynamics-aware, equivariant model that leverages neural ordinary differential equations (ODEs) and message passing neural networks (MPNNs) for predicting the energy relaxation landscape between the initial unrelaxed structure and final relaxed structure. From just the initial structure, which is often the configuration that is fed into DFT simulations, we can accurately recover the energy, forces, and geometric pathways for the trajectory at a competitive prediction accuracy, as evidenced by comprehensive benchmarking with state-of-the-art static models and MLIP-based relaxation methods. Additionally, we provide extensive insights on the use of implicit vs. explicit latent embedding evolution to offer perspectives on optimal strategies for future works that seek to integrate expensive graph-based neural networks and neural ODEs.

Primary Area: applications to physical sciences (physics, chemistry, biology, etc.)

Submission Number: 21302