[Paper Review] Phonon predictions with E(3)-equivariant graph neural networks

Shiang Fang, Mario Geiger+2 more|arXiv (Cornell University)|Mar 17, 2024

Thermal properties of materialsMaterials Science3 citations

TL;DR

This paper introduces Phonax, an E(3)-equivariant graph neural network framework that predicts phonon dispersion and density of states in crystals and molecules by computing second derivatives (Hessians) of a learned energy model. The method enables accurate, symmetry-aware phonon predictions and improves energy model generalization by incorporating Hessian data, directly linking models to experimental vibrational observables.

ABSTRACT

We present an equivariant neural network for predicting vibrational and phonon modes of molecules and periodic crystals, respectively. These predictions are made by evaluating the second derivative Hessian matrices of the learned energy model that is trained with the energy and force data. Using this method, we are able to efficiently predict phonon dispersion and the density of states for inorganic crystal materials. For molecules, we also derive the symmetry constraints for IR/Raman active modes by analyzing the phonon mode irreducible representations. Additionally, we demonstrate that using Hessian as a new type of higher-order training data improves energy models beyond models that only use lower-order energy and force data. With this second derivative approach, one can directly relate the energy models to the experimental observations for the vibrational properties. This approach further connects to a broader class of physical observables with a generalized energy model that includes external fields.

Motivation & Objective

• To develop an efficient, symmetry-aware method for predicting phonon dispersion and density of states in periodic crystals and molecules.
• To enable direct connection between machine learning energy models and experimental vibrational observables through second derivative (Hessian) computations.
• To improve the generalization and accuracy of interatomic potential models by incorporating Hessian data as higher-order training signals.
• To extend the framework to include generalized external fields (e.g., electric fields) for predicting Born effective charges and dielectric responses.
• To provide a unified, differentiable framework for predicting diverse physical observables from a single energy model.

Proposed method

• The method employs E(3)-equivariant graph neural networks (GNNs) to learn a differentiable energy model from atomic positions and forces.
• Phonon properties are derived by computing the Hessian matrix (∂²U/∂xi∂xj) of the learned energy function using automatic differentiation in JAX.
• The Hessian is evaluated at equilibrium geometries to compute phonon band structures and density of states for crystals and molecules.
• Symmetry constraints for IR/Raman activity are derived from the irreducible representations of phonon modes.
• A generalized energy model including external field couplings (e.g., electric fields) is trained up to quadratic order, enabling prediction of Born effective charges and dielectric constants.
• The framework uses a modified NequIP architecture with higher-order derivatives for improved physical consistency and model performance.

Experimental results

Research questions

• RQ1 Can E(3)-equivariant GNNs be used to predict phonon dispersion and density of states with high accuracy and efficiency?
• RQ2 Does incorporating Hessian data into the training of energy models improve their generalization and physical consistency?
• RQ3 Can the framework predict symmetry-resolved vibrational modes (e.g., IR/Raman active) directly from the Hessian and group theory?
• RQ4 How does including higher-order derivatives (e.g., from electric field coupling) enhance the predictive power of the energy model?
• RQ5 Can the same energy model predict diverse physical observables, such as phonons, Born effective charges, and dielectric constants, within a unified differentiable framework?

Key findings

• The Phonax framework enables accurate and efficient prediction of phonon dispersion and density of states in inorganic crystals and molecules using Hessian computations from E(3)-equivariant GNNs.
• Incorporating Hessian data into the training process significantly improves the generalization and accuracy of the underlying energy model beyond models trained only on energy and force data.
• The method correctly predicts symmetry-allowed IR and Raman active modes by analyzing the irreducible representations of phonon modes.
• The framework successfully predicts Born effective charges and dielectric constants by computing second derivatives with respect to external electric fields using automatic differentiation.
• The generalized energy model, including up to quadratic coupling to external fields, improves predictions of energy and forces, especially at higher model orders.
• The approach provides a direct, differentiable link between machine learning potentials and experimentally measurable vibrational properties, enhancing physical interpretability and utility.