Machine-learned, finite temperature Fermi-operator expansions suitable for GPUs and AI-hardware

TL;DR

This paper introduces machine-learned recursive Fermi-operator expansion schemes for finite-temperature electronic structure calculations, optimized for GPU architectures, avoiding diagonalization and achieving significant speedups.

Contribution

It generalizes the SP2 method to finite temperatures using machine learning to optimize expansion coefficients, enabling faster density matrix computations on GPUs.

Findings

Achieves an order-of-magnitude speedup over diagonalization methods.

Eliminates the need for retraining during simulations with changing temperature and chemical potential.

Compatible with modern GPU matrix multiplication hardware.

Abstract

We present several finite-temperature recursive Fermi-operator expansion schemes based on the second-order spectral projection (SP2) method. Our approach builds on a previous observation that the electronic structure problem, as formulated through a recursive SP2 expansion, can be mapped onto the architecture of a deep neural network. Using this perspective, we generalize SP2 to finite electronic temperatures and construct machine learning models to determine optimized expansion coefficients. These coefficients are trained for a specified chemical potential and electronic temperature and are not available in closed analytical form. However, by employing an appropriate affine rescaling strategy to the Hamiltonian matrix, we eliminate the need to retrain the model during a simulation if the temperature and chemical potential change. Our approach avoids explicit diagonalization and relies solely on highly optimized matrix-matrix multiplication kernels. Compared to state-of-the-art diagonalization, we achieve an order-of-magnitude speedup in the single-particle finite-temperature density matrix calculation for small and moderately sized matrices on modern GPUs and dense matrix multiply units.