Learning Smooth and Expressive Interatomic Potentials for Physical Property Prediction

Learning Smooth and Expressive Interatomic Potentials for Physical Property Prediction

Speakers: Xiang Fu, Brandon Wood, Luis Barroso-Luque, Daniel S. Levine, Meng Gao, Misko Dzamba, Larry Zitnick (Brandon Wood and Meng Gao presented)

Conference: ICML 2025

YouTube: https://slideslive.com/39044053

Overview

This article delves into a pivotal talk presented at ICML 2025 by Brandon Wood and Meng Gao from Meta's Fair Chemistry team, highlighting groundbreaking work led by Xiang Fu. The presentation, titled "Learning Smooth and Expressive Interatomic Potentials for Physical Property Prediction," addresses a fundamental challenge in computational materials science: accelerating the discovery and design of novel materials using artificial intelligence. Specifically, it focuses on developing highly accurate and stable AI surrogates for computationally intensive quantum chemistry methods like Density Functional Theory (DFT).

The core of the talk revolves around the critical importance of a smooth and continuous energy landscape for machine learning models that predict material properties. While many AI models can achieve high accuracy in predicting static energies and forces, they often fail to deliver stable molecular dynamics (MD) simulations or accurate predictions of dynamic and vibrational properties. This talk introduces a novel architectural approach, dubbed ESAN, designed to explicitly enforce smoothness in the learned energy function, thereby enabling robust physical property prediction and reliable simulations. This research is crucial for advancing applications in diverse fields, including renewable energy storage, direct air carbon capture, and the development of advanced display materials for technologies like AR glasses.

Background

▶ Watch: Introduction: Accelerating DFT with GNNs (0:00)

The quest for new materials is a cornerstone of scientific and technological progress, yet the vastness of the materials design space presents an intractable challenge for experimental exploration alone. To navigate this immense space, computational methods are indispensable for identifying promising candidates. Among these, Density Functional Theory (DFT) stands out as a widely used quantum chemistry technique that models atoms and electrons quantum mechanically, providing highly accurate insights into material properties. However, DFT calculations are notoriously slow, often requiring hours or even days for a single system, rendering high-throughput screening of millions of candidates impractical.

To circumvent the computational bottleneck of DFT, the machine learning community has increasingly turned to AI surrogates, particularly graph neural networks (GNNs). GNNs offer a natural framework for representing atomic systems, where atoms are nodes and interatomic connections are edges. These models are trained on DFT-generated data to rapidly predict energies and per-atom forces, aiming to reproduce DFT accuracy at a fraction of the computational cost. The promise is to accelerate materials discovery from days to seconds.

However, a significant problem arises when these GNN surrogates are deployed for downstream physical property prediction tasks, such as running molecular dynamics (MD) simulations or predicting vibrational properties like phonons. These tasks critically depend on the derivatives of the energy function. For instance, atomic forces are the negative gradient of the energy with respect to atomic positions, and vibrational modes involve second-order derivatives. While a GNN might achieve excellent accuracy in predicting raw energies and forces, this does not guarantee the underlying energy landscape it has learned is sufficiently smooth and continuous.

The consequence of a non-smooth energy landscape is severe: MD simulations can become unstable, leading to unphysical trajectories where molecules "blow up." Furthermore, predictions of properties reliant on higher-order derivatives become highly inaccurate. The speakers illustrate this with examples where models with low energy and force errors still produce poor physical property predictions. This disconnect stems from fundamental issues in the learned energy landscape, which can exhibit unphysical sharpness, discontinuities, or signs of overfitting, even if the direct energy and force predictions appear good. The challenge, therefore, lies not just in predicting values, but in learning an energy function that faithfully represents the physics, ensuring energy conservation and stability across dynamic processes.

Key Findings

▶ Watch: Problem: Energy/force accuracy ≠ physical property accuracy (3:45)

The central hypothesis presented in the talk is that a smooth energy landscape is not merely a desirable feature, but a necessary condition for accurate physical property prediction tasks and stable molecular dynamics simulations. This hypothesis underpins the entire research effort, leading to several key findings and architectural innovations.

Firstly, the team established a quantifiable metric for assessing the smoothness of a learned energy function: energy drift (E_drift) in simulations. Drawing from a theorem, they demonstrated that the difference in energy (E1 - E0) over a simulated time interval is bounded by a function of the energy function's degree of smoothness. Consequently, a lower energy drift directly correlates with a smoother energy landscape, providing an observable criterion for guiding model design.

This metric allowed for systematic evaluation of architectural choices. A crucial finding highlighted the stark difference between directly predicting forces and deriving them as gradients of a predicted energy function. Models that predict forces separately from energy (direct force prediction) fail to guarantee energy conservation, leading to substantial energy drift. In contrast, deriving forces via autograd from a single, predicted energy value ensures that forces are conservative, resulting in a dramatic reduction in energy drift—by an astonishing seven orders of magnitude over a 100-picosecond simulation. This established gradient force derivation as a necessary, though not entirely sufficient, condition for achieving the desired smoothness.

Another critical design choice identified was the implementation of cutoff functions. Traditional GNNs often employ radial step functions to define the interaction radius between atoms. However, these step functions introduce discontinuities in the energy landscape precisely when an atom crosses the cutoff boundary. The solution proposed is to replace these step functions with smooth envelope functions that cause the energy to decay smoothly to zero, thereby eliminating these unphysical discontinuities and further minimizing energy drift.

These insights culminated in the development of the ESAN architecture, a novel model built upon a previous design (ESEN) that rigorously incorporates these principles of smoothness. The empirical results demonstrate that ESAN, by design, produces energy landscapes that conserve energy, leading to a significantly improved correlation between the model's energy error and various physical property metrics.

Perhaps the most compelling result is ESAN's performance on phonon prediction. Phonons, representing the vibrational modes of atoms in a material, require accurate second-order derivatives of the energy function and thus serve as an excellent proxy for the smoothness of the learned potential. While previous "best" models, despite good energy/force accuracy, completely failed to predict phonon spectra, ESAN achieved "spot on" predictions, accurately reproducing the complex two-dimensional phonon landscape. Furthermore, ESAN demonstrated "very well" performance on the MatBench Discovery benchmark, a widely accepted standard for machine learning force fields in materials science, solidifying its position as a highly accurate and robust model. The speakers also teased that the ESAN architecture serves as the foundation for their subsequent work, the Universal Model for Atoms, indicating its foundational importance.

Technical Deep Dive

▶ Watch: Measuring smoothness: Quantifying energy drift (6:10)

The technical foundation of this work lies in leveraging graph neural networks (GNNs) to model interatomic potentials. In this context, an atomic system is naturally represented as a graph where individual atoms are nodes and their interactions or proximity define edges. The GNN takes 3D atomic positions and atom types as input and is trained to predict total energy and per-atom forces, mimicking the outputs of Density Functional Theory (DFT).

The core technical challenge addressed is the quality of the "energy landscape" learned by the GNN. The energy landscape is a high-dimensional surface that maps atomic configurations to their potential energy. For accurate and stable simulations, this landscape must be smoothly varying and continuous. Unphysical features like sharp changes, discontinuities, or local minima due to overfitting can lead to erroneous force calculations and unstable dynamics.

A crucial concept from mechanics is that forces are conservative and represent the negative gradient of the potential energy with respect to atomic positions. This means the work done by these forces is independent of the path taken, which is fundamental to the conservation of energy. When a GNN directly predicts both energy and forces as separate outputs, there is no inherent guarantee that these predictions will be consistent with the conservative force principle. This inconsistency manifests as energy drift in simulations.

The solution implemented in ESAN is to predict only the total energy of the system using the GNN. The per-atom forces are then derived analytically by taking the negative gradient of this predicted energy with respect to the atomic positions, using automatic differentiation (autograd). This design choice fundamentally ensures that the forces are conservative by construction, directly linking them to the energy function and drastically reducing energy drift. The empirical evidence for this is compelling: a seven-order-of-magnitude reduction in energy drift over 100 picoseconds compared to models that predict forces directly.

Beyond the gradient-derived forces, another critical aspect of GNN design for atomic systems is how interatomic interactions are handled, particularly the concept of a cutoff radius. GNNs typically employ message passing where atoms only interact with neighbors within a certain cutoff distance. Traditionally, this cutoff is implemented using a radial step function: interactions are considered up to the cutoff, and abruptly cease beyond it. This abrupt change, however, introduces a discontinuity in the energy landscape at the cutoff boundary. As an atom moves across this boundary, its interactions with other atoms suddenly appear or disappear, leading to unphysical jerks in the forces and potential energy.

To address this, ESAN incorporates smooth envelope functions (also known as switching functions) instead of step functions for the cutoff. These envelope functions ensure that the interatomic potential and its derivatives smoothly decay to zero as the distance between atoms approaches the cutoff radius. This eliminates the discontinuities, contributing further to the overall smoothness of the learned energy landscape. The talk emphasizes that this is one of "many other examples of this kind of design choices" that, when ablated and tested for energy drift, consistently demonstrate the benefit of smoothness.

While the talk mentions that ESAN is based on a previous design called ESEN and that full architectural details are available at the poster session, it does not delve into the specific layers, message-passing schemes, or embedding techniques unique to ESAN. However, it is clear that ESAN's innovation lies in its principled integration of these smoothness-enforcing mechanisms—gradient-derived forces and smooth cutoff functions—into its architecture, rather than introducing entirely new GNN paradigms. The ability to accurately predict phonons, which are vibrational modes requiring accurate second-order derivatives of the energy, serves as a testament to the high degree of smoothness and accuracy achieved by ESAN.

Experimental Setup & Results

▶ Watch: Design Solution 1: Gradient forces for stability (7:45)

The experimental validation of the ESAN architecture and the underlying smoothness hypothesis relied on a systematic approach centered around quantifying energy drift and evaluating performance on key physical property prediction tasks.

Quantifying Smoothness via Energy Drift:

The primary metric used to assess the smoothness of the learned energy landscape was the energy drift (E_drift) observed in short molecular dynamics (MD) simulations. The principle is that for a sufficiently smooth potential, the total energy of an isolated system should remain constant over time. Deviations from this constancy, or drift, indicate non-conservative forces or discontinuities in the potential. The speakers highlighted a theoretical basis stating that this energy drift is bounded by the degree of smoothness (Nth derivative existence and continuity).

• Force Derivation Comparison: A critical experiment compared models using direct force prediction versus gradient force prediction (forces derived via autograd from the energy). The results showed an astounding difference: gradient forces led to a seven orders of magnitude reduction in energy drift over 100 picoseconds of simulation time. This empirically validated the necessity of deriving forces as gradients for energy conservation.
• Cutoff Function Comparison: Further experiments investigated the impact of cutoff functions. Replacing the traditional discontinuous radial step function with a smooth envelope function was shown to significantly reduce energy drift by eliminating unphysical discontinuities in the energy landscape. This demonstrated the importance of careful design choices beyond just gradient forces.

Physical Property Prediction Results:

The ultimate test for ESAN was its ability to accurately predict complex physical properties.

• Correlation with Physical Properties: A key visualization presented showed a plot where models were categorized by whether they conserved energy (smooth landscapes, represented by diamonds) or not (non-smooth landscapes, represented by X's). The models that conserved energy consistently showed a much better correlation between their energy prediction error and the actual physical property metric of interest. ESAN, falling into the "diamond" category, demonstrated this improved correlation, reinforcing the hypothesis that smoothness is crucial for downstream task performance.
• Phonon Prediction: Phonons, which are vibrational modes in materials, are highly sensitive to the second-order derivatives of the energy function. This makes them an excellent and stringent benchmark for the smoothness and accuracy of an interatomic potential. The talk presented a striking comparison:
• A previous "best" model, despite good energy and force prediction accuracy, yielded "pretty much completely wrong" phonon predictions.
• In contrast, the ESAN model predicted the phonon spectrum "spot on," accurately capturing the complex two-dimensional landscape of vibrational frequencies. This result is a strong testament to ESAN's ability to learn a physically meaningful and smooth energy landscape capable of yielding accurate higher-order derivatives.
• MatBench Discovery Benchmark: ESAN was also evaluated on the MatBench Discovery benchmark, a widely accepted standard for machine learning force fields in materials science. The model performed "very well" on this benchmark, indicating its competitive and robust performance across a range of materials science tasks. While specific numerical scores were not provided, the qualitative assessment suggests state-of-the-art capabilities.

Limitations in Experimental Detail:

The transcript did not specify the exact datasets used for training and evaluation (beyond implying DFT-generated data for various materials), nor did it detail the specific hardware (e.g., GPU types, number of GPUs) or software frameworks (e.g., PyTorch, TensorFlow versions) employed for model training and simulation. Similarly, the exact number of parameters in ESAN or precise throughput numbers were not explicitly stated, focusing more on the qualitative improvements and the underlying principles.

Practical Implications

▶ Watch: Introducing the ESAN architecture design (10:00)

The development and validation of the ESAN architecture carry significant practical implications for researchers and engineers involved in materials science, chemistry, and AI/ML infrastructure.

Firstly, for practitioners and model builders in computational materials science, ESAN provides a robust and reliable tool for accelerating materials discovery. By offering an AI surrogate that not only matches DFT accuracy but also ensures the stability and physical consistency required for dynamic simulations, ESAN drastically reduces the computational cost of screening novel materials. This means faster iteration cycles in designing materials for critical applications such as renewable energy storage, direct air carbon capture, and advanced display technologies like those used in AR glasses. The ability to accurately predict properties like phonons, which are crucial for understanding thermal and mechanical behavior, opens doors for more precise material design.

Secondly, for infra teams and deployers of ML models, the emphasis on a smooth energy landscape directly translates to stable and physically meaningful molecular dynamics (MD) simulations. Previous GNN-based force fields often suffered from instability, leading to unphysical trajectories or "blowing up" molecules, which rendered them unreliable for long-duration simulations. ESAN's design choices, particularly gradient-derived forces and smooth cutoff functions, mitigate these issues, enabling the deployment of AI models for predictive MD simulations that conserve energy and yield trustworthy results. This reduces the need for constant monitoring and intervention, improving the efficiency of computational workflows.

However, it is crucial to acknowledge tradeoffs and limitations. As noted by the speakers during the Q&A, there is an inherent balance between smoothness and accuracy. An "infinitely smooth" model might not be sufficiently accurate in capturing the subtle energy variations crucial for specific chemical or physical phenomena. Therefore, the design of ESAN represents an optimized balance rather than an extreme. Furthermore, while gradient forces are necessary for smoothness, they are not sufficient; other architectural details, like the choice of cutoff functions and potentially implicit constraints (e.g., forces summing to zero), also play a vital role. The talk's discussion on whether to explicitly enforce constraints (like total forces summing to zero) suggests an ongoing area of research where further improvements could be made.

The generalizability of ESAN across the vast chemical space is also an important consideration. While performance on MatBench Discovery and phonon prediction for specific systems is strong, its applicability to entirely novel material classes or extreme conditions would require further rigorous testing. Finally, while ESAN significantly accelerates calculations compared to DFT, the computational cost associated with training and inferring large GNNs still requires substantial computational resources, albeit far less than ab initio methods. The work demonstrates a clear path forward for building more reliable and physically informed AI models for complex scientific tasks, setting a new standard for interatomic potentials.

Key Takeaways

• Smoothness is Critical: A smoothly varying and continuous energy landscape is a necessary condition for accurate physical property prediction and stable molecular dynamics simulations, not just a desirable feature.
• Optimizing for Error is Insufficient: Minimizing only energy and force prediction errors does not guarantee a smooth energy landscape or reliable downstream physical property predictions.
• Gradient-Derived Forces are Essential: Deriving forces as the negative gradient of a single predicted energy value (via autograd) is crucial for ensuring conservative forces and significantly reduces energy drift in simulations by multiple orders of magnitude.
• Architectural Design Matters: Careful architectural choices, such as replacing discontinuous radial step functions with smooth envelope functions for interatomic cutoffs, are vital to eliminate unphysical discontinuities in the energy landscape.
• ESAN Achieves Both Smoothness and Accuracy: The ESAN architecture successfully integrates these principles, leading to models that are both smooth and accurate, enabling "spot on" predictions for complex properties like phonons and strong performance on benchmarks like MatBench Discovery.
• Enabling Accelerated Materials Discovery: This approach provides robust and stable AI surrogates for DFT, significantly accelerating the computational discovery and design of novel materials for a wide range of applications.

About the Speaker(s)

The talk was presented by Brandon Wood and Meng Gao (referred to as Ray Gao in the presentation), both members of the Fair Chemistry team at Meta. The primary lead author of the presented work was Xiang Fu, who unfortunately could not be present at the conference. The broader team involved in this research also includes Luis Barroso-Luque, Daniel S. Levine, Misko Dzamba, and Larry Zitnick, all presumably from Meta. Their collective expertise lies in applying AI and machine learning techniques to address challenges in chemistry and materials science. The Fair Chemistry team at Meta is focused on leveraging AI for the design and discovery of materials for diverse applications, including renewable energy storage, direct air carbon capture technologies, and the development of advanced display materials, such as those utilized in augmented reality (AR) glasses.