# Kronecker-Factored Approximate Curvature for Physics-Informed Neural
Networks

Felix Dangel
Vector Institute
108 College St.
Toronto, Canada
<EMAIL_ADDRESS>
&Johannes Müller∗
Chair of Mathematics of Information Processing
RWTH Aachen University
Aachen, Germany
<EMAIL_ADDRESS>
&Marius Zeinhofer∗
Simula Research Laboratory
Oslo, Norway
<EMAIL_ADDRESS>Equal contribution

###### Abstract

Physics-informed neural networks (PINNs) are infamous for being hard to train.
Recently, second-order methods based on natural gradient and Gauss-Newton
methods have shown promising performance, improving the accuracy achieved by
first-order methods by several orders of magnitude. While promising, the
proposed methods only scale to networks with a few thousand parameters due to
the high computational cost to evaluate, store, and invert the curvature
matrix. We propose Kronecker-factored approximate curvature (KFAC) for PINN
losses that greatly reduces the computational cost and allows scaling to much
larger networks. Our approach goes beyond the established KFAC for traditional
deep learning problems as it captures contributions from a PDE’s differential
operator that are crucial for optimization. To establish KFAC for such losses,
we use Taylor-mode automatic differentiation to describe the differential
operator’s computation graph as a forward network with shared weights. This
allows us to apply KFAC thanks to a recently-developed general formulation for
networks with weight sharing. Empirically, we find that our KFAC-based
optimizers are competitive with expensive second-order methods on small
problems, scale more favorably to higher-dimensional neural networks and PDEs,
and consistently outperform first-order methods and LBFGS.

##  1 Introduction

Neural network-based approaches to numerically solve partial differential
equations (PDEs) are growing at an unprecedented speed. The idea to train
network parameters to minimize the residual of a PDE traces back to at least
Dissanayake & Phan-Thien [15], Lagaris et al. [28], but was only recently
popularized under the name _deep Galerkin method_ (DGM) and _Physics-informed
neural networks_ (PINNs) through the works of Sirignano & Spiliopoulos [53],
Raissi et al. [51]. PINNs are arguably one of the most popular network-based
approaches to the numerical solution of PDEs as they are easy to implement,
seamlessly incorporate measurement data, and promise to work well in high
dimensions. Despite their immense popularity, PINNs are notoriously difficult
to optimize [56] and fail to provide satisfactory accuracy when trained with
first-order methods, even for simple problems [63, 41]. Recently, second-order
methods that use the function space geometry to design preconditioners have
shown remarkable promise in addressing the training difficulties of PINNs [63,
41, 14, 24, 42]. However, these methods require solving a linear system in the
network’s high-dimensional parameter space at cubic computational iteration
cost, which prohibits scaling such approaches. To address this, we build on
the idea of Kronecker-factored approximate curvature (KFAC) and apply it to
Gauss-Newton matrices of PINN losses which greatly reduces the computational
cost:

  * •

We use higher-order forward (Taylor) mode automatic differentiation to
interpret the computation graph of a network’s input derivatives as a larger
net with weight sharing (§ 3.1).

  * •

We use this weight sharing view to propose KFAC for Gauss-Newton matrices of
objectives with differential operators, like PINN losses (§ 3.3 and 14).
Thanks to the generality of Taylor-mode and KFAC for weight sharing layers
[18], our approach is widely applicable.

  * •

We show that, for specific differential operators, the weight sharing in
Taylor-mode can be further reduced by absorbing the reduction of partial
derivatives into the forward propagation, producing a more efficient scheme.
For the prominent example of the Laplace operator, this recovers and
generalizes the _forward Laplacian_ framework [29] (§ 3.2 and 9).

  * •

Empirically, we find that our KFAC-based optimizers are competitive with
expensive second-order methods on small problems, scale more favorably to
higher-dimensional neural networks and PDEs, and consistently outperform
first-order methods and LBFGS (§ 4).

Related work Various approaches were developed to improve the optimization of
PINNs such as adaptive re-weighting of loss terms [56, 55, 58], different
sampling strategies for discretizing the loss [34, 43, 13, 62, 57, 60], and
curriculum learning [26, 57]. While LBFGS is known to improve upon first-order
optimizers [35], recently, other second-order methods that design meaningful
preconditioners that respect the problems geometry have significantly
outperformed it [63, 41, 14, 33, 24, 7, 61]. Müller & Zeinhofer [42] provide a
unified view on these approaches which greatly improve the accuracy of PINNs,
but come with a significant per-iteration cost as one needs to solve a linear
system in the network’s high-dimensional parameter space, which is only
feasible for small networks when done naively. One approach is to use matrix-
free methods to approximately compute Gauss-Newton directions by introducing
an inner optimization loop, see [52, 36] for supervised learning problems and
[63, 4, 24, 61] for PINNs. Instead, our KFAC-based approach uses an explicit
structured curvature representation which can be updated over iterations and
inverted more cheaply.

We build on the literature on Kronecker-factored approximate curvature (KFAC),
which was initially introduced in Heskes [23], Martens [36] as an
approximation of the per-layer Fisher matrix to perform approximate natural
gradient descent. Later, KFAC was extended to convolutional [21], recurrent
[39], attention [47, 44, 22], and recently to general linear layers with
weight sharing [18]. These works do not address preconditioners for loss
functions with contributions from differential operators, as is the case for
PINN losses. Kronecker-factored approximations have been used for variational
Monte-Carlo methods, where despite the objective function involving a
Laplacian term, the preconditioner does not [50, 16]. Our interpretation via
Taylor-mode makes the computation graph of such losses explicit, and allows us
to establish KFAC based on its generalization to linear weight sharing layers
[18].

##  2 Background

For simplicity, we present our approach for multi-layer perceptrons (MLPs)
consisting of fully-connected and element-wise activation layers. However, the
generality of Taylor-mode automatic differentiation and KFAC for linear layers
with weight sharing allow our KFAC to be applied to such layers (e.g. fully-
connected, convolution, attention) in arbitrary neural network architectures.

Flattening & Derivatives We vectorize matrices using the _first-index-varies-
fastest_ convention, i.e. column-stacking (row index varies first, column
index varies second) and denote the corresponding flattening operation by
$\operatorname{vec}$. This allows to reduce derivatives of matrix- or tensor-
valued objects back to the vector case by flattening a function’s input and
output before differentiation. The Jacobian of a vector-to-vector function
${\bm{a}}\mapsto{\bm{b}}({\bm{a}})$ has entries
$[\mathrm{J}_{{\bm{a}}}{\bm{b}}]_{i,j}=\nicefrac{{\partial{b}_{i}}}{{\partial{a}_{j}}}$.
For a matrix-to-matrix function ${\bm{A}}\mapsto{\bm{B}}({\bm{A}})$, the
Jacobian is
$\mathrm{J}_{{\bm{A}}}{\bm{B}}=\mathrm{J}_{\operatorname{vec}{\bm{A}}}\operatorname{vec}{\bm{B}}$.
A useful property of $\operatorname{vec}$ is
$\operatorname{vec}({\bm{A}}{\bm{X}}{\bm{B}})=({\bm{B}}^{\top}\otimes{\bm{A}})\operatorname{vec}{{\bm{X}}}$
for matrices ${\bm{A}},{\bm{X}},{\bm{B}}$ which implies
$\mathrm{J}_{\bm{X}}({\bm{A}}{\bm{X}}{\bm{B}})={\bm{B}}^{\top}\otimes{\bm{A}}$.

Sequential neural nets Consider a _sequential neural network_
$u_{{\bm{\theta}}}=f_{{\bm{\theta}}^{(L)}}\circ
f_{{\bm{\theta}}^{(L-1)}}\circ\ldots\circ f_{{\bm{\theta}}^{(1)}}$ of depth
$L\in\mathbb{N}$. It consists of layers
$f_{{\bm{\theta}}^{(l)}}\colon{\mathbb{R}}^{h^{(l-1)}}\to{\mathbb{R}}^{h^{(l)}}$,
${\bm{z}}^{(l-1)}\mapsto{\bm{z}}^{(l)}=f_{{\bm{\theta}}^{(l)}}({\bm{z}}^{(l-1)})$
with trainable parameters ${\bm{\theta}}^{(l)}\in{\mathbb{R}}^{p^{(l)}}$ that
transform an input ${\bm{z}}^{(0)}\coloneqq{\bm{x}}\in\mathbb{R}^{d\coloneqq
h^{(0)}}$ into a prediction
$u_{\bm{\theta}}({\bm{x}})={\bm{z}}^{(L)}\in{\mathbb{R}}^{h^{(L)}}$ via
intermediate representations ${\bm{z}}^{(l)}\in{\mathbb{R}}^{h^{(l)}}$. In the
context of PINNs, we use networks with scalar outputs ($h^{(L)}=1$) and denote
the concatenation of all parameters by
${\bm{\theta}}=({\bm{\theta}}^{(1)\top},\dots,{\bm{\theta}}^{(L)\top})^{\top}\in{\mathbb{R}}^{P}$.
A common choice is to alternate fully-connected and activation layers. Linear
layers map
${\bm{z}}^{(l-1)}\mapsto{\bm{z}}^{(l)}={\bm{W}}^{(l)}{\bm{z}}^{(l-1)}$ using a
weight matrix
${\bm{W}}^{(l)}=\operatorname{vec}^{-1}{\bm{\theta}}^{(l)}\in{\mathbb{R}}^{h^{(l)}\times
h^{(l-1)}}$ (bias terms can be added as an additional column and by appending
a $1$ to the input). Activation layers map
${\bm{z}}^{(l-1)}\mapsto{\bm{z}}^{(l)}=\sigma({\bm{z}}^{(l-1)})$ element-wise
for a (typically smooth) $\sigma\colon\mathbb{R}\to\mathbb{R}$.

###  2.1 Energy Natural Gradients for Physics-Informed Neural Networks

Let us consider a domain $\Omega\subseteq\mathbb{R}^{d}$ and the partial
differential equation

$\displaystyle\mathcal{L}u=f\quad\text{in }\Omega\,,\qquad u=g\quad\text{on
}\partial\Omega\,,$

with right-hand side $f$, boundary data $g$ and a differential operator
$\mathcal{L}$, e.g. the negative Laplacian
$-\mathcal{L}u=\Delta_{{\bm{x}}}u=\sum_{i=1}^{d}\partial_{{x}_{i}}^{2}u$. We
parametrize $u$ using a neural network and train its parameters
${\bm{\theta}}$ to minimize the loss

$\displaystyle\\!\\!\\!\\!\\!\\!L({\bm{\theta}})$
$\displaystyle=\frac{1}{2N_{\Omega}}\sum_{n=1}^{N_{\Omega}}(\mathcal{L}u_{\bm{\theta}}({\bm{x}}_{n})-f({\bm{x}}_{n}))^{2}+\frac{1}{2N_{\partial\Omega}}\sum_{n=1}^{N_{\partial\Omega}}(u_{\bm{\theta}}({\bm{x}}^{\text{b}}_{n})-g({\bm{x}}^{\text{b}}_{n}))^{2}\eqqcolon
L_{\Omega}({\bm{\theta}})+L_{\partial\Omega}({\bm{\theta}})$ (1)

with points $\\{{\bm{x}}_{n}\in\Omega\\}_{n=1}^{N_{\Omega}}$ from the domain’s
interior, and points
$\\{{\bm{x}}^{\text{b}}_{n}\in\partial\Omega\\}_{n=1}^{N_{\partial\Omega}}$ on
its boundary.111The second regression loss can also include other constraints
like measurement data.

First-order optimizers like gradient descent and Adam struggle at producing
satisfactory solutions when used to train PINNs [9]. Instead, function space-
inspired second-order methods have lately shown promising results [42]. We
focus on _energy natural gradient descent (ENGD[41])_ which—applied to PINN
objectives like (1)—corresponds to the Gauss-Newton method [6, Chapter 6.3].
ENGD mimics Newton’s method _in function space_ up to a projection onto the
model’s tangent space and a discretization error that vanishes quadratically
in the step size, thus providing locally optimal residual updates.
Alternatively, the Gauss-Newton method can be motivated from the standpoint of
operator-preconditioning, where the Gauss-Newton matrix leads to optimal
conditioning of the problem [14].

Natural gradient methods perform parameter updates via a preconditioned
gradient descent scheme
${\bm{\theta}}\leftarrow{\bm{\theta}}-\alpha{\bm{G}}({\bm{\theta}})^{+}\nabla
L({\bm{\theta}})$, where ${\bm{G}}({\bm{\theta}})^{+}$ denotes the pseudo-
inverse of a suitable _Gramian matrix_
${\bm{G}}({\bm{\theta}})\in{\mathbb{R}}^{P\times P}$ and $\alpha$ is a step
size. ENGD for the PINN loss (1) uses the Gramian

$\displaystyle\begin{split}{\bm{G}}({\bm{\theta}})&=\nicefrac{{1}}{{N_{\Omega}}}\textstyle\sum_{n=1}^{N_{\Omega}}\left(\mathrm{J}_{{\bm{\theta}}}\mathcal{L}u_{\bm{\theta}}({\bm{x}}_{n})\right)^{\top}\mathrm{J}_{{\bm{\theta}}}\mathcal{L}u_{\bm{\theta}}({\bm{x}}_{n})+\nicefrac{{1}}{{N_{\partial\Omega}}}\textstyle\sum_{n=1}^{N_{\partial\Omega}}\left(\mathrm{J}_{{\bm{\theta}}}u_{\bm{\theta}}({\bm{x}}_{n}^{\text{b}})\right)^{\top}\mathrm{J}_{{\bm{\theta}}}u_{\bm{\theta}}({\bm{x}}_{n}^{\text{b}})\\\
&\eqqcolon{\bm{G}}_{\Omega}({\bm{\theta}})+{\bm{G}}_{\partial\Omega}({\bm{\theta}})\,.\end{split}$
(2)

(2) is the Gauss-Newton matrix of the residual
${\bm{r}}({\bm{\theta}})=(\nicefrac{{{\bm{r}}_{\Omega}({\bm{\theta}})^{\top}}}{{\sqrt{N_{\Omega}}}},\nicefrac{{{\bm{r}}_{\partial\Omega}({\bm{\theta}})^{\top}}}{{\sqrt{N_{\partial\Omega}}}})^{\top}\in{\mathbb{R}}^{N_{\Omega}+N_{\partial\Omega}}$
with interior and boundary residuals
$r_{\Omega,n}({\bm{\theta}})=\mathcal{L}u_{\bm{\theta}}({\bm{x}}_{n})-f({\bm{x}}_{n})$
and
$r_{\partial\Omega,n}({\bm{\theta}})=u_{\bm{\theta}}({\bm{x}}_{n}^{\text{b}})-g({\bm{x}}_{n}^{\text{b}})$.

###  2.2 Kronecker-factored Approximate Curvature

We review Kronecker-factored approximate curvature (KFAC) which was introduced
by Heskes [23], Martens & Grosse [38] in the context of maximum likelihood
estimation to approximate the per-layer Fisher information matrix by a
Kronecker product to speed up approximate natural gradient descent [1]. The
Fisher associated with the loss $\nicefrac{{1}}{{2N}}\sum_{n=1}^{N}\left\lVert
u_{\bm{\theta}}({\bm{x}}_{n})-y_{n}\right\rVert_{2}^{2}$ with targets
$y_{n}\in{\mathbb{R}}$ is

${\bm{F}}({\bm{\theta}})=\nicefrac{{1}}{{N}}\textstyle\sum_{n=1}^{N}\left(\mathrm{J}_{{\bm{\theta}}}u_{{\bm{\theta}}}({\bm{x}}_{n})\right)^{\top}\mathrm{J}_{{\bm{\theta}}}u_{{\bm{\theta}}}({\bm{x}}_{n})=\nicefrac{{1}}{{N}}\textstyle\sum_{n=1}^{N}\left(\mathrm{J}_{{\bm{\theta}}}u_{n}\right)^{\top}\mathrm{J}_{{\bm{\theta}}}u_{n}\quad\in{\mathbb{R}}^{P\times
P}\,,$ (3)

where $u_{n}=u_{\bm{\theta}}({\bm{x}}_{n})$, and it coincides with the
classical Gauss-Newton matrix [37]. The established KFAC approximates (3).
While the boundary Gramian ${\bm{G}}_{\partial\Omega}({\bm{\theta}})$ has the
same structure as ${\bm{F}}({\bm{\theta}})$, the interior Gramian
${\bm{G}}_{\Omega}({\bm{\theta}})$ does not as it involves derivative rather
than function evaluations of the net.

KFAC tackles the Fisher’s per-layer block diagonal,
${\bm{F}}({\bm{\theta}})\approx\operatorname{diag}({\bm{F}}^{(1)}({\bm{\theta}}),\dots,{\bm{F}}^{(L)}({\bm{\theta}}))$
with
${\bm{F}}^{(l)}({\bm{\theta}})=\nicefrac{{1}}{{N}}\sum_{n=1}^{N}(\mathrm{J}_{{\bm{\theta}}^{(l)}}u_{n})^{\top}\mathrm{J}_{{\bm{\theta}}^{(l)}}u_{n}\in{\mathbb{R}}^{p^{(l)}\times
p^{(l)}}$. For a fully-connected layer’s block, let’s examine the term
$\mathrm{J}_{{\bm{\theta}}^{(l)}}u_{{\bm{\theta}}}({\bm{x}})$ from Equation 3
for a fixed data point. The layer parameters
${\bm{\theta}}^{(l)}=\operatorname{vec}{\bm{W}}^{(l)}$ enter the computation
via ${\bm{z}}^{(l)}={\bm{W}}^{(l)}{\bm{z}}^{(l-1)}$ and we have
$\mathrm{J}_{{\bm{W}}^{(l)}}{\bm{z}}^{(l)}={{\bm{z}}^{(l-1)}}^{\top}\otimes{\bm{I}}$
[e.g. 10]. Further, the chain rule gives the decomposition
$\mathrm{J}_{{\bm{W}}^{(l)}}u=(\mathrm{J}_{{\bm{z}}^{(l)}}u)\mathrm{J}_{{\bm{W}}^{(l)}}{\bm{z}}^{(l)}={{\bm{z}}^{(l-1)}}^{\top}\otimes\mathrm{J}_{{\bm{z}}^{(l)}}u$.
Inserting into ${\bm{F}}^{(l)}({\bm{\theta}})$, summing over data points, and
using the expectation approximation
$\sum_{n}{\bm{A}}_{n}\otimes{\bm{B}}_{n}\approx
N^{-1}(\sum_{n}{\bm{A}}_{n})\otimes(\sum_{n}{\bm{B}}_{n})$ from Martens &
Grosse [38], we obtain the KFAC approximation for linear layers in supervised
square loss regression with a network’s output,

${\bm{F}}^{(l)}({\bm{\theta}})\approx\underbrace{\left(\nicefrac{{1}}{{N}}\textstyle\sum_{n=1}^{N}{\bm{z}}^{(l-1)}_{n}{{\bm{z}}^{(l-1)}_{n}}^{\top}\right)}_{\eqqcolon{\bm{A}}^{(l)}\in{\mathbb{R}}^{h^{(l-1)}\times
h^{(l-1)}}}\otimes\underbrace{\left(\nicefrac{{1}}{{N}}\textstyle\sum_{n=1}^{N}\left(\mathrm{J}_{{\bm{z}}^{(l)}}u_{n}\right)^{\top}\mathrm{J}_{{\bm{z}}^{(l)}}u_{n}\right)}_{\eqqcolon{\bm{B}}^{(l)}\in{\mathbb{R}}^{h^{(l)}\times
h^{(l)}}}\,.$ (4)

It is cheap to store and invert by inverting the two Kronecker factors.

##  3 Kronecker-Factored Approximate Curvature for PINNs

ENGD’s Gramian is a sum of PDE and boundary Gramians,
${\bm{G}}({\bm{\theta}})={\bm{G}}_{\Omega}({\bm{\theta}})+{\bm{G}}_{\partial\Omega}({\bm{\theta}})$.
We will approximate each Gramian separately with a block diagonal matrix with
Kronecker-factored blocks,
${\bm{G}}_{\bullet}({\bm{\theta}})\approx\operatorname{diag}({\bm{G}}^{(1)}_{\bullet}({\bm{\theta}}),\dots,{\bm{G}}^{(L)}_{\bullet}({\bm{\theta}}))$
for $\bullet\in\\{\Omega,\partial\Omega\\}$ with
${\bm{G}}^{(l)}_{\bullet}({\bm{\theta}})\approx{\bm{A}}^{(l)}_{\bullet}\otimes{\bm{B}}^{(l)}_{\bullet}$.
For the boundary Gramian ${\bm{G}}_{\partial\Omega}({\bm{\theta}})$, we can
re-use the established KFAC from Equation 4 as its loss corresponds to
regression over the network’s output. The interior Gramian
${\bm{G}}_{\Omega}({\bm{\theta}})$, however, involves PDE terms in the form of
network derivatives and therefore _cannot_ be approximated with the existing
KFAC. It requires a new approximation that we develop here for the running
example of the Poisson equation and more general PDEs (Equations 9 and 14). To
do so, we need to make the dependency between the weights and the differential
operator $\mathcal{L}u$ explicit. We use Taylor-mode automatic differentiation
to express this computation of higher-order derivatives as forward passes of a
larger net with shared weights, for which we then propose a Kronecker-factored
approximation, building on KFAC’s recently-proposed generalization to linear
layers with weight sharing [18].

###  3.1 Higher-order Forward Mode Automatic Differentiation as Weight Sharing

Here, we review higher-order forward mode, also known as _Taylor-mode_ ,
automatic differentiation [20, 19, 3, tutorial in § C]. Many PDEs only
incorporate first- and second-order partial derivatives and we focus our
discussion on second-order Taylor-mode for MLPs to keep the presentation
light. However, one can treat higher-order PDEs and arbitrary network
architectures completely analogously.

Taylor-mode propagates directional (higher-order) derivatives. We now recap
the forward propagation rules for MLPs consisting of fully-connected and
element-wise activation layers. Our goal is to evaluate first-and second-order
partial derivatives of the form
$\partial_{{x}_{i}}u,\partial^{2}_{{x}_{i},{x}_{j}}u$ for $i,j=1,\dots,d$. At
the first layer, set
${\bm{z}}^{(0)}={\bm{x}}\in\mathbb{R}^{d},\partial_{x_{i}}{\bm{z}}^{(0)}={\bm{e}}_{i}\in\mathbb{R}^{d}$,
i.e., the $i$-th basis vector and
$\partial^{2}_{x_{i},x_{j}}{\bm{z}}^{(0)}={\bm{0}}\in\mathbb{R}^{d}$.

For a linear layer
$f_{{\bm{\theta}}^{(l)}}({\bm{z}}^{(l-1)})={\bm{W}}^{(l)}{\bm{z}}^{(l-1)}$,
applying the chain rule yields the propagation rule

$\displaystyle{\bm{z}}^{(l)}$
$\displaystyle={\bm{W}}^{(l)}{\bm{z}}^{(l-1)}\quad\in{\mathbb{R}}^{h^{(l)}}\,,$
(5a) $\displaystyle\partial_{x_{i}}{\bm{z}}^{(l)}$
$\displaystyle={\bm{W}}^{(l)}\partial_{x_{i}}{\bm{z}}^{(l-1)}\quad\in{\mathbb{R}}^{h^{(l)}}\,,$
(5b) $\displaystyle\partial^{2}_{x_{i},x_{j}}{\bm{z}}^{(l)}$
$\displaystyle={\bm{W}}^{(l)}\partial^{2}_{x_{i},x_{j}}{\bm{z}}^{(l-1)}\quad\in{\mathbb{R}}^{h^{(l)}}\,.$
(5c)

The propagation rule through a nonlinear element-wise activation layer
${\bm{z}}^{(l-1)}\mapsto\sigma({\bm{z}}^{(l-1)})$ is

$\displaystyle{\bm{z}}^{(l)}$
$\displaystyle=\sigma({\bm{z}}^{(l-1)})\quad\in{\mathbb{R}}^{h^{(l)}}\,,$ (6a)
$\displaystyle\partial_{x_{i}}{\bm{z}}^{(l)}$
$\displaystyle=\sigma^{\prime}({\bm{z}}^{(l-1)})\odot\partial_{x_{i}}{\bm{z}}^{(l-1)}\quad\in{\mathbb{R}}^{h^{(l)}}\,,$
(6b) $\displaystyle\partial^{2}_{x_{i},x_{j}}{\bm{z}}^{(l)}$
$\displaystyle=\partial_{x_{i}}{\bm{z}}^{(l-1)}\odot\sigma^{\prime\prime}({\bm{z}}^{(l-1)})\odot\partial_{x_{j}}{\bm{z}}^{(l-1)}+\sigma^{\prime}({\bm{z}}^{(l-1)})\odot\partial^{2}_{x_{i},x_{j}}{\bm{z}}^{(l-1)}\quad\in{\mathbb{R}}^{h^{(l)}}\,.$
(6c)

##### Forward Laplacian

For differential operators of special structure, we can fuse the Taylor-mode
forward propagation of individual directional derivatives in Equations 5 and 6
and obtain a more efficient computation. E.g., to compute not the full Hessian
but only the Laplacian, we can simplify the forward pass, which yields the
_forward Laplacian_ framework of Li et al. [29]. To the best of our knowledge,
this connection has not been pointed out in the literature. Concretely, by
summing (5c) and (6c) over $i=j$, we obtain the Laplacian forward pass for
linear and activation layers

$\displaystyle\Delta_{\bm{x}}{\bm{z}}^{(l)}$
$\displaystyle={\bm{W}}^{(l)}\Delta_{\bm{x}}{\bm{z}}^{(l-1)}\quad\in{\mathbb{R}}^{h^{(l)}}\,,$
(7a) $\displaystyle\Delta_{\bm{x}}{\bm{z}}^{(l)}$
$\displaystyle=\sigma^{\prime}({\bm{z}}^{(l-1)})\odot\Delta_{\bm{x}}{\bm{z}}^{(l-1)}+\textstyle\sum_{i=1}^{d}\sigma^{\prime\prime}({\bm{z}}^{(l-1)})\odot(\partial_{x_{i}}{{\bm{z}}^{(l-1)}})^{\odot
2}\quad\in{\mathbb{R}}^{h^{(l)}}\,.$ (7b)

This reduces computational cost, but is restricted to PDEs that involve
second-order derivatives only via the Laplacian, or a partial Laplacian over a
sub-set of input coordinates (e.g. heat equation, § 4). For a more general
second-order linear PDE operator
${\mathcal{L}}=\sum_{i,j=1}^{d}c_{i,j}\partial^{2}_{{x}_{i},{x}_{j}}$, the
forward pass for a linear layer is
$\mathcal{L}{\bm{z}}^{(l)}={\bm{W}}^{(l)}\mathcal{L}{\bm{z}}^{(l-1)}\in{\mathbb{R}}^{h^{(l)}}$,
generalizing (7a), and similarly for Equation 7b. This is similar to [30],
which transforms the input space using an eigen-decomposition of $c_{i,j}$ to
a basis where it is diagonal with entries $\\{0,\pm 1\\}$, reducing the
computation to two forward Laplacians.

Importantly, the computation of higher-order derivatives for linear layers
boils down to a forward pass through the layer with weight sharing over the
different partial derivatives (Equation 5), and weight sharing can potentially
be reduced depending on the differential operator’s structure (Equation 7a).
Therefore, we can use the concept of KFAC in the presence of weight sharing to
derive a principled Kronecker approximation for Gramians containing
differential operator terms.

###  3.2 KFAC for Gauss-Newton Matrices with the Laplace Operator

Let’s consider the Poisson equation’s interior Gramian block for a linear
layer (suppressing $\Omega$ in $N_{\Omega}$)

$\displaystyle{\bm{G}}^{(l)}_{\Omega}({\bm{\theta}})=\nicefrac{{1}}{{N}}\textstyle\sum_{n=1}^{N}\left(\mathrm{J}_{{\bm{W}}^{(l)}}\Delta_{\bm{x}}u_{n}\right)^{\top}\mathrm{J}_{{\bm{W}}^{(l)}}\Delta_{\bm{x}}u_{n}\,.$

Because we made the Laplacian computation explicit through Taylor-mode
autodiff (§ 3.1, specifically Equation 7a), we can stack all output vectors
that share the layer’s weight into a matrix
${\bm{Z}}_{n}^{(l)}\in{\mathbb{R}}^{h^{(l)}\times S}$ with $S=d+2$ and columns
${\bm{Z}}_{n,1}^{(l)}={\bm{z}}_{n}^{(l)},{\bm{Z}}_{n,2}^{(l)}=\partial_{x_{1}}{\bm{z}}_{n}^{(l)},\dots,{\bm{Z}}_{n,1+d}^{(l)}=\partial_{x_{d}}{\bm{z}}_{n}^{(l)}$,
and ${\bm{Z}}_{n,2+d}^{(l)}=\Delta_{\bm{x}}{\bm{z}}_{n}^{(l)}$ (likewise
${\bm{Z}}_{n}^{(l-1)}\in{\mathbb{R}}^{h^{(l-1)}\times S}$ for the layer
inputs), then apply the chain rule

$\displaystyle\mathrm{J}_{{\bm{W}}^{(l)}}\Delta_{\bm{x}}u_{n}$
$\displaystyle=(\mathrm{J}_{{\bm{Z}}_{n}^{(l)}}\Delta_{\bm{x}}u_{n})\mathrm{J}_{{\bm{W}}^{(l)}}{\bm{Z}}_{n}^{(l)}=\textstyle\sum_{s=1}^{S}{\underbrace{{\bm{Z}}^{(l-1)}_{n,s}}_{\in{\mathbb{R}}^{h^{(l-1)}}}}^{\top}\otimes\underbrace{\mathrm{J}_{{\bm{Z}}^{(l)}_{n,s}}\Delta_{{\bm{x}}}u_{n}}_{\eqqcolon{\bm{g}}_{n,s}^{(l)}\in{\mathbb{R}}^{h^{(l)}}}\,,$

which has a structure similar to the Jacobian in § 2.2, but with an
additional sum over the $S$ shared vectors. With that, we can now express the
exact interior Gramian for a layer as

${\bm{G}}^{(l)}_{\Omega}({\bm{\theta}})=\nicefrac{{1}}{{N}}\textstyle\sum_{n=1}^{N}\sum_{s=1}^{S}\sum_{s^{\prime}=1}^{S}{\bm{Z}}^{(l-1)}_{n,s}{\bm{Z}}^{(l-1)\top}_{n,s^{\prime}}\otimes{\bm{g}}_{n,s}^{(l)}{\bm{g}}_{n,s^{\prime}}^{(l)\top}.$
(8)

Next, we want to approximate Equation 8 with a Kronecker product. To avoid
introducing a new convention, we rely on the KFAC approximation for linear
layers with weight sharing developed by Eschenhagen et al. [18]—specifically,
the approximation called _KFAC-expand_. This drops all terms with $s\neq
s^{\prime}$, then applies the expectation approximation from § 2.2 over the
batch and shared axes:

KFAC for the Gauss-Newton matrix of a Laplace operator
$\displaystyle{\bm{G}}^{(l)}_{\Omega}({\bm{\theta}})\approx\left(\frac{1}{NS}\sum_{n,s=1}^{N,S}{\bm{Z}}^{(l-1)}_{n,s}{{\bm{Z}}^{(l-1)}_{n,s}}^{\top}\right)\otimes\left(\frac{1}{N}\sum_{n,s=1}^{N,S}{\bm{g}}^{(l)}_{n,s}{{\bm{g}}^{(l)}_{n,s}}^{\top}\right)\eqqcolon{\bm{A}}_{\Omega}^{(l)}\otimes{\bm{B}}_{\Omega}^{(l)}$
(9)

###  3.3 KFAC for Generalized Gauss-Newton Matrices Involving General PDE
Terms

To generalize the previous section, let’s consider the general $M$-dimensional
PDE system of order $k$,

$\Psi(u,D_{{\bm{x}}}u,\dots,D^{k}_{{\bm{x}}}u)={\bm{0}}\in{\mathbb{R}}^{M},$
(10)

where $D^{m}_{\bm{x}}u$ collects all partial derivatives of order $m$. For
$m\in\\{0,\dots,k\\}$ there are $S_{m}=\binom{d+m-1}{d-1}$ independent partial
derivatives and the total number of independent derivatives is
$S\coloneqq\sum_{m=0}^{k}S_{m}=\binom{d+k}{k}$. $\Psi$ is a smooth mapping
from all partial derivatives to ${\mathbb{R}}^{M}$,
$\Psi\colon\mathbb{{\mathbb{R}}}^{S}\to\mathbb{{\mathbb{R}}}^{M}$. To
construct a PINN loss for Equation 10, we feed the residual
${\bm{r}}_{\Omega,n}({\bm{\theta}})\coloneqq\Psi(u_{{\bm{\theta}}}({\bm{x}}_{n}),D_{{\bm{x}}}u_{{\bm{\theta}}}({\bm{x}}_{n}),\dots,D^{k}_{{\bm{x}}}u_{{\bm{\theta}}}({\bm{x}}_{n}))\in{\mathbb{R}}^{M}$
where $D^{m}_{{\bm{x}}}u_{{\bm{\theta}}}({\bm{x}}_{n})\in{\mathbb{R}}^{d\times
S_{m}}$ into a smooth convex criterion function
$\ell\colon{\mathbb{R}}^{M}\to{\mathbb{R}}$,

$L_{\Omega}({\bm{\theta}})\coloneqq\nicefrac{{1}}{{N}}\textstyle\sum_{n=1}^{N}\ell({\bm{r}}_{\Omega,n}({\bm{\theta}}))\,.$
(11)

The generalized Gauss-Newton (GGN) matrix [52] is the Hessian of
$L_{\Omega}({\bm{\theta}})$ when the residual is linearized w.r.t.
${\bm{\theta}}$ before differentiation. It is positive semi-definite and has
the form

$\displaystyle{\bm{G}}_{\Omega}({\bm{\theta}})\coloneqq\nicefrac{{1}}{{N}}\textstyle\sum_{n=1}^{N}\left(\mathrm{J}_{\bm{\theta}}{\bm{r}}_{\Omega,n}({\bm{\theta}})\right)^{\top}{\bm{\Lambda}}({\bm{r}}_{\Omega,n})\left(\mathrm{J}_{\bm{\theta}}{\bm{r}}_{\Omega,n}({\bm{\theta}})\right)\,,$
(12)

with
${\bm{\Lambda}}({\bm{r}})\coloneqq\nabla^{2}_{{\bm{r}}}\ell({\bm{r}})\in{\mathbb{R}}^{M\times
M}\succ 0$ the criterion’s Hessian, e.g.
$\ell({\bm{r}})=\nicefrac{{1}}{{2}}\lVert{\bm{r}}\rVert_{2}^{2}$ and
${\bm{\Lambda}}({\bm{r}})={\bm{I}}_{M}$.

Generalizing the second-order Taylor-mode from § 3.1 to higher orders for
the linear layer, we find

$\displaystyle
D_{\bm{x}}^{m}{\bm{z}}^{(l)}={\bm{W}}^{(l)}D_{\bm{x}}^{m}{\bm{z}}^{(l-1)}\qquad\in{\mathbb{R}}^{h^{(l)}\times
S_{m}}$ (13)

for any $m$. Hence, we can derive a forward propagation for the required
derivatives where a linear layer processes at most $S$ vectors222Depending on
the linear operator, one may reduce weight sharing, as demonstrated for the
Laplacian in § 3.1., i.e. the linear layer’s weight is shared over the
matrices
$D^{0}_{\bm{x}}{\bm{z}}^{(l-1)}\coloneqq{\bm{z}}^{(l-1)},D_{\bm{x}}^{1}{\bm{z}}^{(l-1)},\dots,D_{\bm{x}}^{k}{\bm{z}}^{(l-1)}$.
Stacking them into a matrix
${\bm{Z}}^{(l-1)}_{n}=({\bm{z}}^{(l-1)},D_{\bm{x}}^{1}{\bm{z}}^{(l-1)},\dots,D_{\bm{x}}^{k}{\bm{z}}^{(l-1)})\in{\mathbb{R}}^{h^{(l-1)}\times
S}$ (and ${\bm{Z}}^{(l)}_{n}$ for the outputs), the chain rule yields

$\displaystyle\begin{split}{\bm{G}}^{(l)}_{\Omega}({\bm{\theta}})&=\nicefrac{{1}}{{N}}\textstyle\sum_{n=1}^{N}\left(\mathrm{J}_{{\bm{W}}^{(l)}}{\bm{Z}}^{(l)}_{n}\right)^{\top}\left(\mathrm{J}_{{\bm{Z}}^{(l)}_{n}}{\bm{r}}_{\Omega,n}\right)^{\top}{\bm{\Lambda}}({\bm{r}}_{\Omega,n})\left(\mathrm{J}_{{\bm{Z}}^{(l)}_{n}}{\bm{r}}_{\Omega,n}\right)\left(\mathrm{J}_{{\bm{W}}^{(l)}}{\bm{Z}}^{(l)}_{n}\right)\\\
&=\nicefrac{{1}}{{N}}\textstyle\sum_{n,s,s^{\prime}=1}^{N,S,S}\left(\mathrm{J}_{{\bm{W}}^{(l)}}{\bm{Z}}^{(l)}_{n,s}\right)^{\top}\left(\mathrm{J}_{{\bm{Z}}^{(l)}_{n,s}}{\bm{r}}_{\Omega,n}\right)^{\top}{\bm{\Lambda}}({\bm{r}}_{\Omega,n})\left(\mathrm{J}_{{\bm{Z}}^{(l)}_{n,s^{\prime}}}{\bm{r}}_{\Omega,n}\right)\left(\mathrm{J}_{{\bm{W}}^{(l)}}{\bm{Z}}^{(l)}_{n,s^{\prime}}\right)\\\
&=\nicefrac{{1}}{{N}}\textstyle\sum_{n,s,s^{\prime}=1}^{N,S,S}{\bm{Z}}^{(l-1)}_{n,s}{\bm{Z}}^{(l-1)\top}_{n,s^{\prime}}\otimes\left(\mathrm{J}_{{\bm{Z}}^{(l)}_{n,s}}{\bm{r}}_{\Omega,n}\right)^{\top}{\bm{\Lambda}}({\bm{r}}_{\Omega,n})\left(\mathrm{J}_{{\bm{Z}}^{(l)}_{n,s^{\prime}}}{\bm{r}}_{\Omega,n}\right)\end{split}$

where ${\bm{Z}}_{n,s}^{(l-1)}\in{\mathbb{R}}^{h^{(l-1)}}$ denotes the $s$-th
column of ${\bm{Z}}_{n}^{(l-1)}$. Following the same steps as in § 3.2, we
apply the KFAC-expand approximation from [18] to obtain the generalization of
Equation 9:

KFAC for the GGN matrix of a general PDE operator
$\displaystyle\begin{split}{\bm{G}}_{\Omega}^{(l)}({\bm{\theta}})&\approx\left(\\!\\!\frac{1}{NS}\\!\\!\sum_{n,s=1}^{N,S}\\!\\!{\bm{Z}}^{(l-1)}_{n,s}{\bm{Z}}^{(l-1)\top}_{n,s^{\prime}}\\!\\!\right)\otimes\left(\\!\\!\frac{1}{N}\\!\\!\sum_{n,s=1}^{N,S}\\!\\!\left(\mathrm{J}_{{\bm{Z}}^{(l)}_{n,s}}{\bm{r}}_{\Omega,n}\right)^{\top}\\!\\!\\!{\bm{\Lambda}}({\bm{r}}_{\Omega,n})\left(\mathrm{J}_{{\bm{Z}}^{(l)}_{n,s}}{\bm{r}}_{\Omega,n}\right)\\!\\!\right)\\!\\!\\\
&\eqqcolon{\bm{A}}_{\Omega}^{(l)}\otimes{\bm{B}}_{\Omega}^{(l)}\end{split}$
(14)

To bring this expression even closer to Equation 9, we can re-write the second
Kronecker factor using an outer product decomposition
${\bm{\Lambda}}({\bm{r}}_{\Omega,n})=\sum_{m=1}^{M}{\bm{l}}_{n,m}{\bm{l}}_{n,m}$
with ${\bm{l}}_{n,m}\in{\mathbb{R}}^{M}$, then introduce
${\bm{g}}^{(l)}_{n,s,m}\coloneqq(\mathrm{J}_{{\bm{Z}}^{(l)}_{n,s}}{\bm{r}}_{\Omega,n})^{\top}{\bm{l}}_{n,m}\in{\mathbb{R}}^{h^{(l)}}$
and write the second term as
$\nicefrac{{1}}{{N}}\sum_{n,s,m=1}^{N,S,M}{\bm{g}}^{(l)}_{n,s,m}{\bm{g}}^{(l)\top}_{n,s,m}$,
similar to the Kronecker-factored low-rank (KFLR) approach of Botev et al.
[5].

KFAC for variational problems Our proposed KFAC approximation is not limited
to PINNs and can be used for variational problems of the form

$\displaystyle\min_{u}\int_{\Omega}\ell(u,\partial_{{\bm{x}}}u,\dots,\partial^{k}_{{\bm{x}}}u)\mathrm{d}{\bm{x}}\,,$
(15)

where $\ell\colon\mathbb{R}^{K}\to\mathbb{R}$ is a convex function. We can
perceive this as a special case of the setting above with
$\Psi=\operatorname{id}$ and hence the KFAC approximation (14) remains
meaningful. In particular, it can be used for the _deep Ritz method_ and other
variational approaches to solve PDEs [17].

###  3.4 Algorithmic Details

To design an optimizer based on our KFAC approximation, we re-use techniques
from the original KFAC [38] & ENGD [41] algorithms. § B shows pseudo-code
for our method on the Poisson equation.

At iteration $t$, we approximate the per-layer interior and boundary Gramians
using our derived Kronecker approximation (Equations 9 and 14),
${\bm{G}}_{\Omega,t}^{(l)}\approx{\bm{A}}_{\Omega,t}^{(l)}\otimes{\bm{B}}_{\Omega,t}^{(l)}$
and
${\bm{G}}_{\partial\Omega,t}^{(l)}\approx{\bm{A}}_{\Omega,t}^{(l)}\otimes{\bm{B}}_{\Omega,t}^{(l)}$.

Exponential moving average and damping For preconditioning, we accumulate the
Kronecker factors ${\bm{A}}^{(l)}_{\bullet,t},{\bm{B}}^{(l)}_{\bullet,t}$ over
time using an exponential moving average
$\hat{{\bm{A}}}^{(l)}_{\bullet,t}=\beta\hat{{\bm{A}}}^{(l)}_{\bullet,t-1}+(1-\beta){\bm{A}}^{(l)}_{\bullet,t}$
of factor $\beta\in[0,1)$ (identically for
$\hat{{\bm{B}}}^{(l)}_{\bullet,t}$), similar to the original KFAC. Moreover,
we apply the same constant damping of strength $\lambda>0$ to all Kronecker
factors,
$\tilde{{\bm{A}}}^{(l)}_{\bullet,t}=\hat{{\bm{A}}}^{(l)}_{\bullet,t}+\lambda{\bm{I}}$
and
$\tilde{{\bm{B}}}^{(l)}_{\bullet,t}=\hat{{\bm{B}}}^{(l)}_{\bullet,t}+\lambda{\bm{I}}$
such that the curvature approximation used for preconditioning at step $t$ is

$\displaystyle{\bm{G}}_{\bullet,t}\approx\operatorname{diag}\left(\tilde{{\bm{A}}}^{(1)}_{\Omega,t}\otimes\tilde{{\bm{B}}}_{\Omega,t}^{(1)},\dots,\tilde{{\bm{A}}}^{(L)}_{\Omega,t}\otimes\tilde{{\bm{B}}}_{\Omega,t}^{(L)}\right)+\operatorname{diag}\left(\tilde{{\bm{A}}}^{(1)}_{\partial\Omega,t}\otimes\tilde{{\bm{B}}}_{\partial\Omega,t}^{(1)},\dots,\tilde{{\bm{A}}}^{(L)}_{\partial\Omega,t}\otimes\tilde{{\bm{B}}}_{\partial\Omega,t}^{(L)}\right)\,.$

Gradient preconditioning Given layer $l$’s mini-batch gradient
${\bm{g}}_{t}^{(l)}=\nicefrac{{\partial
L({\bm{\theta}}_{t})}}{{\partial{\bm{\theta}}^{(l)}_{t}}}\in{\mathbb{R}}^{p^{(l)}}$,
we obtain an update direction
${\bm{\Delta}}_{t}^{(l)}=-(\tilde{{\bm{A}}}^{(l)}_{\Omega,t}\otimes\tilde{{\bm{B}}}^{(l)}_{\Omega,t}+\tilde{{\bm{A}}}^{(l)}_{\partial\Omega,t}\otimes\tilde{{\bm{B}}}^{(l)}_{\partial\Omega,t})^{-1}{\bm{g}}_{t}^{(l)}\in{\mathbb{R}}^{p^{(l)}}$
using the trick of [38, Appendix I] to invert the Kronecker sum which requires
eigen-decomposing all Kronecker factors.

Learning rate and momentum From the preconditioned gradient
${\bm{\Delta}}_{t}\in{\mathbb{R}}^{P}$, we consider two different updates
${\bm{\theta}}_{t+1}={\bm{\theta}}_{t}+{\bm{\delta}}_{t}$ we call _KFAC_ and
_KFAC*_. KFAC uses momentum over previous updates,
$\hat{{\bm{\delta}}}_{t}=\mu{\bm{\delta}}_{t-1}+{\bm{\Delta}}_{t}$, and $\mu$
is chosen by the practitioner. Like ENGD, it uses a logarithmic grid line
search, selecting ${\bm{\delta}}_{t}=\alpha_{\star}\hat{{\bm{\delta}}}_{t}$
with
$\alpha_{\star}=\operatorname*{arg\,min}_{\alpha}L({\bm{\theta}}_{t}+\alpha\hat{{\bm{\delta}}}_{t})$
where $\alpha\in\\{2^{-30},\dots,2^{0}\\}$. KFAC* uses the automatic learning
rate and momentum heuristic of the original KFAC optimizer. It parametrizes
the iteration’s update as
${\bm{\delta}}_{t+1}(\alpha,\mu)=\alpha{\bm{\Delta}}_{t}+\mu{\bm{\delta}}_{t}$,
then obtains the optimal parameters by minimizing the quadratic model
$m({\bm{\delta}}_{t+1})=L({\bm{\theta}}_{t})+{\bm{\delta}}_{t+1}^{\top}{\bm{g}}_{t}+\nicefrac{{1}}{{2}}{\bm{\delta}}_{t+1}^{\top}({\bm{G}}({\bm{\theta}}_{t})+\lambda{\bm{I}}){\bm{\delta}}_{t+1}$
with the exact damped Gramian. The optimal learning rate and momentum
$\operatorname*{arg\,min}_{\alpha,\mu}m({\bm{\delta}}_{t+1})$ are

$\displaystyle\begin{pmatrix}\alpha_{\star}\\\
\mu_{\star}\end{pmatrix}=-\begin{pmatrix}{\bm{\Delta}}_{t}^{\top}{\bm{G}}({\bm{\theta}}_{t}){\bm{\Delta}}_{t}+\lambda\left\lVert{\bm{\Delta}}_{t}\right\rVert^{2}&{\bm{\Delta}}_{t}^{\top}{\bm{G}}({\bm{\theta}}_{t}){\bm{\delta}}_{t}+\lambda{\bm{\Delta}}^{\top}_{t}{\bm{\delta}}_{t}\\\
{\bm{\Delta}}_{t}^{\top}{\bm{G}}({\bm{\theta}}_{t}){\bm{\delta}}_{t}+\lambda{\bm{\Delta}}^{\top}_{t}{\bm{\delta}}_{t}&{\bm{\delta}}_{t}^{\top}{\bm{G}}({\bm{\theta}}_{t}){\bm{\Delta}}_{t}+\lambda\left\lVert{\bm{\delta}}_{t}\right\rVert^{2}\end{pmatrix}^{-1}\begin{pmatrix}{\bm{\Delta}}_{t}^{\top}{\bm{g}}_{t}\\\
{\bm{\delta}}_{t}^{\top}{\bm{g}}_{t}\end{pmatrix}$

(see [38, Section 7] for details). The computational cost is dominated by the
two Gramian-vector products with ${\bm{\Delta}}_{t}$ and ${\bm{\delta}}_{t}$.
By using the Gramian’s outer product structure [12, 45], we perform them with
autodiff [48, 52] using one Jacobian-vector product each, as recommended in
[38].

##  4 Experiments

Figure 1: Performance of different optimizers on the 2d Poisson equation (16)
measured in relative $L_{2}$ error against wall clock time for architectures
with different parameter dimensions $D$.

We implement KFAC, KFAC*, and ENGD with either the per-layer or full Gramian
in PyTorch [46]. As a matrix-free version of ENGD, we use the Hessian-free
optimizer [36] which uses truncated conjugate gradients (CG) with exact
Gramian-vector products to precondition the gradient. We chose this because
there is a fully-featured implementation from Tatzel et al. [54] which offers
many additional heuristics like adaptive damping, CG backtracking, and
backtracking line search, allowing this algorithm to work well with little
hyper-parameter tuning. As baselines, we use SGD with tuned learning rate and
momentum, Adam with tuned learning rate, and LBFGS with tuned learning rate
and history size. We tune hyper-parameters using Weights & Biases [59] (see
§ A.1 for the exact protocol). For random/grid search, we run an initial
round of approximately 50 runs with generous search spaces, then narrow them
down and re-run for another 50 runs; for Bayesian search, we assign the same
total compute to each optimizer. We report runs with lowest $L_{2}$ error
estimated on a held-out data set with the known solution to the studied PDE.
To be comparable, all runs are executed on a compute cluster with RTX 6000
GPUs (24 GiB RAM) in double precision, and we use the same computation time
budget for all optimizers on a fixed PINN problem. All search spaces and best
run hyper-parameters, as well as training curves over iteration count rather
than time, are in § A.

Pedagogical example: 2d Poisson equation We start with a low-dimensional
Poisson equation from Müller & Zeinhofer [41] to reproduce ENGD’s performance
(Figure 1). It is given by

$\displaystyle\begin{split}-\Delta u(x,y)&=2\pi^{2}\sin(\pi x)\sin(\pi
y)\quad\text{for }(x,y)\in[0,1]^{2}\,\\\ u(x,y)&=0\quad\text{for
}(x,y)\in\partial[0,1]^{2}.\end{split}$ (16)

We choose a fixed data set of same size as the original paper, then use
random/grid search to evaluate the performance of all optimizers for different
$\tanh$-activated MLPs, one shallow and two with five fully-connected layers
of different width (all details in § A.2). We include ENGD whenever the
network’s parameter space is small enough to build up the Gramian.

For the shallow net (Figure 1, left), we can reproduce the results of [41],
where exact ENGD achieves high accuracy. In terms of computation time, our
KFACs are competitive with full-ENGD for a long phase, outperforming the
first-order and quasi-Newton baselines. In contrast to ENGD, which runs out of
memory for networks with more than $10\,000$ parameters, KFAC scales to larger
networks (Figure 1, center and right) and is competitive with other second-
order optimizers like Hessian-free, which uses more sophisticated heuristics.
We make similar observations on a small (1+1)d heat equation with the same
models, see § A.7 and A8.

Figure 2: Performance of different optimizers on the (4+1)d heat equation (17)
measured in relative $L^{2}$ error against wall clock time for architectures
with different parameter dimensions $D$.

An evolutionary problem: (4+1)d heat equation To demonstrate that our methods
can also be applied to other problems than the Poisson equation, we consider a
four-dimensional heat equation

$\displaystyle\begin{split}\partial_{t}u(t,{\bm{x}})-\kappa\Delta_{\bm{x}}u(t,{\bm{x}})&=0\quad\text{for
}t\in[0,1],{\bm{x}}\in[0,1]^{4}\,,\\\
u(0,{\bm{x}})&=\textstyle\sum_{i=1}^{4}\sin(2x_{i})\quad\text{for
}{\bm{x}}\in[0,1]^{4}\,,\\\
u(t,{\bm{x}})&=\textstyle\exp(-t)\sum_{i=1}^{4}\sin(2x_{i})\quad\text{for
}t\in[0,1],{\bm{x}}\in\partial[0,1]^{4}\,,\end{split}$ (17)

with diffusivity constant $\kappa=\nicefrac{{1}}{{4}}$, similar to that
studied in [41] (see § A.6 for the heat equation’s PINN loss). We use the
previous architectures with same hidden widths and evaluate optimizer
performance with random/grid search (all details in § A.8), see Figure 2. To
prevent over-fitting, we use mini-batches and sample a new batch each
iteration. We noticed that KFAC improves significantly when batches are
sampled less frequently and hypothesize that it might need more iterations to
make similar progress than one iteration of Hessian-free or ENGD on a batch.
Consequently we sample a new batch only every 100 iterations for KFAC. To
ensure that this does not lead to an unfair advantage for KFAC, we conduct an
additional experiment for the MLP with $D=$116\,864$$ where we tune batch
sizes, batch sampling frequencies, and all hyper-parameters with generous
search spaces using Bayesian search (§ A.9). We find that this does not
significantly boost performance of the other methods (compare Figures 2 and
A10). Again, we observe that KFAC offers competitive performance compared to
other second-order methods for networks with prohibitive size for ENGD and
consistently outperforms SGD, Adam, and LBFGS. We confirmed these observations
with another 5d Poisson equation on the same architectures, see § A.3 and
A5.

Figure 3:  Optimizer performance on Poisson equations in high dimensions and
different boundary conditions measured in relative $L_{2}$ error against wall
clock time for networks with $D$ parameters.

High-dimensional Poisson equations To demonstrate scaling to high-dimensional
PDEs and even larger neural networks, we consider three Poisson equations
($d=5,10,100$) with different boundary conditions used in [17, 41], which
admit the solutions

$\displaystyle\begin{split}u_{\star}({\bm{x}})&=\textstyle\sum_{i=1}^{5}\cos(\pi
x_{i})\quad\text{for }{\bm{x}}\in[0,1]^{5}\,,\\\
u_{\star}({\bm{x}})&=\textstyle\sum_{k=1}^{5}x_{2k-1}x_{2k}\quad\text{for
}{\bm{x}}\in[0,1]^{10}\,,\\\
u_{\star}({\bm{x}})&=\lVert{\bm{x}}\rVert_{2}^{2}\quad\text{for
}{\bm{x}}\in[0,1]^{100}\,.\end{split}$ (18)

We use the same architectures as before, but with larger intermediate widths
and parameters up to a million (Figure 3). Due to lacking references for
training such high-dimensional problems, we select all hyper-parameters via
Bayesian search, including batch sizes and batch sampling frequencies (all
details in § A.5). We see a similar picture as before with KFAC consistently
outperforming first-order methods and LBFGS, offering competitive performance
with Hessian-free. To account for the possibility that the Bayesian search did
not properly identify good hyper-parameters, we conduct a random/grid search
experiment for the 10d Poisson equation (Figure 3, middle), using similar
batch sizes and same batch sampling frequencies as for the $(4+1)$d heat
equation (all details in § A.4). In this experiment, KFAC also achieved
similar performance than Hessian-free and outperformed SGD, Adam, and LBFGS
(Figure A6).

##  5 Discussion and Conclusion

We extended the concept of Kronecker-factored approximate curvature (KFAC) to
Gauss-Newton matrices of Physics-informed neural network (PINN) losses that
involve derivatives, rather than function evaluations, of the neural net. This
greatly reduces the computational cost of approximate natural gradient
methods, which are known to work well on PINNs, and allows them to scale to
much larger nets. Our approach goes beyond the established KFAC for
traditional supervised problems as it captures contributions from a PDE’s
differential operator that are crucial for optimization. To establish KFAC for
such losses, we use Taylor-mode autodiff to view the differential operator’s
compute graph as a forward net with shared weights, then apply the recently-
developed formulation of KFAC for linear layers with weight sharing.
Empirically, we find that our KFAC-based optimizers are competitive with
expensive second-order methods on small problems and scale to high-dimensional
neural networks and PDEs while consistently outperforming first-order methods
and LBFGS.

Limitations & future directions While our implementation currently only
supports MLPs and the Poisson and heat equations, the concepts we use to
derive KFAC (Taylor-mode, weight sharing) apply to arbitrary architectures and
PDEs, as described in § 3.3. We are excited that our current algorithms show
promising performance when compared to second-order methods with sophisticated
heuristics. In fact, the original KFAC optimizer itself [38] relies heavily on
such heuristics that are said to be crucial for its performance [8]. Our
algorithms borrow components, but we did not explore all bells and whistles,
e.g. adaptive damping and heuristics to distribute damping over the Kronecker
factors. We believe our current algorithm’s performance can further be
improved, e.g. by exploring (1) updating the KFAC matrices less frequently, as
is standard for traditional KFAC, (2) merging the two Kronecker approximations
for boundary and interior Gramians into a single one, (3) removing matrix
inversions [32], (4) using structured Kronecker factors [31], (5) computing
the Kronecker factors in parallel with the gradient [11], (6) using single or
mixed precision training [40], and (7) studying cheaper KFAC flavours based on
the empirical Fisher [27] or input-based curvature [2, 49].

## Acknowledgments and Disclosure of Funding

The authors thank Runa Eschenhagen for insightful discussions on KFAC for
linear weight sharing layers. Resources used in preparing this research were
provided, in part, by the Province of Ontario, the Government of Canada
through CIFAR, and companies sponsoring Vector Institute. JM acknowledges
funding by the Deutsche Forschungsgemeinschaft (DFG, German Research
Foundation) under the project number 442047500 through the Collaborative
Research Center _Sparsity and Singular Structures_ (SFB 1481).

## References

  * Amari [1998] Amari, S.-I.  Natural gradient works efficiently in learning.  _Neural computation_ , 10(2):251–276, 1998.
  * Benzing [2022] Benzing, F.  Gradient Descent on Neurons and its Link to Approximate Second-order Optimization.  In _International Conference on Machine Learning (ICML)_ , 2022.
  * Bettencourt et al. [2019] Bettencourt, J., Johnson, M. J., and Duvenaud, D.  Taylor-mode automatic differentiation for higher-order derivatives in JAX.  In _Advances in Neural Information Processing Systems (NeurIPS); Workhop on Program Transformations for ML_ , 2019.
  * Bonfanti et al. [2024] Bonfanti, A., Bruno, G., and Cipriani, C.  The Challenges of the Nonlinear Regime for Physics-Informed Neural Networks.  _arXiv preprint arXiv:2402.03864_ , 2024.
  * Botev et al. [2017] Botev, A., Ritter, H., and Barber, D.  Practical Gauss-Newton optimisation for deep learning.  In _International Conference on Machine Learning (ICML)_ , 2017.
  * Bottou et al. [2016] Bottou, L., Curtis, F. E., and Nocedal, J.  Optimization methods for large-scale machine learning.  _SIAM Review (SIREV)_ , 60, 2016.
  * Chen et al. [2024] Chen, Z., McCarran, J., Vizcaino, E., Soljačić, M., and Luo, D.  TENG: Time-Evolving Natural Gradient for Solving PDEs with Deep Neural Net.  _arXiv preprint arXiv:2404.10771_ , 2024.
  * Clarke et al. [2023] Clarke, R. M., Su, B., and Hernández-Lobato, J. M.  Adam through a second-order lens.  2023\.
  * Cuomo et al. [2022] Cuomo, S., Di Cola, V. S., Giampaolo, F., Rozza, G., Raissi, M., and Piccialli, F.  Scientific machine learning through physics–informed neural networks: Where we are and what’s next.  _Journal of Scientific Computing_ , 92(3):88, 2022\.
  * Dangel et al. [2020a] Dangel, F., Harmeling, S., and Hennig, P.  Modular Block-diagonal Curvature Approximations for Feedforward Architectures.  In _International Conference on Artificial Intelligence and Statistics (AISTATS)_ , 2020a.
  * Dangel et al. [2020b] Dangel, F., Kunstner, F., and Hennig, P.  BackPACK: Packing more into Backprop.  In _International Conference on Learning Representations (ICLR)_ , 2020b.
  * Dangel et al. [2022] Dangel, F., Tatzel, L., and Hennig, P.  ViViT: Curvature access through the generalized gauss-newton’s low-rank structure.  _Transactions on Machine Learning Research (TMLR)_ , 2022.
  * Daw et al. [2022] Daw, A., Bu, J., Wang, S., Perdikaris, P., and Karpatne, A.  Rethinking the importance of sampling in physics-informed neural networks.  _arXiv preprint arXiv:2207.02338_ , 2022.
  * De Ryck et al. [2023] De Ryck, T., Bonnet, F., Mishra, S., and de Bézenac, E.  An operator preconditioning perspective on training in physics-informed machine learning.  _arXiv preprint arXiv:2310.05801_ , 2023.
  * Dissanayake & Phan-Thien [1994] Dissanayake, M. and Phan-Thien, N.  Neural-network-based approximations for solving partial differential equations.  _communications in Numerical Methods in Engineering_ , 10(3):195–201, 1994.
  * Drissi et al. [2024] Drissi, M., Keeble, J., Sarmiento, J. R., and Rios, A.  Second-order optimisation strategies for neural network quantum states.  _arXiv preprint arXiv:2401.17550_ , 2024.
  * E & Yu [2018] E, W. and Yu, B.  The deep ritz method: a deep learning-based numerical algorithm for solving variational problems.  _Communications in Mathematics and Statistics_ , 6(1):1–12, 2018.
  * Eschenhagen et al. [2023] Eschenhagen, R., Immer, A., Turner, R. E., Schneider, F., and Hennig, P.  Kronecker-Factored Approximate Curvature for Modern Neural Network Architectures.  In _Advances in Neural Information Processing Systems (NeurIPS)_ , 2023.
  * Griewank & Walther [2008] Griewank, A. and Walther, A.  _Evaluating derivatives: principles and techniques of algorithmic differentiation_.  SIAM, 2008.
  * Griewank et al. [1996] Griewank, A., Juedes, D., and Utke, J.  Algorithm 755: ADOL-C: A package for the automatic differentiation of algorithms written in C/C++.  _ACM Transactions on Mathematical Software (TOMS)_ , 22(2):131–167, 1996.
  * Grosse & Martens [2016] Grosse, R. and Martens, J.  A Kronecker-Factored Approximate Fisher Matrix for Convolution Layers.  In _International Conference on Machine Learning (ICML)_ , 2016.
  * Grosse et al. [2023] Grosse, R., Bae, J., Anil, C., Elhage, N., Tamkin, A., Tajdini, A., Steiner, B., Li, D., Durmus, E., Perez, E., et al.  Studying large language model generalization with influence functions.  _arXiv preprint arXiv:2308.03296_ , 2023.
  * Heskes [2000] Heskes, T.  On “natural” learning and pruning in multilayered perceptrons.  _Neural Computation_ , 12(4):881–901, 2000.
  * Jnini et al. [2024] Jnini, A., Vella, F., and Zeinhofer, M.  Gauss-Newton Natural Gradient Descent for Physics-Informed Computational Fluid Dynamics.  _arXiv preprint arXiv:2402.10680_ , 2024.
  * Johnson et al. [2021] Johnson, M. J., Bettencourt, J., Maclaurin, D., and Duvenaud, D.  Taylor-made higher-order automatic differentiation.  2021\.  URL https://github.com/google/jax/files/6717197/jet.pdf.  Accessed January 03, 2024.
  * Krishnapriyan et al. [2021] Krishnapriyan, A., Gholami, A., Zhe, S., Kirby, R., and Mahoney, M. W.  Characterizing possible failure modes in physics-informed neural networks.  _Advances in Neural Information Processing Systems_ , 34:26548–26560, 2021.
  * Kunstner et al. [2019] Kunstner, F., Hennig, P., and Balles, L.  Limitations of the empirical Fisher approximation for natural gradient descent.  _Advances in neural information processing systems_ , 32, 2019.
  * Lagaris et al. [1998] Lagaris, I. E., Likas, A., and Fotiadis, D. I.  Artificial neural networks for solving ordinary and partial differential equations.  _IEEE transactions on neural networks_ , 9(5):987–1000, 1998.
  * Li et al. [2023] Li, R., Ye, H., Jiang, D., Wen, X., Wang, C., Li, Z., Li, X., He, D., Chen, J., Ren, W., et al.  Forward Laplacian: A New Computational Framework for Neural Network-based Variational Monte Carlo.  2023\.
  * Li et al. [2024] Li, R., Wang, C., Ye, H., He, D., and Wang, L.  DOF: Accelerating high-order differential operators with forward propagation.  In _International Conference on Learning Representations (ICLR), Workshop on AI4DifferentialEquations In Science_ , 2024.
  * Lin et al. [2023a] Lin, W., Dangel, F., Eschenhagen, R., Neklyudov, K., Kristiadi, A., Turner, R. E., and Makhzani, A.  Structured Inverse-Free Natural Gradient: Memory-Efficient & Numerically-Stable KFAC for Large Neural Nets.  2023a.
  * Lin et al. [2023b] Lin, W., Duruisseaux, V., Leok, M., Nielsen, F., Khan, M. E., and Schmidt, M.  Simplifying Momentum-based Riemannian Submanifold Optimization.  2023b.
  * Liu et al. [2024] Liu, S., Su, C., Yao, J., Hao, Z., Su, H., Wu, Y., and Zhu, J.  Preconditioning for physics-informed neural networks.  _arXiv preprint arXiv:2402.00531_ , 2024.
  * Lu et al. [2021] Lu, L., Meng, X., Mao, Z., and Karniadakis, G. E.  DeepXDE: A deep learning library for solving differential equations.  _SIAM Review_ , 63(1):208–228, 2021.
  * Markidis [2021] Markidis, S.  The old and the new: Can physics-informed deep-learning replace traditional linear solvers?  _Frontiers in big Data_ , 4:669097, 2021.
  * Martens [2010] Martens, J.  Deep learning via Hessian-free optimization.  In _International Conference on Machine Learning (ICML)_ , 2010.
  * Martens [2020] Martens, J.  New insights and perspectives on the natural gradient method, 2020.
  * Martens & Grosse [2015] Martens, J. and Grosse, R.  Optimizing Neural Networks with Kronecker-factored Approximate Curvature.  In _International Conference on Machine Learning (ICML)_ , 2015.
  * Martens et al. [2018] Martens, J., Ba, J., and Johnson, M.  Kronecker-factored Curvature Approximations for Recurrent Neural Networks.  In _International Conference on Learning Representations_ , 2018.  URL https://openreview.net/forum?id=HyMTkQZAb.
  * Micikevicius et al. [2017] Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., Ginsburg, B., Houston, M., Kuchaiev, O., Venkatesh, G., et al.  Mixed precision training.  2017\.
  * Müller & Zeinhofer [2023] Müller, J. and Zeinhofer, M.  Achieving high accuracy with PINNs via energy natural gradient descent.  In _International Conference on Machine Learning_ , pp. 25471–25485. PMLR, 2023.
  * Müller & Zeinhofer [2024] Müller, J. and Zeinhofer, M.  Optimization in SciML–A Function Space Perspective.  _arXiv preprint arXiv:2402.07318_ , 2024.
  * Nabian et al. [2021] Nabian, M. A., Gladstone, R. J., and Meidani, H.  Efficient training of physics-informed neural networks via importance sampling.  _Computer-Aided Civil and Infrastructure Engineering_ , 36(8):962–977, 2021.
  * Osawa et al. [2023] Osawa, K., Li, S., and Hoefler, T.  Pipefisher: Efficient training of large language models using pipelining and Fisher information matrices.  _Proceedings of Machine Learning and Systems_ , 5, 2023.
  * Papyan [2019] Papyan, V.  Measurements of three-level hierarchical structure in the outliers in the spectrum of deepnet Hessians.  In _International Conference on Machine Learning (ICML)_ , 2019.
  * Paszke et al. [2019] Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S.  PyTorch: An Imperative Style, High-Performance Deep Learning Library.  In _Advances in Neural Information Processing Systems (NeurIPS)_. 2019.
  * Pauloski et al. [2021] Pauloski, J. G., Huang, Q., Huang, L., Venkataraman, S., Chard, K., Foster, I., and Zhang, Z.  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_ , pp. 1–14, 2021.
  * Pearlmutter [1994] Pearlmutter, B. A.  Fast Exact Multiplication by the Hessian.  _Neural Computation_ , 1994.
  * Petersen et al. [2023] Petersen, F., Sutter, T., Borgelt, C., Huh, D., Kuehne, H., Sun, Y., and Deussen, O.  ISAAC Newton: Input-based Approximate Curvature for Newton’s Method.  In _International Conference on Learning Representations (ICLR)_ , 2023.
  * Pfau et al. [2020] Pfau, D., Spencer, J. S., Matthews, A. G., and Foulkes, W. M. C.  Ab initio solution of the many-electron Schrödinger equation with deep neural networks.  _Physical Review Research_ , 2(3):033429, 2020\.
  * Raissi et al. [2019] Raissi, M., Perdikaris, P., and Karniadakis, G. E.  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\.
  * Schraudolph [2002] Schraudolph, N. N.  Fast curvature matrix-vector products for second-order gradient descent.  _Neural Computation_ , 2002.
  * Sirignano & Spiliopoulos [2018] Sirignano, J. and Spiliopoulos, K.  Dgm: A deep learning algorithm for solving partial differential equations.  _Journal of computational physics_ , 375:1339–1364, 2018\.
  * Tatzel et al. [2022] Tatzel, L., Hennig, P., and Schneider, F.  Late-Phase Second-Order Training.  In _Advances in Neural Information Processing Systems (NeurIPS), Workshop Has it Trained Yet?_ , 2022.
  * van der Meer et al. [2022] van der Meer, R., Oosterlee, C. W., and Borovykh, A.  Optimally weighted loss functions for solving PDEs with neural networks.  _Journal of Computational and Applied Mathematics_ , 405:113887, 2022.
  * Wang et al. [2021] Wang, S., Teng, Y., and Perdikaris, P.  Understanding and mitigating gradient flow pathologies in physics-informed neural networks.  _SIAM Journal on Scientific Computing_ , 43(5):A3055–A3081, 2021.
  * Wang et al. [2022a] Wang, S., Sankaran, S., and Perdikaris, P.  Respecting causality is all you need for training physics-informed neural networks.  _arXiv preprint arXiv:2203.07404_ , 2022a.
  * Wang et al. [2022b] Wang, S., Yu, X., and Perdikaris, P.  When and why PINNs fail to train: A neural tangent kernel perspective.  _Journal of Computational Physics_ , 449:110768, 2022b.
  * Weights & Biases [2020] Weights and Biases.  Experiment Tracking with Weights and Biases, 2020.  URL https://www.wandb.ai/.  Software available from wandb.ai.
  * Wu et al. [2023] Wu, C., Zhu, M., Tan, Q., Kartha, Y., and Lu, L.  A comprehensive study of non-adaptive and residual-based adaptive sampling for physics-informed neural networks.  _Computer Methods in Applied Mechanics and Engineering_ , 403:115671, 2023.
  * Zampini et al. [2024] Zampini, S., Zerbinati, U., Turkiyyah, G., and Keyes, D.  PETScML: Second-order solvers for training regression problems in Scientific Machine Learning.  _arXiv preprint arXiv:2403.12188_ , 2024.
  * Zapf et al. [2022] Zapf, B., Haubner, J., Kuchta, M., Ringstad, G., Eide, P. K., and Mardal, K.-A.  Investigating molecular transport in the human brain from MRI with physics-informed neural networks.  _Scientific Reports_ , 12(1):1–12, 2022.
  * Zeng et al. [2022] Zeng, Q., Kothari, Y., Bryngelson, S. H., and Schäfer, F.  Competitive physics informed networks.  _arXiv preprint arXiv:2204.11144_ , 2022.

(Supplementary Material)

[sections] [sections]l1

##  Appendix A Experimental Details and Additional Results

###  A.1 Hyper-Parameter Tuning Protocol

In all our experiments, we tune the following optimizer hyper-parameters and
otherwise use the PyTorch default values:

  * •

SGD: learning rate, momentum

  * •

Adam: learning rate

  * •

Hessian-free: type of curvature matrix (Hessian or GGN), damping, whether to
adapt damping over time (yes or no), maximum number of CG iterations

  * •

LBFGS: learning rate, history size

  * •

ENGD: damping, factor of the exponential moving average applied to the
Gramian, initialization of the Gramian (zero or identity matrix)

  * •

KFAC: factor of the exponential moving average applied to the Kronecker
factors, damping, momentum, initialization of the Kronecker factors (zero or
identity matrix)

  * •

KFAC*: factor of the exponential moving average applied to the Kronecker
factors, damping, initialization of the Kronecker factors (zero or identity
matrix)

Depending on the optimizer and experiment we use grid, random, or Bayesian
search from Weights & Biases to determine the hyper-parameters. Each
individual run is executed in double precision and allowed to run for a given
time budget, and we rank runs by the final $L_{2}$ error on a fixed evaluation
data set. To allow comparison, all runs are executed on RTX 6000 GPUs with 24
GiB of RAM. For grid and random searches, we use a round-based approach.
First, we choose a relatively wide search space and limit to approximately 50
runs. In a second round, we narrow down the hyper-parameter space based on the
first round, then re-run for another approximately 50 runs. We will release
the details of all hyper-parameter search spaces, as well as the hyper-
parameters for the best runs in our implementation.

###  A.2 2d Poisson Equation

##### Setup

We consider a two-dimensional Poisson equation $-\Delta
u(x,y)=2\pi^{2}\sin(\pi x)\sin(\pi y)$ on the unit square $(x,y)\in[0,1]^{2}$
with sine product right-hand side and zero boundary conditions $u(x,y)=0$ for
$(x,y)\in\partial[0,1]^{2}$. We choose a single set of training points with
$N_{\Omega}=900,N_{\partial\Omega}=120$. The $L_{2}$ error is evaluated on a
separate set of $9\,000$ data points using the known solution
$u_{\star}(x,y)=\sin(\pi x)\sin(\pi y)$. Each run is limited to a compute time
of $1\,000$ s. We compare three MLP architectures of increasing size, each of
whose linear layers are Tanh-activated except for the final one: a shallow
$2\to 64\to 1$ MLP with $D=257$ trainable parameters, a five layer $2\to 64\to
64\to 48\to 48\to 1$ MLP with $D=$9\,873$$ trainable parameters, and a five
layer $2\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,097$$ trainable
parameters. For the biggest architecture, full and per-layer ENGD lead to out-
of-memory errors and are thus not tested in the experiments. Figure A4
visualizes the results.

(a)

(b)

Figure A4:  Training loss and evaluation $L_{2}$ error for learning the
solution to a 2d Poisson equation over (4(a)) time and (4(b)) steps. Columns
are different neural networks.

##### Best run details

The runs shown in Figure A4 correspond to the following hyper-parameters:

  * •

$2\to 64\to 1$ MLP with $D=257$

    * –

SGD: learning rate: $1.805\,015\text{$⋅$}{10}^{-02}$, momentum: $0.99$

    * –

Adam: learning rate: $1.692\,339\text{$⋅$}{10}^{-03}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $100$, constant damping:
no, maximum CG iterations: $50$

    * –

LBFGS: learning rate: $0.5$, history size: $150$

    * –

ENGD (full): damping: $1\text{$⋅$}{10}^{-10}$, exponential moving average:
$0.3$, initialize Gramian to identity: yes

    * –

ENGD (layer-wise): damping: $0$, exponential moving average: $0.9$, initialize
Gramian to identity: yes

    * –

KFAC: damping: $1.544\,099\text{$⋅$}{10}^{-12}$, momentum:
$5.117\,575\text{$⋅$}{10}^{-01}$, exponential moving average:
$4.496\,490\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $1.215\,640\text{$⋅$}{10}^{-10}$, exponential moving average:
$9.263\,314\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

  * •

$2\to 64\to 64\to 48\to 48\to 1$ MLP with $D=$9\,873$$

    * –

SGD: learning rate: $3.758\,303\text{$⋅$}{10}^{-03}$, momentum: $0.9$

    * –

Adam: learning rate: $2.052\,448\text{$⋅$}{10}^{-04}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $0.1$, constant damping:
no, maximum CG iterations: $250$

    * –

LBFGS: learning rate: $0.1$, history size: $200$

    * –

ENGD (full): damping: $1\text{$⋅$}{10}^{-10}$, exponential moving average:
$0.6$, initialize Gramian to identity: no

    * –

ENGD (layer-wise): damping: $1\text{$⋅$}{10}^{-06}$, exponential moving
average: $0.3$, initialize Gramian to identity: no

    * –

KFAC: damping: $2.640\,390\text{$⋅$}{10}^{-11}$, momentum:
$9.995\,595\text{$⋅$}{10}^{-02}$, exponential moving average:
$5.556\,664\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $2.989\,247\text{$⋅$}{10}^{-13}$, exponential moving average:
$6.258\,340\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

  * •

$2\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,097$$

    * –

SGD: learning rate: $2.478\,674\text{$⋅$}{10}^{-03}$, momentum: $0.9$

    * –

Adam: learning rate: $6.406\,108\text{$⋅$}{10}^{-04}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $1$, constant damping:
no, maximum CG iterations: $250$

    * –

LBFGS: learning rate: $0.1$, history size: $225$

    * –

KFAC: damping: $1.710\,269\text{$⋅$}{10}^{-13}$, momentum:
$8.484\,996\text{$⋅$}{10}^{-01}$, exponential moving average:
$9.636\,460\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $1.232\,407\text{$⋅$}{10}^{-13}$, exponential moving average:
$9.488\,207\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

##### Search space details

The runs shown in Figure A4 were determined to be the best via a search with
approximately 50 runs on the following search spaces which were obtained by
refining an initially wider search ($\mathcal{U}$ denotes a uniform, and
$\mathcal{LU}$ a log-uniform distribution):

  * •

$2\to 64\to 1$ MLP with $D=257$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.1$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,5$;$0.1$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$100$,$1$,$0.01$,$0.000\,1$,$1\text{$⋅$}{10}^{-06}$\\})$,
constant damping: $\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG
iterations: $\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$75$,$100$,$125$,$150$,$175$,$200$,$225$,$250$\\})$

    * –

ENGD (full): damping:
$\mathcal{U}(\\{$1\text{$⋅$}{10}^{-06}$,$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-10}$,$1\text{$⋅$}{10}^{-12}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

ENGD (layer-wise): damping:
$\mathcal{U}(\\{$0.000\,1$,$1\text{$⋅$}{10}^{-06}$,$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-10}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity: yes

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, exponential
moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to
identity: yes

  * •

$2\to 64\to 64\to 48\to 48\to 1$ MLP with $D=$9\,873$$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.01$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,1$;$0.5$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$1$,$0.1$,$0.01$,$0.001$,$0.000\,1$\\})$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$50$,$75$,$100$,$125$,$150$,$175$,$200$,$225$\\})$

    * –

ENGD (full): damping:
$\mathcal{U}(\\{$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-09}$,$1\text{$⋅$}{10}^{-10}$,$1\text{$⋅$}{10}^{-11}$,$1\text{$⋅$}{10}^{-12}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

ENGD (layer-wise): damping:
$\mathcal{U}(\\{$0.01$,$0.001$,$0.000\,1$,$1\text{$⋅$}{10}^{-05}$,$1\text{$⋅$}{10}^{-06}$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, exponential
moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to
identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

  * •

$2\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,097$$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.01$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,1$;$0.5$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$1$,$0.1$,$0.01$,$0.001$,$0.000\,1$\\})$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$50$,$75$,$100$,$125$,$150$,$175$,$200$,$225$\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, exponential
moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to
identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

###  A.3 5d Poisson Equation

##### Setup

We consider a five-dimensional Poisson equation $-\Delta
u({\bm{x}})=\pi^{2}\sum_{i=1}^{5}\cos(\pi{x}_{i})$ on the five-dimensional
unit square ${\bm{x}}\in[0,1]^{5}$ with cosine sum right-hand side and
boundary conditions $u({\bm{x}})=\sum_{i=1}^{5}\cos(\pi{x}_{i})$ for
${\bm{x}}\in\partial[0,1]^{5}$. We sample training batches of size
$N_{\Omega}=$3\,000$,N_{\partial\Omega}=500$ and evaluate the $L_{2}$ error on
a separate set of $30\,000$ data points using the known solution
$u_{\star}({\bm{x}})=\sum_{i=1}^{5}\cos(\pi{x}_{i})$. All optimizers except
for KFAC sample a new training batch each iteration. KFAC only re-samples
every 100 iterations because we noticed significant improvement with multiple
iterations on a fixed batch. To make sure that this does not lead to an unfair
advantage of KFAC, we conduct an additional experiment where we also tune the
batch sampling frequency, as well as other hyper-parameters; see § A.5. The
results presented in this section are consistent with this additional
experiment (compare the rightmost column of Figure A5 and the leftmost column
of Figure A7). Each run is limited to 3000 s. We compare three MLP
architectures of increasing size, each of whose linear layers are Tanh-
activated except for the final one: a shallow $5\to 64\to 1$ MLP with $D=449$
trainable parameters, a five layer $5\to 64\to 64\to 48\to 48\to 1$ MLP with
$D=$10\,065$$ trainable parameters, and a five layer $5\to 256\to 256\to
128\to 128\to 1$ MLP with $D=$116\,864$$ trainable parameters. For the biggest
architecture, full and layer-wise ENGD lead to out-of-memory errors and are
thus not tested in the experiments. Figure A5 visualizes the results.

(a)

(b)

Figure A5: Training loss and evaluation $L_{2}$ error for learning the
solution to a 5d Poisson equation over (5(a)) time and (5(b)) steps. Columns
are different neural networks.

##### Best run details

The runs shown in Figure A5 correspond to the following hyper-parameters:

  * •

$5\to 64\to 1$ MLP with $D=449$

    * –

SGD: learning rate: $4.829\,757\text{$⋅$}{10}^{-03}$, momentum: $0.9$

    * –

Adam: learning rate: $1.527\,101\text{$⋅$}{10}^{-03}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $1$, constant damping:
no, maximum CG iterations: $250$

    * –

LBFGS: learning rate: $0.05$, history size: $125$

    * –

ENGD (full): damping: $1\text{$⋅$}{10}^{-11}$, exponential moving average:
$0.6$, initialize Gramian to identity: yes

    * –

ENGD (layer-wise): damping: $1\text{$⋅$}{10}^{-06}$, exponential moving
average: $0.6$, initialize Gramian to identity: yes

    * –

KFAC: damping: $3.030\,734\text{$⋅$}{10}^{-13}$, momentum:
$4.410\,155\text{$⋅$}{10}^{-01}$, exponential moving average:
$3.260\,663\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $9.835\,853\text{$⋅$}{10}^{-13}$, exponential moving average:
$7.714\,287\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

  * •

$5\to 64\to 64\to 48\to 48\to 1$ MLP with $D=$10\,065$$

    * –

SGD: learning rate: $1.007\,555\text{$⋅$}{10}^{-03}$, momentum: $0.9$

    * –

Adam: learning rate: $6.999\,994\text{$⋅$}{10}^{-04}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $0.1$, constant damping:
no, maximum CG iterations: $250$

    * –

LBFGS: learning rate: $0.1$, history size: $225$

    * –

ENGD (full): damping: $1\text{$⋅$}{10}^{-08}$, exponential moving average:
$0.3$, initialize Gramian to identity: no

    * –

ENGD (layer-wise): damping: $1\text{$⋅$}{10}^{-06}$, exponential moving
average: $0.3$, initialize Gramian to identity: no

    * –

KFAC: damping: $1.183\,063\text{$⋅$}{10}^{-12}$, momentum:
$9.058\,900\text{$⋅$}{10}^{-01}$, exponential moving average:
$9.588\,846\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $2.829\,461\text{$⋅$}{10}^{-10}$, exponential moving average:
$9.001\,393\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

  * •

$5\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,865$$

    * –

SGD: learning rate: $1.924\,173\text{$⋅$}{10}^{-03}$, momentum: $0.9$

    * –

Adam: learning rate: $5.416\,376\text{$⋅$}{10}^{-04}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $1$, constant damping:
no, maximum CG iterations: $50$

    * –

LBFGS: learning rate: $0.02$, history size: $225$

    * –

KFAC: damping: $1.844\,213\text{$⋅$}{10}^{-11}$, momentum:
$7.528\,559\text{$⋅$}{10}^{-01}$, exponential moving average:
$9.307\,849\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $2.183\,605\text{$⋅$}{10}^{-12}$, exponential moving average:
$9.563\,992\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

##### Search space details

The runs shown in Figure A5 were determined to be the best via a search with
approximately 50 runs on the following search spaces which were obtained by
refining an initially wider search ($\mathcal{U}$ denotes a uniform, and
$\mathcal{LU}$ a log-uniform distribution):

  * •

$5\to 64\to 1$ MLP with $D=449$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.01$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,1$;$0.5$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$1$,$0.1$,$0.01$,$0.001$,$0.000\,1$\\})$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$50$,$75$,$100$,$125$,$150$,$175$,$200$,$225$\\})$

    * –

ENGD (full): damping:
$\mathcal{U}(\\{$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-09}$,$1\text{$⋅$}{10}^{-10}$,$1\text{$⋅$}{10}^{-11}$,$1\text{$⋅$}{10}^{-12}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

ENGD (layer-wise): damping:
$\mathcal{U}(\\{$0.01$,$0.001$,$0.000\,1$,$1\text{$⋅$}{10}^{-05}$,$1\text{$⋅$}{10}^{-06}$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity: yes

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, exponential
moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to
identity: yes

  * •

$5\to 64\to 64\to 48\to 48\to 1$ MLP with $D=$10\,065$$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.01$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,1$;$0.5$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$1$,$0.1$,$0.01$,$0.001$,$0.000\,1$\\})$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$50$,$75$,$100$,$125$,$150$,$175$,$200$,$225$\\})$

    * –

ENGD (full): damping:
$\mathcal{U}(\\{$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-09}$,$1\text{$⋅$}{10}^{-10}$,$1\text{$⋅$}{10}^{-11}$,$1\text{$⋅$}{10}^{-12}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

ENGD (layer-wise): damping:
$\mathcal{U}(\\{$0.01$,$0.001$,$0.000\,1$,$1\text{$⋅$}{10}^{-05}$,$1\text{$⋅$}{10}^{-06}$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity: yes

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-12}$;$1\text{$⋅$}{10}^{-06}$])$, exponential
moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to
identity: yes

  * •

$5\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,865$$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.01$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,1$;$0.5$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$1$,$0.1$,$0.01$,$0.001$,$0.000\,1$\\})$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$50$,$75$,$100$,$125$,$150$,$175$,$200$,$225$\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity: yes

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-12}$;$1\text{$⋅$}{10}^{-06}$])$, exponential
moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to
identity: yes

###  A.4 10d Poisson Equation

##### Setup

We consider a 10-dimensional Poisson equation $-\Delta u({\bm{x}})=0$ on the
10-dimensional unit square ${\bm{x}}\in[0,1]^{5}$ with zero right-hand side
and harmonic mixed second order polynomial boundary conditions
$u({\bm{x}})=\sum_{i=1}^{\nicefrac{{d}}{{2}}}{x}_{2i-1}{x}_{2i}$ for
${\bm{x}}\in\partial[0,1]^{d}$. We sample training batches of size
$N_{\Omega}=$3\,000$,N_{\partial\Omega}=1000$ and evaluate the $L_{2}$ error
on a separate set of $30\,000$ data points using the known solution
$u_{\star}({\bm{x}})=\sum_{i=1}^{\nicefrac{{d}}{{2}}}{x}_{2i-1}{x}_{2i}$. All
optimizers except for KFAC sample a new training batch each iteration. KFAC
only re-samples every 100 iterations because we noticed significant
improvement with multiple iterations on a fixed batch. Each run is limited to
$6\,000$ s. We use a $10\to 256\to 256\to 128\to 128\to 1$ MLP with
$D=$118\,145$$ MLP whose linear layers are Tanh-activated except for the final
one. Figure A6 visualizes the results.

(a)

(b)

Figure A6: Training loss and evaluation $L_{2}$ error for learning the
solution to a 10d Poisson equation over (6(a)) time and (6(b)) steps.

##### Best run details

The runs shown in Figure A6 correspond to the following hyper-parameters:

  * •

SGD: learning rate: $6.488\,751\text{$⋅$}{10}^{-03}$, momentum: $0.9$

  * •

Adam: learning rate: $2.387\,088\text{$⋅$}{10}^{-04}$

  * •

Hessian-free: curvature matrix: GGN, initial damping: $0.001$, constant
damping: no, maximum CG iterations: $250$

  * •

LBFGS: learning rate: $0.2$, history size: $225$

  * •

KFAC: damping: $2.353\,358\text{$⋅$}{10}^{-05}$, momentum:
$9.633\,731\text{$⋅$}{10}^{-01}$, exponential moving average:
$3.908\,933\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

  * •

KFAC*: damping: $6.158\,042\text{$⋅$}{10}^{-10}$, exponential moving average:
$9.221\,738\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

##### Search space details

The runs shown in Figure A6 were determined to be the best via a Bayesian
search on the following search spaces which each optimizer given approximately
the same total computational time ($\mathcal{U}$ denotes a uniform, and
$\mathcal{LU}$ a log-uniform distribution):

  * •

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.01$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$

  * •

Adam: learning rate: $\mathcal{LU}([\text{5e-05};$0.005$])$

  * •

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$1$,$0.1$,$0.01$,$0.001$,$0.000\,1$\\})$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$50$,$250$\\})$

  * •

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$50$,$75$,$100$,$125$,$150$,$175$,$200$,$225$\\})$

  * •

KFAC: damping: $\mathcal{LU}([$1\text{$⋅$}{10}^{-10}$;$0.000\,1$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$

  * •

KFAC*: damping: $\mathcal{LU}([$1\text{$⋅$}{10}^{-10}$;$0.000\,1$])$,
exponential moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker
factors to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

###  A.5 5/10/100-d Poisson Equations with Bayesian Search

##### Setup

Here, we consider three Poisson equations $-\Delta u({\bm{x}})=f({\bm{x}})$
with different right-hand sides and boundary conditions on the unit square
${\bm{x}}\in[0,1]^{d}$:

  * •

$d=5$ with cosine sum right-hand side
$f({\bm{x}})=\pi^{2}\sum_{i=1}^{d}\cos(\pi{x}_{i})$, boundary conditions
$u({\bm{x}})=\sum_{i=1}^{d}\cos(\pi{x}_{i})$ for
${\bm{x}}\in\partial[0,1]^{d}$, and known solution
$u_{\star}({\bm{x}})=\sum_{i=1}^{d}\cos(\pi{x}_{i})$. We assign each run a
budget of $3\,000$ s.

  * •

$d=10$ with zero right-hand side $f({\bm{x}})=0$, harmonic mixed second order
polynomial boundary conditions
$u({\bm{x}})=\sum_{i=1}^{\nicefrac{{d}}{{2}}}{x}_{2i-1}{x}_{2i}$ for
${\bm{x}}\in\partial[0,1]^{d}$, and known solution
$u_{\star}({\bm{x}})=\sum_{i=1}^{\nicefrac{{d}}{{2}}}{x}_{2i-1}{x}_{2i}$. We
assign each run a budget of $6\,000$ s.

  * •

$d=100$ with constant non-zero right-hand side $f({\bm{x}})=-2d$, square norm
boundary conditions $u({\bm{x}})=\left\lVert{\bm{x}}\right\rVert_{2}^{2}$ for
${\bm{x}}\in\partial[0,1]^{d}$, and known solution
$u_{\star}({\bm{x}})=\left\lVert{\bm{x}}\right\rVert_{2}^{2}$. We assign each
run a budget of $10\,000$ s.

We tune the optimizer-hyperparameters described in § A.1, as well as the
batch sizes $N_{\Omega},N_{\partial\Omega}$, and their associated re-sampling
frequencies using Bayesian search. We use five layer MLP architectures with
varying widths whose layers are Tanh-activated except for the final layer.
These architectures are too large to be optimized by ENGD. Figure A7
visualizes the results.

(a)

(b)

Figure A7: Training loss and evaluation $L_{2}$ error for learning the
solution to high-dimensional Poisson equations over (7(a)) time and (7(b))
steps using Bayesian search.

##### Best run details

The runs shown in Figure A7 correspond to the following hyper-parameters:

  * •

5d Poisson equation, $5\to 256\to 256\to 128\to 128\to 1$ MLP with
$D=$116\,865$$

    * –

SGD: learning rate: $2.686\,653\text{$⋅$}{10}^{-04}$, momentum:
$9.878\,243\text{$⋅$}{10}^{-01}$, $N_{\Omega}$: $606$, $N_{\partial\Omega}$:
$2\,001$, batch sampling frequency: $1\,570$

    * –

Adam: learning rate: $6.587\,066\text{$⋅$}{10}^{-05}$, $N_{\Omega}$: $463$,
$N_{\partial\Omega}$: $2\,143$, batch sampling frequency: $5\,964$

    * –

Hessian-free: curvature matrix: GGN, initial damping:
$1.744\,291\text{$⋅$}{10}^{-02}$, constant damping: yes, maximum CG
iterations: $332$, $N_{\Omega}$: $332$, $N_{\partial\Omega}$: $1\,529$, batch
sampling frequency: $1\,412$

    * –

LBFGS: learning rate: $1.749\,124\text{$⋅$}{10}^{-01}$, history size: $339$,
$N_{\Omega}$: $5\,391$, $N_{\partial\Omega}$: $4\,768$, batch sampling
frequency: $155$

    * –

KFAC: damping: $4.251\,462\text{$⋅$}{10}^{-10}$, momentum:
$9.198\,986\text{$⋅$}{10}^{-01}$, exponential moving average:
$9.737\,093\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes, $N_{\Omega}$: $4\,690$, $N_{\partial\Omega}$: $2\,708$, batch sampling
frequency: $2\,369$

    * –

KFAC*: damping: $2.240\,865\text{$⋅$}{10}^{-12}$, exponential moving average:
$8.522\,194\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes, $N_{\Omega}$: $3\,149$, $N_{\partial\Omega}$: $3\,801$, batch sampling
frequency: $1\,393$

  * •

10d Poisson equation, $10\to 256\to 256\to 128\to 128\to 1$ MLP with
$D=$118\,145$$

    * –

SGD: learning rate: $5.805\,516\text{$⋅$}{10}^{-02}$, momentum:
$9.715\,522\text{$⋅$}{10}^{-01}$, $N_{\Omega}$: $537$, $N_{\partial\Omega}$:
$1\,173$, batch sampling frequency: $1\,083$

    * –

Adam: learning rate: $1.337\,679\text{$⋅$}{10}^{-04}$, $N_{\Omega}$: $115$,
$N_{\partial\Omega}$: $1\,960$, batch sampling frequency: $4\,975$

    * –

Hessian-free: curvature matrix: GGN, initial damping:
$3.719\,954\text{$⋅$}{10}^{-01}$, constant damping: no, maximum CG iterations:
$473$, $N_{\Omega}$: $1\,201$, $N_{\partial\Omega}$: $2\,402$, batch sampling
frequency: $2\,749$

    * –

LBFGS: learning rate: $1.868\,041\text{$⋅$}{10}^{-01}$, history size: $491$,
$N_{\Omega}$: $3\,408$, $N_{\partial\Omega}$: $1\,658$, batch sampling
frequency: $1\,362$

    * –

KFAC: damping: $2.093\,830\text{$⋅$}{10}^{-04}$, momentum:
$7.653\,280\text{$⋅$}{10}^{-01}$, exponential moving average:
$9.265\,761\text{$⋅$}{10}^{-02}$, initialize Kronecker factors to identity:
yes, $N_{\Omega}$: $654$, $N_{\partial\Omega}$: $2\,467$, batch sampling
frequency: $1\,484$

    * –

KFAC*: damping: $8.440\,601\text{$⋅$}{10}^{-10}$, exponential moving average:
$8.354\,459\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes, $N_{\Omega}$: $3\,614$, $N_{\partial\Omega}$: $2\,151$, batch sampling
frequency: $292$

  * •

100d Poisson equation, $100\to 768\to 768\to 512\to 512\to 1$ MLP with
$D=$1\,325\,057$$

    * –

SGD: learning rate: $1.450\,764\text{$⋅$}{10}^{-03}$, momentum:
$9.747\,671\text{$⋅$}{10}^{-01}$, $N_{\Omega}$: $177$, $N_{\partial\Omega}$:
$2\,422$, batch sampling frequency: $519$

    * –

Adam: learning rate: $4.172\,759\text{$⋅$}{10}^{-04}$, $N_{\Omega}$: $1\,013$,
$N_{\partial\Omega}$: $1\,171$, batch sampling frequency: $353$

    * –

Hessian-free: curvature matrix: GGN, initial damping:
$7.011\,290\text{$⋅$}{10}^{-05}$, constant damping: no, maximum CG iterations:
$488$, $N_{\Omega}$: $1\,018$, $N_{\partial\Omega}$: $1\,946$, batch sampling
frequency: $755$

    * –

LBFGS: learning rate: $1.797\,096\text{$⋅$}{10}^{-01}$, history size: $112$,
$N_{\Omega}$: $2\,115$, $N_{\partial\Omega}$: $1\,852$, batch sampling
frequency: $23$

    * –

KFAC: damping: $1.357\,141\text{$⋅$}{10}^{-03}$, momentum:
$5.370\,282\text{$⋅$}{10}^{-01}$, exponential moving average:
$7.860\,784\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
no, $N_{\Omega}$: $131$, $N_{\partial\Omega}$: $2\,354$, batch sampling
frequency: $50$

    * –

KFAC*: damping: $3.078\,398\text{$⋅$}{10}^{-04}$, exponential moving average:
$5.796\,979\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes, $N_{\Omega}$: $516$, $N_{\partial\Omega}$: $2\,382$, batch sampling
frequency: $12$

##### Search space details

The runs shown in Figure A7 were determined to be the best via a Bayesian
search on the following search spaces which each optimizer given approximately
the same total computational time ($\mathcal{U}$ denotes a uniform, and
$\mathcal{LU}$ a log-uniform distribution):

  * •

5d Poisson equation, $5\to 256\to 256\to 128\to 128\to 1$ MLP with
$D=$116\,865$$

    * –

SGD: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$,
$N_{\Omega}$: $\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$1$])$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$1$,$2$,\text{\ldots},$500$\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

    * –

LBFGS: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$, history
size: $\mathcal{U}(\\{$5$,$6$,\text{\ldots},$500$\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

    * –

KFAC: damping: $\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$0.01$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

    * –

KFAC*: damping: $\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$0.01$])$, exponential
moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to
identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

  * •

10d Poisson equation, $10\to 256\to 256\to 128\to 128\to 1$ MLP with
$D=$118\,145$$

    * –

SGD: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$, $N_{\partial\Omega}$:
$\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$, batch sampling
frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$5\,000$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$,
$N_{\Omega}$: $\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$5\,000$\\})$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$1$])$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$1$,$2$,\text{\ldots},$500$\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$, $N_{\partial\Omega}$:
$\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$, batch sampling
frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$5\,000$\\})$

    * –

LBFGS: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$, history
size: $\mathcal{U}(\\{$5$,$6$,\text{\ldots},$500$\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$, $N_{\partial\Omega}$:
$\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$, batch sampling
frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$5\,000$\\})$

    * –

KFAC: damping: $\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$0.01$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$, $N_{\partial\Omega}$:
$\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$, batch sampling
frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$5\,000$\\})$

    * –

KFAC*: damping: $\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$0.01$])$, exponential
moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to
identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$, $N_{\partial\Omega}$:
$\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$, batch sampling
frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$5\,000$\\})$

  * •

100d Poisson equation, $100\to 768\to 768\to 512\to 512\to 1$ MLP with
$D=$1\,325\,057$$

    * –

SGD: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$, $N_{\partial\Omega}$:
$\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$, batch sampling
frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$1\,000$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$,
$N_{\Omega}$: $\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$1\,000$\\})$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$1$])$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$1$,$2$,\text{\ldots},$500$\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$, $N_{\partial\Omega}$:
$\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$, batch sampling
frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$1\,000$\\})$

    * –

LBFGS: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$, history
size: $\mathcal{U}(\\{$5$,$6$,\text{\ldots},$500$\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$, $N_{\partial\Omega}$:
$\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$, batch sampling
frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$1\,000$\\})$

    * –

KFAC: damping: $\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$0.01$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$, $N_{\partial\Omega}$:
$\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$, batch sampling
frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$1\,000$\\})$

    * –

KFAC*: damping: $\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$0.01$])$, exponential
moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to
identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$5\,000$\\})$, $N_{\partial\Omega}$:
$\mathcal{U}(\\{$50$,$51$,\text{\ldots},$2\,500$\\})$, batch sampling
frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$1\,000$\\})$

###  A.6 PINN Loss for the Heat Equation

Consider the $(\tilde{d}+1)$-dimensional homogeneous heat equation

$\displaystyle\partial_{t}u(t,\tilde{{\bm{x}}})-\kappa\Delta_{\tilde{{\bm{x}}}}u(t,\tilde{{\bm{x}}})=0$

with spatial coordinates
$\tilde{{\bm{x}}}\in\Omega\subseteq{\mathbb{R}}^{\tilde{d}}$ and time
coordinate $t\in\mathrm{T}\subseteq{\mathbb{R}}$ where $\mathrm{T}$ is a time
interval and $\kappa>0$ denotes the heat conductivity. In this case, our
neural network processes a $(d=\tilde{d}+1)$-dimensional vector
${\bm{x}}=(t,\tilde{{\bm{x}}}^{\top})^{\top}\in{\mathbb{R}}^{d}$ and we can
re-write the heat equation as

$\displaystyle\partial_{{x}_{1}}u({\bm{x}})-\kappa\sum_{d=2}^{d}\Delta_{{x}_{d}}u({\bm{x}})=0\,.$

In the following, we consider the unit time interval $\mathrm{T}=[0;1]$, the
unit square $\Omega=[0;1]^{\tilde{d}}$ and set $\kappa=\nicefrac{{1}}{{4}}$.
There are two types of constraints we need to enforce on the heat equation in
order to obtain unique solutions: initial conditions and boundary conditions.
As our framework for the KFAC approximation assumes only two terms in the loss
function, we combine the contributions from the boundary and initial values
into one term.

To make this more precise, consider the following example solution of the heat
equation, which will be used later on as well. As initial conditions, we use
$u_{0}(\tilde{{\bm{x}}})=u(0,\tilde{{\bm{x}}})=\prod_{i=1}^{\tilde{d}}\sin(\pi\tilde{{x}}_{i})$
for $\tilde{{\bm{x}}}\in\Omega$. For boundary conditions, we use
$g(t,\tilde{{\bm{x}}})=0$ for
$(t,\tilde{{\bm{x}}})\in\mathrm{T}\times\partial\Omega$. The manufactured
solution is

$\displaystyle
u_{\star}(t,\tilde{{\bm{x}}})=\exp\left(-\frac{\pi^{2}\tilde{d}t}{4}\right)\prod_{i=1}^{\tilde{d}}\sin(\pi[\tilde{{x}}_{i}])\,.$

The PINN loss for this problem consists of three terms: a PDE term, an initial
value condition term, and a spatial boundary condition term,

$\displaystyle L({\bm{\theta}})$
$\displaystyle=\frac{1}{N_{\Omega}}\sum_{n=1}^{N_{\Omega}}\left\lVert\partial_{t}u_{{\bm{\theta}}}({\bm{x}}_{n}^{\Omega})-\frac{1}{4}\Delta_{\tilde{{\bm{x}}}_{n}}u_{{\bm{\theta}}}({\bm{x}}_{n}^{\Omega})\right\rVert^{2}_{2}$
$\displaystyle+\frac{1}{N_{\partial\Omega}}\sum_{n=1}^{N_{\partial\Omega}}\left\lVert
u_{{\bm{\theta}}}({\bm{x}}_{n}^{\partial\Omega})-g({\bm{x}}_{n}^{\partial\Omega})\right\rVert^{2}_{2}$
$\displaystyle+\frac{1}{N_{0}}\sum_{n=1}^{N_{0}}\left\lVert
u_{{\bm{\theta}}}(0,{\bm{x}}_{n}^{0})-u_{0}({\bm{x}}_{n}^{0})\right\rVert^{2}_{2}$

with ${\bm{x}}_{n}^{\Omega}\sim\mathrm{T}\times\Omega$, and
${\bm{x}}_{n}^{\partial\Omega}\sim\mathrm{T}\times\partial\Omega$, and
${\bm{x}}_{n}^{0}\sim\\{0\\}\times\Omega$. To fit this loss into our framework
which assumes two loss terms, each of whose curvature is approximated with a
Kronecker factor, we combine the initial value and boundary value conditions
into a single term. Assuming
$N_{\partial\Omega}=N_{0}=\nicefrac{{N_{\text{cond}}}}{{2}}$ without loss of
generality, we write

$\displaystyle L({\bm{\theta}})$
$\displaystyle=\underbrace{\frac{1}{N_{\Omega}}\sum_{n=1}^{N_{\Omega}}\left\lVert\partial_{t}u_{{\bm{\theta}}}({\bm{x}}_{n}^{\Omega})-\frac{1}{4}\Delta_{\tilde{{\bm{x}}}_{n}}u_{{\bm{\theta}}}({\bm{x}}_{n}^{\Omega})-y_{n}^{\Omega}\right\rVert^{2}_{2}}_{L_{\Omega}({\bm{\theta}})}+\underbrace{\frac{1}{N_{\text{cond}}}\sum_{n=1}^{N_{\text{cond}}}\left\lVert
u_{{\bm{\theta}}}({\bm{x}}_{n}^{\text{cond}})-y_{n}^{\text{cond}}\right\rVert^{2}_{2}}_{L_{\text{cond}}({\bm{\theta}})}$

with domain inputs ${\bm{x}}_{n}^{\Omega}\sim\mathrm{T}\times\Omega$ and
targets $y_{n}^{\Omega}=0$, boundary and initial condition targets
$y_{n}^{\text{cond}}=u_{\star}({\bm{x}}_{n}^{\text{cond}})$ with initial
inputs ${\bm{x}}_{n}^{\text{cond}}\sim\\{0\\}\times\Omega$ for
$n=1,\dots,\nicefrac{{N_{\text{cond}}}}{{2}}$ and boundary inputs
${\bm{x}}_{n}^{\text{cond}}\sim\mathrm{T}\times\partial\Omega$ for
$n=\nicefrac{{N_{\text{cond}}}}{{2}}+1,\dots,N_{\text{cond}}$. This loss has
the same structure as the PINN loss in Equation 1.

###  A.7 1+1d Heat Equation

##### Setup

We consider a 1+1-dimensional heat equation
$\partial_{t}u(t,x)-\kappa\Delta_{x}u(t,x)=0$ with
$\kappa=\nicefrac{{1}}{{4}}$ on the unit square and unit time interval,
$x,t\in[0,1]\times[0,1]$. The equation has zero spatial boundary conditions
and the initial values are given by $u(0,x)=\sin(\pi x)$ for
${\bm{x}}\in[0,1]$. We sample a single training batch of size
$N_{\Omega}=$900$,N_{\partial\Omega}=120$
($\nicefrac{{N_{\partial\Omega}}}{{2}}$ points for the initial value and
spatial boundary conditions each) and evaluate the $L_{2}$ error on a separate
set of $9\,000$ data points using the known solution
$u_{\star}(t,x)=\exp(-\nicefrac{{\pi^{2}t}}{{4}})\sin(\pi x)$. Each run is
limited to $1\,000$ s. We compare three MLP architectures of increasing size,
each of whose linear layers are Tanh-activated except for the final one: a
shallow $2\to 64\to 1$ MLP with $D=257$ trainable parameters, a five layer
$2\to 64\to 64\to 48\to 48\to 1$ MLP with $D=$9\,873$$ trainable parameters,
and a five layer $2\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,097$$
trainable parameters. For the biggest architecture, full and layer-wise ENGD
lead to out-of-memory errors and are thus not part of the experiments. Figure
Figure A8 summarizes the results.

(a)

(b)

Figure A8: training loss and evaluation $L_{2}$ error for learning the
solution to a 1+1-dimensional heat equation over (8(a)) time and (8(b)). each
column corresponds to a different neural network.

##### Best run details

The runs shown in Figure A8 correspond to the following hyper-parameters:

  * •

$2\to 64\to 1$ MLP with $D=257$

    * –

SGD: learning rate: $1.752\,752\text{$⋅$}{10}^{-02}$, momentum: $0.99$

    * –

Adam: learning rate: $8.629\,006\text{$⋅$}{10}^{-04}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $0.000\,1$, constant
damping: no, maximum CG iterations: $250$

    * –

LBFGS: learning rate: $0.1$, history size: $125$

    * –

ENGD (full): damping: $1\text{$⋅$}{10}^{-12}$, exponential moving average:
$0.9$, initialize Gramian to identity: no

    * –

ENGD (layer-wise): damping: $1\text{$⋅$}{10}^{-10}$, exponential moving
average: $0.3$, initialize Gramian to identity: no

    * –

KFAC: damping: $1.273\,754\text{$⋅$}{10}^{-08}$, momentum:
$7.562\,617\text{$⋅$}{10}^{-01}$, exponential moving average:
$3.611\,724\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $1.968\,427\text{$⋅$}{10}^{-09}$, exponential moving average:
$9.703\,638\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

  * •

$2\to 64\to 64\to 48\to 48\to 1$ MLP with $D=$9\,873$$

    * –

SGD: learning rate: $9.276\,977\text{$⋅$}{10}^{-02}$, momentum: $0.99$

    * –

Adam: learning rate: $2.551\,515\text{$⋅$}{10}^{-03}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $0.01$, constant
damping: no, maximum CG iterations: $250$

    * –

LBFGS: learning rate: $0.2$, history size: $125$

    * –

ENGD (full): damping: $1\text{$⋅$}{10}^{-06}$, exponential moving average:
$0.9$, initialize Gramian to identity: no

    * –

ENGD (layer-wise): damping: $1\text{$⋅$}{10}^{-08}$, exponential moving
average: $0.6$, initialize Gramian to identity: no

    * –

KFAC: damping: $3.169\,186\text{$⋅$}{10}^{-13}$, momentum:
$7.075\,879\text{$⋅$}{10}^{-01}$, exponential moving average:
$8.860\,410\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $5.035\,695\text{$⋅$}{10}^{-14}$, exponential moving average:
$9.815\,164\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

  * •

$2\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,097$$

    * –

SGD: learning rate: $5.709\,474\text{$⋅$}{10}^{-02}$, momentum: $0.99$

    * –

Adam: learning rate: $6.716\,485\text{$⋅$}{10}^{-04}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $0.01$, constant
damping: no, maximum CG iterations: $250$

    * –

LBFGS: learning rate: $0.2$, history size: $125$

    * –

KFAC: damping: $2.576\,488\text{$⋅$}{10}^{-13}$, momentum:
$2.043\,395\text{$⋅$}{10}^{-02}$, exponential moving average:
$9.727\,829\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $7.343\,493\text{$⋅$}{10}^{-11}$, exponential moving average:
$9.765\,844\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

##### Search space details

The runs shown in Figure A8 were determined to be the best via a search with
approximately 50 runs on the following search spaces which were obtained by
refining an initially wider search ($\mathcal{U}$ denotes a uniform, and
$\mathcal{LU}$ a log-uniform distribution):

  * •

$2\to 64\to 1$ MLP with $D=257$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.1$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,5$;$0.1$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$100$,$1$,$0.01$,$0.000\,1$,$1\text{$⋅$}{10}^{-06}$\\})$,
constant damping: $\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG
iterations: $\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$75$,$100$,$125$,$150$,$175$,$200$,$225$,$250$\\})$

    * –

ENGD (full): damping:
$\mathcal{U}(\\{$1\text{$⋅$}{10}^{-06}$,$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-10}$,$1\text{$⋅$}{10}^{-12}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

ENGD (layer-wise): damping:
$\mathcal{U}(\\{$0.000\,1$,$1\text{$⋅$}{10}^{-06}$,$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-10}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, momentum:
$\mathcal{U}([$0.5$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity: yes

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, exponential
moving average: $\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to
identity: yes

  * •

$2\to 64\to 64\to 48\to 48\to 1$ MLP with $D=$9\,873$$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.1$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,5$;$0.1$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$100$,$1$,$0.01$,$0.000\,1$,$1\text{$⋅$}{10}^{-06}$\\})$,
constant damping: $\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG
iterations: $\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$75$,$100$,$125$,$150$,$175$,$200$,$225$,$250$\\})$

    * –

ENGD (full): damping:
$\mathcal{U}(\\{$1\text{$⋅$}{10}^{-06}$,$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-10}$,$1\text{$⋅$}{10}^{-12}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

ENGD (layer-wise): damping:
$\mathcal{U}(\\{$0.000\,1$,$1\text{$⋅$}{10}^{-06}$,$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-10}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to identity: yes

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$1\text{$⋅$}{10}^{-09}$])$, exponential
moving average: $\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to
identity: yes

  * •

$2\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,097$$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.1$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,5$;$0.1$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$100$,$1$,$0.01$,$0.000\,1$,$1\text{$⋅$}{10}^{-06}$\\})$,
constant damping: $\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG
iterations: $\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$75$,$100$,$125$,$150$,$175$,$200$,$225$,$250$\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-14}$;$1\text{$⋅$}{10}^{-07}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to identity: yes

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-14}$;$1\text{$⋅$}{10}^{-06}$])$, exponential
moving average: $\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to
identity: yes

###  A.8 4+1d Heat Equation

##### Setup

We consider a 4+1-dimensional heat equation
$\partial_{t}u(t,{\bm{x}})-\kappa\Delta_{{\bm{x}}}u(t,{\bm{x}})=0$ with
$\kappa=\nicefrac{{1}}{{4}}$ on the four-dimensional unit square and unit time
interval, ${\bm{x}},t\in[0,1]^{4}\times[0,1]$. The equation has spatial
boundary conditions $u(t,x)=\exp(-t)\sum_{i=1}^{4}\sin(2{x}_{i})$ for
$t,{\bm{x}}\in[0,1]\times\partial[0,1]^{4}$ throughout time, and initial value
conditions $u(0,{\bm{x}})=\sum_{i=1}^{4}\sin(2{x}_{i})$ for
${\bm{x}}\in[0,1]^{4}$. We sample training batches of size
$N_{\Omega}=$3\,000$,N_{\partial\Omega}=500$
($\nicefrac{{N_{\partial\Omega}}}{{2}}$ points for the initial value and
spatial boundary conditions each) and evaluate the $L_{2}$ error on a separate
set of $30\,000$ data points using the known solution
$u_{\star}(t,{\bm{x}})=\exp(-t)\sum_{i=1}^{4}\sin(2{x}_{i})$. All optimizers
except for KFAC sample a new training batch each iteration. KFAC only re-
samples every 100 iterations because we noticed significant improvement with
multiple iterations on a fixed batch. To make sure that this does not lead to
an unfair advantage of KFAC, we conduct an additional experiment where we also
tune the batch sampling frequency, as well as other hyper-parameters; see §
A.9. The results presented in this section are consistent with this additional
experiment (compare the rightmost column of Figure A9 and Figure A10). Each
run is limited to 3000 s. We compare three MLP architectures of increasing
size, each of whose linear layers are Tanh-activated except for the final one:
a shallow $5\to 64\to 1$ MLP with $D=449$ trainable weights, a five layer
$5\to 64\to 64\to 48\to 48\to 1$ MLP with $D=$10\,065$$ trainable weights, and
a five layer $5\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,864$$
trainable weights. For the biggest architecture, full and layer-wise ENGD lead
to out-of-memory errors and are thus not tested. Figure A9 visualizes the
results.

(a)

(b)

Figure A9: Training loss and evaluation $L_{2}$ error for learning the
solution to a 4+1-d heat equation over (9(a)) time and (9(b)) steps. Columns
are different neural networks.

##### Search space details

The runs shown in Figure A9 were determined to be the best via a search with
approximately 50 runs on the following search spaces which were obtained by
refining an initially wider search ($\mathcal{U}$ denotes a uniform, and
$\mathcal{LU}$ a log-uniform distribution):

  * •

$5\to 64\to 1$ MLP with $D=449$

    * –

SGD: learning rate: $7.737\,742\text{$⋅$}{10}^{-03}$, momentum: $0.9$

    * –

Adam: learning rate: $3.708\,460\text{$⋅$}{10}^{-03}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $0.1$, constant damping:
no, maximum CG iterations: $250$

    * –

LBFGS: learning rate: $0.2$, history size: $175$

    * –

ENGD (full): damping: $1\text{$⋅$}{10}^{-10}$, exponential moving average:
$0.6$, initialize Gramian to identity: yes

    * –

ENGD (layer-wise): damping: $1\text{$⋅$}{10}^{-06}$, exponential moving
average: $0$, initialize Gramian to identity: yes

    * –

KFAC: damping: $1.000\,288\text{$⋅$}{10}^{-09}$, momentum:
$9.474\,108\text{$⋅$}{10}^{-01}$, exponential moving average:
$7.783\,519\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $2.965\,060\text{$⋅$}{10}^{-08}$, exponential moving average:
$9.574\,717\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

  * •

$5\to 64\to 64\to 48\to 48\to 1$ MLP with $D=$10\,065$$

    * –

SGD: learning rate: $9.357\,973\text{$⋅$}{10}^{-03}$, momentum: $0.9$

    * –

Adam: learning rate: $7.801\,748\text{$⋅$}{10}^{-04}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $0.001$, constant
damping: no, maximum CG iterations: $250$

    * –

LBFGS: learning rate: $0.1$, history size: $225$

    * –

ENGD (full): damping: $1\text{$⋅$}{10}^{-08}$, exponential moving average:
$0$, initialize Gramian to identity: yes

    * –

ENGD (layer-wise): damping: $1\text{$⋅$}{10}^{-06}$, exponential moving
average: $0.3$, initialize Gramian to identity: no

    * –

KFAC: damping: $4.143\,385\text{$⋅$}{10}^{-14}$, momentum:
$7.660\,303\text{$⋅$}{10}^{-01}$, exponential moving average:
$9.821\,414\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $1.955\,740\text{$⋅$}{10}^{-10}$, exponential moving average:
$9.821\,778\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

  * •

$5\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,865$$

    * –

SGD: learning rate: $7.192\,473\text{$⋅$}{10}^{-03}$, momentum: $0.9$

    * –

Adam: learning rate: $5.266\,284\text{$⋅$}{10}^{-04}$

    * –

Hessian-free: curvature matrix: GGN, initial damping: $0.000\,1$, constant
damping: no, maximum CG iterations: $250$

    * –

LBFGS: learning rate: $0.2$, history size: $200$

    * –

KFAC: damping: $8.581\,322\text{$⋅$}{10}^{-13}$, momentum:
$8.501\,747\text{$⋅$}{10}^{-01}$, exponential moving average:
$9.803\,115\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

    * –

KFAC*: damping: $3.405\,440\text{$⋅$}{10}^{-14}$, exponential moving average:
$8.445\,471\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes

##### Search space details

The runs shown in Figure A9 were determined to be the best via a search with
approximately 50 runs on the following search spaces which were obtained by
refining an initially wider search ($\mathcal{U}$ denotes a uniform, and
$\mathcal{LU}$ a log-uniform distribution):

  * •

$5\to 64\to 1$ MLP with $D=449$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.01$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,5$;$0.1$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$1$,$0.1$,$0.01$,$0.001$,$0.000\,1$\\})$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$50$,$75$,$100$,$125$,$150$,$175$,$200$,$225$\\})$

    * –

ENGD (full): damping:
$\mathcal{U}(\\{$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-09}$,$1\text{$⋅$}{10}^{-10}$,$1\text{$⋅$}{10}^{-11}$,$1\text{$⋅$}{10}^{-12}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

ENGD (layer-wise): damping:
$\mathcal{U}(\\{$0.01$,$0.001$,$0.000\,1$,$1\text{$⋅$}{10}^{-05}$,$1\text{$⋅$}{10}^{-06}$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-12}$;$1\text{$⋅$}{10}^{-06}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to identity: yes

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-13}$;$1\text{$⋅$}{10}^{-07}$])$, exponential
moving average: $\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to
identity: yes

  * •

$5\to 64\to 64\to 48\to 48\to 1$ MLP with $D=$10\,065$$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.01$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,5$;$0.1$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$1$,$0.1$,$0.01$,$0.001$,$0.000\,1$\\})$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$50$,$75$,$100$,$125$,$150$,$175$,$200$,$225$\\})$

    * –

ENGD (full): damping:
$\mathcal{U}(\\{$1\text{$⋅$}{10}^{-08}$,$1\text{$⋅$}{10}^{-09}$,$1\text{$⋅$}{10}^{-10}$,$1\text{$⋅$}{10}^{-11}$,$1\text{$⋅$}{10}^{-12}$,$0$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

ENGD (layer-wise): damping:
$\mathcal{U}(\\{$0.01$,$0.001$,$0.000\,1$,$1\text{$⋅$}{10}^{-05}$,$1\text{$⋅$}{10}^{-06}$\\})$,
exponential moving average: $\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$,$0.99$\\})$,
initialize Gramian to identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-14}$;$1\text{$⋅$}{10}^{-08}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to identity: yes

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-14}$;$1\text{$⋅$}{10}^{-08}$])$, exponential
moving average: $\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to
identity: yes

  * •

$5\to 256\to 256\to 128\to 128\to 1$ MLP with $D=$116\,865$$

    * –

SGD: learning rate: $\mathcal{LU}([$0.001$;$0.01$])$, momentum:
$\mathcal{U}(\\{$0$,$0.3$,$0.6$,$0.9$\\})$

    * –

Adam: learning rate: $\mathcal{LU}([$0.000\,5$;$0.1$])$

    * –

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{U}(\\{$1$,$0.1$,$0.01$,$0.001$,$0.000\,1$\\})$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$50$,$250$\\})$

    * –

LBFGS: learning rate:
$\mathcal{U}(\\{$0.5$,$0.2$,$0.1$,$0.05$,$0.02$,$0.01$\\})$, history size:
$\mathcal{U}(\\{$50$,$75$,$100$,$125$,$150$,$175$,$200$,$225$\\})$

    * –

KFAC: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-14}$;$1\text{$⋅$}{10}^{-08}$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to identity: yes

    * –

KFAC*: damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-14}$;$1\text{$⋅$}{10}^{-08}$])$, exponential
moving average: $\mathcal{U}([$0.5$;$0.99$])$, initialize Kronecker factors to
identity: yes

###  A.9 4+1d Heat Equation with Bayesian Search

##### Setup

We consider the same heat equation as in § A.8 and use the $5\to 256\to
256\to 128\to 128\to 1$ MLP with $D=$116\,865$$. We tune all optimizer hyper-
parameters as described in § A.1 and also tune the batch sizes
$N_{\Omega},N_{\partial\Omega}$, as well as their re-sampling frequencies.
Figure A10 summarizes the results.

(a)

(b)

Figure A10: Training loss and evaluation $L_{2}$ error for learning the
solution to a 4+1-dimensional heat equation over (10(a)) time and (10(b))
using Bayesian search.

##### Best run details

The runs shown in Figure A10 correspond to the following hyper-parameters:

  * •

SGD: learning rate: $1.614\,965\text{$⋅$}{10}^{-02}$, momentum:
$9.899\,167\text{$⋅$}{10}^{-01}$, $N_{\Omega}$: $527$, $N_{\partial\Omega}$:
$2\,157$, batch sampling frequency: $543$

  * •

Adam: learning rate: $2.583\,569\text{$⋅$}{10}^{-04}$, $N_{\Omega}$: $472$,
$N_{\partial\Omega}$: $3\,018$, batch sampling frequency: $177$

  * •

Hessian-free: curvature matrix: GGN, initial damping:
$2.703\,390\text{$⋅$}{10}^{-04}$, constant damping: yes, maximum CG
iterations: $301$, $N_{\Omega}$: $3\,800$, $N_{\partial\Omega}$: $4\,373$,
batch sampling frequency: $4\,319$

  * •

LBFGS: learning rate: $1.029\,194\text{$⋅$}{10}^{-01}$, history size: $488$,
$N_{\Omega}$: $582$, $N_{\partial\Omega}$: $2\,038$, batch sampling frequency:
$315$

  * •

KFAC: damping: $6.569\,698\text{$⋅$}{10}^{-15}$, momentum:
$6.058\,520\text{$⋅$}{10}^{-02}$, exponential moving average:
$9.701\,174\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes, $N_{\Omega}$: $865$, $N_{\partial\Omega}$: $4\,673$, batch sampling
frequency: $728$

  * •

KFAC*: damping: $8.837\,871\text{$⋅$}{10}^{-15}$, exponential moving average:
$9.887\,596\text{$⋅$}{10}^{-01}$, initialize Kronecker factors to identity:
yes, $N_{\Omega}$: $2\,563$, $N_{\partial\Omega}$: $2\,873$, batch sampling
frequency: $9\,647$

##### Search space details

The runs shown in Figure A10 were determined to be the best via a Bayesian
search on the following search spaces which each optimizer given approximately
the same total computational time ($\mathcal{U}$ denotes a uniform, and
$\mathcal{LU}$ a log-uniform distribution):

  * •

SGD: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

  * •

Adam: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$,
$N_{\Omega}$: $\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

  * •

Hessian-free: curvature matrix:
$\mathcal{U}(\\{\text{GGN},\text{Hessian}\\})$, initial damping:
$\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$1$])$, constant damping:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, maximum CG iterations:
$\mathcal{U}(\\{$1$,$2$,\text{\ldots},$500$\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

  * •

LBFGS: learning rate: $\mathcal{LU}([$1\text{$⋅$}{10}^{-06}$;$1$])$, history
size: $\mathcal{U}(\\{$5$,$6$,\text{\ldots},$500$\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

  * •

KFAC: damping: $\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$0.01$])$, momentum:
$\mathcal{U}([$0$;$0.99$])$, exponential moving average:
$\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to identity:
$\mathcal{U}(\\{\text{no},\text{yes}\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

  * •

KFAC*: damping: $\mathcal{LU}([$1\text{$⋅$}{10}^{-15}$;$0.01$])$, exponential
moving average: $\mathcal{U}([$0$;$0.99$])$, initialize Kronecker factors to
identity: $\mathcal{U}(\\{\text{no},\text{yes}\\})$, $N_{\Omega}$:
$\mathcal{U}(\\{$100$,$101$,\text{\ldots},$10\,000$\\})$,
$N_{\partial\Omega}$: $\mathcal{U}(\\{$50$,$51$,\text{\ldots},$5\,000$\\})$,
batch sampling frequency: $\mathcal{U}(\\{$0$,$1$,\text{\ldots},$10\,000$\\})$

##  Appendix B Pseudo-Code: KFAC for the Poisson Equation

MLP $u_{{\bm{\theta}}}$ with parameters
${\bm{\theta}}_{0}=({\bm{\theta}}_{0}^{(1)},\dots,{\bm{\theta}}_{0}^{(L)})=(\operatorname{vec}{\bm{W}}_{0}^{(1)},\dots,\operatorname{vec}{\bm{W}}_{0}^{(L)})$,

interior data $\\{({\bm{x}}_{n},y_{n})\\}_{n=1}^{N_{\Omega}}$,

boundary data
$\\{({\bm{x}}^{\text{b}}_{n},y^{\text{b}}_{n})\\}_{n=1}^{N_{\partial\Omega}}$

exponential moving average $\beta$, momentum $\mu$, Damping $\lambda$, number
of steps $T$

0) Initialization

for $l=1,\dots,L$ do

${\bm{A}}_{\Omega}^{(l)},{\bm{B}}_{\Omega}^{(l)},{\bm{A}}_{\partial\Omega}^{(l)},{\bm{B}}_{\partial\Omega}^{(l)}\leftarrow{\bm{0}}\text{
or }{\bm{I}}$ $\triangleright$ Initialize Kronecker factors

end for

for $t=0,\dots,T-1$ do

1) Compute the interior loss and update its approximate curvature

$({\bm{Z}}_{n}^{(0)}\dots,{\bm{Z}}_{n}^{(L)},\Delta u_{n})\leftarrow\Delta
u_{{\bm{\theta}}_{t}}({\bm{x}}_{n})\quad n=1,\dots,N_{\Omega}$
$\triangleright$ Forward Laplacian wit intermediates

Compute layer output gradients
${\bm{g}}^{(l)}_{n,s}\coloneqq\nicefrac{{\partial\Delta
u_{n}}}{{\partial{\bm{Z}}_{n,s}^{(l)}}}$ with autodiff in one backward pass

$({\bm{g}}_{n,s}^{(1)},\dots,{\bm{g}}_{n,s}^{(L)})\leftarrow\texttt{grad}(\Delta
u_{n},({\bm{Z}}_{n,s}^{(1)},\dots,{\bm{Z}}_{n,s}^{(L)}))\quad
n=1,\dots,N_{\Omega}$, $s=1,\dots,S\coloneqq d+2$

for all $l=1,\dots,L$ do $\triangleright$ Update Kronecker factors of the
interior loss

$\hat{{\bm{A}}}_{\Omega}^{(l)}\leftarrow\beta\hat{{\bm{A}}}_{\Omega}^{(l)}+(1-\beta)\frac{1}{N_{\Omega}S}\sum_{n=1}^{N_{\Omega}}{\bm{Z}}_{n,s}^{(l-1)}{\bm{Z}}_{n,s}^{(l-1)\top}$

$\hat{{\bm{B}}}_{\Omega}^{(l)}\leftarrow\beta\hat{{\bm{B}}}_{\Omega}^{(l)}+(1-\beta)\frac{1}{N_{\Omega}}\sum_{n=1}^{N_{\Omega}}{\bm{g}}^{(l)}_{n,s}{\bm{g}}^{(l)\top}_{n,s}$

end for

$L_{\Omega}({\bm{\theta}}_{t})\leftarrow\frac{1}{2N_{\Omega}}\sum_{n=1}^{N_{\Omega}}(\Delta
u_{n}-y_{n})^{2}$ $\triangleright$ Compute interior loss

2) Compute the boundary loss and update its approximate curvature

$({\bm{z}}_{n}^{(0)}\dots,{\bm{z}}_{n}^{(L)},u_{n})\leftarrow
u_{{\bm{\theta}}_{t}}({\bm{x}}_{n}^{\text{b}})\quad
n=1,\dots,N_{\partial\Omega}$ $\triangleright$ Forward pass with intermediates

Compute layer output gradients ${\bm{g}}_{n}^{(l)}\coloneqq\nicefrac{{\partial
u_{n}}}{{{\bm{z}}^{(l)}_{n}}}$ with autodiff in one backward pass

$({\bm{g}}_{n}^{(1)}\dots,\partial{\bm{g}}_{n}^{(L)})\leftarrow\texttt{grad}(u_{n},({\bm{z}}_{n}^{(0)}\dots,{\bm{z}}_{n}^{(L)}))\quad
n=1,\dots,N_{\partial\Omega}$