Mixed Precision Fermi-Operator Expansion on Tensor Cores from a Machine Learning Perspective

We
present a second-order recursive Fermi-operator expansion scheme
using mixed precision floating point operations to perform electronic
structure calculations using tensor core units. A performance of over
100 teraFLOPs is achieved for half-precision floating point operations
on Nvidia’s A100 tensor core units. The second-order recursive
Fermi-operator scheme is formulated in terms of a generalized, differentiable
deep neural network structure, which solves the quantum mechanical
electronic structure problem. We demonstrate how this network can
be accelerated by optimizing the weight and bias values to substantially
reduce the number of layers required for convergence. We also show
how this machine learning approach can be used to optimize the coefficients
of the recursive Fermi-operator expansion to accurately represent
the fractional occupation numbers of the electronic states at finite
temperatures.