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

Abstract

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 fractional occupation numbers of the electronic states at finite temperatures.

arXiv comment10 pages, 6 figures

References in corpus (13)

• Deep Learning in Neural Networks: An Overview
• The Kernel Polynomial Method
• O(N) methods in electronic structure calculations
• Quantum Chemistry for Solvated Molecules on Graphical Processing Units (GPUs)using Polarizable Continuum Models
• Expansion algorithm for the density matrix
• Accelerating Atomic Orbital-based Electronic Structure Calculation via Pole Expansion and Selected Inversion
• Accelerating Sparse Matrix-Matrix Multiplication with GPU Tensor Cores
• GPGPU Acceleration of All-Electron Electronic Structure Theory Using Localized Numeric Atom-Centered Basis Functions
• Implicit Purification for Temperature-Dependent Density Matrices
• Non-monotonic recursive polynomial expansions for linear scaling calculation of the density matrix
• Generalized canonical purification for density matrix minimization
• Canonical density matrix perturbation theory
• Parameterless stopping criteria for recursive density matrix expansions

Cited by in corpus (3)

• Graph-based Quantum Response Theory and Shadow Born-Oppenheimer Molecular Dynamics
• Subspace recursive Fermi-operator expansion strategies for large-scale DFT eigenvalue problems on HPC architectures
• Sparse approximate matrix-matrix multiplication for density matrix purification with error control