# SLOCC determinant invariants of order $2^{n/2}$ for even $n$ qubits

Xiangrong Li1, Dafa Li2 1 Department of Mathematics, University of California,
Irvine, CA 92697-3875, USA
2 Department of mathematical sciences, Tsinghua University, Beijing 100084
CHINA

###### Abstract

In this paper, we study SLOCC determinant invariants of order $2^{n/2}$ for
any even $n$ qubits which satisfy the SLOCC determinant equations. The
determinant invariants can be constructed by a simple method and the set of
all these determinant invariants is complete with respect to permutations of
qubits. SLOCC entanglement classification can be achieved via the vanishing or
not of the determinant invariants. We exemplify the method for several even
number of qubits, with an emphasis on six qubits.

PACS Number: 03.67.Mn

Quantum entanglement is a key quantum mechanical resource in quantum
computation and information, such as quantum cryptography, quantum dense
coding and quantum teleportation Horodecki . Whereas bipartite entanglement
has been well understood, multipartite entanglement remains largely unexplored
due to the exponential growth of complexity with the number of qubits
involved.

Functions in the coefficients of pure states which are invariant under
stochastic local operations and classical communication (SLOCC) play a vital
role in the study of entanglement classification Dur ; Miyake ; Chterental ;
Cao ; LDFPLA ; LDF07a ; LDF07b ; LDFQIC09 ; Buniy ; Viehmann  as well as
entanglement measures Luque ; Osterloh05 ; LDF09b . Invariants for four qubits
have been presented in Luque , and entanglement measures might be built from
the absolute values of these invariants. The three invariants of order 4,
denoted as $L$, $M$ and $N$, can be expressed in the form of determinants.
Invariants for five qubits have been highlighted in Luque06 ; Osterloh09 . To
date, very few attempts have been made toward the generalization to higher
number of qubits. The SLOCC equations of degree 2 for even $n$ qubits and of
degree 4 for odd $n$ qubits have been recently established for two states
equivalent under SLOCC LDF07a ; LDF09b . More recently, for even $n$ qubits,
the SLOCC determinant equations of degree $2^{n/2}$ has been established and
four determinant invariants of order $2^{n/2}$ have been obtained LDFJPA . In
light of the SLOCC determinant equations, several different genuine entangled
states of even $n$ qubits inequivalent to the $|GHZ\rangle$, $|W\rangle$, and
Dicke states have been constructed.

In this paper, we construct $\binom{n-1}{n/2-1}$ SLOCC determinant invariants
of order $2^{n/2}$ for any even $n$ qubits which satisfy the SLOCC determinant
equations. We also demonstrate the completeness of the set of all these
determinant invariants with respect to permutations of qubits. For six qubits,
we explicitly derive all the ten SLOCC determinants of order $8$. The
determinant invariants can be used for SLOCC classification of any even $n$
qubits. Finally, we illustrate the application of the equations and invariants
by proposing a genuine entangled state of even $n$ qubits and showing that it
is inequivalent to the $|GHZ\rangle$, $|W\rangle$, and Dicke states.

We write the state $|\psi^{\prime}\rangle$ of even $n$ qubits as
$|\psi^{\prime}\rangle=\sum_{i=0}^{2^{n}-1}a_{i}|i\rangle$. We associate to
the state $|\psi^{\prime}\rangle$ a $2^{n/2}$ by $2^{n/2}$ coefficient matrix
$M(a,n)$ whose entries are the coefficients $a_{0},a_{1},\cdots,a_{2^{n}-1}$
arranged in ascending lexicographical order. To illustrate, we list $M(a,4)$
below as:

$M(a,4)=\left(\begin{array}[]{cccc}a_{0}&a_{1}&a_{2}&a_{3}\\\
a_{4}&a_{5}&a_{6}&a_{7}\\\ a_{8}&a_{9}&a_{10}&a_{11}\\\
a_{12}&a_{13}&a_{14}&a_{15}\end{array}\right).$ (1)

Let the state $|\psi\rangle$ of even $n$ qubits be
$|\psi\rangle=\sum_{i=0}^{2^{n}-1}b_{i}|i\rangle$. It is well known that the
states $|\psi\rangle$ and $|\psi^{\prime}\rangle$ are equivalent under SLOCC
if and only if there exist local invertible operators $\mathcal{A}_{1}$,
$\mathcal{A}_{2},\cdots,\mathcal{A}_{n}$ such that Dur

$|\psi^{\prime}\rangle=\underbrace{\mathcal{A}_{1}\otimes\mathcal{A}_{2}\otimes\cdots\otimes\mathcal{A}_{n}}_{n}|\psi\rangle.$
(2)

Let $M(b,n)$ be obtained from $M(a,n)$ by replacing $a$ by $b$ . Then the
following SLOCC determinant equation holds LDFJPA :

$\displaystyle\det M(a,n)$ $\displaystyle=$ $\displaystyle\det
M(b,n)[\det(\mathcal{A}_{1})\cdots\det(\mathcal{A}_{n})]^{2^{(n-2)/2}}.$ (3)

We refer to any determinant that satisfies Eq. (3) as a SLOCC determinant
invariant of order $2^{n/2}$ for even $n$ qubits. In particular, $\det M(a,n)$
is such a determinant invariant. Several other determinant invariants have
recently been obtained in LDFJPA . The aim of this paper is to construct all
the determinant invariants satisfying Eq. (3).

We write $|\psi^{\prime}\rangle$ in terms of an orthogonal basis as
$|\psi^{\prime}\rangle=\sum a_{i_{1}i_{2}\cdots i_{n}}|i_{1}i_{2}$ $\cdots
i_{n}\rangle$, where $i_{1}i_{2}\cdots i_{n}$ is the $n$-bit binary form of
the index $i$. Inspection of the structure of the matrix $M(a,n)$ reveals that
the coefficient $a_{i_{1}\cdots i_{n/2}i_{n/2+1}\cdots i_{n}}$ of the state
$|\psi^{\prime}\rangle$ is the entry in the $(i_{1}\cdots i_{n/2})$th row and
$(i_{n/2+1}$$\cdots$$i_{n})$th column of the matrix (define the topmost row as
the $0$th row and the leftmost column as the $0$th column). In other words,
bits $1,\cdots,{n/2}$ specify the row number, and the rest bits specify the
column number. We observe that using different bits to specify the row number
might result in matrices whose determinants are different from that of
$M(a,n)$. Since $n/2$ bits are needed to specify the row number for square
matrices, this amounts to $\binom{n}{n/2}$ different ways. But, as can easily
be verified, exchanging the row and column bits of a matrix is equivalent to
transposing the matrix. This gives a total of
$\frac{1}{2}\binom{n}{n/2}=\binom{n-1}{n/2-1}$ different determinants. As will
be seen later, these determinants satisfy the SLOCC determinant equations and
form a complete set of determinant invariants of order $2^{n/2}$ of even $n$
qubits with respect to permutations of qubits. We can construct the
determinants in the following way: use bit $n/2$ together with $n/2-1$ other
bits selected from the rest $n-1$ bits to specify the row number and the
remaining $n/2$ bits to specify the column number. We exemplify this for
$n=4$. Using bits 1 and 2 to specify the row number yields

$\displaystyle
D_{4}^{1}=\left|\begin{array}[]{*{4}{>{\centering\arraybackslash$}p{0.5cm}<{$}}}a_{0}$\@add@centering&a_{1}$\@add@centering&a_{2}$\@add@centering&a_{3}$\@add@centering\cr
a_{4}$\@add@centering&a_{5}$\@add@centering&a_{6}$\@add@centering&a_{7}$\@add@centering\cr
a_{8}$\@add@centering&a_{9}$\@add@centering&a_{10}$\@add@centering&a_{11}$\@add@centering\cr
a_{12}$\@add@centering&a_{13}$\@add@centering&a_{14}$\@add@centering&a_{15}$\<EMAIL_ADDRESS>(8)

Using bits 2 and 3 to specify the row number yields

$\displaystyle
D_{4}^{2}=\left|\begin{array}[]{*{4}{>{\centering\arraybackslash$}p{0.5cm}<{$}}}a_{0}$\@add@centering&a_{1}$\@add@centering&a_{8}$\@add@centering&a_{9}$\@add@centering\cr
a_{2}$\@add@centering&a_{3}$\@add@centering&a_{10}$\@add@centering&a_{11}$\@add@centering\cr
a_{4}$\@add@centering&a_{5}$\@add@centering&a_{12}$\@add@centering&a_{13}$\@add@centering\cr
a_{6}$\@add@centering&a_{7}$\@add@centering&a_{14}$\@add@centering&a_{15}$\<EMAIL_ADDRESS>(13)

Using bits 2 and 4 to specify the row number yields

$\displaystyle
D_{4}^{3}=\left|\begin{array}[]{*{4}{>{\centering\arraybackslash$}p{0.5cm}<{$}}}a_{0}$\@add@centering&a_{2}$\@add@centering&a_{8}$\@add@centering&a_{10}$\@add@centering\cr
a_{1}$\@add@centering&a_{3}$\@add@centering&a_{9}$\@add@centering&a_{11}$\@add@centering\cr
a_{4}$\@add@centering&a_{6}$\@add@centering&a_{12}$\@add@centering&a_{14}$\@add@centering\cr
a_{5}$\@add@centering&a_{7}$\@add@centering&a_{13}$\@add@centering&a_{15}$\<EMAIL_ADDRESS>(18)

Each of the above three determinants can be verified to satisfy Eq. (3) by
solving Eq. (2) for the coefficients of $|\psi^{\prime}\rangle$ and then
substituting the coefficients into the corresponding Eqs. (8)-(18), thereby
revealing that all the three determinants are SLOCC determinant invariants of
order 4 for four qubits LDFJPA . It is worth noting that the above three
determinants, ignoring the sign, turn out to be same as the ones given in
Luque .

While the determinants of order 4 for four qubits can be verified to satisfy
Eq. (3), this is an extremely difficult task for large number of qubits. To
solve this problem, we will resort to permutations. Suppose that we use bit
$n/2$ together with $n/2-1$ other bits $\ell_{1}$,
$\ell_{2},\cdots,\ell_{n/2-1}$ selected from the rest $n-1$ bits to specify
the row number and the remaining $n/2$ bits to specify the column number. This
gives a determinant of order $2^{n/2}$. We will show that this determinant can
be obtained by applying a permutation to $\det M(a,n)$. This can be seen as
follows. Let

$\displaystyle
C=\\{\ell_{1},\ell_{2},\cdots,\ell_{n/2-1}\\}\cap\\{1,2,\cdots,n/2-1\\},$ (19)

i.e. $C$ consists of those among the first $n/2-1$ bits which are used to
specify the row number. Consider the following two sets of bits:

$\displaystyle\\{t_{1},t_{2},\cdots,t_{k}\\}$
$\displaystyle=\\{\ell_{1},\ell_{2},\cdots,\ell_{n/2-1}\\}/C,$ (20)
$\displaystyle\\{r_{1},r_{2},\cdots,r_{k}\\}$
$\displaystyle=\\{1,2,\cdots,n/2-1\\}/C,$ (21)

for some $0\leq k\leq n/2-1$. Here $r_{1},\cdots,r_{k}$ are those among the
first $n/2$ bits which are used to specify the column number, and
$t_{1},\cdots,t_{k}$ are those among the last $n/2$ bits which are used to
specify the row number. Define the permutation

$\sigma=(r_{1},t_{1})(r_{2},t_{2})\cdots(r_{k},t_{k}).$ (22)

If $k=0$, we define $\sigma=I$. It is trivial to see that, ignoring the sign,
the determinant constructed above is equal to $\sigma\det M(a,n)$. To find all
the determinants of order $2^{n/2}$, we can simply exhaust all possible values
of $r_{1},\cdots,r_{k}$, $t_{1},\cdots,t_{k}$, and $k$, i.e. for all $1\leq
r_{1}<r_{2}<\cdots<r_{k}\leq n/2-1$, $n/2<t_{1}<t_{2}<\cdots<t_{k}\leq n$, and
$k$ varies from 0 to $n/2-1$. Inspection of the above condition yields
$\sum_{k=0}^{n/2-1}{\binom{n/2-1}{k}}{\binom{n/2}{k}}$ $={\binom{n-1}{n/2-1}}$
different permutations $\sigma$, which give rise to equally as many different
determinants $\sigma\det M(a,n)$ of order $2^{n/2}$.

Now, simply taking the permutations $\sigma$ to both sides of Eq. (3) yield
the following determinant equations:

$\displaystyle\sigma\det M(a,n)$ $\displaystyle=$ $\displaystyle\sigma\det
M(b,n)[\det(\mathcal{A}_{1})\cdots\det(\mathcal{A}_{n})]^{2^{(n-2)/2}}.$ (23)

It follows immediately from Eq. (23) that $\sigma\det M(a,n)$ are determinant
invariants of order $2^{n/2}$.

For the sake of completeness, we first do a simple manipulation of Eq. (22).
If we make use of the fact that

$\displaystyle(r_{k},t_{k})=(1,t_{k})(1,r_{k})(1,t_{k}),$ (24)

we are led to the following equation (ignoring the sign):

$\displaystyle(r_{i},t_{i})\det M(a,n)=(1,r_{i})(1,t_{i})\det M(a,n).$ (25)

That Eq. (25) holds is easily confirmed upon realizing that to take
transposition $(1,t_{i})$ to the determinant $(1,r_{i})(1,t_{i})\det M(a,n)$
is equivalent to interchanging two rows of the determinant. This leads to the
following expression for $\sigma$:

$\displaystyle\sigma=(1,r_{1})(1,t_{1})(1,r_{2})(1,t_{2})\cdots(1,r_{k})(1,t_{k}).$
(26)

Indeed, Eq. (26) is usually more convenient to use than Eq. (22). It can be
demonstrated that applying a transposition in the form $(1,i)$ with
$i=1,\cdots,n$ to the set formed by the determinant invariants constructed
above always yields the same set (ignoring the sign). Since any permutation
can be expressed as a product of transpositions in the form $(1,i)$, this
demonstrates the completeness of the set of all these determinant invariants.
This will be illustrated in discussing the cases $n=4$ and $n=6$ below.

We now proceed to present the determinant invariants for several even number
of qubits.

$n=2$: for two qubits, there is only one determinant invariant
$\left|\begin{array}[]{cc}a_{0}&a_{1}\\\ a_{2}&a_{3}\end{array}\right|$ of
order 2 (see also LDF07a ).

$n=4$: for four qubits, there are three determinant invariants of order 4,
namely, $D_{4}^{1}$, $D_{4}^{2}$, and $D_{4}^{3}$ (see Eqs. (8)-(18)). Note
that $D_{4}^{2}=(1,3)D_{4}^{1}$ and $D_{4}^{3}=(1,4)D_{4}^{1}$.

We now argue that the above three determinants form a complete set of
determinant invariants of order 4 for four qubits with respect to permutations
of qubits. This can be seen as follows. Applying any transposition $(1,i)$
with $i=1,\cdots,4$ to any one of the three determinant invariants always
yields a determinant invariant in the same set. This demonstrates the
completeness of the set formed by these determinant invariants. The results
are summarized in table 1.

Table 1: Completeness of the determinant invariants for four qubits with respect to permutations of qubits trans |  determinant invariants |
---|---|---
$(1,1)$ | $D_{4}^{1}$ | $D_{4}^{2}$ | $D_{4}^{3}$
$(1,2)$ | $D_{4}^{1}$ | $D_{4}^{3}$ | $D_{4}^{2}$
$(1,3)$ | $D_{4}^{2}$ | $D_{4}^{1}$ | $D_{4}^{3}$
$(1,4)$ | $D_{4}^{3}$ | $D_{4}^{2}$ | $D_{4}^{1}$
Table 2: Ten determinant invariants for six qubits

$\sigma$ | $I$ | $(1,4)$ | $(1,5)$ | $(1,6)$ | $(1,2)(1,4)$ | $(1,2)(1,5)$ | $(1,2)(1,6)$ | $(1,4)(1,2)(1,5)^{a}$ | $(1,4)(1,2)(1,6)^{b}$ | $(1,5)(1,2)(1,6)^{c}$
---|---|---|---|---|---|---|---|---|---|---
$\det$ | $D_{6}^{1}$ | $D_{6}^{2}$ | $D_{6}^{3}$ | $D_{6}^{4}$ | $D_{6}^{5}$ | $D_{6}^{6}$ | $D_{6}^{7}$ | $D_{6}^{8}$ | $D_{6}^{9}$ | $D_{6}^{10}$

a $(1,4)(1,2)(1,5)D_{6}^{1}=(1,3)(1,6)D_{6}^{1}=D_{6}^{8}$.

b $(1,4)(1,2)(1,6)D_{6}^{1}=(1,3)(1,5)D_{6}^{1}=D_{6}^{9}$.

c $(1,5)(1,2)(1,6)D_{6}^{1}=(1,3)(1,4)D_{6}^{1}=D_{6}^{10}$.

$n=6$: for six qubits, there are ten determinant invariants of order 8. In
table 2, we list the permutations $\sigma$ and the corresponding determinant
invariants by virtue of Eq. (26). The determinant invariants are explicitly
given as follows:

$D_{6}^{1}=\left|\begin{array}[]{cccccccc}a_{0}&a_{1}&a_{2}&a_{3}&a_{4}&a_{5}&a_{6}&a_{7}\\\
a_{8}&a_{9}&a_{10}&a_{11}&a_{12}&a_{13}&a_{14}&a_{15}\\\
a_{16}&a_{17}&a_{18}&a_{19}&a_{20}&a_{21}&a_{22}&a_{23}\\\
a_{24}&a_{25}&a_{26}&a_{27}&a_{28}&a_{29}&a_{30}&a_{31}\\\
a_{32}&a_{33}&a_{34}&a_{35}&a_{36}&a_{37}&a_{38}&a_{39}\\\
a_{40}&a_{41}&a_{42}&a_{43}&a_{44}&a_{45}&a_{46}&a_{47}\\\
a_{48}&a_{49}&a_{50}&a_{51}&a_{52}&a_{53}&a_{54}&a_{55}\\\
a_{56}&a_{57}&a_{58}&a_{59}&a_{60}&a_{61}&a_{62}&a_{63}\end{array}\right|$
$D_{6}^{2}=\left|\begin{array}[]{cccccccc}a_{0}&a_{1}&a_{2}&a_{3}&a_{32}&a_{33}&a_{34}&a_{35}\\\
a_{4}&a_{5}&a_{6}&a_{7}&a_{36}&a_{37}&a_{38}&a_{39}\\\
a_{8}&a_{9}&a_{10}&a_{11}&a_{40}&a_{41}&a_{42}&a_{43}\\\
a_{12}&a_{13}&a_{14}&a_{15}&a_{44}&a_{45}&a_{46}&a_{47}\\\
a_{16}&a_{17}&a_{18}&a_{19}&a_{48}&a_{49}&a_{50}&a_{51}\\\
a_{20}&a_{21}&a_{22}&a_{23}&a_{52}&a_{53}&a_{54}&a_{55}\\\
a_{24}&a_{25}&a_{26}&a_{27}&a_{56}&a_{57}&a_{58}&a_{59}\\\
a_{28}&a_{29}&a_{30}&a_{31}&a_{60}&a_{61}&a_{62}&a_{63}\end{array}\right|$
$D_{6}^{3}=\left|\begin{array}[]{cccccccc}a_{0}&a_{1}&a_{4}&a_{5}&a_{32}&a_{33}&a_{36}&a_{37}\\\
a_{2}&a_{3}&a_{6}&a_{7}&a_{34}&a_{35}&a_{38}&a_{39}\\\
a_{8}&a_{9}&a_{12}&a_{13}&a_{40}&a_{41}&a_{44}&a_{45}\\\
a_{10}&a_{11}&a_{14}&a_{15}&a_{42}&a_{43}&a_{46}&a_{47}\\\
a_{16}&a_{17}&a_{20}&a_{21}&a_{48}&a_{49}&a_{52}&a_{53}\\\
a_{18}&a_{19}&a_{22}&a_{23}&a_{50}&a_{51}&a_{54}&a_{55}\\\
a_{24}&a_{25}&a_{28}&a_{29}&a_{56}&a_{57}&a_{60}&a_{61}\\\
a_{26}&a_{27}&a_{30}&a_{31}&a_{58}&a_{59}&a_{62}&a_{63}\end{array}\right|$
$D_{6}^{4}=\left|\begin{array}[]{cccccccc}a_{0}&a_{2}&a_{4}&a_{6}&a_{32}&a_{34}&a_{36}&a_{38}\\\
a_{8}&a_{10}&a_{12}&a_{14}&a_{40}&a_{42}&a_{44}&a_{46}\\\
a_{1}&a_{3}&a_{5}&a_{7}&a_{33}&a_{35}&a_{37}&a_{39}\\\
a_{9}&a_{11}&a_{13}&a_{15}&a_{41}&a_{43}&a_{45}&a_{47}\\\
a_{16}&a_{18}&a_{20}&a_{22}&a_{48}&a_{50}&a_{52}&a_{54}\\\
a_{24}&a_{26}&a_{28}&a_{30}&a_{56}&a_{58}&a_{60}&a_{62}\\\
a_{17}&a_{19}&a_{21}&a_{23}&a_{49}&a_{51}&a_{53}&a_{55}\\\
a_{25}&a_{27}&a_{29}&a_{31}&a_{57}&a_{59}&a_{61}&a_{63}\end{array}\right|$
$D_{6}^{5}=\left|\begin{array}[]{cccccccc}a_{0}&a_{1}&a_{2}&a_{3}&a_{16}&a_{17}&a_{18}&a_{19}\\\
a_{4}&a_{5}&a_{6}&a_{7}&a_{20}&a_{21}&a_{22}&a_{23}\\\
a_{8}&a_{9}&a_{10}&a_{11}&a_{24}&a_{25}&a_{26}&a_{27}\\\
a_{12}&a_{13}&a_{14}&a_{15}&a_{28}&a_{29}&a_{30}&a_{31}\\\
a_{32}&a_{33}&a_{34}&a_{35}&a_{48}&a_{49}&a_{50}&a_{51}\\\
a_{36}&a_{37}&a_{38}&a_{39}&a_{52}&a_{53}&a_{54}&a_{55}\\\
a_{40}&a_{41}&a_{42}&a_{43}&a_{56}&a_{57}&a_{58}&a_{59}\\\
a_{44}&a_{45}&a_{46}&a_{47}&a_{60}&a_{61}&a_{62}&a_{63}\end{array}\right|$
$D_{6}^{6}=\left|\begin{array}[]{cccccccc}a_{0}&a_{2}&a_{8}&a_{10}&a_{32}&a_{34}&a_{40}&a_{42}\\\
a_{1}&a_{3}&a_{9}&a_{11}&a_{33}&a_{35}&a_{41}&a_{43}\\\
a_{4}&a_{6}&a_{12}&a_{14}&a_{36}&a_{38}&a_{44}&a_{46}\\\
a_{5}&a_{7}&a_{13}&a_{15}&a_{37}&a_{39}&a_{45}&a_{47}\\\
a_{16}&a_{18}&a_{24}&a_{26}&a_{48}&a_{50}&a_{56}&a_{58}\\\
a_{17}&a_{19}&a_{25}&a_{27}&a_{49}&a_{51}&a_{57}&a_{59}\\\
a_{20}&a_{22}&a_{28}&a_{30}&a_{52}&a_{54}&a_{60}&a_{62}\\\
a_{21}&a_{23}&a_{29}&a_{31}&a_{53}&a_{55}&a_{61}&a_{63}\end{array}\right|$
$D_{6}^{7}=\left|\begin{array}[]{cccccccc}a_{0}&a_{2}&a_{4}&a_{6}&a_{16}&a_{18}&a_{20}&a_{22}\\\
a_{1}&a_{3}&a_{5}&a_{7}&a_{17}&a_{19}&a_{21}&a_{23}\\\
a_{8}&a_{10}&a_{12}&a_{14}&a_{24}&a_{26}&a_{28}&a_{30}\\\
a_{9}&a_{11}&a_{13}&a_{15}&a_{25}&a_{27}&a_{29}&a_{31}\\\
a_{32}&a_{34}&a_{36}&a_{38}&a_{48}&a_{50}&a_{52}&a_{54}\\\
a_{33}&a_{35}&a_{37}&a_{39}&a_{49}&a_{51}&a_{53}&a_{55}\\\
a_{40}&a_{42}&a_{44}&a_{46}&a_{56}&a_{58}&a_{60}&a_{62}\\\
a_{41}&a_{43}&a_{45}&a_{47}&a_{57}&a_{59}&a_{61}&a_{63}\end{array}\right|$
$D_{6}^{8}=\left|\begin{array}[]{cccccccc}a_{0}&a_{2}&a_{4}&a_{6}&a_{8}&a_{10}&a_{12}&a_{14}\\\
a_{1}&a_{3}&a_{5}&a_{7}&a_{9}&a_{11}&a_{13}&a_{15}\\\
a_{16}&a_{18}&a_{20}&a_{22}&a_{24}&a_{26}&a_{28}&a_{30}\\\
a_{17}&a_{19}&a_{21}&a_{23}&a_{25}&a_{27}&a_{29}&a_{31}\\\
a_{32}&a_{34}&a_{36}&a_{38}&a_{40}&a_{42}&a_{44}&a_{46}\\\
a_{33}&a_{35}&a_{37}&a_{39}&a_{41}&a_{43}&a_{45}&a_{47}\\\
a_{48}&a_{50}&a_{52}&a_{54}&a_{56}&a_{58}&a_{60}&a_{62}\\\
a_{49}&a_{51}&a_{53}&a_{55}&a_{57}&a_{59}&a_{61}&a_{63}\end{array}\right|$
$D_{6}^{9}=\left|\begin{array}[]{cccccccc}a_{0}&a_{1}&a_{4}&a_{5}&a_{8}&a_{9}&a_{12}&a_{13}\\\
a_{2}&a_{3}&a_{6}&a_{7}&a_{10}&a_{11}&a_{14}&a_{15}\\\
a_{16}&a_{17}&a_{20}&a_{21}&a_{24}&a_{25}&a_{28}&a_{29}\\\
a_{18}&a_{19}&a_{22}&a_{23}&a_{26}&a_{27}&a_{30}&a_{31}\\\
a_{32}&a_{33}&a_{36}&a_{37}&a_{40}&a_{41}&a_{44}&a_{45}\\\
a_{34}&a_{35}&a_{38}&a_{39}&a_{42}&a_{43}&a_{46}&a_{47}\\\
a_{48}&a_{49}&a_{52}&a_{53}&a_{56}&a_{57}&a_{60}&a_{61}\\\
a_{50}&a_{51}&a_{54}&a_{55}&a_{58}&a_{59}&a_{62}&a_{63}\end{array}\right|$

$D_{6}^{10}=\left|\begin{array}[]{cccccccc}a_{0}&a_{1}&a_{2}&a_{3}&a_{8}&a_{9}&a_{10}&a_{11}\\\
a_{4}&a_{5}&a_{6}&a_{7}&a_{12}&a_{13}&a_{14}&a_{15}\\\
a_{16}&a_{17}&a_{18}&a_{19}&a_{24}&a_{25}&a_{26}&a_{27}\\\
a_{20}&a_{21}&a_{22}&a_{23}&a_{28}&a_{29}&a_{30}&a_{31}\\\
a_{32}&a_{33}&a_{34}&a_{35}&a_{40}&a_{41}&a_{42}&a_{43}\\\
a_{36}&a_{37}&a_{38}&a_{39}&a_{44}&a_{45}&a_{46}&a_{47}\\\
a_{48}&a_{49}&a_{50}&a_{51}&a_{56}&a_{57}&a_{58}&a_{59}\\\
a_{52}&a_{53}&a_{54}&a_{55}&a_{60}&a_{61}&a_{62}&a_{63}\end{array}\right|$

An argument analogous to the one for $n=4$ establishes the completeness of the
set formed by the ten determinant invariants. The results are summarized in
table 3.

Table 3: Completeness of the determinant invariants for six qubits with respect to permutations of qubits trans | determinant invariants
---|---
$(1,1)$ | $D_{6}^{1}$ | $D_{6}^{2}$ | $D_{6}^{3}$ | $D_{6}^{4}$ | $D_{6}^{5}$ | $D_{6}^{6}$ | $D_{6}^{7}$ | $D_{6}^{8}$ | $D_{6}^{9}$ | $D_{6}^{10}$
$(1,2)$ | $D_{6}^{1}$ | $D_{6}^{5}$ | $D_{6}^{6}$ | $D_{6}^{7}$ | $D_{6}^{2}$ | $D_{6}^{3}$ | $D_{6}^{4}$ | $D_{6}^{8}$ | $D_{6}^{9}$ | $D_{6}^{10}$
$(1,3)$ | $D_{6}^{1}$ | $D_{6}^{10}$ | $D_{6}^{9}$ | $D_{6}^{8}$ | $D_{6}^{5}$ | $D_{6}^{6}$ | $D_{6}^{7}$ | $D_{6}^{4}$ | $D_{6}^{3}$ | $D_{6}^{2}$
$(1,4)$ | $D_{6}^{2}$ | $D_{6}^{1}$ | $D_{6}^{3}$ | $D_{6}^{4}$ | $D_{6}^{5}$ | $D_{6}^{8}$ | $D_{6}^{9}$ | $D_{6}^{6}$ | $D_{6}^{7}$ | $D_{6}^{10}$
$(1,5)$ | $D_{6}^{3}$ | $D_{6}^{2}$ | $D_{6}^{1}$ | $D_{6}^{4}$ | $D_{6}^{8}$ | $D_{6}^{6}$ | $D_{6}^{10}$ | $D_{6}^{5}$ | $D_{6}^{9}$ | $D_{6}^{7}$
$(1,6)$ | $D_{6}^{4}$ | $D_{6}^{2}$ | $D_{6}^{3}$ | $D_{6}^{1}$ | $D_{6}^{9}$ | $D_{6}^{10}$ | $D_{6}^{7}$ | $D_{6}^{8}$ | $D_{6}^{5}$ | $D_{6}^{6}$

$n=8$: for eight qubits, there are 35 determinant invariants of order 16 (not
presented here due to space limitations).

Finally, it follows from Eq. (23) that if two $n$-qubit states
$|\psi^{\prime}\rangle$ and $|\psi\rangle$ are SLOCC equivalent, then
$\sigma\det M(a,n)$ vanishes if and only if $\sigma\det M(b,n)$ vanishes. On
the other hand, if one of the determinants $\sigma\det M(a,n)$ and $\sigma\det
M(b,n)$ vanishes while the other does not, then the states
$|\psi^{\prime}\rangle$ and $|\psi\rangle$ belong to different SLOCC
equivalent classes. Thus, each determinant divides the space of the pure
states of even $n$ qubits into two inequivalent subspaces under SLOCC. Let
$c=\binom{n-1}{n/2-1}$. In total, $c$ determinants divide the space into
$2^{c}$ subspaces (or families) under SLOCC. Here each family
$F_{\delta_{1}\cdots\delta_{c}}$ is defined as
$F_{\delta_{1}\cdots\delta_{c}}=S_{\delta_{1}}^{(1)}\cap\cdots\cap
S_{\delta_{c}}^{(c)}$, where $\delta_{i}=0$ or $1$,
$S_{0}^{(i)}=\\{|\psi\rangle|D_{n}^{i}=0\\}$, and
$S_{1}^{(i)}=\\{|\psi\rangle|D_{n}^{i}\neq 0\\}$. Clearly, some families
include infinite SLOCC classes. It is straightforward to see that if two
states are SLOCC equivalent then they belong to the same family. However, the
converse does not hold, i.e. two states belonging to the same family are not
necessarily SLOCC equivalent.

As an application of the SLOCC determinant equations and invariants, consider,
for example, the following genuine entangled state for six qubits:

$|\chi\rangle=(1/\sqrt{8})(|0\rangle+|5\rangle+|18\rangle+|23\rangle+|40\rangle+|45\rangle+|58\rangle-|63\rangle).$

We observe that all the non-zero coefficients of $|\chi\rangle$ lie on the
diagonal of $D_{6}^{10}$. This leads to non-vanishing $D_{6}^{10}$ for
$|\chi\rangle$. However, $D_{6}^{10}$ vanishes for the $|GHZ\rangle$,
$|W\rangle$ and Dicke states. In light of Eq. (23), $|\chi\rangle$ is
inequivalent to the $|GHZ\rangle$, $|W\rangle$ and Dicke states under SLOCC.
For more examples, see LDFJPA .

In summary, we have constructed the set of all determinant invariants of order
$2^{n/2}$ for any even $n$ qubits and showed that the set is complete with
respect to permutations of qubits. We have presented a simple formula for
constructing the determinant invariants and given several examples for even
$n$. The determinant invariants can be used for SLOCC classification of any
even $n$ qubits and the absolute values of the determinant invariants can be
considered as entanglement measures. Finally, a more fundamental problem is
whether the determinant invariants are independent.

###### Acknowledgements.

The paper was supported by NSFC (Grant No.10875061) and Tsinghua National
Laboratory for Information Science and Technology.

## References

  * (1) R. Horodecki, P. Horodecki, M. Horodecki, and K. Horodecki, Rev. Mod. Phys. 81, 865 (2009).
  * (2) W. Dür, G. Vidal, and J.I. Cirac, Phys. Rev. A 62, 062314 (2000).
  * (3) A. Miyake, Phys. Rev. A 67, 012108 (2003).
  * (4) O. Chterental and D.Z. Djoković, Linear Algebra Research Advances (Nova Science, Hauppauge, NY) Chap. 4, p 133 (2007).
  * (5) Y. Cao and A. M. Wang, Eur. Phys. J. D 44, 159 (2007).
  * (6) D. Li, X. Li, H. Huang and X. Li, Phys. Lett. A 359, 428 (2006).
  * (7) D. Li, X. Li, H. Huang, and X. Li, Phys. Rev. A 76, 032304 (2007).
  * (8) D. Li, X. Li, H. Huang, and X. Li, Phys. Rev. A 76, 052311 (2007).
  * (9) D. Li, X. Li, H. Huang, and X. Li, Quantum Inf. Comput. 9, 0778 (2009).
  * (10) R.V. Buniy and T.W. Kephart, arXiv:1012.2630 [quant-ph] (2010).
  * (11) O. Viehmann, C. Eltschka, and J. Siewert, Phys. Rev. A 83, 052330 (2011).
  * (12) J.-G. Luque and J.-Y. Thibon, Phys. Rev. A 67, 042303 (2003).
  * (13) A. Osterloh and J. Siewert, Phys. Rev. A 72, 012337 (2005).
  * (14) D. Li, X. Li, H. Huang, and X. Li, J. Math. Phys. 50, 012104 (2009).
  * (15) J.-G. Luque and J.-Y. Thibon, J. Phys. A 39, 371 (2006).
  * (16) D.Z. Djoković and A. Osterloh, J. Math. Phys. 50, 033509 (2009).
  * (17) X. Li and D. Li, J. Phys. A: Math. Gen. 44, 155304 (2011).


# PETScML: Second-order solvers for training regression problems in Scientific
Machine Learning

Stefano Zampini1, Umberto Zerbinati2, George Turkiyyah1, David Keyes1

(1Extreme Computing Research Center, King Abdullah University of Science and
Technology
2Mathematical Institute, University of Oxford
)

###### Abstract

In recent years, we have witnessed the emergence of scientific machine
learning as a data-driven tool for the analysis, by means of deep-learning
techniques, of data produced by computational science and engineering
applications.

At the core of these methods is the supervised training algorithm to learn the
neural network realization, a highly non-convex optimization problem that is
usually solved using stochastic gradient methods. However, distinct from deep-
learning practice, scientific machine-learning training problems feature a
much larger volume of smooth data and better characterizations of the
empirical risk functions, which make them suited for conventional solvers for
unconstrained optimization.

We introduce a lightweight software framework built on top of the Portable and
Extensible Toolkit for Scientific computation to bridge the gap between deep-
learning software and conventional solvers for unconstrained minimization.

We empirically demonstrate the superior efficacy of a trust region method
based on the Gauss-Newton approximation of the Hessian in improving the
generalization errors arising from regression tasks when learning surrogate
models for a wide range of scientific machine-learning techniques and test
cases. All the conventional second-order solvers tested, including L-BFGS and
inexact Newton with line-search, compare favorably, either in terms of cost or
accuracy, with the adaptive first-order methods used to validate the surrogate
models.

##  1 Introduction

In recent years, there has been a growing interest in incorporating data-
driven approaches into computational science and engineering, inspired by the
advancements in deep-learning methods [36, 62]. This trend has given rise to
the emerging discipline of scientific machine learning (SciML), which aims to
tackle domain-specific data challenges by harnessing the predictive
capabilities, interpretability, and domain knowledge offered by physics-based
models. One of the attractions of neural network models lies in their ability
to handle high-dimensional function approximations [16, 28]; this has led to
the design of numerous techniques as scientific tools for uncovering the
underlying physical laws hidden within experimental data. Examples include the
discovery of partial differential equations (PDEs) [11], the learning of PDEs
[58, 7], and PDE solvers [59, 75]. Noteworthily, deep-learning techniques have
successfully pushed the boundaries of molecular dynamics simulations, enabling
accurate simulations with hundreds of millions of atoms through _ab initio_
methods [31], and they have also been applied to solving the
Hamilton–Jacobi–Bellman equations encountered in deterministic control
problems [50].

###  1.1 Background

Non-convex minimization problems are a fundamental challenge when training
deep-learning models. The loss landscape can be very complicated, given the
presence of multiple minima and saddle points, each possessing distinct
generalization properties. Second-order optimization methods applied to over-
parametrized deep-learning training problems have demonstrated overfitting
tendencies, which hinder their generalization capabilities. This has led to
the widespread preference for stochastic first-order methods, as they
naturally introduce regularization in the stochastic regime, which helps to
mitigate overfitting [64]. Nevertheless, the landscape of deep learning is
ever-evolving, and in the era of vast datasets, we find ourselves in a highly
informative data regime for training these models. The sheer volume of data
can potentially shift the paradigm, creating a setting where overfitting might
no longer severely impact generalization capabilities, challenging
conventional wisdom and opening doors to novel approaches.

Current practice for training deep-learning models is based on stochastic
first-order methods that possess superior capabilities in finding local minima
with better generalization properties for over-parametrized networks and data-
scarce contexts, for example gradient descent with momentum [66] and, to cite
only a few, its adaptive relatives ADAM [32], ADAMW [42], ADAGRAD [21].
Despite the appealing simplicity of these methods, their effectiveness depends
on the training task and it quickly reaches a plateau as larger parameter
sizes are attempted. Hyperparameter tuning, including setting the mini-batch
size, the learning rate, and the annealing schedule, also depends on the
learning task and lacks theoretical support. In practice, grid-searching for a
satisfactory set of hyper-parameters adds substantial expense when training
new models.

On the other hand, given the large volume of data and better characterizations
of the empirical risk functions, in the context of SciML we may expect a
smaller influence of the “approximation” and “generalization” errors on the
approximation properties. This, in turn, shifts the emphasis more towards the
effective minimization of the “training” error, for which stochastic first-
order methods are not the best choice.

###  1.2 Related work

Second-order methods have been extensively used and analyzed in the simulation
and numerical optimization literature and have succeeded in inverse problems
in various settings. Their dimension-independent convergence rate, robustness
with respect to the condition number of the Hessian matrix, and a much smaller
set of hyper-parameters have made them the primary tool for large-scale
optimization in these areas. When it comes to deep-learning training, however,
their straightforward usage is relatively inefficient, and their memory and
computational costs pose fundamental challenges to their successful
application. This has led to several recent research works to adapt these
methods to the context and the extreme scales of current deep-learning.
Stochastic versions of quasi-Newton methods have been explored [57, 24, 68,
72, 6, 47, 55, 26, 2, 37, 46, 71, 52] and have provided improved performance
and robustness over first-order methods. However, all these methods require
modifications of the network implementation to store and operate on
intermediate products.

The rise of deep-learning has paralleled impressive advances in the
development of high-quality and highly efficient open-source software
frameworks. However, most of the focus has been on the design of the various
architectural network components and the efficient implementation of automatic
differentiation for the back-propagation algorithm. While first-order methods
are relatively straightforward to implement, second-order solvers have
received little, if any, attention from the deep-learning framework
communities, and users of these frameworks are thus left with a limited number
of choices to train their models. Much remains to follow early successes in
these directions, above all in the context of matrix-free Newton-like methods,
where recent works have just analyzed their complexity and recently
demonstrated their applicability to academic benchmark problems [41, 70]. For
a recent effort to recast many of the stochastic Quasi-Newton methods cited
above under the same computational framework, see [54].

###  1.3 Contributions

Here, we are interested in matrix-free solvers because of their general
applicability. However, these solvers are characterized by a more complicated
structure requiring a deeper hierarchy of components, like line-search
algorithms, Krylov methods, and preconditioners. These components are at the
core of high-performance libraries designed to solve large-scale minimization
problems; through carefully crafted application programming interfaces, the
scientific community has extensively validated them on various applications,
and they can offer a completely new arsenal of methods to tackle the deep-
learning training problem. In this work, we present PETScML, which is a
lightweight Python interface exposing neural networks written using PyTorch
[56] or JAX [23] to the Portable and Extensible Toolkit for Scientific
Computing (PETSc) [4] and its Python bindings module petsc4py [17], that
allows quick experimentation with the many different optimization solvers
offered by PETSc.

To showcase the efficacy of PETScML, we focus on a particular source of
applications for regression tasks in SciML to empirically demonstrate that
second-order methods can improve trained model accuracies by superior
exploitation of the smoothness of the continuous fields and the richness of
the datasets used, outperforming the hand-tuned adaptive first-order methods
used to validate the techniques. In particular, we focus on the construction
of inexpensive-to-evaluate surrogate models [53, 34, 43, 8] where networks are
trained to learn the realization of the parameter-to-observable maps used to
solve inverse problems, Bayesian inversions, and ultimately tackle the design-
of-experiments framework. Improving the generalization accuracy of surrogate
neural operators can offer a large speed-up in these applications and enable a
qualitatively new level of training performance, accuracy, and robustness
[19].

The paper is organized as follows. Section 2 introduces the minimization
problem in training deep-learning models and describes the matrix-free solvers
studied in the paper. Section 3 gives an overview of the software framework,
while Section 4 contains numerical results testing the solvers on a series of
test cases taken from the recent literature. Section 5 provides further
discussions and concludes the paper.

##  2 Deep-learning training

We are interested in the solution of the supervised learning problem, posed as
the minimization of a non-convex scalar function

$\arg\min_{\Theta}f(\Theta),\quad
f(\Theta)=\frac{1}{N}\sum_{i=1}^{N}\mathcal{L}(y(\Theta,\bar{x}_{i}),\bar{y}_{i})$
(1)

where $N$ is the total number of data points,
$\bar{x}_{i}\in\bar{X},\bar{y}_{i}\in\bar{Y}$ are the training points and
expected outputs, $y(\cdot,\cdot)$ the network prediction,
$\mathcal{L}(\cdot,\cdot)$ a scalar convex “loss” function, and $\Theta$ the
network parameters.

Usually, this minimization problem is initially recast to the so-called “mini-
batch” framework by subsampling the data in $(\bar{X},\bar{Y})$ to obtain a
stochastic estimator of $f$. Training is then performed using first-order
methods, i.e.

$\Theta_{k+1}=\Theta_{k}-\lambda_{k}g_{k},\quad
g_{k}=\frac{1}{b}\sum_{i_{k}}\nabla_{\Theta_{k}}\mathcal{L}(y(\Theta_{k},\bar{x}_{i_{k}}),\bar{y}_{i_{k}}),$

where $\lambda_{k}$ is the learning rate and $i_{k}$ a subset of the dataset
of size $b$, i.e., the mini-batch size. The above update is performed for a
certain number of passes on the entire dataset, where each single pass is
called “epoch” (i.e. an epoch is worth $N/b$ update steps). Depending on the
different variants, the stochastic gradient $g_{k}$ is scaled and adapted
using momentum to accelerate convergence.

Since larger mini-batch sizes improve the arithmetic intensity of the
computations, in an ideal world the only limiting factor for the selection of
the mini-batch size should be the amount of memory available on the computing
devices, a very precious commodity. However, the choice of the mini-batch size
influences the convergence properties and the generalization capabilities of
the stochastic methods. Optimal configurations, from the point of view of the
model accuracies, can lead to suboptimal utilization of the computing
resources [25]. For stochastic optimizers able to deal with large batch sizes,
see [73, 74].

Here, we want to investigate if optimization methods of the type

$\Theta_{k+1}=\Theta_{k}-\lambda_{k}s_{k},\quad
s_{k}=\widetilde{H}_{k}^{-1}(\Theta_{k})g_{k},$

can improve the generalization error and reduce the computational costs needed
to converge to a local minimum. $\widetilde{H}_{k}$ is some approximation of
the symmetric indefinite Hessian $\frac{\partial^{2}f}{\partial\Theta^{2}}$
and $\lambda_{k}$ may be computed by a line-search algorithm. In particular,
in this work, we consider three of the most widely used solvers for
unconstrained minimization: the Limited memory
Broyden–Fletcher–Goldfarb–Shanno Quasi-Newton method (L-BFGS, [39]), inexact
Newton with line-search (LS), and the trust region (TR) method. Although
stochastic variants of these methods have been proposed in the literature [12,
60, 70], here we only consider their deterministic, classical variants and
solve the minimization problem given in (1) directly. A detailed description
of these methods is outside the scope of this manuscript, and in what follows,
we only discuss specific aspects that are relevant to this work; the
interested reader can consult the excellent monographs [51, 15] for additional
details.

The L-BFGS method iteratively constructs a positive definite approximation of
the inverse of the Hessian by storing and combining past iterates
$(\Theta_{k},g_{k})$ using rank-2 updates. The solution of the approximated
problem is then combined with a line-search algorithm to guarantee descent in
the objective function. The major computational costs of the L-BFGS algorithm
thus consist of evaluating the objective function during the line-search
procedure and the gradient to update the L-BFGS approximation. In contrast,
the cost of the linear algebra involved with applying the L-BFGS matrix
inverse is usually negligible for large models and modest history sizes when
using dense compact formulas [13].

The LS and TR algorithms approximatively solve the linearized equations
$\widetilde{H}_{k}^{-1}g_{k}$ by means of a Krylov method; the cost of
computing the action of the Hessian matrix on the vector becomes dominant, and
the memory requirements increase to store the additional computational graphs
needed to evaluate higher-order derivatives. Deciding to which accuracy to
solve the update step is critical for the methods’ efficiency, since they
possess at most linear convergence rate when $\Theta_{k}$ is far from the
solution, and the way they compute the update $s_{k}$ is substantially
different. LS solves the linear systems up to a dynamic relative tolerance

$\|g_{k}+\widetilde{H}_{k}s_{k}\|\leq\nu_{k}\|g_{k}\|,$ (2)

where $\nu_{k}$ is usually computed by looking at successive values of the
norm of the gradient [22], and then it combines the updated solution with a
line-search algorithm, backtracking until a suitable descent in the objective
function is found. On the other hand, TR computes the direction of descent by
solving the constrained quadratic model problem

$\arg\min_{s_{k}}m_{k}(s_{k}):=g^{T}_{k}s_{k}+\frac{1}{2}s^{T}_{k}\widetilde{H}_{k}s_{k},\quad\mbox{s.t.
}\|s_{k}\|\leq\Delta_{k},$ (3)

where the size of the trust region $\Delta_{k}$ is dynamically adjusted
according to the ratio between the reduction in the objective function and the
model approximation, i.e.

$\rho_{k}:=\frac{f(\Theta_{k})-f(\Theta_{k}+s_{k})}{m_{k}(0)-m_{k}(s_{k})}.$
(4)

Usually, the symmetric indefinite linear systems of equations are solved with
a (preconditioned) conjugate gradient (CG) algorithm, which is halted whenever
a negative eigenvalue (i.e., an ascent direction) is encountered. Such
stopping criterion can be evaluated based on the recurrence formulas of CG,
and it can be considered as an extra criterion to prevent over-solving. A
simple modification of the CG algorithm, known as the Steihaug-Toint variant
[65, 67], can solve Eq. (3) approximatively.

A positive semi-definite approximation of the Hessian is also possible [63,
47], since Eq. (1) admits a generalized Gauss-Newton (GN) approximation,
defined as (dropping the $k$ subscript)

$\widetilde{H}_{GN}:=J^{T}H_{\mathcal{L}}J\,\quad J:=\frac{\partial
y}{\partial\Theta}\,\quad
H_{\mathcal{L}}:=\frac{1}{N}\frac{\partial^{2}\mathcal{L}}{\partial y^{2}},$
(5)

which is routinely used to solve nonlinear least-squares problems. Such
approximation results in a positive semi-definite linear operator provided the
loss function is convex. In addition, the memory requirements of the Gauss-
Newton approximation are similar in magnitude to those of the gradients.

##  3 Software architecture

PETScML 111Repository will be made public before publication is a lightweight
Python interface built on top of petsc4py [17], the Python wrapper for PETSc
[3], which is an award-winning software library widely used by the
computational science and engineering community to solve large-scale, time-
dependent, nonlinear equations and optimization problems (via the Toolkit for
Advanced Optimization, TAO). PETSc is a performance portable framework that
uses the Message Passing Interface (MPI) for distributed memory computations
[76] and interfaces with different computational backends, including NVIDIA,
AMD, and Intel GPUs [48]. All the solver configurations available are command-
line customizable and require a small amount of custom user code.

At the highest level, PETScML defines an abstract class that exposes the basic
methods needed by the optimization solvers:

⬇

from MPI import Comm

class function:

def __init__(self, params : Any, comm : Comm, device : str):

”’ Perform the needed initialization.”’

…

def objective(self, params : Any) -> float:

”’ Return the objective value evaluated at params”’

…

def gradient(self, params : Any) -> Vec:

”’ Return the gradient evaluated at params”’

…

def hessian(self, params : Any, gn : bool = False) -> Mat:

”’ Return the Hessian operator evaluated at params,

possibly using the Gauss-Newton approximation.

Uses hessianMult or hessianMultGN”’

…

def hessianMult(self, params: Any, input_vec : Any) -> Any:

”’ Return H(params) @ input_vec. ”’

…

def hessianMultGN(self, params: Any, input_vec : Any) -> Any:

”’ Return GN approximation H(params) @ input_vec. ”’

…

The derived function classes implement these methods for each specific deep-
learning backend, which does not need to be aware of PETSc-specific details.
Parameter values can be in any form accepted by the backends: PETScML
internally uses routines to flatten parameter values into contiguous storage
for vectors and, vice-versa, to fold vector data into the parameters format
needed by the backend. The derived class must only implement the objective,
gradient, and Hessian matrix-vector multiplication using their input data
format for the linearization point and the input vector. Such an approach
allows the separation of the public interface from the internal
implementation.

Distributed memory is supported by specifying the relevant MPI communicator
when initializing the function instance; device configurations supported are
cpu and gpu. NCCL, the NVIDIA collective communication library222Available at
developer.nvidia.com/nccl is internally used for multi-GPU communications,
while the exchange of information via contiguous storage memory is handled
using the no-copy protocol from dlpack333See dmlc.github.io/dlpack/latest/.
The currently supported backends are JAX and PyTorch. Extensions to additional
computational backends, like AMD or Intel GPUs, or Python ML backends can be
easily accomodated.

Models of the type given in (1) are derived from the function class:

⬇

def __init__(self,

params: Any, # parameters

y : Callable, # model: y(params, x)

loss : Callable, # loss(y(params,x), y)

train_loader : Iterable, # training dataset

**kwargs):

…

def update(self, step : int) -> None:

”’Perform the needed actions on the model

and the dataset at the beginning of the step”’

…

These classes automatically compute objective function, gradient, and Hessian
matrix-vector products using the user-provided model y, loss function, and
training dataset train_loader. The update callback is used to load the
subsequent batches and perform the needed model updates at the beginning of a
new optimization step. The loader can be any iterable Python type, such as a
PyTorch DataLoader; distributed loader classes are also provided to handle
multi-device, data-parallel training. Supported neural network module
libraries for the JAX backend include DM-Haiku444Available at dm-
haiku.readthedocs.io/en/latest/ and Flax555Available at
flax.readthedocs.io/en/latest/.

Optimization solvers can be created using the one-liner

⬇

cfg : dict = {},

params : Any = None,

monitor : Callable[[Any, int, float], None] = None,

prefix : str = None)

which takes as input one instance of the function to be minimized and optional
keyword arguments for the solver configuration dictionary cfg, the initial
values for the parameters params, the convergence monitoring routine monitor,
and a string identifying the solver for the command-line customization prefix.
For experimentation purposes, first-order stochastic methods from
OPTAX666Available at optax.readthedocs.io/en/latest/ are wrapped as custom
PETSc solvers; they can be programmatically customized using the configuration
dictionary cfg or by command-line options.

Table 1: Software Library | Version | Link
---|---|---
PETSc | 3.21 | gitlab.com/petsc/petsc
CUDA | 11.7 | developer.nvidia.com/cuda-toolkit
JAX | 0.4.8 | github.com/google/jax
PyTorch | 2.1.0 | github.com/pytorch/pytorch
DeepXDE | 1.9.3 | github.com/lululxvi/deepxde
OPTAX | 0.1.5 | github.com/google-deepmind/optax
DM-Haiku | 0.0.9 | github.com/google-deepmind/dm-haiku

##  4 Numerical results

In this Section, we evaluate the performances of the L-BFGS, trust region, and
inexact Newton with line-search solvers on a series of test cases taken from
the recent literature on SciML techniques for constructing surrogates to be
used in inverse problems governed by PDEs; the solvers’ performances are
always compared against the same adaptive first-order solvers used to validate
the results, exactly matching the configurations reported in the literature.
To keep the focus on solvers’ performances, we describe each of the SciML
techniques in their own subsection; specific first-order solvers are always
referred to as “reference” solvers, and their performances are reported with a
“REF” label in the figures. For trust region and inexact Newton, we either
consider an exact Hessian or its Gauss-Newton approximation (GN), as given in
(5); we remark that these matrices are never computed explicitly, but we only
evaluate their action on given input vectors.

Convergence plots are reported for loss function values (i.e. the training
error) and case-specific metrics (i.e. the generalization error) on the
testing datasets; the minimum value attained throughout the minimization
process is always shown in the legends. As a qualitative comparison, the data
is reported as a function of the epoch number to describe the convergence
histories of the various solvers. We also report the same data against the
number of “oracle calls” performed by each solver as an architecture-agnostic
measure of their costs, based on the number of passes on the computational
graph needed to perform the computations. In particular, an objective function
evaluation costs $1$ oracle call (one forward pass), a gradient costs $2$
oracle calls (one forward and one backward pass), while an exact Hessian
matrix-vector product costs $4$ calls since we need to differentiate the
gradient operation. In contrast, the Hessian matrix-vector product of the
Gauss-Newton approximation (5) requires $2$ oracle calls; one forward pass to
compute the action of $J$ and one backward pass to compute the action of
$J^{T}$, while the computations associated with the Hessian of the loss
function are negligible. These costs are always rescaled, accounting for the
relative size of the mini-batch with respect to the whole dataset.

We stress that these are only rough, memory-bandwith based estimates of the
relative costs of the gradient and the Hessian matrix-vector products; the
actual computational costs depend on many factors including the computing
architecture, the arithmetic intensity of these operations, the implementation
of automatic differentiation, the mini-batch size used, the input pipeline’s
stochasticity, and the computational graph’s complexity. Given the high-degree
of heterogeneity of the various testcases considered (network architectures,
dataset sizes, and ML frameworks), in what follows we will not discuss
computational timings, code performances or scalability aspects. For the
interested reader, we can say that training times are always strongly
correlated with the oracle calls measurements in all the cases studied.

All solvers are initialized with the same set of random model parameters,
separately for each test case. Convergence histories are reported on a log-log
scale to facilitate cross-comparisons among solvers and test cases; to this
end, we also keep fixed the hyper-parameters of the PETScML solvers for all
the numerical experiments. In particular, we always use a simple, Armijo-type
backtracking line-search for the L-BFGS, LS, and LS-GN solvers. The number of
history vectors stored for the L-BFGS approximation is $30$. The dynamic
computation of the relative tolerance $\nu_{k}$ for LS and LS-GN in Eq. (2) is
computed using the procedure outlined in [1] using a initial tolerance of
$\nu_{0}=0.9$. We also experimented with more classical variants [22], but we
found the one proposed in [1] to be more robust. We use standard parameters
for the trust region methods; we accept a step when $\rho_{k}$ given in (4) is
larger than $0.001$. The trust region’s size, initially set to
$\delta_{0}=0.2$, is shrunk by a factor $0.25$ if $\rho_{k}<0.25$ while it is
doubled if $\rho_{k}>0.75$. The maximum allowed trust region size is 10. All
linear systems are solved using the CG method for a maximum of $100$
iterations, using an inexpensive preconditioner constructed with an L-BFGS
update with $5$ history vectors [5]. The minimization procedure is halted
whenever a line search fails, the size of the trust region can no longer be
reduced further, or the prescribed maximum number of iterations is reached.

The relevant software libraries used for the numerical results are reported in
Table 1. The script to reproduce the experiments, figures, and the driver
applications for each specific test are public. In all the test cases
reproduced herein, we follow exactly the experimental settings used in the
literature; unless otherwise stated, the computations are performed in double
precision floating point arithmetic and full-batch size configuration.
Numerical results have been obtained using A100 NVIDIA GPUs with 80 GB of
device memory. The first-order solvers are always run using the OPTAX/JAX
framework.

###  4.1 Fourier Neural Operator

The Fourier Neural Operator (FNO) framework aims to learn mappings between
infinite dimensional spaces using a finite collection of input/output pairs
observations. More specifically, given two Banach spaces
$\mathcal{A},\mathcal{U}$ and a map $\mathcal{G}:\mathcal{A}\to\mathcal{U}$,
the goal is to build an approximation $\mathcal{G}_{\Theta}$ using
observations drawn from some probability measure supported on $\mathcal{A}$
using an empirical cost (or loss) function. Here $\mathcal{G}_{\Theta}$ is a
neural network approximation of the form

$\mathcal{G}_{\Theta}:=\mathcal{Q}\circ\sigma_{L}(\mathcal{K}_{L})\circ\dots\circ\sigma_{1}(\mathcal{K}_{1})\circ\mathcal{P}$

where $\mathcal{Q}$ and $\mathcal{P}$ are the “projection” and “lifting”
operators, $\sigma_{T}$ an activation function acting component-wise on the
output of some kernel operator $\mathcal{K}_{T}$. For additional information,
more rigorous definitions of functional settings, and different choices of the
kernel operator, see [34]. For error analysis, see [33].

####  4.1.1 Burgers’ equations

As a first test case, we consider the one-dimensional Burgers’ equation test
proposed in Section 5.1 of [38], where we want to learn the operator mapping
the initial conditions to the final conditions at time $T=1$ of

$\frac{\partial u}{\partial t}+u~{}\frac{\partial u}{\partial
x}=\nu\frac{\partial^{2}u}{\partial x^{2}},\quad x\in(0,1)$

with $\nu=0.1$ the viscosity and with periodic boundary conditions. We use the
same datasets as in [38], where the initial conditions are generated as
Gaussian random fields while the final conditions are computed by an implicit
method on a fine grid of size $8192$. Following [44], the data generated is
then subsampled on a much coarser grid of size $128$; the training dataset
consists of $1000$ pairs of random initial conditions and the associated final
solution, while the testing dataset contains $100$ of such pairs. For the
numerical results, we used the data provided by the authors777Available at
github.com/neuraloperator/neuraloperator/tree/master. The network has an
initial dense layer of $64$ features, followed by 4 Fourier layers using $16$
modes with a width equal to $64$, and a final dense layer with $128$ features,
for a total of $549,569$ parameters. The activation function is the Gaussian
Error Linear Units (GELU) function [27]. The network is implemented using DM-
Haiku and the JAX backend; all parameters are stored as real numbers, and
complex arithmetic in the Fourier layers is emulated in software.

The loss function is the mean squared $l2$ relative error

$\mathcal{L}(y,\bar{y})=\frac{1}{N}\sum^{N}_{i=1}\frac{\|y_{i}-\bar{y}_{i}\|^{2}}{\|\bar{y}_{i}\|^{2}},$
(6)

while the accuracy metric is the mean relative $l2$ error

$\frac{1}{N}\sum^{N}_{i=1}\frac{\|y_{i}-\bar{y}_{i}\|}{\|\bar{y}_{i}\|}.$ (7)

The solver used in [38] is ADAMW with weight decay $0.0001$, mini-batch size
20, and initial learning rate $0.001$, halved every 100 epochs using an
exponential staircase schedule. Numerical results are reported in Fig. 1.

The solvers tested exhibit different convergence histories regarding loss
function values; the testing metric is always strongly correlated to the loss
value, confirming that no overfitting is happening. The reported test metric
value for the reference solver is $0.0014$, one order of magnitude smaller
than the one reported in [44] but in line with [34] (see Table 3 in the
reference). The reference solver quickly reduces the loss value and the
testing metric before entering a second phase of slow and erratic convergence
to the identified local minimum, a typical behavior of stochastic first-order
solvers. On the other hand, all the other solvers exhibit an initial plateau
phase before converging to a local minimum at different convergence rates
regarding the number of epochs, with L-BFGS and the exact Hessian variants of
TR and LS experiencing longer plateaus. Despite their different convergence
histories, they seem to converge to local minima with similar generalization
errors, always at least one order of magnitude smaller than the reference
solver, with TR-GN improving the generalization error of the reference solver
by almost two orders of magnitude. Analyzing the converge histories from the
point of view of oracle calls, all the solvers have a similar cost, with the
“cheapest” solver being TR-GN.

Figure 1: FNO Burgers’ test case. Convergence histories for loss function
values and testing metrics in terms of epochs (top row, panels A and B) and
oracle calls (bottom row, panels C and D) for different solvers. The values in
the legends denote the minimum metric value achieved. See Section 4.1.1 for
details.

We then consider a “hybrid” solver configuration, where the solvers are
initialized using the parameters obtained by first performing an initial
number of epochs using the reference solver; results are summarized in Figure
2. Performing a small number of epochs can allow further improvement in the
generalization error in some cases, especially for LS and LS-GN; however,
worse generalization performances can be observed if too many epochs are
performed with the reference solver.

Figure 2: FNO Burgers’ test case with a hybrid solver. Convergence histories
for testing metrics starting from checkpointed solutions of the reference
solver. The values in the legends denote the checkpointed epoch. See Section
4.1.1 for details.

####  4.1.2 Navier-Stokes equations

In the second FNO test case, we consider the two-dimensional Navier-Stokes
equations for a viscous, incompressible fluid in streamfunction-vorticity form
as in Section 6.4 of [34],

$\displaystyle\frac{\partial u}{\partial t}+\nabla^{\perp}\psi\cdot\nabla u$
$\displaystyle=\nu\Delta u+f,$ (8) $\displaystyle-\Delta\psi$
$\displaystyle=u,$

where $\nabla^{\perp}=(-\partial_{y},\partial_{x})$ is the rotated gradient
and $\nu$ the viscosity. $u$ and $\psi$ are the vorticity and the stream
function of the fluid velocity field. The spatial domain is $(0,1)^{2}$; the
system is closed with periodic boundary conditions and with a given initial
condition $u(x,0)=u_{0}(x)$. Denoting by $\Psi$ the solution operator of Eq.
(8), here we are interested in learning the map

$\mathcal{G}:\Psi(u_{0},t)|_{t\in[0,10]}\to\Psi(u_{0},t)_{t\in[10,T]},$

for a certain random distribution of the initial conditions $u_{0}$ and a
given final time $T>10$. The datasets are created by solving Eq. (8) with
$\nu=0.001$ for Gaussian random initial conditions on a space-time grid
$256\times 256\times 50000$ up to $T=50$ with a pseudo-spectral method,
details are given in [38]. In our experiments, we used the same data used by
the authors888Available at drive.google.com/drive/folders/1UnbQh2WWc6knEHbLn-
ZaXrKUZhp7pjt-, which is further subsampled to coarser grids: $64\times
64\times 10$ for the input, and $64\times 64\times 40$ for the output. We use
$1000$ input pairs for training and $200$ pairs as testing datasets to
reproduce the experimental setting used to produce the results in the first
column of Table 4 in [34]. The network architecture is similar to the one used
in Section 4.1.1 using an additional normalization step for the input data.
The output of each of the four Fourier layers is further processed by dense
layers and summed to the output of a convolutional layer. The total number of
parameters is $6,558,537$; the network is implemented reusing the original
PyTorch code999Available at
github.com/neuraloperator/neuraloperator/blob/master/fourier_3d.py.

As before, the loss function is the mean squared $l2$ relative error as given
in (6), while the accuracy metric is the mean relative $l2$ error as given in
(7). The stochastic first-order solver used is the same as for the Burgers’
test case, except that here the mini-batch size is 10. Numerical results are
reported in Fig. 3. We do not report the results for the exact Hessian
variants of LS and TR since our PyTorch implementation of the matrix-vector
product ran out of memory.

Figure 3: FNO Navier-Stokes test case. Convergence histories for loss function
values and testing metrics in terms of epochs (top row, panels A and B) and
oracle calls (bottom row, panels C and D) for different solvers. The values in
the legends denote the minimum metric value achieved. See Section 4.1.2 for
details.

The reference solver shows a similar convergence history as for the Burgers’
test case, reaching a final accuracy slightly larger than the one reported in
[34] (0.0099 vs. 0.0089); the other solvers instead converge to local minima
characterized by smaller generalization errors, with TR-GN being able to
further half the accuracy reported in [34]. Overfitting does not happen except
for the LS-GN solver in the final stages of the optimization. The costs of the
solvers are comparable in terms of oracle calls, with TR-GN being the
cheapest.

###  4.2 DeepONet

Like the FNO technique briefly described in Section 4.1, DeepONet [43] tries
to learn maps from input and output spaces. However, differently from FNO, the
output of the operator maps an input function $u$ sampled at $m$ given
“sensors” locations $\\{x_{i}\\}^{m}_{i=1}$ that do not need to be distributed
on equispaced grids. The final goal is to approximate the map as

$\mathcal{G}_{\Theta}(u)(y):=\mathcal{B}(u(x_{1}),\dots,u(x_{m}))\mathcal{T}(y)$

where $\mathcal{B}$ and $\mathcal{T}$ are the “branch” and “trunk” networks
and $y$ any point in the domain of $\mathcal{G}(u)$, not necessarily
distributed at sensors locations. For additional details, see [43]; for
rigorous definitions, error estimates, and convergence rates, see [35, 45]. A
publicly available implementation of these techniques is provided in the
Python package DeepXDE101010Available at github.com/lululxvi/deepxde.

####  4.2.1 Reaction diffusion equations

As a first test case, we consider the reaction-diffusion test case described
as Problem 4 in [43], where the goal is to learn an implicit operator
associated with the one-dimensional nonlinear reaction-diffusion PDE

$\frac{\partial u}{\partial t}=D\frac{\partial^{2}u}{\partial
x^{2}}+ku^{2}+f,\quad x\in(0,1)$

with homogeneous boundary conditions and zero initial conditions. $D$ is the
diffusion coefficient while $k$ the reaction rate, both set to $0.01$ in the
tests. Here, DeepONet is used to learn the operator mapping the forcing term
$f(x)$ to the PDE solution $u(x,t)$ at $m$ random sensor locations. The
DeepONet architecture uses branch and trunk feed-forward networks of depth
$3$, width $100$, and relu [49] activation functions for a total of $40,702$
parameters. The network is implemented using the PyTorch backend and DeepXDE
callbacks.

The data to be learned is generated using $N$ random forcing terms on a space-
time grid $100\times 100$; for each simulation, $P$ random points (different
from the sensor locations) are selected. Here, we consider the configuration
$N=1000$ for training, $N=10000$ for testing, with $P=100$ and $m=100$, which
corresponds to the last data-point of Figure 4b in [43]. We note that for this
testing configuration, we were not able to find the data used to validate the
original results; instead, we regenerate the data using the scripts provided
by the authors111111Available at github.com/lululxvi/deeponet. See the
supplementary material in [43] for additional details. The loss function and
the accuracy metric are the usual mean squared error

$\frac{1}{N}\sum^{N}_{i=1}\|y_{i}-\bar{y}_{i}\|^{2}.$ (9)

The solver used in [43] is ADAM with a constant learning rate of $0.001$.
Numerical results are reported in Fig. 4.

Figure 4: DeepONet reaction-diffusion test case. Convergence histories for
loss function values and testing metrics in terms of epochs (top row, panels A
and B) and oracle calls (bottom row, panels C and D) for different solvers.
The values in the legends denote the minimum metric value achieved. See
Section 4.2.1 for details.

The reference solver, which is used deterministically here, converges more
slowly than in the FNO test cases to a loss function value of $2.7$e-6, one
order of magnitude smaller than the one reported in literature [43]; the
achieved metric value is similar. Train and test errors for the reference
solver are characterized by strong oscillations as the local minimum is
approached. The L-BFGS and LS solvers converge to loss and metric values
similar to the reference solver, while the TR solver and the Gauss-Newton
variants LS-GN and TR-GN are able to reduce the generalization errors further.
LS-GN exhibits a faster convergence rate in the early stages, followed by a
phase of staircase convergence, while instead, the converge rate of TR-GN is
more regular, leading to almost an order of magnitude improvement over the
generalization error of the reference solver. The computational costs of the
solvers are similar, and they all require around $10^{5}$ oracle calls; like
for the FNO test cases, the cheapest solver is TR-GN.

Figure 5: DeepONet reaction-diffusion oracle calls breakdown. Test metric
(black, left-most y-axis), objective function evaluations (shaded blue, right-
most y-axis), and number of Hessian matrix-vector products (shaded red, right-
most y-axis) against epoch number for different solvers (shown on top). See
Section 4.2.1 for details.

For this test case, we also provide a breakdown of the number of objective
function evaluations and the Hessian matrix-vector products needed per epoch
by the different solvers in Figure 5. Analyzing the number of objective
function evaluations, L-BFGS on average requires 3 to 5 objective function
calls per epoch, whereas LS-GN requires fewer. On the other hand, LS requires
many more objective evaluations; this indicates a struggle of the line-search
procedure in identifying optimal descent, which leads to smaller update steps.
The trust region solvers instead utilize the objective function to decide on
the acceptance of the model solution and eventually shrink the trust region;
the epochs in which the objective function is called more than once indicate
the shrinking of the trust region. On this aspect, the TR-GN solver is much
more effective and requires the smallest number of objective function calls.
The major cost of the second-order solvers is in the number of Hessian matrix-
vector multiplications, and both the trust region variants are very effective
in the early stage of the optimization process, with the number of matrix-
vector products progressively ramping up to their limit (100) and then
reducing in the final stages of the solvers. The dynamic selection of the
stopping criterion in the LS and LS-GN solver instead is not effective in
controlling the number of linear iterations; the detection of negative
curvature further reduces the number of linear iterations in the early stages
of the LS and TR solvers.

####  4.2.2 Advection equations

We then consider a variant of DeepONet where the trunk network is replaced by
a precomputed basis for the output functions using the Proper Orthogonal
Decomposition (POD). As a test case, we reproduce the testing configuration
used in Section 5.4.1 in [44]; specifically, we are interested in the one-
dimensional advection equation

$\frac{\partial u}{\partial t}+\frac{\partial u}{\partial x}=0,\quad
x\in(0,1)$

with periodic boundary conditions and learn the mapping from the randomly
generated initial conditions $u_{0}(x)$ to the solution on the space-time
domain $u(x,t)$, with $(x,t)\in[0,1]^{2}$. The data is generated by solving
the advection equation on a space-time grid of size $40\times 40$ and
considering two families of randomly generated initial conditions

Case I: $\displaystyle\,u_{0}(x)=h~{}\mathds{1}_{[c-w/2,c+w/2]},$ Case II:
$\displaystyle\,u_{0}(x)=h~{}\mathds{1}_{[c-w/2,c+w/2]}+\sqrt{\mbox{max}(d^{2}-a^{2}(x-b)^{2},0)},$

where $a,b,c,d,w,h$ are randomly selected (each in a separate interval) and
$\mathds{1}$ is the indicator function. In both cases, the train and test
datasets contain $1000$ pairs of random initialization and the corresponding
space-time solution. We use the data provided by the authors121212Available at
github.com/lu-group/deeponet-fno as training and testing datasets for our
computations; we also note that cases I/II here correspond to cases II/III in
[44] to avoid confusion. The branch network is a feed-forward network with two
layers of widths $512$ and $n$, where $n$ is the number of POD basis
functions: for case I, $n=38$, while for case II, $n=32$. The total number of
parameters is $40,486$ for case I and $37,408$ for case II. The networks are
implemented using the PyTorch backend and DeepXDE callbacks.

The loss function is the mean squared error (9), and the mean squared $l2$
relative error given in (6) is used as the accuracy metric. The solver used in
[44] is ADAM with an initial learning rate of $0.001$, asymptotically decayed
to $0.0001$ as the iteration progresses. The computations are performed in
single precision arithmetic as done in the accompanying software of [44];
numerical results are reported in Fig. 6 for case I and in Fig. 7 for case II.

Figure 6: DeepONet-POD test case I. Convergence histories for loss function
values and testing metric in terms of epochs (top row, panels A and B) and
oracle calls (bottom row, panels C and D) for different solvers. The values in
the legends denote the minimum metric value achieved. See Section 4.2.2 for
details.

We first analyze the solvers’ convergence histories for case I. The reference
solver exhibits a similar convergence history as the one in Section 4.2.1; its
generalization error is half of the one reported in Table 8 (Advection II
column) in [44]. All the other solvers show faster convergence to different
local minima characterized by smaller generalization errors, with the trust
region variants achieving the smallest errors while still being faster than
the other solvers in terms of oracle calls.

Figure 7: DeepONet-POD test case II. Convergence histories for loss function
values and testing metrics in terms of epochs (top row, panels A and B) and
oracle calls (bottom row, panels C and D) for different solvers. The values in
the legends denote the minimum metric value achieved. See Section 4.2.2 for
details.

For test case II, all solvers converge to local minima possessing the same
generalization capabilities. The error achieved by the reference solver is
$0.0039$, perfectly aligned with the one reported in [44] (see Table 8,
Advection III column in the reference). The major difference between solvers
is in the number of oracle calls required to converge, with both the Gauss-
Newton variants requiring a smaller number of calls.

###  4.3 GreenLearning

GreenLearning seeks to discover Green’s functions associated with unknown
linear operators by collecting physical system responses to random excitation
inputs drawn from Gaussian processes [7]. In a more formal setting, the goal
is to learn the operator map by approximating the Green’s function behind a
physical process encoded by a linear operator

$L~{}u(x)=f(x),\quad x\in\Omega,$

from a collection of input/output pairs $\\{f_{i},u_{i}\\}$; the idea is to
construct neural networks $\mathcal{N}_{G}$ and $\mathcal{N}_{h}$ such that

$u_{i}(x)\approx\int_{\Omega}\mathcal{N}_{G}(x,y)f_{i}(y)dy+\mathcal{N}_{h}(x)$

where $\mathcal{N}_{G}$ approximates the Green’s function and
$\mathcal{N}_{h}$ a given homogeneous solution encoding the boundary
conditions.

In this Section, we reproduce the experimental setting described in [8] and
consider the two-dimensional Poisson’s equation on the unit square

$-\Delta u=f$

with zero boundary conditions, and learn the operator that maps forcing terms
$f(x,t)$ to the solution of the PDE. The data is generated by solving the
Poisson equation on a $421\times 421$ grid with different forcing terms
generated as Gaussian random fields and then subsampled on coarser grids of
size $29\times 29$. The training dataset used here contains $50$ elements,
while the test dataset consists of $200$ elements, and they have been obtained
from the data used in [8]131313Available at zenodo.org/records/7701683. The
datasets are first normalized before being fed to the $\mathcal{N}_{G}$
network, characterized by a feed-forward architecture with depth $4$ and width
$50$ that uses rational activation functions [9]. The current approximation of
the Green’s function, $\mathcal{N}_{G}$, is first sampled on a regular grid;
the output obtained is then multiplied with the network input (i.e., an
instance of the randomly generated forcing terms) mimicking the integration of
Green’s function. The total number of model parameters is $7,979$. The network
is implemented using the PyTorch backend.

The loss function and efficiency metric are the mean squared $l2$ relative
error given in (6). The solver used in [8] is ADAMW with weight decay $0.0001$
and initial learning rate $0.001$, scaled by $0.9$ every 100 epochs using an
exponential staircase schedule. The computations are performed in single
precision arithmetic as done in the accompanying software of [8], and
numerical results are reported in Fig. 8; see Figure 1 in [8] for the results
with different training dataset sizes.

Figure 8: GreenLearning test case. Convergence histories for loss function
values and testing metrics in terms of epochs (top row, panels A and B) and
oracle calls (bottom row, panels C and D) for different solvers. The values in
the legends denote the minimum metric value achieved. See Section 4.3 for
details.

As for the other test cases reported in this Section, train and test metrics
are strongly correlated, and no overfitting is observed. The reference solver
shows a very long plateau before converging with minimal oscillations; the
achieved generalization error is aligned with the one reported in [8]. All the
other solvers except for LS appear to converge to similar local minima
characterized by generalization errors smaller than the reference case and
with approximatively the same computational costs. LS never escapes the
initial plateau phase and fails to converge.

##  5 Conclusions

In this paper, we have presented PETScML, a lightweight Python interface on
top of PETSc, designed for solver experimentation in training deep-learning
models. Using PETScML, we have empirically demonstrated the efficacy of
conventional second-order solvers in improving the generalization errors
arising from regression tasks when learning surrogate models, and with smaller
computational costs.

Among the various solvers tested, the trust region method with Gauss-Newton
approximation of the Hessian proved to be the overall best-performing method
for the test cases analyzed, either in terms of final model accuracy or in
terms of computational costs.

Much remains to follow up on these encouraging results, in particular on the
design of linear preconditioning techniques for limiting the number of Krylov
iterations, and on the design of nonlinear preconditioning strategies [14, 10,
40] to limit the number of nonlinear iterations and improve the overall
convergence rate of the solver. The convergence properties of the trust region
method can also be further improved by using different trust region norms and
ellipsoidal scaling techniques. These topics will be the subject of future
work.

Other SciML techniques based on regression tasks can potentially benefit from
the adoption of second-order solvers; possible future directions in this
context include molecular dynamics simulations [69, 18] and continuous
heterogeneous cryo-EM reconstructions [77].

Future work will also focus on the application of second-order solvers to more
conventional deep-learning tasks like image classification [29], sequence
modeling [30], vision tasks [20], and self-supervised learning [61], where we
plan to study the ability of second-order methods to reach state-of-the-art
validation accuracies with minor hyper-parameter tuning, ideally in less wall-
clock time or energy costs.

## Acknowledgments

Umberto Zerbinati acknowledges the support of KAUST, since part of the
research was conducted when he was a student in MSc program of the CEMSE
division. The authors acknowledge the use of the Ibex GPU Cluster of the KAUST
Supercomputing Laboratory and support from the Extreme Computing Research
Center and the AI Initiative at KAUST.

## References

  * [1] Heng-Bin An, Ze-Yao Mo, and Xing-Ping Liu.  A choice of forcing terms in inexact Newton method.  Journal of Computational and Applied Mathematics, 200(1):47–60, 2007.
  * [2] Rohan Anil, Vineet Gupta, Tomer Koren, Kevin Regan, and Yoram Singer.  Second order optimization made practical.  CoRR, abs/2002.09018, 2020.
  * [3] Satish Balay, Shrirang Abhyankar, Mark F. Adams, Steven Benson, Jed Brown, Peter Brune, Kris Buschelman, Emil Constantinescu, Lisandro Dalcin, Alp Dener, Victor Eijkhout, Jacob Faibussowitsch, William D. Gropp, Václav Hapla, Tobin Isaac, Pierre Jolivet, Dmitry Karpeev, Dinesh Kaushik, Matthew G. Knepley, Fande Kong, Scott Kruger, Dave A. May, Lois Curfman McInnes, Richard Tran Mills, Lawrence Mitchell, Todd Munson, Jose E. Roman, Karl Rupp, Patrick Sanan, Jason Sarich, Barry F. Smith, Stefano Zampini, Hong Zhang, Hong Zhang, and Junchao Zhang.  PETSc/TAO users manual.  Technical Report ANL-21/39 - Revision 3.20, Argonne National Laboratory, 2023.
  * [4] Satish Balay, Shrirang Abhyankar, Mark F. Adams, Steven Benson, Jed Brown, Peter Brune, Kris Buschelman, Emil M. Constantinescu, Lisandro Dalcin, Alp Dener, Victor Eijkhout, Jacob Faibussowitsch, William D. Gropp, Václav Hapla, Tobin Isaac, Pierre Jolivet, Dmitry Karpeev, Dinesh Kaushik, Matthew G. Knepley, Fande Kong, Scott Kruger, Dave A. May, Lois Curfman McInnes, Richard Tran Mills, Lawrence Mitchell, Todd Munson, Jose E. Roman, Karl Rupp, Patrick Sanan, Jason Sarich, Barry F. Smith, Stefano Zampini, Hong Zhang, Hong Zhang, and Junchao Zhang.  PETSc Web page, 2023.
  * [5] Luca Bergamaschi, José Marín, and Ángeles Martínez.  Compact quasi-Newton preconditioners for symmetric positive definite linear systems.  Numerical Linear Algebra with Applications, 27(6):e2322, 2020.
  * [6] Léon Bottou, Frank E. Curtis, and Jorge Nocedal.  Optimization methods for large-scale machine learning.  SIAM Review, 60(2):223–311, 2018.
  * [7] Nicolas Boullé, Christopher J Earls, and Alex Townsend.  Data-driven discovery of Green’s functions with human-understandable deep learning.  Scientific reports, 12(1):4824, 2022.
  * [8] Nicolas Boullé, Diana Halikias, and Alex Townsend.  Elliptic PDE learning is provably data-efficient.  Proceedings of the National Academy of Sciences, 120(39):e2303904120, 2023.
  * [9] Nicolas Boullé, Yuji Nakatsukasa, and Alex Townsend.  Rational neural networks.  Advances in neural information processing systems, 33:14243–14253, 2020.
  * [10] Peter R Brune, Matthew G Knepley, Barry F Smith, and Xuemin Tu.  Composing scalable nonlinear algebraic solvers.  SIAM Review, 57(4):535–565, 2015.
  * [11] Steven L Brunton, Joshua L Proctor, and J Nathan Kutz.  Discovering governing equations from data by sparse identification of nonlinear dynamical systems.  Proceedings of the national academy of sciences, 113(15):3932–3937, 2016.
  * [12] Richard H Byrd, Samantha L Hansen, Jorge Nocedal, and Yoram Singer.  A stochastic quasi-Newton method for large-scale optimization.  SIAM Journal on Optimization, 26(2):1008–1031, 2016.
  * [13] Richard H Byrd, Jorge Nocedal, and Robert B Schnabel.  Representations of quasi-Newton matrices and their use in limited memory methods.  Mathematical Programming, 63(1-3):129–156, 1994.
  * [14] Xiao-Chuan Cai and David E Keyes.  Nonlinearly preconditioned inexact newton algorithms.  SIAM Journal on Scientific Computing, 24(1):183–200, 2002.
  * [15] Andrew R Conn, Nicholas IM Gould, and Philippe L Toint.  Trust region methods.  SIAM, 2000.
  * [16] George Cybenko.  Approximation by superpositions of a sigmoidal function.  Mathematics of control, signals and systems, 2(4):303–314, 1989\.
  * [17] Lisandro D Dalcin, Rodrigo R Paz, Pablo A Kler, and Alejandro Cosimo.  Parallel distributed computing using Python.  Advances in Water Resources, 34(9):1124–1139, 2011.
  * [18] Sambit Das, Bikash Kanungo, Vishal Subramanian, Gourab Panigrahi, Phani Motamarri, David Rogers, Paul Zimmerman, and Vikram Gavini.  Large-scale materials modeling at quantum accuracy: Ab initio simulations of quasicrystals and interacting extended defects in metallic alloys.  In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’23, New York, NY, USA, 2023. Association for Computing Machinery.
  * [19] Maarten V de Hoop, Daniel Zhengyu Huang, Elizabeth Qian, and Andrew M Stuart.  The cost-accuracy trade-off in operator learning with neural networks.  Journal of Machine Learning, 2022.
  * [20] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929, 2020.
  * [21] John Duchi, Elad Hazan, and Yoram Singer.  Adaptive subgradient methods for online learning and stochastic optimization.  Journal of machine learning research, 12(7), 2011.
  * [22] Stanley C Eisenstat and Homer F Walker.  Choosing the forcing terms in an inexact Newton method.  SIAM Journal on Scientific Computing, 17(1):16–32, 1996.
  * [23] Roy Frostig, Matthew James Johnson, and Chris Leary.  Compiling machine learning programs via high-level tracing.  Systems for Machine Learning, 4(9), 2018.
  * [24] Donald Goldfarb, Yi Ren, and Achraf Bahamou.  Practical quasi-newton methods for training deep neural networks.  In Proceedings of the 34th International Conference on Neural Information Processing Systems, NeurIPS, pages 201:1–11, 2020.
  * [25] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He.  Accurate, large minibatch sgd: Training imagenet in 1 hour.  arXiv preprint arXiv:1706.02677, 2017.
  * [26] Vineet Gupta, Tomer Koren, and Yoram Singer.  Shampoo: Preconditioned stochastic tensor optimization.  2018\.
  * [27] Dan Hendrycks and Kevin Gimpel.  Gaussian Error Linear Units (GELUs).  arXiv preprint arXiv:1606.08415, 2016.
  * [28] Kurt Hornik, Maxwell Stinchcombe, and Halbert White.  Multilayer feedforward networks are universal approximators.  Neural networks, 2(5):359–366, 1989.
  * [29] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger.  Densely connected convolutional networks.  In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 4700–4708, 2017.
  * [30] Kazuki Irie, Imanol Schlag, Róbert Csordás, and Jürgen Schmidhuber.  Going beyond linear transformers with recurrent fast weight programmers.  Advances in Neural Information Processing Systems, 34:7703–7717, 2021.
  * [31] Weile Jia, Han Wang, Mohan Chen, Denghui Lu, Lin Lin, Roberto Car, E Weinan, and Linfeng Zhang.  Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning.  In SC20: International conference for high performance computing, networking, storage and analysis, pages 1–14. IEEE, 2020.
  * [32] Diederik P Kingma and Jimmy Ba.  ADAM: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980, 2014.
  * [33] Nikola Kovachki, Samuel Lanthaler, and Siddhartha Mishra.  On universal approximation and error bounds for fourier neural operators.  The Journal of Machine Learning Research, 22(1):13237–13312, 2021\.
  * [34] Nikola B Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew M Stuart, and Anima Anandkumar.  Neural Operator: Learning maps between function spaces with applications to PDEs.  J. Mach. Learn. Res., 24(89):1–97, 2023.
  * [35] Samuel Lanthaler, Siddhartha Mishra, and George E Karniadakis.  Error estimates for DeepONets: a deep learning framework in infinite dimensions.  Transactions of Mathematics and Its Applications, 6(1):tnac001, 03 2022.
  * [36] Yann LeCun, Yoshua Bengio, and Geoffrey Hinton.  Deep learning.  nature, 521(7553):436–444, 2015.
  * [37] Xi-Lin Li.  Preconditioned stochastic gradient descent.  IEEE transactions on neural networks and learning systems, 29(5):1454–1466, 2017.
  * [38] Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Fourier Neural Operator for parametric partial differential equations.  arXiv preprint arXiv:2010.08895, 2020.
  * [39] Dong C Liu and Jorge Nocedal.  On the limited memory bfgs method for large scale optimization.  Mathematical programming, 45(1-3):503–528, 1989.
  * [40] Lulu Liu, Weiguo Gao, Han Yu, and David E Keyes.  Overlapping multiplicative schwarz preconditioning for linear and nonlinear systems.  Journal of Computational Physics, 496:112548, 2024.
  * [41] Yang Liu and Fred Roosta.  Convergence of Newton-MR under inexact Hessian information.  SIAM Journal on Optimization, 31(1):59–90, 2021.
  * [42] Ilya Loshchilov and Frank Hutter.  Decoupled weight decay regularization.  arXiv preprint arXiv:1711.05101, 2017.
  * [43] Lu Lu, Pengzhan Jin, Guofei Pang, Zhongqiang Zhang, and George Em Karniadakis.  Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators.  Nature machine intelligence, 3(3):218–229, 2021.
  * [44] Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta Goswami, Zhongqiang Zhang, and George Em Karniadakis.  A comprehensive and fair comparison of two neural operators (with practical extensions) based on fair data.  Computer Methods in Applied Mechanics and Engineering, 393:114778, 2022.
  * [45] Carlo Marcati and Christoph Schwab.  Exponential convergence of deep operator networks for elliptic partial differential equations.  SIAM Journal on Numerical Analysis, 61(3):1513–1545, 2023.
  * [46] James Martens and Roger Grosse.  Optimizing neural networks with kronecker-factored approximate curvature.  In International conference on machine learning, pages 2408–2417. PMLR, 2015.
  * [47] James Martens and Ilya Sutskever.  Learning recurrent neural networks with Hessian-free optimization.  In Proceedings of the 28th international conference on machine learning (ICML-11), pages 1033–1040, 2011.
  * [48] Richard Tran Mills, Mark F. Adams, Satish Balay, Jed Brown, Alp Dener, Matthew Knepley, Scott E. Kruger, Hannah Morgan, Todd Munson, Karl Rupp, Barry F. Smith, Stefano Zampini, Hong Zhang, and Junchao Zhang.  Toward performance-portable PETSc for GPU-based exascale systems.  Parallel Computing, 108:102831, 2021.
  * [49] Vinod Nair and Geoffrey E Hinton.  Rectified linear units improve restricted boltzmann machines.  In Proceedings of the 27th international conference on machine learning (ICML-10), pages 807–814, 2010.
  * [50] Tenavi Nakamura-Zimmerer, Qi Gong, and Wei Kang.  Adaptive deep learning for high-dimensional hamilton–jacobi–bellman equations.  SIAM Journal on Scientific Computing, 43(2):A1221–A1247, 2021.
  * [51] Jorge Nocedal and Stephen J Wright.  Numerical optimization.  Springer, 2006.
  * [52] Thomas O’Leary-Roseberry, Nick Alger, and Omar Ghattas.  Low rank saddle free newton: A scalable method for stochastic nonconvex optimization, 2021.
  * [53] Thomas O’Leary-Roseberry, Peng Chen, Umberto Villa, and Omar Ghattas.  Derivative-informed neural operator: An efficient framework for high-dimensional parametric derivative learning.  Journal of Computational Physics, 496:112555, 2024.
  * [54] Kazuki Osawa, Satoki Ishikawa, Rio Yokota, Shigang Li, and Torsten Hoefler.  ASDL: A unified interface for gradient preconditioning in pytorch.  arXiv preprint arXiv:2305.04684, 2023.
  * [55] Kazuki Osawa, Yohei Tsuji, Yuichiro Ueno, Akira Naruse, Rio Yokota, and S. Matsuoka.  Second-order optimization method for large mini-batch: Training resnet-50 on imagenet in 35 epochs.  ArXiv, abs/1811.12019, 2018.
  * [56] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala.  Pytorch: An imperative style, high-performance deep learning library.  Advances in neural information processing systems, 32, 2019.
  * [57] J. Gregory Pauloski, Qi Huang, Lei Huang, Shivaram Venkataraman, Kyle Chard, Ian Foster, and Zhao Zhang.  KAISA: An adaptive second-order optimizer framework for deep neural networks.  In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC ’21, New York, NY, USA, 2021. Association for Computing Machinery.
  * [58] Maziar Raissi.  Deep hidden physics models: Deep learning of nonlinear partial differential equations.  The Journal of Machine Learning Research, 19(1):932–955, 2018.
  * [59] Maziar Raissi, Paris Perdikaris, and George E Karniadakis.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.  Journal of Computational physics, 378:686–707, 2019.
  * [60] Farbod Roosta-Khorasani and Michael W Mahoney.  Sub-sampled newton methods.  Mathematical Programming, 174:293–326, 2019.
  * [61] Justyna Sarzynska-Wawer, Aleksander Wawer, Aleksandra Pawlak, Julia Szymanowska, Izabela Stefaniak, Michal Jarkiewicz, and Lukasz Okruszek.  Detecting formal thought disorder by deep contextualized word representations.  Psychiatry Research, 304:114135, 2021.
  * [62] Jürgen Schmidhuber.  Deep learning in neural networks: An overview.  Neural networks, 61:85–117, 2015.
  * [63] Nicol N Schraudolph.  Fast curvature matrix-vector products for second-order gradient descent.  Neural computation, 14(7):1723–1738, 2002.
  * [64] Samuel L Smith, Benoit Dherin, David GT Barrett, and Soham De.  On the origin of implicit regularization in stochastic gradient descent.  9th International Conference on Learning Representations, ICLR 2021, 2021.
  * [65] Trond Steihaug.  The conjugate gradient method and trust regions in large scale optimization.  SIAM Journal on Numerical Analysis, 20(3):626–637, 1983.
  * [66] Ilya Sutskever, James Martens, George Dahl, and Geoffrey Hinton.  On the importance of initialization and momentum in deep learning.  In International conference on machine learning, pages 1139–1147. PMLR, 2013.
  * [67] Philippe Toint.  Towards an efficient sparsity exploiting newton method for minimization.  In Sparse matrices and their uses, pages 57–88. Academic press, 1981.
  * [68] Jan van den Brand, Binghui Peng, Zhao Song, and Omri Weinstein.  Training (Overparametrized) Neural Networks in Near-Linear Time.  In James R. Lee, editor, 12th Innovations in Theoretical Computer Science Conference (ITCS 2021), volume 185 of Leibniz International Proceedings in Informatics (LIPIcs), pages 63:1–63:15, Dagstuhl, Germany, 2021. Schloss Dagstuhl–Leibniz-Zentrum für Informatik.
  * [69] Tongqi Wen, Linfeng Zhang, Han Wang, E Weinan, and David J Srolovitz.  Deep potentials for materials science.  Materials Futures, 1(2):022601, 2022.
  * [70] Peng Xu, Fred Roosta, and Michael W Mahoney.  Newton-type methods for non-convex optimization under inexact Hessian information.  Mathematical Programming, 184(1-2):35–70, 2020.
  * [71] Minghan Yang, Dong Xu, Zaiwen Wen, Mengyun Chen, and Pengxiang Xu.  Sketch-based empirical natural gradient methods for deep learning.  Journal of Scientific Computing, 92(3):94, 2022.
  * [72] Zhewei Yao, Amir Gholami, Sheng Shen, Mustafa Mustafa, Kurt Keutzer, and Michael Mahoney.  AdaHessian: An adaptive second order optimizer for machine learning.  Proceedings of the AAAI Conference on Artificial Intelligence, 35(12):10665–10673, May 2021.
  * [73] Yang You, Igor Gitman, and Boris Ginsburg.  Scaling sgd batch size to 32k for imagenet training.  arXiv preprint arXiv:1708.03888, 6(12):6, 2017.
  * [74] Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.  Large batch optimization for deep learning: training bert in 76 minutes.  2020\.
  * [75] Bing Yu et al.  The deep ritz method: a deep learning-based numerical algorithm for solving variational problems.  Communications in Mathematics and Statistics, 6(1):1–12, 2018.
  * [76] Junchao Zhang, Jed Brown, Satish Balay, Jacob Faibussowitsch, Matthew Knepley, Oana Marin, Richard Tran Mills, Todd Munson, Barry F Smith, and Stefano Zampini.  The PetscSF scalable communication layer.  IEEE Transactions on Parallel and Distributed Systems, 33(4):842–853, 2021.
  * [77] Ellen D Zhong, Tristan Bepler, Bonnie Berger, and Joseph H Davis.  CryoDRGN: reconstruction of heterogeneous cryo-EM structures using neural networks.  Nature methods, 18(2):176–185, 2021.