# Quantum perturbation theory using Tensor cores and a deep neural network

Joshua Finkelstein†∗, Emanuel H. Rubensson††∗, Susan M. Mniszewski‡, Christian
F. A. Negre†, Anders M. N. Niklasson†<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>†Theoretical Division, Los Alamos
National Laboratory, Los Alamos, New Mexico 87545  ‡Computer, Computational,
and Statistical Sciences Division, Los Alamos National Laboratory, Los Alamos,
New Mexico 87545  ††Division of Scientific Computing, Department of
Information Technology, Uppsala University, Box 337, SE-751 05 Uppsala, Sweden

###### Abstract

Time-independent quantum response calculations are performed using Tensor
cores. This is achieved by mapping density matrix perturbation theory onto the
computational structure of a deep neural network. The main computational cost
of each deep layer is dominated by tensor contractions, i.e. dense matrix-
matrix multiplications, in mixed precision arithmetics which achieves close to
peak performance. Quantum response calculations are demonstrated and analyzed
using self-consistent charge density-functional tight-binding theory as well
as coupled-perturbed Hartree-Fock theory. For linear response calculations, a
novel parameter-free convergence criterion is presented that is well-suited
for numerically noisy low precision floating point operations and we
demonstrate a peak performance of almost 200 Tflops using the Tensor cores of
two Nvidia A100 GPUs.

††preprint: LA-UR-22-21647

##  I Introduction

A material’s properties are in general characterized by its response to
perturbations. Quantum perturbation theory and response calculations are
therefore critical to our understanding of materials. Rayleigh-Schrödinger
perturbation theory and Green’s function approaches are commonly used in
response calculations Schrödinger (1926); Hylleraas (1930); Sternheimer
(1954); Karplus (1960); Stevens et al. (1963); Pople et al. (1979); Baroni et
al. (1987, 2001); Helgaker et al. (2002); Sakurai and Napolitano (2011), but
they require explicit knowledge of the eigenstates or complex contour
evaluations over residues, which can be a major hindrance for fast
calculations. A less frequently used method is density matrix perturbation
theory (DMPT), which is well-suited for time-independent response
calculations. DMPT was pioneered by McWeeny in the 1960’s McWeeny (1962), but
was reintroduced in a new form based on recursive Fermi-operator expansion
methods in the early 2000’s Niklasson and Challacombe (2004); Weber et al.
(2004); Niklasson et al. (2005); Weber et al. (2005); Niklasson and Weber
(2007); Niklasson et al. (2015); Truflandier et al. (2020); Shang et al.
(2021). This new form of DMPT represents a surprisingly simple methodology for
calculating time-independent quantum response properties at any order. Costly
summations and evaluations of eigenstates, as in Rayleigh-Schrödinger
perturbation theory, or the summations and evaluations of complex residues, as
in Green’s function methods, are replaced by matrix-matrix multiplications of
density matrix response terms and perturbations in the Hamiltonian, leading to
rapidly converging recursive expansions Niklasson and Challacombe (2004);
Truflandier et al. (2020); Shang et al. (2021, 2021) which are well-adapted
for modern hybrid and high-performance computing.

In this article, we explore the use of Nvidia Tensor cores NVIDIA corporation
for increasing the performance of time-independent DMPT calculations. Tensor
cores, and the related Tensor Processing Units (TPUs) Google (2021) by Google
and Matrix cores Advanced Micro Devices  from AMD, are a relatively new type
of accelerated hardware developed to address the explosion in popularity of
machine learning and artificial intelligence (AI) algorithms. These
specialized hardware units are designed to be extraordinarily efficient at
dense tensor contractions, i.e. matrix-matrix multiplications, in low floating
point (mixed) precision, which is the dominating operation required in deep
neural network (DNN) models. To tap into the raw computational power of Tensor
core units, and other similarly accelerated hardware, for more general
scientific applications, we must therefore develop stable computational
algorithms that can take advantage of fast matrix-matrix multiplications and
that are robust to low precision floating point operations. We achieve this by
mapping the recursive DMPT-based response calculations onto the computational
structure of a generalized deep neural network (DNN). This deep neural network
density matrix perturbation theory (DNN-DMPT) framework is ideally suited for
Tensor core computation. Matrix-matrix multiplications, which appear in the
activation functions of our generalized DMPT network, dominate the cost of the
DMPT-based response calculations and are naturally carried out on the Tensor
core hardware. The multiplications can be performed in mixed precision
arithmetics using a dual half-precision matrix representation to enhance the
numerical accuracy Finkelstein et al. (2021). We also derive a robust stopping
criterion that determines at what deep layer the approximate density matrix
response has converged under the numerically noisy conditions of the low
precision.

DMPT and related approaches to time-independent quantum response theory can be
designed to take advantage of matrix sparsity. This makes it possible to
achieve a reduced complexity scaling of the computational cost in quantum
response calculations as a function of system size Weber et al. (2004, 2005);
Niklasson and Weber (2007); Truflandier et al. (2020); Shang et al. (2010,
2021), which is suitable for the study of large systems. However, in this
article, we describe a dense matrix algebra approach, extending previous work
Niklasson and Challacombe (2004); Niklasson et al. (2015); Finkelstein et al.
(2021), which can utilize the substantial computing power offered by the
specialized Tensor core platform. In this way, we show how to achieve high-
performance response calculations of 120 Tflops on the Tensor cores of a
single Nvidia A100 GPU. In this article, a combined multiplication and
addition are counted as a single floating-point operation. Furthermore, by
taking advantage of the computational structure of our DNN-DMPT framework, we
also develop a multi-GPU based approach. We demonstrate a performance of
almost 200 Tflops using two sets of Tensor cores on two separate Nvidia A100
GPUs for a first-order density matrix response calculation.

The development of accelerated quantum response calculations presented in this
article continues a currently ongoing theme of using Tensor cores (or other
similar machine learning inspired hardware) for more general scientific
applications Finkelstein et al. (2021, 2021); Abdelfattah et al. (2019);
Haidar et al. (2020); Li et al. (2021); Lewis et al. (2021); Hauru et al.
(2021); Morningstar et al. (2021); Pederson et al. (2021). This work is
similar in spirit to the transition from CPUs to GPUs for scientific
computations that started over a decade ago Stone et al. (2010); Germann et
al. (2009); Ufimtsev and Martínez (2008, 2009, 2009); Luehr et al. (2011);
Maia et al. (2015); Hacene et al. (2012); Liu et al. (2012); Huhn et al.
(2020); Gordon and Windus (2020); Zhou et al. (2020); Genovese et al. (2009),
and in some sense, represents the next phase of a Darwinian-like computational
evolution driven by new hardware environments. Currently, seven of the ten
most powerful computers in the world utilize chips built with Tensor core
accelerators, and GPU-accelerated architectures are common among the top 500
supercomputers 111https://top500.org. Accessed: 2022-3-10.. Although this
article is focused on response calculations using Tensor cores, our approach
is general and should also apply to other specialized AI hardware platforms
that are designed for high-performance matrix-matrix multiplications with low
precision floating point operations, such as AMD’s Matrix cores or Google’s
TPUs.

In our presentation of time-independent quantum response calculations with
Tensor cores, we first (in Section II) discuss some background on Tensor
cores, electronic structure theory and the density matrix formalism, and how
the unperturbed density matrix can be calculated with a deep neural network
Finkelstein et al. (2021). In Section III we present our approach to time-
independent density matrix perturbation theory, and in Section IV, how we can
formulate the response calculations through a deep neural network. In Section
IV we also derive a robust parameter-free stopping criterion for the layered
network, which is of critical importance for the numerically noisy
environment. In Section V, we show some examples of density matrix linear
response calculations using our deep neural network. We analyze the numerical
accuracy and the performance. These examples highlight the utility and
efficiency of our framework. We conclude with a brief summary and some final
remarks.

##  II Deep Neural Network approximation of the electronic structure using
Tensor cores

Tensor cores can be used to calculate the electronic structure of a molecular
system using the computational framework of a convolutional deep neural
network, where the electronic structure, as in Hartree-Fock or density
functional theory Fock (1930); Roothaan (1951); Hohenberg and Kohn (1964);
Parr and Yang (1989); Dreizler and Gross (1990); Helgaker et al. (2002), is
represented by an effective single-particle density matrix McWeeny (1956).
This technique, described in the subsections below, forms the basis of our
approach to quantum response calculations using DMPT. These response
calculations can then similarly be formulated within the framework of a deep
neural network.

###  II.1 Tensor cores

A Tensor core is a specialized compute unit designed by Nvidia that
exclusively performs a single dense matrix-matrix multiplication each GPU
clock cycle. A schematic drawing of a GPU including four Tensor core units is
shown in Figure 1. In our local cluster environment, an A100 GPU has 432
Tensor cores and each A100 Tensor core can compute a $4\times 8$ matrix times
an $8\times 8$ matrix per clock cycle. Nvidia makes Tensor cores available
through low-level CUDA API commands to individual blocks of $16\times 16$
matrices and additionally through more high-level optimized CUDA libraries
such as cuBLAS NVIDIA corporation (2021), which is what is used in this
article. Since each A100 GPU has a maximum clock rate of 1.41 GHz, this
equates to a peak theoretical flop rate of approximately 156 Tflops, where we
count a combined multiplication and addition as a single floating point
operation. Compared to the standard single-precision A100 GPU compute units,
this is a more than order-of-magnitude increase in performance. Although
impressive, this gain in performance is offset by a reduction in accuracy and
under the condition that the Tensor cores only operate at peak performance
when utilizing half-precision (16-bit) multiplications with single-precision
accumulations. It is therefore paramount that any numerical algorithms used
with the Tensor core hardware are made robust to this introduction of lower
precision arithmetic. We have previously shown how this is possible for
quantum-based molecular dynamics simulations and density-matrix electronic
structure calculations using the computational framework of a generalized
convolutional deep neural network Finkelstein et al. (2021, 2021).

Figure 1: Schematic drawing of an Nvidia A100 GPU streaming multiprocessor
(SM) with 4 Tensor core (TC) units and 64 FP32 CUDA cores. Each A100 GPU on
our local cluster has 108 SM’s. INT and FP64 cores are not displayed.

###  II.2 Density matrix construction using a deep neural network

In Kohn-Sham density functional theory Hohenberg and Kohn (1964); Parr and
Yang (1989); Dreizler and Gross (1990) the electron density, $\rho({\bf r})$,
is given as the sum of densities from occupied effective single-particle
orbitals, $\\{\psi_{i}\\}$, i.e. $\rho({\bf r})=\sum_{i\in{\rm
occ}}|\psi_{i}({\bf r})|^{2}$. The single-particle molecular orbitals that
represent the electronic ground state density are obtained from the solution
of the non-linear Kohn-Sham eigenvalue equation,

$\left(-\frac{1}{2}\nabla^{2}+V_{\rm
KS}\left[\mathbf{R},\rho\right](\mathbf{r})\right)\psi_{i}(\mathbf{r})=\varepsilon_{i}\psi_{i}(\mathbf{r})\;.$
(1)

Here $V_{\rm KS}\left[\mathbf{R},\rho\right](\mathbf{r})$ is the density-
dependent effective single-particle Kohn-Sham potential (see Appendix for
details). The single-particle orbitals, $\\{\psi_{i}({\bf r})\\}$, can be
expanded in some finite basis set,

$\displaystyle\psi_{i}({\bf r})=\sum_{l}C_{i,l}\phi_{l}({\bf r}),$ (2)

so that the ground state density is given by

$\displaystyle\rho({\bf r})=\sum_{i}|\psi_{i}|^{2}$
$\displaystyle=\sum_{i}\sum_{ll^{\prime}}C_{i,l}^{*}C_{i,l^{\prime}}\phi_{l}^{*}({\bf
r})\phi_{l^{\prime}}({\bf r})$ (3)
$\displaystyle\equiv\sum_{l,l^{\prime}}D_{l,l^{\prime}}\phi_{l}^{*}({\bf
r})\phi_{l^{\prime}}({\bf r})\;,$ (4)

which defines the density matrix, $D$. In this density matrix formulation, the
non-linear Kohn-Sham eigenvalue equation is replaced by the construction of
the _density matrix_ (see Appendix for details), which is given implicitly by
the Heaviside matrix step function,

$\displaystyle D^{\perp}=\theta\left(\mu I-H^{\perp}[D]\right)\;.$ (5)

Here $H^{\perp}[D]\equiv H^{\perp}[\rho]$ is the orthogonalized matrix
representation, $H^{\perp}=Z^{T}HZ$, of the density-matrix dependent or
density dependent Kohn-Sham Hamiltonian matrix,

$\displaystyle H_{ij}[\rho]=\int\phi_{i}^{*}({\bf
r})\left(-\frac{1}{2}\nabla^{2}+V_{\rm
KS}\left[\mathbf{R},\rho\right](\mathbf{r})\right)\phi_{j}(\mathbf{r})d{\bf
r},$ (6)

and $Z$ is an inverse factor of the overlap matrix,
$O_{ij}=\int\phi_{i}^{*}({\bf r})\phi_{j}({\bf r})d{\bf r}$, such that
$Z^{T}OZ=I$. There are a number of efficient techniques to calculate $Z$ Negre
et al. (2016); Niklasson (2004); Rubensson et al. (2008); Benzi et al. (1996);
Rubensson et al. (2021). A favorite of ours is based on a matrix recursion
technique that is also well-suited for Tensor core multiplications (see
Appendix for details). The Hamiltonian matrix, $H[\rho]\equiv H[D]$, depends
on the density matrix, $D$, through the density, $\rho({\bf r})$ in the Kohn-
Sham potential, $V_{\rm KS}\left[\mathbf{R},\rho\right](\mathbf{r})$. The
solution of the electronic structure problem is therefore equivalent to a
direct self-consistent calculation of the Heaviside matrix step function of
$H^{\perp}[D]$, with the step formed at the chemical potential $\mu$, and
where $\mu$ is set such that the trace of the density matrix equals the number
of occupied states, i.e. ${\rm Tr}[D^{\perp}]=N_{\rm occ}$ .

A practical way of calculating a matrix step function is by expressing
$\theta$ as a recursive expansion Niklasson and Challacombe (2004); McWeeny
(1959); Niklasson (2011); Truflandier et al. (2020); Higham (2008). A
particularly simple and efficient technique for constructing such an expansion
is the recursive second-order spectral projection scheme (SP2) Niklasson
(2002); Rubensson (2011); Rubensson and Niklasson (2014). In this method,
repeated compositions of the polynomials $g(x)=x^{2}$ and $g(x)=2x-x^{2}$,
applied to a re-scaled Hamiltonian matrix (so that the spectrum is inside
[0,1]), leads to an approximation of the Heaviside step function in Eq. 5. In
each iteration, the polynomial is chosen based on which polynomial gives a
matrix trace closest to the true occupation number. In this way, the step is
formed automatically at the correct chemical potential.

The SP2 scheme can easily be mapped onto the computational structure of a
convolutional deep neural network (DNN) Finkelstein et al. (2021), where the
orthogonalized density matrix, given by the step function of $H^{\perp}$, is
generated through a recursive DNN expansion,

$D_{\perp}=f\left(\ldots f(W_{1}f(W_{0}X_{0}+B_{0})+B_{1})\ldots\right).$ (7)

Here $\\{W_{n}\\}$ and $\\{B_{n}\\}$ are the weight and bias values of the
network, with the matrix square activation function, $f(S)=S^{2}$. This DNN
formulation of the SP2 scheme is formally the same scheme as the original SP2
method, though it has some conceptual advantages Finkelstein et al. (2021).
The DNN formulation may also appear more familiar to many readers thanks to
the widespread popularity of machine learning methods using neural networks.
In general, we can expect Tensor cores to accelerate any recursive expansion
that is dominated by matrix-matrix multiplications. This includes matrix
inversions using Schulz iterations Schulz (1933) or the inverse factorization
of the overlap matrix Niklasson (2004), as well as higher-order recursive
Fermi-operator expansion methods Niklasson (2011); Bowler and Miyazaki (2012).
However, in most cases these schemes do not naturally fit into a DNN
structure.

After an initial normalization layer, the orthogonalized density matrix
approximation at the $n$-th deep layer of the DNN-SP2 scheme is constructed
recursively by the network and is given by $S_{n}$, where

$\displaystyle S_{n}=W_{n}X_{n}+B_{n}\;,$ (8) $\displaystyle
X_{n}=f(S_{n-1})=S_{n-1}^{2}\;,$ (9) $\displaystyle
W_{n}=\sigma_{n}I,~{}~{}\sigma_{n}=\pm 1\;,$ (10) $\displaystyle
B_{n}=(1-\sigma_{n})S_{n}\;.$ (11)

A schematic diagram illustrating the deep neural network structure is given in
Figure 2.

Figure 2:  The DNN-SP2 electronic structure method represented as a deep
neural network. Given an orthogonalized Hamiltonian matrix, $H^{\perp}$, as
the input layer, the neural network outputs the (orthogonalized) density
matrix, $D^{\perp}$, given by the Heaviside matrix step function
$D^{\perp}=\theta(\mu I-H^{\perp})$, with step formed at the Fermi level,
$\mu$. The dominating computational cost for the network is the activation
function, $f(S_{n-1})=S_{n-1}^{2}$, a matrix square operation, which Tensor
cores are optimized for. After an initial normalization layer, the weights
($W_{n}$), biases ($B_{n}$) and layers ($X_{n}$) are given by
$W_{n}=\sigma_{n}I,B_{n}=(1-\sigma_{n})S_{n}$, $S_{n}=W_{n}X_{n}+B_{n}$, and
$X_{n}=f(S_{n-1})=S_{n-1}^{2}$ for the deep layers. The number $\sigma_{n}$ is
chosen to be $\pm 1$ and $S_{n}$ represents the $n$-th density matrix
approximation.

The recursive deep neural network procedure in Eq. 7 starts with an input
layer, $X_{0}$, which is given by the orthogonalized Hamiltonian,
$X_{0}=H^{\perp}$. Next, specific weight and bias values, $W_{0}$ and $B_{0}$,
are chosen so that the eigenvalue spectrum of $H^{\perp}$ is brought to within
the interval $[0,1]$ in reverse order, i.e. with the lowest lying eigenvalues
closer to 1 and the highest closer to 0. This initial normalization step
requires some approximate upper and lower bounds of the eigenvalue spectrum of
$H^{\perp}$, which can be estimated, for example, by Gershgorin circles.
Application of the activation function (the matrix square) to this linearly
transformed initial layer moves us to the next layer, $X_{1}$. This sequence
of steps represents the initial step $f(W_{0}X_{0}+B_{0})$ in Eq. 7. After
this first layer, at each deep layer of the DNN-SP2 network in Eq. (8), the
weights are chosen as $W_{n}=\sigma_{n}I$ and biases as
$B_{n}=(1-\sigma_{n})S_{n-1}$, with $\sigma_{n}=\pm 1$. The signs are selected
such that the density matrix converges to the correct occupation by projecting
the eigenvalues of the occupied states to 1 and the unoccupied states to 0
Niklasson (2002, 2011). The layers are repeated until all eigenvalues of the
output matrix are sufficiently close to 1 or 0, i.e. $f(S_{n-1})\approx
S_{n-1}$, in which case we have an idempotent approximation of the density
matrix, i.e. $D^{\perp}=D^{\perp}D^{\perp}$ or $D=DOD$. This idempotency
condition can be used to check the convergence of the DNN-SP2 scheme Niklasson
(2011); Finkelstein et al. (2021); Kruchinina et al. (2016).

###  II.3 The activation function and its dual matrix representation

To make full use of the Tensor core hardware, the DNN-SP2 algorithm needs to
be re-purposed for low precision arithmetic, as the Tensor cores are at peak
performance only when multiplying matrices in half-precision (FP16) and
accumulating the products in single-precision (FP32). This can be accomplished
by approximating a given matrix $X$ as a sum of two low precision parts
Finkelstein et al. (2021); Markidis et al. (2018); Ootomo and Yokota (2022),

$X\approx X_{\rm high}+X_{\rm low},$ (12)

where, in the case of half-precision representations, the two matrices are
given by

$\begin{array}[]{ll}X_{\rm high}={\rm FP16}\left[X\right]\\\ X_{\rm low}={\rm
FP16}\left[X-X_{\rm high}\right].\end{array}$ (13)

Here “${\rm FP16}[X]$” denotes the half-precision (16 bit) floating point (FP)
representation of $X$. Because we represent the error between $X$ and its
half-precision representation, $X_{\rm high}$, also in half-precision, we
expect the overall approximation in Eq. 12 to have an accuracy of about
$10^{-6}$. Machine epsilon for FP16 arithmetic is approximately $10^{-3}$. The
matrix square activation function, $f$, can then be approximated by two
separate half-precision Tensor core matrix-matrix multiplications with their
product accumulated in single-precision (32 bit), that is,

$\displaystyle\begin{split}f(X)&\approx{\rm FP32}[(X_{\rm high}+X_{\rm
low})^{2}]\\\ &\approx{\rm
FP32}\left[\mathcal{A}+\mathcal{B}+\mathcal{B}^{T}\right],\end{split}$ (14)

where

$\displaystyle\mathcal{A}$ $\displaystyle=X_{\rm high}\times X_{\rm
high}~{}({\rm Tensor~{}core~{}mult.})$ (15) $\displaystyle\mathcal{B}$
$\displaystyle=X_{\rm high}\times X_{\rm low}~{}~{}({\rm
Tensor~{}core~{}mult.}),$ (16)

and the $X_{\rm low}\times X_{\rm low}$ has been discarded. To further enhance
accuracy, a refinement, or purification layer may be used as the final output
layer of the network in Eq. 7, where all matrix algebra is carried out in
double-precision, FP64 Finkelstein et al. (2021, 2021). This purification step
is equivalent to a final double flip, e.g. $\sigma=1$ followed by $\sigma=-1$
or vice versa, and was shown to improve accuracy by second order for ground-
state density matrix calculations Finkelstein et al. (2021). However,
experience has shown this additional refinement step to be unnecessary in DMPT
and other certain cases, e.g. for QMD simulations Finkelstein et al. (2021),
and we will not consider any refinement layers in this article. In the case of
DMPT, this ultimately stems from there being two main sources of error: the
idempotency error and the commutation error. In a variational total energy
expression, the energy error is first order in idempotency and second order in
commutation. This is not true for the response calculations. The additional
refinement therefore does not help.

##  III Density matrix perturbation theory as a Deep Neural Network

###  III.1 Density matrix perturbation theory

In density functional theory, a time-independent external perturbation, e.g.
from a static electric or a magnetic field, can be introduced to the
electronic structure through the Kohn-Sham Hamiltonian matrix, $H$. If
$\lambda$ is the magnitude of the perturbation we expand the Hamiltonian in
powers of $\lambda$ as,

$\displaystyle H(\lambda)=H^{(0)}+\lambda
H^{(1)}+\lambda^{2}H^{(2)}+\cdots\;,$ (17)

where $H^{(1)},H^{(2)},\ldots$ are the perturbations to various order in
$\lambda$. An example might be where $H^{(1)}$ is the dipole moment operator
that couples the system to an external field in some direction, and $\lambda$
is the field strength in that direction. The corresponding expansion of
$H(\lambda)$ in an orthogonalized representation is given by

$\displaystyle H^{\perp}(\lambda)=H^{(0)}_{\perp}+\lambda
H^{(1)}_{\perp}+\lambda^{2}H^{(2)}_{\perp}+\cdots\;,$ (18)

where $H^{(k)}_{\perp}=Z^{T}H^{(k)}Z$.

The expression in Eq. 17 will induce a similar $\lambda$-expansion in the
density matrix, $D$, i.e. describing the response in the electronic structure
due to the perturbation. In general (as described in Section II.2) this is a
non-linear matrix function of the Hamiltonian and the perturbation, where the
density matrix and its response $D^{\perp}(H(\rho,\lambda))=\theta(\mu
I-H^{\perp}(\rho,\lambda))$, has to be calculated self-consistently, both with
respect to the unperturbed ground-state density (or density matrix) and its
response to the original perturbation in $H$. This DMPT approach is a density
matrix formulation of what is usually referred to as the coupled-perturbed
self-consistent field or density functional perturbation theory Stevens et al.
(1963); Pople et al. (1979); Baroni et al. (1987, 2001); Ochsenfeld and Head-
Gordon (1997).

Expanding $D$ in $\lambda$,

$\displaystyle D(\lambda)=D^{(0)}+\lambda
D^{(1)}+\lambda^{2}D^{(2)}+\cdots\;,$ (19)

we can identify $D^{(k)}$ from

$\displaystyle D^{(k)}_{\perp}=\frac{1}{k!}\frac{\partial^{k}\theta(\mu
I-H^{\perp}(\rho,\lambda))}{\partial\lambda^{k}}{\bigg{|}}_{\lambda=0}\;,$
(20)

through a Taylor series expansion for $D(\lambda)$, where
$D^{(k)}=ZD^{(k)}_{\perp}Z^{T}$. This response in the density matrix leads to
the corresponding response in the electron density,

$\displaystyle\rho({\bf r})=\rho^{(0)}({\bf r})+\lambda\rho^{(1)}({\bf
r})+\lambda^{2}\rho^{(2)}({\bf r})+\ldots\;,$ (21)

where

$\displaystyle\rho^{(k)}({\bf
r})=\sum_{l,l^{\prime}}D_{l,l^{\prime}}^{(k)}\phi_{l}({\bf
r})\phi_{l^{\prime}}({\bf r}).$ (22)

The response in the density matrix using a recursive expansion of the matrix
step function forms the basis of DMPT and together with the electron density
couples to the perturbation in the Hamiltonian. In general, the response
equations therefore have to be solved self-consistently. In its formulation
above, the DMPT is only valid for static, time-independent perturbations of
non-degenerate systems at zero electronic temperatures.

The recursive nature of the DNN-SP2 algorithm allows us to map DMPT onto the
computational structure of a DNN. In the same way as in Ref. Niklasson and
Challacombe, 2004, we can determine how perturbations of the Hamiltonian of
each order in $\lambda$ propagate, layer by layer, through the DNN-SP2
network. An approximate density matrix and its response to various orders in
$\lambda$ are then given as the final output layer. To understand how this
works, we can first look at the non-perturbative case in Eqs. (8) and (9). The
$n$-th layer of the DNN-SP2 algorithm, $X_{n}$, is given by
$X_{n}=f(S_{n-1})=S_{n-1}^{2}$, and the $n$-th density matrix approximation,
$S_{n}$, is given by

$\displaystyle\begin{split}S_{n}&=W_{n}X_{n}+B_{n}\\\
&=\sigma_{n}S_{n-1}^{2}+(1-\sigma_{n})S_{n-1}\;.\end{split}$ (23)

If we use this update rule in Eq. (23) with an approximate density matrix
expanded up to some order in $\lambda$ at layer $n$, i.e.

$\displaystyle S_{n}(\lambda)=S_{n}^{(0)}+\lambda
S_{n}^{(1)}+\lambda^{2}S_{n}^{(2)}+\cdots~{},$ (24)

we are able to calculate the updated response order by order. For $n>0$, the
first three orders in $\lambda^{k}$ ($k=0,1,2)$, of the approximate density
matrix and its response are given by,

$\displaystyle\begin{split}S_{n}^{(0)}&=W_{n}(S_{n-1}^{(0)})^{2}+B_{n}^{(0)}\\\
S_{n}^{(1)}&=W_{n}(S_{n-1}^{(0)}S_{n-1}^{(1)}+S_{n-1}^{(1)}S_{n-1}^{(0)})+B_{n}^{(1)}\\\
S_{n}^{(2)}&=W_{n}(S_{n-1}^{(0)}S_{n-1}^{(2)}+S_{n-1}^{(1)}S_{n-1}^{(1)}\\\
&\qquad+S_{n-1}^{(2)}S_{n-1}^{(0)})+B_{n}^{(2)},\\\ \end{split}$ (25)

and more generally, for all $k$ by,

$\displaystyle
S_{n}^{(k)}=W_{n}\sum_{j=0}^{k}S_{n-1}^{(j)}S_{n-1}^{(k-j)}+B_{n}^{(k)},$ (26)

where the bias functions, $B_{n}^{(k)}$, is given by

$\displaystyle B_{n}^{(k)}=(1-\sigma_{n})S_{n-1}^{(k)}\;.$ (27)

We can then expect that $D^{(k)}_{\perp}=\lim_{n\to\infty}S_{n}^{(k)}$. At
$n=0$, the initial transformation is given via the expansion

$\displaystyle S_{0}(\lambda)$
$\displaystyle=\frac{\varepsilon_{n}I-H^{\perp}(\lambda)}{\varepsilon_{n}-\varepsilon_{1}}$
$\displaystyle=\frac{\varepsilon_{n}I-H^{(0)}_{\perp}}{\varepsilon_{n}-\varepsilon_{1}}-\lambda\frac{1}{\varepsilon_{n}-\varepsilon_{1}}H^{(1)}_{\perp}-\lambda^{2}\frac{1}{\varepsilon_{n}-\varepsilon_{1}}H^{(2)}_{\perp}-\cdots$
$\displaystyle=S_{0}^{(0)}+\lambda
S_{0}^{(1)}+\lambda^{2}S_{0}^{(2)}+\cdots\;,$

where $\varepsilon_{1},\varepsilon_{n}$ the estimates of the smallest and
largest eigenvalue bounds of $H^{(0)}_{\perp}$, respectively. The recursive
expansion of the density matrix response in Eqs. 26, 27 and 25 forms the basis
of DMPT in its DNN-SP2 form.

Alternatively, we can formulate this DMPT algorithm purely in terms of the
network parameters $X$, $W$ and $B$ which may be more desirable for direct
machine learning applications Finkelstein et al. (2021). In this case, we use
the relation

$\displaystyle X_{n+1}=(W_{n}X_{n}+B_{n})^{2}\;,$ (28)

so that, assuming $W_{n}$ is a scaled identity matrix,

$\displaystyle X_{n+1}^{(k)}$
$\displaystyle=(W_{n})^{2}\sum_{j=0}^{k}X_{n}^{(j)}X_{n}^{(k-j)}$ (29)
$\displaystyle\qquad+W_{n}\sum_{j=0}^{k}\\{X_{n}^{(j)},B_{n}^{(k-j)}\\}$ (30)
$\displaystyle\qquad\qquad+\sum_{j=0}^{k}B_{n}^{(j)}B_{n}^{(k-j)},$ (31)

where we use the anti-commutator notation, $\\{A,B\\}=AB+BA$. Notice, that for
symmetric and commuting matrices, $(AB)^{T}=B^{T}A^{T}=BA=AB$, which is used
to speed up the calculations.

In the next section, we adapt the DNN-SP2 method to calculate $D^{(k)}$, for
any $k$, using a blocked matrix representation with the same convolutional
deep neural network as the original DNN-SP2 scheme. This block-matrix
formulation is valuable because it allows a more straightforward mathematical
analysis of DMPT, while at the same time, generates an efficient way of
calculating the matrix function derivatives in Eq. (20), using dense matrix
algebra as in Eq. (26), which is well-suited to the Tensor core design.

###  III.2 Block matrix representation

We now re-express DMPT as a DNN using upper-triangular block matrices through
the DNN-SP2 method. We call this technique the DNN-SP2 Perturbation Theory
(DNN-SP2-PT) method. Given a maximum response order $m$, we map the $m$-th
order perturbative expansion for the Hamiltonian, $H=H^{(0)}+\lambda
H^{(1)}+\lambda^{2}H^{(2)}+\cdots+\lambda^{m}H^{(m)}$, onto an
$(m+1)N\times(m+1)N$ block upper-triangular matrix,

$\displaystyle{\bf H}$
$\displaystyle=\begin{pmatrix}H^{(0)}&H^{(1)}&H^{(2)}&\cdots&H^{(m)}\\\
&H^{(0)}&H^{(1)}&\cdots&H^{(m-1)}\\\ &&H^{(0)}&\cdots&H^{(m-2)}\\\
&&&\ddots&\vdots\\\ &&&&H^{(0)}\\\ \end{pmatrix}\;.$ (32)

The bold notation indicates block matrices and the “$\perp$” symbols
indicating an orthogonalized representation are dropped for convenience. All
the lower triangular blocks of ${\bf H}$ are zero matrices. Similarly, the
$m$-th order expansion for the density matrix can be mapped onto the matrix

$\displaystyle{\bf D}$
$\displaystyle=\begin{pmatrix}D^{(0)}&D^{(1)}&D^{(2)}&\cdots&D^{(m)}\\\
&D^{(0)}&D^{(1)}&\cdots&D^{(m-1)}\\\ &&D^{(0)}&\cdots&D^{(m-2)}\\\
&&&\ddots&\vdots\\\ &&&&D^{(0)}\\\ \end{pmatrix}\;.$ (33)

The idempotency conditions for the $m$-th order density matrix expansion are
then given by the blocks of ${\bf D}^{2}={\bf D}$. Similarly, the commutation
conditions for DMPT can be captured for each order from the blocks using the
block matrix equation ${\bf H}{\bf D}-{\bf D}{\bf H}=0$.

This block representation can then be used in a straightforward manner to
generalize the DNN-SP2 scheme. The input layer is set to be ${\bf X}_{0}={\bf
H}$, in the same way as for DNN-SP2, where the initial weight and bias are
also chosen to yield a shifting and rescaling of the spectrum for ${\bf H}$
(or, equivalently, $H^{(0)}$):

$\displaystyle{\bf S}_{0}$ $\displaystyle=\frac{\varepsilon_{N}{\bf I}-{\bf
H}}{\varepsilon_{N}-\varepsilon_{1}}\;.$ (34)

The values $\varepsilon_{1},\varepsilon_{N}$ are the smallest and largest
eigenvalue estimates for the non-perturbed system, $H^{(0)}$, i.e. for
$H(\lambda)$ in the limit $\lambda=0$. In the same way as before, the $n$-th
layer of the network, ${\bf X}_{n}$, is then defined through

$\displaystyle{\bf S}_{n}=W_{n}{\bf X}_{n}+{\bf B}_{n}$ (35)
$\displaystyle{\bf X}_{n}=f({\bf S}_{n-1})$ (36)

where $f$ is the matrix square activation function, i.e., $f({\bf X})={\bf
X}^{2}$, the weights, $W_{n}$, are also defined as before through
$W_{n}=\sigma_{n}{\bf I}$ (with $\sigma_{n}=\pm 1$) and the biases, ${\bf
B}_{n}$, are defined by

$\displaystyle{\bf B}_{n}$ $\displaystyle=(1-\sigma_{n}){\bf S}_{n-1}\;.$ (37)

The upper-triangular block matrix

$\displaystyle{\bf S}_{n}$
$\displaystyle=\begin{pmatrix}S_{n}^{(0)}&S_{n}^{(1)}&S_{n}^{(2)}&\cdots&S_{n}^{(m)}\\\
&S_{n}^{(0)}&S_{n}^{(1)}&\cdots&S_{n}^{(m-1)}\\\
&&S_{n}^{(0)}&\cdots&S_{n}^{(m-2)}\\\ &&&\ddots&\vdots\\\ &&&&S_{n}^{(0)}\\\
\end{pmatrix}\;,$ (38)

corresponds to the approximate density matrix expansion at the $n$-th layer,
$S_{n}^{(0)}+\lambda
S_{n}^{(1)}+\lambda^{2}S_{n}^{(2)}+\cdots+\lambda^{m}S_{n}^{(m)}$.

Figure 3 presents a visual of the deep neural network for this generalized
block matrix formulation of density matrix perturbation theory. As in the case
of regular DNN-SP2, it is possible to append an additional double-precision
refinement step to the final output layer, but again, we do not consider this
here. No significant benefit in accuracy of the response matrices was observed
with this extra refinement step.

Figure 3: Deep neural network diagram for the generalized block matrix
formulation of the DNN-SP2-PT method with $f({\bf X})={\bf X}^{2}$, a matrix
square activation function. The network is exactly the same as the network for
the DNN-SP2 method we previously presented in Ref. Finkelstein et al., 2021
except that here ${\bf B}_{k},{\bf S}_{k},{\bf
X}_{k}\in\mathbb{R}^{(m+1)N\times(m+1)N}$ are block matrices with the
different response orders on the superdiagonals, where $m$ is the highest-
order response considered and $N$ is the number of basis orbitals.

###  III.3 Activation function

Applying the matrix square activation function, $f$, to ${\bf S}_{n}$ from the
previous section leads to

$\displaystyle f({\bf S}_{n})$
$\displaystyle=\begin{pmatrix}f^{(0)}&f^{(1)}&f^{(2)}&\cdots&f^{(m)}\\\
&f^{(0)}&f^{(1)}&\cdots&f^{(m-1)}\\\ &&f^{(0)}&\cdots&f^{(m-2)}\\\
&&&\ddots&\vdots\\\ &&&&f^{(0)}\\\ \end{pmatrix}\;,$ (39)

where

$\displaystyle
f^{(k)}=f^{(k)}(S^{(0)},S^{(1)},...,S^{(k)})\equiv\sum_{j=0}^{k}S^{(j)}S^{(k-j)}\;.$
(40)

It is straightforward to verify that this matrix square activation definition
and the upper-triangular block matrix representation of the DNN-SP2 scheme
yields the exact same update rules for each matrix component, $S_{n}^{(k)}$,
as shown in the previous section in Eq. 26.

It is also worth noting in particular that $f^{(m)}$ depends only on
$S^{(0)},S^{(1)},...,S^{(m)}$. So, to obtain the $m$-th order response
approximation, $S_{n}^{(m)}$, at layer $n$, information from only the first
$m$ order terms are needed from the previous layer, $n-1$. This observation
provides a natural way of parallelizing the matrix multiplications over
multiple GPU devices. Given $m$ GPU devices, the $\ell$-th device,
$1\leq\ell\leq m$, can calculate the $f^{(\ell)}$ activation function,
requiring only that $S^{(0)},S^{(1)},...,S^{(\ell-1)}$ be sent from the
previous $\ell-1$ devices unidirectionally. Because the information flow only
moves one-way, efficiencies can be recovered as latencies are hidden more
easily. This multi-GPU approach will be demonstrated in the case of a linear
response calculation in Section IV.

Though we now need $m+1$ matrices to calculate the $m$-th order response (as
we see in Eq. 40), the additional storage requirement will turn out to not be
the principal limiting factor for system sizes that can be treated compared to
the straightforward ground state calculations Finkelstein et al. (2021) of the
density matrix. It is the Tensor core to GPU I/O that remains the bottleneck
and limits system size to around 20,000 orbitals for optimal performance.

###  III.4 Dual matrix representations for the activation function

The computational cost for each deep layer is dominated by the matrix
activation function, which requires repeated matrix-matrix multiplications. To
enhance the accuracy of the low precision floating-point operations from the
Tensor cores, we follow the same recipe as our original DNN-SP2 scheme
Finkelstein et al. (2021), and use a dual matrix representation as in Eqs.
(12) and (13). In the activation functions, $f^{(m)}$, we use this dual matrix
representation for each factor $S^{(j)}$, where

$\displaystyle\begin{split}\
f^{(m)}(S^{(0)},...,S^{(m)})&\equiv\sum_{j=0}^{m}S^{(j)}S^{(m-j)}\\\
&=\sum_{j=0}^{m}(S^{(j)}_{\rm high}+S^{(j)}_{\rm low})\\\
&\qquad\times(S^{(m-j)}_{\rm high}+S^{(m-j)}_{\rm low})\;.\end{split}$ (41)

In the same way as before, the half-precision representations of the matrices
$S^{(j)}$ are given by

$\displaystyle S^{(j)}_{\rm high}$ $\displaystyle={\rm
FP16}\left[S^{(j)}\right]$ (42) $\displaystyle S^{(j)}_{\rm low}$
$\displaystyle={\rm FP16}\left[S^{(j)}-S^{(j)}_{\rm high}\right],$ (43)

where “${\rm FP16}[S]$” denotes the half-precision (16 bit) floating point
(FP) representation of $S$ Niklasson (2011); Finkelstein et al. (2021);
Kruchinina et al. (2016). We then approximate the $m$-th order activation
function, $f^{(m)}$, by

$\displaystyle f^{(m)}\approx\sum_{j=1}^{m}\big{(}S^{(j)}_{\rm
high}S^{(m-j)}_{\rm high}+S^{(j)}_{\rm low}S^{(m-j)}_{\rm high}+S^{(j)}_{\rm
high}S^{(m-j)}_{\rm low}\big{)}\;,$ (44)

and discard the less significant low $\times$ low terms. Each summand in the
activation function $f^{(m)}$ is then approximated by separate half-precision
Tensor core matrix-matrix multiplications with their product accumulated in
single-precision (FP32), i.e.

$\displaystyle S^{(j)}S^{(m-j)}\approx{\rm FP32}\left[A+B+C\right]$ (45)
$\displaystyle\begin{split}A&=S^{(j)}_{\rm high}S^{(m-j)}_{\rm
high}~{}~{}({\rm Tensor~{}core~{}mult.}),\\\ B&=S^{(j)}_{\rm
high}S^{(m-j)}_{\rm low}~{}~{}({\rm Tensor~{}core~{}mult.}),\\\
C&=S^{(j)}_{\rm low}S^{(m-j)}_{\rm high}~{}~{}~{}({\rm
Tensor~{}core~{}mult.}).\end{split}$ (46)

This is, once again, analogous to the original DNN-SP2 scheme. Symmetry of the
Hamiltonian matrix $H$ and its perturbations $H^{(l)}$ implies

$S^{(m-j)}S^{(j)}=(S^{(j)}S^{(m-j)})^{T}$

for all $0\leq j<m$, so that we actually only need $\sim m/2$ matrix
multiplications instead of $m$ (see Python code in Supplemental Information).

##  IV Linear Response Calculations

To illustrate and analyze our DNN-SP2 DMPT framework we focus on first-order
linear response calculations, for time- and basis-set-independent
perturbations. Higher orders in the response are given by a straightforward
generalization and extension to basis-set dependent response calculations are
also possible Niklasson et al. (2005). First we look at calculations using
self-consistent density functional-based tight-binding theory and then
Hartree-Fock theory. A pseudocode for a first-order response calculation is
shown in Algorithms 1, 2 and 3 and a Python translation of it is given in the
Supporting Information.

Before analyzing the performance of Algorithm 1 we need to decide how to
determine when convergence is reached for the response. This is far from
trivial as the numerically noisy low precision Tensor core arithmetics make it
hard to determine convergence and previous stopping criteria for the
unperturbed DNN-SP2 scheme may no longer apply Finkelstein et al. (2021).

###  IV.1 Parameter-free stopping criterion

The stopping criterion that we develop in this section is based on a
generalization of the previous parameter-free convergence measures developed
for the ground state density matrix Finkelstein et al. (2021); Kruchinina et
al. (2016). Parameter-free stopping criteria are based on the idea of
identifying an analytical convergence rate that is valid under exact
arithmetics. Whenever this convergence rate is broken, which eventually occurs
because of the finite precision in the numerical calculations, we expect that
we have the best possible answer using the available numerical precision and
we terminate the calculation.

In exact arithmetic,

$\displaystyle\lim_{n\to\infty}{\bf
S}_{n}=\lim_{n\to\infty}\begin{pmatrix}S_{n}^{(0)}&S_{n}^{(1)}\\\
0&S_{n}^{(0)}\end{pmatrix}$ $\displaystyle=\begin{pmatrix}D^{(0)}&D^{(1)}\\\
0&D^{(0)}\end{pmatrix}$ (47) $\displaystyle={\bf D}\;,$ (48)

where $D^{(1)}$ is the matrix directional derivative (or Gâteaux derivative)
in the direction of the linear perturbation, $H^{(1)}$, i.e.

$\displaystyle D^{(1)}=\lim_{\lambda\to 0}\frac{D(H^{(0)}+\lambda
H^{(1)})-D(H^{(0)})}{\lambda}\;.$ (49)

We may write

$g\begin{pmatrix}H^{(0)}&H^{(1)}\\\ 0&H^{(0)}\end{pmatrix}={\bf
S}_{n}=\begin{pmatrix}S_{n}^{(0)}&S_{n}^{(1)}\\\
0&S_{n}^{(0)}\end{pmatrix}\;,$ (50)

where $g=g_{n}\circ g_{n-1}\circ\cdots\circ g_{1}$ is the composition of all
prior activation functions, weight and bias applications of the DNN-SP2 method
(as in Eq. 7), or equivalently, polynomials $x^{2}$ and $2x-x^{2}$, using the
original SP2 framework. Consider a $\sigma$–flip iteration, that is, a double
iteration choosing $\sigma_{n+1}=1$ and $\sigma_{n+2}=-1$, or vice versa, for
the next two iterations in the DNN-SP2-PT method. This is equivalent to
applying the function $h(x)=2x^{2}-x^{4}$ or $h(x)=(2x-x^{2})^{2}$ (the two
possible function compositions with $x^{2}$ and $2x-x^{2}$) to ${\bf S}_{n}$
in Eq. (50) to get,

$\displaystyle{\bf S}_{n+2}=h\begin{pmatrix}S_{n}^{(0)}&S_{n}^{(1)}\\\
0&S_{n}^{(0)}\end{pmatrix}=\left\\{\begin{array}[]{l}2{\bf S}_{n}^{2}-{\bf
S}_{n}^{4}\\\ \left(2{\bf S}_{n}-{\bf
S}_{n}^{2}\right)^{2}\end{array}\right.\;.$ (53)

The evaluation of $h$ in Eq. 53 yields the approximate density matrix and its
response at layer $n+2$, ${\bf S}_{n+2}$. Due to the particular upper-
triangular block form of ${\bf S}_{n}$, results by Mathias Mathias (1996);
Higham (2008) say that the evaluation of $h$ at the matrix ${\bf S}_{n}$ in
Eq. (53) can be written as

$\displaystyle h({\bf S}_{n})=\begin{pmatrix}h(S_{n}^{(0)})&\frac{\partial
h(S_{n}^{(0)}+\lambda S_{n}^{(1)})}{\partial\lambda}\big{|}_{\lambda=0}\\\
0&h(S_{n}^{(0)})\end{pmatrix}\;,$ (54)

where the upper right block, $S_{n+2}^{(1)}$, is the directional derivative of
$h$ at $S_{n}^{(0)}$ in the direction of $S_{n}^{(1)}$. Furthermore, assuming
an eigendecomposition $S_{n}^{(0)}=V\Lambda V^{T}$, where $\
\Lambda=\textrm{diag}(\lambda_{i})$ and $\lambda_{1}\geq\dots\geq\lambda_{N}$,
the theorem of Dalecki-Krein Higham (2008); Daletskii and Krein (1965) tells
us that this derivative can be computed analytically and is given by

$\displaystyle\left.\frac{\partial h(S_{n}^{(0)}+\lambda
S_{n}^{(1)})}{\partial\lambda}\right|_{\lambda=0}=V(G\odot(V^{T}S_{n}^{(1)}V))V^{T}\;,$
(55)

where in this expression the symbol $\odot$ represents the elementwise, or
Hadamard, product and the entries of the matrix $G$ are defined by
$G_{ij}=h[\lambda_{i},\lambda_{j}]$, the divided differences quotient for $h$
using eigenvalues $\lambda_{i}$ and $\lambda_{j}$. The divided differences
quotients are defined as

$\displaystyle
h[\lambda_{i},\lambda_{j}]=\begin{cases}\frac{h(\lambda_{i})-h(\lambda_{j})}{\lambda_{i}-\lambda_{j}},&\lambda_{i}\neq\lambda_{j}\\\
h^{\prime}(\lambda_{i}),&\lambda_{i}=\lambda_{j}\end{cases}\;.$ (56)

The expression in Eq. (53) then becomes

$\displaystyle{\bf
S}_{n+2}=\begin{pmatrix}h(S_{n}^{(0)})&V(G\odot(V^{T}S_{n}^{(1)}V))V^{T}\\\
0&h(S_{n}^{(0)})\end{pmatrix}\;.$ (57)

As a first attempt to deriving a stopping criterion, we assume that the zeroth
order approximation, $S_{n}^{(0)}$, has perfectly converged after $n$
iterations to $D^{(0)}$. Then, all of its eigenvalues, $\lambda_{i}$, will be
$0$ or $1$, so that the divided differences quotients,
$h[\lambda_{i},\lambda_{j}]$, are also $0$ or $1$ because:

$\displaystyle
h[\lambda_{i},\lambda_{j}]=\begin{cases}0=h^{\prime}(0),&\lambda_{i}=0,\lambda_{j}=0\\\
1=\frac{h(0)-h(1)}{0-1},&\lambda_{i}=0,\lambda_{j}=1\\\
1=\frac{h(1)-h(0)}{1-0},&\lambda_{i}=1,\lambda_{j}=0\\\
0=h^{\prime}(1),&\lambda_{i}=1,\lambda_{j}=1\\\ \end{cases}\;.$ (58)

Therefore we have that

$\displaystyle\Lambda=\begin{pmatrix}I&0\\\ 0&0\end{pmatrix}$ (59)

and

$\displaystyle G=\begin{pmatrix}0&&{\bf 1}\\\ {\bf 1}&&0\\\ \end{pmatrix}\;,$
(60)

where ${\bf 1}$ is a matrix of ones and $I$ is the identity matrix. The form
of $G$ then implies that the operation $G\odot$ is idempotent, so that
application of $h$ a second time to the expression in Eq. (53) does not change
it. In other words, applying another $\sigma$–flip iteration does not improve
or worsen the approximation to $D^{(1)}$. This would suggest that a single
$\sigma$-flip iteration following the convergence of $S_{n}^{(0)}$ should be
sufficient for the convergence of $S_{n}^{(1)}$. However it was sometimes
observed numerically that the first order idempotency error,
$\text{Idemp}_{n}^{(1)}=S_{n}^{(1)}-S_{n}^{(0)}S_{n}^{(1)}-S_{n}^{(1)}S_{n}^{(0)}$,
continues to decrease if one does additional $\sigma$–flip iterations.

To further understand this behavior let us again assume that $S_{n}^{(0)}$ has
converged exactly. The first order idempotency error is given by

$\displaystyle\text{Idemp}_{n}^{(1)}$
$\displaystyle=S_{n}^{(1)}-S_{n}^{(0)}S_{n}^{(1)}-S_{n}^{(1)}S_{n}^{(0)}$ (61)
$\displaystyle=VV^{T}S_{n}^{(1)}VV^{T}-V\Lambda V^{T}S_{n}^{(1)}VV^{T}$ (62)
$\displaystyle\hskip 56.9055pt-VV^{T}S_{n}^{(1)}V\Lambda V^{T}$ (63)
$\displaystyle=V((I-\Lambda)Y-Y\Lambda)V^{T}$ (64)
$\displaystyle=V\begin{pmatrix}-Y_{1,1}&0\\\ 0&Y_{2,2}\end{pmatrix}V^{T},$
(65)

where $Y=V^{T}S_{n}^{(1)}V$ and $Y_{i,j},\ i,j\in\\{1,2\\}$ are the
submatrices of $Y$ conforming to the previous block partitions of $\Lambda$
and $G$. This means that the purpose of any additional iterations should be
only to zero out the diagonal blocks of $Y$, ideally leaving the off-diagonal
blocks, which do not contribute to the idempotency error, unchanged. Any
changes to the off-diagonal blocks are the result of spurious fluctuations at
the level of the numerical accuracy. Failure to converge using a single
$\sigma$-flip iteration is due to the supposedly zero diagonal blocks of $G$
not being exactly zero. Numerical errors in those blocks play a key role for
the final convergence of $S_{n}^{(1)}$.

We therefore model the convergence in this final phase by including a block-
diagonal perturbation of $G$:

$G=\begin{pmatrix}\delta_{1}&{\bf 1}\\\ {\bf 1}&\delta_{2}\end{pmatrix}$ (66)

with

$(\delta_{1})_{i,j}=h[\lambda_{i},\lambda_{j}],\quad
i,j=1,\dots,N_{\text{occ}}$ (67)

and

$(\delta_{2})_{i,j}=h[\lambda_{N_{\text{occ}}+i},\lambda_{N_{\text{occ}}+j}],\quad
i,j=1,\dots,N-N_{\text{occ}},$ (68)

where $\lambda_{1}\approx\dots\approx\lambda_{N_{\text{occ}}}\approx 1$,
$\lambda_{N_{\text{occ}}+1}\approx\dots\approx\lambda_{N}\approx 0$ are the
actual eigenvalues of $S_{n}^{(0)}$. Given this model, we derive a worst case
reduction of the first order idempotency error by a $\sigma$-flip iteration.
The first order idempotency error after a $\sigma$-flip iteration is then
given by

$\displaystyle\text{Idemp}_{n+2}^{(1)}$
$\displaystyle=S_{n+2}^{(1)}-S_{n}^{(0)}S_{n+2}^{(1)}-S_{n+2}^{(1)}S_{n}^{(0)}$
(69) $\displaystyle=V(G\odot Y)V^{T}-V\Lambda V^{T}V(G\odot Y)V^{T}$ (70)
$\displaystyle\hskip 56.9055pt-V(G\odot Y_{n})V^{T}V\Lambda V^{T}$ (71)
$\displaystyle=V((G-\Lambda G-G\Lambda)\odot Y)V^{T}$ (72)
$\displaystyle=V\begin{pmatrix}-\delta_{1}\odot Y_{1,1}&0\\\ 0&\delta_{2}\odot
Y_{2,2}\end{pmatrix}V^{T}.$ (73)

Thus, instead of immediate convergence after a single $\sigma$-flip iteration,
we get an iteration with fast linear convergence and a small prefactor
determined by $\delta_{1}$ and $\delta_{2}$. Note that we have used
$S_{n}^{(0)}$ rather than $S_{n+2}^{(0)}$ in (69) since we have assumed our
algorithm stops updating $S_{n}^{(0)}$ once it has converged.

The absolute values of the entries of $\delta_{1}$ can be bounded as

$\displaystyle|(\delta_{1})_{i,j}|$
$\displaystyle=|h[\lambda_{i},\lambda_{j}]|$ (74)
$\displaystyle\leq\max(|h^{\prime}(\lambda_{i})|,|h^{\prime}(\lambda_{j})|)$
(75) $\displaystyle\leq\max_{k}|h^{\prime}(\lambda_{k})|$ (76)
$\displaystyle\leq C_{\text{pt}}\max_{k}|\lambda_{k}-\lambda_{k}^{2}|$ (77)
$\displaystyle=C_{\text{pt}}\|S_{n}^{(0)}-(S_{n}^{(0)})^{2}\|_{2}$ (78)
$\displaystyle\leq C_{\text{pt}}\|S_{n}^{(0)}-(S_{n}^{(0)})^{2}\|_{F}.$ (79)

The inequality in (75) is based on the assumption that all eigenvalues
$\lambda_{1},\dots,\lambda_{N_{\text{occ}}}$ are larger than the inflection
point of $h(x)$ so that $h$ is concave on their convex hull. The inflection
points of $h(x)=(2x-x^{2})^{2}$ and $h(x)=2x^{2}-x^{4}$ are
$1-1/\sqrt{3}\approx 0.42$ and $1/\sqrt{3}\approx 0.58$, respectively. The
constant introduced in (77) can be chosen as

$C_{\text{pt}}\geq 6+4\sqrt{0.25+\|S_{n}^{(0)}-(S_{n}^{(0)})^{2}\|_{F}}\;,$

which ensures that the inequality is valid on an interval containing all
eigenvalues; including the one that gives the maximum. In practice, one may
for example set $C_{\text{pt}}=9$ which makes the inequality valid on the
$[-0.25,1.25]$ interval, which should always include all eigenvalues.

The corresponding $|(\delta_{2})_{i,j}|\leq
C_{\text{pt}}\|S_{n}^{(0)}-(S_{n}^{(0)})^{2}\|_{F}$ bound can be derived
similarly.

Using all of this, and that the Frobenius norm is unitarily invariant, we get
that the reduction of the first order idempotency error in the Frobenius norm
is bounded as follows:

$\displaystyle\begin{split}\|\text{Idemp}_{n+2}^{(1)}\|_{F}&=\sqrt{\|\delta_{1}\odot
Y_{1,1}\|_{F}^{2}+\|\delta_{2}\odot Y_{2,2}\|_{F}^{2}}\\\ &\leq
C_{\text{pt}}\|S_{n}^{(0)}-(S_{n}^{(0)})^{2}\|_{F}\\\ &\hskip
28.45274pt\times\sqrt{\|Y_{1,1}\|_{F}^{2}+\|Y_{2,2}\|_{F}^{2}}\\\
&=C_{\text{pt}}\|S_{n}^{(0)}-(S_{n}^{(0)})^{2}\|_{F}\\\ &\hskip
56.9055pt\times\|\text{Idemp}_{n}^{(1)}\|_{F}.\end{split}$ (80)

Our proposed parameter-free convergence criteria is then given by this bound
no longer holding. We also remark that there is negligible cost to evaluating
the Frobenius norm and idempotency error for determining the convergence of
$S_{n}^{(1)}$, because the matrix square activation function of $S_{n}^{(0)}$
is no longer computed once $S_{n}^{(0)}$ has converged. The parameter-free
convergence criteria often leads to a few extra (around three or four)
iterations beyond what is required to reach convergence for $S_{n}^{(0)}$, and
no additional accuracy is achieved by performing more iterations once the
convergence in Eq. (80) has been reached. The convergence criterion is given
in Algorithm 3 as part of the DNN-SP2-PT scheme in Algorithms 1, 2 and 3.

Algorithm 1 Pseudocode for the DNN-SP2-PT formulation of the recursive Fermi-
operator expansion algorithm for linear response calculations.

$N$,  Number of orbitals

$N_{\rm occ}$,  Number of occupied orbitals

$H^{(0)}$ ,  Hamiltonian matrix

$H^{(1)}$ ,  First order perturbation of the Hamiltonian

$\varepsilon_{1},~{}\varepsilon_{N},~{}\mbox{Spectral bound estimates of
}H^{(0)}$

$X_{0}^{(0)}=Z^{T}H^{(0)}Z$,  Orthogonlized input

$X_{0}^{(1)}=Z^{T}H^{(1)}Z$,  Orthogonlized input response

$W_{0}=-(\varepsilon_{N}-\varepsilon_{1})^{-1},~{}~{}B_{0}^{(0)}=\varepsilon_{N}(\varepsilon_{N}-\varepsilon_{1})^{-1}I,~{}~{}B_{0}^{(1)}=0$

$S_{0}^{(0)}=W_{0}X_{0}^{(0)}+B_{0}^{(0)}$,  Initial transform, zeroth order

$N_{S}={\rm Tr}[S_{0}^{(0)}]$,  Occupation of $S_{0}^{(0)}$

$S_{0}^{(1)}=W_{0}X_{0}^{(1)}+B_{0}^{(1)}$,  Initial transform, first order

$n=1$,  Number of layers

$D^{(0)},D^{(1)}$ converged = false

while $D^{(0)}$ and $D^{(1)}$ converged = false do

$X_{n}^{(1)}=S_{n-1}^{(0)}S_{n-1}^{(1)}+S_{n-1}^{(1)}S_{n-1}^{(0)}$,
Activation function $f^{(1)}$

if $D^{(0)}$ converged = false then

$X_{n}^{(0)}=({S_{n-1}^{(0)}})^{2}$,  Activation function $f^{(0)}$

$N_{X}^{(0)}={\rm Tr}[X_{n}^{(0)}]$

${\rm IdErr}_{n}^{(0)}=N_{S}^{(0)}-N_{X}^{(0)}$,  Idemp. error estimate

$\sigma_{n}=\mbox{sign}\left(|2N_{S}^{(0)}-N_{X}^{(0)}-N_{\rm
occ}|-|N_{X}^{(0)}-N_{\rm occ}|\right)$

Call Layer Update $D^{(0)}$ (Alg. 2)

else

$S_{n}^{(0)}=S_{n-1}^{(0)}$

$\sigma_{n}=(-1)\times\sigma_{n-1}$,  Do the $\sigma$ flip

$W_{n}=\sigma_{n}I$

end if

Call Layer Update $D^{(1)}$ (Alg. 3)

$n=n+1$

end while

Algorithm 2 Parameter-free convergence criteria and deep layer update for
$D^{(0)}$. The convergence criteria shown here for ${\rm IdErr}_{n}^{(0)}$ was
derived in a previous work Finkelstein et al. (2021) by establishing an
analytical bound on ${\rm IdErr}_{n}^{(0)}$.

if ${\rm IdErr}_{n}^{(0)}<=0$ then

$D^{(0)}$ converged =
true$\text{FroErr}^{(0)}=||S_{n-1}^{(0)}-(S_{n-1}^{(0)})^{2}||_{F}$,  Idemp.
error

else if $n>2$ and $\sigma_{n-1}\neq\sigma_{n-2}$ and ${\rm
IdErr}_{n}^{(0)}>4.5\times({\rm IdErr}_{n-2}^{(0)})^{2}$ then

$D^{(0)}$ converged =
true$\text{FroErr}^{(0)}=||S_{n-1}^{(0)}-(S_{n-1}^{(0)})^{2}||_{F}$,  Idemp.
error

else

$W_{n}=\sigma_{n}I,~{}B_{n}^{(0)}=(I-W_{n})S_{n-1}^{(0)}$

$S_{n}^{(0)}=W_{n}X_{n}^{(0)}+B_{n}^{(0)}$,  Update deep layer, zeroth-order

$N_{S}^{(0)}=W_{n}N_{X}^{(0)}+(1-\sigma_{n})N_{S}^{(0)}$,  Update occ. of
$S_{n}^{(0)}$

end if

Algorithm 3 Parameter-free convergence criteria and deep layer update for
$D^{(1)}$.

if $D^{(0)}$ converged = true then

$\text{FroErr}_{n}^{(1)}=||S_{n-1}^{(1)}-(S_{n-1}^{(0)}S_{n-1}^{(1)}+S_{n-1}^{(1)}S_{n-1}^{(0)})||_{F},$
First order idemp. error

if
$\text{FroErr}_{n}^{(1)}>9\times\text{FroErr}^{(0)}\times\text{FroErr}_{n-2}^{(1)}$
then

$D^{(1)}$ converged = true

end if

end if

if $D^{(1)}$ converged = false then

$B_{n}^{(1)}=(I-W_{n})S_{n-1}^{(1)}$

$S_{n}^{(1)}=W_{n}X_{n}^{(1)}+B_{n}^{(1)},$ Update deep layer, first order

end if

###  IV.2 Uniform electric field perturbations in DFTB theory

To demonstrate and evaluate the DNN-SP2-PT scheme in Algorithms 1, 2 and 3
using Tensor cores, we will first use second-order self-consistent charge
density functional tight-binding (SCC-DFTB) theory Porezag et al. (1995);
Elstner et al. (1998); Finnis et al. (1998); Frauenheim et al. (2000);
Hourahine et al. (2020), which is an approximation of Kohn-Sham DFT. SCC-DFTB
theory is based on a second-order expansion of the Kohn-Sham DFT energy
functional around a reference electron charge distribution of overlapping
neutral atomic electron densities, where the electrostatic interaction is
approximated by overlapping atom-centered electron charge distributions at
short distances and monopole net partial atomic charges at large distances. In
all our SCC-DFTB examples we use periodic boundary conditions, but without any
k-points or Bloch functions (Gamma-point only). Our implementation is based on
the SCC-DFTB software package LATTE Bock et al. (2008). As a perturbation we
use an external electric field, ${\boldsymbol{\cal E}}=[{\cal E}_{x},{\cal
E}_{y},{\cal E}_{z}]$ that interacts with the system through the dipole
moment. Because of the periodic boundary conditions and the real-space
representation, the external field interaction term has a simple sawtooth form
and our response calculations do not use the geometric (Berry) phase approach
that plays a crucial role in the modern theory of polarization King-Smith and
Vanderbilt (1993); Resta (1993). This theory is required to better estimate
the polarizability measured in experiments for systems with periodic cells. In
this way our calculations only represent an approximate model problem, but
still fully illustrate the efficiency and accuracy of our approach to quantum
response calculations.

In our model the perturbed DFTB approximation of the Kohn-Sham Hamiltonian
matrix is,

$H=H^{(0)}+\sum_{\alpha}{\cal E}_{\alpha}H_{\alpha}^{(1)},~{}~{}\alpha=x,y,z,$
(81)

where

$H^{(1)}_{\alpha}=\frac{1}{2}\left(OR_{\alpha}+R_{\alpha}O\right),$ (82)

and $H^{(0)}$ is the DFTB Hamiltonian of the unperturbed system. Here
$O_{kl}=\int\phi_{k}^{*}({\bf r})\phi_{l}({\bf r})d{\bf r}$ is the basis-set
overlap matrix and $R_{\alpha}$ is the diagonal matrix with the atomic
coordinates (in the $\alpha=x,y,z$ direction) for each atom-centered basis-set
orbital at position ${\bf R}$. Orthogonalization by the square root of the
inverse overlap matrix, $Z=O^{-1/2}$, leads to the transformed input
Hamiltonian matrices,

$\displaystyle{H^{(0)}}^{\perp}=Z^{T}H^{(0)}Z,\quad{H^{(1)}}^{\perp}=Z^{T}H^{(1)}Z\;,$
(83)

that can be used to demonstrate and evaluate the DNN-SP2-PT calculations with
Tensor cores. In the first calculations below we will only investigate non-
self-consistent perturbation calculations of the linear response in the
electron density represented by the first-order perturbation in the density
matrix $D^{(0)}$. No coupled perturbed self-consistent field optimizations
Stevens et al. (1963); Pople et al. (1979); Baroni et al. (1987, 2001); Weber
et al. (2004, 2005) are performed for the DFTB examples. The computational
cost of repeated coupled-perturbed self-consistent-field iterations would
depend linearly on the number of iterations required to reach some chosen
convergence tolerance, which here can be ignored.

The use of Tensor cores makes a substantial difference in performance. Figure
4 displays the computational performance of the DNN-SP2-PT method when
calculating the first order response for water systems of various sizes
(periodic boxes containing varying number of water molecules randomly
distributed) using Nvidia A100 GPUs. For the water system with Hamiltonian
matrix size of $19,008\times 19,008$ (corresponding to 3,168 water molecules),
the observed peak performance of the DNN-SP2-PT calculations on the Tensor
cores of a single A100 GPU (not including memory data transfers or I/O) is
approximately 120 Tflops, consistent with previous results for the DNN-SP2
algorithm Finkelstein et al. (2021, 2021). Using the Tensor cores increases
the flop rate by an order of magnitude in comparison to the GPU-only
performance. Though, notice that the GPU calculations are single-precision
(FP32) and the dual matrix representation that is needed to achieve a
comparable level of accuracy approximately doubles the cost for the half-
precision (FP16) calculations on the Tensor cores.

To further increase performance, we take advantage of the computational
structure inherent in the DNN-DMPT formalism discussed above in Section III.3.
There, it was suggested that a multi-GPU approach could be used to exploit the
unidirectional dependency in the response calculation. The two activation
functions, $f^{(0)}$ and $f^{(1)}$, can utilize the Tensor cores of two
separate A100 GPUs to calculate $D^{(0)}$ and $D^{(1)}$. The algorithm only
requires that $S_{n}^{(0)}$ be passed, once, unidirectionally between the two
GPUs at each step. Using the Tensor cores of two GPUs, in parallel, our peak
performance is approximately 195 Tflops, an almost 63% gain in performance
over a single A100.

The flop rate for both implementations, that is the single GPU and the multi-
GPU configuration, was estimated using the formula:

$\displaystyle\text{Est. flop rate}=\frac{5\times N^{3}\times\text{number of
iterations}}{\text{time for main loop [s]}}\;$ (84)

where timings for Eq. 84 were measured using CUDA event timers for a single
run. The variation in time over consecutive runs is small, less than 5% for
the smallest system and less than 1% for the largest one. The factor 5 comes
from the fact that five total Tensor core matrix multiplications in half-
precision are needed in each iteration using the dual matrix representation,
as can be seen from Eq. 14 and Eq. 46. The update to $S^{(0)}$ requires two
Tensor core multiplications and the update to $S^{(1)}$ requires three Tensor
core multiplications.

To test the numerical accuracy of the DNN-SP2-PT method, a reference
calculation using double-precision arithmetic instead of mixed precision is
produced for each system of water. We then compare the density matrix response
generated by the DNN-SP2-PT method using Tensor cores and the dual matrix
representation with that of the reference calculation by computing the
relative 2-norm error between the two. The result is shown in the bottom panel
of Figure 4. These relative errors appear to be independent of system size and
hover around $5\times 10^{-5}$, which is about a factor 20 less than machine
epsilon for FP16 arithmetic. With 10 bits in the mantissa, FP16 provides about
3 digits accuracy, because $1/2^{10}\approx 10^{-3}$. Ideally we would
therefore have an accuracy of 6 digits with the dual matrix representation.
Even if we don’t fully reach this limit, we get an accuracy of $\sim 10^{-5}$.
A thorough investigation of Tensor core FP16/FP32 precision can be found in
Refs. Ootomo and Yokota, 2022 and Fasi et al., 2021.

590410266127201651219008$50$$100$$150$$200$Matrix size ($N$)TflopsGPU-only
(FP32) Tensor core (FP16/FP32)2x Tensor core (FP16/FP32)

500010000150002000005e-51e-4Matrix size ($N$)Relative errorDNN-SP2-PT + Tensor
cores

Figure 4: Top figure displays the flop rate for periodic box water systems
with 66% occupation, $N_{\rm occ}=(2/3)\times N$, as a function of number of
orbitals, $N$, for the DNN-SP2-PT algorithm using only one Nvidia A100 GPU
with Tensor cores (Tensor core) and two separate sets of A100 Tensor cores (2x
Tensor core), along with a standard single-precision version using one A100
GPU’s single-precision compute cores only (GPU-only). The theoretical flop
rate for a single A100 GPU using the Tensor cores is 156 Tflops and for the
A100 only using its FP32 compute cores is 9.75 Tflops nvd . Bottom figure
displays $||\cdot||_{2}$ norm error in the first-order density matrix response
produced by the DNN-SP2-PT method using Tensor cores on the test water boxes.
The reference was obtained from a double-precision DNN-SP2-PT response
calculation.

Because the DNN-SP2-PT method calculates density matrix derivatives in the
direction of the perturbation, we could, alternatively, compute these
derivatives through finite differencing. How does such a direct approach
compare to the DNN-SP2-PT scheme? To investigate this we consider two and four
point finite differencing stencils, with an accuracy of second and fourth
order in the step size, $h$. The first order density matrix $D^{(1)}$, can be
then approximated using a finite difference approximation by

$\begin{array}[]{ll}{\displaystyle
D^{(1)}=\frac{D(H^{(0)}+hH^{(1)})-D(H^{(0)}-hH^{(1)})}{2h}+\mathcal{O}(h^{2})}\;,\end{array}$
(85)

or by

$\begin{array}[]{l}{\displaystyle
D^{(1)}=\frac{-D(H^{(0)}+2hH^{(1)})+8D(H^{(0)}+hH^{(1)})}{12h}}\\\ {}{}\\\
{\displaystyle\hskip
14.22636pt+\frac{-8D(H^{(0)}-hH^{(1)})+D(H^{(0)}-2hH^{(1)})}{12h}+\mathcal{O}(h^{4})\;,}\end{array}$
(86)

for a suitably chosen step sizes $h$ and where each evaluation of $D$ uses
regular DNN-SP2 with Tensor core acceleration. The second order approximation
costs about the same as a first-order DNN-SP2-PT calculation (4 instead of 5
matrix multiplications in each layer), whereas the four-point stencil would
take about twice as long (8 instead of 5 multiplications). A water box with
$N=5904$ orbitals is used for testing. Figure 5 shows the relative error in
approximating $D^{(1)}$ by using both finite differencing stencils for several
values of step size. The second order 2-point stencil approximation yields an
almost order-of-magnitude decrease in accuracy, for the best possible case,
but requires finding the optimal value of step size, $h$, which depends on
machine epsilon for the numerical precision of the algorithm. Using the higher
order 4-point stencil yields better accuracy, though for a different $h$
value, but is still a factor 2 less accurate than our DMPT framework even for
the most optimal step size displayed. The four-point stencil requires 4
separate density matrix calculations which represents a considerable overhead
compared to the DNN-SP2-PT calculation. From this analysis it is clear that
our DNN-SP2-PT method offers significant advantages. No search for an optimal
step size $h$ is needed (which is non-trivial due to the mixed FP16/FP32
precision), the numerical accuracy is higher, and the cost is significantly
lower compared to the more accurate four-point stencil finite difference
approximation.

Figure 5: Standard finite differencing 2-point and 4-point stencil for
approximating the first-order term in the density matrix response, Eqs. (85)
and (86). The dotted line shows the relative error obtained from using our
Tensor core based DNN-SP2-PT approach.

###  IV.3 Hartree-Fock based polarizability calculations

$N$ | $N_{\rm occ}$ | DMPT | Ergo Reference
---|---|---|---
|  | TC FP16 | CPU FP64
1128 | 235 | -286.118378567 | -286.104481906
1824 | 380 | -463.263240067 | -463.26617634
2400 | 500 | -610.91459163 | -610.92738903
4560 | 950 | -1169.0446183 | -1169.104407943
7224 | 1505 | -1856.85780567 | -1856.9148556467
Table 1: Isotropic polarizabilities,
$(\alpha_{xx}+\alpha_{yy}+\alpha_{zz})/3$, of water clusters using DNN-SP2-PT
with Tensor cores. Double-precision Ergo calculations on CPUs are given as
reference values, which use the regular SP2 Fermi-operator expansion scheme
for the density matrix construction. Here, $N$ is the number of basis
functions and $N_{\rm occ}$ is the occupation number. The 6-31G∗∗ basis set is
used for the Ergo calculations.

In the demonstration and analysis of the DNN-SP2-PT approach above we used
SCC-DFTB theory with periodic boundary conditions. In this section, we use
restricted Hartree-Fock theory, where the response is given from the solution
of the coupled-perturbed self-consistent field equations Stevens et al.
(1963); Pople et al. (1979); Weber et al. (2004, 2005). Molecular clusters
without periodic boundary conditions were used. To generate the relevant self-
consistent Fockians (i.e. the effective single-particle Hamiltonians) and
their first-order perturbations we have used the electronic structure software
package Ergo Rudberg et al. (2018). Our calculations only include the last
step of this coupled-perturbed self-consistent field calculation, with the
already converged self-consistent Fockian, and its response obtained from
double-precision arithmetic. As a scalable testbed system we have used water
clusters of approximately spherical shape and the Gaussian 6-31G∗∗ basis set.
The water clusters were extracted from a molecular dynamics simulation of bulk
water at standard temperature and pressure by including all water molecules
within spheres of varying radius. The water cluster geometries are publicly
available at ergoscf.org. The coupled perturbed self-consistent field
iterations were converged to within a tolerance $\|Y-Y^{T}\|_{F}<10^{-3}$,
$Y=(H^{(0)}D^{(1)}+H^{(1)}D^{(0)})O$, for the first order commutation
conditions.

Table 1 shows the calculated isotropic polarizabilities using the DNN-SP2-PT
algorithm with Tensor cores in low mixed precision, FP16/FP32 floating-point
arithmetics, in comparison to the Ergo reference CPU calculations in double-
precision. The accuracy (4 to 5 digits) is about the same as for the SCC-DFTB
calculations with the same system-size independence. This is the expected
accuracy for the dual matrix representations using mixed half-precision
calculations with single-precision accumulation as in Eqs. (45) and (46).

##  V Conclusions

We have shown how time-independent quantum response calculations based on
density matrix perturbation theory can be generated through a deep neural
network. Tensor cores, a new type of accelerated hardware developed for
machine learning, are ideally suited for computations using this network
structure. Our results add to the nascent, yet growing, body of work that
utilizes specialized machine learning hardware for non-AI scientific purposes,
e.g. see Refs. Morningstar et al., 2021; Lewis et al., 2021; Hauru et al.,
2021; Finkelstein et al., 2021, 2021. With this article, we have further
broadened the applicability of Tensor cores to quantum response calculations,
which represents yet another example of a more general non-AI science
application. This further demonstrates their, mostly untapped, potential for
accelerating numerical methods in chemistry and materials science.

Despite the low precision arithmetics of Tensor cores, we maintain sufficient
accuracy for DFTB-based linear response calculations while obtaining
significant performance gains, more than an order of magnitude, over what can
be achieved with a high-performance GPU-only implementation. The computational
structure of density matrix perturbation theory, and our corresponding deep
neural network formulation, is uniquely suited toward leveraging Tensor cores.
This structural framework further allows us to design a natural multi-GPU
approach to perturbation theory. With two separate intra-node A100 devices, we
achieved close to 200 Tflops when computing the density matrix and its first
order response.

Though we have focused our applications to Nvidia Tensor core computations,
our DNN-SP2-PT scheme is quite general and similar gains in performance should
be expected from other accelerated hardware architectures such as TPUs and
Matrix cores Google (2021); Advanced Micro Devices .

##  VI Acknowledgments

This work is supported by the LANL LDRD-ER program, and by the U.S. Department
of Energy through the Los Alamos National Laboratory, as well as the Swedish
national strategic e-science research program (eSSENCE). We thank the CCS-7
group and Darwin cluster at Los Alamos National Laboratory for computational
resources. Darwin is funded by the Computational Systems and Software
Environments (CSSE) subprogram of LANL’s ASC program (NNSA/DOE).

##  VII Supporting Information

Complete Python implementation of linear response calculation using DNN-
SP2-PT; includes Python implementation of psuedocode in Algorithms 1-3.

##  Appendix A Kohn-Sham density-functional theory and the density matrix

The ground-state electron density, $\rho_{\rm min}(\mathbf{r})$, in Hohenberg-
Kohn density functional theory Hohenberg and Kohn (1964); Parr and Yang
(1989); Dreizler and Gross (1990) is obtained from a constrained energy
minimization over all $v$-representable densities, $\rho\in v$ (i.e. over all
physically relevant densities that can be generated by some underlying
potential),

${\displaystyle\rho_{\rm min}(\mathbf{r})=\arg\min_{\rho\in
v}\left\\{E(\mathbf{R},\rho)\left|~{}\int\rho(\mathbf{r})d\mathbf{r}=N_{e}\right.\right\\}},$
(87)

with energy functional,

${\displaystyle E(\mathbf{R},\rho)=F[\rho]+\int\rho(\mathbf{r})v_{\rm
ext}(\mathbf{R},\mathbf{r})d\mathbf{r}}.$ (88)

Here $F[\rho]$ is a system-independent universal functional, $v_{\rm
ext}(\mathbf{R},\mathbf{r})$ is the external potential, typically from
interaction with the atomic nuclei, and $N_{e}$ is the number of electrons. In
orbital-dependent Kohn-Sham Kohn and Sham (1965) density functional theory the
representation of $F[\rho]$ is given by a sum of a kinetic, Hartree (H), and
an exchange-correlation (xc) term (in atomic units),

$F[\rho]=\sum_{i}\langle\psi_{i}|-\frac{1}{2}\nabla^{2}|\psi_{i}\rangle+E_{\rm
H}[\rho]+E_{\rm xc}[\rho],$ (89)

where the electron density, $\rho$, is represented as the sum of squared
single-particle Kohn-Sham orbitals weighted by the orbital occupation factors,
$\mathfrak{f}_{i}$,

$\displaystyle\rho(\mathbf{r})=2\sum_{i}\mathfrak{f}_{i}|\psi_{i}(\mathbf{r})|^{2}\;,$
(90) $\displaystyle\int|\psi_{i}(\mathbf{r})|^{2}d{\bf r}=1\;.$ (91)

The factor of 2 is included separately, because we assume a double occupation
(a spin up and spin down electron) of each single-particle orbital $\psi_{i}$.
Here we will assume only zero electronic temperature with integer occupation
factors where $\\{\mathfrak{f}_{i}\\}$ are either 1 or 0. The molecular
orbitals $\\{\psi_{i}\\}$ are the eigenstates obtained from the self-
consistent solution to the non-linear Kohn-Sham (KS) eigenvalue equation,

$\left(-\frac{1}{2}\nabla^{2}+V_{\rm
KS}\left[\mathbf{R},\rho\right](\mathbf{r})\right)\psi_{i}(\mathbf{r})=\varepsilon_{i}\psi_{i}(\mathbf{r})\;,$
(92)

which comes from the constrained functional minimization of Eq. 87. The
equation is non-linear, because the effective single-particle Kohn-Sham
potential,

$\displaystyle V_{\rm KS}\left[\mathbf{R},\rho\right]({\bf
r})\equiv\frac{\delta E_{\rm H}[\rho]}{\delta\rho}+\frac{\delta E_{\rm
xc}[\rho]}{\delta\rho}+v_{\rm ext}({\bf R,r})\;,$ (93)

depends on the electron density, $\rho({\bf r})$, formed by the eigenstates
$\\{\psi_{i}\\}$ in Eq. (90). In our calculations we use a finite (real-
valued) basis set, $\\{\phi_{\mu}\\}$, to represent the single-particle
eigenstates,

$\displaystyle\psi_{i}({\bf r})=\sum_{l}C_{i,l}\phi_{l}({\bf r}).$ (94)

With this basis-set representation, the electron density in Eq. 90 is
described by a _density matrix_ , $D$, defined by

$\displaystyle\rho({\bf r})$
$\displaystyle=2\sum_{i}\mathfrak{f}_{i}\left(\sum_{l}C_{i,l}\phi_{l}({\bf
r})\right)\left(\sum_{m}C_{i,m}\phi_{m}({\bf r})\right)$ (95)
$\displaystyle=2\sum_{lm}D_{lm}\phi_{l}({\bf r})\phi_{m}({\bf r})$ (96)

or as

$\displaystyle D$ $\displaystyle=C\Theta C^{T}\;,$ (97)

with $\Theta$ being a diagonal matrix containing the occupation factors
$\\{\mathfrak{f}_{i}\\}$. With the density matrix representation of the
electron density, the minimization in Eq. 87 can be performed over all
components, $\\{C_{i,j}\\}$ of the density matrix, which leads to the
generalized eigenvalue problem

$\displaystyle HC=OCE\;.$ (98)

Here $O$ is the overlap matrix with matrix element $O_{ij}=\int\phi_{i}({\bf
r})\phi_{j}({\bf r})d{\bf r}$ 222The variable $S$ is usually reserved for the
overlap matrix, however we have chosen to use $S$ in a different context
throughout the article., $E$ is a diagonal matrix of the eigenvalues
$\\{\varepsilon_{i}\\}$, and $H\equiv h[\rho]\equiv H[D]$, is the charge or
density-matrix dependent Kohn-Sham Hamiltonian with matrix elements,

$H_{ij}=\int\phi_{i}({\bf r})\left(-\frac{1}{2}\nabla^{2}+V_{\rm
KS}\left[\mathbf{R},\rho\right](\mathbf{r})\right)\phi_{j}(\mathbf{r})d{\bf
r}.$ (99)

The generalized Kohn-Sham eigenvalue equation in Eq. (98) is typically
transformed into a regular eigenvalue problem,

$\displaystyle H^{\perp}C^{\perp}=C^{\perp}E,$ (100)

through a congruence transformation, where

$\displaystyle H^{\perp}=Z^{T}HZ,$ (101) $\displaystyle C^{\perp}=Z^{-1}C,$
(102)

with $Z$ defined by an inverse factorization of the overlap matrix, i.e. by

$\displaystyle Z^{T}OZ=I.$ (103)

The particular choice of $Z=O^{-1/2}$ is usually referred to as the Löwdin
orthogonalization. A number of efficient techniques can be used to calculate
$Z$ Negre et al. (2016); Niklasson (2004); Rubensson et al. (2008); Benzi et
al. (1996); Rubensson et al. (2021) of which some are also well-suited for
Tensor core calculations utilizing matrix-matrix multiplications. The
corresponding density matrix transformation is given by

$\displaystyle D=ZD^{\perp}Z^{T}.$ (104)

In a favorite inverse factorization algorithm of ours, $Z$ is given by the
following recursion, which easily can take advantage of Tensor core
calculations,

$\displaystyle X_{n}=Z_{n}^{T}SZ_{n},~{}~{}n=0,1,2,\ldots$ (105)
$\displaystyle Z_{n+1}=Z_{n}\sum_{k=0}^{2}a_{k}X_{n}^{k},$ (106)
$\displaystyle a_{k}=\\{1.875,-1.25,0.375\\}.$ (107)

Here $Z_{0}$ is some approximation such that $\|Z_{0}^{T}SZ_{0}-I\|<1$ and
this error estimate decays cubically in each iteration Niklasson (2004).

Because the occupation factors $\\{\mathfrak{f}_{i}\\}$ are either 0 or 1 at
zero electronic temperature, the density matrix in Eq. (95) and (97) is given
by a matrix step function,

$\displaystyle D^{\perp}=\theta(\mu I-H^{\perp}[D]),$ (108) $\displaystyle
D=ZD^{\perp}Z^{T},$ (109)

of the Kohn-Sham Hamiltonian in its orthogonalized representation,
$H^{\perp}[D]$, with the step formed at the chemical potential, $\mu$.

## References

  * Schrödinger (1926) Schrödinger, E. Quantisierung als Eigenwertproblem. _Annalen der Physik_ 1926, _385_ , 437–490.
  * Hylleraas (1930) Hylleraas, E. A. Über den Grundterm der Zweielektronenprobelme. _Z.S. f Physik_ 1930, _65_ , 209–225.
  * Sternheimer (1954) Sternheimer, R. M. Electronic Polarizabilities of Ions from the Hartree-Fock Wave Functions. _Phys. Rev._ 1954, _96_ , 951–968.
  * Karplus (1960) Karplus, M. Weak Interactions in Molecular Quantum Mechanics. _Rev. Mod. Phys._ 1960, _32_ , 455–460.
  * Stevens et al. (1963) Stevens, R. M.; Pitzer, R. M.; Lipscomb, W. N. Perturbed Hartree—Fock Calculations. I. Magnetic Susceptibility and Shielding in the LiH Molecule. _The Journal of Chemical Physics_ 1963, _38_ , 550–560.
  * Pople et al. (1979) Pople, J. A.; Krishnan, R.; Schlegel, H. B.; Binkley, J. S. Derivative studies in Hartree-Fock and Møller-Plesset theories. _Int. J. Quant. Chem._ 1979, _16_ , 225–241.
  * Baroni et al. (1987) Baroni, S.; Giannozzi, P.; Testa, A. Green’s-function approach to linear response in solids. _Phys. Rev. Lett._ 1987, _58_ , 1861–1864.
  * Baroni et al. (2001) Baroni, S.; de Gironcoli, S.; Dal Corso, A.; Giannozzi, P. Phonons and related crystal properties from density-functional perturbation theory. _Rev. Mod. Phys._ 2001, _73_ , 515–562.
  * Helgaker et al. (2002) Helgaker, T.; Jorgensen, P.; Olsen, J. _Molecular Electronic-Structure Theory_ , 1st ed.; John Wiley and Sons Ltd: New York, 2002.
  * Sakurai and Napolitano (2011) Sakurai, J.; Napolitano, J. _Modern Quantum mechanics, 2nd edition_ ; Person, 2011.
  * McWeeny (1962) McWeeny, R. Perturbation Theory for the Fock-Dirac Density Matrix. _Phys. Rev._ 1962, _126_ , 1028–1034.
  * Niklasson and Challacombe (2004) Niklasson, A. M. N.; Challacombe, M. Density Matrix Perturbation Theory. _Phys. Rev. Lett._ 2004, _92_ , 193001.
  * Weber et al. (2004) Weber, V.; Niklasson, A. M. N.; Challacombe, M. Ab Initio Linear Scaling Response Theory: Electric Polarizability by Perturbed Projection. _Phys. Rev. Lett._ 2004, _92_ , 193002.
  * Niklasson et al. (2005) Niklasson, A. M. N.; Weber, V.; Challacombe, M. Non-Orthogonal Density Matrix Perturbation Theory. _J. Chem. Phys._ 2005, _123_ , 44107\.
  * Weber et al. (2005) Weber, V.; Niklasson, A. M. N.; Challacombe, M. Higher Order Response in $O(N)$ by Perturbed Projection. _J. Chem. Phys._ 2005, _123_ , 44106.
  * Niklasson and Weber (2007) Niklasson, A. M. N.; Weber, V. Linear scaling density matrix perturbation theory for basis-set-dependent quantum response calculations: An orthogonal formulation. _J. Chem. Phys._ 2007, _127_ , 064105.
  * Niklasson et al. (2015) Niklasson, A. M. N.; Cawkwell, M. J.; Rubensson, E. H.; Rudberg, E. Canonical Density Matrix Perturbation Theory. _Phys. Rev. E_ 2015, _92_ , 063301.
  * Truflandier et al. (2020) Truflandier, L. A.; Dianzinga, R. M.; Bowler, D. R. Notes on density matrix perturbation theory. _J. Chem. Phys._ 2020, _153_ , 164105\.
  * Shang et al. (2021) Shang, H.; Duan, X.; Li, F.; Zhang, L.; Xu, Z.; Liu, K.; Luo, H.; Ji, Y.; Zhao, W.; Xue, W.; Chen, L.; Zhang, Y. Many-core acceleration of the first-principles all-electron quantum perturbation calculations. _Comput. Phys. Commun._ 2021, _267_ , 108045.
  * Shang et al. (2021) Shang, H.; Li, F.; Zhang, Y.; Zhang, L.; Fu, Y.; Gao, Y.; Wu, Y.; Duan, X.; Lin, R.; Liu, X.; Liu, Y.; Chen, D. Extreme-scale ab initio quantum Raman spectra simulations on the leadership HPC system in China. Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis. 2021; pp 1–13.
  * (21) NVIDIA corporation, Nvidia Tensor Cores. https://www.nvidia.com/en-us/data-center/tensor-cores/, Accessed: 2021-6-8.
  * Google (2021) Google, Tensor Processing Units. https://cloud.google.com/tpu/docs/system-architecture, 2021; Accessed: 2021-6-15.
  * (23) Advanced Micro Devices, AMD CDNA™ 2 Architecture. https://www.amd.com/en/technologies/cdna2, Accessed: 2022-3-10.
  * Finkelstein et al. (2021) Finkelstein, J.; Smith, J. S.; Mniszewski, S. M.; Barros, K.; Negre, C. F. A.; Rubensson, E. H.; Niklasson, A. M. N. Mixed Precision Fermi-Operator Expansion on Tensor Cores from a Machine Learning Perspective. _J. Chem. Theory Comput._ 2021, _17_ , 2256–2265.
  * Shang et al. (2010) Shang, H.; Xiang, H.; Li, Z.; Yang, J. Linear scaling electronic structure calculations with numerical atomic basis set. _Int. Rev. Phys. Chem._ 2010, _29_ , 665–691.
  * Shang et al. (2021) Shang, H.; Liang, W.; Zhang, Y.; Yang, J. Efficient parallel linear scaling method to get the response density matrix in all-electron real-space density-functional perturbation theory. _Comput. Phys. Commun._ 2021, _258_ , 107613.
  * Finkelstein et al. (2021) Finkelstein, J.; Smith, J. S.; Mniszewski, S. M.; Barros, K.; Negre, C. F.; Rubensson, E. H.; Niklasson, A. M. Quantum-Based Molecular Dynamics Simulations Using Tensor Cores. _J. Chem. Theory Comput._ 2021, _17_ , 6180–6192.
  * Abdelfattah et al. (2019) Abdelfattah, A.; Tomov, S.; Dongarra, J. Fast Batched Matrix Multiplication for Small Sizes Using Half-Precision Arithmetic on GPUs. 2019 IEEE International Parallel and Distributed Processing Symposium (IPDPS). 2019; pp 111–122.
  * Haidar et al. (2020) Haidar, A.; Bayraktar, H.; Tomov, S.; Dongarra, J.; Higham, N. J. Mixed-precision iterative refinement using tensor cores on GPUs to accelerate solution of linear systems. _Proc. R. Soc. A_ 2020, _476_ , 20200110.
  * Li et al. (2021) Li, B.; Cheng, S.; Lin, J. tcFFT: Accelerating Half-Precision FFT through Tensor Cores. _arXiv preprint arXiv:2104.11471_ 2021,
  * Lewis et al. (2021) Lewis, A. G. M.; Beall, J.; Ganahl, M.; Hauru, M.; Mallick, S. B.; Vidal, G. Large Scale Distributed Linear Algebra With Tensor Processing Units. 2021\.
  * Hauru et al. (2021) Hauru, M.; Morningstar, A.; Beall, J.; Ganahl, M.; Lewis, A.; Vidal, G. Simulation of quantum physics with Tensor Processing Units: brute-force computation of ground states and time evolution. 2021.
  * Morningstar et al. (2021) Morningstar, A.; Hauru, M.; Beall, J.; Ganahl, M.; Lewis, A. G. M.; Khemani, V.; Vidal, G. Simulation of quantum many-body dynamics with Tensor Processing Units: Floquet prethermalization. 2021.
  * Pederson et al. (2021) Pederson, R.; John Kozlowski, R. S.; Beall, J.; Ganahl, M.; Hauru, M.; Lewis, A. G.; Mallick, S. B.; Blum, V.; Vidal, G. Tensor Processing Units as Quantum Chemistry Supercomputers. 2021.
  * Stone et al. (2010) Stone, J. E.; Hardy, D. J.; Ufimtsev, I. S.; Schulten, K. GPU-accelerated molecular modeling coming of age. _J. Mol. Graph. Model._ 2010, _29_ , 116 – 125.
  * Germann et al. (2009) Germann, T. C.; Kadau, K.; Swaminarayan, S. 369 Tflop/s molecular dynamics simulations on the petaflop hybrid supercomputer ‘Roadrunner’. _Concurr. Comput._ 2009, _21_ , 2143–2159.
  * Ufimtsev and Martínez (2008) Ufimtsev, I. S.; Martínez, T. J. Quantum chemistry on graphical processing units. 1. Strategies for two-electron integral evaluation. _J. Chem. Theory Comput._ 2008, _4_ , 222–231.
  * Ufimtsev and Martínez (2009) Ufimtsev, I. S.; Martínez, T. J. Quantum Chemistry on Graphical Processing Units. 2. Direct Self-Consistent-Field Implementation. _J. Chem. Theory Comput._ 2009, _5_ , 1004–1015.
  * Ufimtsev and Martínez (2009) Ufimtsev, I. S.; Martínez, T. J. Quantum chemistry on graphical processing units. 3. Analytical energy gradients, geometry optimization, and first principles molecular dynamics. _J. Chem. Theory Comput._ 2009, _5_ , 2619–2628.
  * Luehr et al. (2011) Luehr, N.; Ufimtsev, I. S.; Martínez, T. J. Dynamic Precision for Electron Repulsion Integral Evaluation on Graphical Processing Units (GPUs). _J. Chem. Theory Comput._ 2011, _7_ , 949–954.
  * Maia et al. (2015) Maia, J. D. C.; Carvalho, G. A. U.; Mangueira, C. P.; Santana, S. R.; Cabral, L. A. F.; Rocha, G. B. GPU Linear Algebra Libraries and GPGPU Programming for Accelerating MOPAC Semiempirical Quantum Chemistry Calculations. _J. Chem. Theory Comput._ 2015, _11_ , 3131–3144.
  * Hacene et al. (2012) Hacene, M.; Anciaux-Sedrakian, A.; Rozanska, X.; Klahr, D.; Guignon, T.; Fleurat-Lessard, P. Accelerating VASP electronic structure calculations using graphic processing units. _J. Chem. Theory Comput._ 2012, _33_ , 2581–2589.
  * Liu et al. (2012) Liu, F.; Luehr, N.; Kulik, H. J.; Martínez, T. J. Quantum chemistry for solvated molecules on graphical processing units using polarizable continuum models. _Comput. Phys. Commun._ 2012, _8_ , 3072–3081.
  * Huhn et al. (2020) Huhn, W. P.; Lange, B.; Yu, V. W.; Yoon, M.; Blum, V. GPU acceleration of all-electron electronic structure theory using localized numeric atom-centered basis functions. _Comput. Phys. Commun._ 2020, _254_ , 107314.
  * Gordon and Windus (2020) Gordon, M. S.; Windus, T. L. Modern Architectures and Their Impact on Electronic Structure Theory. _Chem. Rev._ 2020, _120_ , 9015–9020.
  * Zhou et al. (2020) Zhou, G.; Nebgen, B.; Lubbers, N.; Malone, W.; Niklasson, A. M. N.; Tretiak, S. Graphics Processing Unit-Accelerated Semiempirical Born Oppenheimer Molecular Dynamics Using PyTorch. _J. Chem. Theory Comput._ 2020, _16_ , 4951–4962, PMID: 32609513.
  * Genovese et al. (2009) Genovese, L.; Ospici, M.; Deutsch, T.; Méhaut, J.-F.; Neelov, A.; Goedecker, S. Density functional theory calculation on many-cores hybrid central processing unit-graphic processing unit architectures. _J. Chem. Phys._ 2009, _131_ , 034103.
  * Fock (1930) Fock, V. Näherungsmethode zur Lösung des quantenmechanischen Mehrkörperproblems. _Z.S. f Physik_ 1930, _61_ , 126–148.
  * Roothaan (1951) Roothaan, C. _Rev. Mod. Phys._ 1951, _23_ , 69.
  * Hohenberg and Kohn (1964) Hohenberg, P.; Kohn, W. Inhomgenous electron gas. _Phys. Rev._ 1964, _136_ , B:864–B871.
  * Parr and Yang (1989) Parr, R. G.; Yang, W. _Density-functional theory of atoms and molecules_ ; Oxford University Press: Oxford, 1989.
  * Dreizler and Gross (1990) Dreizler, R.; Gross, K. _Density-functional theory_ ; Springer Verlag: Berlin Heidelberg, 1990.
  * McWeeny (1956) McWeeny, R. The density matrix in self-consistent field theory .1. Iterative construction of the density matrix. _Proc. R. Soc. London Ser. A-Math_ 1956, _235_ , 496.
  * NVIDIA corporation (2021) NVIDIA corporation, cuBLAS. https://developer.nvidia.com/cuBLAS, 2021; Accessed: 2021-4-15.
  * Negre et al. (2016) Negre, C. F. A.; Mniszewski, S. M.; Cawkwell, M. J.; Bock, N.; Wall, M. E.; Niklasson, A. M. N. Recursive Factorization of the Inverse Overlap Matrix in Linear-Scaling Quantum Molecular Dynamics Simulations. _J. Chem. Theory Comput._ 2016, _12_ , 3063.
  * Niklasson (2004) Niklasson, A. M. N. Iterative refinement method for the approximate factorization of a matrix inverse. _Phys. Rev. B_ 2004, _70_ , 193102.
  * Rubensson et al. (2008) Rubensson, E. H.; Bock, N.; Holmtrömm, E.; Niklasson, A. M. N. Recursive inverse factorization. _J. Chem. Phys._ 2008, _128_ , 104105\.
  * Benzi et al. (1996) Benzi, M.; Meyer, C. D.; Tuma, M. _SIAM J. Sci. Comput._ 1996, _17_ , 1135–1149.
  * Rubensson et al. (2021) Rubensson, E. H.; Artemov, A. G.; Kruchinina, A.; Rudberg, E. Localized inverse factorization. _IMA J. Numer. Anal._ 2021, _41_ , 729–763.
  * McWeeny (1959) McWeeny, R. The density matrix in many-electron quantum mechanics I. Generalized product functions. Factorization and physical interpretation of the density matrices. _Proc. Roy. Soc. A_ 1959, _253_ , 242–259.
  * Niklasson (2011) Niklasson, A. M. N. In _Linear-Scaling Techniques in Computational Chemistry and Physics: Methods and Applications_ ; Zalesny, R., Papadopoulos, M. G., Mezey, P. G., Leszczynski, J., Eds.; Springer Netherlands: Dordrecht, 2011; pp 439–473.
  * Higham (2008) Higham, N. J. _Functions of matrices: theory and computation_ ; SIAM, 2008\.
  * Niklasson (2002) Niklasson, A. M. N. Expansion algorithm for the density matrix. _Phys. Rev. B_ 2002, _66_ , 155115.
  * Rubensson (2011) Rubensson, E. H. Nonmonotonic recursive polynomial expansion for linear scaling calculation of the density matrix. _J. Chem. Theory Comput._ 2011, _7_ , 1233.
  * Rubensson and Niklasson (2014) Rubensson, E. H.; Niklasson, A. M. N. Interior eigenvalues from density matrix expansions in quantum mechanical molecular dynamics. _SIAM J. Sci. Comput._ 2014, _36_ , 148.
  * Schulz (1933) Schulz, G. _Z. Angew. Math. Mech._ 1933, _13_ , 57.
  * Bowler and Miyazaki (2012) Bowler, D. R.; Miyazaki, T. O(N) methods in electronic structure calculations. _Rep. Prog. Phys._ 2012, _75_ , 036503–036546.
  * Kruchinina et al. (2016) Kruchinina, A.; Rudberg, E.; Rubensson, E. H. Parameterless Stopping Criteria for Recursive Density Matrix Expansions. _J. Chem. Theory Comput._ 2016, _12_ , 5788–5802.
  * Markidis et al. (2018) Markidis, S.; Chien, S. W. D.; Laure, E.; Peng, I. B.; Vetter, J. S. NVIDIA Tensor Core Programmability, Performance amp; Precision. 2018 IEEE International Parallel and Distributed Processing Symposium Workshops (IPDPSW). 2018; pp 522–531.
  * Ootomo and Yokota (2022) Ootomo, H.; Yokota, R. Recovering single precision accuracy from Tensor Cores while surpassing the FP32 theoretical peak performance. _arXiv preprint arXiv:2203.03341_ 2022,
  * Ochsenfeld and Head-Gordon (1997) Ochsenfeld, C.; Head-Gordon, M. A reformulation of the coupled perturbed self-consistent field equations entirely within a local atomic orbital density matrix-based scheme. _Chem. Phys. Lett._ 1997, _270_ , 399–405.
  * Mathias (1996) Mathias, R. A Chain Rule for Matrix Functions and Applications. _SIAM J. Matrix Anal. A._ 1996, _17_ , 610–620.
  * Daletskii and Krein (1965) Daletskii, J. L.; Krein, S. Integration and differentiation of functions of Hermitian operators and applications to the theory of perturbations. _AMS Translations (2)_ 1965, _47_ , 10–1090.
  * Porezag et al. (1995) Porezag, D.; Frauenheim, T.; Köhler, T.; Seifert, G.; Kaschner, R. Construction of tight-binding-like potentials on the basis of density-functional theory: Application to carbon. _Phys. Rev. B_ 1995, _51_ , 12947–12957.
  * Elstner et al. (1998) Elstner, M.; Porezag, D.; Jungnickel, G.; Elsner, J.; Haugk, M.; Frauenheim, T.; Suhai, S.; Seifert, G. Self-consistent-charge density-functional tight-binding method for simulations of complex materials properties. _Phys. Rev. B_ 1998, _58_ , 7260–7268.
  * Finnis et al. (1998) Finnis, M. W.; Paxton, A. T.; Methfessel, M.; van Schilfgarde, M. Crystal structures of zirconia from first principles and self-consistent tight binding. _Phys. Rev. Lett._ 1998, _81_ , 5149.
  * Frauenheim et al. (2000) Frauenheim, T.; Seifert, G.; aand Z. Hajnal, M. E.; Jungnickel, G.; Poresag, D.; Suhai, S.; Scholz, R. A self-consistent charge density-functional based tight-binding method for predictive materials simulations in physics, chemistry and biology. _Phys. Stat. sol._ 2000, _217_ , 41.
  * Hourahine et al. (2020) Hourahine, B. et al. DFTB+, a software package for efficient approximate density functional theory based atomistic simulations. _J. Chem. Phys._ 2020, _152_ , 124101.
  * Bock et al. (2008) Bock, N.; Cawkwell, M. J.; Coe, J. D.; Krishnapriyan, A.; Kroonblawd, M. P.; Lang, A.; ; Liu, C.; Saez, E. M.; Mniszewski, S. M.; Negre, C. F. A.; Niklasson, A. M. N.; Sanville, E.; Wood, M. A.; Yang, P. LATTE. 2008; https://github.com/lanl/LATTE.
  * King-Smith and Vanderbilt (1993) King-Smith, R. D.; Vanderbilt, D. Theory of polarization of crystalline solids. _Phys. Rev. B_ 1993, _47_ , 1651–1654.
  * Resta (1993) Resta, R. Macroscopic Electric Polarization as a Geometric Quantum Phase. _Europhysics Letters (EPL)_ 1993, _22_ , 133–138.
  * Fasi et al. (2021) Fasi, M.; Higham, N. J.; Mikaitis, M.; Pranesh, S. Numerical behavior of NVIDIA tensor cores. _PeerJ Computer Science_ 2021, _7_ , e330.
  * (83) NVIDIA A100. https://www.nvidia.com/content/dam/en-zz/Solutions/Data-Center/nvidia-ampere-architecture-whitepaper.pdf, Accessed: 2021-4-28.
  * Rudberg et al. (2018) Rudberg, E.; Rubensson, E. H.; Sałek, P.; Kruchinina, A. Ergo: An open-source program for linear-scaling electronic structure calculations. _SoftwareX_ 2018, _7_ , 107–111.
  * Kohn and Sham (1965) Kohn, W.; Sham, L. J. Self-Consistent Equations Including Exchange And Correlation Effects. _Phys. Rev._ 1965, _140_ , 1133.