Equivariant attention for thermal transport simulations: Proof of concept
Title: Equivariant attention for thermal transport simulations: Proof of concept
DNr: Berzelius-2023-28
Project Type: LiU Berzelius
Principal Investigator: Florian Knoop <florian.knoop@liu.se>
Affiliation: Linköpings universitet
Duration: 2023-02-25 – 2023-09-01
Classification: 10304


We propose to evaluate a new neural network potential, So3krates [1], for the simulation of thermal transport properties of complex materials. Thermal transport is relevant in many technological applications such as waste heat management and recovery. Qualitative and quantitive understanding of this process requires a microscopic analysis of the atomic motion (= heat) which can be gained from long molecular dynamics (MD) simulations on large system sizes. These simulations require an accurate description of the interatomic interactions which can be achieved with density functional theory (DFT) [2,3], but reaching the necessary time scales and system sizes in pure ab initio MD simulations is still a big challenge [2-4]. Machine learning potentials trained on DFT reference data promise to remove this computational bottleneck, and their performance in benchmark settings is constantly increasing [5]. However, besides these impressive developments, their actual performance in real-world problems such as thermal transport simulations remains to be investigated. So3krates is a newly developed neural network potential that uses a sparse equivariant graph representation of the material [5], and features a self-attention mechanism coupled with message passing to describe interactions beyond local neighbordhoods [1], yielding fast, stable, and accurate molecular dynamics simulations. It is implemented in the automatic differentiation framework jax [6] which allows for efficient training on modern GPU hardware. Preliminary tests show a smooth training process and feasibility for thermal transport simulations. In this proposal, we want to evaluate whether So3krates can take full advantage of Berzelius' hardware, and run 1-2 test systems with converged reference results [3]. We will then proceed by testing more systems, in particular structurally complex ones that pose a bigger challenge for the potential. We can build on previous work in terms of training data [3] and simulation infrastructure [7], so that we can focus on the machine learning aspects of the problem. The goal of the study is to evaluate and potentially demonstrate the possibility to replace ab initio molecular dynamics for thermal transport simulations, and thereby reduce the computational costs of this method by up to three orders of magnitude. We ask for a default allocation sufficient to thoroughly test the described framework, and prepare follow-up proposals for more application-oriented projects. [1] https://openreview.net/forum?id=tlUnxtAmcJq [2] https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.118.175901 [3] https://arxiv.org/abs/2209.12720 [4] https://arxiv.org/abs/2209.01139 [5] https://www.nature.com/articles/s41467-022-29939-5 [6] http://github.com/google/jax [7] https://joss.theoj.org/papers/10.21105/joss.02671